显存明明够用，但用DeepSpeed跑Llama微调总是OOM，求大佬指点

检查下自定义forward里有没有创建临时变量没清干净，我之前被这个坑过，显存泄漏了。

这问题我折腾过挺久，你提到跑demo能过但自己代码不行，大概率是数据加载或collator那块出了问题。可以试试把dataloader的num_workers调成0，或者检查下是不是在forward里不小心保留了中间变量没释放。另外A100 40G跑7B模型，ZeRO-3开offload加上gradient checkpointing，batch size 16确实不该爆，但如果你用了自定义的loss计算或者padding策略，可能会让显存占用翻倍。我之前遇到过类似情况，最后发现是自己写了个很长的序列拼接逻辑，导致实际序列长度远超max_length。建议你把trainer的日志调成debug级别，看看每一步显存占用变化，或者干脆先跑一个极小的数据集（比如32条），排除数据量问题。还有个小技巧：用torch.cuda.memory_summary()打印详细分配，能定位到是哪一层在吃显存。不过你提到model.to('那里没写完，我猜你是想设置device？如果是，记得offload到CPU后不要手动设device，让DeepSpeed接管gpu分配。

offload时注意一下CPU内存够不够，我之前就是CPU内存爆了导致假OOM。

检查下自定义forward里有没有创建临时tensor没释放，我之前也被这个坑过。

试试把trainer里的gradient_accumulation_steps调小点，或者检查下自定义forward里有没有没释放的中间变量。

遇到同样的问题，A100 80G跑7B有时也会莫名其妙的OOM。你提到的自定义forward函数确实可能是坑，DeepSpeed对模型结构有隐含要求，比如某些自定义的激活函数或者层会破坏它的内存分配逻辑。建议你试试先把demo脚本里的模型结构一步步替换成你自己的代码，定位到具体哪一步开始报错。另外model.to('cuda')之后才初始化DeepSpeed也可能导致显存碎片，试试先初始化再load模型。

这情况我遇到过类似的，感觉你问题大概率出在自定义forward函数上。DeepSpeed ZeRO-3对模型结构要求挺苛刻的，它会把参数分片到不同设备，如果你的forward里有某些操作直接引用了未包装的原始参数，或者自己写了显存管理，就可能跟DeepSpeed的内存追踪机制打架。我上次就是有个自定义的attention mask计算没用torch原生的方法，一跑就OOM，换成HuggingFace的demo反而稳得很。另外你说的offload到CPU和gradient checkpointing，其实这两个配合不好的话也会炸，比如offload的阈值设得太低，CPU和GPU之间来回倒腾反而占更多临时显存。建议你把所有自定义层都换成HuggingFace兼容的实现，或者试试先用ZeRO-2跑一遍排除问题。还有你那个model.to('cuda')是不是放在DeepSpeed初始化之后了？顺序不对会导致重复分配显存。

碰到过类似问题，很多时候是自定义forward里有些临时变量没释放，比如中间tensor没显式del或者detach，ZeRO-3的显存管理会更敏感。建议先用torch.cuda.memory_summary()看看峰值出现在哪个环节，另外检查下trainer里有没有不小心把模型又to('cuda')了一次，导致offload失效。数据加载那边也可以试试把num_workers设成0或者pin_memory关掉，有时候多进程会预占显存。

我也碰到过类似的坑，最后发现是自定义的forward里有个变量没注册到模型参数里，导致内存分配异常。建议你检查下代码里有没有动态创建tensor或者没用nn.Module包装的操作。另外model.to(device)之后，确保所有输入也都正确放在同一设备上，有时候一个小疏忽就会让ZeRO3的显存调度炸掉。

我遇到过类似的坑，问题很可能出在data collator或者自定义dataset的显存占用上。你那个能跑通的demo大概率是用了最简单的padding方式，而你的代码可能无意中把整个序列都pad到了最大长度，导致实际batch占用的显存远高于预期。建议用torch.cuda.memory_summary()看一下具体哪块在涨，另外offload到CPU后记得检查一下optimizer和gradient的显存分配，有时候DeepSpeed的配置里offload参数没写全也会导致它失效。

八成是自定义forward里用了绝对position embedding，ZeRO-3对动态图支持有坑，试试关掉allgather的overlap。

我之前也踩过类似的坑，ZeRO-3下显存够不够真不能光看batch size，它会把参数、梯度和优化器状态全部分片，但如果你代码里有任何地方触发了全量参数收集（比如自定义forward里用了model.parameters()或者调用了torch.no_grad下的完整推理），那内存就会瞬间爆掉。我猜你那个demo脚本能跑通，很可能是因为它没做额外的张量拼接或index操作，而你的数据加载里如果用了动态padding或者把多个样本塞成一个batch，那显存占用会比理论值高好几倍。offload到CPU其实也有坑，尤其是当你有频繁的CPU-GPU张量拷贝时，通信开销反而会让显存碎片化更严重，建议你监控一下nvidia-smi看是不是有峰值尖刺。另外，你model.to('后面没打完，我猜你是不是想放device map？如果是的话，记得跟DeepSpeed的配置里zero_optimization.stage3_gather_16bit_weights_on_model_save要配合，不然保存时也会OOM。还有个排查技巧：把你自己的trainer换成HuggingFace原版Trainer，但保留你的dataset和collator，如果这样能跑通，那问题大概率在自定义的training step里，比如loss.backward前有没有手动调zero_grad或者梯度累积步数没对齐。最后，检查下有没有不小心把eval模式下的模型也塞进同一张卡，有时候验证集forward没包在torch.no_grad里，内存就悄悄翻倍了。

八成是自定义forward里用了绝对位置编码或者缓存，ZeRO-3对动态图特别敏感，试试把offload改给NVMe。

我之前也栽过，后来把trainer的remove_unused_columns设False就好了，你查查是不是数据collator返回了多余字段。

说实话看到你这个配置我第一反应是batch size 16确实有点激进，A100 40G跑7B模型就算ZeRO-3+offload，单卡极限也就12-14左右，而且你说显存明明够用，这个“够用”是不是看的是nvidia-smi的显存占用而不是实际分配？很多情况下DeepSpeed会预分配内存池，nvidia-smi显示的不是真实峰值。我之前也遇到过类似情况，后来发现是自定义forward函数里有个中间变量没释放，导致计算图一直挂着，你检查一下是不是有类似的问题，特别是用了梯度检查点的话，有些算子会强制保留激活值。另外你说HuggingFace的demo能跑通，那多半不是配置问题，而是你代码里某个地方隐式创建了新的计算图，比如在loss计算里用了Python原生循环而不是torch的算子，这会让ZeRO-3的分片策略失效。还有个坑是如果你在trainer外面手动调了model.to('cuda')，但DeepSpeed内部又自己重新分配设备，两套地址空间会打架，建议把设备管理完全交给DeepSpeed。最后你可以试试在训练循环里加torch.cuda.empty_cache()，虽然治标不治本，但能帮你判断是不是碎片化问题，如果加了之后能跑更久，那就说明是内存碎片而不是容量不够。

八成是自定义forward里没用model.inputs，或者数据没pad到同一长度，查查collator吧，我之前就这么栽过。

我之前也踩过类似的坑，A100 40G跑7B理论够但实际玄学很多。你试试把trainer的batch size直接砍到1，然后gradient accumulation设成16，这样能排除是不是显存碎片化的问题。另外自定义forward里如果有动态shape或者临时tensor没释放，DeepSpeed的显存规划会直接崩，建议在关键层加torch.cuda.empty_cache()看看。还有你model.to('cuda')之后有没有把input也显式放到device上？有时候数据在CPU上会触发隐式拷贝，直接爆显存。

八成是自定义forward里没走deepspeed的梯度挂钩，试试把模型包进deepspeed.initialize后别手动调device。

其实我之前也被这玩意儿坑过，A100 40G跑7B理论算力确实够，但问题往往出在“显存够”这个判断上。你试试把batch size直接降到1，然后开gradient accumulation，哪怕多攒几步，也比直接爆掉强。另外你提到offload到CPU了，但确认下是不是把optimizer states和param都offload了，有时候只offload了gradient，峰值照样炸。

更可疑的是你说demo能跑通，自己的代码不行，那八成是你自定义forward里有一些中间变量没释放，或者你用了什么额外的activation缓存。我之前就是因为在forward里存了个attention mask的副本，结果显存悄悄翻倍。你可以试着用torch.cuda.max_memory_allocated()打一下峰值，对比demo脚本，看看差在哪。

还有你怀疑DeepSpeed和自定义forward不兼容，这个确实有，特别是你如果在forward里用了Python原生的list或者dict存tensor，ZeRO-3的分片逻辑有时候会把这些当普通对象，不参与显存管理。建议把所有中间量都改成显式tensor，或者干脆包成ModuleList。

最后检查一下你数据加载是不是用了pin_memory=True，有时候这个和DeepSpeed的异步加载撞了，也会导致临时显存暴涨。我上次就是关掉pin_memory就好了，莫名其妙。

我之前也踩过类似的坑，ZeRO-3对显存的计算方式跟你直观感觉不太一样，它会把参数分片后的临时buffer也算进去，所以光看batch size和模型大小没用。你那个demo能跑通很可能是因为它用了更简单的数据collator或者序列长度短，你检查下自己代码里有没有把输入pad到固定长度，这个影响很大。另外自定义forward如果是用inplace操作或者有中间变量没释放，也会让显存峰值暴涨，建议用torch.cuda.max_memory_allocated打点看看具体哪一步爆的。还有你model.to('后面是不是漏写了设备？有时候offload配置没生效就会这样。

我之前也踩过类似的坑，A100 40G跑7B其实挺极限的，尤其是ZeRO-3本身会预留一部分显存做参数分片和通信缓冲。你试试把batch size降到8，然后看看是不是自定义forward里用了什么会额外创建大tensor的操作，比如attention mask的拼接或者无用的中间变量。另外，demo脚本能跑通但你的不行，大概率是trainer的datacollator或者input_ids的padding方式不一样，导致实际序列长度比预期长很多。还有个冷门点，检查下是不是把model.to('cuda:0')写在了deepspeed初始化之后，顺序错了会白占一份显存。