BMTrain 是 OpenBMB 推出的高效大模型训练工具包，用于训练参数量达数百亿级别的大模型，支持预训练与微调。其核心目标是让分布式训练代码保持接近单机训练的简洁程度。
用途 BMTrain 面向大模型（可含数百亿参数）的训练，涵盖预训练和微调场景，通过分布式方式完成训练。
主要能力
- 以接近单机训练的代码风格进行分布式训练。
- 支持 ZeRO 优化：将 torch.nn.Module 替换为 bmtrain.DistributedModule，将 torch.nn.Parameter 替换为 bmtrain.DistributedParameter，并用 bmtrain.Block 包装 transformer 模块。
- 支持通信优化：使用 bmtrain.TransformerBlockList 替代 torch.nn.ModuleList，并调整模块调用方式，以降低通信额外开销并将通信与计算时间重叠。
- 支持所有 PyTorch 原生优化器与损失函数，并提供 BMTrain 的融合优化器用于混合精度训练。
- 提供 bmtrain.lr_scheduler 模块中的常用 LRScheduler。
- 提供 OptimManager 统一管理优化器，由其执行 zero_grad()、backward()、clip_grad_norm() 和 step()。混合精度训练时可通过 loss scale 防止梯度下溢，loss_scale 在训练中根据梯度自适应调整；不使用混合精度时可设为 None（默认）。
- 支持张量并行（1.0.0 版本引入），并进行了代码重构。
使用方式 安装：推荐使用 pip install bmtrain；也可下载源码后运行 pip install .。安装需编译 c/cuda 源码，可能耗时几分钟到十几分钟，建议直接在训练环境中编译，以避免环境差异带来的问题。 使用步骤：
- 初始化：在代码开头调用 init_distributed，类似 PyTorch 分布式模块开头使用 init_process_group。注意使用 BMTrain 时不要使用 PyTorch 的分布式模块及其相关通信函数。
- 启用 ZeRO 优化：进行上述模块替换并用 bmtrain.Block 包装 transformer 模块。
- 启用通信优化：使用 TransformerBlockList 并调整模块调用。
- 启动分布式训练：使用与 PyTorch 分布式模块相同的启动命令，命令中涉及 MASTER_ADDR（主节点 IP）、MASTER_PORT（主节点端口）、NNODES（节点总数）、GPU_PER_NODE（每节点 GPU 数）、NODE_RANK（当前节点序号）。
示例与训练流程 官方提供基于 BMTrain 训练 GPT-2 的示例，主要包括：模型定义（使用 DistributedModule 与 DistributedParameter 启用 ZeRO 优化）、BMTrain 初始化（init_distributed 设置随机种子，init_parameters 初始化分布式参数）、优化器与学习率调度器初始化、训练循环（流程与普通训练循环相似，额外通过 optim_manager 的 add_optimizer 管理优化器）。
性能与支持模型 官方以 4 台服务器、每台 8 张 V100 训练了 13B 参数 GPT-2 模型，并测量每 GPU 吞吐（每 GPU 每秒样本数），对比了 DeepSpeed + Megatron 的 ZeRO 各阶段与模型并行配置。GPT-2 模型结构为：128 个注意力头、5120 隐藏维度、512 序列长度。BMTrain 已将 NLP 中多数常见模型迁移支持，列表见 ModelCenter 仓库。
适用场景 需要以较低代码改造成本进行大模型分布式预训练或微调的研究与工程场景；已使用 PyTorch 且希望启用 ZeRO 与通信优化的训练流程。
限制与说明
- 使用 BMTrain 时不应同时使用 PyTorch 分布式模块及其通信函数。
- BMTrain 对 PyTorch 底层做了改动，若程序输出异常，可通过 issue 提交信息。
- 安装需要编译 c/cuda 源码，耗时且建议在训练环境中进行。
- 具体性能数据与对比细节以官方文档与仓库为准。
许可与社区 该包以 Apache 2.0 许可发布。可通过 QQ 群 735930538、网站 https://www.openbmb.cn、微博 http://weibo.cn/OpenBMB、Twitter https://twitter.com/OpenBMB 联系。
