Vaibhaav-Tiwari/deepseek-v3-jax

What this is

This repository reimplements the core DeepSeek-V3 architecture from the technical report in 190 non-blank, non-comment lines of plain JAX — no Flax, Equinox, or PyTorch. Parameters are dictionaries of arrays, and the model is written as functions in deepseek_v3.py. It covers Multi-head Latent Attention with decoupled RoPE, DeepSeekMoE (1 shared + 16 routed experts, top-4), auxiliary-loss-free bias balancing, a sequence-wise balance loss, multi-token prediction, and cached prefill/decode generation.

The published run is a small pretraining experiment, not a competitive language model: 36.2M parameters (13.6M active per token), 6 layers, 256-token context, trained on ~21.7M tokens of public-domain detective fiction. The author states validation loss dropped from 5.02 to 3.51 over 12,000 steps and that the routing bias evens expert load roughly without changing validation loss at step 2000. A companion write-up explains the code piece by piece.

The project suits people who want to read or modify a faithful, compact implementation of DeepSeek-V3's architectural ideas, or who want a small JAX training example. It is explicitly an architecture demo, not a usable chat model.

Quick start

Requirements: an Apple Silicon Mac (the published run used 24 GB), Python 3.12, and uv. Clone and install pinned dependencies:

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

Every version is pinned; jax-metal 0.1.1 requires jax==0.4.34, so don't upgrade JAX. Verify the setup with JAX_PLATFORMS=cpu .venv/bin/pytest tests.py -q (39 tests). To run the pretrained checkpoint instead of training, download the v1.0 release checkpoint and tokenizer, then sample:

mkdir -p data runs/main
curl -L -o runs/main/ckpt.pkl  <release ckpt.pkl URL>
curl -L -o data/tokenizer.json <release tokenizer.json URL>
.venv/bin/python sample.py --prompt "The fog came in off the" --num 3

The checkpoint is a Python pickle, which can execute code when loaded; only load checkpoints from sources you trust. Full training requires building the corpus with prepare_data.py, which downloads ~120 MB of books and trains the tokenizer, then running train.py (~2.6 hours for the main run). Exact release URLs, data hashes, and the ablation commands are in the reproduction steps.

Constraints

• The full training run is tuned for Apple Silicon with jax-metal; a CPU fallback exists via JAX_PLATFORMS=cpu but the smaller --config small preset has only been checked for a couple of steps, not trained to completion.
• The architecture is deliberately simplified: dense MoE compute (every expert runs, unpicked ones zeroed), explicit latent decompression instead of weight absorption, and no FP8, DualPipe, node-limited routing, YaRN, or post-training. Scale is 36.2M parameters versus 671B in the real model.
• Re-running prepare_data.py will likely produce different dataset hashes because Project Gutenberg files cannot be pinned; the README suggests downloading the published data for byte-identical reproduction.
• Training text is public domain in the US only; status elsewhere may differ. Code is MIT-licensed.

HotRepo's take

The project appeared on the GitHub new-project list for Python on 2026-11-15, gaining 9 stars in a one-hour window with 12 forks. Its Show HN thread drew 180 points and 9 comments the day before, where commenters praised the from-first-principles approach (armcat) while one reader asked for a roadmap to build intuition (mattfrommars); replies pointed to 3blue1brown, Karpathy's tutorials, JAXformer, and nanochat. Another commenter noted that using datasets others have used makes results easier to compare (nico), which is relevant given the single small run here. The technical value is the readable reference implementation and reproducible ablation, not benchmark results.