用DeepSpeed跑Llama微调，ZeRO-3总是OOM，是我配置姿势不对吗？

ZeRO-3上来就OOM，第一步就崩，大概率不是batch size的问题，而是模型初始化阶段就把显存打满了。你offload开了没错，但有个坑很多人踩：stage3下模型参数是分片加载的，如果没用zero.Init()包住模型构造过程，每个rank还是会完整实例化一份7B权重，80G直接见底。另外检查下stage3_param_persistence_threshold和stage3_max_live_parameters，默认值偏大，调小一点能明显缓解峰值。我当初也是卡在这，后来把sub_group_size降到1e6、reduce_bucket_size压到5e8才跑通。说ZeRO-2够用的其实有道理，7B单卡80G用ZeRO-2加offload，batch size开个2到4问题不大，还省了stage3那堆通信开销。建议先拿ZeRO-2验证一遍流程能不能跑，确认没问题再回头调stage3，不然你连是配置错还是框架问题都分不清。

单卡80G跑7B模型用ZeRO-3一启动就OOM，大概率不是配置姿势的问题，而是ZeRO-3本身在初始化阶段就会把完整模型参数分摊到所有可用设备上，你只有一张卡，等于没分摊，反而多了一层通信和分区开销。而且stage3会在forward之前就all-gather参数，峰值显存比stage2还高，很多人第一次用都会踩这个坑。你offload到CPU之后OOM依旧，说明瓶颈可能不在optimizer状态，而是模型加载时就已经把fp32权重全塞进显存了。可以试试先用ZeRO-2加offload跑，7B在80G上单卡基本是够的，除非你序列长度拉得特别大。另外检查下有没有设置stage3_gather_16bit_weights_on_model_save和sub_group_size，这两个参数配不好也会让stage3在初始化时炸掉。真要上stage3，建议至少两张卡起步，单卡玩stage3有点自找麻烦。

7B模型单卡80G上ZeRO-3还OOM，大概率不是显存真不够，而是stage3的param partition在初始化时会把完整模型先加载一遍。你可以试试加"stage3_param_persistence_threshold": 0和"stage3_max_live_parameters": 0，或者直接用from_pretrained时指定low_cpu_mem_usage=True。我之前也踩过这个坑，后来发现是transformers版本和deepspeed对不上，换了个组合就正常了。

单卡80G跑7B模型用ZeRO-3确实有点用力过猛了，stage 3会把参数、梯度、优化器状态全切片，通信开销大不说，配置稍微不对就容易炸。你offload全开了反而可能更糟，CPU offload会让第一步的all-gather压力集中到某张卡上，如果没配好pin_memory和NVMe路径，反而比ZeRO-2更容易OOM。我之前也踩过类似的坑，后来发现是transformers加载模型时没加low_cpu_mem_usage=True，加上device_map="auto"和DeepSpeed冲突了，模型在CPU和GPU之间反复拷贝导致显存峰值爆掉。建议你先试试ZeRO-2加offload，7B单卡80G其实完全够，batch size开个4到8都没问题。另外检查一下deepspeed config里的zero_optimization是不是漏了stage字段，或者stage设成了3但contiguous_gradients没开，这个不打开梯度会碎片化。还有个小细节是gradient_accumulation_steps和你实际dataloader的batch size要对齐，不然第一步就会因为shape不匹配直接崩。

刚入门，这个对我帮助很大。

写得挺好，建议补充一些性能数据。

单卡A100 80G跑7B还OOM确实不太正常，我第一反应是你模型是不是用float32加载的？很多人忘了设torch_dtype=torch.bfloat16，光加载权重就吃掉28G，ZeRO-3分片再一折腾显存直接见底。另外ZeRO-3启动时会先做一次完整的参数收集，第一步的峰值显存往往比稳态高不少，stage3_gather_16bit_weights_on_model_save这类参数没配对也容易炸。你可以先试试换成ZeRO-2跑一下，7B模型单卡80G用ZeRO-2加bf16基本是够的，别被网上的说法吓到。还有个容易忽略的点是transformers版本和deepspeed版本不匹配，老版本对Llama的auto tp支持有问题，会偷偷把整个模型铺到一张卡上。建议你把ds_config完整贴出来看看，特别是zero_optimization下面有没有漏掉stage3_param_persistence_threshold之类的细节。实在不行先用accelerate加fsdp跑通做个对照，排除是DeepSpeed配置还是模型加载本身的问题。