Colab运行Disco Diffusion触发CUDA out of memory报错

运行文本生成图像算法Disco Diffusion的Colab Notebook时，执行渲染操作触发运行时错误，报错信息如下：

CUDA out of memory. Tried to allocate 960.00 MiB (GPU 0; 15.78 GiB total capacity; 14.11 GiB already allocated; 158.75 MiB free; 14.14 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

已知复现条件：

• 小尺寸参数渲染可正常运行
• 升级GCE VM配置后仍触发相同报错

1. 修复CUDA显存碎片问题 #

从报错信息看，当前GPU总剩余显存足够，但没有960MiB的连续显存块，是典型的显存碎片问题。

• 在Notebook最顶部、所有引入PyTorch的代码单元之前，新增代码块写入以下配置：

import os
os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "max_split_size_mb:128"

• 写入后必须完全重启Colab运行时再重新执行全流程，不要在已加载过模型的旧会话中直接重试，旧会话的显存碎片不会被自动清理。

2. 核对实际分配的GPU硬件 #

Colab的GPU资源为动态调度，升级GCE VM配置不代表一定能分配到高显存GPU：

• 新增代码块执行!nvidia-smi命令，查看当前实例实际挂载的GPU型号与显存大小。运行大尺寸Disco Diffusion渲染至少需要16GiB及以上显存的GPU（如T4、A100），如果实际分配到的是12GiB及以下显存的型号，重置运行时重新申请资源即可。

3. 调整渲染配置降低显存峰值 #

• 初始渲染分辨率不要超过当前显存支持的上限：16GiB显存的GPU建议初始渲染分辨率最高设置为1280*720，需要更高分辨率输出时，开启内置的分块渲染、后期超分功能，不要直接渲染2K/4K级别的原始图。
• 将批次生成参数batch_size固定为1，关闭多图批量生成。
• 关闭非必要的附加模型：多余的CLIP引导模型、人脸修复模型、二次扩散模型都会额外占用显存，仅保留生成必需的模型即可。
• 开启设置中的梯度检查点（gradient checkpointing）开关，该功能会牺牲10%-20%的渲染速度，降低约30%的显存峰值占用。

4. 规范运行时操作避免无效显存占用 #

• 每次渲染报错后必须重启运行时再重试，运行失败残留的张量数据不会被自动回收，会持续占用显存导致后续任务启动即OOM。
• 同一个运行时实例不要同时运行其他AI模型任务，避免多进程抢占显存。

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