分布式训练与显存优化¶

本章导读

大模型训练的第一约束通常不是算力，而是显存：参数、梯度、优化器状态、激活和临时缓冲区共同决定单卡能放多大的模型。
数据并行（Data Parallel，DP）、张量并行（Tensor Parallel，TP）、流水线并行（Pipeline Parallel，PP）、序列并行（Sequence Parallel，SP）、专家并行（Expert Parallel，EP）解决的是不同维度的切分问题，通信量和适用拓扑完全不同。
零冗余优化器（Zero Redundancy Optimizer，ZeRO）与完全分片数据并行（Fully Sharded Data Parallel，FSDP）把“每张卡都存全量状态”改成“按需 all-gather”，用通信换显存；理解分片对象和通信次数，才能判断某个配置是否可行。
训练吞吐用模型浮点利用率（Model FLOPs Utilization，MFU）衡量；估算训练时间只需 \(6ND\) 与峰值算力，但真实瓶颈常在通信、重计算与故障恢复。

核心要点

混合精度 + AdamW 全参训练约需 16 字节/参数：bf16 权重 2、bf16 梯度 2、fp32 master 4、Adam 一阶矩 4、二阶矩 4。
7B 模型全参状态约 112GB，70B 约 1.12TB；单卡放不下时必须组合 TP/PP/ZeRO-3。
TP 每层有多次 all-reduce，通信密集，通常限制在单机 8 卡内；PP 只在 stage 边界传激活，适合跨机；DP 通信量随参数量线性增长。
ZeRO-1 分片优化器状态，ZeRO-2 再分片梯度，ZeRO-3 再分片参数；FSDP 是 ZeRO-3 的工程实现，可按层 all-gather。
激活重计算用约 33% 额外前向计算换显存，是长序列训练能跑起来的常用手段。
训练 FLOPs 约 \(6ND\)（前向 \(2ND\)、反向 \(4ND\)）；MFU 把实测吞吐折算成峰值算力的比例，预训练常见 35%–50%。
通信瓶颈由“通信量/计算量”与拓扑共同决定：NVLink 域内做 TP，InfiniBand 域间做 DP/PP。
容错的核心是可恢复的分布式检查点（checkpoint）：保存分片参数、优化器、RNG、数据加载器状态，并定期校验。

1. 显存构成¶

1.1 五类显存¶

设模型参数量为 \(P\)，参数精度字节数为 \(b_{\text{param}}\)，则：

\[M_{\text{param}}=P\cdot b_{\text{param}},\qquad M_{\text{grad}}=P\cdot b_{\text{grad}}\]

其中 \(M_{\text{param}}\) 是权重显存，\(M_{\text{grad}}\) 是梯度显存。bf16 下两者各为 \(2P\) 字节。

Adam/AdamW 为每个参数维护一阶矩 \(m\) 与二阶矩 \(v\)，通常用 fp32 存储：

\[M_{\text{Adam}}=8P\ \text{bytes}\quad(m:4P,\ v:4P)\]

混合精度训练还需要一份 fp32 的 master weights（主权重），用于吸收小梯度更新，避免 bf16/fp16 直接更新时被舍入：

\[M_{\text{master}}=4P\ \text{bytes}\]

因此“bf16 权重 + bf16 梯度 + fp32 master + Adam m/v”的总状态为：

\[M_{\text{state}}=2P+2P+4P+4P+4P=16P\ \text{bytes}\]

若用纯 fp32 训练，则是 \(4P\) 参数 + \(4P\) 梯度 + \(8P\) 优化器 = \(16P\)，总字节数相同，但纯 fp32 没有 master 副本，前向也更慢。混合精度的收益主要在计算速度与激活显存，而不是优化器状态。

模型	参数量	bf16 权重	bf16 梯度	fp32 master	Adam m/v	合计（估算）
7B	7×10⁹	14 GB	14 GB	28 GB	56 GB	112 GB
13B	1.3×10¹⁰	26 GB	26 GB	52 GB	104 GB	208 GB
70B	7×10¹⁰	140 GB	140 GB	280 GB	560 GB	1120 GB
175B	1.75×10¹¹	350 GB	350 GB	700 GB	1400 GB	2800 GB

表中按十进制 GB 估算，未计激活、临时缓冲与显存碎片；实际至少再留 20%–30% 余量。可以看到 7B 全参训练在单张 80GB 卡上放不下，必须分片或卸载。

