企业官网建设流程全解析

很多朋友做大模型训练，第一关不是代码能不能跑，而是显存够不够。这个问题听着简单，但如果不把显存估计和混合精度训练放在一起算，十有八九会在第一个OOM面前懵掉。我最近在A100上配一个13B的全参数训练任务，估算表写了三页纸，结果第一版配置跑了一百多个step就炸了显存。回头复盘时发现，问题不在工具，而在对显存构成的理解。

这篇整理就是我当时最想看的东西：训练一个模型，显存到底花在哪几块；混合精度训练到底省了什么、没省什么；以及真的被OOM卡住时，按什么顺序排查最有效。适合准备在单卡或多卡上做全参数训练的同学，也适合刚入大模型训练、想搞清楚"一张卡到底能塞多大模型"的读者。

1. 显存到底被谁吃掉了：四大占用的测算模型

1.1 一张公式表，把训练态说清楚

大模型训练显存和推理显存完全是两码事。推理只需要权重和一部分临时激活，训练却要同时容纳权重、梯度、优化器状态以及反向传播所需的大量中间结果。把这几块分开算，才能定位OOM的真正来源。

先定义几个符号：P 表示模型参数量，b 表示单个参数的字节数。FP32 是4字节，FP16/BF16 是2字节，INT8/FP8 是1字节。

全FP32训练时，用AdamW优化器，显存的"权重相关部分"如下：

也就是说，一个 7B 模型，全FP32加AdamW，纯状态就要 16 × 7 = 112 GB。一张 A100 80G 根本装不下，更别提反向传播还要吃激活。

1.2 拿7B模型实际算一遍

再说混合精度。主流做法是 FP16/BF16 前向和反向，同时保留一份 FP32 主权重，更新时用主权重做累积，避免低精度浮点的舍入误差。这样权重相关部分就变成：

• FP32主权重：4P
• 低精度权重副本：2P
• 梯度：4P（在反向传播中通常按FP32累积）
• 优化器状态：8P
• 合计：18P

注意，这里比全FP32多了 2P。这是因为混合精度同时保留了两份权重。7B 就是 126 GB，比全FP32反而多了14GB。很多教程一张嘴就说混合精度省显存，这里得立刻打个问号。

那混合精度到底省在哪？答案是计算速度和激活值，不是模型状态。这一点理解透了，后面所有配置决策都不会跑偏。

1.3 为什么说"混合精度省显存"是个需要拆开看的说法

如果只看权重、梯度、优化器这三块，混合精度的显存收益几乎为零，甚至还有开销。但实际工程中，混合精度训练通常能让你把 batch size 开得更大，这是因为：

1. 计算图的中间张量大部分是低精度。autocast 作用下的线性层、注意力分数、MLP 中间结果以 FP16/BF16 存在，相比 FP32 直接少一半。
2. 激活值存储减半。反向传播依赖前向的中间结果，这些结果以低精度保存，长序列大batch下节省非常可观。
3. 通信量减少。梯度是低精度传输的话，多卡训练带宽压力降低。

所以"混合精度省显存"这个结论，只有在激活值占比高的场景才严格成立。如果序列短、batch小，模型状态是大头，混合精度的显存优势不明显。但如果序列长度4096、batch凑到4以上，激活值轻松超过20GB，混合精度和checkpointing带来的减幅就是决定性的。

2. 混合精度训练：省下的不是状态，是计算和激活

2.1 FP16/BF16与FP32的取舍

混合精度不是简单地把模型切成float16，它背后是浮点数数值分布的取舍。

FP16有5位指数、10位尾数。它能表示的最大值约65504，精度大约千分之一到万分之一级别。问题是梯度经常小到1e-7甚至1e-9，在FP16下直接变成0，也就是下溢。梯度一旦为0，权重根本不会更新，训练看起来loss不动，其实就是被精度坑了。

