使用Optuna时突发CUDA Out of Memory问题的解决方案问询

• 完全隔离每个Trial的资源
  所有与训练相关的对象（模型、优化器、数据加载器等）必须在objective函数内部初始化，禁止复用全局变量或外部定义的对象。trial结束前明确销毁资源：

  def objective(trial):
      # 所有初始化逻辑放在函数内
      model = MyModel(trial.suggest_int("hidden_size", 64, 256)).to(device)
      optimizer = torch.optim.Adam(model.parameters())
      dataloader = get_dataloader(batch_size=32)
  
      # 训练逻辑...
  
      # trial结束后强制清理
      model.cpu()
      del model, optimizer, dataloader
      gc.collect()
      torch.cuda.empty_cache()
      return score
  

• 限制并发Trial数量
  检查study.optimize的n_jobs参数，单GPU场景下务必设置为n_jobs=1，避免多个trial同时占用显存。如果使用剪枝器（Pruner），触发剪枝时立即执行资源清理，不要等到函数末尾：

  def objective(trial):
      # 训练过程中检查剪枝
      for epoch in range(epochs):
          loss = train_one_epoch(model, dataloader)
          trial.report(loss, epoch)
          if trial.should_prune():
              # 剪枝时立即清理
              model.cpu()
              del model, optimizer
              gc.collect()
              torch.cuda.empty_cache()
              raise optuna.TrialPruned()
      return final_score
  

• 排查内存泄漏根源
  在trial前后打印显存使用情况，定位未释放的内存：

  def objective(trial):
      print(f"Initial GPU memory: {torch.cuda.memory_allocated()/1e9:.2f} GB")
      # 训练逻辑...
      print(f"Post-training GPU memory: {torch.cuda.memory_allocated()/1e9:.2f} GB")
      
      # 清理后再次检查
      del model, optimizer
      gc.collect()
      torch.cuda.empty_cache()
      print(f"Post-cleanup GPU memory: {torch.cuda.memory_allocated()/1e9:.2f} GB")
      return score
  

  如果清理后显存仍未回落，说明存在未被正确销毁的全局对象或引用，需要进一步排查代码中的变量引用关系。

• 优化Optuna的存储与日志
  不要在trial.set_user_attr中存储大体积对象（如模型权重、张量数据），仅保留必要的调优参数和指标。若使用SQLite作为存储后端，确保每个trial结束后相关数据库连接被正确关闭，避免内存累积。

• 强化显存清理逻辑
  在清理步骤中加入torch.no_grad()上下文，并重置显存统计：

  def objective(trial):
      # 训练逻辑...
      
      with torch.no_grad():
          model.cpu()
          del model, optimizer
          gc.collect()
          torch.cuda.empty_cache()
          torch.cuda.reset_max_memory_allocated()
      return score
  

  这能确保梯度张量、中间计算张量被彻底释放，避免残留占用显存。

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