1.2 激活显存¶

激活（activation）是前向过程中保存、反向传播要用的中间张量。以 Transformer 每层每样本的粗略公式（Megatron-LM 常用，单位字节，fp16）：

\[M_{\text{act per layer}}\approx s\cdot b\cdot h\left(34+5\frac{a\cdot s}{h}\right)\]

其中 \(s\) 是序列长度，\(b\) 是 micro-batch，\(h\) 是隐藏维度，\(a\) 是注意力头数。括号中的 \(34\) 来自 QKV、MLP 中间层、LayerNorm 等；\(5as/h\) 来自注意力分数矩阵（若 FlashAttention 不物化分数，这一项会大幅下降）。

以 7B（\(h=4096\)、\(a=32\)、32 层）、\(s=4096\)、\(b=1\) 估算：每层每样本约 \(4096\times4096\times(34+160)\approx3.26\) GB，32 层约 104GB。这是“不用 FlashAttention、不做重计算”的最坏情况；实际使用 FlashAttention 后分数不落盘，激活可降到 20–40GB；再做全层激活重计算，只保留每层输入，约 \(2sbh\) 字节/层，总计 1–2GB，另加反向重算时的工作区（量级估算，随实现变化）。

激活显存与 batch、序列长度近似线性（注意力分数项对 \(s\) 是平方），因此长上下文训练首先撞上的是激活墙，而不是参数墙。

1.3 临时缓冲区与碎片¶

通信缓冲区：NCCL 的 all-reduce/all-gather 需要连续显存；bucket 越大越省启动开销，但峰值越高。
cuBLAS/cuDNN workspace：GEMM、FlashAttention 反向会申请临时空间，长序列时可达数 GB。
显存碎片：反复分配不同大小张量会造成碎片，PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True 可缓解。
CUDA context 与 kernel：每张卡约 0.5–1GB 固定开销（估算）。
offload 的 pinned memory：CPU 端固定页内存不计入 GPU 显存，但会占用主机内存与 PCIe 带宽。

2. 并行策略¶

2.1 数据并行（DP）¶

每张卡保存完整模型，把全局 batch 切成 \(N\) 份；各卡独立前向/反向，最后对梯度做 all-reduce 求平均。设参数量 \(P\)、精度 2 字节，则每步梯度通信量约 \(2P\) 字节（ring all-reduce 的流量为 \(2P(N-1)/N\)，接近 \(2P\)）。

优点：实现简单、计算通信比高、适合单机多卡与多机扩展。缺点：每卡都要存全量参数/梯度/优化器状态，7B 以上很快放不下；global batch 随卡数线性增大，需要调整学习率与 warmup；梯度 all-reduce 在跨机时受 IB 带宽限制。

梯度累积（gradient accumulation）用小 micro-batch 模拟大 batch，不增加显存中的激活峰值，但会增加单步时间；它与 DP 正交，常一起使用。

2.2 张量并行（TP）¶

TP 把单层矩阵乘按列/行切分到多卡。以 MLP 的 \(Y=XW\)、\(W\in\mathbb{R}^{h\times4h}\) 为例：

列并行：\(W=[W_1,W_2]\)，各卡算 \(XW_i\)，输出按列拼接；无需通信。
行并行：\(W=[W_1;W_2]\)，各卡算 \(X_iW_i\)，输出需要 all-reduce 求和。

Transformer 一层通常包含 attention 的 QKV 列并行、输出投影行并行、MLP 上投影列并行、下投影行并行，因此前向有 2 次 all-reduce，反向有 2 次，共 4 次/层；每次通信量约 \(b\cdot s\cdot h\) 个元素，合计约 \(8bsh\) 字节/层（估算，随切分方式变化）。TP 度越大，单卡计算越少、通信占比越高，通常不超过单机 8 卡，且必须走 NVLink。

TP 的收益是单层就能切，适合“层数不多但单层很宽”的模型；缺点是通信频繁、对拓扑敏感，跨机 TP 往往得不偿失。

2.3 流水线并行（PP）¶

PP 按层把模型切成 \(p\) 个 stage，每个 stage 放在不同设备上；micro-batch 依次流过 stage。经典 GPipe 先做完全部前向再反向，气泡（bubble）比例为：

\[\text{bubble fraction}=\frac{p-1}{m+p-1}\approx\frac{p-1}{m}\]

