大语言模型训练核心技巧与优化策略

3.1 优化器状态划分（Stage 1）

将优化器状态划分成 Nd 份，每一份存到不同的 GPU 上。每个 GPU 只需要存储和更新总优化器状态的 1/Nd。

• 显存占用：假设标准数据并行中优化器消耗 KX，ZeRO Stage 1 将优化器显存降低至 KX/Nd。

3.2 梯度划分（Stage 2）

在优化器状态划分的基础上，将梯度划分成 Nd 份，每一份存到不同的 GPU 上。

• 显存占用：降低至 2X + (2X + KX)/Nd。当 Nd 很大时，梯度和优化器状态占比可忽略不计。

3.3 参数划分（Stage 3）

在前两者的基础上，将参数划分成 Nd 份，每一份存到不同的 GPU 上。在前向和反向传播时，通过广播（Broadcast）获取完整参数。

• 显存占用：降低至 (4X + KX)/Nd。只要有足够数量的显卡，ZeRO Stage 3 能够适应任意大的模型。

4. 模型并行与加速策略

4.1 数据并行（Data Parallelism, DP）

不同设备执行相同的模型，处理不同的数据批次。这是最基础的并行方式，但受限于单卡显存大小。

4.2 朴素模型并行（Pipeline Parallelism）

当一个模型大到单个 GPU 无法训练时，最直接的想法是对模型层进行划分，将划分后的部分放置在不同的 GPU 上。

• 流程：GPU1 执行前向传播，缓存激活值发送给 GPU2；GPU2 完成前向和 Loss 计算后，开始反向传播，将梯度返还给 GPU1。
• 缺点：低 GPU 利用率（任意时刻仅一个 GPU 工作），计算和通信没有重叠，高显存占用（需保存整个 minibatch 的激活）。

4.3 GPipe

GPipe 将 minibatch 划分为更小且相等尺寸的 microbatch 来提高效率。前一个计算设备计算出该 microbatch 对应的结果会马上传给下一个计算设备，同时接着对下一个 microbatch 进行计算。这样能同时进行通信和计算。

• Bubble：尽管提高了效率，设备仍会有一段闲置时间，被称为 Bubble。最终会以 mini-batch 为单位将各个设备上的梯度汇总在一起进行参数更新（梯度累积）。

4.4 张量并行（Tensor Parallelism, TP）

张量并行的核心是将矩阵乘法进行拆分，分配到多个 GPU 上计算，降低对单个 GPU 的计算需求。TP 需要大量通讯，因此不建议跨多个节点进行张量并行。实际中，若一个节点有 4 个 GPU，最高的张量并行度通常为 4。

• 一维张量并行：列并行将通信的结果进行拼接，行并行则是对通信结果相加。
• Megatron-LM：针对 Transformer 的 MLP 和 Attention 结构提出了一种高效的张量并行方法。全连接层（MLP）和自注意力层（Self-Attention）的张量并行通过特定的切分策略实现。

4.5 3D 并行

基于流水线并行将模型按 stage 划分到不同的管道，每个管道处理一个模型的 stage；基于张量并行将模型的每个 stage 按张量切分，划分成不同块；最后数据并行可以将这种 2D 组合扩展到任意数量的 GPU 上。

示例配置：

• 模型分成 4 个 stage（PP=4）。
• 每台机器有 8 张 GPU，张量并行度为 4（TP=4）。
• 数据并行度为 2（DP=2）。
• 基于 ZeRO 的 3D 并行允许每个 GPU 只保存训练步骤所需的一小部分数据（参数、梯度和优化器状态）。

显存估算： 已知 Transformer encoder 的参数为：embedding(E)，sequence(s)，attention head(ah)，vocabulary(v)，hidden size(h)，layer(n)。

• 自注意力层 = h * h * 4
• 全连接层 = h * 4h * 2
• 词表 = v * h
• 输入 = s * h

设 DP=8，TP=8，PP=16，使用基于 ZeRO 的 3D 并行，单张 GPU 的模型参数量将大幅降低，具体取决于 ZeRO Stage 的设置。

5. FLOPs 计算与分析

FLOPs（Floating Point Operations）意指浮点运算数，用来衡量算法/模型的复杂度。基于标准 Transformer decoder 结构的模型的 FLOPs 计算方法如下：

5.1 详细计算方法

• Embeddings: 2 × seq_len × vocab_size × d_model
• Attention (Single Layer):
  • Key, query and value projections: 2 × 3 × seq_len × d_model × (key_size × num_heads)
  • Key @ Query logits: 2 × seq_len × seq_len × (key_size × num_heads)
  • Softmax: 3 × num_heads × seq_len × seq_len
  • Softmax @ query reductions: 2 × seq_len × seq_len × (key_size × num_heads)
  • Final Linear: 2 × seq_len × (key_size × num_heads) × d_model
• Dense Block (Single Layer): 2 × seq_len × (d_model × ffw_size + d_model × ffw_size)
• Final Logits: 2 × seq_len × d_model × vocab_size

Total forward pass FLOPs = embeddings + num_layers × (total_attention + dense_block) + logits Total backward pass FLOPs = 2 × Total forward pass FLOPs Total FLOPs = Total forward pass FLOPs + Total backward pass FLOPs

5.2 近似估算公式

Total FLOPs ≈ 6DN，其中 D 是总的训练 tokens 数，N 是模型的参数量。这个公式提供了快速评估训练计算成本的依据。

6. 总结

大语言模型的训练是一个系统工程，需要在显存、通信和计算之间寻找平衡。选择合适的优化策略组合至关重要：

1. 小批量训练：优先使用混合精度和 ZeRO 技术最大化单卡显存利用率。
2. 超大规模模型：必须结合流水线并行和张量并行，并配合 ZeRO Stage 3 进行参数切分。
3. 通信敏感场景：优化 Ring AllReduce 实现，减少跨节点通信开销。
4. 成本评估：利用 FLOPs 公式预估算力需求，合理规划集群规模。

通过上述技术的综合应用，可以在现有的硬件条件下实现更大规模模型的训练与微调，推动人工智能技术的发展。