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

我之前也踩过类似的坑，问题往往不在显存总量，而是碎片化或临时buffer爆掉。你试试把batch size调到1，然后梯度累积设成16，看还OOM不，要是能跑，基本就是峰值内存的问题。另外检查下dataloader的num_workers，有时候多进程加载数据会额外吃不少显存，尤其你开offload的时候，CPU和GPU之间搬运很容易卡在某个瞬时峰值上。自定义forward的话，最好确认下有没有隐式创建大tensor，比如用了mask或者position_ids没及时释放，这种小细节最容易漏。我上次就是被一个view操作搞出的临时张量坑了整整两天。

遇到过类似的坑，最后发现是data collator在搞鬼。你的demo脚本能跑通但自己的代码不行，大概率不是模型本身的问题，而是数据在进模型之前就被pad到了奇怪的shape。我上次就是自定义了一个collate函数，结果padding策略和DeepSpeed的梯度累积对不上，显存直接翻倍。你可以试试在trainer里把remove_unused_columns设为False，有时候自定义forward会触发这玩意儿，导致额外的tensor被传进设备。

另外你提到model.to('后面没打完，我猜你是不是把模型放到CPU上了然后再offload？这个顺序其实很关键，ZeRO-3要求先初始化模型再往GPU上放，如果你在DeepSpeed引擎初始化之前就调了model.to('cuda')，参数会被复制一份到默认设备，显存直接炸。正确做法是让DeepSpeed自己管理设备，别手动干预。

还有个小细节，gradient checkpointing和ZeRO-3一起用的时候，需要确保checkpoint的输入是tuple而不是list，不然会触发多次反向传播，显存占用成倍增长。我这边之前就因为一个return dict的问题卡了一整天，改成tuple就秒过。

如果实在排查不出来，建议用torch.cuda.memory_summary()看下峰值分配在哪一层，或者直接跑一个空的dataloader看是不是数据加载器的问题。这种问题往往不在模型，而在你没想到的角落，比如tokenizer的padding_side设置不对，或者attention_mask没传对，都会让显存波动很大。祝早日跑通。

八成是自定义forward里用了绝对位置编码或者缓存了张量，ZeRO-3对动态图支持很差，试试把input_ids显式pad到固定长度。

你这情况大概率是数据collator返回了变长张量，DeepSpeed的显存规划就崩了，把padding策略改成max_length试试。

我之前也踩过类似的坑，最后发现是自己dataloader里num_workers设太大，每个worker都复制了一份模型state，直接挤爆显存。你可以先试试把num_workers调成0或者1，然后pin_memory关掉看看。另外，自定义forward里如果有张量没走模型的device管理，比如硬编码了.cuda()，也会让DeepSpeed的显存统计失效，得确保所有中间变量都显式放到当前device上。还有，offload到CPU不是万能的，有时候反而会因为频繁的CPU-GPU拷贝触发碎片化，建议把ZeRO-3的stage3_param_persistence关掉试试。最后检查一下你的trainer有没有手动调了model.to('cuda')，DeepSpeed会用自己的device管理，你手动搬一次可能就重复占用了。

我之前也踩过类似的坑，后来发现是自定义forward里用了绝对位置编码的tensor，导致activation显存翻倍了。你可以试试把offload改成offload_optimizer+offload_param分开设，或者直接关掉offload跑一遍，看是不是offload本身在跟gradient checkpointing打架。另外检查下dataloader有没有把整个batch一次性搬到GPU上，有时候pin_memory和num_workers设置不当会额外占显存。

还有个小技巧，用torch.cuda.max_memory_allocated()打印一下峰值，对比demo脚本和你代码的差距在哪个模块。如果是trainer配置问题，试试把DeepSpeed的zero_force_ds_cpu_optimizer设为false，某些自定义层会跟这个选项冲突。最后建议你把模型换成随机初始化的同样结构跑一次，排除数据侧的影响。

我猜问题大概率出在你自己代码的显存分配上，而不是DeepSpeed本身。ZeRO-3的offload是参数和优化器状态，但激活值、梯度以及临时张量还是留在GPU里的，如果你的forward函数里不小心创建了大尺寸的中间变量，比如注意力矩阵或者重复拼接的操作，那即使batch size是16也可能瞬间把40G挤爆。我之前遇到过类似情况，最后发现是自定义的layer norm实现里用了float64计算，直接让显存翻倍。你可以先用torch.cuda.max_memory_allocated()对比一下demo脚本和你自己代码的峰值占用，看看差距到底在哪一步拉开的。另外，model.to('cuda')之后如果有些子模块被重复指定设备，或者数据加载时pin_memory和non_blocking设置不当，也会导致隐性显存碎片。还有个排查方向，试试把batch size降到2，如果还OOM那就不是你算力预算的问题，而是代码里有显存泄漏，比如每次iter都新建一个张量没释放。至于自定义forward不兼容，我建议你把所有自定义层先替换成原生HuggingFace实现跑一次，排除嫌疑，再逐个加回来。

我之前也遇到过类似玄学问题，后来发现是自定义forward里有个中间变量没删，显存碎片把空间吃死了，你可以试试在backward后手动清一下cache。另外offload到CPU不一定管用，如果数据加载用了pin_memory或者collate_fn里有重复引用，反而会加剧内存交换。建议先用一个极小batch逐步加大，看峰值显存在哪一步爆的，同时对比demo脚本和你的代码在dataloader的num_workers上有没有差异。

我之前也踩过类似的坑，A100 40G跑7B理论上确实够，但DeepSpeed ZeRO-3的显存分配策略和你想的不太一样，它会把参数、梯度和优化器状态分片，但如果你在forward里用了activations或者中间变量没释放，照样会爆。你提到offload到CPU了，但注意ZeRO-3的offload是分层的，如果CPU内存不够或者NVMe路径没配好，反而会卡在显存和内存的交换上，特别是batch size 16的时候，每步的通信开销和临时张量会突然涨上去。我觉得问题可能不在trainer配置，而是在你的自定义forward函数里，比如有没有用到torch.no_grad()包住推理部分，或者有没有把某些tensor显式detach掉，这些都会影响显存的峰值占用。你试过用DeepSpeed的zero.Init上下文包裹模型定义吗？那个能延迟参数初始化，省下不少显存。还有一个笨办法，把batch size降到8跑一下，看看是不是线性增长，如果是，那多半是数据加载时pin_memory或者num_workers开太多导致内存碎片化。你那个能跑通的demo脚本，它是不是用了官方的Trainer而不是你自己搭的training loop？如果是的话，差距可能就在梯度累积和loss缩放的处理上，建议你打印一下每一步的显存峰值，先定位是前向还是反向爆的。

我之前也踩过类似的坑，A100 40G跑7B微调理论上确实够，但OOM往往不是显存总量的问题，而是峰值显存瞬间爆炸。你开了offload和gradient checkpointing还挂，大概率是某个中间变量没被ZeRO-3的partition机制覆盖到，比如自定义forward里用了全量参数做计算，或者显式调用了.cpu()、.detach()这类操作打断了ZeRO的tensor追踪。另外，你提到同样的配置跑HF demo能通，那差别很可能出在dataloader的collate逻辑上——如果自己写了padding到固定长度，或者动态batch里某个sample特别长，那显存占用会忽高忽低，刚好在你观察的“够用”边缘翻车。建议你把trainer的batch_size调到1，先跑一个step看看峰值，再用torch.cuda.max_memory_allocated打出来对比，基本能定位是模型层还是数据层的问题。还有，你说的model.to('cuda')那个帖子内容被截断了，如果是把模型整体to到device上，在ZeRO-3里反而可能破坏sharding，最好交给DeepSpeed的engine自己管理device。我最后是靠把自定义forward里的中间tensor全部显式contiguous()，然后重新make_partition解决的，你可以试试。

我之前也踩过类似的坑，后来发现是data collator里有个tensor没转成cuda，导致它一直留在CPU上，DeepSpeed一offload就在边界上卡爆了。你可以试试在trainer里打印一下每个batch的device分布，大概率不是forward的问题。另外，你检查过max_length吗？有时候padding到超长序列会让显存瞬间翻倍，demo脚本可能刚好绕过了这个。

我之前也踩过类似的坑，最后发现是dataloader的num_workers开太多，每个worker都会复制一份模型状态，显存直接翻倍。你可以先试试把num_workers调到0或者1，排除这个因素。另外你自己的代码里如果用了model.to('cuda')之后再包DeepSpeed，可能会触发重复显存分配，建议检查一下包装顺序。还有那个demo脚本能跑通，很可能是因为它用了官方推荐的固定随机种子和batch sampler，你这边数据长度不一致导致padding后有效batch变大。实在不行就开个nsys trace看看显存峰值出现在哪个step，比瞎猜快多了。

遇到过一模一样的坑，最后发现根本不是显存的事儿，是显存碎片化加通信缓冲区的锅。你A100 40G跑7B加ZeRO-3理论上确实够，但DeepSpeed默认会为每个layer预分配一块连续显存做all-gather的buffer，这个buffer大小是按max_seq_len和hidden_size算的，如果你自定义的forward里用了比模型默认更长的序列或者更大的中间tensor，它就会额外多占一大块。我之前是自己写了个带cross-attention的模块，结果DeepSpeed不知道这个模块的显存需求，prefetch机制直接炸了。

还有你说的offload到CPU，得确认下是不是把optimizer state和gradient都offload了，但参数还在GPU上，ZeRO-3的partitioned parameters在每次forward/backward前要全量收集，这个收集过程本身就要临时占用跟模型一样大的显存，如果你数据加载时用了pin_memory=True，又没设num_workers=0，CPU那边也会积压显存。建议你跑的时候开一下deepspeed --num_gpus=1加--autotuning看看，或者直接在代码里打印torch.cuda.memory_summary()，看下是哪个阶段涨上去的。

另外你说同样的demo能跑，多半是那个demo没开gradient checkpointing或者用的序列长度短，你检查下是不是自己的数据里batch内padding太多，导致实际seq_len比max_seq_len大很多。反正我最后是靠手动设置zero_allow_untested_optimizer和把reduce_bucket_size调小解决的，你可以试试把zero_force_ds_cpu_optimizer设成false，然后显式指定offload_optimizer_device='cpu'，有时候默认配置会在CPU和GPU之间反复搬运，反而增加峰值占用。

八成是你自定义forward里手动做了张量切片或临时变量没释放，试试关掉offload把batch调小到8排除干扰。

我之前也被这个坑过，A100 40G跑7B按理说ZeRO-3加offload是够的，但你这个问题大概率不是显存容量本身，而是碎片化或者临时峰值。你想想，offload到CPU之后，GPU显存里只留当前层参数，但反向传播时如果某个tensor特别大，比如logits或者loss那块没处理好，峰值一下就飙上去了。我猜你自定义forward里是不是有什么中间变量没显式del，或者用了大型的padding mask矩阵？这种临时tensor在ZeRO-3下会被每层都复制一份，很吃显存。

另外你提到demo能跑通，自己的代码不行，我怀疑是trainer的配置差异。比如HuggingFace那个demo很可能默认用了gradient_accumulation_steps来摊薄batch，或者把optimizer的state也offload了，而你自己的代码可能把optimizer留在GPU上，那每个参数要占8字节的Adam状态，7B模型光这就要28G，加上参数和梯度，40G卡直接爆。你检查下ds_config里optimizer的offload开了没，还有zero_force_32bit_weights或者reduce_scatter这些选项，有时候默认值在自定义模型上会异常。

还有个细节，你确认自己的模型没用model.half()或者bfloat16吗？如果和DeepSpeed的混合精度设置冲突，会导致某些层以fp32跑，显存直接翻倍。我建议你先把batch size降到1，开着zero_allow_untested_optimizer跑一下，看是不是还OOM，如果不爆了，那就是batch size乘序列长度导致的峰值问题，这时候可以试试用torch.utils.checkpoint手动包住那几个大算子，别依赖全局的gradient checkpointing。

最后，你那个model.to('cuda')的写法如果出现在DeepSpeed初始化之后，很容易出问题，因为ZeRO-3要求模型先留在CPU上，用deepspeed.initialize接管分布，你再手动to cuda反而打乱了它的分片逻辑。我之前就是在这上面卡了一整天，把初始化顺序调对之后，同样的配置直接跑起来了。

我之前也踩过类似的坑，问题往往不在显存总量，而是碎片化或者某个中间激活值瞬间暴涨。你试试把batch size直接调成1，然后梯度累积开大点，看能不能跑通，能跑通的话就是峰值显存计算偏差。另外检查下是不是dataloader的num_workers开太多，或者pin_memory=True导致额外显存占用，这个很容易被忽略。还有你那个自定义forward里如果有动态shape的操作，比如根据长度padding，可能会让DeepSpeed的显存规划完全失效，建议先用静态shape的跑一遍对照。

我之前也踩过类似的坑，A100 40G跑7B按理说ZeRO-3加offload是能塞下的，但问题往往出在你自己写的forward里——比如中间变量没释放，或者用了不支持的张量操作导致显存碎片化。你可以试试把batch size调到1，先排除是不是数据维度的锅，再在训练循环里手动清一下cache，看峰值显存到底涨到哪一步。另外你的trainer配置里有没有设find_unused_parameters？自定义模型里有些参数没参与loss计算，DeepSpeed会默认全量梯度同步，也可能拖爆显存。最后建议对比下demo脚本和你代码的model.to('cuda')时机，有时候offload到CPU的权重没正确回传，也会引发诡异的OOM。

我之前也踩过类似的坑，后来发现是自己dataloader里有个自定义collate_fn把序列padding到了固定长度，导致实际显存占用比理论翻倍。建议你先打印一下实际输入shape和token数，再看看是不是这个原因。另外，ZeRO-3的offload有时候和某些自定义forward里的显存分配方式冲突，可以试试把offload参数调细一点，比如只offload优化器状态。你那个demo能跑通，很可能就是它的数据长度和模型forward更规整。

我之前也踩过类似的坑，后来发现是自定义forward里有个中间变量没清干净，显存碎片被撑爆了。你可以试试把batch size调到1跑通，再逐步加，顺便看下nvidia-smi的显存分配是不是突然跳变。另外，offload到CPU有时候会跟gradient checkpointing打架，两个一起开反而更吃显存，可以只留一个试试。

我之前也踩过类似的坑，后来发现是dataloader的num_workers开太多，显存被数据预取占了，你试试把workers降到2或者0，顺便pin_memory关掉。另外你说自定义forward不兼容，这个确实有可能，ZeRO-3对模型结构里的动态shape或者非tensor返回值特别敏感，检查下有没有在forward里返回loss以外的额外变量。还有个笨办法，把batch size先降到1跑通，再逐步加上去，定位是显存峰值问题还是累积泄漏。

我之前也踩过类似的坑，最后发现是dataloader的num_workers开太多导致显存被额外占用，你可以试试把worker数调到0或者2看看。另外自定义forward里如果有临时tensor没及时释放，ZeRO-3的显存计算会不准确，建议用torch.cuda.empty_cache()在step间隔清一下。还有个思路是直接对比demo脚本和你的代码在模型实例化后的显存占用，如果demo跑通时显存还剩很多，那大概率是dataset的collate函数里有什么东西被重复拷贝到GPU了。