其中 \(m\) 是 micro-batch 数；\(m\) 远大于 \(p\) 时近似为 \((p-1)/m\)。1F1B（one forward one backward）在稳态下交错执行前向与反向，气泡约为 \((p-1)/m\)，并减少激活占用；交错式 1F1B 把每个 stage 再切成 \(v\) 个虚拟 stage，气泡进一步降到约：

\[\frac{p-1}{m\cdot v}\]

PP 只在 stage 边界传激活，每个 micro-batch 前向传 \(bsh\) 个元素、反向传 \(bsh\) 个梯度元素（bf16 下合计约 \(4bsh\) 字节），适合跨机；难点是 stage 负载均衡（按 FLOPs 而非层数切）、micro-batch 数要足够多才能掩盖气泡、以及反向依赖导致的调度复杂。

2.4 序列并行（SP）¶

SP 把序列维度切到多卡，主要针对 LayerNorm、Dropout、残差等“按 token 独立”的操作，避免 TP 后每张卡重复存完整序列的激活。典型做法（Megatron 的 sequence parallel）：

进入 attention/MLP 前对序列做 all-gather，得到完整 QKV 或输入；
计算完 attention/MLP 后对输出做 reduce-scatter，回到序列分片状态。

通信模式与 TP 类似，但作用在序列维，能显著降低长序列激活；对超长上下文（32k–1M），还会用 Ring Attention、Ulysses 等等价切分注意力的方法，把 \(s^2\) 的注意力计算分散到多卡。

2.5 专家并行（EP）¶

MoE（Mixture of Experts）把 \(E\) 个专家分布到不同卡，每个 token 经路由选 top-\(k\) 个专家。EP 的核心通信是 all-to-all：dispatch 把 token 发到对应专家所在卡，combine 把结果收回，每层两次，通信量约 \(O(bsh\cdot k)\)。

难点：负载不均导致“热点专家”等待；需要 capacity factor（每专家最多接收 token 数）与负载均衡辅助损失；all-to-all 对网络带宽和延迟敏感；EP 通常与 DP/TP 组合，专家参数量大但每 token 只激活一部分，适合“总参数量大、单 token 计算量可控”的稀疏模型。

2.6 并行策略对比¶

策略	切分维度	每步通信	通信频率	适用范围	主要瓶颈
DP	batch	梯度 all-reduce，约 \(2P\) 字节	每步一次	全局，跨机	参数量、跨机带宽
TP	隐藏维/头	每层 4 次 all-reduce，约 \(8bsh\) 字节	每层	单机 8 卡内	NVLink 带宽、延迟
PP	层	stage 边界激活，前向+反向约 \(4bsh\) 字节/micro-batch	每 micro-batch	跨机	气泡、负载均衡
SP	序列	all-gather + reduce-scatter	每层	长序列、与 TP 配合	通信量、实现复杂度
EP	专家	每层 2 次 all-to-all	每层	MoE、多机	负载不均、all-to-all

3. ZeRO 与 FSDP¶

3.1 三个 stage 分片什么¶

ZeRO（Zero Redundancy Optimizer）把 DP 中每卡重复存储的状态切开，需要时再收集：

阶段	分片对象	每卡显存（混合精度，估算）	额外通信
基线 DP	无	\(16P\)	梯度 all-reduce，\(2P\)
ZeRO-1	优化器状态（master + m + v）	\(4P+12P/N\)	reduce-scatter 梯度 + all-gather 参数，共约 \(2P\)
ZeRO-2	优化器状态 + 梯度	\(2P+14P/N\)	reduce-scatter 梯度 + all-gather 参数，共约 \(2P\)
ZeRO-3	优化器状态 + 梯度 + 参数	\(16P/N\)	2 次参数 all-gather + 梯度 reduce-scatter，约 \(3P\)

其中 \(N\) 是 DP 组大小。以 7B、8 卡为例：基线 112GB，ZeRO-1 约 38.5GB，ZeRO-2 约 26.3GB，ZeRO-3 约 14GB（均不含激活）。

70B、64 卡：ZeRO-1 约 293GB，ZeRO-2 约 155GB，二者都超过 80GB；ZeRO-3 约 17.5GB 才能放下。这说明“加卡”不能解决 ZeRO-2 的参数复制问题——参数和梯度仍然每卡一份，必须靠 TP/PP 或 ZeRO-3 真正切开。

