开头先从一次真实翻车现场说起。去年我把一个 13B 模型放到单卡上做微调，盯着nvidia-smi看显存从 12GB 往上涨，然后眼睁睁看它撞上 80GB 的墙——OOM 报错弹出来那一刻，我才意识到自己对大模型训练显存估计的理解有多肤浅。后来换了混合精度训练，又踩了 BF16 和 FP16 的坑，才把这一整套逻辑理顺。这篇就把大模型训练里显存估算的方法和混合精度训练的底层机制一次讲透，包括怎么算、怎么配、踩过的坑怎么排。
1. 训练显存的五个去向：先搞清楚钱花在了哪
1.1 参数、梯度、优化器状态：三种最直接的“大头”
问一个实际问题：训练一个模型，显存到底被谁吃掉了？很多人的第一反应是“模型参数”。这话对了一半。真正吃显存的其实是三份数据：参数本身（weight）、反向传播算出来的梯度（gradient）、以及优化器内部维护的状态（optimizer states）。
参数所占空间最直观——模型有多少个参数，每个参数几个字节，一乘就出来。梯度呢，模型反向传播时需要把 loss 对每个参数的导数暂存下来，形状和参数一模一样，所以它占的空间也和参数等量。优化器状态就容易被忽略了，但它往往才是最大的开销，尤其当你用 AdamW 这类自适应优化器时。
为什么优化器状态这么大？以 AdamW 为例，它给每个参数额外保存两样东西：一阶动量 m 和二阶动量 v，都是 FP32 格式。对于混合精度训练，还得再保存一份 FP32 的 master weight（主权重）。这三项加起来每参数 12 字节，而参数本身用 BF16 存才 2 字节，梯度 2 字节。一对比你就知道，优化器状态有多大分量了。
1.2 激活值：随 batch size 和序列长度膨胀的“隐形开销”
第四个大头是激活值（activation values）。前向传播每一层的输出需要保存在内存里，供反向传播计算梯度时使用。这一部分和模型参数量的关系不大，而是和你的 batch size、序列长度、隐藏层维度直接挂钩。
它的增长模式很吓人：batch size 翻一倍，激活内存也差不多翻一倍；序列长度翻一倍，激活内存同样可能翻数倍。你可以把它理解成一次性的“过路货”——算完就扔，但在算完之前必须完整存着。对于长序列场景，激活内存完全可能追平甚至超过参数内存。
1.3 通信缓冲区与显存碎片：容易被低估的杂项
第五类是杂项开销。多卡训练时，梯度同步需要临时存储通信缓冲区，像 NCCL 的 all-reduce 操作，每张卡都要预留一块发送和接收数据的空间。此外还有显存碎片：动态分配和释放过程中产生的空洞，不会因为你操作完就自动消失。
显存碎片在长训练任务中尤其烦人，明明nvidia-smi显示还有 10GB 空闲，程序却告诉你 OOM。原因就是大块连续内存被切碎了，无法满足模型运行时对连续显存的请求。后面我会专门说怎么处理碎片问题。
2. 显存估算公式：从参数规模直接推算出“得用几块卡”
2.1 不同优化器下的单位参数量成本
搞清楚了显存去向，估算就有章可循。先把“每多少个字节每参数量”这个基础数字记熟，这套体系建立后，任何模型都能快速估算。
| 训练配置 | 参数 | 梯度 | 优化器状态 | 合计（字节/参数） |
|---|---|---|---|---|
| FP32 + SGD | 4 | 4 | 0（纯 SGD 无动量） | 8 |
| FP32 + SGD（带动量） | 4 | 4 | 4 | 12 |
| FP32 + AdamW | 4 | 4 | 8（m + v） | 16 |
| BF16/FP16 + AdamW | 2 | 2 | 12（master + m + v） | 16 |
| BF16 + AdamW + ZeRO-1 | 2 | 2 | 12/N（N 卡分片） | 随时 N 缩小 |
| BF16 + AdamW + ZeRO-3 | 2/N | 2/N | 12/N | 随时 N 缩小 |
注意一个关键细节：混合精度训练省显存，重点并不在参数和梯度那几字节，而是省了激活值（可用 FP16 半精度存储），同时保证优化器状态依然用 FP32 维持训练稳定性。如果你用 BF16 + AdamW 但不开 ZeRO，光参数+梯度+优化器状态每参数还是要 16 字节，和纯 FP32 + AdamW 几乎一样。这是很多人对混合精度的第一个误解。
2.2 实际计算公式与示例：13B 模型到底需要多少显存
用 13B 模型做例子算一笔账。模型 130 亿参数，混合精度 + AdamW，每参数 16 字节，那么运行权重、梯度、优化器状态的静态开销就是：
13 × 10^9 × 16 字节 ≈ 208GB
是的，光这三样就要 208GB，还没算激活值、通信缓冲和中间碎片。这就是为什么 13B 单卡微调，即便用 BF16 也要把 batch 调很小，否则直接爆显存。
加上激活值，我给一个工程上的粗略公式：
总显存需求 ≈ 参数量 × 16字节（静态部分） + batch_size × seq_len × hidden_size × num_layers × 激活系数
激活系数通常在 2 到 20 之间，取决于是否开梯度检查点、是否存 FP16、实现细节等。开梯度检查点能把这个系数压到接近 1 到 2。保守估算时，我一般把激活部分按静态部分的 20% 到 40% 算，然后再加上 2GB 的通信冗余和碎片余量。
算出来的结果，再比对照你手头 GPU 的显存，用整除确定需要几张卡、要不要上 ZeRO。例如 13B 模型静态 208GB，四张 80GB 卡总共 320GB 显存，那大概率够用，但如果是两张 80GB 卡只有 160GB，就必须开 ZeRO-1 或 ZeRO-2 来分片优化器状态了。
2.3 脚本实测：用 PyTorch 快速清点模型显存
纸上算完，还要实测验证。你在本地用一段小代码，就能量出模型权重到底占多少字节。
import torch def count_params_and_bytes(model): total_params = 0 total_bytes = 0 for name, param in model.named_parameters(): if param.requires_grad: n = param.numel() total_params += n total_bytes += n * param.element_size() if total_params < 5_000_000: print(f"{name}: {n} params, {param.element_size()} bytes/elem") return total_params, total_bytes model = get_your_model() total_params, total_bytes = count_params_and_bytes(model) print(f"Total params: {total_params:,}") print(f"Model weights memory: {total_bytes / 1024**3:.2f} GB")跑完后打印出来的权重块，配合上面那张每参数成本表，能反推更精细的需求。还有两个 API 在训练时值得盯：torch.cuda.memory_allocated()显示当前实际分配的显存，torch.cuda.max_memory_allocated()显示到目前为峰值。把这两行代码放在一个训练 step 的首尾，每次打印，会对显存随 batch 变化的趋势有很直观的感受。
3. 混合精度训练的原理与选型：为什么 BF16 是训练首选
3.1 FP16、BF16、FP32：表示范围和精度的区别
混合精度训练的核心，是让计算和存储用低精度格式，同时又避免精度过低导致训练崩溃。从底层原理看，FP16 和 BF16 都是 2 字节数据类型，但分配方式完全不同。
FP16 有 5 位指数位和 10 位尾数位，表示范围大约在 6e-8 到 65504 之间。范围小是它的硬伤，一旦数值超出 65504 就会溢出为无穷大（inf）。BF16 则是 8 位指数位加 7 位尾数位，指数范围和 FP32 几乎一致（因为 FP32 也是 8 位指数），但尾数只有 7 位——精度低得多，可范围要安全得多。
训练场景中，梯度的大小往往起伏很大，FP16 的窄范围几乎是天然劣势，必须靠 loss scaling 硬撑。而 BF16 因为范围和 FP32 一致，不需要太多额外保命机制就能稳定训练。这就是为什么当 A100、H100、以及新出的消费级显卡支持 BF16 之后，训练社区迅速转向 BF16 的原因。
3.2 master weight 与 loss scaling：混合精度的两个核心机制
混合精度训练不是“所有数值都用半精度”这么简单。整个机制里，最容易被忽视也最重要的两件事：master weight，以及针对 FP16 的 loss scaling。
master weight 是说，模型需要保留一份 FP32 格式的权重副本，训练过程中用它来更新参数，再把它转成半精度用于前向和反向计算。为什么不能直接在半精度权重上更新？因为半精度的尾数太短，一次学习率的微小增量可能比尾数能表示的最小步长还要小，更新了几百上千个 step 后误差累积起来，loss 就完全不收敛了。所以 master weight 相当于一个高精度“账本”，每轮计算完再以低精度副本参与训练。
loss scaling 则是 FP16 训练专属手段。反向传播算出的梯度普遍数值很小，FP16 能表达的最小正数有限，极小梯度会直接变成 0。处理方法是把 loss 乘一个大系数（比如 1024），梯度整体放大到 FP16 可表示的范围，再反向传播；等优化器取到梯度后，除以同样的系数恢复真实大小。PyTorch 的GradScaler会在训练过程中动态调整这个缩放系数——当发现某个 step 的梯度溢出为 inf，就把系数调小；持续若干 step 没溢出，再尝试调大。
3.3 什么场景用 FP16、什么用 BF16
做选择之前，先看你的 GPU 支持情况。RTX 3090、V100 这类老一点的卡原生支持 FP16，但对 BF16 不友好或根本加速不了。A100 及以后的服务器卡基本都原生支持 BF16。消费级 RTX 4090 也支持 BF16，但也要确认是原生计算还是模拟。
如果显卡不支持 BF16，那就用 FP16 + loss scaling，靠 GradScaler 动态维护数值范围。支持 BF16 时，我强烈优先选 BF16。理由很简单：省心。你不用整天为 loss 爆炸、loss 变 NaN 发愁，可以把精力放到模型本身的问题上。BF16 尾数少这个缺点，在大多数 LLM 训练任务中并不致命，因为训练的收敛主要仰仗优化器的累加和 master weight 的纠偏。
一句话总结选型标准：能用 BF16 就用 BF16，原生不支持再退到 FP16，千万别用半精度把所有值都直接降了——那是性能灾难，也是炼丹事故的源头。
4. 实战：PyTorch/DeepSpeed 混合精度训练配置与显存优化组合拳
4.1 最小可用的 PyTorch AMP 训练示例
在实际工程中，PyTorch 的torch.autocast搭配GradScaler，是上手最快的混合精度方案。我通常这样写训练循环：
import torch model = model.cuda() optimizer = torch.optim.AdamW(model.parameters(), lr=1e-5) scaler = torch.cuda.amp.GradScaler() use_bf16 = torch.cuda.is_bf16_supported() for epoch in range(epochs): for input_ids, labels in dataloader: input_ids, labels = input_ids.cuda(), labels.cuda() optimizer.zero_grad() dtype = torch.bfloat16 if use_bf16 else torch.float16 with torch.autocast(device_type="cuda", dtype=dtype): loss = model(input_ids, labels=labels).loss if not use_bf16: scaler.scale(loss).backward() scaler.step(optimizer) scaler.update() else: loss.backward() optimizer.step()关键点在于：BF16 模式下不需要 GradScaler，直接普通 backward 和 optimizer.step() 即可。FP16 模式下则必须使用 scaler，否则极小的梯度会被下溢掉，训练直接停在原地。这个差异刚接触混合精度时很容易写错。
4.2 DeepSpeed 与混合精度、ZeRO 的配合方案
单卡能跑动 7B、13B 吗？能，但把 batch 压到 1 之后，显存可能依然不够，这时候就要上 ZeRO。ZeRO 的核心思想一句话说就是把模型训练过程的数据从“每张卡都存一份副本”变成“分片存储、集体通信”，ZeRO-1/2 主要切优化器状态，ZeRO-3 把参数、梯度、优化器状态全部切分。
DeepSpeed 下的配置我直接给一份能跑的 JSON：
{ "train_batch_size": 32, "gradient_accumulation_steps": 4, "gradient_clipping": 1.0, "bf16": { "enabled": true }, "zero_optimization": { "stage": 1, "allgather_partitions": true, "reduce_scatter": true }, "optimizer": { "type": "AdamW", "params": { "lr": 1e-5, "betas": [0.9, 0.999], "eps": 1e-8 } }, "scheduler": { "type": "WarmupDecayLR", "params": { "warmup_min_lr": 0, "warmup_max_lr": 1e-5, "warmup_num_steps": 100, "total_num_steps": 10000 } } }这条配置的含义是：bf16.enabled = true开启混合精度；zero_optimization.stage = 1只切分优化器状态，适合显存差点意思但不多的情况；gradient_accumulation_steps = 4把 4 个小 batch 的梯度累加后再更新一次，相当于扩大了有效 batch size，还不需要额外增加显存。
注意一点：DeepSpeed 里fp16和bf16两个配置只能选一个，不能同时开启。用 FP16 时fp16.initial_scale_power一般设成 32，loss_scale_window设成 1000 左右，让动态 loss scaling 在一个合理范围内波动。
4.3 梯度检查点与激活重计算：牺牲速度换显存
开完混合精度 + ZeRO-1，模型可能还是差一口气，而你想省显存，最简单的手段是开梯度检查点（gradient checkpointing，又叫 activation checkpointing 或 activation recomputation）。原理一句话：前向传播时不保存每一层的激活值，只保存一小组“检查点”；反向传播需要某一层激活时，临时把之前的前向路径重新算一遍。
代价是大约 10% 到 30% 的训练速度损失，换来的是激活内存的直接“减半再减半”。实际跑大模型时，我的经验是优先开它，尤其在序列长度长、batch size 又压缩不了的场景里，比降 batch 更划算。
PyTorch 里的打开方式也很直接：
model.gradient_checkpointing_enable()HuggingFace 的 Transformer 模型基本都内置该方法。DeepSpeed 配置里也可以通过设置activation_checkpointing片段启用。实测一个 7B 模型，显存峰值可能从 60GB 掉到 40GB 左右，7B 就拥有了在单张 48GB 卡上训练的空间。这个“速度换显存”的 trade-off，在大模型时代的收益非常可观。
4.4 显存实测案例：7B 模型三种方案对比
分享一组我实测过的显存占用数据，模型是 7B，batch size 2，序列长度 2048，单卡训练（A100 80GB 环境）。不同方案组合下，峰值显存差异明显。
| 方案 | 权重+梯度+优化器状态 | 激活内存 | 实测峰值显存 | 每 step 耗时 |
|---|---|---|---|---|
| 全 FP32，不开 ZeRO | 约 112GB（OOM） | 无法评估 | OOM | 无数据 |
| BF16 + ZeRO-1 | 约 28GB | 约 30GB | 约 65GB | 约 21s |
| BF16 + ZeRO-1 + 梯度检查点 | 约 28GB | 约 8GB | 约 42GB | 约 26s |
| BF16 + ZeRO-3 + 梯度检查点 | 约 12GB | 约 8GB | 约 26GB | 约 33s |
从这组数据能看出，单纯开混合精度对静态数据（权重+梯度+优化器状态）的削减有限，真正的“显存杀手”往往在激活值和优化器状态上。想要极致的显存控制，就得混合精度 + ZeRO + 梯度检查点三管齐下。
5. 训练中的显存与精度问题排查实录
5.1 Loss 变 NaN/Inf，八成跟混合精度有关
训练到一半 loss 突然变成 NaN，这大概是混合精度训练里出现频率最高的事故。很多人第一反应是调学习率，但我建议先按下面顺序排查一圈。
先看是不是 FP16 的溢出问题。检查scaler.get_scale()输出的缩放系数，如果它一直在自动下降，说明梯度频繁溢出。对治手段是把loss_scale_window调大，或者手动把initial_scale_power从 32 降到 24，也能换取更保守的范围。第二个常见根源是学习率过大，大模型训练里混合精度会导致每个 step 的有效更新量变大，原来 FP32 下能跑的 1e-4 在混合精度下可能就崩了，先降一半再继续观察。第三个可能是数据里有异常值，比如 label 出现 inf，或者 embedding 层输入没归一化。
如果你已经切到 BF16，那 NaN 的概率本来就低很多，真出现了，优先查模型结构或数据，别再怀疑是混合精度框架的锅。
5.2 OOM 了怎么办：三步排查法
OOM 报错人人都遇到过，排查也有套路，别一上来就降 batch。
第一步看torch.cuda.memory_allocated()和torch.cuda.memory_reserved()。前者是真正在用的显存，后者是缓存池保留的大小。如果reserved远超allocated，说明主要问题是碎片或者缓存不释放。试试设置环境变量：
export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True这个选项让 PyTorch 用可扩展内存段来分配，能显著缓解碎片问题。第二步把梯度检查点打开，把激活重计算的收益吃下来。第三步才轮到调 batch size，配合gradient_accumulation_steps把有效 batch 补回来。这三步走完，绝大多数单卡和单节点 OOM 都能解决。
还有一个细节：多卡训练时，通信缓冲区也会占显存。减少通信峰值的方法是把梯度同步操作拆小，或者换allreduce的通信后端；但最简单的办法其实是再开一个zero_optimization.stage。通常 stage 1 升到 stage 2 能多省一点一块，stage 2 升到 stage 3 能省更多，但通信量会明显增加，训练会变慢，得权衡。
5.3 常见问题速查表
把我在实际训练中遇到最多的几个问题整理成一张速查表，熟读这张表能帮你省下大把 debug 时间。
| 症状 | 可能原因 | 解决手段 |
|---|---|---|
| Loss 在某个 step 突然变 NaN/Inf | FP16 loss scale 过高导致梯度溢出 | 调大loss_scale_window，或改用 BF16，或降低学习率 |
| 训练速度慢但显存充足 | 梯度检查点重计算开销过大 | 只对部分层开启 checkpoint，或减少 checkpoint 密度 |
| step 一开始就直接 OOM | 激活值计算量巨大 | 开梯度检查点，或者把序列长度先缩短到一半验证 |
reserved多但allocated少 | 显存碎片化缓存 | 设置expandable_segments，或定期重启训练进程 |
| 多卡训练时速度和显存都不理想 | 梯度同步通信量大 | 优先开 ZeRO-2/3，或用梯度累积延长同步周期 |
| FP16 训练比 BF16 loss 抖动严重 | FP16 动态范围窄 | 换成 BF16，或调大initial_scale_power并加强 clipping |
5.4 一个容易忽略的细节：控制变量比调参更重要
做这些排查时最重要的一条，是每次只改变一个变量。我在调显存优化时吃过亏：同时开了梯度检查点、换了 ZeRO stage、又调了 batch size，结果训练速度暴跌，根本分不清是哪一项的影响。后来学乖了，每次只动一个开关，记录max_memory_allocated()和每 step 耗时，测试四五轮后，再根据数据决定保留哪些优化项。这个习惯看起来笨，但在大模型训练这种纯试错成本极高的场景里，反而最节约时间。
回归到混合精度和显存本身，我个人在实际操作中的体会是：显存估计永远要有 20% 的安全余量。理论上算出来 62GB，你别真拿 64GB 的卡去跑，因为 PyTorch 缓存池、CUDA context 和 NCCL 都会额外吃掉几个 GB。预先留好余量、把估算公式做成习惯，再配合混合精度和 ZeRO 这把组合拳，大模型训练就算搬到一张消费级显卡上，也不是什么不可能完成的任务。接下来我大概率会写一写多卡分布式训练里的通信开销怎么优化，那又是一个全新的坑。
从零组建指南：设备选型、布线配置与故障排查)
