腾讯云NLP大模型预训练最佳实践

一、TI-Deepspeed介绍

根据当前人工智能的趋势，越大的自然语言模型可以提供越好的准确性，目前GPT-3的模型参数达到175B。但是由于成本、时间和代码集成的障碍，较大的模型难以训练。

模型参数量

微软开源了Deepspeed深度学习训练优化库，它通过优化transformer kernel性能、极致的显存优化及节省、提升模型scale能力等多个层面对大模型训练做了详细的分析以及极致的性能优化，已经成为了超大NLP模型预训练的“利器“。

通过公司内外团队及客户交流发现，越来越多的技术团队、深度学习工程师、高校科研工作者对NLP大模型预训练有需求，腾讯云作为国内主要云厂商之一，需要在NLP领域迅速补齐短板，为公司内外部客户赋能、提速。据此腾讯云TI平台团队在对Deepspeed调研和实践的基础上，从性能和易用性两方面对Deepspeed框架进行了相关优化，并根据NLP大模型不同的参数规模沉淀出了完整且高性能的分布式训练方案。

腾讯云TI平台团队旨在通过“一套框架”+“三套最佳实践”，提供一套完整的：GPU机器+需求带宽+分布式训练的完整解决方案，与IAAS团队合作，致力于更好的服务外部有NLP预训练需求的客户。

二、TI-Deepspeed框架简介

1、TI-Deepspeed单机性能提升

1）transformer kernel优化

性能优化结果

2）optimizer算子优化

TI-Deepspeed对优化器也做了非常多的优化，如adam，adamw，lamb优化器等。以adamw优化器为例，torch原生的adamw对每个要更新的tensor，都要去launch对应的m，v， weight cuda kernel执行对应计算，导致需要频繁的launch kernel，如果tensor比较小的话，kernel执行时间要比kernel launch时间要短，实际timeline观察optimizer step阶段行为显示也是如此。

TI-Deepspeed通过算子融合优化，采用multi_tensor_apply逻辑，把所有tensor的m/v/weight kernel分别加到一个vector_list里面形成m/v/weight的的vector list，只需要launch几次kernel，即可完成weight更新。

2、TI-Deepspeed显存节省技术

1）模型显存占用分析

模型的显存占用主要包括两个部分，一部分为Model States，主要包括优化器的状态和梯度以及参数等。另一部分为其他memory占用，如1）activation memory，用于保存中间结果，backward计算用到；2）临时buffer；3）临时buffer的频繁创建和释放导致的内存碎片，导致明明还有显存但是cudamalloc申请失败。

我们以nlp领域常用的优化器adamw结合混合精度训练为例，分析一下model states的显存占用。

针对FP16精度下进行权重更新的问题，apex是以FP32 维护和更新权重。训练过程中，将各层权重保存为FP32格式(FP32 Master Weights)，每次迭代时，制作这些权重的FP16副本并使用它们用于前向计算和反向计算，更新时将梯度再转换为FP32并用于更新FP32的权重。使用这种方法可以解决权重相较于其更新值过大的问题。因此Adamw优化器保存了每个参数的 fp32 master weight, momentum, and variance ，同时保存了每个参数的fp16 的weights和gradients。

以一个83亿大小的模型为例，model states的显存占用需要133GB，远大于V00机型32GB的显存上限。

2）ZeRO优化器显存节省技术

开源Deepspeed框架对显存占用做了优化，其核心亮点就是“ZeRO”零冗余优化器，TI-Deepspeed对ZeRO做了深入研究并进行吸收。

zero优化器主要分为stage1，stage2，stage3。stage1对optimizer states做partition，避免每个数据并行DP进程保存一份完整的optimizer states。stage2对optimizer states和gradients做partition，stage3对optimizer states、gradients、weights都做partition。

优化器显存优化

