SD3.5-medium内存占用过高求助：AWS实例运行失败排查

我使用AWS g4dn.xlarge GPU实例，配置为16GB内存+48GB交换空间，搭载16GB显存的Tesla T4 GPU。根据Stability官方说明，该配置应该能运行SD3.5 Medium模型。

从Hugging Face下载模型并启动测试程序后，初始阶段内存+交换空间总计占用约30GB，随后系统内存占比下降，Nvidia GPU内存占用缓慢增长，但在GPU内存分配到约15GB时出现内存分配失败的问题。

我的疑问：

1. 系统与GPU层面的内存占用量是否正常？
2. 我的程序存在什么问题？

附上源代码：

import os
import json
import torch

from diffusers import DiffusionPipeline

pipe = DiffusionPipeline.from_pretrained("./stable-diffusion-3.5-medium/")
if torch.cuda.is_available():
    print('use cuda')
    pipe = pipe.to("cuda")
elif torch.mps.is_available():
    print('use mps')
    pipe = pipe.to('mps')
else:
    print('use cpu')

data = []
with open('data.json', 'r') as f:
    data = json.load(f)

os.makedirs('output', exist_ok=True)
for row in data:
    prompt   = row['prompt']
    filename = 'output/%s.png' % (row['uuid'])
    height   = 1280
    width    = 1280
    
    if row['aspect_ratio'] == '16:9':
        width = 720
    elif row['aspect_ratio'] == '9:16':
        width = 720
        height = 1280
    
    print('saving', filename)
    image = pipe(prompt, height=height, width=width).images[0]
    image.save(filename)

1. 内存占用是否正常？ #

• 系统内存初始占用30GB属于正常：加载SD3.5 Medium模型时，权重会先加载到系统内存，再转移到GPU显存，这个阶段系统内存占用较高，后续权重转移完成后内存占比下降符合预期。
• GPU内存到15GB时失败也符合实际：Tesla T4的16GB显存并非全部可用于模型推理，系统会预留几百MB到1GB左右的显存给驱动、CUDA上下文等开销，实际可用显存约15GB以内，接近上限时自然触发OOM错误。

2. 程序存在的问题及优化方案 #

你的代码未启用任何显存优化策略，导致推理时显存占用超出可用上限，具体问题和修复方式如下：

• 未使用低精度加载模型：默认FP32精度加载的SD3.5 Medium权重接近15GB，再加上推理时的中间张量，直接占满可用显存。
• 未优化注意力与中间张量：推理过程中生成的中间特征张量会额外占用显存，无优化情况下进一步挤压可用空间。
• 宽高设置错误：16:9/9:16的宽高计算不符合比例，错误的分辨率也会增加不必要的显存消耗。

优化后的代码示例：

import os
import json
import torch

from diffusers import DiffusionPipeline

# 启用FP16精度加载，使用safetensors格式加速加载
pipe = DiffusionPipeline.from_pretrained(
    "./stable-diffusion-3.5-medium/",
    torch_dtype=torch.float16,
    use_safetensors=True,
    variant="fp16"
)
# 启用注意力切片、梯度检查点减少显存占用
pipe.enable_attention_slicing()
pipe.enable_gradient_checkpointing()

if torch.cuda.is_available():
    print('use cuda')
    pipe = pipe.to("cuda")
elif torch.mps.is_available():
    print('use mps')
    pipe = pipe.to('mps')
else:
    print('use cpu')

data = []
with open('data.json', 'r') as f:
    data = json.load(f)

os.makedirs('output', exist_ok=True)
for row in data:
    prompt   = row['prompt']
    filename = 'output/%s.png' % (row['uuid'])
    height   = 1280
    width    = 1280
    
    # 修正宽高比例计算
    if row['aspect_ratio'] == '16:9':
        width = 1920
        height = 1080
    elif row['aspect_ratio'] == '9:16':
        width = 1080
        height = 1920
    
    print('saving', filename)
    # 推理时用FP16精度
    image = pipe(
        prompt,
        height=height,
        width=width,
        dtype=torch.float16
    ).images[0]
    image.save(filename)
    # 每次推理后清理显存
    torch.cuda.empty_cache()

额外优化建议：

• 用nvidia-smi命令检查实例是否有其他进程占用显存，关闭不必要的进程释放空间。
• 若仍有显存压力，可启用pipe.enable_model_cpu_offload()，将部分模型层动态卸载到CPU，以少量推理时间为代价进一步降低显存占用。

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