DeepSeek V3 in 190 Lines of Pure JAX: How It Works

⏱️ Reading time: 8 min

A 190-line pure JAX implementation reproduces the entire DeepSeek V3 architecture: multi-head latent attention, auxiliary-loss-free mixture-of-experts, and multi-token prediction included. Developer Vaibhaav Tiwari trained it end to end on a MacBook Air M5, using public-domain detective novels as the only training corpus, and published the full code on GitHub, with the last recorded push on November 14, 2026.

📑 En este artículo

1. What’s Inside DeepSeek V3’s 190 Lines of JAX
2. Who Should Read This Code
3. What We Tested: Installation Fails Outside macOS
4. The Numbers From the Author’s Own Training Run
5. DeepSeek V3 in JAX vs. Other Reimplementations
6. Limitations: What Not to Expect From This Code
7. What This Means for You if You Work With AI
8. References

What’s Inside DeepSeek V3’s 190 Lines of JAX

The file uses no Flax, no Equinox, no PyTorch: parameters are flat dictionaries of arrays, and every piece of the model is a function that reads top to bottom in deepseek_v3.py. That sets it apart from most educational reimplementations, which usually lean on a high-level framework that hides the internal mechanics. Tiwari also published a technical write-up with diagrams and a function-by-function explanation on his personal site.

The table the repository itself publishes maps every idea from the DeepSeek V3 technical report to a concrete function:

• Multi-head latent attention (MLA): instead of storing the full keys and values for every token, it caches a single smaller latent vector (480 down to 144 numbers per token per layer in this implementation), in the mla and init_cache functions. It’s like keeping a compressed summary of the conversation instead of the full history, decompressing it only when the model needs to attend to that token.
• Decoupled RoPE: a single rotary key shared across all heads, kept outside the compressed latent, in rope and apply_rope.
• DeepSeekMoE: a shared expert that always runs, plus 16 routed experts of which 4 activate per token, in route and moe. It works similarly to an on-call doctor who’s always available, plus up to four specialists called in depending on the case.
• Auxiliary-loss-free balancing: a per-expert bias that nudges routing toward less-used experts without touching their weights, in update_bias and route.
• Multi-token prediction (MTP): an extra module that also predicts the token after the next one, with a weight of 0.3 in the loss function, in forward and loss_fn.

Who Should Read This Code

The natural audience isn’t someone looking to train a production model, but a developer or student who has already read about DeepSeek V3 and wants to see, line by line, how each idea (MLA, expert routing, load balancing) translates into code that actually runs. It’s useful, for example, for a backend engineer who wants to understand what an MoE router really does before deciding whether the complexity is worth it in their own project, or for someone preparing an internal talk who needs an example that fits entirely on one screen.

For an ML engineer with no access to a GPU cluster, it’s also the cheapest way to get hands-on with these concepts: the entire example training run happened on a consumer laptop, not cloud infrastructure. It isn’t useful, on the other hand, for anyone looking for a production-ready library: the README itself states this is an architecture demonstration. The author says as much after showing a sample of generated text: it isn’t meant to compete as a language model.

What We Tested: Installation Fails Outside macOS

We cloned the repository in an isolated Linux sandbox and followed the steps in the README. The clone took 1.208 s and the repository was 16.4 MB; creating the virtual environment took 1.812 s. The dependency installation failed after 0.716 s:

$ .venv/bin/python -m pip install --no-input -r requirements.txt
Collecting jax==0.4.34 (from -r requirements.txt (line 1))
  Downloading jax-0.4.34-py3-none-any.whl.metadata (22 kB)
ERROR: Could not find a version that satisfies the requirement jax-metal==0.1.1 (from versions: none)
ERROR: No matching distribution found for jax-metal==0.1.1

The reason isn’t a bug: jax-metal is only published for macOS on Apple Silicon, and requirements.txt pins the exact versions jax==0.4.34 and jax-metal==0.1.1, with no alternative. The README documents the installation like this:

git clone https://github.com/Vaibhaav-Tiwari/deepseek-v3-jax.git && cd deepseek-v3-jax
uv venv -p 3.12 .venv
uv pip install --python .venv/bin/python -r requirements.txt

