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.<p>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 <a href="https://jaxformer.com" rel="nofollow">https://jaxformer.com</a>, <a href="https://github.com/karpathy/nanochat" rel="nofollow">https://github.com/karpathy/nanochat</a> and, of course, the original "Attention is all you need" paper at <a href="https://arxiv.org/abs/1706.03762" rel="nofollow">https://arxiv.org/abs/1706.03762</a>.
Intuition: there are YouTube videos from 3blue1brown [1] and Welch Labs [2] on neural networks, LLMs, and related tech that are good.<p>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.<p>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.<p>[1] <a href="https://www.youtube.com/playlist?list=PLZHQObOWTQDNU6R1_67000Dx_ZCJB-3pi" rel="nofollow">https://www.youtube.com/playlist?list=PLZHQObOWTQDNU6R1_6700...</a><p>[2] <a href="https://www.youtube.com/@WelchLabs/videos" rel="nofollow">https://www.youtube.com/@WelchLabs/videos</a><p>[3] <a href="https://karpathy.ai/zero-to-hero.html" rel="nofollow">https://karpathy.ai/zero-to-hero.html</a><p>[4] <a href="https://news.ycombinator.com/item?id=50245138">https://news.ycombinator.com/item?id=50245138</a>
> to asking a language model to create a 3D end-to-end working model for me to capture the nuances<p>Love this! Yes, having a visual of the system makes it a lot easier to reason about and understand<p>I wish all AI models came with interactive animated workflows of their architecture<p>> so I wanted it to be different just for the sake of it and used various books by…<p>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<p>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.<p>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!<p>No CUDA meant ROCm surprisingly just worked with JAX.<p>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.<p>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.<p>Container setup is in a PR if it helps anyone: <a href="https://github.com/Vaibhaav-Tiwari/deepseek-v3-jax/pull/1" rel="nofollow">https://github.com/Vaibhaav-Tiwari/deepseek-v3-jax/pull/1</a>
[flagged]
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.<p>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.<p>Full write-up with diagrams and results: <a href="https://vaibhaav.in/writing/deepseek-v3-jax/" rel="nofollow">https://vaibhaav.in/writing/deepseek-v3-jax/</a><p>The trained checkpoint and exact data are in the v1.0 release, so you can generate text without training.<p>Would love corrections if I got anything wrong :)