分布式训练

3D Parallelism 的原理，具体实现 掌握

3D并行训练是大规模模型训练的核心范式，它有机组合了数据并行、流水线并行和张量并行三种维度，在万卡级集群上实现高效训练。其设计哲学是通过多维度的模型切分，将计算、内存和通信开销分布到大量设备上，使训练超大规模模型（如万亿参数）成为可能。

基本原理
数据并行是最基础的维度，将训练数据分片到不同设备组，每个设备持有完整的模型副本，独立进行前向和反向传播，最后同步梯度。这种并行方式的优点是实现简单、扩展性好，但受限于单设备内存容量，无法训练超过设备内存的模型。

流水线并行将模型按层切分到不同设备，形成计算流水线。每个设备只负责模型的一部分层，数据以微批次形式在设备间流动。这种切分能训练超出单设备内存的模型，但会引入流水线气泡——设备等待数据到达的空闲时间。通过精心调度微批次的执行顺序（如1F1B调度），可以显著减少气泡时间。

张量并行是最精细的并行维度，将单个层的计算切分到多个设备。例如在Transformer的线性层中，将权重矩阵按行或列分割；在注意力层中，将注意力头分布到不同设备。这种并行需要设备间频繁通信，但对内存的节省最为直接。

具体实现
在Megatron-LM框架中，3D并行的实现展现了精妙的设计。假设有P个GPU，可以分解为P = D × P × T，其中D是数据并行度，P是流水线并行度，T是张量并行度。

张量并行通常部署在高速互联（如NVLink）的设备组内，因为其通信密集，需要频繁进行All-Reduce操作。在Transformer层中，MLP部分采用列并行和行并行的组合策略，注意力头的计算也被分割到不同设备，通过通信聚合结果。

流水线并行跨越设备组，通常部署在不同节点间。每个设备负责一组连续的层，数据像流水线一样在设备间传递。优化关键在于气泡时间的最小化，通过调整微批次大小和流水线阶段数来平衡内存使用和计算效率。梯度累积技术常用于解决微批次大小与全局批次大小不匹配的问题。

数据并行作为最外层的并行维度，每个数据并行组内的设备运行完整的流水线。梯度同步在数据并行组间进行，通过Ring-AllReduce算法高效实现。

通信模式分析
3D并行中的通信呈现层次化特征。张量并行组内通信最为频繁但数据量较小，适合高速互联；流水线阶段间通信量中等，需要良好带宽；数据并行组间通信量最大但频次较低。现代实现通过计算与通信重叠、梯度累积后同步等技术隐藏通信延迟。

实践中，典型的配置可能是：节点内8个GPU组成一个张量并行组，节点间16个节点组成流水线并行，全局256个这样的组合进行数据并行。这种配置能够训练万亿参数模型，同时保持较高的硬件利用率。

混合精度 重要

混合精度训练是通过协调使用不同精度的数值格式（通常是FP16和FP32）来加速训练并减少内存占用的关键技术。其核心洞察是神经网络训练对数值精度的需求具有不对称性：前向传播和梯度计算可以使用较低精度，而权重更新和某些敏感操作需要较高精度。

工作原理
典型实现中，权重通常以FP32格式存储（主权重），在前向传播时转换为FP16进行计算，得到FP16的激活值。反向传播使用FP16计算梯度，然后这些梯度被转换为FP32用于权重更新。这种设计既利用了FP16的计算速度和内存效率，又通过FP32主权重保持了数值稳定性。

损失缩放是混合精度训练的关键技术。由于FP16的动态范围较小，梯度值可能因下溢而变为零。通过在反向传播前对损失值乘以一个缩放因子（通常为1024-65536），将梯度“放大”到FP16可表示的有效范围内，在优化器更新前再将梯度除以相同因子。