BMTrain is an open-source toolkit designed for efficient large-scale model training, including pre-training and fine-tuning. It targets researchers and engineers working with models containing tens of billions of parameters, simplifying distributed training to feel like stand-alone development.

How It Works

BMTrain integrates with PyTorch, enabling distributed training through its init_distributed function, replacing PyTorch's native distributed module. It implements ZeRO optimization by requiring users to replace torch.nn.Module with bmtrain.DistributedModule and torch.nn.Parameter with bmtrain.DistributedParameter. Transformer blocks can be further optimized by wrapping them in bmtrain.Block with specified ZeRO levels. Communication overhead is reduced by using bmtrain.TransformerBlockList for sequential blocks.

Quick Start & Requirements

• Installation: pip install bmtrain (compiles C/CUDA source code, may take 5-10 minutes).
• Prerequisites: PyTorch, CUDA (implied by compilation).
• Usage: Initialize with bmt.init_distributed(), replace PyTorch modules with BMTrain equivalents, and launch using torch.distributed.launch or torchrun.
• Documentation: https://www.openbmb.org/ (linked via website).

Highlighted Details

• Supports ZeRO-2 and ZeRO-3 optimizations.
• Offers bmtrain.optim.AdamOffloadOptimizer and bmtrain.lr_scheduler for optimized training.
• Provides an OptimManager to handle optimizer zero-grad, backward, clipping, and step operations.
• Claims significant throughput improvements over standard ZeRO implementations in benchmarks.

Maintenance & Community

• Active development with releases up to v1.0.0 (Feb 2024) adding Tensor parallel support.
• Community channels include a QQ Group (735930538).
• Website: https://www.openbmb.org/
• Twitter: https://twitter.com/OpenBMB

Licensing & Compatibility

• Apache 2.0 License.
• Permits commercial use and linking with closed-source projects.

Limitations & Caveats

BMTrain makes deep modifications to PyTorch's internals, potentially leading to unexpected behavior. Users are advised to submit issues for any observed problems.