ModelCenter
高效低资源大型模型实现
概述 • 文档 • 安装 • 快速开始 • 支持的模型 • 简体中文
最新动态
- 2023/06/13 ModelCenter 1.0.3 ModelCenter 支持 T5 的束搜索生成。
- 2023/05/28 ModelCenter 1.0.2 ModelCenter 支持 LLaMA 及其生成。
- 2023/02/28 ModelCenter 1.0.1 ModelCenter 支持 FLAN-T5(fp32)版本。
- 2022/11/21 ModelCenter 1.0.0 ModelCenter 支持 BMTrain>=0.2.0。
- 2022/07/14 ModelCenter 0.1.5 ModelCenter 支持 Mengzi、GLM、Longformer 和 KV_PLM。
- 2022/07/05 ModelCenter 0.1.3 ModelCenter 支持 mT5、T5v1.1、ViT 和 Wenzhong。
- 2022/04/27 ModelCenter 0.1.1 ModelCenter 支持 RoBERTa。
- 2022/04/06 ModelCenter 0.1.0 ModelCenter 发布首个稳定版本,修复了模型性能和 GPU 内存使用方面的一些问题。
- 2022/03/21 ModelCenter 0.0.1-beta ModelCenter 发布首个公开测试版。
概述
ModelCenter 基于 OpenBMB/BMTrain 后端实现了预训练语言模型(PLMs)。ModelCenter 支持高效、低资源、可扩展的模型使用和分布式训练。
我们的主要优势是:
- 易于使用。与 Deepspeed 和 Megatron 相比,我们有更好更灵活的代码封装和易于配置的 Python 环境,训练代码与 PyTorch 风格一致。
- 更高效的内存利用。内存占用大的模型可能在 GPU 计算能力充分利用之前就导致内存溢出(OOM)。我们的实现将内存占用减少了数倍,允许使用更大的批量大小来更有效地利用 GPU 的计算能力。
- 低资源高效分布式训练。在 OpenBMB/BMTrain 的支持下,我们能够轻松地将 ZeRO 优化扩展到任何 PLMs,并为更快的分布式训练优化了通信和时间调度。
文档
我们的文档提供了有关该软件包的更多信息。
安装
1. 从 PyPI 安装(推荐)
$ pip install model-center
2. 从源代码安装
$ git clone https://github.com/OpenBMB/ModelCenter.git
$ cd ModelCenter
$ pip install -r requirements.txt
$ python3 setup.py install
快速开始
在快速开始中,您将了解如何在分类任务上微调 BERT 模型。
1. 初始化 bmtrain 后端
首先,您需要在代码开头导入 bmtrain
并使用 bmtrain.init_distributed()
,这可以初始化分布式环境。
import bmtrain as bmt
bmt.init_distributed(seed=0)
2. 准备模型
接下来,您可以从 model_center
中简单地获取一个预训练的 BERT 模型,例如 bert-base-uncased。在分类任务上微调 BERT 时,需要在最后一层后添加一个前馈层。
import torch
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 = BertConfig.from_json_file("your/path/to/config.json")
model = Bert(config)
bmt.init_parameters(model)
# bmt.load(model, "your/path/to/pytorch_model.pt")
3. 准备数据集
下一步是准备用于训练和评估的数据集。在这里,我们使用 SuperGLUE benchmark 中的 BoolQ 数据集。您需要下载数据集并将解压后的文件夹放在 your_path_to_dataset
。
from model_center.dataset.bertdataset import 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)
4. 训练模型
现在,选择优化器、学习率调度器、损失函数,然后开始训练模型!这里我们训练 BERT 5 个 epoch,并在每个 epoch 结束时进行评估。
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)
optim_manager = bmt.optim.OptimManager(loss_scale=1024)
optim_manager.add_optimizer(optimizer, lr_scheduler)
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']
# 模型前向传播
logits = model(input_ids, attention_mask)
# 计算损失
loss = loss_func(logits.view(-1, logits.shape[-1]), labels.view(-1))
# 使用 bmt.sum_loss(loss) 收集所有分布式进程的损失信息
global_loss = bmt.sum_loss(loss).item()
# 梯度清零
optim_manager.zero_grad()
# 在反向传播之前缩放损失以避免 fp16 精度下溢
optim_manager.backward(loss)
# 梯度裁剪
grad_norm = optim_manager.clip_grad_norm(optimizer.param_groups, max_norm=10.0, scale = optimizer.scale, norm_type = 2)
# 对 optim_manager 中的所有优化器进行步进
optim_manager.step()
# 在分布式训练时仅在 rank 0 上打印信息
bmt.print_rank(
"loss: {:.4f} | lr: {:.4e}, scale: {:10.4f} | grad_norm: {:.4f} |".format(
global_loss,
lr_scheduler.current_lr,
int(optimizer.scale),
grad_norm,
)
)
# 评估模型
model.eval()
with torch.no_grad():
pd = [] # 预测
gt = [] # 真实标签
for data in dev_dataloader:
input_ids = data["input_ids"]
attention_mask = data["attention_mask"]
labels = data["labels"]
logits = model(input_ids, attention_mask)
loss = loss_func(logits.view(-1, logits.shape[-1]), labels.view(-1))
logits = logits.argmax(dim=-1)
pd.extend(logits.cpu().tolist())
gt.extend(labels.cpu().tolist())
# 从所有分布式进程收集结果
pd = bmt.gather_result(torch.tensor(pd).int()).cpu().tolist()
gt = bmt.gather_result(torch.tensor(gt).int()).cpu().tolist()
# 计算指标
from sklearn.metrics import accuracy_score
acc = accuracy_score(gt, pd)
bmt.print_rank(f"accuracy: {acc*100:.2f}")
5. 运行代码
您可以使用与 PyTorch 的分布式模块相同的启动命令运行上述代码。
根据您的 PyTorch 版本选择以下命令之一。
${MASTER_ADDR}
表示主节点的 IP 地址。${MASTER_PORT}
表示主节点的端口。${NNODES}
表示节点总数。${GPU_PER_NODE}
表示每个节点的 GPU 数量。${NODE_RANK}
表示此节点的排名。
torch.distributed.launch(更适合 torch < 1.10)
$ 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(更适用于 torch >= 1.10)
$ torchrun --nnodes=${NNODES} \
--nproc_per_node=${GPU_PER_NODE} \
--rdzv_id=1 \
--rdzv_backend=c10d \
--rdzv_endpoint=${MASTER_ADDR}:${MASTER_PORT} \
train.py
更多信息可以在文档中找到。
支持的模型
-
CPM-1[论文]。我们目前支持通过
CPM1.from_pretrained(identifier)
加载以下检查点:- cpm1-large
-
CPM-2[论文]。我们目前支持通过
CPM2.from_pretrained(identifier)
加载以下检查点:- cpm2-large
-
BERT[论文]。我们目前支持通过
Bert.from_pretrained(identifier)
加载以下检查点:- bert-base-cased
- bert-base-uncased
- bert-large-cased
- bert-large-uncased
- bert-base-chinese
- bert-base-multilingual-cased
- kv-plm
-
RoBERTa[论文]。我们目前支持通过
Roberta.from_pretrained(identifier)
加载以下检查点:- roberta-base
- roberta-large
-
T5[论文]。我们目前支持通过
T5.from_pretrained(identifier)
加载以下检查点:- t5-small
- t5-base
- t5-large
- t5-3b
- t5-11b
- t5-v1_1-small
- t5-v1_1-base
- t5-v1_1-large
- t5-v1_1-xl
- t5-v1_1-xxl
- mt5-small
- mt5-base
- mt5-large
- mt5-xl
- mt5-xxl
- mengzi-t5-base
- flan-t5-small
- flan-t5-base
- flan-t5-large
- flan-t5-xl
- flan-t5-xxl
-
GPT-2[论文]。我们目前支持通过
GPT2.from_pretrained(identifier)
加载以下检查点:- gpt2-base
- gpt2-medium
- gpt2-large
- gpt2-xl
- wenzhong-gpt2-3.5b
-
GPT-J[论文]。我们目前支持通过
GPTj.from_pretrained(identifier)
加载以下检查点:- gptj-6b
-
Longformer[论文]。我们目前支持通过
Longformer.from_pretrained(identifier)
加载以下检查点:- lawformer
-
GLM[论文]。我们目前支持通过
GLM.from_pretrained(identifier)
加载以下检查点:- glm-10b-zh
-
ViT[论文]。我们目前支持通过
ViT.from_pretrained(identifier)
加载以下检查点:- vit-base-patch16-224
-
LLaMA[论文]。通过
transfer/hugLLaMa_bmtrainLLaMa.py
转换检查点。
性能
您可以在OpenBMB/BMTrain仓库中找到更多性能指标。
社区
我们欢迎每个人按照我们的贡献指南贡献代码。
您也可以在其他平台上找到我们:
- QQ群:735930538
- 网站:https://www.openbmb.org
- 微博:http://weibo.cn/OpenBMB
- Twitter:https://twitter.com/OpenBMB
许可证
该软件包根据Apache 2.0许可证发布。