TencentPretrain 使用 DeepSpeed ZeRO-3 并行训练

李煜东

多模态预训练框架 TencentPretrain 现已支持 LoRA 和 DeepSpeed ZeRO-3 并行，单机 8 * A100 可训练千亿参数量大模型。

DeepSpeed 是一个GPU分布式训练加速框架，它的核心在于通过显存优化来训练更大的模型，DeepSpeed的详细原理介绍可以参考这篇文章。在之前版本的 TencentPretrain 中已经支持 ZeRO stage1和2，最近我们的更新添加了对 ZeRO stage 3 的支持。经过测试，单机 8 * A100 40G 训练 LLaMA 30B 在 512 seq length 时能达到 256 batch size，大幅度提高了我们训练大模型的上限。

基于 ZeRO-3，我们训练了中文 LLaMA 基础模型和 ChatLLaMA 对话模型，数据、代码和模型权重均公开，仓库正在持续更新中：

本文主要介绍如何在 TencentPretrain 中使用 DeepSpeed ZeRO-3 训练和推理大模型，以及如何使用 LoRA 进一步降低计算开销快速微调模型。

DeepSpeed ZeRO-3

在预训练阶段，使用 --enable_zero3 参数以及对应的配置文件（例如 deepspeed_zero3_config.json）启动 ZeRO-3 训练 LLaMA 模型。数据集的制作方法参考这里

deepspeed pretrain.py --deepspeed --deepspeed_config models/deepspeed_config.json \
                      --pretrained_model_path models/llama-7b.bin \
                      --dataset_path $OUTPUT_DATASET_PATH --spm_model_path $LLaMA_PATH/tokenizer.model \
                      --config_path models/llama/7b_config.json \
                      --output_model_path models/llama_zh_7b \
                      --world_size 8 --data_processor lm --enable_zero3

使用DeepSpeed进行语言模型推理：

deepspeed scripts/generate_lm_deepspeed.py --load_model_path models/llama-7b.bin --spm_model_path $LLaMA_PATH/tokenizer.model \
                               --deepspeed_config models/deepspeed_zero3_config.json \
                               --test_path beginning.txt --prediction_path generated_sentence.txt \
                               --config_path models/llama/7b_config.json  --enable_zero3

LoRA

本次更新还加入了对 Low-Rank Adaptation (LoRA) 的支持。使用 --use_lora 参数启用 LoRA 训练。LoRA的具体参数可以通过 --lora_r、 --lora_alpha、--lora_dropout 来配置。当启用 use_lora 时，模型仅保存 LoRA 部分的参数。

python3 pretrain.py --pretrained_model_path models/llama-7b.bin \
                    --dataset_path $OUTPUT_DATASET_PATH --spm_model_path $LLaMA_PATH/tokenizer.model \
                    --config_path models/llama/7b_config.json \
                    --output_model_path models/llama-7b_loar.bin \
                    --world_size 8  --learning_rate 1e-4 \
                    --data_processor lm --use_lora --lora_dropout 0.05

加载预训练模型权重和 LoRA 权重进行增量训练：

python3 pretrain.py -pretrained_model_path models/llama-7b.bin \