CUDA内存不足报错求助：已尝试常见方案仍未解决

问题背景 #

遇到RuntimeError: CUDA out of memory报错，已尝试过Stack Overflow上的常规显存不足解决方案，包括调整max_split_size_mb、排查显存充足仍报错的情况、通用显存优化指南等，还尝试设置环境变量：

!export 'PYTORCH_CUDA_ALLOC_CONF=max_split_size_mb:4000'

但上述方法均无效。现寻求可行解决方案，并询问是否可以拆分以下Unet推理代码来缓解内存压力：

with autocast(config.DEVICE):
      for i, t in tqdm(enumerate(scheduler.timesteps)):
        
        latent_model_input = torch.cat([latents] * 2)
        sigma = scheduler.sigmas[i]
        latent_model_input = latent_model_input / ((sigma**2 + 1) ** 0.5)

        with torch.no_grad():
              noise_pred = unet(latent_model_input, t, encoder_hidden_states=text_embeddings).sample

        noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
        noise_pred = noise_pred_uncond + config.GUIDANCE_SCALE * (noise_pred_text - noise_pred_uncond)

        latents = scheduler.step(noise_pred, i, latents).prev_sample

报错堆栈信息 #

---------------------------------------------------------------------------
RuntimeError                              Traceback (most recent call last)
<ipython-input-50-c064e16dfce6> in <module>
      7 
      8         with torch.no_grad():
----> 9               noise_pred = unet(latent_model_input, t, encoder_hidden_states=text_embeddings).sample
     10 
     11         noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)

11 frames
/usr/local/lib/python3.7/dist-packages/torch/functional.py in einsum(*args)
    358         return einsum(equation, *_operands)
    359 
--> 360     return _VF.einsum(equation, operands)  # type: ignore[attr-defined]
    361 
    362 

RuntimeError: CUDA out of memory. Tried to allocate 35.60 GiB (GPU 0; 39.59 GiB total capacity; 7.00 GiB already allocated; 30.62 GiB free; 7.24 GiB reserved in total by PyTorch) If reserved memory is >> allocated memory try setting max_split_size_mb to avoid fragmentation.  See documentation for Memory Management and PYTORCH_CUDA_ALLOC_CONF

可行解决方案 #

1. 拆分Unet推理代码（直接缓解显存压力） #

你的代码中torch.cat([latents] * 2)会将输入张量直接翻倍，导致显存占用瞬间增加一倍。可以拆分推理步骤，分别处理无条件和条件输入，避免一次性加载双倍张量：

修改后的代码：

with autocast(config.DEVICE):
      for i, t in tqdm(enumerate(scheduler.timesteps)):
        
        sigma = scheduler.sigmas[i]
        # 先处理无条件文本嵌入推理
        latent_uncond = latents / ((sigma**2 + 1) ** 0.5)
        with torch.no_grad():
              noise_pred_uncond = unet(latent_uncond, t, encoder_hidden_states=text_embeddings_uncond).sample
        
        # 再处理目标文本嵌入推理
        latent_text = latents / ((sigma**2 + 1) ** 0.5)
        with torch.no_grad():
              noise_pred_text = unet(latent_text, t, encoder_hidden_states=text_embeddings).sample
        
        # 合并推理结果
        noise_pred = noise_pred_uncond + config.GUIDANCE_SCALE * (noise_pred_text - noise_pred_uncond)

        latents = scheduler.step(noise_pred, i, latents).prev_sample
        # 手动释放临时张量显存
        del latent_uncond, latent_text, noise_pred_uncond, noise_pred_text
        torch.cuda.empty_cache()

注意：需要提前准备好单独的text_embeddings_uncond（即空文本对应的嵌入张量），如果之前是通过cat合并到同一张量中的，现在要单独生成并保存。

2. 其他显存优化手段 #

• 降低输入分辨率：如果latent张量尺寸过大（比如对应1024x1024的原始图像），可以降到512x512，显存占用会直接降到原来的1/4。
• 启用梯度检查点：给Unet启用梯度检查点，牺牲少量速度换取显存节省：

  unet.enable_gradient_checkpointing()
  

• 调整max_split_size_mb为更小值：之前设置的4000可能过大，尝试设为128或256，更细粒度的内存拆分能减少碎片：

  !export 'PYTORCH_CUDA_ALLOC_CONF=max_split_size_mb:128'
  

• 强化混合精度：确保autocast正常生效，若GPU支持（如Ampere及以上架构），可尝试FP8精度进一步降低显存占用。
• 定期清空缓存：在循环内定期调用torch.cuda.empty_cache()，及时释放临时张量占用的显存。

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