Peak memory for the entire test (clone, environment, and failed install) was 86.1 MB. The README mentions an alternative path, “No Mac? See Running on CPU,” but following the default steps breaks before reaching that option: you have to manually edit requirements.txt to remove the Metal dependency, and the excerpt we had on hand doesn’t detail how long that CPU path takes.

The Numbers From the Author’s Own Training Run

We didn’t test this ourselves: these are the figures the author published about his own run, on a 24 GB MacBook Air M5. The model has 36.2 million total parameters, of which only 13.6 million activate per token thanks to expert routing. He trained on 207 books (21.7 million training tokens, with Dashiell Hammett’s novels repeated 3 times) for 12,000 steps, about 98 million tokens seen in total (~4.5 epochs), at a median of 10,700 tokens per second, over 2.6 hours.

Validation loss dropped from 5.02 at step 250 to 3.51 by the end of training. Auxiliary-loss-free balancing produces its clearest data point at step 2,000. With the bias enabled, per-expert load stayed between 5.7% and 6.9% (an even split would be 6.25%); without the bias, it swung between 0.09% and 18.6%. Despite that huge difference in distribution, validation loss at that same step was nearly identical with and without balancing (3.921 versus 3.923), confirming the original paper’s argument: the bias spreads work across experts without costing the model anything (source).

The model itself, generated from the prompt “The fog came in off the”: “The fog came in off the bay, and the great trees were spreading and the stars were shining” (source).

flowchart TD
  A["Input token"] --> B["Shared expert (always active)"]
  A --> C["Router: top-4 of 16 experts"]
  C --> D["Expert 1"]
  C --> E["Expert 2"]
  C --> F["Expert N"]
  B --> G["Weighted sum"]
  D --> G
  E --> G
  F --> G
  G --> H["MoE layer output"]

DeepSeek V3 in JAX vs. Other Reimplementations

It isn’t the only way to read a modern LLM architecture without depending on a heavy framework. It’s worth comparing it to two well-known references in the same space:

Project	Framework and scale	What it’s for
deepseek-v3-jax	Pure JAX, no Flax or Equinox; built for a single laptop	Reading end to end how MLA, MoE, and MTP are implemented together
Official DeepSeek-V3 (deepseek-ai)	PyTorch, built for distributed production inference	Running the real model, not studying the architecture step by step
MaxText (Google)	JAX, geared toward training at scale on TPU and GPU	A starting point for training large models, not for reading as a tutorial

The underlying difference is one of purpose, not quality: the official repository and MaxText are built to scale, this one to be understood. None of the three replaces the other two.

Limitations: What Not to Expect From This Code

• It’s an architecture demonstration. The author states in the README that it isn’t meant to be a competitive language model.
• It only runs directly on macOS with Apple Silicon, due to the pinned jax-metal==0.1.1 dependency; on Linux, the default installation fails, as we confirmed in our test.
• The published checkpoint is a Python pickle, a format that can execute arbitrary code when loaded; the repository itself recommends loading only checkpoints from trusted sources.
• Training data hashes may not match across different runs, because some books come from Project Gutenberg and that archive changes over time (the Hammett novels, by contrast, are pinned to an exact Standard Ebooks commit and are reproducible).

What This Means for You if You Work With AI

If you work with language models and have never read the source code of a modern MoE architecture, this repository is faster to go through than a technical report and more transparent than a library that hides the details from you. To actually run it you need a Mac with Apple Silicon: our test confirms that in a standard Linux environment, the kind typical of a CI pipeline or a cloud instance, the default installation breaks on jax-metal, and the README leaves no documented shortcut to skip that dependency without manually editing requirements.txt.

If your goal is different (running a real DeepSeek model or training something at scale), this code isn’t the starting point; that’s what the official DeepSeek-V3 repository or a framework like MaxText are for.

References

• Vaibhaav-Tiwari/deepseek-v3-jax: source code, README, and original training results.
• Release v1.0: checkpoint and exact data from the published run, dated November 14, 2026.
• Vaibhaav Tiwari’s technical write-up: diagrams and a function-by-function explanation of the 190 lines.

Featured image: Foto de CDC en Unsplash

How each article is made (sources, testing and editorial criteria): How we work.