数据并行(DP):每个GPU持有完整的模型副本,处理不同的小批量数据,梯度在GPU之间取平均。简单高效,适用于能装入单个GPU的模型。张量并行(TP):将各层分割到多个GPU上,每个GPU计算每个矩阵乘法的一部分。当单层的权重无法装入一个GPU时需要使用。流水线并行(PP):不同的层在不同的GPU上运行,微批量数据在流水线中流动。
完全分片数据并行(FSDP,来自PyTorch)和DeepSpeed ZeRO(来自微软)将模型参数、梯度和优化器状态分片到各GPU上。每个GPU只存储模型的一小部分,参数在需要计算时按需聚合,然后释放。这使得训练远超单个GPU内存的模型成为可能。DeepSpeed ZeRO有三个阶段:第一阶段分片优化器状态,第二阶段增加梯度分片,第三阶段增加参数分片。
分布式训练的根本挑战是通信:GPU必须同步梯度(在数据并行中)或交换激活值(在模型/流水线并行中)。这些通信通过NVLink(节点内,900 GB/s)或InfiniBand(节点间,400 Gb/s)进行。当GPU花在等待通信上的时间多于计算时间时,训练效率下降。最优配置通过将紧密耦合的操作(如张量并行)保持在节点内,松散耦合的操作(如数据并行)跨节点分布来最小化跨节点通信。