ModelCenter

高效低资源的大模型实现

总览 • 文档 • 安装 • 快速上手 • 模型支持 • English

总览

ModelCenter 基于 OpenBMB/BMTrain 实现了一系列经典的预训练语言模型。 ModelCenter 在模型实现上的宗旨是高效、低资源与高可用性, 并且能够支持分布式的训练.

我们的主要优势有：

易用性。相比 Deepspeed, Megatron, 我们拥有更好更灵活的封装，且配置 python 环境容易, 训练代码与 PyTorch 风格统一。
更高效的显存利用。模型占用显存较大时，可能会导致 GPU 的计算能力未被充分使用时显存占用就已经跑满。我们的实现可以将显存占用降低数倍，进而使用更大的 batch-size 对 GPU 的计算能力进行更充分的利用。
低资源的高效分布式训练。在 OpenBMB/BMTrain 的支持下，我们能够将 ZeRO3 的优化轻易地扩展至各大预训练语言模型，并在分布式训练的通信和调度上作出优化。

文档

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

安装

1. 用 pip 安装 (推荐)

$ 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 来构建模型，而不需要现成的预训练参数，可以参考下面的方法：

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轮，并且在每轮训练结束后验证模型的性能。

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)

        # step all optimizers inside 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} 为本节点的 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 官方文档。

模型支持

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[paper]. 我们支持使用 Longformer.from_pretrained(identifier) 来加载下列模型：
- lawformer
GLM[paper]. 我们支持使用 GLM.from_pretrained(identifier) 来加载下列模型：
- glm-10b-zh
ViT[paper]. 我们支持使用 ViT.from_pretrained(identifier) 来加载下列模型：
- vit-base-patch16-224
LLaMA[paper]. 通过 transfer/hugLLaMa_bmtrainLLaMa.py 转换 checkpoint。

运行性能

你可以在 OpenBMB/BMTrain 仓库中找到更多的性能测试效果.

开源社区

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

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

QQ群: 735930538
官方网站: https://www.openbmb.org
微博: http://weibo.cn/OpenBMB
Twitter: https://twitter.com/OpenBMB

开源许可

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

References

Footnotes

CPM: A Large-scale Generative Chinese Pre-trained Language Model. Zhengyan Zhang, Xu Han, Hao Zhou, Pei Ke, Yuxian Gu, Deming Ye, Yujia Qin, Yusheng Su, Haozhe Ji, Jian Guan, Fanchao Qi, Xiaozhi Wang, Yanan Zheng, Guoyang Zeng, Huanqi Cao, Shengqi Chen, Daixuan Li, Zhenbo Sun, Zhiyuan Liu, Minlie Huang, Wentao Han, Jie Tang, Juanzi Li, Xiaoyan Zhu, Maosong Sun. ↩
CPM-2: Large-scale Cost-efficient Pre-trained Language Models. Zhengyan Zhang, Yuxian Gu, Xu Han, Shengqi Chen, Chaojun Xiao, Zhenbo Sun, Yuan Yao, Fanchao Qi, Jian Guan, Pei Ke, Yanzheng Cai, Guoyang Zeng, Zhixing Tan, Zhiyuan Liu, Minlie Huang, Wentao Han, Yang Liu, Xiaoyan Zhu, Maosong Sun. ↩
BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding. Jacob Devlin, Ming-Wei Chang, Kenton Lee and Kristina Toutanova. ↩
RoBERTa: A Robustly Optimized BERT Pretraining Approach Yinhan Liu, Myle Ott, Naman Goyal, Jingfei Du, Mandar Joshi, Danqi Chen, Omer Levy, Mike Lewis, Luke Zettlemoyer, Veselin Stoyanov. ↩
T5: Exploring the Limits of Transfer Learning with a Unified Text-to-Text Transformer Colin Raffel and Noam Shazeer and Adam Roberts and Katherine Lee and Sharan Narang and Michael Matena and Yanqi Zhou and Wei Li and Peter J. Liu. ↩
GPT2: Language Models are Unsupervised Multitask Learners. Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, and Ilya Sutskever. ↩
GPT-J (from EleutherAI) released in the repo mesh-transformer-jax by Ben Wang and Aran Komatsuzaki. ↩

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

README-ZH.md

README-ZH.md

ModelCenter

最新动态

总览

文档

安装

1. 用 pip 安装 (推荐)

2. 从源代码安装

快速上手

1. 初始化 BMTrain 后端

2. 准备模型

3. 准备数据集

4. 训练模型

5. 运行代码

torch.distributed.launch

torchrun

模型支持

运行性能

开源社区

开源许可

References

Files

README-ZH.md

Latest commit

History

README-ZH.md

File metadata and controls

ModelCenter

最新动态

总览

文档

安装

1. 用 pip 安装 (推荐)

2. 从源代码安装

快速上手

1. 初始化 BMTrain 后端

2. 准备模型

3. 准备数据集

4. 训练模型

5. 运行代码

torch.distributed.launch

torchrun

模型支持

运行性能

开源社区

开源许可

References

Footnotes