Accelerators are getting faster but Python remains slow. The investment into a true JIT for Python (CPython, the only one folks use for ML) seems to be far away.

Projects like JAX, Triton and Taichi are awesome but remain somewhat brittle to use. We believe this is because there isn't a true JIT in Python (like the one found in JavaScript) that these projects can leverage. On top of full endorsement from the language's community, building a true JIT takes millions of lines of heavily interconnected code and likely hundreds of thousands if not millions of engineer hours.[1][2] There is very little to indicate this extremely expensive type of investment will come from CPython.

As a hedge against CPython never becoming fast, we're creating a project called Shumai that attempts to deeply integrate with a new JavaScript runtime (Bun[3]).

The Bun JS runtime (and other JS runtimes) have had true JIT compilation for over a decade and are fiercely competitive with each other. This has lead to a prosperous and efficiency-driven language that rapidly adopts new features and is very open to improvement. JavaScript has dropped semicolons, added `import` statements (learning from Python), made asynchronous programming native (this is very important for scaling models, but may not be relevant for small training), added a much faster FFI than Python (for C extensions, 10-100x less overhead), and even has things like importing directly from a URL (no need for `pip install`).

This type of ecosystem, in our minds, is a very natural fit for research.

It's also markedly faster than Python for anything that isn't large matrix multiplications that saturate the GPU[4]. The work being done in kernel fusion (e.g. AITemplate, Triton, XLA) is still very relevant and necessary for the best performance. However, *other* work, such as doing complicated bookkeeping, simulating games for RL, running very dynamic networks, interleaving complex network code with computation, etc. would all benefit from a faster host language.

There are other more obvious benefits to using JS for ML (which TF.js has done a great job of capturing), but theses are not the current focus of the Shumai project. Instead, we are focusing on single node performance and multi-node scaling (with JS-style functional + reactive programming).

If you're interested in checking out the current state of the code, here it is: https://github.com/facebookresearch/shumai

~~~

[1] V8, a JIT engine and nothing else: 2.3 million lines of code https://www.openhub.net/p/v8-js/analyses/latest/languages_summary

[2] WebKit: 26 million lines of code https://www.openhub.net/p/WebKit/analyses/latest/languages_summary (note that this includes a browser not just a JIT)

[3] Bun: https://github.com/oven-sh/bun

[4] Bun's JS is 41x faster than the newest Python, ~120x faster than 3.8: https://jott.live/markdown/py3.11_vs_3.8