前言
显存的挑战:当前常见的GPU 显存如 A100 是 80G 显存,而如 175B 的 GPT-3 模型需要 175B * 4bytes 即 700GB 模型参数空间,连加载都困难更不要说训练了。 计算的挑战:175B模型所需要的计算量也是非常庞大的,并且叠加预训练数据量的增大,所需的计算量与BERT时代完全不可同日而语。
三大并行策略
并行技术可以整体划分为两类:数据并行(DP),流水线并行(PP)以及张量并行(TP)。
1 数据并行(DP)
原理:数据并行简单来说,就是每个GPU worker都有一个完整模型的副本,然后获取 per_gpu_batch_size 的数据,每个worker在自己那批数据上计算梯度。最后所有worker 对梯度汇总进行平均最后更新模型权重。 举例:假设有 8 张卡,per_gpu_batch_size=4,那么 global batch size=8*4=32,此时每张卡都保存一个模型,然后每张卡通过 4 个数据来分别计算梯度,之后计算梯度均值,最后每张卡再进行参数更新。
2 流水线并行(PP)与张量并行(TP)
2.1 流水线并行
思想:将模型的不同层放到不同的 GPU 上。比如:一个6层的trasformer ,那么就可以将前3层放在一个GPU上,后3层放在另一个GPU上。如下图所示:
缺点:每个GPU在处理相同batch时都需要等待前一个GPU完成计算才能开始处理,这样容易产生GPU空闲时间,即流水线气泡。为了改善这种情况,通常使用流水线并行(PP)来通过让不同的batch数据的计算重叠。
2.2 张量并行
思想:对层内进行分割,实际上是将Tensor 矩阵运算分配到不同的 GPU 上。如下图就是将每一层分为两部分:
优点:张量并行存储效率较高, 缺点:每次前向和后向计算时会引入额外通信开销。
DeepSpeed 与 Zero
什么是Deepspeed?
Deepspeed Zero 的核心思想
模型状态:模型参数(fp16/bf16,2A),模型梯度(fp16/bf16,2A),Adam状态(fp32的模型参数备份,fp32的momentum和fp32的variance,4A+4A+4A)。假设模型参数量为 A,则共需要:2A + 2A + (4A + 4A + 4A) = 16A 字节存储,如一个7B的模型,模型状态就需要 16 * 7 = 112GB 剩余状态:除了模型状态之外的显存占用,包括激活值(activation)、各种临时缓冲区(buffer)以及无法使用的显存碎片(fragmentation)。
速度: Zero 1 > Zero 2 > Zero 2 + offload > zero 3 > zero 3 + offload 显存:阶段 0 (DDP) < 阶段 1 < 阶段 2 < 阶段 2 + 卸载 < 阶段 3 < 阶段 3 + 卸载
Deepspeed 的优化细节
Deepspeed Zero Stage 1
Deepspeed Zero Stage 2
Deepspeed Zero Stage 3
Zero-offload