3.2 通信开销¶

ZeRO-1/2 的通信量与普通 DP 同阶（reduce-scatter 代替 all-reduce），但可以把梯度分片通信与反向计算重叠；ZeRO-3 每层前向要 all-gather 参数、反向要再次 all-gather 并 reduce-scatter 梯度，通信量约为 DP 的 1.5 倍（\(3P\) vs \(2P\)，估算），对跨机带宽更敏感。

ZeRO-3 的显存收益来自“参数只在用时存在”：一层算完就释放，峰值约为单层参数量加分片状态；代价是频繁的小 all-gather，需要把 bucket 调大以摊薄延迟。ZeRO-Offload 进一步把优化器状态放到 CPU，ZeRO-Infinity 放到 NVMe，用 PCIe/NVMe 带宽换显存。

3.3 FSDP 与 ZeRO-3 的关系¶

完全分片数据并行（Fully Sharded Data Parallel，FSDP）是 PyTorch 原生实现，思想与 ZeRO-3 一致：参数按 flat parameter 分片，前向按层 all-gather，计算后释放，反向再次 all-gather 并 reduce-scatter 梯度。区别主要在工程细节：

FSDP 以 FlatParameter 为单位分片，支持 FULL_SHARD、SHARD_GRAD_OP（类似 ZeRO-2）、HYBRID_SHARD（节点内全分片、节点间复制，即 HSDP）。
FSDP 可与 TP/PP 组合；device_mesh 让多维并行配置更清晰。
FSDP 的 checkpoint 需要 FullStateDictConfig 或分片保存，不能直接 torch.save(model.state_dict()) 后按单卡加载。
ZeRO 是 DeepSpeed 的实现，生态里还有 ZeRO-Infinity、MoE 支持、通信压缩等。

3.4 3D 并行组合¶

典型配置：70B 模型、512 张 80GB H100，TP=8、PP=4、DP=16。每 stage 约 17.5B 参数，TP=8 后每卡约 2.19B 参数，按 16 字节/参数估算约 35GB；再加激活与通信缓冲约 45–50GB，可落在 80GB 内。global batch = DP 16 × micro-batch 1 × gradient accumulation 16 = 256 条序列（示例）。若再叠加 ZeRO-1，优化器状态还能按 DP 继续分片，显存更宽裕。

组合原则：

先让单层放得下：TP 度覆盖单层参数与激活；
再让层数分得均：PP 按 FLOPs 均衡，stage 数不超过 micro-batch 数；
最后用 DP 扩吞吐：DP 度决定 global batch 与梯度 all-reduce 规模；
长序列额外开 SP/CP（context parallel），MoE 额外开 EP。

4. 混合精度¶

4.1 FP16、BF16、FP8¶

格式	位宽	指数位	尾数位	动态范围	是否需要 loss scaling	典型用途
FP32	32	8	23	约 \(10^{\pm38}\)	否	master weights、优化器状态
TF32	19	8	10	同 FP32	否	A100 上替代 FP32 做 GEMM
FP16	16	5	10	约 \(6\times10^4\) 上限	是	推理、部分训练
BF16	16	8	7	约 \(10^{\pm38}\)	否	训练默认
FP8 E4M3	8	4	3	约 \(\pm448\)	是（缩放）	前向权重/激活
FP8 E5M2	8	5	2	约 \(\pm57344\)	是（缩放）	反向梯度

FP16 尾数多、精度高，但最大可表示数只有 65504，小梯度容易下溢，需要损失缩放（loss scaling）：前向算出的 loss 乘一个缩放因子 \(S\)，反向得到放大 \(S\) 倍的梯度，更新前再除回去。动态损失缩放会根据是否出现 inf/NaN 调整 \(S\)：连续若干步成功则增大，溢出则减半并跳过该步。

# PyTorch 原生 AMP：自动维护 loss scale（示意）
amp_dtype = torch.bfloat16                      # fp16 时改为 torch.float16
scaler = torch.cuda.amp.GradScaler(init_scale=2**16, growth_interval=2000,
                                   enabled=(amp_dtype == torch.float16))
for x, y in loader:
    opt.zero_grad(set_to_none=True)
    with torch.autocast(device_type="cuda", dtype=amp_dtype):
        loss = model(x, labels=y).loss
    scaler.scale(loss).backward()          # 仅 fp16 需要；bf16 时 scaler 已禁用
    scaler.unscale_(opt)                   # 裁剪梯度前必须先 unscale
    torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
    scaler.step(opt)
    scaler.update()

