多显卡训练大模型OOM排查与DDP/FSDP/LoRA选型实战

你在一个小会议室里盯着训练终端跑大语言模型微调模型加载正常loss 还在往下走结果没几分钟直接抛红torch.OutOfMemoryError: CUDA out of memory。你第一反应是抬头看nvidia-smi发现四张卡里三张闲着只有第一张卡被塞得满满的。明明显卡总显存绰绰有余为什么会 OOM这个问题我实打实踩过很多次今天就把多显卡训练大模型时的显存排查、方案选型和最终跑通的完整思路写出来。这篇文章不是什么高深论文就是一次真实 OOM 排查记录。适合刚接触大模型微调、准备用多卡但被显存折磨的人。文章会先讲清楚“显卡够用”为什么是个错觉然后给出一套从日志到代码的诊断方法再对比 DDP、FSDP、LoRA 等常见路线最后用一个模拟的 7B 模型训练场景演示怎么从单卡 OOM 改到四卡稳定跑。所有步骤都能直接参考不需要你有深厚的分布式训练经验。1. 先说清楚“显卡够用”和“训练不OOM”不是一回事1.1 默认情况下多张显卡只有一张在真正跑模型很多第一次做多卡训练的开发者会有一个默认假设模型和数据到了 GPU 环境里PyTorch 会自动把计算分散到所有卡上。这个假设是错的。如果你写的是下面这种代码# 看起来是“在GPU上训练”实际上只用了 cuda:0 model load_model(...) model model.cuda() optimizer torch.optim.AdamW(model.parameters()) for batch in loader: batch {k: v.cuda() for k, v in batch.items()} out model(**batch) loss out.loss loss.backward() optimizer.step()那么不管机器上有 4 张卡还是 8 张卡cuda()都只会把模型放到默认设备cuda:0数据也被搬到了cuda:0。一个 7B 模型用 bf16 加载光权重就是 14GB 左右。如果你的单卡只有 12GB那模型一放上去就会 OOM如果单卡是 24GB可能模型放得下但一开训练激活值、梯度、优化器状态一起涌上来照样 OOM。你说“四张卡加起来 48GB、96GB怎么还会 OOM”因为 PyTorch 默认按单卡分配不按“总显存”分配。多卡只是摆在那里你不写分布式代码它就不会帮你分活。这个误判是最常见的 OOM 原因没有之一。另外DataParallel虽然会“用”多卡但它在主卡上做 loss 计算和梯度同步主卡显存压力远大于其他卡而且存在线程竞争问题。早期项目里可以用但大模型训练我建议直接跳过它转向进程级的DistributedDataParallel简称 DDP或 FSDP。1.2 训练时占显存的不只是模型权重很多人用“模型文件多大”来判断显存够不够这会把问题严重低估。模型权重只是训练显存的一部分甚至不是最大的一部分。以 7B 全参数微调为例假设使用 bf16 权重模型权重约 14GB反向传播时需要保存的梯度约 14GBAdamW 优化器状态这个是大头通常需要保存 fp32 的主权重、一阶动量、二阶动量加起来接近每参数 12 字节也就是 7B 参数约 70~80GB前向传播过程中各层保留的激活值这个取决于 batch size、序列长度、模型层数长文本下完全可以超过模型权重本身。你可以把模型权重理解为“仓库里的货物”训练过程还需要为货物周转搭建临时货架梯度、优化器状态、激活值就是这些临时货架。仓库总空间再大货架本身没处放照样装不下。所以当你说“显卡够用”的时候建议先算一下自己到底在训练什么规模、用什么优化器、开多大的 batch。很多人的 24GB 单卡明明只够放权重却试图全参数微调 7B 模型那 OOM 是必然结果不是偶然崩溃。1.3 多卡片训练不等于显存“拼起来”这是第二个典型误解以为四张 24GB 显卡加起来 96GB就可以跑一个需要 80GB 显存的任务。DDP 的工作方式是在每张卡上都放一份完整的模型副本、一份完整的梯度每张卡各自维护一份完整的优化器状态。梯度通过集合通信做全局同步但显存并没有被“分片”。也就是说DDP 的显存需求约等于“单卡需要的显存 × 卡数”而且每张卡都必须满足单卡需求。如果单卡跑这个训练任务需要 30GB四张卡都是 24GB那么 DDP 跑起来的真实显存需求是 4 × 30GB 120GB超过总显存 96GB。结果就是四张卡几乎同时报 OOM而不是“总显存足够所以没事”。只有把参数、梯度、优化器状态真正切到不同卡上也就是使用 ZeRO/FSDP 这类分片策略显存才会按卡的规模“聚沙成塔”。这个概念我后面会展开先记住DDP 解决的是速度和吞吐不解决单卡放不下的问题。1.4 还有一个隐蔽原因显存碎片和 PyTorch 缓存还有一种情况更气人报错显示 CUDA out of memory但看nvidia-smi明明还有好几个空闲显存块加起来也不少为什么就不能用这通常不是真实容量不够而是显存碎片化。PyTorch 为了保证性能会预分配一大块显存作为缓存用完后不立即还给系统。频繁的分配、释放、再分配后缓存会被切成很多不连续的小块。如果你突然要申请一个大张量比如某个中间激活而剩余空间虽然总量够但找不到能容纳整块张量的连续区域分配器就会抛出 OOM。针对这种场景PYTORCH_CUDA_ALLOC_CONF环境变量很有用# 限制大块显存被切碎预留连续空间 export PYTORCH_CUDA_ALLOC_CONFmax_split_size_mb:128 # 或者允许显存段扩展减少碎片 export PYTORCH_CUDA_ALLOC_CONFexpandable_segments:True如果显存碎片严重这两个配置能救你一命。但注意如果真实占用已经接近 95%它们只是缓解不能解决容量不足。具体怎么判断是碎片还是真不够下一段写。2. 动手前先做三步诊断别一上来就改模型2.1 先看多卡是真的都在跑还是只有主卡在拼命我最常的处理方式是先开一个实时监控让 OOM 现场自己开口说话watch -n 1 nvidia-smi然后观察几轮。这里有几个典型画面只有0号卡显存满其他卡空闲基本可以断定代码没有做分布式初始化模型和优化器状态全堆在一张卡上四张卡都满了但某张卡先 OOM大概率是 DDP 场景下每卡都持有完整副本单卡容量不够显存占用忽高忽低OOM 发生在某个峰值瞬间要么是 batch 太大导致激活峰值要么是显存碎片导致峰值申请失败。如果程序已经在 OOM 前卡住还可以在训练循环里临时加一段打印print(torch.cuda.memory_summary(deviceNone, abbreviatedFalse))它会把 PyTorch 显存分配器的缓存状态、最大分配量、当前保留量都打出来。很多人一看到nvidia-smi有空闲显存就觉得奇怪但memory_summary里的reserved和allocated能更准确反映问题。reserved是 PyTorch 向显卡预留的缓存allocated是其中真正有张量占用的部分。如果reserved已经接近上限而allocated不高说明大量显存被缓存和碎片占住了。2.2 用峰值统计量化“到底缺多少”不要靠猜直接把峰值显存打出来torch.cuda.reset_peak_memory_stats(devicelocal_rank) # 跑一个完整的训练 step optimizer.zero_grad() out model(**batch) loss out.loss loss.backward() optimizer.step() torch.cuda.synchronize() peak_gb torch.cuda.max_memory_allocated(devicelocal_rank) / (1024 ** 3) print(f进程 {local_rank} 峰值显存: {peak_gb:.2f}GB)如果峰值已经接近 24GB 甚至超过 24GB那就是真容量不够。如果峰值只有 18GB 却 OOM则优先考虑显存碎片、缓存和分配策略问题。这里多说一句训练过程中有些临时张量在反向完成后会被释放但 PyTorch 不急着把缓存还给 GPU所以nvidia-smi上看到的“占用”可能一直很满。这不影响实际可用的新张量除非是碎片化。2.3 做一个最小冒烟测试再逐步加量我自己的习惯是先在同一个脚本里加一个开关强制用小 batch、短序列跑一步batch_size1seq_len64如果用的是微调框架把它官方的最小配置跑通如果最小配置能跑再逐步把seq_len加到 512、1024、2048或者把 batch 从 1 加到 2、 4。每一步记录崩溃前的峰值显存。一般会看到一个明显拐点加载模型时就 OOM说明模型权重加 CUDA 上下文已经超了单卡需要 FSDP/ZeRO 或模型并行前向阶段 OOM通常是激活值太大反向阶段 OOM通常是梯度或自动混合精度相关优化器 step 阶段 OOM通常是优化器状态太大。冒烟测试跑不通后面的花哨配置都不要谈。2.4 区分真 OOM 和假 OOM假 OOM 指的是“显存其实够但分配器找不到连续块”除了用PYTORCH_CUDA_ALLOC_CONF调节外还可以在代码里主动做一次整理torch.cuda.empty_cache()但请注意empty_cache()只释放 PyTorch 缓存中尚未分配给张量的部分对已经被张量占用的显存没有任何作用。所以如果你在每个 step 里都调用empty_cache()会发现训练慢得离谱OOM 也没解决。它适合放在“内存碎片已经很严重”的节点不适合救火。我遇到过一个案例模型、优化器、batch 怎么算都够但每隔几百步就 OOM最后加上expandable_segments:True后就稳定了。这类问题不改代码只改环境变量就能解决但如果不做诊断很容易被当成“需要换更大的卡”。3. 多显卡训练到底怎么选DDP、FSDP、还是 LoRA3.1 DDP简单易用但不减少单卡显存DDP 的正确使用姿势是每个 GPU 跑一个独立进程通过初始化进程组来同步梯度。启动时用torchrunexport CUDA_VISIBLE_DEVICES0,1,2,3 torchrun --nproc_per_node4 train.py脚本里import os import torch import torch.distributed as dist from torch.nn.parallel import DistributedDataParallel as DDP dist.init_process_group(backendnccl) local_rank int(os.environ[LOCAL_RANK]) torch.cuda.set_device(local_rank) model load_model(...).to(local_rank) model DDP(model, device_ids[local_rank])DDP 的好处是代码改动小、稳定性高、加速效果明显。它的局限也很直接每张卡上还是一份完整模型、一份完整梯度、一份完整优化器状态。如果单卡本来放不下DDP 只会让四张卡同时 OOM而不是把显存需求降到四分之一。所以 DDP 适合什么场景适合单卡刚好能跑但训练太慢想把速度提上去的情况。不适合单卡已经 OOM 的场景。3.2 FSDP/ZeRO真正让显存“分片”FSDPFully Sharded Data Parallel是 PyTorch 提供的基础能力思想与 DeepSpeed 的 ZeRO 类似不再每张卡保存完整参数和优化器状态而是把参数、梯度、优化器状态切成 N 份每张卡只保存 1/N。训练过程大致是前向时如果需要某个 transformer block 的参数通过 all-gather 把完整参数临时收集回当前卡用完整参数做前向计算完激活后把别人的参数分片释放反向时把梯度 reduce-scatter只保留本卡负责的那一段梯度优化器 step 只更新本地分片对应的参数。这样一来优化器状态和梯度的显存占用几乎可以按卡数等比例下降。7B 全参数微调在 4 × 24GB 上就有机会跑但要注意通信开销显著增加训练速度会变慢。FSDP 的坑也不少最关键的是包装粒度。要把模型按 transformer layer 级别切分而不是把整个模型包成一个整体from torch.distributed.fsdp import FullyShardedDataParallel as FSDP from torch.distributed.fsdp.wrap import transformer_auto_wrap_policy # 假设你的模型有 MyTransformerDecoderLayer policy transformer_auto_wrap_policy( transformer_layer_cls{MyTransformerDecoderLayer} ) model FSDP(model, auto_wrap_policypolicy)如果你把整个模型直接塞进 FSDP没有对 layer 做分割那它本质上还是一个大块通信和显存都很难看。另外FSDP 配合混合精度时要小心mixed_precision需要显式配置否则容易出现精度和显存两个问题同时出现。3.3 LoRA成本最低的“绕行”方案如果你的目标不是研究全参数微调而是把某个开源大模型适配到自己的数据上那么 LoRA 往往是最快路径。LoRA 的思路是冻结原模型权重只在若干层旁边插入小规模的 adapter 参数。训练时只更新 adapter所以优化器状态和梯度的大小从 7B 级降到几百万参数级。模型权重虽然还要加载到显存里参与前向但不再需要为它保存 Adam 状态显存压力瞬间小一大截。以 7B 模型 bf16 为例加载权重需要 14GBLoRA 训练动态显存可能只要几 GB。这样在单张 24GB 卡上都可能跑起来多卡场景下更是轻松不少。代码层面你可以手动只让 adapter 参数需要梯度for name, param in model.named_parameters(): if lora not in name: param.requires_grad False训练优化器只接收requires_gradTrue的参数即可。当然实际项目里一般会用现成的 LoRA 工具但核心就是这么一回事冻结主体只训练局部轻量参数。3.4 选型速查场景推荐方案原因单卡能跑想提速DDP改动小速度快单卡差一点batch 压一压能跑混合精度 梯度检查点先把峰值降下来再考虑多卡单卡放不下权重/优化器状态FSDP/ZeRO真正的显存分片只想快速微调预算有限LoRA优化器状态大幅下降超长序列导致激活爆炸梯度累积 激活检查点降低激活峰值几十 B 以上超大模型FSDP 模型并行单靠分片通信压力过大选型时先回到问题本身你是卡在权重、卡在激活还是卡在优化器状态对症下药比什么都重要。4. 实战记录一个 7B 模拟项目从单卡 OOM 到四卡跑通4.1 现场还原看起来在跑多卡其实没有我处理过一个模拟项目目标是微调一个 7B 级别的语言模型机器是 4 张 24GB 显存的显卡。最初脚本写得非常“朴实”python train.py模型加载顺利loss 也正常下降但跑了一会儿就在第一张卡上报torch.OutOfMemoryError: CUDA out of memory。打开nvidia-smi画面很典型0 号卡显存爆满1/2/3 号卡显存占用几乎为 0。这其实就说明整个模型、优化器状态、梯度全部堆在 0 号卡上其他卡完全是摆设。原因不是“显存不够”而是没有启动分布式。于是我把启动命令改为export CUDA_VISIBLE_DEVICES0,1,2,3 export PYTORCH_CUDA_ALLOC_CONFexpandable_segments:True torchrun --nproc_per_node4 train_fsdp.py这一步至少让多卡真正参与训练了。但问题没有结束四张卡开始同时 OOM。4.2 逐个降低显存峰值把打印峰值显存的代码加进去后发现每个 rank 的峰值都在 28GB 到 32GB 之间明显超过单卡 24GB。这就确认了不是碎片问题是真实容量不够。我当时做的第一件事是开启梯度检查点也叫激活检查点model.gradient_checkpointing_enable()它的原理是不保存前向过程的全部中间激活反向时重新计算一部分。用计算时间换显存效果立竿见影。训练流程里的激活峰值能降一大截代价是训练变慢。接着我把 batch size 降下来并配合梯度累积维持整体训练 batch 不变accum_steps 8 for micro_step, batch in enumerate(dataloader): out model(**batch) loss out.loss / accum_steps loss.backward() if (micro_step 1) % accum_steps 0: optimizer.step() optimizer.zero_grad()这里要理解为什么梯度累积能救 OOM它把一个大 batch 拆成多个 micro batch每个 micro batch 的前向激活在下一个 micro batch 开始前就被释放所以激活峰值下降到 micro batch 级别。但它不会显著降低模型权重、梯度和优化器状态的占用。如果峰值来自激活梯度累积很有效如果峰值来自优化器状态它帮不了你。4.3 重新组合方案FSDP 低精度 LoRA由于 7B 全参数微调即使有梯度检查点优化器状态依然占大头我最终采用了 FSDP 配合低精度并挂了 LoRA 进一步降低优化器压力。FSDP 包装时按 transformer 层自动分片使用 bf16 混合精度from torch.distributed.fsdp import FullyShardedDataParallel as FSDP from torch.distributed.fsdp import MixedPrecision from torch.distributed.fsdp.wrap import transformer_auto_wrap_policy mixed_precision_policy MixedPrecision( param_dtypetorch.bfloat16, reduce_dtypetorch.bfloat16, buffer_dtypetorch.bfloat16, ) wrap_policy transformer_auto_wrap_policy( transformer_layer_cls{MyTransformerDecoderLayer} ) model load_model_with_lora(...) model FSDP( model, auto_wrap_policywrap_policy, mixed_precisionmixed_precision_policy, )需要注意的是带 LoRA 参数的模型在 FSDP 包装顺序上要格外小心不同框架的兼容性也不同。我的原则是先确认单个进程内的参数分配正常再做 FSDP 包装。如果你发现包装后 LoRA 参数没有被训练先检查requires_grad的设定是否在 FSDP 外层保持一致。4.4 显存日志和最终效果改造后再次打印峰值显存四张卡的负载分布大致变成rank0: max_allocated 21.8GB / reserved 23.1GB rank1: max_allocated 21.2GB / reserved 22.6GB rank2: max_allocated 20.9GB / reserved 22.3GB rank3: max_allocated 21.5GB / reserved 22.9GB虽然离 24GB 上限不远但已经能稳定训练。注意这里有个容易误解的点FSDP 不是把所有显存等分成四份。因为在每个 transformer block 的前向过程中FSDP 需要临时把所有分片参数 all-gather 回当前卡所以当前块的完整参数和激活仍然会暂时出现在同一张卡上。峰值会高于理论 1/4但相比 DDP 每卡都存全量已经缓解太多。后续我又进一步压缩了序列长度、调整了max_split_size_mb并把不必要的评估步骤从训练循环里拿掉训练总算稳定跑完了。4.5 关于性能的取舍FSDP 能跑不代表一定要用。如果你的机器网络带宽一般FSDP 的通信开销会让训练明显变慢。我在实测中单步时间比 DDP 模式增加了不少但因为 DDP 模式根本跑不起来所以这个慢是值得的。如果只是做指令微调LoRA 是一个更轻的选择。LoRA 让优化器状态变得极小很多时候甚至不需要 FSDP单卡就能跑。但如果你想了解多卡训练机制、排查奇怪 OOMFSDP 这套逻辑值得亲手跑通一遍。5. 常见报错和排查经验速查5.1 常见问题速查表现象可能原因建议处理只有 0 号卡 OOM其他卡空闲没有初始化分布式使用torchrun DDP/FSDP四张卡同时 OOM每张卡都持有完整模型和优化器状态改用 FSDP/ZeRO 或 LoRAnvidia-smi还有空闲显存但程序 OOM显存碎片或 PyTorch 缓存问题设置expandable_segments/max_split_size_mbbatch_size1仍然 OOM激活值或权重/优化器状态本身超卡开梯度检查点、LoRA、低精度torch.cuda.empty_cache()无效显存被实际张量占用不是缓存不要依赖它查真实峰值训练中途 OOM不是加载时激活值峰值 / 动态 shape 导致碎片加梯度累积固定序列长度5.2 调优过程中的三个经验第一显存峰值日志一定要加。不要只在 OOM 后看nvidia-smi因为崩溃时很多临时张量已经释放你很难还原现场。最好在训练循环里定期打印torch.cuda.reset_peak_memory_stats() # 一个完整的 forward/backward/optimizer step loss.backward() torch.cuda.synchronize() peak torch.cuda.max_memory_allocated() / (1024**3)第二混合精度要确认真的生效。很多人只用了autocast以为就已经减半但 autucast 只影响部分计算权重和优化器状态不一定会变成低精度。用 FSDP 的MixedPrecision或确认参数dtype真的变成了torch.bfloat16。如果平台支持 bf16它通常比 fp16 更稳省去手动调节 loss scaling 的麻烦。第三环境变量可以在不改变代码的情况下解决不少问题。PYTORCH_CUDA_ALLOC_CONF是我每次排查 OOM 必试的选项。但要注意不是所有 PyTorch 版本都支持expandable_segments如果报错可以退回max_split_size_mb。建议先查当前 PyTorch 版本对这几个配置的兼容性再决定用哪个。5.3 一个容易被忽略的坑多卡负载不均匀即便用了 DDP 或 FSDP如果 batch size 很小每张卡算的任务量不一定均衡如果模型里某层有序列级别的全局操作也可能造成某张卡比其他卡多占一部分显存。多卡训练时不要只看总显存要看每张卡的最高占用。某张卡先 OOM 而其他卡还有余量说明负载倾斜可以尝试调整 batch 分配或减少该卡的额外任务。写在最后的实操体会我处理这类 OOM 问题最深的体会是先量化再动手。不要看到报错就先怀疑模型结构、降低学习率、换优化器这些和显存峰值往往没有直接关系。真正有用的工具是峰值统计、nvidia-smi卡位观察、以及最小配置冒烟测试。一旦定位到瓶颈是哪一块方案就很清晰激活高就减 batch、开梯度检查点优化器高就上 LoRA 或 FSDP碎片高就用环境变量调分配策略。如果你现在也卡在同样的报错上建议按文章的顺序走一遍先确认多卡真的在跑再打印峰值最后决定要不要上 FSDP。多显卡训练大模型不是把代码里加一行cuda()就能解决的事但只要知道显存到底被谁吃掉了绝大部分 OOM 都能找到出路。