ModelCenter

<div align="center"> <h1><img src="https://yellow-cdn.veclightyear.com/835a84d5/e7a04945-8bb5-41b8-8c9c-9009c15f3afd.png" height="32px"/> ModelCenter</h1>

高效低资源大型模型实现

</div> <p align="center"> <a href="#概述">概述</a> • <a href="#文档">文档</a> • <a href="#安装">安装</a> • <a href="#快速开始">快速开始</a> • <a href="#支持的模型">支持的模型</a> • <a href="./README-ZH.md" target="_blank">简体中文</a> </p> <p align="center"> <a href='https://modelcenter.readthedocs.io/en/latest/?badge=latest'> <img src='https://readthedocs.org/projects/modelcenter/badge/?version=latest' alt='文档状态' /> </a> <a href="https://github.com/OpenBMB/ModelCenter/releases"> <img alt="GitHub release (最新版本包括预发布)" src="https://img.shields.io/github/v/release/OpenBMB/ModelCenter?include_prereleases"> </a> <a href="https://github.com/OpenBMB/ModelCenter/blob/main/LICENSE"> <img alt="GitHub" src="https://img.shields.io/github/license/OpenBMB/ModelCenter"> </a> </p>

概述

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转换检查点。