大模型技术实践（五）｜支持千亿参数模型训练的分布式并行框架-CSDN博客

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

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

01

数据并行与模型并行

1.1  数据并行

数据并行可分为DPData Parallelism数据并行、DDPDistributed Data Parallelism分布式数据并行、ZeROZero Redundancy Optimizer零冗余优化器3种方式。

其中DP的做法是在每个设备上复制1份模型参数在每个训练步骤中一个小批量数据均匀地分配给所有数据并行的进程以便每个进程在不同子集的数据样本上执行前向和反向传播并使用跨进程的平均梯度来局部更新模型。
DP通常用参数服务器Parameters Server来实现其中用作计算的GPU称为Worker用作梯度聚合的GPU称为ServerServer要和每个Worker传输梯度那么通信的瓶颈就在Server上。受限于通信开销大和通信负载不均的因素DP通常用于单机多卡场景。
DDP为了克服Server的带宽制约计算效率采用Ring-AllReduce的通信方式。在Ring-AllReduce中所有的GPU形成1个环形通信拓扑在1轮环形传递后所有GPU都获得了数据的聚合结果带宽利用率高。在Ring-AllReduce的每1轮通信中各个GPU都与相邻GPU通信而不依赖于单个GPU作为聚合中心有效解决了通信负载不均的问题。DDP同时适用于单机和多机场景。
ZeRO采用的是数据并行结合张量并行的方式后面将详细展开讲解。

1.2 模型并行

模型并行MPModel Parallelism可分为流水线并行PPPipeline Parallelism和张量并行TPTensor Parallesim都是解决当GPU放不下一个模型时而将模型拆分到不同的GPU上的方法。
 

流水线并行

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

张量并行

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

02

分布式框架

2.1 Megatron

2019年英伟达发布的Megatron是一个基于PyTorch的分布式训练框架实现了一种简单高效的层内模型并行方法TP是切分矩阵的形式实现的可以训练具有数十亿参数的Transformer模型。Megatron不需要新的编译器或库更改可以通过在PyTorch中插入几个通信操作来完全实现。当然Megatron目前支持TP、PP、SPSequence 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-DPZeRO-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个GPUNd = 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 和FSDPPyTorch 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-LMhttps://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 Facehttps://huggingface.co/docs/accelerate/index