BF16 指数位与 FP32 相同，动态范围足够，通常不需要 loss scaling；代价是尾数只有 7 位，舍入误差比 FP16 大，但在深度学习中被证明足够。BF16 已成为预训练默认，FP16 更多用于推理或旧硬件。

4.2 FP8 训练¶

FP8 训练的核心是缩放（scaling）：把张量按 per-tensor 或 per-block 的 amax 映射到 FP8 的可表示范围。前向的权重/激活用 E4M3（精度高、范围小），反向梯度用 E5M2（范围大、精度低），并维护 amax 历史做延迟缩放（delayed scaling）。H100 的 Transformer Engine 把缩放、cast、GEMM 融合，FP8 理论吞吐约为 BF16 的 2 倍。

风险：离群值（outlier）会拉大 amax，导致其余值量化过粗；需要 per-channel/per-block 缩放、残差与 LayerNorm 保持高精度、定期用 BF16 校准。FP8 不是“免费加速”，在中小模型或短序列上，缩放开销与精度回退可能抵消收益。

4.3 master weights 为什么需要¶

假设权重是 bf16，学习率 \(10^{-5}\)，参数值约 \(0.1\)，更新量约 \(10^{-6}\)。bf16 在 \(0.1\) 附近的间隔约 \(10^{-3}\)，更新会被直接舍入掉。保留一份 fp32 master，更新在 fp32 上累加，再把结果 cast 回 bf16 用于前向，可以避免“权重被冻住”。这也是为什么混合精度训练的状态字节数仍是 16P。

5. 激活重计算与 offload¶

5.1 梯度检查点¶

激活重计算（gradient checkpointing / activation recomputation）在前向时只保存每个 checkpoint 段的输入，不保存段内中间张量；反向时重新前向该段，重建中间量再算梯度。全层检查点把激活从 \(O(L)\) 降到 \(O(1)\) 量级（只存层输入），代价是每个 step 多一次前向：

\[C_{\text{train}}\approx 6ND\ \rightarrow\ C_{\text{ckpt}}\approx 8ND\]

因为前向是 \(2ND\)，反向是 \(4ND\)，重算一次前向增加 \(2ND\)，总计算增加约 33%。若用“每 \(\sqrt{L}\) 层设一个 checkpoint”的 sqrt 策略，显存为 \(O(\sqrt{L})\)，重算量更小。

from torch.utils.checkpoint import checkpoint

class TransformerBlock(nn.Module):
    def forward(self, x):
        # 训练时只存输入，反向重算 attention/MLP 内部激活
        if self.training and self.grad_checkpointing:
            return checkpoint(self._forward, x, use_reentrant=False)
        return self._forward(x)

    def _forward(self, x):
        x = x + self.attn(self.norm1(x))
        x = x + self.mlp(self.norm2(x))
        return x

选择策略：全层检查点最省显存但最慢；选择性检查点只重算注意力（或只重算 MLP），通常能省 50%–70% 激活而只增加 10%–15% 计算（估算）。use_reentrant=False 对新版 PyTorch 的自动混合精度、非张量输入更友好。

5.2 Offload¶

优化器 offload：把 master weights、m、v 放 CPU，GPU 只保留 bf16 参数与梯度；更新时把梯度搬到 CPU 算 Adam，再把参数拷回。DeepSpeed ZeRO-Offload 是典型实现。PCIe Gen4 x16 双向约 64GB/s，远低于 HBM 的 3TB/s，频繁搬运会让 GPU 空等。
参数/激活 offload：ZeRO-Infinity 用 NVMe 存优化器状态，容量大但带宽只有数 GB/s，适合“显存极度紧张、吞吐不敏感”的场景。
计算与搬运重叠：用独立 CUDA stream 预取下一层参数，能把部分带宽隐藏；但 offload 的训练吞吐通常只有全 GPU 的 30%–60%（估算）。

offload 的取舍是“用时间换显存”：当模型只差一点放不下时很有效；当差距是数倍时，吞吐会不可接受，应优先考虑 TP/PP/ZeRO-3。

6. 计算量估算与 MFU¶

6.1 6ND 从哪来¶

