规模化：分布式训练、FSDP、DeepSpeed
你的 124M 模型在一张 GPU 上训练完了。现在试试 70 亿参数。模型塞不进显存。数据在单机上要跑好几周。到了规模上，分布式训练不是可选项。它是唯一的出路。
类型： Build 语言： Python 前置要求： 阶段 10，第 04 课（预训练一个 Mini GPT） 预计时间： ~120 分钟
学习目标
- 解释三类并行（数据、张量、流水线），以及在模型和集群规模下各自何时必要
- 用 PyTorch DDP 实现数据并行训练，跨多张 GPU 做梯度同步
- 为给定模型规模计算显存预算（权重 + 优化器状态 + 梯度 + 激活），定出最低硬件需求
- 配置 FSDP 或 DeepSpeed ZeRO 各阶段，把模型状态分片到多张 GPU，让超出单卡显存的模型也能塞下
问题背景
一个 7B 参数模型在 FP16 下，光权重就要 14GB。Adam 优化器为每个参数再存两份拷贝（一阶和二阶矩估计）。那又是 28GB。反向传播时的梯度再加 14GB。还没存下任何一个激活，你就到 56GB 了。
一张 NVIDIA A100 有 80GB 显存。
80GB 里吃掉 56GB。剩下 24GB 留给激活——前向传播时计算出、必须保留到反向传播的中间值。对一条 2048-token 序列、4096 维的模型，单层的激活大约用 64MB。32 层的话，每个样本要 2GB。batch size 8 需要 16GB。你有 24GB。batch size 12 就爆了。
现在试试 70B 参数。光权重：FP16 下 140GB。一张 GPU 装不下。你至少需要 2 张 A100（2 x 80GB = 160GB）才能放下权重。加上优化器状态和梯度，你需要的远不止：最少 3+ 张 GPU，按分片策略来看实际要 8-16 张。
Llama 3 405B 在 16,384 张 NVIDIA H100 上训练。这次训练的算力成本估计为 1 亿美元。DeepSeek V3 训出一个可比的模型只花了约 560 万美元，靠的是在架构上的聪明（Mixture of Experts 意味着每个 token 只激活一小部分参数）和训练效率。
本节课覆盖让大规模训练成为可能的四种策略：数据并行、张量并行、流水线并行和完全分片数据并行。你会用纯 Python 模拟每一种，先理解机制，再去碰任何分布式训练框架。
核心概念
为什么必须分布
下面是真实模型的显存账。每个数字都是算出来的，不是估的。
| 模型 | 参数 | 权重（FP16） | Adam 状态 | 梯度（FP16） | 合计（不含激活） |
|---|---|---|---|---|---|
| GPT-2 Small | 124M | 248 MB | 992 MB | 248 MB | 1.5 GB |
| Llama 3 8B | 8B | 16 GB | 64 GB | 16 GB | 96 GB |
| Llama 3 70B | 70B | 140 GB | 560 GB | 140 GB | 840 GB |
| Llama 3 405B | 405B | 810 GB | 3,240 GB | 810 GB | 4,860 GB |
"Adam 状态" 这一列是要命的。Adam 为每个参数存一个滑动均值（m）和一个滑动方差（v），都用 FP32。对一个 70B 模型，那是 70B x 4 字节 x 2 = 560GB。光优化器就要七张 A100。
单张 H100 有 80GB。Llama 3 405B 至少要 61 张 H100 才能放下权重、优化器和梯度。加上激活数字还会涨。Meta 用 16,384 张 GPU 不是因为他们想用——是因为他们不得不用。
数据并行
最简单的分布式策略。把整个模型拷贝到 N 张 GPU。把每个训练批次切成 N 等份。每张 GPU 在自己那份数据上跑一次前向和反向传播。反向传播后，把梯度在所有 GPU 间求平均。每张 GPU 用相同的平均梯度更新自己那份权重，让所有拷贝保持同步。
好处： 吞吐线性扩展。N 张 GPU 每步处理 N 倍的数据。通信仅限于梯度求平均，而它能和计算重叠。
坏处： 每张 GPU 都持有模型、优化器状态和梯度的完整拷贝。对一个 70B 模型，每张 GPU 需要 840GB。数据并行完全不减少每卡显存。它只减少训练时间。
算式： 有效 batch size = per_gpu_batch_size x N。对 N=64 张 GPU、每卡 batch 16，有效 batch 是 1,024。Llama 3 用的有效 batch size 是每步 1600 万 token。
graph TD
subgraph DataParallel["数据并行 (N=4 GPU)"]
B["完整 Batch\n(1024 个样本)"] --> S["拆分"]
S --> G1["GPU 1\n完整模型副本\n256 个样本"]
S --> G2["GPU 2\n完整模型副本\n256 个样本"]
S --> G3["GPU 3\n完整模型副本\n256 个样本"]
S --> G4["GPU 4\n完整模型副本\n256 个样本"]
G1 --> AR["AllReduce\n梯度平均"]
G2 --> AR
G3 --> AR
G4 --> AR
AR --> U["更新\n（所有 GPU 上一致）"]
end
style B fill:#1a1a2e,stroke:#e94560,color:#fff
style G1 fill:#1a1a2e,stroke:#0f3460,color:#fff
style G2 fill:#1a1a2e,stroke:#0f3460,color:#fff
style G3 fill:#1a1a2e,stroke:#0f3460,color:#fff
style G4 fill:#1a1a2e,stroke:#0f3460,color:#fff
style AR fill:#1a1a2e,stroke:#51cf66,color:#fff
style U fill:#1a1a2e,stroke:#51cf66,color:#fff张量并行
把单个层切到多张 GPU。一次矩阵乘法被分摊到多张 GPU，每张算出结果的一部分。
考虑一个前馈层里形状为 (8192, 8192) 的权重矩阵。用 4 路张量并行，每张 GPU 持有一个 (8192, 2048) 的分片。每张 GPU 把输入乘上自己的分片，产出一个部分结果。这些部分结果被合并（通过 all-reduce 或 all-gather）成完整输出。
好处： 减少每卡的模型权重显存。一个 70B 模型切到 8 张 GPU，意味着每张 GPU 只持有约 8.75B 参数量的权重。
坏处： 每层之后都需要快速的 GPU 间通信。每次 matmul 后的 all-reduce 增加延迟。这在 NVLink（同节点 GPU 间 900 GB/s）上效果好，但跨 InfiniBand 连接的节点（400 Gb/s，约 50 GB/s）就很差。张量并行几乎总是限制在单个节点内（8 张 GPU）。
真实用法： Megatron-LM 开创了张量并行。Llama 3 405B 在每个节点内用 8 路张量并行。
流水线并行
按层切分模型。GPU 1 跑第 1-8 层。GPU 2 跑第 9-16 层。GPU 3 跑第 17-24 层。GPU 4 跑第 25-32 层。数据流过流水线：GPU 1 算完自己的层，把激活发给 GPU 2，GPU 2 算完自己的层发给 GPU 3，依此类推。
好处： GPU 间通信极少——只有层边界处的激活，相比梯度或权重要小得多。能跨节点工作，因为带宽需求低。
坏处： 流水线气泡。当 GPU 4 在对微批次 1 做前向时，GPU 1、2、3 闲着（它们已经前向完自己那部分了）。反向传播时，模式反过来。朴素流水线下，N 个流水线阶段的 GPU 利用率只有 1/N。
GPipe 和 PipeDream 通过把批次切成微批次来解决气泡问题。GPU 1 一前向完微批次 1，就立刻开始微批次 2。这让计算在流水线阶段间重叠。有 M 个微批次、N 个阶段时，气泡占比降到 (N-1)/M。用 M=16 个微批次、N=4 个阶段，气泡是 3/16 = 18.75% 的空闲时间。
FSDP：完全分片数据并行
FSDP 把数据并行的可扩展性和分片的显存效率结合起来。每张 GPU 不再持有模型的完整拷贝，而只持有 1/N 的参数、梯度和优化器状态。
在一层的前向传播之前，FSDP 跑一次 all-gather，把所有 GPU 上的完整参数收集到每张 GPU 的显存里。前向之后，每张 GPU 丢掉非本地的参数。反向时，all-gather 再跑一次来重建参数以计算梯度。反向之后，一次 reduce-scatter 分发梯度分片，让每张 GPU 只存 1/N 的梯度。
一个 70B 模型在 8 张 GPU 上的算式：
| 组件 | 不用 FSDP | 用 FSDP |
|---|---|---|
| 权重（FP16） | 每卡 140 GB | 每卡 17.5 GB |
| Adam 状态（FP32） | 每卡 560 GB | 每卡 70 GB |
| 梯度（FP16） | 每卡 140 GB | 每卡 17.5 GB |
| 合计 | 每卡 840 GB | 每卡 105 GB |
不用 FSDP，你没法把 70B 模型塞进单张 80GB GPU。在 8 张 GPU 上用 FSDP，每张 GPU 用 105GB——等等，这还是装不下。你至少需要 16 张 GPU 才能让每卡降到 80GB 以下，或者把 FSDP 和激活检查点结合（反向时重算激活而不是存它们）。
通信成本比普通数据并行更高，因为每层前都有 all-gather。但这点显存节省，让以前不可能的训练成为可能。
graph TD
subgraph FSDP["FSDP：完全分片数据并行 (4 GPU)"]
direction TB
S["模型：4 层，已分片"]
subgraph GPU1["GPU 1"]
G1S["分片：1/4 参数\n1/4 优化器\n1/4 梯度"]
end
subgraph GPU2["GPU 2"]
G2S["分片：1/4 参数\n1/4 优化器\n1/4 梯度"]
end
subgraph GPU3["GPU 3"]
G3S["分片：1/4 参数\n1/4 优化器\n1/4 梯度"]
end
subgraph GPU4["GPU 4"]
G4S["分片：1/4 参数\n1/4 优化器\n1/4 梯度"]
end
AG["All-Gather\n（每层前重建完整参数）"]
FW["前向传播\n（暂时持有完整参数）"]
RS["Reduce-Scatter\n（反向后分发梯度分片）"]
S --> GPU1
S --> GPU2
S --> GPU3
S --> GPU4
GPU1 --> AG
GPU2 --> AG
GPU3 --> AG
GPU4 --> AG
AG --> FW
FW --> RS
end
style G1S fill:#1a1a2e,stroke:#0f3460,color:#fff
style G2S fill:#1a1a2e,stroke:#0f3460,color:#fff
style G3S fill:#1a1a2e,stroke:#0f3460,color:#fff
style G4S fill:#1a1a2e,stroke:#0f3460,color:#fff
style AG fill:#1a1a2e,stroke:#e94560,color:#fff
style FW fill:#1a1a2e,stroke:#51cf66,color:#fff
style RS fill:#1a1a2e,stroke:#e94560,color:#fffDeepSpeed ZeRO
DeepSpeed 的 ZeRO（Zero Redundancy Optimizer，零冗余优化器）在概念上和 FSDP 完全一致，但由微软独立开发。它定义了三个阶段，每个分片得更激进：
| 阶段 | 分片 | 显存节省 | 通信 |
|---|---|---|---|
| ZeRO-1 | 仅优化器状态 | ~4 倍削减 | 和数据并行相同 |
| ZeRO-2 | + 梯度 | ~8 倍削减 | 略多一点 |
| ZeRO-3 | + 参数 | ~N 倍削减（N 张 GPU） | 每层 all-gather |
ZeRO-3 等价于 FSDP。命名不同，机制相同。DeepSpeed 证明了这个概念后，PyTorch 把 FSDP 加进来作为原生实现。
DeepSpeed 还引入了 ZeRO-Offload（把优化器状态卸载到 CPU 内存，更便宜也更大）和 ZeRO-Infinity（卸载到 NVMe SSD）。这些用计算速度换显存容量——被卸载的操作更慢，但腾出了 GPU 显存。
混合精度训练
现代训练同时使用多种浮点格式：
- 前向传播：FP16 或 BF16（16 位）。FP32 一半的显存。matmul 在 tensor core 上快 2 倍。
- 主权重：FP32（32 位）。由优化器维护，用于权重更新时的数值精度。
- 损失缩放：反向传播前把损失乘以一个大常数，防止 FP16 梯度下溢到零。优化器步骤前再除以同一个常数。
BF16（Brain Float 16）和 FP32 有相同的指数范围（8 个指数位），但精度更低（7 个尾数位 vs FP32 的 23 个）。它很少需要损失缩放，因为它能表示相同范围的值。FP16 有 5 个指数位和 10 个尾数位——它能表示细粒度的值，但在极端量级上会溢出/下溢。
Google 的 TPU 原生用 BF16。NVIDIA 的 A100 和 H100 同时支持 FP16 和 BF16。业界基本转向了 BF16，因为它消除了损失缩放的麻烦。
一个 7B 模型的显存对比：
| 精度 | 权重 | 优化器 | 梯度 | 合计 |
|---|---|---|---|---|
| 处处 FP32 | 28 GB | 56 GB | 28 GB | 112 GB |
| 混合（BF16 + FP32 主权重） | 14 GB | 56 GB | 14 GB | 84 GB |
混合精度在这个模型上省了 28GB。优化器状态无论如何都留在 FP32——这是大部分显存的去处。
Megatron-LM 和 3D 并行
真正的大规模训练把三种并行全用上：
- 数据并行跨节点组（扩展 batch size）
- 张量并行在节点内（把层切到 8 张 GPU）
- 流水线并行跨节点（把层组切到不同机器）
Llama 3 405B 在 16,384 张 H100 上：
- 每个节点内 8 路张量并行（每节点 8 张 GPU）
- 跨节点 16 路流水线并行（16 个流水线阶段）
- 在剩下的维度上 128 路数据并行（16,384 / 8 / 16 = 128）
这种 3D 分解（8 x 16 x 128 = 16,384）就是你扩展到上千张 GPU 的方式。每张 GPU 看到不同的数据分片（数据并行），持有每层的一个切片（张量并行），并计算一组不同的层（流水线并行）。
DeepSeek V3 走了另一条路。他们的 Mixture of Experts 架构每个 token 只激活 671B 参数里的 37B。这意味着每张 GPU 只需计算（并为之存激活）那些被激活的参数。他们在 2,048 张 H800 GPU 上训练——不到 Meta GPU 数量的 1/8——花了 560 万美元，对比 Meta 估计的 1 亿美元。
graph TD
subgraph ThreeD["3D 并行 (Llama 3 405B)"]
direction TB
subgraph DP["数据并行 (128 路)\n将 batch 分到 128 组"]
subgraph PP["流水线并行 (16 路)\n将层分到 16 个阶段"]
subgraph TP["张量并行 (8 路)\n将每层分到 8 张 GPU"]
G1["GPU 1\n第 1-N 层的切片"]
G2["GPU 2\n第 1-N 层的切片"]
G8["GPU 8\n第 1-N 层的切片"]
end
end
end
end
N1["总计：8 x 16 x 128 = 16,384 张 GPU"]
style G1 fill:#1a1a2e,stroke:#0f3460,color:#fff
style G2 fill:#1a1a2e,stroke:#0f3460,color:#fff
style G8 fill:#1a1a2e,stroke:#0f3460,color:#fff
style N1 fill:#1a1a2e,stroke:#e94560,color:#fffpaged-kv-cache动手构建
第 1 步：模拟数据并行
把一个批次切到模拟的多张 GPU 上。每张 GPU 在自己那份上跑前向。把 "梯度"（我们用损失值模拟）求平均。
import numpy as np
def simulate_data_parallelism(data, num_gpus, model_fn):
batch_size = len(data)
shard_size = batch_size // num_gpus
remainder = batch_size % num_gpus
gpu_losses = []
gpu_gradients = []
offset = 0
for gpu_id in range(num_gpus):
extra = 1 if gpu_id < remainder else 0
shard = data[offset:offset + shard_size + extra]
offset += shard_size + extra
loss, grad = model_fn(shard)
gpu_losses.append(loss)
gpu_gradients.append(grad)
avg_loss = np.mean(gpu_losses)
avg_gradient = np.mean(gpu_gradients, axis=0)
return avg_loss, avg_gradientall-reduce 操作（梯度求平均）是数据并行里唯一的通信。实践中，这在 NVIDIA GPU 上用 NCCL 库，它实现环形 all-reduce：每张 GPU 把自己 1/N 的梯度发给邻居，从另一个邻居收 1/N，N-1 步之后每张 GPU 都有完整的平均值。总通信量：2 x 梯度大小 x (N-1)/N，对大 N 趋近于梯度大小的 2 倍。
第 2 步：模拟张量并行
把一个权重矩阵切到多张 GPU。每张 GPU 算一次部分矩阵乘法。合并结果。
def simulate_tensor_parallelism(input_data, weight_matrix, num_gpus):
d_in, d_out = weight_matrix.shape
assert d_out % num_gpus == 0, f"d_out {d_out} not divisible by num_gpus {num_gpus}"
shard_size = d_out // num_gpus
partial_results = []
for gpu_id in range(num_gpus):
start = gpu_id * shard_size
end = start + shard_size
weight_shard = weight_matrix[:, start:end]
partial = input_data @ weight_shard
partial_results.append(partial)
full_output = np.concatenate(partial_results, axis=-1)
direct_output = input_data @ weight_matrix
error = np.abs(full_output - direct_output).max()
return full_output, error误差应该恰好是零（或机器精度）。张量并行在数学上是精确的——它产出的结果和在单张 GPU 上算完整 matmul 一样。切分沿着输出维度，所以每张 GPU 产出不同的列块，拼接就重建出完整结果。
对列并行的线性层（切输出维度），你做拼接。对行并行（切输入维度），你做求和。在 transformer FFN 里，第一个线性层（扩展）用列并行，第二个线性层（收缩）用行并行。这避免了两层之间的一次 all-reduce。
第 3 步：模拟流水线并行
把一个模型的层切到虚拟 GPU 上。展示气泡问题——早期阶段闲着、后期阶段在算。
def simulate_pipeline_parallelism(num_layers, num_stages, num_microbatches):
layers_per_stage = num_layers // num_stages
timeline = {}
clock = 0
for mb in range(num_microbatches):
for stage in range(num_stages):
start_time = max(
timeline.get((stage, mb - 1, "fwd"), (0, 0))[1] if mb > 0 else 0,
timeline.get((stage - 1, mb, "fwd"), (0, 0))[1] if stage > 0 else 0,
)
end_time = start_time + layers_per_stage
timeline[(stage, mb, "fwd")] = (start_time, end_time)
last_fwd_end = max(v[1] for v in timeline.values())
for mb in range(num_microbatches - 1, -1, -1):
for stage in range(num_stages - 1, -1, -1):
deps = [last_fwd_end]
if mb < num_microbatches - 1 and (stage, mb + 1, "bwd") in timeline:
deps.append(timeline[(stage, mb + 1, "bwd")][1])
if stage < num_stages - 1 and (stage + 1, mb, "bwd") in timeline:
deps.append(timeline[(stage + 1, mb, "bwd")][1])
start_time = max(deps)
end_time = start_time + layers_per_stage
timeline[(stage, mb, "bwd")] = (start_time, end_time)
total_time = max(v[1] for v in timeline.values())
compute_time = num_microbatches * num_stages * layers_per_stage * 2
bubble_fraction = 1.0 - compute_time / (total_time * num_stages)
return timeline, total_time, bubble_fraction4 个阶段、1 个微批次时，气泡占比是 75%——任意时刻四张 GPU 里有三张闲着。用 16 个微批次，它降到约 19%。消除气泡的代价是显存：你必须同时存下所有在途微批次的激活。
第 4 步：显存计算器
为任意模型规模算出确切的显存需求。
def memory_calculator(
params_billions,
precision_bytes=2,
optimizer="adam",
num_gpus=1,
sharding="none",
sequence_length=2048,
batch_size_per_gpu=1,
hidden_dim=None,
num_layers=None,
):
params = params_billions * 1e9
weight_memory = params * precision_bytes
if optimizer == "adam":
optimizer_memory = params * 4 * 2
elif optimizer == "sgd":
optimizer_memory = params * 4
else:
optimizer_memory = 0
gradient_memory = params * precision_bytes
total_no_activation = weight_memory + optimizer_memory + gradient_memory
if hidden_dim and num_layers:
activation_per_layer = (
sequence_length * batch_size_per_gpu * hidden_dim * precision_bytes * 4
)
activation_memory = activation_per_layer * num_layers
else:
activation_memory = params * precision_bytes * 0.5
if sharding == "fsdp" or sharding == "zero3":
weight_memory /= num_gpus
optimizer_memory /= num_gpus
gradient_memory /= num_gpus
elif sharding == "zero2":
optimizer_memory /= num_gpus
gradient_memory /= num_gpus
elif sharding == "zero1":
optimizer_memory /= num_gpus
per_gpu_total = weight_memory + optimizer_memory + gradient_memory + activation_memory
return {
"params_billions": params_billions,
"weights_gb": weight_memory / 1e9,
"optimizer_gb": optimizer_memory / 1e9,
"gradients_gb": gradient_memory / 1e9,
"activations_gb": activation_memory / 1e9,
"per_gpu_total_gb": per_gpu_total / 1e9,
"total_across_gpus_gb": per_gpu_total * num_gpus / 1e9,
"fits_on_80gb": per_gpu_total / 1e9 <= 80,
"num_gpus": num_gpus,
"sharding": sharding,
}这个计算器回答每个 ML 工程师都会问的问题："我需要多少张 GPU？" 喂给它模型规模，看它能不能装下。调整分片策略，直到每卡总量降到 80GB 以下。
第 5 步：混合精度模拟
对比 FP32、FP16 和混合精度训练之间的显存用量。
def mixed_precision_comparison(params_billions):
params = params_billions * 1e9
fp32_weights = params * 4
fp32_optimizer = params * 4 * 2
fp32_gradients = params * 4
fp32_total = fp32_weights + fp32_optimizer + fp32_gradients
fp16_weights = params * 2
fp16_master = params * 4
fp16_optimizer = params * 4 * 2
fp16_gradients = params * 2
fp16_total = fp16_weights + fp16_master + fp16_optimizer + fp16_gradients
mixed_weights = params * 2
mixed_optimizer = params * 4 * 2
mixed_gradients = params * 2
mixed_total = mixed_weights + mixed_optimizer + mixed_gradients
return {
"fp32_total_gb": fp32_total / 1e9,
"fp16_with_master_gb": fp16_total / 1e9,
"mixed_bf16_gb": mixed_total / 1e9,
"savings_vs_fp32": 1 - mixed_total / fp32_total,
}对大多数人最大的意外：混合精度并不把显存减半。优化器状态（Adam 的 m 和 v）无论精度如何都留在 FP32。对一个 7B 模型，FP32 训练用 112GB。混合精度用 84GB。那是削减 25%，不是 50%。优化器才是大头。
实际使用
跑完所有模拟
def run_all_demos():
print("=" * 70)
print("DATA PARALLELISM SIMULATION")
print("=" * 70)
np.random.seed(42)
data = np.random.randn(64, 32)
weight = np.random.randn(32, 16)
def model_fn(batch):
output = batch @ weight
loss = np.mean(output ** 2)
grad = 2 * batch.T @ (batch @ weight) / len(batch)
return loss, grad
for n_gpus in [1, 2, 4, 8]:
loss, grad = simulate_data_parallelism(data, n_gpus, model_fn)
print(f" {n_gpus} GPUs: loss={loss:.4f}, grad_norm={np.linalg.norm(grad):.4f}")
print()
print("=" * 70)
print("TENSOR PARALLELISM SIMULATION")
print("=" * 70)
x = np.random.randn(4, 8192)
W = np.random.randn(8192, 8192)
for n_gpus in [1, 2, 4, 8]:
output, error = simulate_tensor_parallelism(x, W, n_gpus)
print(f" {n_gpus} GPUs: output_shape={output.shape}, max_error={error:.2e}")
print()
print("=" * 70)
print("PIPELINE PARALLELISM SIMULATION")
print("=" * 70)
for n_mb in [1, 4, 8, 16, 32]:
_, total_t, bubble = simulate_pipeline_parallelism(32, 4, n_mb)
print(f" {n_mb:2d} micro-batches: total_time={total_t:4d}, bubble={bubble:.1%}")
print()
print("=" * 70)
print("MEMORY CALCULATOR")
print("=" * 70)
configs = [
(7, "none", 1),
(7, "fsdp", 8),
(70, "none", 1),
(70, "fsdp", 8),
(70, "fsdp", 16),
(405, "fsdp", 64),
(405, "fsdp", 128),
]
print(f" {'Model':>8} {'Sharding':>8} {'GPUs':>5} {'Per-GPU':>10} {'Fits 80GB':>10}")
print(" " + "-" * 50)
for params, shard, gpus in configs:
result = memory_calculator(params, num_gpus=gpus, sharding=shard)
fits = "Yes" if result["fits_on_80gb"] else "No"
print(f" {params:>6}B {shard:>8} {gpus:>5} {result['per_gpu_total_gb']:>8.1f}GB {fits:>10}")
print()
print("=" * 70)
print("MIXED PRECISION COMPARISON")
print("=" * 70)
for params_b in [7, 13, 70, 405]:
result = mixed_precision_comparison(params_b)
print(f" {params_b}B: FP32={result['fp32_total_gb']:.0f}GB, "
f"Mixed BF16={result['mixed_bf16_gb']:.0f}GB, "
f"Savings={result['savings_vs_fp32']:.0%}")拿去用
本节课产出 outputs/prompt-distributed-training-planner.md——一个 prompt，接收一个模型规模和可用硬件，然后产出一份完整的分布式训练计划：并行策略、显存预算、通信开销和预期吞吐。
练习
改造显存计算器，把激活检查点纳入。用了检查点，只在每 K 层存一次激活（典型 K=1，意味着全部重算）。展示显存-计算的权衡：检查点能省多少显存，又让训练慢多少（全检查点大约多 33% 计算）？
把流水线并行模拟扩展为实现 PipeDream 用的 1F1B（一前向一反向）调度。对 4 个阶段、8 个微批次，把气泡占比和朴素调度对比。1F1B 调度的峰值显存应该更小，因为它更早开始反向传播。
实现一个梯度累积模拟器。不在每个微批次后做 all-reduce，而是本地累积梯度 K 步，然后再 all-reduce。展示这如何把通信减少 K 倍，但产出完全相同的最终梯度（因而训练也完全相同）。
做一个成本估算器。给定模型规模、目标 token 数、GPU 类型（A100 每小时 2 美元，H100 每小时 3.50 美元）和并行策略，估出总训练成本（以美元计）。和已知成本对照验证：据报道 Llama 3 405B 花了约 1 亿美元，DeepSeek V3 花了约 560 万美元。
给显存计算器加上 ZeRO-Offload。假设每节点 CPU 内存 512GB、NVMe 2TB。展示把优化器状态卸载到 CPU 如何让一个 70B 模型在 4 张 GPU 上训练（而不是 16 张），代价是优化器步骤慢 30-50%。
关键术语
| 术语 | 人们怎么说 | 它实际是什么 |
|---|---|---|
| 数据并行 | "把模型拷到每张 GPU" | 每张 GPU 处理不同的数据分片；每步后梯度通过 all-reduce 求平均 |
| 张量并行 | "把一层切到多张 GPU" | 切分权重矩阵让每张 GPU 算 matmul 的一部分；需要快速的 NVLink 互联 |
| 流水线并行 | "把层切到多张 GPU" | 每张 GPU 跑不同的一组层；数据用微批次流过流水线以减少气泡 |
| FSDP | "把所有东西分片" | 完全分片数据并行——每张 GPU 持有 1/N 的权重、梯度和优化器状态；计算前 all-gather |
| ZeRO | "DeepSpeed 版的 FSDP" | 零冗余优化器，三个阶段：分片优化器（阶段 1）、+ 梯度（阶段 2）、+ 参数（阶段 3） |
| All-reduce | "跨 GPU 求平均" | 一种集合操作，每张 GPU 最终得到所有 GPU 输入的总和（或平均）——通常实现为环形 all-reduce |
| All-gather | "从所有 GPU 收集" | 一种集合操作，每张 GPU 最终得到所有 GPU 数据的拼接——在 FSDP 里用于重建完整参数 |
| Reduce-scatter | "求和并分发" | 一种集合操作，归约（求和）数据并把不同块分散到不同 GPU——在 FSDP 里用于梯度分片 |
| 混合精度 | "用半精度训练" | 前向/反向用 FP16/BF16、优化器状态用 FP32——省约 25% 显存而非 50%，因为优化器是大头 |
| 流水线气泡 | "流水线里的空闲时间" | GPU 闲着等上一阶段数据的时间占比——用更多微批次来减少 |
延伸阅读
- Rajbhandari et al., 2020 -- "ZeRO: Memory Optimizations Toward Training Trillion Parameter Models" -- 定义了三个分片阶段的 DeepSpeed ZeRO 论文
- Shoeybi et al., 2020 -- "Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism" -- NVIDIA 为 transformer 设计的张量并行
- Narayanan et al., 2021 -- "Efficient Large-Scale Language Model Training on GPU Clusters Using Megatron-LM" -- 结合数据、张量、流水线的 3D 并行
- Zhao et al., 2023 -- "PyTorch FSDP: Experiences on Scaling Fully Sharded Data Parallel" -- PyTorch 的原生 FSDP 实现
- Llama 3 Technical Report -- 16,384 张 GPU 训练的 3D 并行细节
- DeepSeek-V3 Technical Report -- MoE 架构如何把训练成本降低一个数量级