大模型技术实践(五)|支持千亿参数模型训练的分布式并行框架

在上一期的大模型技术实践中，我们介绍了增加式方法、选择式方法和重新参数化式方法三种主流的参数高效微调技术（PEFT）。微调模型可以让模型更适合于我们当前的下游任务，但当模型过大或数据集规模很大时，单个加速器（比如GPU）负载和不同加速器之间的通信是值得关注的问题，这就需要关注并行技术。

并行化是大规模训练中训练大型模型的关键策略，本期内容UCloud将为大家带来“加速并行框架”的技术科普和实践。在学习Accelerate、DeepSpeed、Megatron加速并行框架之前，我们先来了解一下数据并行和模型并行。

01

数据并行与模型并行

1.1 数据并行

数据并行可分为DP（Data Parallelism，数据并行）、DDP（Distributed Data Parallelism，分布式数据并行）、ZeRO（Zero Redundancy Optimizer，零冗余优化器）3种方式。

其中，DP的做法是在每个设备上复制1份模型参数，在每个训练步骤中，一个小批量数据均匀地分配给所有数据并行的进程，以便每个进程在不同子集的数据样本上执行前向和反向传播，并使用跨进程的平均梯度来局部更新模型。

DP通常用参数服务器（Parameters Server）来实现，其中用作计算的GPU称为Worker，用作梯度聚合的GPU称为Server，Server要和每个Worker传输梯度，那么通信的瓶颈就在Server上。受限于通信开销大和通信负载不均的因素，DP通常用于单机多卡场景。

DDP为了克服Server的带宽制约计算效率，采用Ring-AllReduce的通信方式。在Ring-AllReduce中，所有的GPU形成1个环形通信拓扑，在1轮环形传递后，所有GPU都获得了数据的聚合结果，带宽利用率高。在Ring-AllReduce的每1轮通信中，各个GPU都与相邻GPU通信，而不依赖于单个GPU作为聚合中心，有效解决了通信负载不均的问题。DDP同时适用于单机和多机场景。

ZeRO采用的是数据并行结合张量并行的方式，后面将详细展开讲解。

1.2 模型并行

模型并行（MP，Model Parallelism）可分为流水线并行（PP，Pipeline Parallelism）和张量并行（TP，Tensor Parallesim），都是解决当GPU放不下一个模型时，而将模型拆分到不同的GPU上的方法。

流水线并行

如上图所示，PP将模型在Layer层面上进行水平切分，不同切分部分在不同GPU上运行，并使用微批处理（Micro-Batching，在1个Batch上再划分得到Micro-Batch）和隐藏流水线泡沫（Pipeline Bubble，GPU空转的部分）。由于水平分割和微批处理，模型功能（如权重共享和批量归一化）的实现都很困难。

张量并行

如上图所示，TP在模型的Layer内部进行切分，所有GPU计算所有层的不同部分在单个节点之外无法高效扩展，这是由于细粒度计算和昂贵的通信所致。

02

分布式框架

2.1 Megatron

2019年英伟达发布的Megatron是一个基于PyTorch的分布式训练框架，实现了一种简单高效的层内模型并行方法（TP，是切分矩阵的形式实现的），可以训练具有数十亿参数的Transformer模型。Megatron不需要新的编译器或库更改，可以通过在PyTorch中插入几个通信操作来完全实现。当然Megatron目前支持TP、PP、SP（Sequence Parallelism）和Selective Activation Recomputation，此处对TP进行讲解。

切分矩阵的方式

切分矩阵的方式可按行或者按列，其目的是为了当1个模型无法完整放入1个GPU时，让这个模型能塞到多个GPU中[1]。

1.对于MLP层的切分

上图中的f和g分别表示2个算子，每个算子都包含一组Forward + Backward操作，也就是前向传播和后向传播操作。具体[2]可表示为：

左侧为输入数据X按列切开，将权重矩阵A按行切开。右侧为将输入数据X复制到2个GPU上，将权重矩阵A按列切开。当然，这里是2个GPU的例，同理可类推到多个GPU的情况。

2.对于Multi-Head Attention层的切分

Multi-Head Attention，也就是多头注意力，其结构如下图所示：

其每个Head本身就是独立计算，再将结果Concat起来，就可以把每个Head的权重放到1个GPU上。当然，1个GPU上可以有多个Head。

如上图所示，先把输入数据X复制到多个GPU上，此时每个Head分别在1个GPU上。对每个Head的参数矩阵Q、K、V按列切开，切分原理、算子f、算子g与MLP层切分章节中的描述一致。之后，各个GPU按照Self-Attention的计算方式得到结果，再经过权重B按行切开的线性层计算。

MLP层和Multi-Head Attention层的通信量

1.MLP层的通信量

