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

我之前也踩过类似的坑，问题多半不在显存理论占用上，而是你自定义的forward里可能有隐式的显存碎片化，或者某个中间变量没被释放。建议先用torch.cuda.max_memory_allocated打点看看峰值到底在哪个阶段爆的，别只盯着batch size算。另外，ZeRO-3下model.to('cuda')会干扰参数分片，最好让DeepSpeed自己管理设备，你手动to反而可能触发全量拷贝。我之前是把trainer里的model.to去掉，改成用deepspeed.initialize的model，问题就解决了。你那个demo能跑通，大概率就是因为它没做这些多余操作。

我之前也踩过类似的坑，A100 40G跑7B按理说ZeRO-3加offload确实够，但问题往往不在显存总量，而在碎片化和显存峰值。你那个demo能跑通自己的不行，我赌八成是数据collator或者attention mask的shape没对齐，导致某些batch的sequence特别长，触发了一次性分配超大临时tensor。另外自定义forward里如果用了torch.where或者masked_fill这类操作，DeepSpeed的partitioned parameters在反向传播时可能会产生额外的activation副本，这个很隐蔽。建议你先把batch size降到1，开一个profiler看每步的峰值显存，重点盯一下optimizer step前后的变化。还有，model.to('cuda')之后最好再调一次deepseed.initialize()，顺序反了会导致offload配置失效。你试试在dataloader里显式设置drop_last=True，有时候最后那个batch不够整除也会让ZeRO-3的梯度规约出现奇怪的临时缓冲。要是还不行，直接开NCCL_DEBUG=INFO跑一遍，看是不是有节点间通信时同步了未分片的参数。

是不是你自定义forward里返回了多余tensor导致activation没释放？之前我踩过类似的坑，检查下loss计算那块。

八成是数据collator没drop_last，最后那个batch不够16显存反而炸了，试试padding到固定长度。

我之前也踩过类似的坑，多半不是显存不够，而是你的数据加载器在每次step时把整个batch的input都塞进了GPU，加上你自定义forward里可能有临时变量没释放，导致峰值内存暴涨。建议你把trainer的batch size先降到4试试，如果还爆就查一下是不是dataloader的pin_memory或者num_workers设置太高，把CPU内存也挤爆了。另外，ZeRO-3的offload对某些自定义算子支持不好，你试试把offload改成只offload optimizer state，或者干脆用ZeRO-2对比一下，排除是offload策略的问题。最后，model.to('cuda')之后记得检查一下有没有把embedding或norm层也放回device，有时候这俩会被DeepSpeed自动挪走，你手动再放回去反而冲突。

另一个思路：你确定“显存够用”是看nvidia-smi的占用吗？有时候预留显存和实际可用显存差很多，尤其开了CUDA graph或者显存碎片化的时候。我遇到过类似情况，最后发现是自己在forward里用了torch.no_grad()但没关，导致中间激活值没有被正确释放，白白占了几G。你可以在每个step前后打印一下torch.cuda.max_memory_allocated()，看看峰值到底出现在哪一段代码，这样比瞎猜效率高多了。如果峰值确实在数据加载阶段

大概率是数据加载时每个step的显存峰值没控制住，试试把dataloader的num_workers调成0或pin_memory关掉。

这报错看着像自定义forward里有隐式建图没进ZeRO管理，建议检查下有没有用torch.no_grad或临时tensor没释放。

大概率是你自定义forward里没用model inputs，导致ZeRO-3的partition没生效，检查下inputs的device和requires_grad。

试试把offload换成NVMe或者调小batch，A100 40G跑7B全参微调本来就紧巴巴的。

我之前也踩过类似的坑，A100 40G跑7B按理说ZeRO-3加offload应该够，但问题多半不在显存总量，而是碎片化或者某个中间变量突然暴涨。你换个思路试试，用torch.cuda.max_memory_allocated()打个断点看看到底是哪一步峰值上去的，我那次就是自定义loss里有个超大临时张量没释放。另外你说demo能跑自己代码不行，建议把trainer的remove_unused_columns设为False，有时候数据预处理会在forward前多复制一份到GPU，这坑特别隐蔽。至于model.to()那个截断，是不是想写device_map？如果是手动分配的话，检查下是否和ZeRO-3的partition逻辑冲突了，我之前就因为这俩打架导致OOM。

八成是自定义forward里用了绝对位置编码或者中间变量没走DeepSpeed的显存优化，试试把offload的cpu_offload改成nvme_offload再开zero_quantized。

你那个demo脚本是不是没开evaluation？我上次就是validation时多算了几个logits直接爆显存，把eval_accumulation_steps调小就好了。

我之前也踩过类似的坑，最后发现是dataloader的num_workers开太多，每个worker都复制了一份模型状态，显存直接翻倍。你试试把num_workers改成0或者pin_memory关掉，说不定就好了。另外自定义forward里如果有临时tensor没释放，也可能导致显存碎片化，建议加torch.cuda.empty_cache()在step之间手动清一下。

还有个小细节，ZeRO-3下model.to('cuda')之后要确保所有参数都通过zero.GatheredParameters访问，如果你在自定义loss里直接用了model.parameters()，会强制把全部参数拉回GPU，那offload就白设了。我之前就是在这上面卡了一整天。

试试把trainer里的remove_unused_columns设成False，八成是你自定义数据集里多了些没用的键导致显存炸了。

