分布式深度学习

分布式训练将计算分散到多个GPU和机器上，以训练单个设备无法容纳或训练太慢的模型。本文件涵盖混合精度、数据并行、模型并行、流水线并行、ZeRO、FSDP、张量并行以及全规约等通信原语——这些对于大规模训练LLM至关重要。

• 在单个GPU上训练大型神经网络最终会遇到瓶颈。模型可能无法放入内存，或者训练可能需要数月。分布式训练将工作分散到多个设备（GPU、TPU或整台机器）上，以更快地训练和训练更大的模型。本文件涵盖了实现这一目标的技术。

• 要理解为何分布式重要，从训练的计算成本开始。在一个包含 \(d_{\text{in}}\) 个输入和 \(d_{\text{out}}\) 个输出的密集层上，对一批 \(B\) 个样本进行一次前向传播需要大约 \(2 \cdot B \cdot d_{\text{in}} \cdot d_{\text{out}}\) 次FLOP（浮点运算）：对输出矩阵的每个元素进行一次乘法和一次加法。反向传播的成本大约是前向传播的两倍（计算相对于输入和权重的梯度），因此一个密集层的一个训练步骤约为 \(6 \cdot B \cdot d_{\text{in}} \cdot d_{\text{out}}\) 次FLOP。

• 对于隐藏维度为 \(d\) 的Transformer层，自注意力块涉及四个投影（Q、K、V和输出），每个的成本为 \(O(B \cdot n \cdot d^2)\) 次FLOP（其中 \(n\) 是序列长度），加上注意力矩阵计算 \(O(B \cdot n^2 \cdot d)\)。前馈块有两个密集层，通常扩展到 \(4d\) 再回来：\(O(B \cdot n \cdot 8d^2)\)。每层总计：大约 \(O(B \cdot n \cdot 12d^2 + B \cdot n^2 \cdot d)\)。乘以层数，你就会明白为什么训练GPT规模的模型需要数千个GPU小时。

• 内存墙通常是更严格的约束。在训练期间，GPU内存必须同时容纳四样东西：

• 参数：模型权重。一个70亿参数的模型在FP32中（每个参数4字节）仅权重就需要28 GB。
• 梯度：与参数大小相同。又是28 GB。
• 优化器状态：Adam维护两个额外的缓冲区（一阶和二阶矩估计），每个与参数大小相同。即使模型使用较低精度，这些也以FP32格式保存以确保数值稳定性。对于我们的7B模型，那就是 \(2 \times 28 = 56\) GB。

• 激活值：在前向传播过程中保存下来供反向传播使用的中间值。大小取决于批量大小、序列长度和模型宽度。这通常是最主要的组成部分，并随批量大小线性增长。

• 对于使用FP32 Adam的7B模型：28（参数）+ 28（梯度）+ 56（优化器）= 112 GB，这还没算激活值。单个80 GB的A100 GPU无法容纳。这就是分布式策略至关重要的原因。

• 混合精度训练是第一道防线。不是将所有内容存储在FP32（32位浮点）中，而是使用FP16或BF16（16位）进行前向和反向传播，同时将权重的FP32主副本保留给优化器更新。

• FP16具有高精度（10位尾数），但范围有限，可能导致上溢/下溢。损失缩放（在反向传播前将损失乘以一个大因子，然后将梯度除以相同因子）缓解了这个问题。

• BF16（脑浮点）具有与FP32相同的指数范围（8位指数），但精度较低（7位尾数）。它几乎从不溢出，很少需要损失缩放，因此使用更简单。BF16是现代Transformer训练的默认选择。

• 混合精度大致将激活值和梯度的内存减半（前向/反向传播期间的主要成本），同时将优化器状态保留在FP32中以确保数值稳定性。

• 数据并行是最简单的分布式策略。你在 \(N\) 个GPU上复制整个模型，将每个小批量分成 \(N\) 个相等的块，并将一个块发送到每个GPU。每个GPU在其块上独立运行前向和反向传播。然后梯度在所有GPU上平均（使用全规约操作），每个GPU更新其本地模型副本。

• 从模型的角度来看，这相当于使用大了 \(N\) 倍的小批量进行训练。如果每个GPU处理一个大小为 \(B\) 的批次，则有效批量大小为 \(N \cdot B\)。

• 梯度平均可以同步或异步进行。同步SGD等待所有GPU完成后再进行平均，确保与使用更大批量的单GPU训练数学上等价。缺点是，最慢的GPU（"掉队者"）会拖慢所有人。

• 异步SGD让每个GPU独立地更新一个共享的参数服务器，无需等待。这消除了掉队者问题，但引入了"陈旧梯度"：一个GPU可能基于略微过时的参数计算梯度。陈旧梯度增加了噪声，可能减缓收敛。在实践中，带高效通信的同步SGD更受青睐。

