SD 3.5 Medium 在 M1 上：优化至 11.6 GB

苹果M1芯片上Stable Diffusion 3.5中型模型的内存优化指南

拥有25亿参数的Stable Diffusion 3.5中型模型在苹果硅M1芯片上运行时，由于T5-XXL编码器（fp16格式下占10.75 GB）的存在，最高可能消耗高达21 GB的统一内存。通过优化加载方式，峰值内存使用量可降至仅11.6 GB，使配备16 GB内存的M1 Mac也能流畅运行。本文将深入剖析问题根源，并提供分步加载策略。

内存峰值的根本原因：

• 权重重复：模型先加载到CPU，再复制到MPS。
• enable_model_cpu_offload() 效果不佳——CPU与MPS争夺统一内存资源。
• Accelerate的device_map钩子持有强引用，阻止垃圾回收机制释放内存。

标准代码：

pipe = StableDiffusion3Pipeline.from_pretrained(
    "stabilityai/stable-diffusion-3.5-medium",
    torch_dtype=torch.float16, variant="fp16")
pipe.to("mps")

无论num_inference_steps如何设置，峰值内存始终为21.4 GB。

解决方案：分阶段加载法

将流程拆分为三个阶段：提示编码（T5）、卸载、生成（DiT + VAE）。

1. 使用 device_map="balanced" 加载

直接将模型加载至MPS，跳过CPU中间步骤：

pipe = StableDiffusion3Pipeline.from_pretrained(
    path, torch_dtype=torch.float16, device_map="balanced")

T5峰值占用约11 GB。

2. 正确移除钩子

from accelerate.hooks import remove_hook_from_submodules

for attr in list(vars(pipe)):
    comp = getattr(pipe, attr, None)
    if isinstance(comp, torch.nn.Module):
        remove_hook_from_submodules(comp)
        setattr(pipe, attr, None)

del pipe
gc.collect()
torch.mps.synchronize()
torch.mps.empty_cache()

若不执行此操作，T5将长期驻留内存（约10.75 GB）。

3. 完整的分阶段工作流

阶段一：编码器

enc_pipe = StableDiffusion3Pipeline.from_pretrained(path,
    transformer=None, vae=None,
    text_encoder=None, text_encoder_2=None,
    tokenizer=None, tokenizer_2=None,
    variant="fp16", torch_dtype=torch.float16,
    device_map="balanced")

prompt_embeds, neg_embeds, pooled, neg_pooled = enc_pipe.encode_prompt(
    prompt="a warrior with a sword",
    prompt_2="a warrior with a sword",
    prompt_3="a warrior with a sword",
    device="mps", num_images_per_prompt=1)

完成编码后立即卸载。

阶段二：生成器

gen_pipe = StableDiffusion3Pipeline.from_pretrained(path,
    text_encoder=None, text_encoder_2=None, text_encoder_3=None,
    tokenizer=None, tokenizer_2=None, tokenizer_3=None,
    variant="fp16", torch_dtype=torch.float16)
gen_pipe.to("mps")

image = gen_pipe(
    prompt_embeds=prompt_embeds,
    negative_prompt_embeds=neg_embeds,
    pooled_prompt_embeds=pooled,
    negative_pooled_prompt_embeds=neg_pooled,
    num_inference_steps=40).images[0]

.to("mps") 可复用T5缓存，使DiT内存占用从11.5 GB降至6 GB。

内存使用对比表

| 阶段 | 优化前 | 优化后 |

|-------------------|--------|--------|

| T5加载 | 15.3 GB | 11.5 GB |

| T5卸载后 | 10.75 GB | 0.01 GB |

| DiT + 扩散模型 | 8.9 GB | 8.9 GB |

| 峰值内存 | 21.4 GB | 11.6 GB |

硬件性能表现

• M1 16 GB：耗时4分54秒（T5：1分06秒，DiT：3分34秒），峰值内存11–13 GB。
• M3 Pro 36 GB：耗时1分47秒（T5：14秒，DiT：1分27秒），峰值内存9–11 GB。

分辨率：512x512，推理步数：40步。

核心要点总结：

• T5-XXL是主要内存占用源（约占总内存的71%）。
• 结合device_map="balanced"与钩子移除，是实现有效卸载的关键。
• 分阶段加载可使峰值内存降低45%。
• 在M1 16 GB Mac上无需交换内存即可稳定运行。
• 对DiT调用.to("mps")能高效复用缓存。

— Editorial Team