我之前也踩过类似的坑，A100 40G跑7B按理说ZeRO-3加offload是够的，但你换成自己的代码就OOM，大概率不是显存问题，而是你自定义forward里某个中间变量没被释放，或者有个隐式的超大tensor被存下来了。建议你先开一下DeepSpeed的显存日志，看看是哪个step爆的，同时检查下你数据加载时是不是偷偷把整个batch的input_ids都搬到了GPU上。另外那个demo能跑通，说明配置没问题，你对比下两边在dataloader和trainer的collate_fn上有没有差异，我上次就是死在自定义loss里多存了个logits。如果查不到，试试把batch size调成1跑一遍，如果还爆，那就是代码里有显存泄漏，跟DeepSpeed本身关系不大。

我之前也踩过类似的坑，A100 40G跑7B按理说ZeRO-3加offload是够的，但问题往往出在你自己代码里那些隐藏的tensor上。你试试在训练循环里加个torch.cuda.max_memory_allocated()看看峰值到底在哪一步爆的，八成是forward里某个中间变量没释放，或者你用了自定义loss导致activation被额外保留了。另外DeepSpeed和HF的Trainer集成其实挺挑版本的，你确认下transformers和deepspeed的版本匹配吗？我之前就是版本不兼容导致offload根本没生效，换成官方推荐的组合立马就好了。还有个骚操作，你把batch size降到1先跑通，然后逐步往上加，同时盯紧nvidia-smi的显存变化，这样定位起来比瞎猜快多了。

我前两天刚踩过类似的坑，最后发现是trainer里没关掉模型本身的gradient_checkpointing，跟DeepSpeed的配置重复了，显存直接翻倍。你那个demo能跑通很可能是它默认用了一个干净的training_args，而你自己的代码里可能隐藏着一些冗余的配置。另外你说的自定义forward函数，我怀疑问题不在forward本身，而是你在里面创建了新的tensor没及时释放，比如用了什么中间变量没删，导致计算图一直没被清掉，ZeRO-3对这部分特别敏感。建议你先把offload关掉，batch size降到4跑一次，看还能不能OOM，如果还是爆，那基本就是代码里有显存泄漏，而不是DeepSpeed的问题。还有个思路，你试试在每次step之后手动torch.cuda.empty_cache()，虽然治标不治本，但能帮你确认到底是不是缓存碎片的问题。我自己最后是发现数据加载器里有个collate_fn把整个batch的attention mask都复制了一份到GPU，改掉之后显存从39G直接降到20G，你可以往这个方向查查。

八成是自定义forward里没用model.inputs或position_ids这些，DeepSpeed钩子没接上，试试关掉offload纯ZeRO-3跑一下看报错栈。

我之前也踩过类似的坑，最后发现是dataloader里num_workers设太高，CPU内存被撑爆后反过来拖垮了显存分配。你试试把workers调成0或者2，顺便检查下自己的forward里有没有重复建张量没释放。另外ZeRO-3和自定义forward确实容易有兼容问题，可以试试关掉zero_optimization里的stage3_gather_16bit_weights_on_update，或者改成ZeRO-2对比下。还有个小细节，你model.to('cuda')之后记得设一下device_map，有时候不指定会重复分配。

我也遇到过类似情况，最后发现是dataloader里num_workers开太多，每个worker都复制了一份模型状态，显存直接翻倍。建议你把worker设成0或者1试试，另外检查下是不是有地方不小心调用了model.cuda()把offload的tensor又拉回GPU了。自定义forward的话，如果里面有不必要的中间变量保存，也可能导致activation显存暴涨，可以试着用torch.no_grad包一下不需要梯度的部分。

我之前也踩过类似的坑，最后发现是自定义forward里有个临时tensor没显式释放，导致ZeRO-3的显存统计和实际分配对不上。你试试在每次step后清一下cache，或者把offload的optimizer状态单独开一个分区，有时候默认配置会漏掉某些参数。另外检查下dataloader是不是有num_workers泄漏，那个也会悄悄吃显存。如果demo能跑而你的不能，建议把trainer的save_strategy关掉试试，我之前就是checkpoint保存时突然申请大块显存炸了。

我之前也踩过类似的坑，最后发现是dataloader的num_workers开太多，加上自定义collate_fn里偷偷把数据pad到max_length，显存一下子就被吃满了。建议你把batch size先降到1，然后开--profile看看实际峰值显存，大概率会发现比预期高不少。另外，如果你在forward里手动创建了大的中间tensor（比如attention mask扩展），就算offload了也没用，这些临时变量不会进CPU。至于model.to('cuda')后面那句被截断了，猜测你是不是想写device_map='auto'？那玩意儿跟ZeRO-3一起用有时候会互相打架，建议直接删掉试试。

同样配置demo能跑自己的不行，八成问题出在data collator或者模型输入上，检查下有没有把labels传到不该传的地方。我之前遇到过类似情况，是自定义loss里对logits做了reshape导致显存峰值暴涨，ZeRO-3下这个放大效应特别明显。建议你开一下显存监控，把每个step的峰值打出来，对比demo和自己代码的差异，基本能定位到是forward还是backward爆的。另外可以试试把batch size降到4跑一次，如果还爆就大概率不是batch size的问题，而是某个中间变量没释放。

你这代码八成是偷偷把模型又拷了一份到显存，查查dataloader里有没有意外的tensor累积吧，之前被这坑过。