• 梯度累积是一种软件技巧，用于在有限硬件上模拟更大的批量大小。不必每个小批量做一次更新，而是运行多次前向/反向传播并累积梯度，然后做一次更新。这与更大批量得到相同的结果，而无需更多GPU内存用于激活值（一次只有一个小批量的激活值在内存中）。

• 当模型本身太大无法放入单个GPU时，需要模型并行。有两种主要形式。

• 张量并行将单个层分割到多个GPU上。一个大的矩阵乘法 \(Y = XW\) 可以按列分割：将 \(W\) 分区为 \([W_1, W_2]\) 分布在两个GPU上，并行计算 \(Y_1 = XW_1\) 和 \(Y_2 = XW_2\)，然后拼接。这适用于注意力投影和前馈层。它需要GPU之间快速通信（通常是节点内的NVLink），因为每层都必须组合部分结果。

• 流水线并行将不同的层分配到不同的GPU上。GPU 0运行第1-4层，GPU 1运行第5-8层，依此类推。数据像流水线一样流经整个管道。朴素的方法有一个"流水线气泡"：当GPU 0处理微批次1的前向传播时，GPU 1-3处于空闲状态。微批处理通过将小批量分割成更小的微批次来缓解这个问题，这些微批次按顺序流经流水线，使所有GPU大部分时间保持忙碌。

• 混合并行结合了数据并行、张量并行和流水线并行。一个典型的大模型设置可能使用节点内的张量并行（8个GPU通过快速NVLink连接）、跨节点的流水线并行以及跨节点组的数据并行。这就是GPT-4和Llama等模型的训练方式。

• 分布式训练的效率在很大程度上取决于通信。关键操作是全规约（all-reduce）：给定 \(N\) 个GPU上各有一个值，计算总和（或平均值）并将结果分发给所有GPU。

• 朴素的全规约将所有数据发送到一个GPU，求和，然后广播回来。通信量为 \(O(N)\)，并在根节点造成瓶颈。

• 环全规约（Ring all-reduce） 要高效得多。将 \(N\) 个GPU排列成一个环。每个GPU将其数据分割成 \(N\) 块。在 \(N - 1\) 步中，每个GPU向邻居发送一块，并从另一个邻居接收一块，累加部分和。再经过 \(N - 1\) 步后，完整的总和传播到所有GPU。每个GPU的总数据传输量：数据大小的 \(2(N-1)/N\) 倍，随着 \(N\) 的增长趋近于 \(2\times\)。关键在于，这不随 \(N\) 增加，使其带宽最优。

• 参数服务器是一种替代架构，其中专用服务器节点保存模型参数。工作节点计算梯度并将其发送到服务器，服务器更新参数并将其发送回来。这更简单，但可能在服务器处造成通信瓶颈。

• NCCL（NVIDIA集合通信库）是GPU间通信的标准库。它提供了全规约、全收集、广播和其他集合操作的高效实现，自动为网络拓扑选择最佳算法。

• 缩放定律描述了模型性能如何随计算量、数据量和模型大小而提升。原始的Kaplan等人（2020）缩放定律发现，损失随每个因素以幂律方式下降：

\[L(N) \propto N^{-\alpha_N}, \quad L(D) \propto D^{-\alpha_D}, \quad L(C) \propto C^{-\alpha_C}\]

• 其中 \(N\) 是参数数量，\(D\) 是数据集大小，\(C\) 是计算预算。

• Chinchilla缩放定律（Hoffmann等人，2022）表明大多数模型训练不足：对于给定的计算预算，应该训练一个更小的模型，使用比以前认为的更多的数据。最优比例大约是每参数20个token。一个7B模型应该看到大约140B个token，而不是Llama 1在65B模型上使用的300B个token。这一发现将领域转向了"计算最优"训练。

• 混合专家（MoE） 是一种在不按比例增加计算量的情况下扩展模型容量的架构。每个Transformer层不是使用一个前馈网络，而是有 \(N\) 个"专家"网络（每个都是一个标准FFN）。一个门控网络（路由器）检查每个token并将其发送到top-\(K\)个专家（通常 \(K = 1\) 或 \(K = 2\)）。

• 总参数量要大得多（因为有 \(N\) 个专家），但每个token的FLOPs大致保持不变（因为每个token只有 \(K\) 个专家激活）。例如，Mixtral 8x7B共有47B个参数，但每次前向传播只用大约13B，以较小模型的代价获得更大模型的性能。

• MoE带来了挑战。负载均衡：如果路由器将大多数token发送到同一个专家，其他专家就被浪费了。辅助损失鼓励均匀路由。通信：不同的专家可能位于不同的GPU上，因此路由token需要全对全通信，这很昂贵。

