> 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.
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 its awesome to have a stepping off point now!
No CUDA meant ROCm suprisingly just worked with JAX. Given recent events at NVDA its also very timely :)
Gotchas for anyone with a 7600 / 7600 XT:
1. ROCm 10.1 supports gfx1102, but JAX's ROCm plugin rejects it ("no supported devices found"), so we need to set HSA_OVERRIDE_GFX_VERSION=11.0.0
2. XLA's Triton GEMM fusions exceed Navi 33's dispatch limits at the default config and hangs the GPU, so set XLA_FLAGS=--xla_gpu_enable_triton_gemm=false
3. JAX ≥0.5 changed the default PRNG, so set JAX_THREEFRY_PARTITIONABLE=false if you are looking to replicate precisely
My setup: 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
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
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 :)
Respect and thanks.
The whole reason for me buying the card was to play around with deep learning and models, so its awesome to have a stepping off point now!
No CUDA meant ROCm suprisingly just worked with JAX. Given recent events at NVDA its also very timely :)
Gotchas for anyone with a 7600 / 7600 XT:
1. ROCm 10.1 supports gfx1102, but JAX's ROCm plugin rejects it ("no supported devices found"), so we need to set HSA_OVERRIDE_GFX_VERSION=11.0.0
2. XLA's Triton GEMM fusions exceed Navi 33's dispatch limits at the default config and hangs the GPU, so set XLA_FLAGS=--xla_gpu_enable_triton_gemm=false
3. JAX ≥0.5 changed the default PRNG, so set JAX_THREEFRY_PARTITIONABLE=false if you are looking to replicate precisely
My setup: 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
PR with the container setup, a comparison script, and plot: https://github.com/Vaibhaav-Tiwari/deepseek-v3-jax/pull/1