对参数量 \(P\)、训练 token 数 \(D\) 的稠密 Transformer，忽略 attention 的 \(s^2\) 项，前向每个 token 约 \(2P\) FLOPs（每个参数一次乘加），反向约为前向的 2 倍（计算对权重和对输入的梯度），因此：

\[C_{\text{train}}\approx 2ND+4ND=6ND\]

其中 \(N\) 是参数量（习惯用 \(N\) 而非 \(P\)），\(D\) 是训练 token 数。长序列时注意力项不可忽略：

\[C\approx 6ND+12\,n_{\text{layer}}\,s^{2}d\]

其中 \(n_{\text{layer}}\) 是层数，\(s\) 是序列长度，\(d\) 是隐藏维度；系数 12 来自前向两次 \(s^2d\) 矩阵乘、反向约两倍，causal 注意力可近似减半。\(s=4096\)、\(d=4096\)、32 层时，注意力项约占 6ND 的 15%（估算），序列到 128k 时会成为主导。

6.2 MFU¶

\[\text{MFU}=\frac{C_{\text{achieved}}}{\text{peak FLOPs}\times T}=\frac{6ND}{P_{\text{peak}}\cdot T}\]

其中 \(T\) 是训练时间，\(P_{\text{peak}}\) 是硬件峰值算力（bf16 dense）。MFU 把“实际吞吐”折算成理论峰值的比例，便于跨硬件比较。

硬件	BF16 峰值（dense，估算）	显存/带宽	典型 MFU
A100 SXM 80GB	312 TFLOPS	80GB HBM2e，2.0 TB/s	35%–45%（预训练）
H100 SXM 80GB	989 TFLOPS	80GB HBM3，3.35 TB/s	40%–50%（预训练）
H800 SXM	989 TFLOPS	80GB，NVLink 400GB/s	35%–45%（受互联限制）
消费级 RTX 4090	165 TFLOPS（bf16，估算）	24GB，1.0 TB/s	10%–25%（单卡/小集群）

以 7B、1T token、1024 张 H100、MFU 40% 估算：

\[T=\frac{6\times7\times10^9\times10^{12}}{1024\times989\times10^{12}\times0.4}\approx1.04\times10^5\ \text{s}\approx29\ \text{小时}\]

70B、15T token 在同样集群上约 \(180\) 天（估算），所以实际需要数千到上万张卡，或更长时间。数量级比精确值更重要：它决定了“这个训练计划是否现实”。

6.3 通信与 roofline¶

一个算子是否受通信限制，看它的算术强度（FLOPs/byte）与机器的“算力/带宽”比值。H100 的 bf16 峰值/显存带宽约 \(989\text{ TFLOPS}/3.35\text{ TB/s}\approx295\) FLOPs/byte；NVLink 900GB/s 对应约 265 TFLOPS 的有效算力，IB 400Gb/s（50GB/s）只有约 15 TFLOPS。这解释了：

TP 的 all-reduce 走 NVLink，仍可能达到 50% 以上峰值；跨 IB 做 TP 会被带宽拖死。
DP 的梯度 all-reduce 量随 \(P\) 增长，70B 模型每步 140GB 量级，跨机必须用分层 all-reduce、reduce-scatter、通信压缩或 ZeRO。
PP 的边界激活只有约 \(4bsh\) 字节/micro-batch，跨机成本低，但气泡会浪费算力。

7. 拓扑、通信与容错¶

7.1 常见拓扑¶

NVLink/NVSwitch：单机 8 卡全互联。H100 每卡 NVLink 4 约 900GB/s 双向，NVSwitch 让 8 卡任意两两通信不经过 CPU；A100 为 600GB/s。TP/SP 应限制在这个域内。
InfiniBand（IB）：跨机主力，NDR 400Gb/s ≈ 50GB/s/端口，HDR 200Gb/s ≈ 25GB/s。大规模集群用 rail-optimized 胖树（fat-tree），保证任意两卡间有对等带宽；也有 DragonFly 等低直径拓扑。
以太网 RoCEv2：成本低，但需要 PFC/ECN 调优，尾延迟与拥塞控制不如 IB。
PCIe：GPU-CPU/GPU-GPU（无 NVLink 时）通道，Gen4 x16 约 64GB/s，Gen5 约 128GB/s；offload 与跨卡通信都受它限制。

