MosaicML Composer

MosaicML Composer is a Python library for training neural networks on your own GPUs or multi-node clusters. Built on PyTorch, it's aimed at researchers and engineers who want distributed training with room to customize their experiments. It's open source under Apache 2.0, and prebuilt Docker images are available.

Its Trainer handles the training loop, distributed data loading and memory optimization. It supports large language models, diffusion models, BERT-style embedding models and CNNs. MosaicML uses it to train MPT models.

Composer supports PyTorch Fully Sharded Data Parallel (FSDP) to spread large models across GPUs, as well as Distributed Data Parallel (DDP). Sharded checkpoints can resume with a different GPU count. Automatic resumption picks up interrupted runs from the latest saved checkpoint, while automatic microbatch sizing selects a size that fits GPU memory.

For custom experiments, callbacks let researchers insert their own logic throughout training. Composer also includes training speedup algorithms that can be combined into model-specific recipes, plus callbacks for memory monitoring, image logging and training-time estimates.

Training runs on the hardware you provide. Optional integrations support remote dataset streaming, cloud checkpoint storage and external experiment tracking. They are separate from running training on your own hardware.