总结：大模型分布式训练技术原理

原文地址：分布式训练技术原理

• 数据并行
  • DataParallel(DP)
    • 基于单进程多线程实现，它使用一个进程来计算模型权重，在每个批处理期间将数据分发到每个GPU，每个GPU处理一部分输入，然后进行前向传播和反向传播，将所有卡中的梯度进行汇聚(规约)来更新主卡中的模型参数，最后将主卡中的模型参数再次分发到每一块GPU中，进行下一轮循环。问题：主卡的通信很快成为瓶颈，GPU利用率通常很低。不能使用apex进行混合精度训练（是apex实现上的问题）。要求所有的GPU都在同一个节点上（不支持分布式）

      

  • DistributedDataParallel (DDP)
    • 基于多进程多线程实现，在每个设备上维护一个模型副本，并通过后向传递的集体AllReduce操作同步梯度，从而确保在训练期间跨副本的模型一致性。为了加快训练速度，DDP将梯度通信与后向传播计算重叠，促进在不同资源上并发执行工作负载。

      

    • DDP需要额外的建立进程组阶段（Construction），Construction阶段需要首先明确通信协议和总进程数。总进程数就是指有多少个独立的并行进程，被称为worldsize。根据需求每个进程可以占用一个或多个GPU，但并不推荐多个进程共享一个GPU，这会造成潜在的性能损失。并行组建立之后，每个GPU上会独立的构建模型，然后GPU-1中模型的状态会被广播到其它所有进程中以保证所有模型都具有相同的初始状态。推理、损失函数计算、梯度计算都是并行独立完成的。DDP实现并行训练的核心在于梯度同步。通常每个GPU也会建立独立的优化器。由于模型具有同样的初始状态和后续相同的梯度，因此每轮迭代后不同进程间的模型是完全相同的。为了优化性能，DDP中针对allreduce操作进行了更深入的设计。通过将全部模型参数划分为无数个小的bucket，在bucket级别建立allreduce。当所有进程中bucket0的梯度计算完成后就立刻开始通信，此时bucket1中梯度还在计算。这样可以实现计算和通信过程的时间重叠。这种设计能够使得DDP的训练更高效。

      

  • ZeRO
    • Model state
      • Optimizer->ZeRO1
        • 将optimizer state分成若干份，每块GPU上各自维护一份
        • 每块GPU上存一份完整的参数W,做完一轮foward和backward后，各得一份梯度,对梯度做一次AllReduce（reduce-scatter + all-gather），得到完整的梯度G,由于每块GPU上只保管部分optimizer states，因此只能将相应的W进行更新,对W做一次All-Gather
      • Gradient+Optimzer->ZeRO2
        • 每个GPU维护一块梯度
        • 每块GPU上存一份完整的参数W,做完一轮foward和backward后，算得一份完整的梯度,对梯度做一次Reduce-Scatter，保证每个GPU上所维持的那块梯度是聚合梯度,每块GPU用自己对应的O和G去更新相应的W。更新完毕后，每块GPU维持了一块更新完毕的W。同理，对W做一次All-Gather，将别的GPU算好的W同步到自己这来
      • Parameter+Gradient+Optimizer->ZeRO3
        • 每个GPU维护一块模型状态
        • 每块GPU上只保存部分参数W，做forward时，对W做一次All-Gather，取回分布在别的GPU上的W，得到一份完整的W，forward做完，立刻把不是自己维护的W抛弃，做backward时，对W做一次All-Gather，取回完整的W，backward做完，立刻把不是自己维护的W抛弃. 做完backward，算得一份完整的梯度G，对G做一次Reduce-Scatter，从别的GPU上聚合自己维护的那部分梯度,聚合操作结束后，立刻把不是自己维护的G抛弃。用自己维护的O和G，更新W。由于只维护部分W，因此无需再对W做任何AllReduce操作
    • Residual state
      • activation->Partitioned Activation Checkpointing
        • 每块GPU上只维护部分的activation，需要时再从别的地方聚合过来就行。需要注意的是，activation对显存的占用一般会远高于模型本身，通讯量也是巨大的
      • temporary buffer->Constant Size Buffer
        • 优点：1. 提升带宽利用率。当GPU数量上升，GPU间的通讯次数也上升，每次的通讯量可能下降（但总通讯量不会变）。数据切片小了，就不能很好利用带宽了。如果等数据积攒到一定大小，再进行通讯，可以提高带宽的利用率。2. 使得存储大小可控。在每次通讯前，积攒的存储大小是常量，是已知可控的，更方便使用者对训练中的存储消耗和通讯时间进行预估。
      • unusable fragment->Memory Defragmentation
        • 对碎片化的存储空间进行重新整合，整出连续的存储空间。防止出现总存储足够，但连续存储不够而引起的存储请求失败。
    • offload
      • ZeRO-Offload
        • 原理：forward和backward计算量高，因此和它们相关的部分，例如参数W（fp16），activation，就全放入GPU；update的部分计算量低，因此和它相关的部分，全部放入CPU中。例如W(fp32)，optimizer states（fp32）和gradients(fp16)等
        • 实现：ZeRO-Offload 分为 Offload Strategy 和 Offload Schedule 两部分，前者解决如何在 GPU 和 CPU 间划分模型的问题，后者解决如何调度计算和通信的问题
      • ZeRO