很多人对分布式训练的最初想象是"多插几张卡，训练变快"。真上手算过账才知道，单卡的窘境往往不是慢，是根本放不下。这篇把显存这本账拆开算一遍，给你一套开训前就能做的预算方法。
显存里到底装了什么
训练态显存有四大户头：模型参数、梯度、优化器状态、激活值。最容易被低估的是优化器状态——以 Adam 加混合精度为例，每个参数除了 fp16 权重和 fp16 梯度，还要养着 fp32 主权重副本以及一阶、二阶动量。粗算下来：
每参数显存占用（混合精度 + Adam）：
fp16 权重 2 bytes
fp16 梯度 2 bytes
fp32 主权重 4 bytes
一阶动量 fp32 4 bytes
二阶动量 fp32 4 bytes
合计 16 bytes / 参数
代入实际规模感受一下：7B 模型光这四项就是 7×16 = 112GB——单张 80GB 的 A100 放不下，这还没算激活值。这就是为什么"模型参数量"和"训练所需显存"从来不是一比一的关系。
激活值：序列长度一翻倍它先爆
激活值的占用随 batch size 和序列长度线性增长，而且Transformer 的中间层张量很多。粗估公式：激活显存 ≈ 层数 × batch × seq_len × hidden_dim × 常数因子。实践中序列长度翻倍，最先溢出的往往就是这一项。
对应的省钱手段及各自代价：
| 手段 | 省什么 | 付什么 |
|---|---|---|
| 混合精度 | 权重/激活减半 | 数值稳定性要盯着 loss scale |
| 梯度检查点 | 激活大头 | 重算前向，速度降 20-30% |
| 优化器分片 | 优化器状态摊到多卡 | 通信量增加 |
理解了"每种手段省的是哪一户头"，看 DeepSpeed 的 ZeRO 配置就不会一脸懵——ZeRO-1 分片优化器状态，ZeRO-2 加分梯度，ZeRO-3 连参数也分，逐级用通信换显存。
扩卡不等于线性加速
加卡解决容量，但通信会反噬速度。端到端训练时间可以拆成四段：
总时长 = 有效计算 + 数据准备 + 通信同步 + 等待/故障恢复
这不是精确公式，是排查清单。优化前先测每段占比——通信占了三成就先解决拓扑和并行策略，而不是继续加卡。
实测各段耗时的命令：
# GPU 利用率与显存实时观察
nvidia-smi dmon -s um -d 5
# PyTorch 显存明细：分配/保留/峰值一目了然
python -c "
import torch
print(torch.cuda.memory_summary())
"
# 重点看 peak memory 和 reserved 段，两者差距大说明碎片化严重
判定标准：nvidia-smi 里 SM 利用率长期低于 50%，大概率瓶颈在数据加载或通信，加卡收益有限。
可复现性也是成本
一次跑完不算完。代码版本、数据版本、随机种子、并行配置、checkpoint 策略都要落档。大规模任务里硬件故障是常态而非意外：
- checkpoint 间隔太疏，故障后重算代价大
- 保存太密，训练卡在 I/O 上
- 稳妥做法：先用短任务验证保存、加载、断点续训全链路，再定正式任务的间隔
开训前的预算表
把账算在前面，能提前暴露"这套配置根本跑不起来"：
| 项目 | 估算 | 可压缩手段 |
| 参数 | 参数量 × 2 (fp16) | — |
| 梯度 | 参数量 × 2 | — |
| 优化器状态 | 参数量 × 12 | ZeRO 分片 |
| 激活 | 依 seq/batch 估算 | 检查点 |
| 通信缓冲 | 预留 1-2GB/卡 | — |
预算不必精确到最后一 GB，但至少要能回答："序列长度翻倍后，哪一项先成为瓶颈？"
收尾
分布式训练的本质是用通信换容量和吞吐。最终衡量工程质量的不是启动命令多长，而是资源利用率、故障恢复时间和可观测性这三样。下次动多卡之前，先把这本账在纸上过一遍——很多"必须上集群"的问题，算完发现换个精度或分片策略单机就能扛。
评论0