Home

Awesome

CPM-Pretrain

版本更新记录

为了促进中文自然语言处理研究的发展,本项目提供了大规模预训练语言模型的预训练代码。项目主要基于DeepSpeed、Megatron实现,可以支持数据并行、模型加速、流水并行的代码。

安装

1、首先安装pytorch等基础依赖,再安装APEX以支持fp16。

pip install torch
git clone https://github.com/NVIDIA/apex
cd apex
pip install -v --no-cache-dir --global-option="--cpp_ext" --global-option="--cuda_ext" ./

2、安装DeepSpeed,这部分更多细节请参照DeepSpeed项目安装说明。

pip install deepspeed

3、本项目也需要TDS方可运行,此部分更多细节可以见TDS。TDS插件只需要将代码拷贝到当前项目的根目录下,再使用下列代码载入DeepSpeed即可(当前项目里此步骤已经完成)。

import tds as deepspeed

模型

当前项目主要支持GPT-2、T5两类模型,未来会有更多模型放出。

预训练与设置脚本

1、多机多卡设置

#每台机器的显卡数量
GPUS_PER_NODE=4

#支持机器之间通信的IP地址
MASTER_ADDR=localhost

#支持机器之间通信的端口
MASTER_PORT=8888

#本次训练一共涉及几台机器
NNODES=1

#当前机器的序号,编号从0开始
NODE_RANK=0

2、数据并行、模型并行、流水并行设置

# 整体模型划分为几层,进行流水并行
pp_size=2

# 每层模型划分为几块,进行模型并行
mp_size=2

3、设置模型尺寸

#层数
NLAYERS=2

#输入层大小
NHIDDEN=128

#单张卡每次接触到的batch大小
BATCHSIZE=4

#梯度累积步数
GAS=16

4、DeepSpeed脚本设置,为DeepSpeed设置训练的流水并行、混合精度等细节,只有在流水并行被启用时,此脚本方被使用。如果pp_size=0或pp_size=1,意味着流水并行不被激活,此时无论脚本里的设置为何,均不会发生作用,代码只会使用Megatron-LM进行模型并行。

config_json="ds_config_t5.json"

脚本形式:
{
    "train_batch_size": 128, 
    "train_micro_batch_size_per_gpu": 4, 
    "steps_per_print": 10, 
    "gradient_clipping": 1.0, 
    "fp16": {
        "enabled": true, 
        "loss_scale": 0, 
        "loss_scale_window": 1000, 
        "hysteresis": 2, 
        "initial_scale_power": 16
    }, 
    "zero_allow_untested_optimizer": true, 
    "wall_clock_breakdown": true
}

注意:脚本里的 train_batch_size = BATCHSIZE * GAS * GPUS_PER_NODE * NNODES / pp_size / mp_size

5、也可以自己自由设置训练参数,以run_t5.sh为例

    --model-parallel-size 2 \
    --pipe-parallel-size 2 \
    --num-layers 12\
    --hidden-size 128 \
    --kv-hidden-size 16 \
    --ff-hidden-size 256 \
    --num-attention-heads 8 \
    --enc-seq-length 1024\
    --dec-seq-length 384\
    --max-position-embeddings 1024 \
    --batch-size 4\
    --gas 16 \
    --lr 1.5e-4 \
    --lr-decay-style cosine \
    --min-lr 1.0e-5 \
    --weight-decay 1e-2 \
    --clip-grad 1.0 \
    --warmup 0.01 \

6、运行脚本

bash run_t5.sh 或 bash run_gpt2.sh

实现自己的模型

我们以pretrain_gpt2.py为例子,介绍下如何实现自己的模型

引用

如果您使用了我们的代码,请您引用下面的文章。

@article{cpm-v1,
  title={CPM: A Large-scale Generative Chinese Pre-trained Language Model},
  author={Zhang, Zhengyan and Han, Xu, and Zhou, Hao, and Ke, Pei, and Gu, Yuxian and Ye, Deming and Qin, Yujia and Su, Yusheng and Ji, Haozhe and Guan, Jian and Qi, Fanchao and Wang, Xiaozhi and Zheng, Yanan and Zeng, Guoyang and Cao, Huanqi and Chen, Shengqi and Li, Daixuan and Sun, Zhenbo and Liu, Zhiyuan and Huang, Minlie and Han, Wentao and Tang, Jie and Li, Juanzi and Sun, Maosong},
  year={2020}
}