$ show salesforce/jaxformer
Jaxformer is a JAX-based library for training large language models (LLMs) on TPUs with data and model parallelism. It provides a complete training loop supporting distributed training across TPU cores, checkpointing, and resumable data loading from TFRecord files. The codebase includes model definitions for transformer-based architectures (like CodeGen), utilities for learning rate scheduling and optimization, and integration with Hugging Face for model conversion and sampling. The system is designed for scalability, supporting models up to 6B parameters via sharding strategies, and includes tools for preprocessing large datasets into TFRecord format. While the core training functionality is tightly coupled to the TPU environment, several components like the resumable data iterator and checkpointing mechanism demonstrate reusable patterns.
The LLM identified 3 files worth highlighting for reuse:
Provides a resumable iterator for TFRecord datasets that can resume training from a specific file and batch. Handles sharding, repeating, and batching across multiple files in a distributed setting. Self-contained with minimal dependencies (TensorFlow, NumPy, standard library) and showcases robust data loading patterns for large-scale training.
Implements checkpoint saving and loading for model and optimizer states in a distributed JAX environment. Each process saves its own shard to disk (e.g., GCS), enabling fault-tolerant training. The code is isolated, uses standard libraries (JSON, NumPy) with only minor dependencies on JAX and utils, and demonstrates a clean pattern for reusable checkpointing in ML systems.
Contains reusable training utilities including a GPT-3 style learning rate scheduler (linear warmup + cosine decay) and a weight decay function that excludes layer normalization and bias terms. These functions are self-contained, rely only on JAX and Optax, and represent standard practices in transformer training that can be easily extracted for use in other projects.