BF16则是8位指数、7位尾数。指数范围跟FP32一致，能表示的最小非零数约1e-38，几乎不会下溢。但尾数只有7位，相当于小数点后两位半的十进制精度。所以BF16不会让梯度直接消失，但每个数值的相对误差比FP16大。

实践中我把这两种精度总结为：FP16适合对精度要求中等、动态范围可控的场景，必须配梯度缩放；BF16适合训练大模型，动态范围够用，省去了loss scaling，绝大多数框架默认都选它。

2.2 master weight / loss scaling / 动态scale机制

这三个机制是混合精度能稳定训练的根基。

master weight 指的是保留一份FP32权重副本。前向传播和反向传播用低精度权重计算，但每次优化器更新时，是在FP32主权重上完成的。为什么要绕这么一圈？因为 Adam 更新是很小的增量，如果直接在FP16权重上加，这次增量很可能小于FP16的最小分辨率，权重纹丝不动。保留主权重，相当于给权重穿了一件防弹衣。

loss scaling 是为FP16准备的。反向传播得到的梯度除以scale之前是放大的，这样梯度就不会因为过小而变成0。PyTorch 的 GradScaler 会在检测到 inf/NaN 时自动调整 scale，这就是动态缩放。

from torch.cuda.amp import GradScaler scaler = GradScaler() for step, batch in enumerate(dataloader): with torch.autocast(device_type='cuda', dtype=torch.float16): loss = model(batch) scaler.scale(loss).backward() scaler.unscale_(optimizer) torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0) scaler.step(optimizer) scaler.update()

如果用 BF16，基本不需要 GradScaler，autocast 的 dtype 传 torch.bfloat16 就行。这能省不少调参精力，也是我推荐BF16的一个重要原因。

2.3 在训练循环里如何落地

混合精度不是"把所有张量转成FP16"这么简单。你如果自己手动把 input 和 model 都 .half()，大概率训练一两步就 NaN。原因在于：LayerNorm、Softmax、损失函数这些计算对精度极其敏感，低精度会导致数值漂移。

PyTorch 的 autocast 做得比较聪明的地方是维护了一张算子白名单：支持低精度的算子用FP16，数值敏感的算子自动保留FP32。实战里我建议不要手动到处 .half()，直接用 autocast 包裹前向和 loss 计算，让框架来决策。

with torch.autocast(device_type='cuda', dtype=torch.bfloat16): logits = model(input_ids) loss = loss_fn(logits, labels)

这里再提醒一句：如果你的模型里自定义了某些 op，比如用纯 Python 手写的 attention，autocast 不会自动帮你处理。需要手动把中间计算的张量降到合适的精度，或者干脆用 torch.bfloat16 显式转换。这是很多混合精度"跑起来很慢"或"不省显存"的隐藏原因。

2.4 从13B模型的角度看收益

我们用13B、seq=4096、batch=1这个配置来算一笔账。

模型状态这块，混合精度 18P = 18 × 13 = 234 GB，这是分布到多卡上的。激活值方面，按Transformer的粗略估算（后文会展开），一条4096的序列要吃将近28GB激活。如果不用混合精度，激活区是FP32，要56GB；用了混合精度，激活区降到28GB，直接省28GB。模型状态虽然多了26GB，但激活占的是每卡本地显存，模型状态可以靠多卡摊掉，两者性质完全不同。

所以混合精度对大模型训练的收益，主要体现在本地激活和中间张量上。这也是为什么长上下文场景对混合精度需求极其迫切。序列越长，激活越大，混合精度的优势就越明显。

3. 激活显存：那个公式最容易被低估的"隐形杀手"

3.1 激活为什么比参数更吃显存

很多人第一次看到激活值估算结果都会吓一跳。原因是反向传播需要前向时每一层的输出，这些中间结果必须保留到backward，否则没法算梯度。Transformer层数越多、序列越长，保留的中间张量就越恐怖。

业界有个很实用的经验公式，来自Korthikanti等人的论文：对于Decoder-only的Transformer，每个token大约要占 34 × L × H 字节的激活，其中L是层数，H是隐藏维度。这个公式假设了标准注意力加MLP的结构，实测误差大概在10%-20%以内，足够做方案评估。

