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.

> 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

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!

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 :)

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