NCCL 会根据拓扑选择 ring 或 tree；大规模 all-reduce 常用分层算法（节点内 reduce-scatter + 节点间 all-reduce + 节点内 all-gather）。通信与计算重叠（overlap）是提升 MFU 的关键：把 all-gather 放在上一层的反向计算中预取，把梯度 reduce-scatter 与下一层反向重叠。

7.2 容错与断点续训¶

万卡训练中，硬件故障是常态：GPU 掉卡、ECC 错误、NVLink 降速、网卡抖动、静默数据损坏（silent data corruption，SDC）。工程要点：

分片 checkpoint：每个 rank 只写自己的分片，避免全量 gather 到 rank0 造成内存峰值；保存 optimizer、scheduler、随机数生成器（RNG）、data loader 的 epoch/step。
异步 checkpoint：先把分片拷到 CPU/远端存储，再继续训练；checkpoint 期间用双缓冲避免阻塞。
原子写与版本号：先写临时目录，完成后 rename；保留最近 \(k\) 个 checkpoint，防止写坏。
故障检测：NCCL watchdog、心跳、loss/梯度范数异常检测；发现 NaN 或 loss 突增时回滚到最近 checkpoint，并隔离可疑节点。
弹性训练：torchelastic/JobSet 支持 worker 重启后重新 rendezvous，改变 world size 后要重新切分数据与并行组。
checkpoint 间隔：按平均无故障时间（Mean Time Between Failures，MTBF）与 checkpoint 成本权衡，大集群常 30–60 分钟一次（估算）；间隔太长，故障后重算损失大；太短，I/O 拖慢训练。

8. 代码示例¶

8.1 DDP 最小示例¶

import os, torch, torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP
from torch.utils.data import DataLoader, DistributedSampler

# 启动：torchrun --nproc_per_node=8 ddp.py
dist.init_process_group("nccl")
rank = int(os.environ["LOCAL_RANK"]); torch.cuda.set_device(rank)
model = DDP(MyModel().to(rank), device_ids=[rank])
ds = MyDataset(); sampler = DistributedSampler(ds, shuffle=True)
loader = DataLoader(ds, batch_size=8, sampler=sampler)
opt = torch.optim.AdamW(model.parameters(), lr=3e-4)
for epoch in range(10):
    sampler.set_epoch(epoch)                 # 每个 epoch 的切分不同
    for x, y in loader:
        x, y = x.to(rank), y.to(rank)
        opt.zero_grad(set_to_none=True)
        loss = model(x, labels=y).loss
        loss.backward(); opt.step()          # DDP 反向中自动 all-reduce 梯度
dist.destroy_process_group()

8.2 激活检查点¶

from torch.utils.checkpoint import checkpoint

class CheckpointedBlock(torch.nn.Module):
    def __init__(self, attn, mlp, norm1, norm2):
        super().__init__()
        self.attn, self.mlp, self.norm1, self.norm2 = attn, mlp, norm1, norm2
        self.grad_checkpointing = True

    def _forward(self, x):
        x = x + self.attn(self.norm1(x))
        x = x + self.mlp(self.norm2(x))
        return x

    def forward(self, x):
        if self.training and self.grad_checkpointing:
            return checkpoint(self._forward, x, use_reentrant=False)
        return self._forward(x)

9. 面试高频问题¶

1. 7B 模型全参训练需要多少显存？如何拆解？¶

混合精度 + AdamW 约 16 字节/参数：bf16 权重 2、梯度 2、fp32 master 4、m/v 各 4。
7B 约 112GB，70B 约 1.12TB；还要加激活、通信缓冲和碎片，实际预留 20%–30%。
单卡 80GB 放不下 7B，需要 ZeRO-3/FSDP、TP 或 offload。
纯 fp32 也是 16 字节/参数，混合精度省的是计算和激活，不是优化器状态。

2. DP、TP、PP 的通信量和适用场景有什么区别？¶

DP 每步 all-reduce 约 \(2P\) 字节，计算通信比高，适合跨机扩吞吐，但每卡存全量状态。
TP 每层 4 次 all-reduce、约 \(8bsh\) 字节，通信密集，适合单机 NVLink 内 8 卡。
PP 只在 stage 边界传激活，约 \(4bsh\) 字节/micro-batch，适合跨机，但有气泡和负载均衡问题。
实际用 3D 并行：TP 放单机、PP 跨机切层、DP 扩全局 batch。

3. ZeRO-1/2/3 分别分片什么？通信开销如何变化？¶

