Full DeepSeek-V3 architecture pretrained on a MacBook Air in 190 lines of JAX

A project shared on Hacker News by developer Vaibhaav Tiwari implements the DeepSeek-V3 architecture in 190 lines of pure JAX — no Flax, no Equinox, no PyTorch — and pretrains the model from scratch on a MacBook Air. According to the repository, parameters are plain dictionaries of arrays and every part of the model is an ordinary function in a single file, deepseek_v3.py, that can be read top to bottom. The line count covers non-blank, non-comment lines after the file's docstring.

What the 190 lines contain

The code brings in each idea from the DeepSeek-V3 technical report:

• Multi-head Latent Attention (MLA), which stores one small compressed latent per token instead of full keys and values, shrinking the cache from 480 to 144 numbers per token per layer at this scale.
• Decoupled rotary position embeddings, where a single rotary key is shared across all heads and kept outside the compressed latent.
• DeepSeekMoE routing: one shared expert that always participates plus 16 routed experts, with the top four chosen per token.
• Auxiliary-loss-free load balancing, in which a per-expert bias steers routing toward idle experts and affects only which experts are picked, not how their outputs are weighted.
• A small sequence-wise balance loss (α = 1e-4) complementing the bias.
• Multi-token prediction (MTP), an extra module that also forecasts the token after next, weighted at λ = 0.3 in the loss.
• Training and generation: bf16 compute over f32 master weights, with prefill and decode paths that use the MLA cache.

Training run and results

The default configuration trained on a MacBook Air M5 with 24 GB of memory. The model has 36.2 million parameters, 13.6 million of them active per token — a tiny fraction of DeepSeek-V3's 671 billion. The corpus was 207 public-domain detective novels totalling 21.7 million training tokens, with held-out crime fiction reserved for validation. Training ran 12,000 steps, about 98 million tokens (~4.5 epochs), in roughly 2.6 hours at a median 10,700 tokens per second. Validation loss fell from 5.02 at step 250 to 3.51.

The repository also reports an ablation of the bias-based balancing. With the bias enabled, each expert handled 5.7–6.9% of tokens at step 2000, close to the perfectly even 6.25%; without it, loads ranged from 0.09% to 18.6%. Validation loss at that step was effectively identical either way (3.921 versus 3.923), which the author points out is exactly the paper's claim: the bias evens out expert use at no cost.

Sampled text is grammatical but generic, and the author is explicit that this is a demonstration of the architecture rather than a capable language model.

What was simplified

The project makes deliberate trade-offs for readability. Every expert computes on every token and non-selected outputs are zeroed, so the routing math is exact but there is no sparse dispatch or expert parallelism. MLA skips weight absorption and decompresses the latent explicitly, which is clearer but slower at inference. Entirely absent are FP8 training, the DualPipe schedule, node-limited routing, YaRN context extension and any post-training. Scale is heavily reduced: 6 layers against the original 61 and a 256-token context window.

Reproducing the run

Reproduction requires an Apple Silicon Mac, Python 3.12 and the uv package manager, with every dependency pinned — jax-metal 0.1.1, for example, only works with jax 0.4.34. A data-preparation script downloads the books and trains an 8,192-token BPE tokenizer; because Project Gutenberg has re-issued 28 of its 203 books in the corpus with small edits, freshly built datasets hash differently from the published ones, though token counts stay the same and the exact published bytes are available separately. A 39-test suite covers masking, cache consistency, routing, MTP targets and numerical handling, and a smaller CPU-only configuration exists for 8–16 GB machines. Runs on Metal are not bit-for-bit deterministic, so reproduced numbers land near rather than exactly on the published ones. One caution from the repository: the checkpoint is a Python pickle, which can execute code when loaded.

Why it matters

Frontier architectures usually reach developers only through large, framework-heavy codebases where the interesting ideas sit behind layers of infrastructure. A single readable file that contains all of DeepSeek-V3's mechanisms — and actually trains end to end on a laptop — turns a 671-billion-parameter design into something a person can study, modify and verify in an afternoon. The balancing ablation is a small, concrete check of the paper's central routing claim, and the whole exercise shows how conceptually compact modern LLM engineering becomes once distributed-training machinery is stripped away. For anyone learning JAX or mixture-of-experts design, it doubles as a working reference implementation.