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!
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
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.
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 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!
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.
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.
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
> 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
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
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.