I don’t think this is straightforward but it may be a skill issue on my part. It would require dockerizing headless Chrome with WebGPU support and dynamically injecting custom bundled JavaScript into the page, then extracting the results with Chrome IPC
Yes, we are actively working on it! The goal is to be a full ML research library, not just a model inference runtime. You can join the Discord to follow along
Hm okay, seems like an interesting set of benchmarks — let me know if there’s anything I can do to help make jax-js more compatible with your docker setup
Since ONNX is just a model data format, you can actually parse and run ONNX files in jax-js as well. Here’s an example of running DETR ResNet-50 from Xenova’s transformers.js checkpoint in jax-js
I don’t think I intend to support everything in ONNX right now, especially quant/dequant, but eventually it would be interesting to see if we can help accelerate transformers.js with a jax-js backend + goodies like kernel fusion
jax-js is more trying to explore being an ML research library, rather than ONNX which is a runtime for exported models
I don’t think tf.tidy() is a sound API under jvp/grad transformations, also it prevents you from using async which makes it incompatible with GPU backends (or blocks the page), a pretty big issue. https://github.com/tensorflow/tfjs/issues/5468
Thanks for the feedback though, just explaining how we arrived at this API. I hope you’d at least try it out — hopefully you will see when developing that the refs are more flexible than alternatives.
For sure! It looks like this is benchmarking the autodiff cpu time, not the actual kernels though, which (correct me if I’m wrong) isn’t really relevant for an ML library — it’s more for if you have a really complex scientific expression
Wanted to share this library that I'm open-sourcing now after working on it for the past year.
You write code like in JAX/NumPy, but it’s fully interactive on the frontend and compiles down to shaders on the user’s GPU (with WebGPU). So far I’ve used it for purely frontend-only ML demos. https://jax-js.com/mobileclip
Thanks! I think that's a great idea. Containers are nice, but there's nothing quite like a "real" VM system in terms of learning + developing, and they can be very cheap if you're OK with the performance hit from nested virtualization :)
Another idea I had is to customize the rootfs. It boots Alpine right now, but should be trivial to change it to anything else (and it only takes ~5 seconds to build!).
Hello! Thrilled to be sharing something I've worked on for the past 8 months at Modal — a cloud, GPU notebooks product.
We wanted to improve the Jupyter Notebook experience while offering high-performance, cloud compute. So we sought out to fix some of the problems in cloud offerings for notebooks, and we ended up rebuilding a lot of the interface ourselves. This means you can do things like:
- Within the same notebook, swap from CPU to GPU, and the kernel will restart in seconds while keeping your work.
- Use _lots_ of hardware: up to 8x H100/B200 GPUs, not just one.
- Automatic idle shutdown (configurable timeout) so you don't get a huge bill / forget to stop your notebook.
- Real-time collaborative editing.
- Integrations with Pyright and Ruff for a better IDE-like experience.
Unlike Google Colab, we are a full cloud offering, so you can integrate with the rest of the Modal platform — use custom images, attach persistent volumes and so on. We've heard these are the actual, boring but really important features you need to do experimental work in notebooks.
Happy to hear any feedback! Would love for you to try it out.
Hello - I wanted to share a bit about an open-source project I’ve been working on, building a fundamental library for ML and numerical computing in the browser.
There are different tradeoffs here, but I’d like to replicate the speed and ergonomics of Python libraries with NumPy/PyTorch/JAX without the difficulty of distribution. jax-js is lightweight enough to be embedded in a website, but optimized to take advantage of tech like Wasm and WebGPU.
Let me know if you have any feedback on this compiler design. It’s working well and faster than tfjs at matrix multiplication, while having kernel fusion as well.
Just want to mention that these standalone python builds have been super important for simplifying a lot of dev tooling, even if users may not see it directly — I work for a cloud infrastructure company, and they’re what allow us to give users a one-line way of adding Python to any Docker image they desire. It’s helpful to have reproducible, standalone Python builds of this quality (and release transparency). Thanks y’all for taking this on.
Sorry, just realized the misunderstanding. To clarify Modal still uses the kernel WireGuard module. The userspace part that’s in Go and not in other languages that we use is the wgctrl library.
Right, vprox servers act as multiplexed NAT instances with a VPN attached. You do still need the VPN part though since our containers run around the world, in multiple regions and availability zones. Setting the gateway to a machine running fck-nat would only work if that machine is in the same subnet (e.g., for AWS, in one availability zone).
The other features that were hard requirements for us were multi-tenancy and high availability / failover.
By the way, fck-nat is just a basic shell script that sets the `ip_forward` and `rp_filter` sysctls and adds an IP masquerade rule. If you look at vprox, we also do this but build a lot on top of it. https://github.com/modal-labs/vprox
Yep! This is something we have internal tests for haha, you have good instincts that it can be tricky. Here's an example of using that for multi-GPU training https://modal.com/docs/examples/llm-finetuning
Thanks. We did check out Tailscale, but they didn't quite have what we were looking for: some high-availability custom component that plugs into a low-level container runtime. (Which makes sense, it's pretty different from their intended use case.)
Modal is actually a happy customer of Tailscale (but for other purposes). :D