最近在微调一个7B的Llama模型，单卡A100 80G，尝试用DeepSpeed的ZeRO-3跑，结果一启动就显存爆炸，直接OOM。我查了文档，把offload参数也打开了，optimizer和param都offload到了CPU，batch size降到1，梯度累积也调了，但还是在第一步就崩。
我怀疑是不是我的模型加载方式有问题？或者ZeRO-3需要特殊的模型并行配置？看网上有人说ZeRO-2就够用，但我怕显存不够。有没有大佬遇到过类似情况？是不是我漏了什么关键参数？真诚求教，实在不想为了省钱白嫖半天还跑不起来……
楼主
22小时前
用DeepSpeed跑Llama微调，ZeRO-3总是OOM，是我配置姿势不对吗？
请 登录 后发表回复
全部回复
共 1 条
2楼
17小时前
这情况我也踩过坑，7B模型单卡A100 80G跑ZeRO-3确实容易OOM，尤其是一启动就崩。建议你检查下zero_optimization.stage3_gather_16bit_weights_on_model_save是不是设成了false，另外试试把gradient_accumulation_steps提到8以上，同时关掉offload_optimizer只offload param，我这样改完就没崩过了。你用的是transformers的from_pretrained加载吗？可以加个torch_dtype=torch.float16试试。