分布式训练是指将机器学习或深度学习任务分解成多个子任务,并在多个计算设备上并行训练。分布式训练的总体目标就是加快总的训练速度,减少模型训练的总体时间。
总训练速度正比于单设备计算速度*计算设备总量*多设备加速比
其中,单设备计算速度由单块加速芯片的运算速度和数据I/O能力(简单来说,数据 I/O 能力就是指计算机的存储系统向计算核心(如 GPU/CPU)搬运数据的速度。)决定,主要的技术手段有混合精度训练、算子融合、梯度累加等;由于受通信效率的影响,计算设备数量增多会造成加速比急剧降低;多设备加速比是由计算和通信效率决定的,需要结合算法和网络拓扑结构进行优化,分布式训练并行策略的主要目标就是提升分布式训练系统中的多设备加速比。
多设备加速比(Speedup)是衡量并行计算性能的核心指标,它反映了通过增加处理器或计算设备(如CPU核心、GPU等)数量,能让一个任务的执行速度提升多少倍。
加速比的计算非常直观,其基本公式为: 加速比 (Sp) = T1 / Tp * T1:使用单个设备完成该任务所需的串行运行时间。 * Tp:使用 P 个设备并行完成相同任务所需的运行时间。 例如,如果一个程序在单台机器上需要运行 100 秒,而将其分配到 4 台机器上后仅需 30 秒,那么加速比就是 100 / 30 ≈ 3.33 倍。
随着设备数量的增加,加速比往往会逐渐降低,主要受限于以下几个因素:
1. 通信开销:多个设备之间需要频繁交换数据,网络传输和数据同步会消耗大量时间。
2. 串行部分限制:任何程序都有一部分代码是无法并行化的(必须按顺序执行)。根据著名的阿姆达尔定律(Amdahl's Law),系统的最大加速比受到这部分串行代码的严格限制。即使你有无限多的处理器,整体速度也无法超越串行部分的瓶颈。
3. 负载不均衡:如果任务分配不均,有的设备早早完成了工作处于空闲状态,而有的设备还在忙碌,这会拉低整体的效率。
4. 资源竞争与同步:多个设备同时访问内存、缓存或进行线程同步(如加锁/解锁)时会产生竞争和等待。
分布式训练系统需要克服计算墙(单个设备的计算能力和模型所需计算能力之间的差异)、显存墙(模型的完整参数)、通信墙(各设备之间的传输和同步)等挑战。
如果进行并行加速,可以对数据进行切分,将同一个模型复制到多个设备中,称为数据并行(DP);对模型进行划分,称为模型并行(MP);对数据和模型同时切分,称为混合并行(HP)
数据并行:每个设备根据本地样本计算误差,得到梯度,然后对梯度进行广播,最后所有设备聚合梯度值后使用平均梯度对模型进行更新,完成该批次的训练。数据并行系统可以通过增加计算设备,有效提升每秒全局批次数。
模型并行:一种是按模型的层切分到不同设备,即层间并行或流水线并行;一种是将同一层的参数切分到不同设备,即层内并行或张量并行。
流水线并行使得前后阶段能够流水式、分批工作。但是下游设备需要长时间空闲状态,等待上游设备计算完成,形成了流水线气泡。
GPipe方法:将mini-batch划分为更小的micro-batch,每次处理一个micro-batch的数据。缺点是只有当一个小批次中的所有前向计算都完成后,才能执行后向计算。保存所有micro-batch的前向中间结果,导致OOM。
解决方法是1F1B流水线并行策略,包括非交错式和交错式两种模式。
1F1B非交错式调度模式的过程分为三个阶段,第一个阶段是热身阶段,比如说先让所有设备都被填满并完成至少一次前向计算。第二个阶段是稳定阶段,每台设备进入做一次新数据前向,再做一次旧数据后向的循环状态。第三个阶段是结尾阶段,当所有前向计算结束后,完成剩余的后向计算的过程。
优势包括两个,首先极大的减少了前向计算产生的中间变量,排除了OOM显存溢出;第二是更早的反向传播,训练反馈延迟更低。
1F1B交错式让设备负责多段不连续的层,每个设备被分配到多个阶段,可以进一步压缩流水线中的空闲等待时间,但是会增加通信开销。
* 模型层的交错分配: 传统的做法是把连续的层分给同一个设备(例如4个设备跑16层模型,设备0负责1-4层,设备1负责5-8层...)。而交错式调度会将模型切成更多、更细的“模型块”,并交错地分配给设备。 * 举个例子:同样是4个设备跑16层模型,设备0可能负责第1、2层和第9、10层;设备1负责第3、4层和第11、12层,以此类推。这样,每个物理设备上都运行着多个不连续的“虚拟阶段”。
* 极高的资源利用率: 由于每个设备负责了多个虚拟阶段,数据流的旅程变短了,粒度更细。这使得设备可以在执行其中一个模型块的前向计算时,并行处理另一个模型块的反向计算。这种高度交错的安排,让流水线中的“空泡”时间大幅缩短,GPU几乎时刻都在高效运转。
不过1F1B非交错式和GPipe完成一轮计算的时间是相同的,因为流水线空闲时间的总量是等价的。
张量并行将参数切分到不同设备,主要由嵌入式表示、矩阵乘和交叉熵损失构成。
嵌入式表示将词维度切分,分别存储,当每个节点查询各自的词表时,如果无法查到,那么该词表示为0,各设备查询后得到结果张量,最后通过ALLREDUCESUM通信,跨设备求和,得到完整的全量结果。
矩阵乘的张量并行要充分利用矩阵的分块乘法原理。
交叉熵损失可以按照类别维度切分,同时通过中间结果通信,得到最终的全局交叉熵损失。
混合并行将多种并行策略如数据并行、流水线并行和张量并行等混合使用
ZeRO(零冗余优化器)包含三种方法:对优化器状态进行分区,对模型梯度进行分区,对模型参数进行分区。
分布式训练集群属于高性能计算集群,目标是提供海量的计算能力,两种常见架构包括参数服务器架构和去中心化架构。
参数服务器架构包括两种服务器角色:训练服务器和参数服务器。参数服务器提供充足的内存资源和通信资源,训练服务器提供计算资源。训练服务器都拥有完整的模型,对数据集进行切片并计算,再将梯度推送到相应的参数服务器,参数服务器会等待训练服务器完成梯度推送,再计算平均梯度并更新参数。参数服务器架构包括同步训练和异步训练两种模式,同步训练会等到收到所有的梯度后进行参数更新;异步训练不会等收到所有的梯度,而是直接基于已收到的梯度进行参数更新。
去中心化架构中,没有中央服务器或控制节点,由节点之间进行直接通信和协调。好处是减少通信瓶颈,提高系统的可扩展性,可以通过集合通信技术实现分布式训练,常用通信原语包括 broadcast, scatter, reduce, all reduce, gather, all gather, reduce scatter, all to all 等。
broadcast 主节点把自身的数据发送到集群中的其他节点
scatter 主节点对数据进行划分并散布至其他指定的节点
reduce 将不同节点上的计算结果进行聚合到一个节点上
all reduce 将聚合后的节点再广播到各个计算设备上
gather 将多个节点上的数据收集到单个节点上
all gather 收集数据后的节点在进行广播
reduce scatter 将每个节点中的张量切分为多个块,再把每个块分配到不同的节点并进行特定操作
all to all 将每个节点中的张量切分为多个块,再把每个块分配到不同的节点
参考书籍:
大规模语言模型理论与实践__第二版