Writing an LLM from scratch, part 34a -- building a JAX training loop for an LLM training run

15 hours ago

The author built an LLM training loop from scratch using JAX, NNX, and Optax, based on notes from Sebastian Raschka's book.
An 'A-to-A' model (identity mapping) was used to test the training loop before building a full LLM.
Key challenges included OOM errors, JIT decorator issues, and learning how to use JAX for gradient accumulation.
Gradient accumulation was implemented using Optax's MultiSteps wrapper, and gradient clipping used clip_by_global_norm.
Checkpointing and restarting from checkpoints were enabled using Safetensors and Orbax, with a dispatcher in the learning rate extraction code.
The training loop achieved 110k tokens/second throughput and a final loss of 0.001 on the A-to-A task.
Learning rate scheduling used Optax's warmup_cosine_decay_schedule, with automatic charting of the learning rate over time.
The next post will cover building and training the actual small LLM using this training harness.