Lyra is an open-source stack for training large language models in JAX on Google TPUs. It’s an end-to-end stack implementing the full infrastructure behind pretraining, from data preparation and model implementation to the optimizer, training loop, and custom TPU kernels.
A central engineering challenge was making the computations inside the training loop use TPU compute and memory efficiently. The project combines a complete pretraining system with targeted optimization of its most expensive operations.
Training stack
The model follows a GPT-OSS-style decoder architecture, combining local and global attention with Mixture-of-Experts layers that select a subset of experts for each token.
The surrounding stack includes a hybrid Muon + Adam optimizer, configurable sharding across devices, resumable checkpoints, training metrics, evaluation, and text generation.
A 500-million-token run of Lyra Small (1.6B parameters, 1.0B active per token) on a single TPU v6e exercised the complete stack before full pretraining. It ran 7,630 updates in 4h 40m end to end at about 34,000 tokens/s (24% MFU), and validation loss fell from 12.62 to 3.06. Diagnostics tracked throughput, gradient norms, and expert routing, which stayed balanced throughout. The run notes include the full training plots, and the checkpoint is on Hugging Face.
Faster training with Pallas
Custom kernels written in Pallas, JAX’s interface for programming accelerators, optimize attention, linear cross-entropy, and the matrix multiplications used by Mixture-of-Experts layers.
One example is cross-entropy, which measures prediction error during training. The fused kernel computes prediction scores and loss in blocks, avoiding the need to store a large matrix of scores for every token across the full vocabulary.
On one TPU v6e at lyra-small shapes (batch 4, 4,096-token context), the selected kernels achieve up to 5.27× speedups over matched JAX/XLA reference implementations, alongside substantial reductions in temporary memory:
| Operation | Pallas | XLA | Speedup | Temp. memory saved |
|---|---|---|---|---|
| Cross-entropy forward | 5.07 ms | 7.33 ms | 1.45× | >99.9% |
| Cross-entropy forward + backward | 21.42 ms | 24.82 ms | 1.16× | >99.9% |
| Sliding attention forward | 1.50 ms | 5.97 ms | 3.98× | >99.9% |
| Sliding attention forward + backward | 4.25 ms | 22.40 ms | 5.27× | 98.5% |
| MoE GEMM up + SwiGLU | 1.51 ms | 7.26 ms | 4.80× | >99.9% |
| MoE GEMM down | 0.86 ms | 3.24 ms | 3.76× | >99.9% |
The full benchmark report includes the training-step profile, exact shapes, memory measurements, and additional kernels.
Preparing 30 billion tokens
The data preparation pipeline ran in a container on GCP, tokenizing a
30-billion-token slice of a shuffled copy of NVIDIA’s ClimbMix dataset with
o200k_harmony and writing the results as 100 ArrayRecord shards of 300 million
tokens each.
The prepared dataset is published on Hugging Face, so training runs can download ready-to-use shards without repeating the preprocessing step.
Closing
Lyra was a lot of fun to build, and it deepened my understanding of how model code executes on hardware. Exploring fusion boundaries, tiling strategies for multi-step reductions, array materialization, and control flow showed me how much these details can affect performance.
Working across the full stack let me trace those effects from the training loop down to the individual hardware instruction. I hope to return to TPUs soon and go deeper into distributed training and multi-host setups.
Thank you for reading!