Skip to content

Latest commit

 

History

History
369 lines (261 loc) · 12.6 KB

README-ZH.md

File metadata and controls

369 lines (261 loc) · 12.6 KB
Raw

BMTrain

大模型高效训练工具包

总览 • 文档 • 安装 • 使用说明 • 性能 • English

Documentation Status GitHub release (latest by date including pre-releases) GitHub

最新动态

  • 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 版本中的一些问题。
  • 2022/02/11 BMTrain 0.0.15 公开发布了第一个 beta 版本。

总览

BMTrain 是一个高效的大模型训练工具包，可以用于训练数百亿参数的大模型。BMTrain 可以在分布式训练模型的同时，能够保持代码的简洁性。

文档

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

安装

  • 用 pip 安装（推荐）: pip install bmtrain

  • 从源代码安装: 下载工具包，然后运行 pip install . (setup.py的安装方式将会在未来被setuptools弃用)

安装 BMTrain 可能需要花费数分钟的时间，因为在安装时需要编译 c/cuda 源代码。 我们推荐直接在训练环境中编译 BMTrain，以避免不同环境带来的潜在问题。

使用说明

步骤 1: 启用 BMTrain

首先，你需要在代码开头初始化 BMTrain。正如在使用 PyTorch 的分布式训练模块需要在代码开头使用 init_process_group 一样，使用 BMTrain 需要在代码开头使用 init_distributed。

import bmtrain as bmt
bmt.init_distributed(
    seed=0,
    zero_level=3,   # 目前支持2和3
    # ...
)

注意： 使用 BMTrain 时请不要使用 PyTorch 自带的 distributed 模块，包括 torch.distributed.init_process_group 以及相关通信函数。

步骤 2: 使用 ZeRO 优化

使用ZeRO优化需要对模型代码进行简单替换：

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

并在 transformer 模块上使用 bmtrain.CheckpointBlock。

下面是一个例子：

原始代码

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.CheckpointBlock(SomeTransformerBlock()), # 修改这里
            bmt.CheckpointBlock(SomeTransformerBlock()), # 修改这里
            bmt.CheckpointBlock(SomeTransformerBlock())  # 修改这里
        ])
    
    def forward(self):
        x = self.param
        for module in self.module_list:
            x = module(x, 1, 2, 3)
        return x

步骤 3: 通信优化

为了进一步缩短通信额外开销，将通信与运算时间重叠，可以使用 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.CheckpointBlock(SomeTransformerBlock()),
            bmt.CheckpointBlock(SomeTransformerBlock()),
            bmt.CheckpointBlock(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.CheckpointBlock(SomeTransformerBlock()),
            bmt.CheckpointBlock(SomeTransformerBlock()),
            bmt.CheckpointBlock(SomeTransformerBlock())
        ])
    
    def forward(self):
        x = self.param
        x = self.module_list(x, 1, 2, 3) # 修改这里
        return x

步骤 4: 运行分布式训练代码

BMTrain 使用 PyTorch 原生的分布式训练启动器，你可以根据 PyTorch 版本选择下列命令中的一个。

  • ${MASTER_ADDR} 为主节点的 IP 地址
  • ${MASTER_PORT} 为主节点的端口
  • ${NNODES} 为节点数量
  • ${GPU_PER_NODE} 为每个节点的 GPU 数量
  • ${NODE_RANK} 为本节点的 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

更多信息请参考 PyTorch 官方文档。

样例

我们提供了一个使用 BMTrain 训练 GPT-2 的样例。 代码主要包含以下几个部分。

第 1 部分: 模型定义

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

上面是代码的目录结构。

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

第 2 部分: 初始化 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`加载checkpoint

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

bmtrain.init_distributed(seed=0) 用于初始化分布式训练环境，并设置随机数种子便于复现。

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

第 3 部分: 初始化优化器和学习率调整策略

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 提供的融合（fused）优化器用于混合精度训练。

此外，在 bmtrain.lr_scheduler 中 BMTrain 也提供了常见的学习率调整策略。

第 4 部分: 训练

# 新建优化器管理器实例
optim_manager = bmtrain.optim.OptimManager(loss_scale=1024)
# 将所有的 optimzer 及(可选)其对应的 lr_scheduler 收入优化器管理器管理。
optim_manager.add_optimizer(optimizer, lr_scheduler)
# 可以再次调用 add_optimizer 加入其他优化器

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

    # 前向传播并计算梯度
    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() # 聚合所有rank上的损失, 仅用于输出训练日志

    # 梯度清零
    optim_manager.zero_grad() # 为每个 optimizer 调用 zero_grad

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

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

    # 更新参数
    optim_manager.step()

    # ... 保存checkpoint、打印日志 ...

这部分代码略有些长，但写起来就像常见的训练代码一样，你不需要为分布式训练调整太多的代码。

你可以根据代码中的注释来了解各部分代码的作用。

唯一需要说明的是 optim_manager。在使用 BMTrain 后，优化器的部分相关操作需要有一些细节上的调整。我们在 optim_manager 帮你实现了这些细节, 你只需要通过 add_optimizer 将优化器和学习率调整策略收入 optim_manager 管理，并由 optim_manger 代为执行 zero_grad(), backward(), clip_grad_norm() 和 step() 等操作。

如果你没有使用混合精度训练，你可以不用损失缩放，只需要将 OptimManger(loss_scale=None) 构造函数中 loss_scale 置为 None 即可, 这也是 OptimManager 的默认构造参数。

如果你使用了混合精度训练，损失缩放是混合精度训练中的一项常用技术，我们在 optim_manager.backward(loss) 帮你对 loss 进行了放缩，用于避免梯度下溢。只需要将 OptimManger 构造函数中 loss_scale 置为一个浮点数即可。 loss_scale 会在训练过程中根据梯度进行自适应的调整。

性能

我们训练了一个有130亿参数的 GPT-2 模型，使用了4台服务器，每台服务器有8张V100显卡。我们测试了训练过程中每个GPU的吞吐量（每个GPU每秒处理的样本数），结果见下表。

模型结构：

  • 40层
  • 128个注意力头
  • 5120的隐藏层维数
  • 512的序列长度
batch size 8 16 24 32
BMTrain 24.15 26.94 29.42 28.28
ZeRO3(mp=1) 14.88 21.69 24.38 -
ZeRO3(mp=4) 15.51 - - -
ZeRO3(mp=8) 15.51 - - -
ZeRO2(mp=1) - - - -
ZeRO2(mp=4) 22.85 - - -
ZeRO2(mp=8) 21.33 - - -

ZeROa(mp=b) 表示 DeepSpeed + Megatron ZeRO stage a 和 model parallelism = b。

表格中的 - 表示超出显存。

模型支持

我们已经将大多数常见的 NLP 模型移植到了 BMTrain 中。你可以在 ModelCenter 项目中找到支持模型的列表。

开源社区

欢迎贡献者参照我们的贡献指南贡献相关代码。

您也可以在其他平台与我们沟通交流：

开源许可

该工具包使用Apache 2.0开源许可证。

其他说明

BMTrain 工具包对 PyTorch 进行了底层修改，如果你的程序输出了意料之外的结果，可以在 issue 中提交相关信息。