这里要特别注意：这个值是"未开激活重计算"的量。开了gradient checkpointing，激活占用会大幅下降，后面专门说。

3.2 7B/13B/70B在seq=4096下的预算表

我把常见模型规格代入公式，算出来更有体感。

注意，这只是 batch=1 的数量。如果 batch=4，13B的激活直接来到114GB，单张80G卡光激活就爆。很多OOM就是死在这里——模型状态通过了多卡分片塞进去了，结果batch里每多一条序列都在疯狂吃激活，而你没提前算。

3.3 checkpointing、激活partition与计算换显存的取舍

gradient checkpointing 不是把所有前向结果都丢掉，而是只保留每个Transformer block的输入。反向传播时如果需要某层的中间结果，就临时用保存的输入再算一遍前向。这样激活占用通常能减到原来的1/4甚至更低，代价是大约30%的额外前向计算。

在HuggingFace Transformers里开启非常方便：

model.gradient_checkpointing_enable()

我实际跑过一个13B的LoRA任务，开checkpointing之后，激活从28GB降到不到6GB，训练速度只慢了大约25%。对于显存紧张但不那么追求极限吞吐的场景，这笔买卖很划算。

再往下走还有 activation offload，把激活张量搬到CPU或NVMe，换显存但不增加计算。代价是PCIe传输延迟，速度损耗比checkpointing更明显。常规训练里，我建议优先用checkpointing，offload留给推理而不是训练。

4. 从OOM到摸清底细：我的显存排查与验证流程

4.1 先量后调：两条命令看清峰值

遇到OOM，第一反应不该是狂调batch size，而是先量化。PyTorch给了两个非常好用的接口：

torch.cuda.reset_peak_memory_stats() # 训练几个step print(f"峰值显存: {torch.cuda.max_memory_allocated() / 1024**3:.1f} GiB") torch.cuda.memory_summary(abbreviated=True)

max_memory_allocated 给出的是这个进程在GPU上分配过的最大显存，比 nvidia-smi 的"进程占用"更准确。nvidia-smi 显示的是驱动角度的reserved memory，包含CUDA缓存和碎片，不代表你真正创建的张量大小。只看 nvidia-smi 往往会高估实际需求。

4.2 一次完整的OOM定位过程

说一个我最近的真实案例。13B全参数训练，4卡A100 40G，开BF16混合精度和ZeRO-2。按公式，模型状态126GB，ZeRO-2分片后每卡约63GB，这已经超过40G，理论上就该预警了。但我当时想省事，没仔细算激活就开跑了。结果前100个step还好，第120步开始OOM，而且每次都在同一个位置崩。

后来打了 max_memory_allocated，发现峰值38.5GB，离40G只差一点点。问题很明确：ZeRO-2把梯度和优化器分片在4卡上，但每卡仍然持有全量主权重和低精度权重，这部分就有42GB，光这已经超了，激活再压上来必死。

我的调整路径是：

1. 从ZeRO-2切到ZeRO-3，让参数也分片，模型状态每卡降到约31.5GB。
2. batch size保持4不动，开gradient checkpointing，激活从28GB降到约7GB。
3. 峰值降到36GB左右，稳定训练，速度损失在可接受范围内。

整个过程最强的教训是：先用公式估算，再上卡验证。如果一开始就算清楚63GB装不进40G，也就不用白跑两小时等OOM了。

4.3 常见误判与兜底手段

排查显存问题要会看现象，我整理了几个踩过的坑：

还有个容易被忽略的坑：torch.cuda.empty_cache() 并不会把显存还给你，它只是把缓存池清空、归还给驱动，但下一次分配又要重新建立缓存。OOM时调用empty_cache通常是自欺欺人，该排查的峰值排查，别指望清理能解决紧张局面。

4.4 各种配置的对显存的影响