由上述章节可知，MLP层进行Forward和Backward操作都有一次All-Reduce操作。All-Reduce操作包括Reduce-Scatter操作和All-Gather操作，每个操作的通讯量都相等，假设这2个操作的通讯量都为φ，则进行一次All-Reduce的通讯量为2φ，MLP层的总通讯量为4φ。

2.Multi-head Attention层的通信量

由上图可知，Self-Attention层在Forward和Backward中都要做一次All-Reduce，总通讯量也是。

张量并行与数据并行的结合

1.MP+DP混合的结构

由上图可知，中间每个虚线框表示一台机器，每台机器有8个GPU，合计512个GPU。同一个机器内的1个或多个GPU构成1个模型并行组，不同机器上同一个位置的GPU构成1个数据并行组，图中有8路模型并行组和64路数据并行组。

2.MP与MP+DP的通信量对比

左图是MP模式，可以看到单个GPU的计算效率为100%。随着GPU个数的增加，通信量增大，GPU的计算效率有一定的下降。右图是MP+DP模式，64个GPU（可理解为64台机器，1台机器1个GPU，此时相当于DP模式）的计算效率有96%之高，是由于DP在计算梯度时，可一边继续往下做Backward，一边把梯度发送出去和DP组内其他GPU做All-Reduce。同理，当GPU个数增多，GPU的计算效率也会下降。

2.2 DeepSpeed

2020年微软发布了分布式训练框DeepSpeed和一种新型内存优化技术ZeRO-1，极大地推进了大模型训练的进程。后续，微软又陆续推出ZeRO-2、ZeRO-3技术等，ZeRO这3个阶段称为ZeRO-DP（ZeRO-Powered Data Parallelism）。另外，DeepSpeed还支持自定义混合精度训练处理，一系列基于快速CUDA扩展的优化器，ZeRO-Offload到CPU和磁盘/NVMe。DeepSpeed支持PP、DP、TP这3种范式，支持万亿参数模型的训练。

其中，ZeRO-DP用来克服数据并行性和模型并行性的限制，通过将模型状态（参数、梯度和优化器状态）在数据并行进程之间进行分片，使用动态通信调度实现在分布式设备之间共享必要的状态。ZeRO-R技术可减少剩余的显存消耗。在模型训练过程中产生的激活值（Activations）、存储中间结果的临时缓冲区、显存碎片，我们称之为剩余状态。

ZeRO-DP

1.ZeRO-DP的3个阶段

ZeRO-DP[3]的3个阶段，可在参数zero_optimization中设置。比如：

{
"zero_optimization": { 
"stage": stage_number, 
}
}

其中，stage_number可写1、2、3，当然也可以写0。ZeRO-1只对优化器状态进行分片，ZeRO-2在ZeRO-1的基础上还对梯度分片，ZeRO-3在ZeRO-2的基础上还对模型参数分片。当stage_number为0时，不做任何分片处理，此时相当于DDP。

在ZeRO-DP优化的3个阶段下，1个模型状态在各个GPU上内存消耗情况，如下图所示：

其中Ψ表示模型大小（参数数量），K表示优化器状态的内存倍增器。Nd表示数据并行度，也就是GPU的个数。

1. Pos对应优化器状态分片，也就是ZeRO-1，内存减少4倍。

2. Pos+g对应添加梯度分片，也就是ZeRO-2， 内存减少8倍。

3. Pos+g+p对应模型参数分片，内存减少与数据并行性程度Nd线性相关。

论文中提到，使用64个GPU（Nd = 64）将导致内存减少64倍、通信量略有增加约50%。

2.ZeRO-DP的通信量

由上图可知，ZeRO-DP在使用Pos和Pg时不会增加额外的通信，同时可以实现高达8倍的内存减少。使用PP，除了Pos和Pg之外，ZeRO-DP最多会引入1.5倍的通信开销，同时进一步降低内存占用Nd倍。

ZeRO-R

1.ZeRO-R的思路

将Activations Checkpoints分片到各个GPU上，并使用All-Gather操作按需重构它们，消除了模型并行中的内存冗余。对于非常大的模型，甚至可以选择将激活分区移动到CPU内存。

2.ZeRO-R的通信量

分片Activations Checkpoints（记为Pa）的通信量权衡取决于模型大小、Checkpoints策略和模型并行策略。对于模型并行来说，Pa的总通信开销不到原始通信量的10％。

当模型并行与数据并行结合使用时，Pa可以用来将数据并行的通信量降低一个数量级，代价是模型并行的通信量增加了10％，并在数据并行通信成为性能瓶颈时显著提高效率。

2.3 Accelerate

Accelerate[4]由Huggingface于2021年发布，是一个适用于Pytorch用户的简单封装的库，其简化了分布式训练和混合精度训练的过程。Accelerate基于torch_xla和torch.distributed，只需要添加几行代码，使得相同的PyTorch代码可以在任何分布式配置下运行！简而言之，它使得大规模训练和推理变得简单、高效和适应性强。Accelerate 可与DeepSpeed、Megatron-LM 和FSDP（PyTorch Fully Sharded Data Parallel）等扩展一起使用。

