最新动态

• 2024/02/26 BMTrain 1.0.0 发布。代码重构并支持张量并行。详见 更新日志
• 2023/08/17 BMTrain 0.2.3 发布。详见 更新日志
• 2022/12/15 BMTrain 0.2.0 发布。详见 更新日志
• 2022/06/14 BMTrain 0.1.7 发布。支持 ZeRO-2 优化！
• 2022/03/30 BMTrain 0.1.2 发布。适配 OpenPrompt 和 OpenDelta
• 2022/03/16 BMTrain 0.1.1 公开发布首个稳定版本，修复了 beta 版本中的许多 bug
• 2022/02/11 BMTrain 0.0.15 公开发布首个 beta 版本

概述

BMTrain 是一个高效的大模型训练工具包，可用于训练数十亿参数的大模型。它能以分布式方式训练模型，同时保持代码像单机训练一样简单。

文档

我们的文档提供了关于该包的更多信息。

安装

• 通过 pip（推荐）： pip install bmtrain
• 从源码安装：下载包后运行 pip install .

安装 BMTrain 可能需要几分钟到十分钟，因为它需要在安装时编译 c/cuda 源代码。我们建议直接在训练环境中编译 BMTrain，以避免因环境不同而可能产生的问题。

使用

第一步：初始化 BMTrain

在使用 BMTrain 之前，需要在代码开头初始化它。就像使用 PyTorch 的分布式模块需要在代码开头使用 init_process_group 一样，使用 BMTrain 需要在代码开头使用 init_distributed。

import bmtrain as bmt
bmt.init_distributed(
    seed=0,
    # ...
)

注意： 使用 BMTrain 时，不要使用 PyTorch 的分布式模块及其相关的通信函数。

第二步：启用 ZeRO 优化

要启用 ZeRO 优化，需要对原始模型的代码进行一些简单的替换。

• torch.nn.Module -> bmtrain.DistributedModule
• torch.nn.Parameter -> bmtrain.DistributedParameter

并用 bmtrain.Block 包装 transformer 块。

以下是一个示例。

原始代码