如果显存只差一点点，优先考虑以下三个微调手段，而不是粗暴降低batch：

• 把优化器换成功率更低的：Adafactor或8bit Adam，优化器状态从8P降到2P左右，13B能省78GB的分布压力。
• 调整allocator策略：PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True 在有些场景能减少碎片，但不是所有卡都生效，自己测了才知道。
• 检查模型是否加载了不必要的fp32副本：比如 .half() 之后又 hot swap，导致同一份参数在显存里留了两份。打印峰值后对比估算值，很容易看出来。

5. 多卡训练下的显存账单：ZeRO各阶段的取舍

5.1 DP到ZeRO-3：每一层分片省了什么

单卡放不下，就要靠多卡把模型状态摊开。PyTorch DDP并不省显存，它每张卡都持有全量模型状态，只做梯度同步。真正省显存的是ZeRO系列，DeepSpeed和FSDP都能实现。

把7B混合精度的126GB模型状态拿出来看：

ZeRO-3看起来最省显存，但每个step都要把参数从多卡聚合到当前卡上做前向，通信量比ZeRO-2高一个量级。所以选型不是越高级越好，而是够用就好。

5.2 DeepSpeed/FSDP配置要点

我最常给到用户的建议：单机4卡以下、模型不超过13B，ZeRO-2基本够；模型到了30B以上或者目标是单卡能塞更多，才上ZeRO-3。DeepSpeed的配置长这样：

zero_optimization: stage: 2 offload_optimizer: device: cpu pin_memory: true

这里 offload_optimizer 是把优化器状态放到CPU内存，能进一步降低GPU占用，代价是每次更新都要走PCIe。实测情况下，offload优化器会让step时间增加30%到50%，如果GPU显存不缺，别开。

5.3 显存省了，通信和时间去哪了

ZeRO省显存，本质是拿通信换显存：前向要all-gather参数，反向要reduce-scatter梯度，optimizer step前还可能再来一轮通信。如果你把batch size剁到2来迁就ZeRO-3，那GPU利用率通常很低，因为大量时间卡在等待通信完成。

我一直坚持的原则是：先用激活手段（混合精度、checkpointing）把每卡峰值降到80%以下，再考虑分布式分片；不要一上来就ZeRO-3+offload。前者牺牲的是固定比例的计算，后者牺牲的是通信延迟，通信瓶颈在多卡扩展下远比计算代价难补。

6. 实测参考表与我这段时间跑模型的经验

6.1 常用规模/卡型/配置参考表

最后给一张我常拿来当起点的配置参考表，基于BF16混合精度，未开激活offload：

这些数字不是精确到字节的，只是我用来做方向判断的基准。每换一个模型，我一定会花二十分钟把公式重新填一遍：16P或18P的模型状态、34LH × seq × batch的激活、checkpointing的折扣、ZeRO阶段的分摊，然后再上卡验证峰值。这半小时能省下后面一整晚的修复时间。

6.2 几个我不会再犯的操作习惯

1. 不要用梯度累积代替显存优化。batch太小再累积梯度，训练曲线会显得特别不稳定，而且累积步数变长，debug起来很受罪。先把batch和seq调到一个不OOM的稳妥值，再根据收敛曲线上调。

2. 开混合精度后，batch size翻倍不一定掉精度，但一定要盯loss曲线。FP16或BF16改变的是数值表示，偶尔会在微调任务上出现损失函数变吵的情况，这时候优先排除LossScale问题，而不是急着把batch缩回去。

3. 用显存换时间时要节制。checkpointing能开，但别同时把activation offload、CPU offload、ZeRO-3全开了，每一步都在PCIe上爬，训练速度会惨不忍睹。优先砍计算，缓存让给通讯。

4. 排查OOM时先把峰值打出来再动配置。我一直觉得，算一次显存预算，胜过在卡上焦虑一晚上。如果你也经常被OOM打断，试试把这张表贴在你工位上，下次改配置前先花几分钟填满它，你会发现自己省下很多不必要的返工。