训练不稳定的显存因素
💥 训练过程中出现 Loss Spike,与显存有什么关系?¶
Loss Spike(损失尖峰)本身是一个数值现象,但它经常与显存问题交织在一起,形成“数值不稳定↔显存压力”的恶性循环。两者的关系可以从以下几个层面理解:
- 瞬时激活值膨胀,直接推高显存峰值
当 Loss Spike 发生时,通常意味着某些层的输出出现了异常大的值(例如注意力 logits 突然变成上千,而不是正常的个位数)。这些大值在 Tensor 中存储时,虽然占用的字节数不变(FP16/BF16 的位宽固定),但如果框架为了后续计算而分配了额外的临时缓冲区(例如为了 safe softmax 保存的 max 值),或者大激活值导致后续算子产生更大的中间结果(如注意力分数矩阵因数值溢出而未被 FlashAttention 优化),就可能造成瞬间的显存尖峰。例如,标准注意力中,如果 QK^T 的值域因 Spike 而变得极大,softmax 后可能产生 NaN,而某些框架的 fallback 机制可能尝试分配更大工作空间重新计算,导致 OOM。
- 梯度异常放大,反向传播时显存需求激增
Loss Spike 必然导致梯度异常大(梯度爆炸)。在反向传播中,每一层的梯度张量大小等于该层参数大小,但这个张量的数值范围变大会带来两个显存问题:
-
混合精度训练中,梯度可能超出 FP16 的表示范围(65504),变成 inf/NaN。框架会尝试恢复或跳过更新,但在这之前,梯度张量已经分配并占用了显存。如果频繁出现,可能造成碎片或额外缓冲累积。
-
某些优化器(如 LAMB、Adafactor)会根据梯度统计量动态调整内部状态,这些状态张量也可能随异常梯度而膨胀?虽然不会自动放大占用字节,但若因为数值问题导致优化器创建了新的状态副本,可能会暂时增加显存占用。
-
Loss Scaling 的动态调整引发显存波动
在混合精度训练中,为应对 Spike,Loss Scale 会自动缩小。这个过程需要保持一份 FP32 的主权重副本。如果 Scale 频繁调整,框架可能会进行额外的拷贝或分配,导致短时的显存波动。更关键的是,如果 Scale 变得非常小,后续梯度可能会下溢为零,使得训练无效,但显存方面不会直接节省——那些零梯度仍然占据同样的显存空间。
-
日志、调试信息、checkpoint 等附加操作 有些训练脚本会在检测到 Loss Spike 时自动保存 checkpoint 或打印详细张量统计,这些操作会临时分配 CPU/GPU 内存,可能导致显存进一步紧张。尤其当 Spike 触发
torch.cuda.empty_cache()后,如果后续立即又分配大块,可能因为碎片而失败。 -
碎片化加剧
不稳定的训练导致频繁的分配-释放模式(例如,重计算、异常处理),可能加剧显存碎片,使得即使总空闲足够,下一次正常分配时却找不到连续块而 OOM。
📊 实际案例:训练 LLaMA 7B 时,第 1000 步突然 Loss Spike 至 100.0,同时 OOM。分析发现,该步遇到一批超长文档,注意力分数矩阵未优化完全,分配了 8GB 临时张量,而碎片化使此次分配失败。根本原因是超长序列导致激活膨胀和 Loss Spike,两者叠加。
✅ 应对策略:使用梯度裁剪(max_grad_norm)限制梯度大小;开启 FlashAttention 消除注意力矩阵;强制数据截断;启用混合精度和 loss scaling 的自动调整;以及监控每步的显存变化,避免 Spike 传递为 OOM。
💣 梯度爆炸是否会导致显存溢出?为什么?¶
直接回答:梯度爆炸本身不直接改变梯度张量的尺寸(字节数),因此不会直接增加显存占用。但它可以通过引发数值异常(inf/NaN)、触发框架的保护机制、造成额外分配等间接途径导致显存溢出。
🔍 为什么不直接增加:
梯度张量的形状和数据类型在反向传播前已经确定。例如,一个线性层的权重梯度形状为 [out_features, in_features],FP16 下固定占用 2 * out * in 字节。即使梯度数值从 0.001 爆炸到 10000,它仍然是 2 字节,所以显存占用本身不随数值大小变化。因此,单纯的梯度数值变大不会像增大 batch size 那样直接增加显存占用。
🔧 间接导致 OOM 的机制:
-
混合精度中的溢出与重计算 在 FP16 训练中,梯度的表示范围有限(最大 65504)。一旦梯度超过此值,就变成 inf/NaN。AMP(自动混合精度)检测到梯度溢出后,会跳过本次权重更新,并降低 Loss Scale。在这个过程中,为了恢复 FP32 的主权重用于下一次正确的更新,框架可能会进行额外的张量拷贝或保留之前的 master weights,但通常不显著增加显存。然而,某些自定义实现可能会在溢出时回退到 FP32 的梯度重新计算一遍(例如 Megatron 的某些版本),这会导致该层梯度显存瞬时加倍(同时存在 FP16 和 FP32 的副本),从而可能 OOM。
-
触发重计算或 Checkpointing 的副作用 如果梯度爆炸导致中间激活被标记为无效,梯度检查点机制在重计算时可能会重新产生大激活值,占用更多临时显存。
-
优化器状态异常增长 某些自适应优化器(如 LAMB)在梯度中出现 inf 时,可能会在内部缓冲区产生极大值,尽管不会扩大缓冲区尺寸,但若框架实现不当(如动态分配临时存储),可能引发额外分配。
-
日志与调试输出的显存开销 程序检测到梯度爆炸后,可能打印 grad norm 或保存所有梯度的直方图,这些操作可能会实例化额外的张量(例如
torch.norm会创建中间张量),消耗一定显存,在临界情况下诱发 OOM。 -
通信缓冲区的瞬间压力 分布式数据并行中,每个 rank 在 All-Reduce 梯度时,会将本地梯度拷贝到通信缓冲区。若梯度值巨大(尽管尺寸不变),不会增加通信量。但如果 NCCL 内部为处理溢出而进行额外的校验或重传,可能会临时分配更多缓冲(但罕见)。
📈 典型场景:训练初期,学习率设置过高,前几步就出现梯度爆炸,控制台打印“Gradient overflow, skipping step”,但随后 OOM。排查发现,跳过步骤后,框架错误地保留了上一次的梯度张量未释放,同时新梯度又分配了,导致梯度内存翻倍。
✅ 总结:梯度爆炸主要从“数值异常导致框架行为异常”这个路径间接影响显存。规避方法包括:梯度裁剪、合适的学习率调度、使用 BF16 代替 FP16(动态范围更大),以及确保框架正确处理溢出后的资源释放。
🛡️ 混合精度训练的 Loss Scaling 如何防止梯度下溢,并间接保护显存?¶
💡 Loss Scaling 的核心目的是在 FP16 训练中放大损失值,使反向传播产生的小梯度不落入 FP16 的亚正常(subnormal)范围或直接变为零,从而保持梯度精度。它间接保护显存的方式是:通过维持有效的梯度流,避免了因梯度过小导致优化器状态膨胀或无效更新累积而引发的额外显存开销。
🔬 机制详解:
- 防止梯度下溢(直接作用)
FP16 的最小正规数约为 6.1e-5,小于此值的梯度将被冲刷为零。训练后期,许多梯度的量级在 1e-6 甚至更小。如果没有 Scaling,这些梯度变为零,权重无法更新,模型停滞。
Loss Scaling 将 Loss 乘以一个大常数(如 65536),那么反向传播计算出的梯度也会被同比放大,从而使小梯度移入 FP16 的可表示范围。优化器在更新权重前,再将梯度除以相同的 Scale 恢复原始尺度。这样,小梯度得以保留,训练正常进行。
-
间接保护显存——避免无效积累与碎片化
-
避免无效的优化器状态膨胀:某些自适应优化器(如 Adam)会维护梯度的一阶、二阶矩估计。如果大量梯度变为零,这些矩估计可能会错误地放大(因为零梯度导致矩估计趋向于零,学习率自适应部分变大?其实不会,但若梯度完全为零,优化器状态中的动量会逐渐衰减,但不会引起显存增加)。更关键的间接保护是:如果大量梯度下溢导致训练无效,开发者可能会误以为是模型容量不足而增大 batch size 或增加层数,从而增加显存。Scaling 保持训练有效,避免了这种不必要的模型扩张尝试。
-
减少重计算与异常处理开销:当梯度过小导致 loss 保持不变或 NaN 时,一些框架的监控机制可能会触发额外的诊断操作(如 dump 张量),这会消耗显存。Scaling 降低了这些异常的发生频率。
-
维护稳定的内存分配模式:正常的梯度流保证了每次迭代的张量分配/释放规律一致,避免因异常跳过步骤而导致的内存泄漏(如未释放的梯度)。
📊 实际作用链路:
无 Scaling → 大量梯度变为 0 → 优化器更新无效 → 训练无法收敛 → 用户可能增大 batch size 或序列长度试图改善 → 显存需求增加 → 可能 OOM。
有 Scaling → 梯度保持 → 正常训练 → 无需盲目增大超参 → 显存可控。
⚙️ 间接保护的具体例子:在 QLoRA 训练中,基础模型冻结,可训练参数极少,但若没有 Scaling,LoRA 的梯度可能下溢为零,导致训练完全失败。用户可能会错误地认为需要增加 LoRA 的 rank 甚至解冻更多层,从而增加大量显存。而正确的 Scaling 使小梯度有效,避免了这些错误决策。
✅ 因此,Loss Scaling 虽不直接减少字节,但通过保持梯度健康,维护了训练的正常内存行为和收敛路径,间接防止了因训练无效而采取的增加显存措施。
🧪 为什么使用 FP16 训练时,出现 NaN 有时会伴随 OOM?¶
💡 FP16 训练中,NaN 经常是数值不稳定的症状,而这些不稳定往往由激活或梯度过大引起。大激活值在产生 NaN 的同时,可能已经占用了大量显存(例如注意力矩阵),或触发框架分配额外缓冲区用于异常处理,两者叠加导致 OOM。
🔍 具体关联:
-
爆炸的激活值导致显存突增 NaN 的一个常见来源是注意力 logits 变得极大,softmax 后产生 NaN。在产生 NaN 的前向过程中,已经计算出了超大的注意力分数矩阵,若没有 FlashAttention 优化,该矩阵尺寸为
[B, heads, S, S],显存占用与序列长度平方成正比。当 S 很大时,这个矩阵本身就消耗大量显存,可能直接导致 OOM。所以 NaN 和 OOM 常同时出现,但 OOM 是因(或者并存),NaN 是果。 -
梯度溢出导致的额外内存分配 反向传播检测到梯度 inf/NaN 后,PyTorch 的 AMP 可能尝试回退或重算某些部分。例如,
scaler.step(optimizer)发现溢出,会跳过optimizer.step(),但之前反向传播分配的梯度张量仍在。如果跳过步骤后,用户的代码没有正确清理或覆盖这些梯度,下一次迭代又会分配新的梯度,造成梯度内存泄漏,逐 step 累积最终 OOM。 -
异常处理引发的张量拷贝 一些训练框架在遇到 NaN 时,会自动保存当前模型状态或打印调试信息。例如,DeepSpeed 可能会 dump 模型参数或梯度到 CPU,过程中可能在 GPU 上产生副本。如果原已接近显存上限,这些额外操作会触发 OOM。
-
通信库的异常重试 分布式训练中,All-Reduce 梯度时若某个 rank 的梯度包含 inf/NaN,NCCL 可能会报告错误或重试,但不会增加显存。但某些包装器会先 gather 一小部分检查,额外分配缓冲,可能造成临时峰值。
-
激活值检查点的重计算风暴 如果 NaN 出现在梯度检查点重计算阶段,重计算可能分配原前向没有的临时张量,引发瞬时 OOM。
📊 真实案例:使用 FP16 训练 GPT-2 时,某个 batch 出现 NaN loss,随后 OOM。分析发现,该 batch 文本极长,注意力分数矩阵分配了 15GB,已接近显存极限,而后续 softmax 产生 NaN 时又尝试分配一个同等大小的矩阵用于调试日志(作者代码中 torch.save 了注意力矩阵),直接撑爆显存。解决方案:移除调试代码,开启 FlashAttention。
✅ 结论:FP16 训练中的 NaN 和 OOM 是一对“难兄难弟”,通常由同一个根源——数值过大导致张量膨胀——引发。强化数值稳定性(如使用 BF16、梯度裁剪、更好的初始化)可同时减少两者。
📊 训练不稳定时,如何通过检查激活值和梯度的统计量来排查显存问题?¶
💡 训练不稳定往往表现为激活值或梯度的均值、方差、最大值等统计量异常。通过监控这些统计量,可以提前发现会导致显存问题的异常膨胀,并定位是哪些层或操作引发了过大的张量分配。
🔬 具体排查方法:
- 激活值统计监控
- 记录每层的输出张量的 min、max、mean、std。如果某层的 max 值突然飙升(例如从 10 变到 1000),或者 std 异常增大,说明该层产生了大激活值。大激活值直接增加后续层激活存储,尤其在注意力层可能导致 O(L²) 的矩阵膨胀。
- 关注注意力 logits 的 scale:在 softmax 前打印 QK^T 的最大值。如果该值达到数千,极易产生溢出和显存峰值。
-
使用 PyTorch hooks:注册 forward hook 来收集统计信息,仅记录标量,开销小。
-
梯度统计监控
- 梯度范数 (total norm, per-layer norm):如果某层的梯度范数远大于其他层,或 total norm 突然暴增,表明该层可能发生梯度爆炸。虽然梯度大小不直接增加显存,但爆炸通常与异常大的激活相关联,而激活已占据大量显存。
- 梯度的 max/min:观察是否有 inf/NaN。一旦出现,立即检查显存使用情况,可能存在泄漏或碎片。
-
结合
torch.cuda.memory_allocated():在记录梯度统计的同时记录显存,寻找相关性。 -
定位显存分配热点
- 使用 PyTorch Profiler 的 memory 视图,结合统计量异常的时间点,找出是哪个算子分配了大块显存。
-
例如,若发现第 12 层 attention 的 max 激活值特别大,同时 profiler 显示该层的
aten::scaled_dot_product_attention分配了大量内存,那么就需要对该层进行数值稳定处理(如换用 FlashAttention、降低学习率、初始化调整)。 -
利用 TensorBoard 或 WandB 绘制曲线
- 将
loss、grad_norm、activation_max、allocated_memory画在同一图表上,观察是否在 allocated memory 尖峰之前,先有 grad_norm 或 activation_max 的异常攀升。这有助于建立因果关系。
📊 实用代码片段:
python
def forward_hook(module, input, output, name): if isinstance(output, torch.Tensor): print(f"{name} activation: max={output.max().item():.2f}, std={output.std().item():.2f}")for name, module in model.named_modules(): if 'attention' in name: module.register_forward_hook(partial(forward_hook, name=name))
✅ 通过持续追踪数值统计,可以在 OOM 之前发出预警,并快速锁定问题层,避免盲目调整超参。
🧬 数据加载的某些样本异常长,如何导致显存尖峰?¶
💡 异常长的样本直接增加序列长度 L,而 Transformer 中的激活显存与 L 成线性或平方关系(注意力矩阵)。因此,当某个 batch 包含超长样本时,激活显存瞬间膨胀,形成一个尖锐的峰值,若超过剩余显存则 OOM。
🔍 机制分解:
-
注意力矩阵的平方增长 在标准自注意力中,需要计算
QK^T并存储 softmax 的结果,形状为[B, heads, L, L]。当 L 从 2048 跳变到 8192 时,该矩阵的显存占用量增加了 16 倍。即使使用 FlashAttention 不存储完整矩阵,前向/反向过程中仍会临时需要与 L 成线性关系的 Q、K、V 张量,其大小也增加了 4 倍。 -
线性层激活的线性增长 每个 Transformer 层的输入输出、FFN 中间激活都是
[B, L, H]形状。L 增大时,这些张量同比例增大。对于一个 32 层的 7B 模型,L 从 2048→8192,仅 Q、K、V 张量每层就从约 16MB 涨到 64MB,总计增加量可能达数 GB。 -
padding 造成计算浪费 如果使用静态 padding,一个 batch 内所有样本会被 pad 到该 batch 最长样本的长度。一个超长样本会导致整个 batch 的序列长度都被提升到该长度,激活显存成倍放大。
-
动态形状带来的额外开销 框架可能在遇到不同序列长度时重新分配内部缓冲区,这个操作本身需要临时显存,且可能留下碎片。
📈 实例:
假设显存还剩 5GB。训练 L=2048 时,激活占用 3GB,一切正常。突然遇到一个 L=10000 的样本,注意力矩阵(无 FlashAttention)大小 ≈ 32 heads * (10000^2) * 2 bytes ≈ 6.4 GB,仅此一项就超出剩余显存,立即 OOM。
🔧 应对措施:
-
强制截断:
max_seq_length=4096,超出部分截断或滑动窗口。 -
使用 FlashAttention 消除 O(L²) 矩阵。
-
动态 batch:检测到长序列时自动减小 batch size,保证总 token 数大致恒定。
-
预计算数据长度分布,采样时避免极端组合。
✅ 因此,数据异常长样本是显存尖峰的最常见触发器,控制序列长度是关键。
⚖️ 使用动态序列长度训练时,如何避免显存超限?¶
💡 动态序列长度训练旨在通过不统一 pad 到固定最大长度来节省计算,但这也导致每个 batch 的显存需求波动。避免超限的核心思想是:控制每个 batch 的总 token 数,而非样本数,并预留足够的显存余量应对波动。
🔧 具体策略:
-
基于 token 数的 batch 策略 不设置固定的
batch_size,而是设置max_tokens_per_batch。DataLoader 在采样时会不断累积样本,直到 batch 中 token 总数达到上限。这样,无论样本长度如何变化,单 batch 的总 token 数保持恒定,激活显存的线性部分(Q、K、V、FFN)基本稳定,但注意力矩阵的 O(L²) 部分对于个别长样本仍可能尖峰。这需要用 FlashAttention 消除二次影响。 -
设置最大序列长度硬限制 即使使用动态长度,也必须设定一个
max_seq_length,超出的文本进行截断。这避免了极端值。可根据显存容量、batch token 上限和模型结构反推出max_seq_length的安全值。 -
梯度检查点与重计算 动态长度下,激活峰值难以预测。开启梯度检查点可将驻留激活大幅降低,使系统对长度变化更鲁棒。即使偶然出现较长 batch,也只重计算部分激活,峰值可控。
-
动态调整 gradient accumulation 根据当前 batch 的实际总 token 数,动态决定是否立即进行 optimizer step。如果 token 数超标,可以拆分或减少累积步数,但实现较复杂。更简单的是固定 token 上限,由 dataloader 保证。
-
使用 PagedAttention 思想(如果适用) 训练时 KV Cache 不需分页,但激活可以类似地分块管理?目前训练框架较少用。但可考虑 Megatron 的序列并行,将长序列切分到多卡,降低单卡激活。
-
监控与预判 在训练循环中,每个 batch 开始前检查样本长度,估算所需显存。若估算超标,则跳过该 batch 或将其进一步拆分(但可能影响收敛)。虽然开销大,但对某些关键任务可行。
📊 参数建议:
-
启用
FlashAttention。 -
使用
--max_seq_length 4096。 -
设置
max_tokens_per_batch使max_tokens_per_batch * hidden * layers * factor远小于剩余显存。 -
通过
nvidia-smi持续监控,找到安全 token 上限。
✅ 总结:动态序列长度 + token-based batching + FlashAttention + 梯度检查点,可在灵活性和显存安全之间取得平衡。
🧠 训练 MoE 模型时,专家负载不均衡如何导致显存问题?¶
💡 专家负载不均衡指不同专家处理的 token 数量差异巨大。这会导致某些专家所在的 GPU(专家并行时)或该专家的优化器状态/激活集中爆发,显存消耗严重不均,进而导致负载高的设备 OOM,而其他设备显存闲置。
🔍 详细机制:
-
专家并行下的显存不对称:在 MoE 训练中,常采用专家并行(Expert Parallelism)将不同专家分布到不同 GPU。如果路由策略使得绝大部分 token 都被分配到少数几个专家(例如 90% 的 token 都去了专家 1 和 2),那么持有这些专家的 GPU 将需要存储和计算大量的中间激活、梯度和优化器状态,而其他专家的 GPU 则相对空闲。这导致单卡显存峰值极高,容易出现 OOM。
-
激活值与临时缓冲区膨胀:专家接收的 token 数决定了其 FFN 层的批量大小。负载高的专家需要处理的 token 数量可能是平均的数十倍,那么该专家的输入张量大小(
[num_tokens, hidden])急剧增加,前向和反向产生的激活值以及梯度缓冲区随之线性甚至超线性增长。例如,top-2 路由下,若某个专家被 10 倍 token 选中,其对应的中间激活占用可能达到 10 倍,瞬间撑爆显存。 -
优化器状态分布不均:由于专家参数是分片存储的,但优化器状态(如 Adam 的动量、方差)与参数一一对应,同样会集中在热门专家的 GPU 上。如果某个专家的参数梯度更新频繁,其优化器状态占用不变,但该专家的参数分片所在的 GPU 还需要存储大量其他专家的状态吗?不,专家并行下每卡只有部分专家,但热门专家的卡需要处理大量 token 的梯度同步,可能导致通信缓冲区堆积。
-
通信负载不均导致缓冲积压:All-to-All 通信用于 token 路由,如果负载不均,某些 GPU 需要发送/接收远超其他卡的数据量,通信缓冲区可能被占满,甚至触发反压,使得临时张量不能及时释放,造成显存堆叠。
📊 实例:
假设 8 个专家分布在 4 张 GPU 上(每卡 2 个专家)。由于路由 collapse,专家 0 和 1 处理了 70% 的 token,它们都在 GPU0 上。正常每卡激活约 5GB,现在 GPU0 需要处理 3.5 倍 token,激活飙升到 17.5GB,而其他卡仅 2.5GB。GPU0 OOM,训练崩溃。
✅ 应对方案:
-
负载均衡损失 (Load Balancing Loss):添加辅助损失鼓励路由器平均分配 token,如 Switch Transformer 的负载均衡损失。
-
专家容量限制 (Expert Capacity):限制每个专家最多处理的 token 数,超出的 token 通过残差连接直接传到下一层或丢弃(会导致精度损失)。
-
动态重新分配专家(resampling):在检测到不均衡时,动态将热门专家的参数迁移到空闲 GPU 上?实现复杂。
-
使用 ZeRO 分片专家状态:即使专家并行,进一步用 ZeRO 对优化器状态进行分片,缓解单卡压力。
-
监控与自动调整:实时监控每张卡的显存占用和 token 分配,及时调整容量因子或路由策略。
🔁 显存问题本质:负载不均打破了分布式训练中精心设计的显存平衡,使得某些设备成为瓶颈,导致全局 OOM。
🤖 在 RLHF 的 PPO 阶段,为什么显存容易爆?有何应对?¶
💡 PPO(Proximal Policy Optimization)微调阶段需要同时维护策略模型、参考模型、价值模型和奖励模型等多个大模型,并且要存储大量的生成样本、logprobs、优势值等中间结果,显存需求是 SFT 阶段的数倍,极易 OOM。
🔍 显存爆炸的根源:
- 多模型常驻显存 典型的 RLHF PPO 训练需要四个模型:
- Actor(策略模型,被训练)
- Reference(参考模型,冻结,计算 KL 惩罚)
- Critic(价值模型,通常与 Actor 同架构,被训练)
-
Reward Model(奖励模型,冻结) 每个模型都需要加载权重、存储激活(至少 Actor 和 Critic 需要反向传播)。如果这些模型全部驻留在 GPU 显存中,所需显存是单模型 SFT 的 3-4 倍。
-
生成阶段的显存消耗 PPO 每个 step 需要根据当前策略生成一批回复(rollout),这是一个自回归解码过程,会产生完整的 KV Cache。如果生成序列较长,KV Cache 会占据大量显存,并且这些 Cache 在计算奖励和优势时可能需要保留。
-
存储 rollout 数据 生成的序列、对应的 log probabilities、values、rewards、advantages 等都需要保存在显存或频繁在 CPU/GPU 间搬运。尤其在计算 PPO 损失时,需要用到这些旧策略下的概率,必须保持张量存活。
-
多次前向/反向 为了计算 PPO 的 clipped 目标,需要对同一批数据多次前向计算 Actor(新旧策略),再加上 Critic 的更新,反向传播的激活和梯度数量庞大。
-
混合精度训练的额外开销 多个模型可能导致 Loss Scaling 各自独立,优化器状态数目庞大。
⚙️ 应对策略:
-
模型卸载(Offload):将冻结的 Reference 和 Reward Model 卸载到 CPU,仅在需要计算 KL 惩罚或奖励时临时加载到 GPU。DeepSpeed-Chat 就使用这种策略,大幅度降低 GPU 显存。
-
共享模型组件:Actor 和 Critic 可以共享底层 Transformer,仅头部不同,减少权重冗余。
-
使用 LoRA/QLoRA:对 Actor 和 Critic 仅训练低秩适配器,冻结主模型,大量减少可训练参数和优化器状态。
-
梯度检查点与 FlashAttention:开启后降低激活值占用。
-
减小生成长度和 batch size:适当缩短 rollout 序列,减少 KV Cache 和储存。
-
利用 ZeRO 分片:跨多卡分片优化器状态和梯度。
-
流水线并行:将模型切分到多张 GPU,但会增加通信。
📊 实际案例:使用 8×A100 80G 训练 LLaMA-7B 的 RLHF,若不做任何优化,单卡显存超 70GB。开启 DeepSpeed-Chat 的 Hybrid Engine 和 Offload 后,单卡显存降至 30GB 左右。
✅ 因此,PPO 阶段显存容易爆是由于多模型共存和大量中间数据,必须综合运用卸载、PEFT 和分片技术。
🔄 训练中如果修改了并行策略,如何快速评估新显存需求?¶
💡 修改并行策略(如从 TP=2 改 TP=4,或增加 PP 深度)后,显存需求的变化可以通过解析公式和内存估算工具快速预估,核心是计算单卡权重、优化器、激活的新分配量。
🔍 评估步骤:
- 权重与优化器状态显存
- 对于 ZeRO:显存 ≈ 参数量 × 精度 / DP 度(分片程度)。修改 ZeRO stage 或 DP 度直接影响。
- 对于 TP:权重显存 ≈ 原权重 / TP 度。优化器状态同样被切分。
- 对于 PP:单卡权重 ≈ 总权重 × (该卡层数 / 总层数)。
-
组合时:例如 TP+ZeRO,权重显存 = 总权重 / (TP × DP_ZeRO)。
-
激活值显存
- 激活大小与 batch size、序列长度、隐藏维度、层数相关。TP 切分层内激活,PP 切分层数,但 PP 同时增加 micro-batch 数会影响激活驻留。粗略估算激活峰值 ≈
micro_batch_size × seq_len × hidden × 层数 × (34~40) 字节,再除以 TP 度(如果 SP 启用)。 -
序列并行(SP)会进一步将序列维度的激活切分。
-
通信缓冲区
-
TP 需要 All-Reduce 缓冲,大小约等于各层激活尺寸。PP 需要发送/接收中间激活缓冲区。
-
快速计算工具
- DeepSpeed 的
ds_estimate:可以输入模型参数和并行配置,估算训练所需显存。 - 手动公式:总结自己的经验公式,例如对于 GPT 类模型:
total_mem = (params * 16 + 2 * params * TP_degree) / (DP * PP) + activation。 - Megatron-LM 的性能模型:可根据配置输出内存预估。
-
huggingface accelerate estimate-memory:
accelerate estimate-memory命令可以加载模型并模拟不同并行策略下的显存分布。 -
小规模模拟
- 在单卡上用极小的模型(同架构、缩小层数)模拟新并行策略的显存占比,然后等比放大到实际模型。
- 或者使用
torch.cuda.memory_summary()在旧配置下观察分布,再按比例计算新配置。
📊 实例:原本用 4 卡 ZeRO-3 训练 13B,想改为 2 卡 TP=2 + ZeRO-2。估算:
-
ZeRO-3 单卡权重 = 26GB/4 = 6.5GB
-
新方案:TP=2 权重减半,ZeRO-2 不分片权重,所以单卡权重 = 26GB/2 = 13GB,优化器分片减半?需要逐项计算。总显存可能反而增加,导致 OOM,因此不可行。
✅ 关键:快速评估依赖于平时积累的显存计算模板,以及利用现成工具验证。修改前务必计算,避免盲目改参数导致 OOM。
💥 损失函数在某些特定输入下产生极大值,会影响显存吗?¶
💡 损失值本身只是一个标量,不直接占用大量显存。但产生这个极大值的过程(如计算 logits、softmax)可能已经分配了大量张量,或者导致梯度爆炸,从而间接影响显存。
🔍 详细分析:
-
损失计算过程中的显存分配 例如交叉熵损失
CrossEntropyLoss需要 logits 和 labels。如果 logits 因异常输入而数值范围极大,softmax 算子可能会在内部使用额外的稳定计算(如减去最大值),这个过程可能分配临时张量。但这通常只是几十 MB,不会导致 OOM。但如果使用了自定义损失,里面有大量矩阵运算(如对比学习中的成对距离),可能因输入异常而触发大矩阵分配。 -
损失极大值伴随的激活异常 损失很大通常意味着模型输出异常,这往往是因为某一层的激活值爆炸。那个爆炸的激活张量(如注意力矩阵)才是显存杀手。损失极大值是症状,根源在大激活。
-
梯度爆炸与优化器状态 损失极大导致梯度极大,虽然梯度张量尺寸固定,但如果梯度包含 inf/NaN,可能触发框架的异常处理流程(如跳过更新时未释放旧梯度),造成梯度内存泄漏。更常见的是,如果开启了
max_grad_norm裁剪,计算梯度范数需要遍历所有梯度,产生中间张量,但开销可忽略。 -
动态计算图扩展 如果损失函数中包含了非固定形状的操作(如基于输入长度的动态 mask),异常输入可能导致计算图变得庞大,分配更多显存。
📌 结论:损失极大值本身不占显存,但它揭示了模型内部的数值不稳定,而正是这种不稳定(如爆炸的激活)导致了显存问题。因此,当观察到 loss spike 时,应立即检查激活和梯度的统计量,并着手解决数值稳定性,而不是担心损失标量。
💾 训练时检查点保存为何可能引起显存波动?¶
💡 保存检查点时,需要将模型参数、优化器状态等聚合到 CPU 或磁盘,这个过程会在 GPU 上创建额外的连续缓冲区,导致显存占用瞬间上升。若原有显存已接近极限,这一额外分配可能触发 OOM。
🔍 波动原因:
-
参数收集与拷贝 在分布式训练中,保存检查点通常需要将分片在各卡的参数(或完整模型)收集到 rank0 或每卡保存自己的分片。DeepSpeed 在保存合并的 HF 格式模型时,会通过 All-Gather 将参数分片聚合成完整张量,这会在 GPU 上创建一个全局参数的副本,显存瞬时增加一倍(对于 ZeRO-3 尤为明显)。
-
状态序列化缓冲区 将张量转移到 CPU 或写入磁盘前,可能需要连续化内存(
contiguous()),这会产生新的张量副本。如果参数很大,这个副本会占用可观的显存。 -
优化器状态保存 保存优化器状态同样需要聚合分片,尤其是当保存完整优化器状态时,每个参数的动量、方差都要汇总,显存开销与模型大小相当。
-
内存碎片的影响 保存时的临时分配请求可能会因为碎片而失败,尽管总空闲足够。
📊 实例:使用 ZeRO-3 训练 13B,每卡显存占用约 18GB/32GB。当保存合并 checkpoint 时,需要聚合全部 13B 参数(FP16 26GB),单卡显存瞬间飙升至 44GB,OOM。
🔧 缓解措施:
-
使用分片保存(每卡保存自己的分片),避免聚合。
-
将保存操作放在 CPU 上执行(DeepSpeed 的
save_checkpoint默认使用 CPU 内存聚合)。 -
在保存前调用
torch.cuda.empty_cache()释放缓存碎片。 -
降低保存频率,或使用异步保存,将数据先移到 CPU 再写入磁盘。
-
监控显存峰值,在显存余量不足时,采用更轻量的保存方式(如只保存模型权重,不保存优化器状态)。
✅ 因此,检查点保存是训练中显存的定时波动源,需通过分片保存和异步化来规避 OOM。
🧪 为什么训练初期最好用小 batch 测试显存上限?¶
💡 训练刚开始时,模型、数据、优化器等组件初次分配,显存碎片少,此时最能测得各组件实际占用的“冷启动”峰值。用小 batch 逐渐增加可以安全地探索显存上限,找到不会 OOM 的最大可行 batch size。
🔍 理由:
-
避免直接 OOM 导致不可恢复 如果一开始就用大 batch,一旦 OOM,进程崩溃,无法获得任何有效信息。用小 batch 跑通一个 step 后,逐步增加,可以在 OOM 前记录下极限值。
-
测量真实的静态占用 训练初期没有碎片,分配器给出的
memory_allocated()较为准确。运行几个 step 后,显存碎片会使得实际可用容量小于初期,导致后期 OOM。但初期测试可以给出理想上限,为后续留出余量。 -
快速收敛到安全配置 通过二分法增加 batch size:先跑通 batch=1,记录显存;倍增到 OOM,然后在前一个值与 OOM 值之间取折中,再结合梯度累积达到目标全局 batch。
-
评估各项开销占比 初期只加载模型和优化器,可以分别测量权重的显存、优化器显存,然后跑一个 step 得到激活显存。这有助于分析哪部分是瓶颈,从而针对性地选择优化方案(如是否需要梯度检查点、量化等)。
📊 具体操作:
-
启动训练脚本,设置一个很小的
per_device_train_batch_size=1,max_steps=3。 -
在第一步前后打印
torch.cuda.memory_allocated(),记录激活开销。 -
逐步提高 batch size,观察显存线性增长情况,直到接近显存上限的 90%。
-
乘以一个安全系数(如 0.8),作为最终使用的 batch size。
✅ 因此,初期小 batch 测试是一种无损伤的显存容量探测,为后续稳定训练建立安全基线。
🔢 梯度累积步数突然改变,会影响显存吗?¶
💡 梯度累积步数(gradient_accumulation_steps)的改变本身不会显著改变峰值显存(因为梯度一直累加,占用不变),但会改变训练中的显存分布节奏,可能因为优化器更新的频率变化而间接影响碎片化,或在某些实现中导致额外缓冲区堆积。
🔍 详细分析:
-
梯度显存不变:无论累积步数是 1 还是 8,反向传播后梯度一直存在于
.grad中,占用固定大小(等于模型参数量)。直到optimizer.step()才清零。因此,梯度累积步数不影响梯度张量的总占用。 -
激活显存可能变化:如果配合梯度检查点,某些实现会在每个 micro-batch 后释放激活,所以激活峰值只取决于单个 micro-batch 的大小。累积步数增加只意味着重复多次前向/反向,并不会让激活叠加。
-
优化器状态更新频率:累积步数越大,优化器 step 频率越低。这不会改变优化器状态显存。但在某些分布式设置中,如果通信与计算重叠,较大的累积步数可能使得通信缓冲区长时间不被释放,造成一种“慢性”堆积?通常不会。
-
潜在的碎片效应:频繁的 micro-batch 迭代会产生更多的分配和释放周期,可能加速碎片化。但碎片化与累积步数没有单调关系。
-
实现层面的陷阱:如果代码在每次 backward 后错误地调用了
optimizer.zero_grad(),那么梯度不会累积,那也就不存在 OOM 问题了,但训练会错。如果不小心在累积期间保留了计算图(如retain_graph=True),会导致激活累积,累积步数越多 OOM 风险越大。
📈 唯一可能直接关联:当梯度累积步数设置过大,且全局 batch size 固定,micro-batch size 就必须非常小。极小的 micro-batch 可能导致 GPU 利用率极低,但显存占用(权重+优化器+激活)几乎不变,所以不会 OOM。所以一般不会因此 OOM。
⚠️ 特别注意:如果从无梯度累积改为有梯度累积,而保持相同的 per_device_train_batch_size,那么显存占用毫无变化。若想通过增加累积步数来模拟大 batch,同时减小 micro-batch size 以降低激活,是可以省显存的。所以“改变梯度累积步数”如果伴随 micro-batch size 调整,则会影响显存;若单纯改变累积次数而不动 micro-batch,则无影响。
✅ 结论:单纯改变梯度累积步数而不改变 micro-batch size,不会影响显存占用。显存变化是因为与之配合的 batch size 或其它参数变化。
🎛️ 如何利用 DeepSpeed 的 autotuning 自动寻找最优显存配置?¶
💡 DeepSpeed 的 Autotuner 通过自动搜索不同 ZeRO stage、offload 策略和并行度组合,运行微型训练任务并测量显存与吞吐,最终输出满足显存约束下吞吐最高的配置。
🔧 使用步骤:
- 编写 DeepSpeed 配置文件,启用 autotuning
在配置中设置
"autotuning": {"enabled": true},并指定搜索空间参数,如 ZeRO stage 范围、offload 选项等。
{
"autotuning": {
"enabled": true,
"arg_mappings": {
"train_batch_size": "--per_device_train_batch_size",
"model_parallel_size": "--tp"
}
}
}
- 启动训练,触发 autotuning
正常执行
deepspeed train.py,Autotuner 会首先运行一系列微型训练(通常几步)来收集内存和速度数据。它会尝试不同的组合,例如: - ZeRO stage 1, 2, 3
- Offload optimizer / param to cpu or nvme
- 调整 micro-batch size
-
可能的 TP/PP 度(如果指定)
-
分析结果 Autotuner 会生成一个日志,记录每种配置下的显存使用、吞吐(samples/sec)、通信开销等。最终会给出推荐配置,并解释原因。
-
直接使用或手动调整 可将推荐配置导出为最终的 DeepSpeed 配置文件,或根据自己需求微调。
📊 Autotuner 的搜索原理:
-
它利用一个分析模型,结合实测数据点,预测其他配置的显存。不需要跑遍所有组合,而是通过插值和外推。
-
对于 ZeRO stage,它知道各阶段的理论节省比例,并通过实际内存分配校正。
-
对于 offload,它会模拟 PCIe 带宽影响速度。
🧰 高级功能:
-
可以与 hf Trainer 集成。
-
支持自定义搜索空间。
-
在 autotuning 过程中,如果 OOM,它会自动跳过该配置并记录,不会崩溃。
✅ 优势:避免了手动猜测参数组合,一键找到在给定硬件上能运行的最大 batch size 和最快配置,特别适用于新模型或新硬件环境。
📌 注意:autotuning 需要额外的启动时间和少量显存余量来运行测试。它给出的是在当前硬件和任务下的最优,迁移到不同数据长度可能还需微调。但它是科学评估显存需求、快速部署的强大工具。