使用Accelerate搭配Deepspeed Zero Stage 2时如何释放GPU内存？

问题根源 #

DeepSpeed Zero Stage 2采用分布式分片存储模型参数与优化器状态，Accelerator封装的底层分布式状态并非仅靠Python层面的del或垃圾回收就能清理——DeepSpeed内部存在大量未被Python引用的显存分配和全局状态，常规方法无法触达这些资源。

可行解决方案 #

1. 调用DeepSpeed专属清理接口 #

在执行常规清理前，先清理DeepSpeed引擎持有的底层资源：

# 先解除模型的分布式包装
model = accelerator.unwrap_model(model)
# 清理模型模块引用
if hasattr(model, 'module'):
    model.module = None
# 清理优化器引用
if 'optimizer' in locals():
    optimizer = None
# 调用DeepSpeed的零优化显存释放接口
accelerator.deepspeed_engine.zero_optimization_partitioned_param_deallocate()
# 销毁DeepSpeed引擎实例
accelerator.deepspeed_engine.destroy()

# 再执行常规清理步骤
del model
del accelerator
gc.collect()
torch.cuda.empty_cache()

2. 拆分阶段为独立进程 #

如果上述方法仍无效，将三个模型的运行逻辑拆分为三个独立脚本，用accelerate launch分别启动。进程结束后系统会自动回收所有显存，彻底避免跨阶段的显存残留：

• run_model1.py：处理第一个模型的训练/推理
• run_model2.py：处理第二个模型的训练/推理
• run_model3.py：处理第三个模型的训练/推理

依次执行命令：

accelerate launch run_model1.py
accelerate launch run_model2.py
accelerate launch run_model3.py

3. 排查隐藏引用 #

检查是否有其他变量持有模型或Accelerator的引用，比如：

• 全局变量中的模型指针
• 自定义回调、钩子函数中保留的模型引用
• Dataloader的collate_fn等逻辑中绑定的模型相关状态

可以通过gc.get_referrers(model)查看所有引用模型的对象，逐个清理后再执行内存释放。

关键注意事项 #

• DeepSpeed Zero Stage 2的显存分配由底层引擎管理，Python垃圾回收无法处理未被Python对象引用的显存块，必须调用专属清理接口。
• 不要在同一个Accelerator实例下反复加载模型，Accelerator初始化时绑定的分布式环境容易残留状态，导致显存无法释放。

内容的提问来源于stack exchange，提问作者Llmw123