Stage 1 分片优化器状态（master + m + v），每卡 \(4P+12P/N\)；Stage 2 再分片梯度，\(2P+14P/N\)；Stage 3 再分片参数，\(16P/N\)。
Stage 1/2 通信与 DP 同阶（reduce-scatter），Stage 3 每层 all-gather 参数，通信约为 DP 的 1.5 倍。
Stage 3 显存最省，但小 all-gather 频繁，需要 bucket 与重叠；跨机时可能被带宽拖慢。
70B 在 64 张 80GB 上，ZeRO-2 仍需每卡 140GB 参数+梯度，根本放不下，必须 ZeRO-3 或 TP/PP。

4. FSDP 和 ZeRO-3 是什么关系？¶

思想相同：参数分片、前向按需 all-gather、反向 reduce-scatter 梯度。
FSDP 是 PyTorch 原生实现，支持 FULL_SHARD、SHARD_GRAD_OP、HYBRID_SHARD；ZeRO 是 DeepSpeed 实现。
FSDP 以 FlatParameter 为单位，按层 wrap 很重要；wrap 太粗会增大峰值，太细会增加通信。
checkpoint 要保存分片状态，或用 FullStateDictConfig 聚合，不能直接单卡加载。

5. 混合精度为什么需要 master weights？BF16 还需要 loss scaling 吗？¶

bf16 在参数附近的量化间隔大于小学习率更新量，直接更新会被舍入掉。
fp32 master 累积更新，再把结果 cast 回 bf16 前向，保证更新不丢失。
BF16 指数位与 FP32 相同，动态范围足够，通常不需要 loss scaling。
FP16 需要动态 loss scale；FP8 需要 per-tensor/per-block scaling 与 amax 历史。

6. 激活重计算的代价是什么？什么时候不该用？¶

前向只存段输入，反向重算中间激活；计算从 \(6ND\) 增到约 \(8ND\)，多 33%。
显存从 \(O(L)\) 降到 \(O(1)\) 或 \(O(\sqrt{L})\)，是长序列训练的关键手段。
计算已经受限于算力、显存充足时，重计算会白白降低吞吐。
选择性重算（只重算 attention）通常性价比最高；优先用 FlashAttention 降低激活，再考虑重算。

7. 训练 FLOPs 和 MFU 怎么算？常见值是多少？¶

稠密模型约 \(6ND\)：前向 \(2ND\)、反向 \(4ND\)；长序列加 \(12n s^2 d\) 的注意力项。
MFU = 实测 FLOPs /（峰值算力 × 时间），用 bf16 dense 峰值作分母。
预训练常见 35%–50%，微调 20%–35%，小集群/消费卡更低；FP8 峰值不同，不能直接横比。
估算训练时间时先算 \(6ND\)，再除以“卡数 × 峰值 × MFU”。

8. 为什么 TP 通常不超过单机 8 卡？¶

TP 每层多次 all-reduce，通信量与激活大小成正比，计算量却随 TP 度下降。
跨机 TP 走 IB，带宽比 NVLink 低一个数量级，通信时间会超过计算时间。
TP 度超过单机 GPU 数后，层内通信必须跨机，延迟和拥塞都会放大。
需要更大模型时，优先 PP/DP 跨机，TP 保持在 NVLink 域内。

9. 大规模训练如何做容错和断点续训？¶

分片保存参数、优化器、scheduler、RNG、数据加载器状态；异步写、原子 rename、保留多版本。
心跳与 NCCL watchdog 检测掉卡/挂起；loss/梯度异常触发回滚与节点隔离。
checkpoint 间隔按 MTBF 与写入成本权衡，大集群常 30–60 分钟。
弹性训练要处理 world size 变化后的并行组重建与数据重切分；SDC 需要定期校验。

10. 长上下文训练为什么需要序列并行？它和 TP 的区别是什么？¶

长序列的激活与注意力计算随 \(s\) 增长，TP 只切隐藏维，序列维仍然每卡一份。
SP 切序列维，LayerNorm/Dropout 等按 token 操作可并行，attention 前 all-gather、后 reduce-scatter。
更长的上下文还需要 Ring Attention/Ulysses 把 \(s^2\) 的注意力分散到多卡，通信模式不同。
SP 降低激活而不减参数；TP 降低参数与计算；两者常组合使用。

本文档为面试复习资料，公式与数字建议结合原始论文复核。