其他
计算成本节省9成：大模型高效训练工具BMTrain
针对上述难题，我们推出了大模型高效训练工具包BMTrain与模型仓库ModelCenter。
在编程难度方面，我们致力做最简洁最有效的封装，仅使用少量的代码替换，即可达到与原生 PyTorch 一致的编程体验，一键安装工具包降低配置难度，让大模型真正飞入千家万户。
bmtrain.DistributedParameter 替换 torch.nn.Parameter bmtrain.DistributedModule 替换 torch.nn.Module bmtrain.CheckpointBlock 替换 torch.nn.ModuleList 中的模块 BMTrainModelWrapper 自动包装模型，实现分布式加速。# Automatically wrap a model in a BMTrain model
bmt_model = BMTrainModelWrapper(model) # model: torch.nn.module
▶ 消费级算力（单卡2080Ti）
▶ 入门级算力（单卡V100）
▶ 中等算力（单机8卡A100）
在中等规模算力条件下（显卡为A100 40GB，NVLink），BMTrain可以训练大规模的 GPT-13B (130亿参数，样本长度 512）。
▶ 高级算力（多机8卡A100）
在较高规模算力条件下（显卡为A100 40GB，NVLink，400Gbps IB），BMTrain可以训练超大规模的 GPT-3（1750亿参数）。
STEP 01 · 快速安装
▶ BMTrain安装
· 也可以选择从GitHub进行源码安装
$ cd BMTrain
$ python3 setup.py install
▶ ModelCenter安装
$ cd ModelCenter
$ pip install -r requirements.txt
$ python3 setup.py install
STEP 02 · 启动BMTrain
bmt.init_distributed()
注意： 使用 BMTrain 时请不要使用 PyTorch 自带的 distributed模块包括torch.distributed.init_process_group以及相关通信函数
STEP 03 · 性能优化
▶ ZeRO-3优化
torch.nn.Module 替换为 bmtrain.DistributedModule torch.nn.Parameter 替换为 bmtrain.DistributedParameter ▶ Checkpointing优化
bmtrain.CheckpointBlock 即可▶ 通信优化
TransformerBlockList 来进一步优化。在使用时需要对代码进行简单替换： torch.nn.ModuleList 替换为 bmtrain.TransformerBlockList for module in self.module_list: x= module(x, ...) 替换为 x = self.module_list(x, ...) BMTrain 使用 PyTorch 原生分布式训练启动器，你可以根据 PyTorch 版本选择下列命令中的一个。
$ 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 ${ARGS}
${MASTER_ADDR} 为主节点的 IP 地址，只有一个节点可以写 localhost 或 127.0.0.1 ${MASTER_PORT} 为主节点的端口 ${NNODES} 为节点数量（一般为机器数量） ${GPU_PER_NODE} 为每个节点的 GPU 数量 ${NODE_RANK} 为本节点的 RANK ${ARGS} 为代码输入的其他参数▶ 01 准备模型
接下来，你可以从 model_center 中获取预训练好的 BERT 模型，例如 bert-base-uncased。由于我们是在一个分类任务上微调 BERT 模型，所以需要在最后一层后添加一个全连接层。
from model_center.model import Bert, BertConfig
from model_center.layer import Linear
class BertModel(torch.nn.Module):
def __init__(self, config):
super().__init__()
self.bert = Bert.from_pretrained("bert-base-uncased")
self.dense = Linear(config.dim_model, 2)
bmt.init_parameters(self.dense)
def forward(self, input_ids, attention_mask):
pooler_output = self.bert(input_ids=input_ids, attention_mask=attention_mask).pooler_output
logits = self.dense(pooler_output)
return logits
config = BertConfig.from_pretrained("bert-base-uncased")
model = BertModel(config)
如果只需要 config 来构建模型，而不需要现成的预训练参数，可以参考下面的方法：
model = Bert(config)
bmt.init_parameters(model)
# bmt.load(model, "your/path/to/pytorch_model.pt")
▶ 02 准备数据集
下一步是准备数据集，用于训练和验证模型。这里，我们使用 SuperGLUE benchmark 中的 BoolQ 数据集。你需要下载该数据集，并将解压后的文件夹放在 your_path_to_dataset 路径下。
from model_center.dataset import DistributedDataLoader
from model_center.tokenizer import BertTokenizer
tokenizer = BertTokenizer.from_pretrained("bert-base-uncased")
splits = ['train', 'dev']
dataset = {}
for split in splits:
dataset[split] = DATASET['BoolQ']('your_path_to_dataset', split, bmt.rank(), bmt.world_size(), tokenizer, max_encoder_length=512)
batch_size = 64
train_dataloader = DistributedDataLoader(dataset['train'], batch_size=batch_size, shuffle=True)
dev_dataloader = DistributedDataLoader(dataset['dev'], batch_size=batch_size, shuffle=False)
▶ 03 模型训练
现在，在设置优化器、学习率调整策略和损失函数后，就可以开始训练模型了！在示例代码（https://modelcenter.readthedocs.io/en/latest/notes/quickstart.html）中，我们训练 BERT 模型5轮，并且在每轮训练结束后验证模型的性能。
optimizer = bmt.optim.AdamOffloadOptimizer(model.parameters())
lr_scheduler = bmt.lr_scheduler.Noam(
optimizer,
start_lr = 1e-5,
warmup_iter = 100,
end_iter = -1)
loss_func = bmt.loss.FusedCrossEntropy(ignore_index=-100)
for epoch in range(5):
model.train()
for data in train_dataloader:
input_ids = data['input_ids']
attention_mask = data['attention_mask']
labels = data['labels']
optimizer.zero_grad()
# model forward
logits = model(input_ids, attention_mask)
# calculate loss
loss = loss_func(logits.view(-1, logits.shape[-1]), labels.view(-1))
# scale loss to avoid precision underflow of fp16
loss = optimizer.loss_scale(loss)
# model backward
loss.backward()
# clip gradient norm
_ = bmt.optim.clip_grad_norm(optimizer.param_groups, max_norm=10.0, scale = optimizer.scale, norm_type = 2)
bmt.optim_step(optimizer, lr_scheduler)
▶ 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）中给出的实现方式。