Writing an LLM from scratch, part 34a -- building a JAX training loop for an LLM training run
Building a JAX training loop for an LLM from scratch using Flax NNX and Optax. First validated the harness with a minimal "A-to-A" model (embedding then projection) before adding Transformer layers. Key challenges included fixing JAX's GPU memory default and slow data iteration by committing data to CPU memory.
Background
- The author is documenting a personal challenge: build and train an LLM from scratch using only his own notes, referencing no book or previous code. He previously trained a GPT-2-small-class model in PyTorch and now rebuilds it in JAX to test his understanding.
- This is part 34a of a series; it covers the training harness (data loading, optimization loop, device management). Part 34b will cover the actual LLM architecture.
- JAX is Google's framework for high-performance numerical computing. Flax NNX provides neural network building blocks (like PyTorch), and Optax provides optimizers (e.g. AdamW).
- The author starts with an "A-to-A" model: an embedding layer plus a linear layer, trained to reproduce its input tokens. This minimal model validates the training loop before adding Transformer layers.
- A key lesson: JAX data must be explicitly placed on the correct device (CPU for storage, GPU for computation); failing to do so causes severe slowdowns.