昇腾910B上LoRA微调OOM：不是模型太大，是优化器状态和梯度检查点没对齐

在华为云 ModelArts 上用昇腾 910B（64GB）对 7B 模型做 LoRA 微调，按经验，LoRA 的可训练参数只有几百万，64GB 显存跑 7B 应该绰绰有余。但实际启动训练后，不到 200 step 就 OOM，npu-smi info 显示显存占用持续爬升而非平稳。排查后发现，问题不在 LoRA 本身，而在优化器状态的精度配置和梯度检查点（gradient checkpointing）在 torch_npu 上的行为上。

一、 第一个坑：优化器状态的 dtype 没有跟随模型精度

7B 模型以 FP16 加载，LoRA 参数也是 FP16，但 PyTorch 默认的 AdamW 优化器会为每个可训练参数维护 FP32 的 exp_avg 和 exp_avg_sq。LoRA 参数虽然少，但如果代码中错误地把优化器作用在了全部模型参数上（比如忘记只传 filter(lambda p: p.requires_grad, model.parameters())），就会为 7B 个参数各维护两份 FP32 状态，显存直接爆掉。检查代码后发现，我们的优化器初始化确实写成了 AdamW(model.parameters(), lr=...)，把冻结的基座参数也纳入了优化器管理。修正为只传 LoRA 参数后，优化器状态显存从 50GB+ 降到不足 1GB。

python
# 错误：把所有参数传给优化器
optimizer = AdamW(model.parameters(), lr=2e-4)

# 正确：只优化 requires_grad=True 的 LoRA 参数
trainable_params = [p for p in model.parameters() if p.requires_grad]
optimizer = AdamW(trainable_params, lr=2e-4)
二、 第二个坑：梯度检查点与 torch_npu 的激活重计算

梯度检查点通过“不保存中间激活、反向时重新计算”来省显存，是微调大模型的标配。但在 torch_npu 上，如果模型启用了 gradient_checkpointing_enable() 但没有配合 use_reentrant=False，重计算路径可能走的是旧版实现，与昇腾的算子融合策略不兼容，导致显存不降反升。另外，检查点本身会增加约 20%~30% 的计算时间，如果 cachedCount 类的“预取”逻辑没有关闭，前向激活在重计算前被额外保留，显存优化效果大打折扣。

我们的修正方式是显式指定非重入模式：

python
model.gradient_checkpointing_enable(gradient_checkpointing_kwargs={"use_reentrant": False})
同时确认 DataLoader 的 prefetch_factor 没有设置过大（设为 2 即可），避免数据加载侧占用不必要的 NPU 内存。

三、 第三个坑：LoRA 的 rank 和 target_modules 选择

排查中还发现，最初为了“效果好”，把 lora_r 设成了 64，target_modules 覆盖了 q_proj、k_proj、v_proj、o_proj、gate_proj、up_proj、down_proj 全部七个模块。虽然 LoRA 参数总量仍然不大，但每个模块的 LoRA 适配器在前向时都会产生额外的中间激活，在梯度检查点模式下这些激活的保存和重计算开销叠加。把 lora_r 降到 16、target_modules 收敛到 q_proj 和 v_proj 后，单步训练的峰值显存再降约 15%，而验证集上的 loss 曲线在 500 step 内与全模块配置的差距在 0.02 以内。

四、 总结

昇腾上 LoRA 微调的 OOM，多数不是“显存真不够”，而是优化器误管了冻结参数、梯度检查点配置不兼容、LoRA 配置过于激进三者之一。排查顺序：先确认 optimizer 只接收 requires_grad=True 的参数，再检查 gradient_checkpointing_kwargs 是否指定了 use_reentrant=False，最后按任务实际效果收敛 lora_r 和 target_modules。三步做完，64GB 的 910B 跑 7B LoRA 微调完全够用，不必急着换更大规格的算力。