计算成本节省9成：大模型高效训练工具BMTrain
背景介绍
2018 年，预训练语言模型技术横空出世并引发了人工智能领域的性能革命。研究表明，增大参数量与数据规模是进一步提升语言模型性能的有效手段，对十亿、百亿乃至千亿级大模型的探索成为业界的热门话题。这引发了国内外研究机构与互联网企业的激烈竞争，将模型规模与性能不断推向新的高度。除 Google、OpenAI 等国外知名机构外，近年来国内相关研究机构与公司也异军突起，形成了大模型的研究与应用热潮，人工智能由此进入“大模型时代”。
然而在“大模型时代”，模型巨大的参数规模和算力需求，带来了很大的 训练和微调难题：
01 算力成本高
模型巨大的参数量无法在单张显卡中完成存储与计算。OpenAI 训练 GPT-3 （约 1750 亿参数）使用了上千张 GPU，Google 训练 PaLM （约 5000 亿参数）使用了六千片 TPU ，带来了高昂的训练成本，配置规模庞大的计算集群也十分困难。
02 编程难度大
为了利用分布式算力加速大模型的训练和微调，程序员需要编写复杂的分布式程序来驱动大模型。现有的框架（如 DeepSpeed ）已经能够较好地支持模型的分布式训练，但依然需要用户进行较为复杂的编程与配置。
03 计算速度低
在分布式训练框架中，不同计算节点之间需要频繁地进行通信，此时通信带宽往往成为模型的训练瓶颈，若不能合理设计通信策略并进行有效且正确的代码实现，训练模型所需的时间将被大大延长，不能充分发挥计算设备的潜力。
针对上述难题，我们推出了大模型高效训练工具包 BMTrain 与模型仓库 ModelCenter 。
BMTrain & ModelCenter
在 OpenBMB 全流程加速工具框架中，致力于解决大模型训练难题的 BMTrain 是极为重要的一环。在设计中，我们追求以下特点：
01 高效率
作为大模型训练的 “发动机”，BMTrain 能够在任意数量的 GPU 上进行高效的大模型预训练与微调，最优化分布式框架的通信开销，在超大规模模型训练场景下与 DeepSpeed 等框架相比可以节省 90% 的算力成本。
02 低资源
不同于 OpenAI、Google 等机构的算力条件，为了让更多实验室和企业也能够训练大模型，我们致力于在保持高效计算前提下最小化大模型的算力需求，实现单张消费级显卡全参数微调 BERT-Large，8 台 A100 小集群训练 GPT-3，进一步降低大模型的算力门槛。
03 可扩展
在编程难度方面，我们致力做最简洁最有效的封装，仅使用少量的代码替换，即可达到与原生 PyTorch 一致的编程体验，一键安装工具包降低配置难度，让大模型真正飞入千家万户。
与此同时，我们还基于 BMTrain 实现了一系列大规模预训练语言模型，集成于 ModelCenter 仓库。BMTrain 与 ModelCenter 共同组成了高效分布式预训练框架，适用于任意的 Transformer 结构，可以在较少数量的 GPU 上训练，且高度兼容 PyTorch、Transformers 库，学习成本极低。目前我们已经支持了常用的英文模型如 BERT、GPT、T5、RoBERTa 以及中文模型如 CPM-1、CPM-2 等。
简单易用
BMTrain 基于简单易用的设计原则，致力做最简洁最有效的封装，仅使用少量的代码替换，即可达到与原生 PyTorch 一致的编程体验，快速实现大模型的训练加速。我们还实现了模型仓库 ModelCenter，方便使用者快速使用主流架构大模型。
01 无缝衔接
贴合 PyTorch 使用习惯，上手门槛更低，仅需 简单替换 即可完成训练提速：
下面将提供一个简单的对比图来直观展示 BMTrain 的便捷易用（左边是原始代码，右边是替换后的代码）。
02 一键转换
对于 PyTorch 实现的原始模型，直接使用 BMTrain 提供的 BMTrainModelWrapper 自动包装模型，实现分布式加速。
03 模型仓库
我们基于 BMTrain 构建了模型仓库 ModelCenter，可以直接加载进行使用。ModelCenter 的使用方式与 Transformers 类似，且兼容 Transformers 中各类数据处理接口。
效果拔群
BMTrain在大模型训练上效果出色，在不同规模的算力条件下均有较好的性能表现。
01 消费级算力（单卡2080Ti）
在消费级显卡 2080Ti 上，BMTrain 可以实现 BERT-Large 的微调（3 亿参数，样本长度 512）。
02 入门级算力（单卡V100）
在入门级算力条件下（显卡为 V100 32GB），BMTrain 可以实现 BERT-Large 的高效训练（3 亿参数，样本长度 512）。
03 中等算力（单机8卡A100）
在中等规模算力条件下（显卡为 A100 40GB，NVLink），BMTrain 可以训练大规模的 GPT-13B (130亿参数，样本长度 512）。
04 高级算力（多机8卡A100）
在较高规模算力条件下（显卡为A100 40GB，NVLink，400Gbps IB），BMTrain 可以训练超大规模的 GPT-3（1750亿参数）。
使用 BMTrain 或者 ColossalAI ，64 卡 A100 跑完 GPT-3 的 300B token大概需要2年，服务器与显卡租金大约 900 万左右。根据我们的实验估算，使用 128 张 A100 时，单卡吞吐量可以提升 2.5 倍以上，6 个月可以跑完 GPT-3，服务器租金大约 500 万左右。虽然训练出 GPT-3 的成本依然高昂，但与 GPT-3 的 1200 万美元相比，成本仍然节约了 90% 以上。
上手教程
BMTrain 和 ModelCenter 提供了丰富的文档，方便使用者快速上手，便捷地体验大模型的魅力。
STEP 01 · 快速安装
01 BMTrain安装
· 可以通过 pip 直接安装（推荐）
· 也可以选择从 GitHub 进行源码安装
02 ModelCenter安装
· 可以通过 pip 直接安装（推荐）
· 也可以选择从 GitHub 进行源码安装
STEP 02 · 启动 BMTrain
首先，你需要在代码开头初始化 BMTrain。正如在使用 PyTorch 的分布式训练模块需要在代码开头使用 init_process_group 一样，使用 BMTrain 需要在代码开头使用 init_distributed。
注意： 使用 BMTrain 时请不要使用 PyTorch 自带的 distributed 模块包括 torch.distributed.init_process_group 以及相关通信函数。
STEP 03 · 性能优化
01 ZeRO-3优化
使用ZeRO-3优化只需要对模型代码进行简单替换：· torch.nn.Module 替换为 bmtrain.DistributedModule · torch.nn.Parameter 替换为 bmtrain.DistributedParameter
02 Checkpointing优化
在模块上套用 bmtrain.CheckpointBlock 即可
03 通信优化
为了进一步缩短通信额外开销，将通信与运算时间重叠，可以使用 TransformerBlockList 来进一步优化。在使用时需要对代码进行简单替换：· torch.nn.ModuleList 替换为 bmtrain.TransformerBlockList · for module in self.module_list: x= module(x, ...) 替换为 x = self.module_list(x, ...)
STEP 04 · 运行分布式训练代码
BMTrain 使用 PyTorch 原生分布式训练启动器，你可以根据 PyTorch 版本选择下列命令中的一个。
· torch.distributed.launch
· torchrun
其中：
${MASTER_ADDR} 为主节点的 IP 地址，只有一个节点可以写 localhost 或 127.0.0.1
${MASTER_PORT} 为主节点的端口
${NNODES} 为节点数量（一般为机器数量）
${GPU_PER_NODE} 为每个节点的 GPU 数量
${NODE_RANK} 为本节点的 RANK
${ARGS} 为代码输入的其他参数
STEP 05 · 使用ModelCenter
本节将首先以如何在一个分类数据集上微调 BERT 模型为例，介绍如何使用 ModelCenter 的现有模型。其次将简单介绍如何使用 ModelCenter 实现一个新的模型。
01 准备模型
接下来，你可以从 model_center 中获取预训练好的 BERT 模型，例如 bert-base-uncased。由于我们是在一个分类任务上微调 BERT 模型，所以需要在最后一层后添加一个全连接层。
如果只需要 config 来构建模型，而不需要现成的预训练参数，可以参考下面的方法：
02 准备数据集
下一步是准备数据集，用于训练和验证模型。这里，我们使用 SuperGLUE benchmark 中的 BoolQ 数据集。你需要下载该数据集，并将解压后的文件夹放在 your_path_to_dataset 路径下。
03 模型训练
现在，在设置 优化器、学习率调整策略和损失函数 后，就可以开始训练模型了！在 示例代码（https://modelcenter.readthedocs.io/en/latest/notes/quickstart.html）中，我们训练 BERT 模型 5 轮，并且在每轮训练结束后验证模型的性能。
04 运行代码
运行代码的方式和上节所给出的一样，你可以根据 PyTorch 版本选择其中的一个，具体命令可以参照上节：运行分布式训练代码 中所给出的示例。
05 如何创建新的模型
在 ModelCenter 的
model_center/layer中，我们给出了包括像Linear、LayerNorm、Embedding等等 常用模块 的实现。且这些模块都是基于bmtrain.DistributedParameter和bmtrain.DistributedModule实现的，支持分布式训练。因此在创建新的模型时，我们可以直接调用这些模块，完成网络的搭建。在 ModelCenter 的
model_center/layer中，我们同样实现了一些 常用的模块的组合，被称作block，如组合了Attention和Add&Norm的SelfAttentionBlock。每一个block都有不同的选择， 且支持pre-layernorm和post-layernorm。我们使用
bmtrain.CheckpointBlock和bmtrain.TransformerBlockList来包装我们的tranformer block。这些在不增加大量计算时间的情况下大量减少了 GPU 内存使用量。在 ModelCenter 所提供的 常用模型 的基础上，通过 加入相应的模块，你可以很容易地构建一个新的模型。你只需要将模型特定功能添加到通用的模型结构中，就可以轻而易举地实现一个新模型。
可以参考示例代码（https://modelcenter.readthedocs.io/en/latest/notes/write_model.html）中给出的实现方式。
未来展望
在未来，我们将在 BMTrain 中进一步完善训练的工具链，并达到更好的 PyTorch 兼容性；在 ModelCenter 中支持更多的预训练大模型，并完善模型对推理的支持。我们诚邀更多的开发者加入我们的开发队伍，共同维护 OpenBMB 开源社区。也欢迎大家使用 BMTrain 实现更多类型的大模型，共同维护大模型构建与应用生态，为大模型时代贡献一份力量。
🔗BMTrain：https://github.com/OpenBMB/BMTrain
关注我们
微信搜索关注 “OpenBMB开源社区”
获取更多大模型干货知识和前沿资讯！
🚪传送门｜相关链接
🔗 官方网站：https://www.openbmb.org
🔗 GitHub：https://github.com/OpenBMB
🔗 交流QQ群：735930538
🔗 启智社区：https://git.openi.org.cn/OpenBMB
🔗 微博：http://weibo.cn/OpenBMB
🔗 知乎：https://www.zhihu.com/people/OpenBMB
🔗 Twitter：https://twitter.com/OpenBMB
