DeepSeek V3 Architecture Distilled to 190 Lines of Pure JAX

A developer has created a minimal implementation of the DeepSeek-V3 language model architecture using pure JAX, with all non-blank code lines totaling approximately 190. The implementation uses only plain JAX functions and array dictionaries, eliminating dependencies on heavier frameworks like PyTorch, Flax, or Equinox. Every component reads sequentially in a single file, designed to be fully transparent and understandable.

The implementation was trained from scratch on a MacBook Air M5 using public-domain detective fiction, completing 12,000 training steps in roughly 2.6 hours on 24GB of memory. While dramatically smaller than the original DeepSeek-V3 (36.2M versus 671B parameters), it successfully demonstrates key architectural features including mixture-of-experts routing and expert balancing. Validation loss metrics confirmed that the expert balancing mechanism performs as the original paper describes, with similar results whether bias correction was enabled or disabled.

This is explicitly positioned as an architecture demonstration rather than a production system. The project omits various optimizations from the original, including FP8 training, sparse expert dispatch, and advanced context extensions. Complete setup instructions, training code, and checkpoints are included for reproducibility under an MIT license, using public-domain training data.