Developer Trains DeepSeek V3 Architecture in 190 Lines of Pure JAX on MacBook Air

The repository, released as v1.0, implements DeepSeek-V3 ideas in deepseek_v3.py without Flax, Equinox, or PyTorch. It covers Multi-head Latent Attention, Decoupled RoPE, DeepSeekMoE with one shared and 16 routed experts, auxiliary-loss-free balancing, sequence-wise balance loss, and multi-token prediction. The code stores parameters as plain dictionaries of arrays and trains with bf16 compute on f32 master weights. The author describes it as an architecture demo, not a strong language model.

The published run used 207 public-domain detective-fiction books, 21.7 million training tokens, and a held-out validation set of crime fiction. It trained a 36.2 million-parameter model, with 13.6 million active parameters per token, for 12,000 steps over about 98 million tokens, roughly 4.5 epochs, in about 2.6 hours on an Apple Silicon Mac with 24 GB memory. Median throughput was about 10,700 tokens per second. Validation loss fell from 5.02 at step 250 to 3.51. With load balancing enabled, expert usage was 5.7% to 6.9%, versus 0.09% to 18.6% without it; at step 2000, validation loss was 3.921 versus 3.923.

The implementation compresses several DeepSeek-V3 mechanisms into a small file. MLA caches a latent vector per token, reducing cached values from 480 to 144 numbers per token per layer in this setup. The MoE layer always runs one shared expert and routes top-4 of 16 experts. Multi-token prediction adds a module that also predicts the token after next. The auxiliary-loss-free bias changes routing choices rather than expert weights. The repository notes that jax-metal 0.1.1 is pinned and that a known Metal gradient path is avoided by the model code.

The result matters because it gives developers a readable, reproducible path into DeepSeek-V3 ideas without a large training stack. It shows that a small language model using these mechanisms can be trained and sampled on consumer hardware, while the checkpoint and tokenizer are released for inspection. The author states the output is not a strong language model, and the release warns that checkpoints are Python pickles that can run code, so users should load only trusted files.