如何清理Deepspeed引擎占用的GPU内存？解决重复训练内存泄漏问题

我尝试用Llama-7B在不同数据集样本上微调并评估，因为Llama-7B无法在我的GPU上原生训练，所以采用Deepspeed进行微调，代码如下：

def train(tokenizer, model, device, loader, optimizer, model_params):
    print('GPU MEMORY USAGE BEFORE ENGINE CREATED')
    print(get_gpu_memory_usage())
    model.train()
    model, optimizer, _, _ = deepspeed.initialize(model=model, model_parameters=model.parameters(), config=model_params['DEEPSPEED_CONFIG'])
    print('GPU MEMORY USAGE AFTER INITIALIZING ENGINE')
    print(get_gpu_memory_usage())
    outputs = model(...)
    loss = outputs[0]
    print('GPU MEMORY USAGE DURING STEP')
    print(get_gpu_memory_usage())
    # optimizer.zero_grad()
    model.backward(loss)
    # loss.backward()
    model.step()
    # optimizer.step()
    optimizer = None
    torch.cuda.empty_cache()
    print('GPU MEMORY USAGE AFTER CLEARING CACHE')
    print(get_gpu_memory_usage())

    return losses

多次调用train()函数时，首次训练可成功，但第二次就会出现内存不足。通过get_gpu_memory_usage()打印发现，每次调用train()的初始内存（对应GPU MEMORY USAGE BEFORE ENGINE CREATED）持续增长（第一次13MB，第二次7000MB，第三次14000MB）。而使用非Deepspeed版本的代码（注释Deepspeed相关逻辑，启用optimizer.zero_grad()、loss.backward()和optimizer.step()）在小模型上测试时，内存每次调用都保持稳定。推测是Deepspeed初始化的对象未被正确清理出GPU内存，已尝试清空CUDA缓存、删除对象等操作但无效，求彻底释放Deepspeed对象占用内存的方法。

以下是几个可彻底释放Deepspeed占用内存的方法，按优先级尝试：

1. 显式销毁Deepspeed引擎并卸载分布式进程组
   Deepspeed初始化后会创建分布式进程组和引擎对象，仅删除optimizer或清空缓存不足以释放内存，必须显式调用引擎销毁方法并清理进程组：

def train(tokenizer, model, device, loader, optimizer, model_params):
    print('GPU MEMORY USAGE BEFORE ENGINE CREATED')
    print(get_gpu_memory_usage())
    model.train()
    model, optimizer, _, _ = deepspeed.initialize(model=model, model_parameters=model.parameters(), config=model_params['DEEPSPEED_CONFIG'])
    print('GPU MEMORY USAGE AFTER INITIALIZING ENGINE')
    print(get_gpu_memory_usage())
    outputs = model(...)
    loss = outputs[0]
    print('GPU MEMORY USAGE DURING STEP')
    print(get_gpu_memory_usage())
    model.backward(loss)
    model.step()

    # --- 新增内存清理步骤 ---
    # 销毁Deepspeed引擎实例
    if hasattr(model, 'destroy'):
        model.destroy()
    # 卸载分布式进程组
    torch.distributed.destroy_process_group()
    # 将模型移回CPU后删除，避免GPU残留引用
    model = model.module.to('cpu') if hasattr(model, 'module') else model.to('cpu')
    del model
    optimizer = None
    # 强制触发Python垃圾回收
    import gc
    gc.collect()
    torch.cuda.empty_cache()
    # --- 清理步骤结束 ---

    print('GPU MEMORY USAGE AFTER CLEARING CACHE')
    print(get_gpu_memory_usage())

    return losses

2. 每次训练重新初始化模型与优化器
   不要复用同一个model和optimizer对象多次调用train()，每次训练前重新加载模型、创建优化器：

# 调用train的逻辑修改为：
for dataset_sample in dataset_samples:
    # 重新加载预训练模型
    model = AutoModelForCausalLM.from_pretrained("meta-llama/Llama-2-7b-hf")
    model.to(device)
    # 重新创建优化器
    optimizer = AdamW(model.parameters(), lr=5e-5)
    # 执行当前样本的训练
    losses = train(tokenizer, model, device, dataset_sample_loader, optimizer, model_params)

3. 配置Deepspeed ZeRO优化自动释放内存
   如果使用ZeRO优化策略，在Deepspeed配置文件中开启权重自动回收配置，训练后自动释放显存：

{
  "zero_optimization": {
    "stage": 3,
    "stage3_gather_16bit_weights_on_model_save": true,
    "stage3_param_persistence_threshold": 1e4,
    "contiguous_gradients": true
  }
}

4. 排查全局变量引用
   检查代码中是否存在全局变量保存了model、optimizer或Deepspeed引擎的引用，这类引用会阻止Python垃圾回收清理对象，导致显存持续占用。

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