2.4 小结

Accelerate更加稳定和易于使用，适合中小规模的训练任务。DeepSpeed和Megatron支持更大规模的模型。通过Accelerate库，可轻松在单个设备或多个设备上结合Megatron、DeepSpeed进行分布式训练。当然，Megatron、Deepspeed 也可以结合使用，比如Megatron-DeepSpeed，这是NVIDIA的Megatron-LM的DeepSpeed版本。

03

在UCloud云平台选择A800

进行Baichuan2大模型的微调实验

首先参照UCloud文档中心（https://docs.ucloud.cn），登录UCloud控制台（https://console.ucloud.cn/uhost/uhost/create）。

在UCloud云平台上创建云主机，选择显卡为A800，配置如下：

实验项目代码获取：https://github.com/hiyouga/LLaMA-Efficient-Tuning

模型下载方式：Git Clone https://huggingface.co/baichuan-inc/Baichuan2-13B-Chat

数据集来源于：

1.https://huggingface.co/datasets/neuclir/csl/viewer/default/csl 包含多篇论文摘要、标题、关键词、学科， 可用来做标题总结，关键词提取，学科分类。

2.https://huggingface.co/datasets/hugcyp/LCSTS/viewer/default/train?p=24002微博新闻缩写成标题。

3.https://huggingface.co/datasets/csebuetnlp/xlsum 中文/英文 BBC新闻、摘要和标题。

本实验的微调任务是从一段文字中提取出概括性的标题。选8万条数据去训练，1000条数据作为测试集，数据样例为：

[
{
"instruction": "为以下内容生成一个概括性的标题：\n",
"input": "随着IT巨头们将触角逐渐伸向移动中的汽车产品，鲶鱼效应推动了“车联网”的迅速发展，领军人物苹果与谷歌率先在这场跨界之争中形成了针锋相对的格局，继手机屏幕、电视屏幕等领域之后，又展开了新一轮的“入口”抢滩战。",
"output": "汽车联网苹果已获进展谷歌紧随其后开发",
"source": "LCSTS"
},
{
"instruction": "为以下内容生成一个概括性的标题：\n",
"input": "调查显示，近半数俄罗斯人一年内未读过一本书。曾在苏联时期受到极度喜爱的书籍为何在今天遭遇几乎无人问津的尴尬境地，实体书店又将如何应对危机？《透视俄罗斯》记者带您一探究竟。",
"output": "阅读形式多元化俄罗斯或将告别纸质图书时代",
"source": "LCSTS"
}
]

根据以下配置去执行训练（此处并未做参数调优）：

deepspeed --num_gpus 8 --master_port=9901 src/train_bash.py \
--deepspeed ds_config.json \
--stage sft \
--model_name_or_path /data/text-generation-webui/models/Baichuan2-13B-Chat \
--do_train True \
--overwrite_cache False \
--finetuning_type lora \
--template baichuan2 \
--dataset_dir data \
--dataset summary_instruction_train \
--cutoff_len 1024 \
--learning_rate 1e-05 \
--num_train_epochs 5.0 \
--max_samples 100000 \
--per_device_train_batch_size 8 \
--gradient_accumulation_steps 4 \
--lr_scheduler_type cosine \
--max_grad_norm 1.0 \
--logging_steps 5 \
--save_steps 800 \
--warmup_steps 0 \
--flash_attn False \
--lora_rank 8 \
--lora_dropout 0.1 \
--lora_target W_pack \
--resume_lora_training True \
--output_dir saves/Baichuan2-13B-Chat/lora/2023-10-08-18-20-07 \
--fp16 True \
--plot_loss True

再进行测试，测试结果为：

{ "predict_bleu-4": 15.863159399999999,
"predict_rouge-1": 29.348522,
"predict_rouge-2": 10.655794799999999,
"predict_rouge-l": 26.600239000000002,
"predict_runtime": 120.571,
"predict_samples_per_second": 8.294,
"predict_steps_per_second": 1.037 }

上述结果中：

BLEU-4是一种用于评估机器翻译结果的指标；ROUGE-1、ROUGE-2和ROUGE-L是ROUGE系统中的三个常用变体，一组用于评估文本摘要和生成任务的自动评估指标；predict_runtime是模型进行预测所花费的总运行时间，单位为秒；predict_samples_per_second是模型每秒钟处理的样本数量；predict_steps_per_second是模型每秒钟执行的步骤数量。

【参考文献】

[1]《Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism》

[2] Megatron-LM：https://developer.download.nvidia.com/video/gputechconf/gtc/2020/presentations/s21496-megatron-lm-training-multi-billion-parameter-language-models-using-model-parallelism.pdf

[3]《ZeRO: Memory Optimizations Toward Training Trillion Parameter Models》

[4] accelerate Hugging Face：https://huggingface.co/docs/accelerate/index