import torch
class MyModule(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.param = torch.nn.Parameter(torch.empty(1024))
        self.module_list = torch.nn.ModuleList([
            SomeTransformerBlock(),
            SomeTransformerBlock(),
            SomeTransformerBlock()
        ])

    def forward(self):
        x = self.param
        for module in self.module_list:
            x = module(x, 1, 2, 3)
        return x

替换后

import torch
import bmtrain as bmt
class MyModule(bmt.DistributedModule): # 修改处
    def __init__(self):
        super().__init__()
        self.param = bmt.DistributedParameter(torch.empty(1024)) # 修改处
        self.module_list = torch.nn.ModuleList([
            bmt.Block(SomeTransformerBlock(), zero_level=3), # 修改处，目前支持 2 和 3
            bmt.Block(SomeTransformerBlock(), zero_level=3), # 修改处，目前支持 2 和 3
            bmt.Block(SomeTransformerBlock(), zero_level=3)  # 修改处，目前支持 2 和 3
        ])

    def forward(self):
        x = self.param
        for module in self.module_list:
            x = module(x, 1, 2, 3)
        return x

第三步：启用通信优化

为了进一步减少通信的额外开销，并将通信时间与计算时间重叠，可以使用 TransformerBlockList 进行优化。

可以通过对代码进行以下替换来启用它们：

• torch.nn.ModuleList -> bmtrain.TransformerBlockList
• for module in self.module_list: x = module(x, ...) -> x = self.module_list(x, ...)

原始代码

import torch
import bmtrain as bmt
class MyModule(bmt.DistributedModule):
    def __init__(self):
        super().__init__()
        self.param = bmt.DistributedParameter(torch.empty(1024))
        self.module_list = torch.nn.ModuleList([
            bmt.Block(SomeTransformerBlock()),
            bmt.Block(SomeTransformerBlock()),
            bmt.Block(SomeTransformerBlock())
        ])

    def forward(self):
        x = self.param
        for module in self.module_list:
            x = module(x, 1, 2, 3)
        return x

替换后

import torch
import bmtrain as bmt
class MyModule(bmt.DistributedModule):
    def __init__(self):
        super().__init__()
        self.param = bmt.DistributedParameter(torch.empty(1024))
        self.module_list = bmt.TransformerBlockList([ # 修改处
            bmt.Block(SomeTransformerBlock()),
            bmt.Block(SomeTransformerBlock()),
            bmt.Block(SomeTransformerBlock())
        ])

    def forward(self):
        x = self.param
        for module in self.module_list:
            x = module(x, 1, 2, 3)
        return x

第四步：启动分布式训练

BMTrain 使用与 PyTorch 分布式模块相同的启动命令。

您可以根据 PyTorch 版本选择其中一种。

• ${MASTER_ADDR} 表示主节点的 IP 地址
• ${MASTER_PORT} 表示主节点的端口
• ${NNODES} 表示节点总数
• ${GPU_PER_NODE} 表示每个节点的 GPU 数量
• ${NODE_RANK} 表示当前节点的排名

torch.distributed.launch

$ python3 -m torch.distributed.launch --master_addr ${MASTER_ADDR} --master_port ${MASTER_PORT} --nproc_per_node ${GPU_PER_NODE} --nnodes ${NNODES} --node_rank ${NODE_RANK} train.py

torchrun

$ torchrun --nnodes=${NNODES} --nproc_per_node=${GPU_PER_NODE} --rdzv_id=1 --rdzv_backend=c10d --rdzv_endpoint=${MASTER_ADDR}:${MASTER_PORT} train.py

更多信息，请参考文档。

示例

我们提供了一个基于 BMTrain 训练 GPT-2 的示例。代码主要包括以下部分。

第一部分：模型定义

├── layers
│   ├── attention.py
│   ├── embedding.py
│   ├── feedforward.py
│   ├── __init__.py
│   ├── layernorm.py
│   └── linear.py
└── models
    ├── gpt.py
    └── __init__.py

以上是模型定义部分的代码目录结构。

我们定义了 GPT-2 中所需的所有层，并使用 BMTrain 的 DistributedModule 和 DistributedParameter 来启用 ZeRO 优化。

第二部分：BMTrain 初始化

bmtrain.init_distributed(seed=0)

model = GPT(
    num_layers=8,
    vocab_size=10240, 
    dim_model=2560,
    dim_head=80,
    num_heads=32,
    dim_ff=8192,
    max_distance=1024,
    bias=True,
    dtype=torch.half
)

bmtrain.init_parameters(model) # 或使用 `bmtrain.load` 加载检查点

# ... 其他初始化（数据集）...

bmtrain.init_distributed(seed=0) 用于初始化分布式训练环境并设置随机种子以确保可重复性。

bmtrain.init_parameters(model) 用于初始化模型的分布式参数。

第三部分：优化器和学习率调度器初始化

loss_func = torch.nn.CrossEntropyLoss(ignore_index=-100)
optimizer = bmtrain.optim.AdamOffloadOptimizer(model.parameters(), weight_decay=1e-2)
lr_scheduler = bmtrain.lr_scheduler.Noam(optimizer, start_lr=1e-3, warmup_iter=40, end_iter=1000, num_iter=0)

BMTrain 支持 所有 PyTorch 原生的优化器和损失函数，您也可以使用 BMTrain 提供的融合优化器进行混合精度训练。

此外，BMTrain 还在 bmtrain.lr_scheduler 模块中提供了常见的 LRScheduler。

第四部分：训练循环

# 创建一个新的优化器管理器实例
optim_manager = bmtrain.optim.OptimManager(loss_scale=1024)
# 让 optim_manager 管理所有优化器及（可选的）对应的学习率调度器
optim_manager.add_optimizer(optimizer, lr_scheduler)
# add_optimizer 可以多次调用来添加其他优化器

for iteration in range(1000):
    # ... 为每个进程加载数据 ...

    # 前向传播并计算损失
    pos = torch.arange(enc_input.size(1)).long().cuda().repeat(enc_input.size(0), 1)
    logits = model(
        enc_input,
        pos,
        pos < enc_length[:, None]
    )
    batch, seq_len, vocab_out_size = logits.size()

    loss = loss_func(logits.view(batch * seq_len, vocab_out_size), targets.view(batch * seq_len))

    global_loss = bmtrain.sum_loss(loss).item() # 对所有进程的损失求和，仅用于训练日志

    # 梯度清零
    optim_manager.zero_grad() # 调用每个优化器的 zero_grad

    # 损失缩放和反向传播
    optim_manager.backward(loss)

    # 梯度裁剪
    grad_norm = optim_manager.clip_grad_norm(optimizer.param_groups, max_norm=1.0)

    # 优化器步进
    optim_manager.step()

    # ... 保存检查点或打印日志 ...

训练循环部分会稍长一些，但就像普通训练循环一样，您不需要为分布式训练做太多适配。

您可以按照代码中的注释了解每段代码的作用。

唯一的额外说明是 optimizer。使用 BMTrain 后，优化器中的一些细节需要调整。我们在 optim_manager 中实现了所有这些所需的细节。您只需要让 optim_manager 通过 add_optimizer 管理所有优化器，并让 optim_manager 代替执行 zero_grad()、backward()、clip_grad_norm() 和 step() 即可。

如果您不使用混合精度训练，则可以在没有 loss_scale 的情况下进行训练。只需在 OptimManager 的 __init__ 函数中将 loss_scale 设置为 None（这也是默认值）。

如果您使用混合精度训练，loss scale 是混合精度训练中广泛使用的防止梯度下溢的技术。通过在 optim_manager.backward(loss) 中在反向传播前缩放 loss，并在 OptimManager 的 __init__ 函数中将 loss_scale 设置为某个浮点数。loss_scale 会根据训练过程中的梯度自适应调整。

性能

我们使用 4 台服务器，每台服务器配备 8 个 V100，训练了一个 13B 参数的 GPT-2 模型，并测量了训练过程中每个 GPU 的吞吐量（每个 GPU 每秒处理的样本数）。

模型结构：
* 40 层
* 128 个注意力头
* 5120 隐藏维度
* 512 序列长度

ZeROa(mp=b) 表示 DeepSpeed + Megatron ZeRO 阶段 a 且模型并行度 = b。

- 表示内存溢出。

支持的模型

我们已经将 NLP 中大多数常见模型迁移到了 BMTrain。您可以在 ModelCenter 仓库中找到支持的模型列表。

社区

我们欢迎所有人按照我们的贡献指南贡献代码。

您也可以在其他平台上找到我们：
- QQ 群：735930538
- 网站：https://www.openbmb.org
- 微博：http://weibo.cn/OpenBMB
- Twitter：https://twitter.com/OpenBMB

许可证

该包根据 Apache 2.0 许可证发布。

其他说明

BMTrain 对 PyTorch 进行了底层更改，因此如果您的程序输出意外结果，您可以在 issue 中提交相关信息。