• 容错在训练运行持续数周或数月、涉及数千个GPU时至关重要。如果单个GPU失效，你不想丢失所有进度。检查点定期将模型权重、优化器状态和训练状态（学习率、步数、数据位置）保存到磁盘。如果发生故障，你可以从最近的检查点重新开始。

• 梯度检查点（也称为激活重计算）是一种内存优化，而非容错机制。在前向传播过程中，不是保存所有激活值供反向传播使用，而是只在某些检查点保存激活值。在反向传播过程中，从检查点重新计算缺失的激活值。这以计算换取内存：它使前向传播成本增加约33%，但可以将激活内存减少 \(\sqrt{L}\) 倍（其中 \(L\) 是层数）。

• 综合起来，训练前沿模型结合了所有这些技术：BF16混合精度、使用环全规约在数千个GPU上进行数据并行、节点内的张量并行、跨节点的流水线并行、减少内存的梯度检查点、提高参数效率的MoE，以及用于容错的定期检查点。系统工程与算法设计一样具有挑战性。

• 总结分布式训练工具包：

技术	作用	权衡
混合精度 (BF16)	将激活值/梯度的内存减半	轻微数值差异
数据并行	在GPU间扩展批量大小	梯度同步的通信开销
张量并行	在GPU间分割层	需要快速互联
流水线并行	在GPU间分割模型阶段	流水线气泡（计算浪费）
梯度累积	模拟大批量	更慢（多次前向/反向传播）
梯度检查点	减少激活内存	约多33%计算
环全规约	高效的梯度平均	大模型受限于带宽
MoE	更多容量，相同FLOPs	负载均衡、路由复杂性
缩放定律	指导计算分配	经验公式，未必在所有规模都成立

编程任务（使用CoLab或笔记本）

1. 计算Transformer层的FLOPs和内存需求。给定隐藏维度 \(d\)、序列长度 \(n\)、批量大小 \(B\) 和层数，估计总训练成本。

   import jax.numpy as jnp
   
   def transformer_layer_flops(d, n, B):
       """一个Transformer层前向传播的近似FLOPs。"""
       # QKV投影：3 * (B * n * d * d) * 2（乘法-加法）
       qkv_flops = 3 * 2 * B * n * d * d
       # 注意力：(B * n * n * d) * 2 用于QK^T，(B * n * n * d) * 2 用于attn*V
       attn_flops = 2 * 2 * B * n * n * d
       # 输出投影：(B * n * d * d) * 2
       out_flops = 2 * B * n * d * d
       # FFN：两层，d->4d 和 4d->d：2 * (B * n * d * 4d) * 2
       ffn_flops = 2 * 2 * B * n * d * 4 * d
       return qkv_flops + attn_flops + out_flops + ffn_flops
   
   def transformer_layer_memory(d, n, B, dtype_bytes=2):
       """一个层的近似激活内存（字节）。"""
       # QKV：3 * B * n * d
       qkv_mem = 3 * B * n * d * dtype_bytes
       # 注意力权重：B * heads * n * n（近似 B * n * n * sizeof）
       attn_mem = B * n * n * dtype_bytes
       # FFN中间值：B * n * 4d
       ffn_mem = B * n * 4 * d * dtype_bytes
       return qkv_mem + attn_mem + ffn_mem
   
   # 示例：GPT-2规模
   d, n, B, L = 1024, 1024, 8, 24
   fwd_flops = transformer_layer_flops(d, n, B)
   total_flops = 3 * L * fwd_flops  # 前向+反向的3倍
   act_mem = L * transformer_layer_memory(d, n, B)
   param_count = L * (12 * d * d + 13 * d)  # 近似
   
   print(f"模型：d={d}, n={n}, B={B}, L={L}")
   print(f"参数：{param_count / 1e6:.0f}M")
   print(f"每步FLOPs：{total_flops / 1e12:.2f} TFLOPs")
   print(f"激活内存：{act_mem / 1e9:.2f} GB (BF16)")
   print(f"参数内存 (FP32)：{param_count * 4 / 1e9:.2f} GB")
   print(f"Adam优化器内存：{param_count * 8 / 1e9:.2f} GB")
   print(f"总训练内存：{(param_count * 16 + act_mem) / 1e9:.2f} GB")
   