zero优化器相当于对实际负责参数更新的optimizer的上层封装，zero优化器负责对相关params做partition，负责维护相关通信，内部调用底层optimizer执行实际的参数更新。其中stage1和stage2的通信思路为：每个数据并行进程负责通过reduce-scatter规约部分梯度，然后更新这部分梯度对应的权重。反向传播完毕后，所有数据并行进程通过allgather获取全量权重信息。因此stage1与stage2并没有为了节省显存而增加通信量。stage3的通信思路为：每个数据并行进程负责通过reduce-scatter规约部分梯度，然后更新这部分梯度对应的权重。由于每个数据并行进程只保留部分权重，因此在实际forward中进行allgather操作获取当前submodule需要的权重信息，因此stage3的通信量为stage2/stage1的1.5倍。

值得注意的是，TI平台团队通过对zero优化器的详细分析，发现目前的zero优化显存节省方案并不支持lamb优化器，原因是它对所有梯度tensor做了flatten后，根据DP进程数做平均切分，不再维护每个weights 的shape信息，这一点在深度学习开发过程中要注意，目前zero优化支持NLP常用的adam、adamw优化器。

3）其他显存节省技术

activation checkpointing作为大模型训练节省显存的常用手段，TI-Deepspeed在此基础上做了更多的优化，包括对activation checkpointing做了partition，使得每个数据并行的进程只保存了一部分activation checkpointing，只有用到的时候通过allgather去获取。

通过分配固定size的临时buffer，避免模型增大时，临时buffer无限增大；此外对碎片化内存进行管理，尽可能把耗费内存的变量存储到一块连续buffer中。

3、TI-Deepspeed多机可扩展性优化

TI-Deepspeed结合腾讯云iaas基础设施对开源deepspeed框架多机可扩展性做了深度优化，也从实践和调研中积累了如何在腾讯云机器上更好的发挥框架性能优势的经验。

1）TI-Deepspeed+Ti-Horovod结合通信优化

开源Deepspeed框架中，不使用zero优化器时，模型的反向计算和梯度通信是串行的，在带宽有限或者模型大通信量场景下，梯度通信的开销十分巨大，由于反向计算和梯度通信串行，导致梯度通信无法被反向计算隐藏，TI-Deepspeed针对这种情况进行了优化，使得反向计算和梯度通信可以并行进行。同时在梯度通信中，引入了Ti-Horovod。

Ti-Horovod是腾讯云TI平台团队结合腾讯云底层基础设施在开源Horovod基础上定制化优化过的分布式训练通信优化库，TI-Horovod包含以下优化特性:

1. 在保留原生horovod的易用性上，增加了更好的性能通信方式，HorovodAllgather支持ncclAllgather通信
2. 实现了2D-allreduce通信，相对于全局allreduce，能够更加充分的利用云上带宽
3. 实现了top-k梯度压缩通信，进一步减小通信量，支持精度补偿，降低对精度影响
4. 针对tf框架，增加了多种梯度融合方式，比horovod原生的梯度融合机制更加高效
5. 支持fp16梯度通信和混合精度训练
6. 支持梯度累积

2）TI-Deepspeed大模型参数通信优化

1、大模型参数多轮通信allgather

随着深度学习模型复杂度和数据集规模的增大，计算效率成为了不可忽视的问题，GPU 显存一直是训练先进深度学习模型的最大瓶颈——大规模训练经常会碰到模型参数太多，显存无法容纳的情况。针对各种情况TI-Deepspeed自动对大模型通信参数进行切分，通过for循环的多次调用完成多轮发送，通过allgather_bucket_size控制每轮通信的参数量。

2、AllGather优化，规避内存不对齐情况导致的通信性能差问题

Nvidia的nccl作为业内最成熟的GPU通信库并没有做到各种case场景下的极致的性能优化，实际使用中的一个问题是，NcclAllgather接口针对某些数据包的大小，通信效率十分低下，而gpu通信效率直接影响了训练速度，带来的是高额的成本支出。

