将模型初始化移入函数为何能解决GPU OOM问题？

问题场景 #

在评估大模型多个checkpoint时，初始使用以下循环代码：

for model_name in model_list:
    model = LLM(model_name, trust_remote_code=True)
    results = evaluate_model(model, task)
    del model
    gc.collect()
    torch.cuda.empty_cache()

尽管显式执行了删除模型、调用gc.collect()和torch.cuda.empty_cache()清理CUDA缓存的操作，仍频繁触发GPU内存不足（OOM）错误，错误信息如下：

ValueError: No available memory for the cache blocks. Try increasing `gpu_memory_utilization` or switching to eager mode. You can also reduce the `max_num_seqs` as needed to decrease memory usage.

将模型实例化逻辑移入独立函数后，OOM问题完全解决：

def do_lm_eval(model_name: str, task: str) -> dict:
    gc.collect()
    torch.cuda.empty_cache()
    model = LLM(model_name, trust_remote_code=True)
    results = evaluate_model(model, task)
    del model
    gc.collect()
    torch.cuda.empty_cache()
    return results

for model_name in model_list:
    results = do_lm_eval(model_name, task)

疑问 #

1. 为何移入函数能解决OOM，原代码已做清理操作？
2. Python内存管理或PyTorch与CUDA的交互特性，是否让函数作用域的清理更有效？
3. 迭代大模型时，还有哪些内存管理的最佳实践？

解答 #

1. 函数封装解决OOM的核心原因 #

原循环代码中，即使执行了del model和gc.collect()，仍可能存在隐式的残留引用：比如evaluate_model内部可能持有模型的间接引用（闭包、全局变量缓存、PyTorch的autograd梯度缓存残留），或者循环作用域中，results等变量可能间接关联模型的GPU张量引用，导致GC无法彻底回收这些资源。

而函数作用域是独立的栈帧，函数执行完毕后，栈帧会被销毁，该作用域内的所有局部变量（包括模型实例、中间张量）都会被强制标记为可回收，没有额外的隐式引用干扰，GC能更彻底地清理所有关联内存，进而让torch.cuda.empty_cache()真正释放GPU显存。

2. 函数作用域在内存清理上的特性 #

• Python内存管理层面：函数的局部变量存放在栈帧中，函数执行结束后，栈帧会被弹出，局部变量的引用计数直接归0，无需依赖GC的后续扫描，避免了循环作用域中可能存在的引用残留（比如循环变量复用导致的引用未完全切断）。
• PyTorch与CUDA交互层面：PyTorch的张量对象在CPU侧的Python对象和GPU侧的显存块是一一绑定的，如果CPU侧的Python对象未被彻底回收，GPU显存块会一直被占用。函数作用域结束后，所有局部张量的Python对象被销毁，PyTorch能更及时地释放对应的GPU显存，配合torch.cuda.empty_cache()就能彻底清理显存碎片。

原代码中的gc.collect()可能因为存在未被发现的强引用而无法生效，而函数作用域的销毁是强制的，不存在这类问题。

3. 迭代大模型时的内存管理最佳实践 #

• 强制作用域隔离：延续案例中的思路，把模型加载、推理、清理的逻辑封装到独立函数中，利用函数栈帧的自动销毁特性避免引用残留。
• 禁用不必要的autograd：在评估/推理场景下，使用torch.no_grad()上下文管理器关闭梯度计算，避免存储大量梯度张量占用显存：

  def do_lm_eval(model_name: str, task: str) -> dict:
      gc.collect()
      torch.cuda.empty_cache()
      model = LLM(model_name, trust_remote_code=True)
      with torch.no_grad():
          results = evaluate_model(model, task)
      del model
      gc.collect()
      torch.cuda.empty_cache()
      return results
  

• 模型量化加载：使用load_in_8bit/load_in_4bit量化方式加载模型（比如通过transformers的BitsAndBytesConfig），大幅降低单模型的显存占用。
• 及时清理中间结果：在评估函数内部，主动删除不需要的中间张量；如果使用流水线推理，确保每轮推理后清理流水线缓存。
• 监控显存状态：用torch.cuda.memory_allocated()和torch.cuda.memory_reserved()实时监控显存使用，定位内存泄漏点；也可以通过nvidia-smi命令行工具查看显存占用变化。
• 减少显存碎片：频繁加载卸载大模型时，可在每轮清理后调用torch.cuda.empty_cache()主动整理显存，避免因碎片导致的OOM。

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