2. 模拟数据并行训练。将数据集分割到多个"虚拟GPU"上，独立计算梯度，平均它们，并验证结果与单GPU训练匹配。

   import jax
   import jax.numpy as jnp
   
   # 简单线性模型：y = wx + b
   key = jax.random.PRNGKey(0)
   X = jax.random.normal(key, (64, 4))
   w_true = jnp.array([1.0, -2.0, 3.0, 0.5])
   y = X @ w_true + 0.1 * jax.random.normal(key, (64,))
   
   def loss_fn(w, X, y):
       return jnp.mean((X @ w - y) ** 2)
   
   grad_fn = jax.grad(loss_fn)
   
   # 单GPU：全批量梯度
   w = jnp.zeros(4)
   grad_single = grad_fn(w, X, y)
   
   # 数据并行：分割到4个"GPU"上
   n_gpus = 4
   chunk_size = len(X) // n_gpus
   grads = []
   for i in range(n_gpus):
       X_chunk = X[i*chunk_size:(i+1)*chunk_size]
       y_chunk = y[i*chunk_size:(i+1)*chunk_size]
       grads.append(grad_fn(w, X_chunk, y_chunk))
   
   # 全规约：平均梯度
   grad_parallel = jnp.mean(jnp.stack(grads), axis=0)
   
   print("单GPU梯度：", grad_single)
   print("数据并行梯度（平均）：", grad_parallel)
   print(f"匹配：{jnp.allclose(grad_single, grad_parallel, atol=1e-5)}")
   
   # 训练两者并比较
   w_single, w_parallel = jnp.zeros(4), jnp.zeros(4)
   lr = 0.1
   for step in range(100):
       w_single = w_single - lr * grad_fn(w_single, X, y)
   
       grads = [grad_fn(w_parallel, X[i*chunk_size:(i+1)*chunk_size],
                        y[i*chunk_size:(i+1)*chunk_size]) for i in range(n_gpus)]
       avg_grad = jnp.mean(jnp.stack(grads), axis=0)
       w_parallel = w_parallel - lr * avg_grad
   
   print(f"\n100步之后：")
   print(f"单GPU权重：{w_single}")
   print(f"数据并行权重：{w_parallel}")
   print(f"最大差异：{jnp.max(jnp.abs(w_single - w_parallel)):.2e}")
   

3. 实现一个简单的混合专家层。创建一个门控网络，将token路由到top-K个专家并组合它们的输出。

   import jax
   import jax.numpy as jnp
   
   def expert_fn(x, W1, b1, W2, b2):
       """简单的2层FFN专家。"""
       h = jnp.maximum(0, x @ W1 + b1)  # ReLU
       return h @ W2 + b2
   
   def moe_layer(x, gate_W, experts_params, top_k=2):
       """
       MoE前向传播。
       x: (batch, d_model)
       gate_W: (d_model, n_experts)
       experts_params: 每个专家的 (W1, b1, W2, b2) 列表
       """
       n_experts = len(experts_params)
   
       # 门控：计算路由分数
       gate_logits = x @ gate_W  # (batch, n_experts)
       gate_probs = jax.nn.softmax(gate_logits, axis=-1)
   
       # Top-K选择
       top_k_indices = jnp.argsort(-gate_probs, axis=-1)[:, :top_k]
       top_k_probs = jnp.take_along_axis(gate_probs, top_k_indices, axis=-1)
       # 重新归一化
       top_k_probs = top_k_probs / jnp.sum(top_k_probs, axis=-1, keepdims=True)
   
       # 计算专家输出（简化：运行所有专家，稍后掩码）
       expert_outputs = jnp.stack([
           expert_fn(x, *experts_params[i]) for i in range(n_experts)
       ], axis=1)  # (batch, n_experts, d_model)
   
       # 收集top-K专家输出并加权
       batch_idx = jnp.arange(x.shape[0])[:, None]
       selected_outputs = expert_outputs[batch_idx, top_k_indices]  # (batch, top_k, d_model)
       output = jnp.sum(selected_outputs * top_k_probs[:, :, None], axis=1)
   
       return output, gate_probs
   
   # 设置
   key = jax.random.PRNGKey(42)
   batch, d_model, d_ff, n_experts = 8, 16, 32, 4
   
   # 初始化专家
   experts_params = []
   for i in range(n_experts):
       k1, k2, key = jax.random.split(key, 3)[0], jax.random.split(key, 3)[1], jax.random.split(key, 3)[2]
       experts_params.append((
           jax.random.normal(k1, (d_model, d_ff)) * 0.1,
           jnp.zeros(d_ff),
           jax.random.normal(k2, (d_ff, d_model)) * 0.1,
           jnp.zeros(d_model),
       ))
   
   key, subkey = jax.random.split(key)
   gate_W = jax.random.normal(subkey, (d_model, n_experts)) * 0.1
   x = jax.random.normal(key, (batch, d_model))
   
   output, gate_probs = moe_layer(x, gate_W, experts_params, top_k=2)
   
   print(f"输入形状：{x.shape}")
   print(f"输出形状：{output.shape}")
   print(f"门控概率（第一个样本）：{gate_probs[0]}")
   print(f"专家使用率（批量平均）：")
   for i in range(n_experts):
       usage = jnp.mean(gate_probs[:, i])
       print(f"  专家 {i}: {usage:.3f}")