大语建模布式训练时的量化分分析与最佳实践 ， 以 GPT - 175B 为例

大型训练模型的挑战与解决方案

大型训练模型的挑战

1. 高计算成本

   • 每次迭代的计算成本公式为 \( h(1 + 6h^{16B}) \)，其中 \( h \) 为隐藏大小，\( B \) 为批量大小，\( V \) 为词汇大小。
   • GPT-3 175B 模型训练需约 2150 ZettaFLOP（1 ZettaFLOP = 1024 ExaFLOP），使用 128 个 DGX A100 服务器需 170 天，计算效率约 50%。

2. 高内存成本

   • 模型状态总内存约 3.5 TB，其中参数 350GB，梯度 350GB，优化器 2800GB。
   • 单个 GPU 无法容纳模型，需跨节点并行。

大型变压器训练技术概述

• 并行技术：管道并行、张量并行、序列并行、专家并行。
• 内存优化：分布式优化器（类似 DeepSpeed ZeRO-1）、检查点激活、选择性激活检查点。
• 其他技术：FP16/BF16 训练、优化内核、通信重叠。

GPT-3 模型的分布式训练

1. 模型并行性

   • 张量模型并行：层内拆分，实现简单，适合大矩阵，但通信频繁。
   • 管道模型并行：层间拆分，通信粗粒度，可推广至所有 DNN，但需大批量数据。

2. 张量并行示例

   • MLP 层：前向全归约，反向 all-reduce。
   • 自注意力层：前向 all-gather，反向 reduce-scatter。

3. 序列并行

   • Megatron v3 新增，前向 reduce-scatter，反向 all-gather。

4. 交错管线模型并行

   • 结合张量并行和管道并行，优化通信开销。

存储要求与优化

1. 激活内存优化

   • 完整检查点：存储输入激活，需额外前向计算，内存减少但计算开销约 36%。
   • 选择性检查点：仅重计算部分激活，开销降至 4%。

2. 分布式优化器

   • Megatron v3 新增，分区优化器状态，不牺牲性能。

3. 存储分析

   • 模型内存：每个设备约 1.37TB（175B 模型）。
   • 模型状态：AMP 优化后内存成本显著降低。
   • 激活内存：完整检查点需前向重计算，选择性检查点结合序列并行效果更佳。

数据并行化与通信成本

• 数据并行性：\( T1 = Mg_{BW} \ imes 2(d-1)/d \)，其中 \( BW \) 为总线带宽，\( d \) 为数据并行规模。
• 张量并行通信成本：与并行规模和通信模式相关。
• 管道并行通信：采用散射-聚集机制。

计算分析

• GPT-3 175B 训练：每个 GPU 模型 FLOPS > 150，训练效率需 > 120 Model TFLOPS/A100。

建议

• 内存不足时：采用选择性激活检查点、分布式优化器，逐步引入张量并行（序列并行+1024/2048 GPU），再逐步管道并行，最后使用完整激活检查点。
• 内存充足时：增大 batch size，采用数据并行。
• 混合精度建议：BF16 适用于 > 20B 模型。