$ show mosaicml/llm-foundry
LLM Foundry is a production-grade framework for training, finetuning, evaluating, and serving large language models, with a focus on the MPT architecture and HuggingFace integration. It provides a full pipeline from data preparation (MDS streaming), to model definition (custom MPT, HF wrappers, embedding models), to distributed training (Composer, FSDP, TP), and to evaluation (ICL benchmarks, gauntlets). The architecture is modular, using YAML-driven configuration and a central registry pattern for extensibility. The codebase includes dozens of callbacks for monitoring and training stabilization, custom optimizers, and efficient data collation (bin packing). While tightly integrated with MosaicML's Composer and platform, the core registry and layer abstractions are reusable. The repository is actively maintained with comprehensive test suites, though some features remain experimental and several dependencies (flash attention, megablocks) require careful version management.
The LLM identified 3 files worth highlighting for reuse:
Provides TypedRegistry class and construct_from_registry function for decoupled component registration and instantiation. Supports type validation, entry point discovery, and keyword argument injection. This pattern is used throughout the codebase for models, data loaders, optimizers, callbacks, and more, and could be extracted into a standalone project for any scenario requiring plugin-style component management.
Contains BinPackCollator and FastBinPackCollator classes for efficient sequence packing to minimize padding tokens. Includes auto-profiling to determine optimal packing ratio. This is a reusable solution to a common problem in LLM training (sequence length variation) and could be adapted for any seq2seq project.
Implements KillLossSpike callback that detects divergence in training loss using a rolling window and terminates training early to save resources. Also supports resuming from a checkpoint before the spike. This is a novel, self-contained component for robust training pipelines.