真实面经题目 · 原创解析
这道题考察大模型训练显存构成和 ZeRO 的边界。回答要先把 GPU 显存里的主要训练状态拆清楚：参数、梯度、优化器状态、激活、通信/临时 buffer 和碎片；再说明 ZeRO 主要分片的是模型状态，Stage 1/2/3 分别处理优化器状态、梯度和参数；最后解释 activation 为什么通常用 checkpoint/recompute：它是 batch/sequence 相关的短生命周期中间结果，反向依赖层级顺序，分页到 CPU/NVMe 会引入高带宽低延迟瓶颈，而重计算能用额外 FLOPs 换显存，通常更可控。
我会先区分训练显存里的几类东西。第一类是模型参数，通常有 FP16/BF16 参与前向反向，也可能有 FP32 master weight。第二类是梯度，反向传播时每个参数对应梯度。第三类是优化器状态，比如 Adam 的一阶动量和二阶动量，很多实现还保留 FP32 参数副本，这部分常常比参数本身还大。第四类是 activation，也就是前向中为了反向求梯度保存的中间输出。除此之外还有通信 bucket、临时 workspace、CUDA kernel buffer、KV/attention 中间张量和显存碎片。 DeepSpeed ZeRO 主要解决的是模型状态冗余。普通数据并行里每张卡都有完整参数、完整梯度和完整优化器状态，显存随模型规模重复占用。ZeRO Stage 1 分片 optimizer states，Stage 2 进一步分片 gradients，Stage 3 再分片 parameters，需要时通过通信临时 all-gather。ZeRO-Offload 或更激进方案还能把部分模型状态 offload 到 CPU/NVMe，但核心对象仍然是参数、梯度、优化器状态这类相对稳定、可分片、跨 step 存在的模型状态。 Activation 的性质不一样。它不是固定模型状态，而是每个 micro-batch、sequence length、layer、训练配置动态产生的中间结果。前向产生，反向按依赖顺序消费，生命周期短但峰值很高，尤其长序列、大 batch、attention 和 MLP 中间张量会拉高显存。Activation 也不能像参数那样简单按 rank 持久分片，因为反向需要对应样本、对应层的中间值，且形状和生命周期跟计算图强绑定。 为什么通常用 checkpoint/recompute，而不是像参数或优化器状态一样分页？核心是带宽和调度。把 activation 写到 CPU 或 NVMe，再在反向时读回来，会把训练关键路径变成频繁的大量 H2D/D2H 或存储 IO，PCIe/NVMe 带宽和延迟通常跟不上 GPU 算力，容易让 step time 被数据搬运支配。activation 数量多、粒度碎、生命周期短，分页管理本身也复杂，很容易出现同步等待和显存碎片。 Checkpoint/recompute 的思路是只保存少量关键边界激活，反向时重新跑一部分前向来恢复中间值。它用额外计算换显存，适合大模型训练的原因是 GPU 算力相对充足，而跨设备搬运 activation 的代价常常更高且更不稳定。重计算还能保持数据在 GPU 内部流动，调度更简单，和 pipeline/tensor parallel 也更容易组合。 当然这不是说 activation 永远不能 offload。某些系统可以做 activation offloading、CPU checkpoint 或分层存储，但那通常是内存极端不足时的补充，代价是更复杂的调度和更慢的训练。实际工程会组合 ZeRO 分片模型状态、activation checkpoint、混合精度、sequence parallel、flash attention 和 micro-batch 调整。验证时要看 peak memory、step time、MFU、通信等待、recompute overhead、OOM 率和长序列稳定性，而不是只看理论显存公式。
训练显存主要包括参数、梯度、优化器状态、activation、通信 bucket、临时 workspace 和碎片。大模型训练不能只盯参数量。
ZeRO 主要分片模型状态。Stage 1 分片优化器状态，Stage 2 分片梯度，Stage 3 分片参数，目标是减少数据并行中的冗余副本。
参数、梯度和优化器状态与模型结构绑定，跨 step 存在，形状稳定，适合分片、通信聚合和必要时 offload。
Activation 是每个 micro-batch 前向产生、反向消费的中间结果，生命周期短、数量多、形状依赖 batch/sequence/layer，峰值显存很高。
Activation 分页会把大量碎片化中间结果搬到 CPU/NVMe 并在反向读回，容易受 PCIe/NVMe 带宽、延迟、同步和调度复杂度限制。
Checkpoint/recompute 只保存关键边界，反向时重跑部分前向，用更多 FLOPs 换更少显存。对 GPU 训练来说通常比频繁跨设备搬运更可控。
实际训练常把 ZeRO、activation checkpoint、混合精度、FlashAttention、sequence parallel、micro-batch 调整和 offload 策略组合使用。
需要同时观察 peak memory、step time、tokens/s、MFU、通信等待、IO 等待、recompute overhead、OOM 率和不同 sequence length 下的稳定性。
Stage 1 分片优化器状态，Stage 2 在此基础上分片梯度，Stage 3 进一步分片参数。Stage 越高显存越省，但通信、参数 gather 和实现复杂度也会上升。
Adam 通常为每个参数维护一阶动量和二阶动量，很多混合精度训练还保留 FP32 master weight。这样优化器相关状态可能达到参数本身的数倍。
取决于 checkpoint 粒度、模型结构和并行策略。粗略来说会增加一部分前向重算 FLOPs，但换来明显显存下降。实际要看 step time、MFU 和是否让 batch/sequence 扩大后总体吞吐更好。
不是。极端显存不足时可以做 activation offload 或 CPU checkpoint，但它通常会引入数据搬运和调度成本。工程上要把它当补充手段，而不是默认替代 recompute。
参数是模型固定状态，可以按 rank 持久分片并在计算前 gather；activation 与样本、层、sequence 和计算图绑定，反向需要精确的中间值，生命周期短且形状动态，简单持久分片并不能直接满足反向依赖。