TI-Deepspeed通过详细的问题定位和论证发现，某些数据包的大小，通信效率十分低下是由于内存不对齐导致，即通信size不是16的整数倍。TI-Deepspeed通过封装nccl的allgather接口，在封装接口中对发送size进行判断，若发送size不是16的整数倍，则对发送buffer进行padding补零，新建pad size大小的zero tensor，将该tensor与原发送tensor做concat。

3）cpu offload增大batchsize，隐藏通信开销

cpu offload可以和zero优化器的stage2和stage3使用。cpu offload 保存了zero优化器对相关model states的切分，只是把对应model states都offload到了cpu 内存中。在反向计算过程中，梯度经过reduce scatter规约push到cpu上，在cpu上执行weight的更新后再move到gpu中。cpu offload利用了单独的流来实现计算和cpu<->gpu之间通信的overlap，保证训练效率。

cpu offload技术

Deepspeed专门提供了经过优化的cpu adam optimizer，使用cpu offload时，可以使用该optimizer提高训练效率。

经过在腾讯云机器上实际测试，cpu offload在超过百亿级，接近千亿规模的模型上发挥的性能优势更大，通过把model states offload到cpu memory上，节省了GPU的显存以支持更大的batchsize，可以更好的隐藏通信开销，提升多机可扩展性。

4）3D并行结合optimizer states切分

Deepspeed 3D并行包括数据并行，模型并行和流水线并行。通过不同的Group维护数据并行组，模型并行组和流水线并行组。其中dp group做allreduce同步梯度，同一个pipe group通过broadcast 去send activation（前向）或者send grads（反向）。同一个pipe group只有pipeline的第一个stage需要读取train数据，pipeline的最后一个stage需要读取label数据，算出loss。

流水线并行技术

TI-Deepspeed通过梯度累积实现不同节点间无数据流依赖关系的并行，充分发挥3D并行的优势，充分权衡计算效率和内存效率。

由于模型并行和数据并行参数量比较大，对带宽的要求比较高，所以TI-Deepspeed尽可能在节点内做模型并行和数据并行，充分利用节点内高带宽。而流水线并行因为切分的stage比较多，而且也有梯度累计去做无数据流依赖的不同stage之间的并行，对数据传输效率要求没有那么高，所以可以做节点间的pipeline并行。假设是64卡，本着节点内模型并行，节点内数据并行，节点间pipeline并行的原则，假设mp=4 dp=2 pp=4，那么就形成了dp_groups=[[0, 4], [1, 5], [2, 6], [3, 7], [8, 12], [9, 13], [10, 14], [11, 15], [16, 20], [17, 21], [18, 22], [19, 23], [24, 28], [25, 29], [26, 30], [27, 31]]。

3D并行技术

三、TI-Deepspeed大模型训练最佳实践

腾讯云TI平台团队在对Deepspeed调研和实践的基础上，从性能和易用性两方面对Deepspeed框架进行了相关优化，并根据NLP大模型不同的参数规模沉淀出了完整且高性能的分布式训练方案。腾讯云TI平台团队旨在通过“一套框架”+“三套最佳实践”，更好的服务外部有NLP预训练需求的客户。

1、训练平台介绍

2、3亿参数规模

目前腾讯云TI平台团队Bert Large预训练已对标业界SOTA，使用黑石2.0 8机64卡在13.65小时可以完成Bert Large训练，比NV dgx-2h同样卡数快12.5%。

3、百亿参数规模

腾讯云TI平台团队使用黑石2.0 64卡训练83亿和100亿参数规模 GPT-2, 通过使用zero-2纯数据并行，每卡可以达到近40TFLOPs的性能。

4、四百亿参数规模

腾讯云TI平台团队通过3D并行结合梯度累积，对3D并行进行改造以支持切分optimizer states，使用64卡训练400亿参数规模模型，可以达到40TFlops/GPU。

5、千亿参数规模

腾讯云TI平台团队使用zero-stage3结合cpu offload，在有限的64卡情况下训练千亿模型，通过增大batchsize，隐藏通信开销，提升训练效率。