I am just starting Raschka's book to write an llm from scratch in C++ with zero dependencies with avx-512 on a 9950x, and then I'd like to go through the literature for more advanced stuff.
I don't know what Jax is exactly but while I am aware that the algorithm implementation is not something terribly complicated, 200 lines still sounds impressive.
If anyone wants to share any tips, they're welcome.
you can think of JAX as numpy with automatic gradients and a compiler so no backward pass or going into kernels was needed.
also, Raschka's book is excellent! you should also check out his model comparison primer that's awesome + you can go on to https://jax-ml.github.io/scaling-book/ next
I love this, there is nothing like recreating a model from first principles, plus makes me want to revisit Jax again (I'll do it after I spend enough time with Zig). Thanks for this!
Thanks for this! I've been trying to get local training working on Fedora with ROCm for a while now, and this is the first project that ran end-to-end on my RX 7600 XT 16GB.
The whole reason for me buying the card was to play around with deep learning and models, so it's awesome to have a stepping off point now!
No CUDA meant ROCm surprisingly just worked with JAX.
Gotchas for anyone with a 7600 XT: it needed two env vars, HSA_OVERRIDE_GFX_VERSION=11.0.0 because JAX's plugin rejects gfx1102, and XLA_FLAGS=--xla_gpu_enable_triton_gemm=false to stop a GPU hang.
I'm on Fedora 45 with the in-kernel amdgpu driver, rootless podman and AMD's rocm/jax image (ROCm 10.1, JAX 0.11.1). Nothing ROCm is installed on the host, which is nice.
> Gotchas for anyone with a 7600 XT: it needed two env vars, HSA_OVERRIDE_GFX_VERSION=11.0.0 because JAX's plugin rejects gfx1102, and XLA_FLAGS=--xla_gpu_enable_triton_gemm=false to stop a GPU hang.
It's sad that after so many years AMD's software support is still atrocious for non-datacenter hardware, and their short-sighted management still doesn't see it as a priority to fix. This is costing them sales, but I guess they only care about selling to hyperscalers.
Post like these always drives a motivation to understand how inference or what OP did.
For folks who are more familiar in this space, is there a good road map how to build and understand the intuition behind this? I’m comfortable in Java and Python and do vanilla SWE work.
Intuition: there are YouTube videos from 3blue1brown [1] and Welch Labs [2] on neural networks, LLMs, and related tech that are good.
Building: the videos/tutorials from Andrej Karpathy in Python [3] are considered to be very good at building and understanding the basics, from a small/toy neural network engine to transformers and tokenizers.
There was a post here a few weeks ago [4] on reasoning models using Python/PyTorch with links to an older series on non-reasoning transformers.
> to asking a language model to create a 3D end-to-end working model for me to capture the nuances
Love this! Yes, having a visual of the system makes it a lot easier to reason about and understand
I wish all AI models came with interactive animated workflows of their architecture
> so I wanted it to be different just for the sake of it and used various books by…
That’s cool and great for your own project. At the same time it helps to use datasets that others have used, because then you can compare results more directly
Excellent write up, I really appreciate all the detail and the sample outputs at the end. Excited to see what I could apply for this to the small <10M models I’ve been training
Surely that can be tried! of course at a small scale if you're short on GPUs. what I understand is that V4 replaces MLA with a hybrid of two compressed attention types and changes the residual connections, while keeping the mixture-of-experts design, but I haven't gone in-depth architecture :)
I dont think one's available as of now but you can check out llama.cpp, though they run them for real weights rather than small build-from-scratch versions like this!
I wanted to understand what actually makes DeepSeek V3 different, so I rebuilt its three main ideas from scratch in pure JAX: Multi-head Latent Attention, mixture-of-experts with bias-based load balancing instead of an auxiliary loss, and multi-token prediction. The model is about 190 lines.
It’s a 36.2M-parameter version trained on a MacBook Air M5 (24 GB) with jax-metal, on 207 old detective novels. 12,000 steps took about 2.6 hours.
I am just starting Raschka's book to write an llm from scratch in C++ with zero dependencies with avx-512 on a 9950x, and then I'd like to go through the literature for more advanced stuff.
I don't know what Jax is exactly but while I am aware that the algorithm implementation is not something terribly complicated, 200 lines still sounds impressive.
If anyone wants to share any tips, they're welcome.
you can think of JAX as numpy with automatic gradients and a compiler so no backward pass or going into kernels was needed. also, Raschka's book is excellent! you should also check out his model comparison primer that's awesome + you can go on to https://jax-ml.github.io/scaling-book/ next
Thanks, I will definitely check these out!
I love this, there is nothing like recreating a model from first principles, plus makes me want to revisit Jax again (I'll do it after I spend enough time with Zig). Thanks for this!
Thanks for this! I've been trying to get local training working on Fedora with ROCm for a while now, and this is the first project that ran end-to-end on my RX 7600 XT 16GB.
The whole reason for me buying the card was to play around with deep learning and models, so it's awesome to have a stepping off point now!
No CUDA meant ROCm surprisingly just worked with JAX.
Gotchas for anyone with a 7600 XT: it needed two env vars, HSA_OVERRIDE_GFX_VERSION=11.0.0 because JAX's plugin rejects gfx1102, and XLA_FLAGS=--xla_gpu_enable_triton_gemm=false to stop a GPU hang.
I'm on Fedora 45 with the in-kernel amdgpu driver, rootless podman and AMD's rocm/jax image (ROCm 10.1, JAX 0.11.1). Nothing ROCm is installed on the host, which is nice.
Container setup is in a PR if it helps anyone: https://github.com/Vaibhaav-Tiwari/deepseek-v3-jax/pull/1
Thanks Joe, this is great. Reviewed and tested what I could without an AMD card, and merged the PR!
> Gotchas for anyone with a 7600 XT: it needed two env vars, HSA_OVERRIDE_GFX_VERSION=11.0.0 because JAX's plugin rejects gfx1102, and XLA_FLAGS=--xla_gpu_enable_triton_gemm=false to stop a GPU hang.
It's sad that after so many years AMD's software support is still atrocious for non-datacenter hardware, and their short-sighted management still doesn't see it as a priority to fix. This is costing them sales, but I guess they only care about selling to hyperscalers.
Post like these always drives a motivation to understand how inference or what OP did.
For folks who are more familiar in this space, is there a good road map how to build and understand the intuition behind this? I’m comfortable in Java and Python and do vanilla SWE work.
Intuition: there are YouTube videos from 3blue1brown [1] and Welch Labs [2] on neural networks, LLMs, and related tech that are good.
Building: the videos/tutorials from Andrej Karpathy in Python [3] are considered to be very good at building and understanding the basics, from a small/toy neural network engine to transformers and tokenizers.
There was a post here a few weeks ago [4] on reasoning models using Python/PyTorch with links to an older series on non-reasoning transformers.
[1] https://www.youtube.com/playlist?list=PLZHQObOWTQDNU6R1_6700...
[2] https://www.youtube.com/@WelchLabs/videos
[3] https://karpathy.ai/zero-to-hero.html
[4] https://news.ycombinator.com/item?id=50245138
I'm not an expert but there's https://jaxformer.com, https://github.com/karpathy/nanochat and, of course, the original "Attention is all you need" paper at https://arxiv.org/abs/1706.03762.
you should definitely start with Stanford CS336, even im re-doing it
> to asking a language model to create a 3D end-to-end working model for me to capture the nuances
Love this! Yes, having a visual of the system makes it a lot easier to reason about and understand
I wish all AI models came with interactive animated workflows of their architecture
> so I wanted it to be different just for the sake of it and used various books by…
That’s cool and great for your own project. At the same time it helps to use datasets that others have used, because then you can compare results more directly
Excellent write up, I really appreciate all the detail and the sample outputs at the end. Excited to see what I could apply for this to the small <10M models I’ve been training
that's true, I'll try to incorporate standard datasets when the goal would be comparison. the detective dataset was just for fun :)
Could you build and train a V4.1 the same way?
Surely that can be tried! of course at a small scale if you're short on GPUs. what I understand is that V4 replaces MLA with a hybrid of two compressed attention types and changes the residual connections, while keeping the mixture-of-experts design, but I haven't gone in-depth architecture :)
anyone got a c++ version of this?
I dont think one's available as of now but you can check out llama.cpp, though they run them for real weights rather than small build-from-scratch versions like this!
thank you so so much for writing this, I will recreate it for learning
thats great :)
I wanted to understand what actually makes DeepSeek V3 different, so I rebuilt its three main ideas from scratch in pure JAX: Multi-head Latent Attention, mixture-of-experts with bias-based load balancing instead of an auxiliary loss, and multi-token prediction. The model is about 190 lines.
It’s a 36.2M-parameter version trained on a MacBook Air M5 (24 GB) with jax-metal, on 207 old detective novels. 12,000 steps took about 2.6 hours.
Full write-up with diagrams and results: https://vaibhaav.in/writing/deepseek-v3-jax/
The trained checkpoint and exact data are in the v1.0 release, so you can generate text without training.
Would love corrections if I got anything wrong :)
This is very illustrative. I'm sure a great learning experience for you but so much more for so many others.
Respect and thanks.