{"slug": "show-hn-deepseek-v3-in-200-lines-of-pure-jax", "title": "Show HN: DeepSeek V3 in ~200 lines of pure Jax", "summary": "Vaibhaav Tiwari released a 190-line pure JAX implementation of the DeepSeek-V3 architecture, pretrained from scratch on a MacBook Air M5 with 24 GB of memory on public-domain detective fiction. The 36.2M-parameter model (13.6M active per token) trained for 12,000 steps over roughly 98M tokens in about 2.6 hours at ~10.7k tokens/s, cutting validation loss from 5.02 at step 250 to 3.51. The v1.0 release includes the exact 138 MB checkpoint and tokenizer from the published run, and the author reports that DeepSeek's auxiliary-loss-free bias balancing held per-expert load to 5.7–6.9% versus 0.09–18.6% without it, with validation loss essentially unchanged at step 2000 (3.921 vs 3.923).", "body_md": "The architecture from the [DeepSeek-V3 technical report](https://arxiv.org/abs/2412.19437) in **190 lines of pure JAX**, pretrained from scratch on a MacBook Air on public-domain detective fiction.\n\nNo Flax, no Equinox, no PyTorch. Params are plain dicts of arrays and every part of the model is a function you can read top to bottom in [`deepseek_v3.py`](https://github.com/Vaibhaav-Tiwari/deepseek-v3-jax/blob/main/deepseek_v3.py).\n\n**Read the write-up:** [DeepSeek V3 in 190 lines of pure JAX](https://vaibhaav.in/writing/deepseek-v3-jax/), with diagrams, the code explained piece by piece, and the results.\n\n| Idea from the paper | What it does | Where | \n|---|---|---|\n| Multi-head Latent Attention (MLA) | Caches one small latent per token instead of full keys and values: **480 → 144** numbers per token per layer here | `mla` ,`init_cache` | \n| Decoupled RoPE | One rotary key shared by all heads, kept outside the compressed latent | `rope` ,`apply_rope` ,`mla` | \n| DeepSeekMoE | 1 shared expert that always runs + 16 routed experts, top-4 per token | `route` ,`moe` | \n| Auxiliary-loss-free balancing | A per-expert bias nudges routing toward idle experts; it changes *which* experts are picked, not their weights | `update_bias` ,`route` | \n| Sequence-wise balance loss | The small complementary loss from the paper (α = 1e-4) | `seq_balance_loss` | \n| Multi-token prediction (MTP) | An extra module also predicts the token after next, weighted λ = 0.3 in the loss | `forward` ,`loss_fn` | \n| Training and generation | bf16 compute on f32 master weights; prefill + decode with the MLA cache | `make_train_step` ,`generate` | \n\nOne run of the default config on a MacBook Air M5 (24 GB):\n\n| Model | 36.2M parameters, 13.6M active per token | \n| Data | 207 books, 21.7M training tokens (Hammett repeated 3×); validation is held-out crime fiction | \n| Training | 12,000 steps, ~98M tokens (~4.5 epochs), ~2.6 h | \n| Speed | ~10.7k tokens/s (median) | \n| Validation loss | 5.02 (step 250) → **3.51** | \n| Expert load, with bias | 5.7–6.9% per expert (even is 6.25%), at step 2000 | \n| Expert load, without bias | 0.09–18.6% per expert, at step 2000 | \n\nWith balancing on and off, validation loss at step 2000 was the same (3.921 vs 3.923). The bias evens out the experts without costing anything, which is the paper's claim.\n\nIts own words, from the prompt \"The fog came in off the\":\n\nThe fog came in off the bay, and the great trees were spreading and the stars were shining.\n\nThis is an architecture demo, not a strong language model. See the write-up for what worked, what didn't, and what's missing.\n\nThe [v1.0 release](https://github.com/Vaibhaav-Tiwari/deepseek-v3-jax/releases/tag/v1.0) has the exact checkpoint and data from the published run. After [installing](#reproduce-it) (step 1):\n\n```\nmkdir -p data runs/main\ncurl -L -o runs/main/ckpt.pkl      https://github.com/Vaibhaav-Tiwari/deepseek-v3-jax/releases/download/v1.0/ckpt.pkl          # 138 MB\ncurl -L -o data/tokenizer.json     https://github.com/Vaibhaav-Tiwari/deepseek-v3-jax/releases/download/v1.0/tokenizer.json\n.venv/bin/python sample.py --prompt \"The fog came in off the\" --num 3\n```\n\nThe checkpoint is a Python pickle, which can run code when loaded. Only load checkpoints from sources you trust.\n\nYou need an Apple Silicon Mac, Python 3.12 and [uv](https://docs.astral.sh/uv/). The published run used 24 GB of memory. No Mac? See [Running on CPU](#running-on-cpu).\n\n**1. Install**\n\n```\ngit clone https://github.com/Vaibhaav-Tiwari/deepseek-v3-jax.git && cd deepseek-v3-jax\nuv venv -p 3.12 .venv\nuv pip install --python .venv/bin/python -r requirements.txt\n```\n\nEvery version in `requirements.txt` is pinned. `jax-metal 0.1.1` only works with `jax==0.4.34`, so don't upgrade JAX.\n\n**2. Check your setup**\n\n```\nJAX_PLATFORMS=cpu .venv/bin/pytest tests.py -q    # 39 passed\n.venv/bin/python smoke_metal.py                   # optional: what jax-metal supports on your Mac\n```\n\nIn `smoke_metal.py`, \"grad through 2-axis contraction\" fails on jax-metal 0.1.1. That's expected, and the model is written to avoid it.\n\n**3. Build the dataset**\n\n```\n.venv/bin/python prepare_data.py\n```\n\nThis downloads the 207 books in [`sources.txt`](https://github.com/Vaibhaav-Tiwari/deepseek-v3-jax/blob/main/sources.txt) (~120 MB, about 7 minutes), cleans them, trains an 8,192-token BPE tokenizer and writes `data/train.bin`, `data/val.bin` and `data/tokenizer.json`. It should print:\n\n```\ntrain: 211 docs, 83.4M chars, 21.69M tokens\nval: 4 docs, 1.1M chars, 0.30M tokens\n```\n\nThe published run's data had these hashes:\n\n```\nshasum -a 256 data/train.bin data/val.bin data/tokenizer.json\nd37aca97e707bb81b6e4d266cc1adfd1f3e3556457c6eb724e495118a734e0bb  data/train.bin\nb7ac285b115e6793b5d36aa5b4bc6e81a5fccfc7af0908a306d3828ebca9c31c  data/val.bin\nc5a142b4006a967c880dda44ad792719d53c16d947da3b6c2332d6132d9967cc  data/tokenizer.json\n```\n\nYou will most likely get different hashes. The Hammett books are pinned to exact Standard Ebooks commits and come out identical, but Project Gutenberg files can't be pinned, and Gutenberg has since re-issued 28 of the 203 books with small edits (mostly accents stripped, like \"José\" → \"Jose\"). The token counts come out the same, so training should behave the same.\n\n**To train on the exact same bytes**, skip `prepare_data.py` and download the published data instead:\n\n```\nmkdir -p data\nfor f in train.bin val.bin tokenizer.json; do curl -L -o data/$f https://github.com/Vaibhaav-Tiwari/deepseek-v3-jax/releases/download/v1.0/$f; done\nshasum -a 256 data/*.bin data/tokenizer.json    # should match the hashes above\n```\n\n**4. Train**\n\n```\n.venv/bin/python train.py                                         # the main run: 12k steps, ~2.6 h\n.venv/bin/python train.py --steps 2000 --out runs/bias_on         # balancing ablation, with the bias\n.venv/bin/python train.py --steps 2000 --out runs/bias_off --no-bias   # ...and without it\n```\n\nEach run writes `metrics.jsonl` (loss, validation loss, per-expert load, biases, tokens/s) and `ckpt.pkl` (saved at every evaluation) into its `--out` folder, which defaults to `runs/main`. Seeds are fixed, but GPU math on Metal isn't bit-for-bit deterministic, so expect your numbers to land close to the ones above rather than match them exactly.\n\n**5. Generate**\n\n```\n.venv/bin/python sample.py --prompt \"The fog came in off the\" --num 3\n```\n\nDefaults: temperature 0.8, top-k 50, 300 new tokens per sample (`--temp`, `--top-k`, `--n`).\n\n**6. Make the figures**\n\n```\n.venv/bin/python plots.py                    # needs runs/main, runs/bias_on, runs/bias_off\nJAX_PLATFORMS=cpu .venv/bin/python charts.py # the rest of the write-up's charts\n```\n\nBoth write into `post/figures/`.\n\nTo redraw the figures without training, download the logs from the release: `metrics-main.jsonl`, `metrics-bias_on.jsonl` and `metrics-bias_off.jsonl` go into `runs/main/`, `runs/bias_on/` and `runs/bias_off/` as `metrics.jsonl`. `charts.py` also needs the checkpoint and the data.\n\nEverything runs without jax-metal. It's much slower, so use the smaller config:\n\n```\nJAX_PLATFORMS=cpu .venv/bin/python train.py --config small\n```\n\n`--config small` (4 layers, 8 experts, top-2, 128-token context) is meant for 8–16 GB Macs too. It has only been checked for a couple of steps, not trained to the end.\n\n| File | What it is | \n|---|---|\n| [`deepseek_v3.py`](https://github.com/Vaibhaav-Tiwari/deepseek-v3-jax/blob/main/deepseek_v3.py) | **The model** : config, MLA, MoE, balancing, MTP, loss, train step, cached generation. The 190 lines. | \n| [`train.py`](https://github.com/Vaibhaav-Tiwari/deepseek-v3-jax/blob/main/train.py) | Training loop: AdamW with warmup and cosine decay, evaluation, checkpoints, logging | \n| [`prepare_data.py`](https://github.com/Vaibhaav-Tiwari/deepseek-v3-jax/blob/main/prepare_data.py) ,[`sources.txt`](https://github.com/Vaibhaav-Tiwari/deepseek-v3-jax/blob/main/sources.txt) | Download, clean and tokenize the corpus | \n| [`sample.py`](https://github.com/Vaibhaav-Tiwari/deepseek-v3-jax/blob/main/sample.py) | Text generation from a checkpoint | \n| [`tests.py`](https://github.com/Vaibhaav-Tiwari/deepseek-v3-jax/blob/main/tests.py) | 39 tests: causal masking, cache vs full forward, routing, balancing, MTP targets, NaN handling, bf16 master weights, and more | \n| [`smoke_metal.py`](https://github.com/Vaibhaav-Tiwari/deepseek-v3-jax/blob/main/smoke_metal.py) | Which JAX operations work on Apple's GPU | \n| [`plots.py`](https://github.com/Vaibhaav-Tiwari/deepseek-v3-jax/blob/main/plots.py) ,[` charts.py`](https://github.com/Vaibhaav-Tiwari/deepseek-v3-jax/blob/main/charts.py) | The figures, all drawn from the logs, checkpoint and corpus | \n| [`post/`](https://github.com/Vaibhaav-Tiwari/deepseek-v3-jax/blob/main/post) | Diagrams, figures and screenshots from the write-up | \n\n\"190 lines\" counts the non-blank, non-comment lines of `deepseek_v3.py` after its docstring:\n\n```\ntail -n +5 deepseek_v3.py | grep -cvE '^\\s*(#|$)'    # 190\n```\n\n- **Dense MoE compute.** Every expert runs on every token and the unpicked ones are zeroed. The routing math is exact, but there's no sparse dispatch or expert parallelism.\n- **No weight absorption in MLA.** The latent is decompressed explicitly, which is easier to read and slower at inference.\n- **Not included:** FP8 training, DualPipe, node-limited routing, YaRN context extension, and any post-training.\n- **Scale.** 36.2M parameters vs 671B, 6 layers vs 61, a 256-token context vs 128K.\n\nAll training text is public domain in the United States: Dashiell Hammett's *Red Harvest*, *The Dain Curse*, *The Maltese Falcon* and Continental Op stories (CC0 editions from [Standard Ebooks](https://standardebooks.org)), plus about 200 detective and crime novels from [Project Gutenberg](https://www.gutenberg.org). Every source is listed in [`sources.txt`](https://github.com/Vaibhaav-Tiwari/deepseek-v3-jax/blob/main/sources.txt). Public-domain status outside the US may differ.\n\nDeepSeek-AI, [DeepSeek-V3 Technical Report](https://arxiv.org/abs/2412.19437), 2024.\n\nCode: [MIT](https://github.com/Vaibhaav-Tiwari/deepseek-v3-jax/blob/main/LICENSE). The training text is public domain in the US (see above).", "url": "https://wpnews.pro/news/show-hn-deepseek-v3-in-200-lines-of-pure-jax", "canonical_source": "https://github.com/Vaibhaav-Tiwari/deepseek-v3-jax", "published_at": "2026-11-14 15:04:15+00:00", "updated_at": "2026-11-14 15:23:11.742285+00:00", "lang": "en", "topics": ["large-language-models", "machine-learning", "ai-research", "developer-tools"], "entities": ["DeepSeek", "Vaibhaav Tiwari", "JAX", "DeepSeek-V3", "MacBook Air M5", "GitHub"], "also_reported_by": [], "alternates": {"html": "https://wpnews.pro/news/show-hn-deepseek-v3-in-200-lines-of-pure-jax", "markdown": "https://wpnews.pro/news/show-hn-deepseek-v3-in-200-lines-of-pure-jax.md", "text": "https://wpnews.pro/news/show-hn-deepseek-v3-in-200-lines-of-pure-jax.txt", "jsonld": "https://wpnews.pro/news/show-hn-deepseek-v3-in-200-lines-of-pure-jax.jsonld"}}