企业官网建设流程全解析

Diffusers 推理优化实战：fp16、注意力切片与 CPU 卸载的完整技术指南

【免费下载链接】diffusers🤗 Diffusers: State-of-the-art diffusion models for image, video, and audio generation in PyTorch.项目地址: https://gitcode.com/GitHub_Trending/di/diffusers

本文围绕 Diffusers 官方文档《메모리와 속도》（docs/source/ko/optimization/fp16.md）中讲解的推理优化技术展开，覆盖 cuDNN autotuner、TF32、float16 半精度权重、注意力切片、sliced VAE 解码、CPU 卸载（sequential / model offload）、Channels Last 内存格式、UNet 追踪（tracing）以及 memory-efficient attention 共 9 类技术。读完本文后，你将掌握如何为受限 VRAM 环境挑选合适的组合策略，并能对照源码理解每个优化开关（如enable_attention_slicing、enable_sequential_cpu_offload）在底层的具体实现机制。

为什么扩散模型推理需要优化，以及各项技术的收益总览

扩散模型的生成是一个迭代过程：噪声在若干步（steps）内被逐步细化为图像，因此推理天然是计算密集型的。官方文档给出了这些优化设置在单卡上的实测收益（测试条件：NVIDIA TITAN RTX，50 步 DDIM，提示词 "a photo of an astronaut riding a horse on mars"，生成 512x512 单张图像）：

从这张表可以读出清晰的优先级：半精度权重（fp16）是性价比最高的单点优化，memory-efficient attention 收益最大，而 cuDNN autotuner 几乎不改变延迟但属于零成本设置。文档还建议在优先使用 memory-efficient attention 的前提下，参考 xFormers 安装指南 完成依赖安装。

零成本设置：cuDNN autotuner 与 TF32

启用 cuDNN auto-tuner

NVIDIA cuDNN 针对卷积提供了多种候选算法。Autotuner 会对候选内核跑一次简短基准测试，然后为给定输入尺寸和硬件选择最优内核。由于 Stable Diffusion 类模型以卷积网络为主体（其他网络类型当前不在该设置的适用范围内），在推理前加入以下两行即可启用：

import torch torch.backends.cudnn.benchmark = True

注意该设置仅在输入尺寸固定时才稳定发挥收益；若每次推理的输入形状都不同，autotuner 会反复重新选核，反而引入开销。

在 Ampere 及之后的 GPU 上用 TF32 替代 fp32

从源码结构看，这一节的适用前提是硬件为 NVIDIA Ampere（RTX 30 系、A 系数据中心卡等）及之后的 CUDA 设备。在这些设备上，矩阵乘和卷积可以在 TensorFloat32（TF32）模式下运行，速度更快但精度略低。PyTorch 的默认行为是：卷积默认开启 TF32，而矩阵乘默认不启用。官方文档的建议是：除非网络确实需要完整 float32 精度，否则应当把矩阵乘也切到 TF32，因为精度损失通常可忽略，而计算速度提升明显。只需在推理前追加：

import torch torch.backends.cuda.matmul.allow_tf32 = True

TF32 与 fp16/bfloat16 正交：TF32 加速的是 float32 权重下的运算路径，二者可叠加使用。

半精度权重：用 fp16 换取约 2.6 倍加速

从总览表中可以看到，fp16 是单点收益最大的设置（9.50s → 3.61s，约 2.63 倍）。其原理很直接：以 float16 加载并运行权重，既能把显存占用近似减半，又能利用 GPU 上吞吐更高的半精度算力。在from_pretrained时通过dtype参数（早期版本写作torch_dtype）直接指定即可：

from diffusers import StableDiffusionPipeline import torch pipe = StableDiffusionPipeline.from_pretrained( "stable-diffusion-v1-5/stable-diffusion-v1-5", dtype=torch.float16, ) pipe = pipe.to("cuda") prompt = "a photo of an astronaut riding a horse on mars" image = pipe(prompt).images[0]

文档中还有一条重要警告：不要在任意管线中使用torch.autocast—— 混用 autocast 可能导致生成黑色图像，而且始终比纯 float16 更慢。如果你的目标硬件支持 bfloat16（大多数现代 NVIDIA GPU 均支持），也可将dtype=torch.float16换成dtype=torch.bfloat16；bfloat16 具有更大的动态范围，数值上更稳健，这一点在英文版同名文档（docs/source/en/optimization/fp16.md）的 "Model data type" 一节中也有对应说明。

注意力切片：用约 10% 的耗时换下 3.2GB 显存

当显存不够时，第一件可以做的事是启用切片注意力（sliced attention）：不再一次性计算完整的注意力，而是按切片分步计算。

• 对 batch size 为 1 的场景同样有效：只要模型使用一个以上的注意力头，QK^T 注意力矩阵就可以按头逐个顺序计算，从而节省可观显存；
• 代价是推理时间约慢 10%；
• 收益是约 3.2GB 的小显存设备即可运行 Stable Diffusion。

import torch from diffusers import StableDiffusionPipeline pipe = StableDiffusionPipeline.from_pretrained( "stable-diffusion-v1-5/stable-diffusion-v1-5", dtype=torch.float16, ) pipe = pipe.to("cuda") prompt = "a photo of an astronaut riding a horse on mars" pipe.enable_attention_slicing() image = pipe(prompt).images[0]

源码视角：slice_size参数的三种取值

在实现层面，enable_attention_slicing是定义在DiffusionPipeline基类上的方法，接受slice_size: str | int = "auto"参数，文档字符串明确了三种取值语义：

• "auto"（默认）：把注意力头的输入减半，分两步计算注意力；
• "max"：一次只跑一个切片，节省最多显存（对应文中enable_attention_slicing(1)的写法——slice_size=1时按attention_head_dim // slice_size划分切片数，即每个切片只处理 1 个维度单位，要求attention_head_dim能被slice_size整除）；
• 正整数：使用attention_head_dim // slice_size个切片。

该方法内部调用set_attention_slice，它会遍历管线签名中的所有模块，把切片配置下发给每个实现了set_attention_slice的子模块（UNet 中的各注意力层）。基类源码中还有一个值得注意的警告：如果已经在用 PyTorch 2.0+ 的scaled_dot_product_attention（SDPA）或 xFormers，就不要开启注意力切片——这些后端本身已经很省显存，叠加切片反而会造成显著减速。

大批量生成的 sliced VAE 解码

在受限 VRAM 下解码大批量图像（尤其是 32 张及以上的批次）时，瓶颈往往出现在 VAE 解码阶段。sliced VAE 解码的策略是：把批次中的 latent 图像逐个送入解码器，而不是一次性全部解码，从而把峰值显存压下来。

调用方式是推理前对pipe.vae调用enable_slicing()（对应AutoencoderKL上的 [~AutoencoderKL.enable_slicing] API）：

import torch from diffusers import StableDiffusionPipeline pipe = StableDiffusionPipeline.from_pretrained( "stable-diffusion-v1-5/stable-diffusion-v1-5", dtype=torch.float16, ) pipe = pipe.to("cuda") prompt = "a photo of an astronaut riding a horse on mars" pipe.vae.enable_slicing() images = pipe([prompt] * 32).images

性能特征：多图像批次下 VAE 解码会有轻微性能损耗；单图像批次则完全没有性能影响，可以无条件开启。它也可以与注意力切片或 xFormers memory-efficient attention 叠加，进一步压低内存使用。

最省显存的 sequential CPU offload：把权重搬到 CPU 上换

enable_sequential_cpu_offload()的思路是：权重平时全部留在 CPU，某个子模块执行前向时才临时搬到 GPU，用完再搬走。

import torch from diffusers import StableDiffusionPipeline pipe = StableDiffusionPipeline.from_pretrained( "stable-diffusion-v1-5/stable-diffusion-v1-5", dtype=torch.float16, ) prompt = "a photo of an astronaut riding a horse on mars" pipe.enable_sequential_cpu_offload() image = pipe(prompt).images[0]

这样可以把内存消耗压到 3GB 以下。但必须理解它的代价来源：该机制工作在子模块（submodule）粒度而非整模型粒度，而管线中的 UNet 会被执行num_inference_steps次——每一步都会触发不同子模块的"上载 → 计算 → 卸载"循环，频繁的 CPU↔GPU 内存搬运使推理速度显著变慢。

从源码看，enable_sequential_cpu_offload的 docstring 准确描述了这一机制：所有torch.nn.Module组件（_exclude_from_cpu_offload中的组件除外）的 state dict 被保存到 CPU，模块本体移到torch.device('meta')，只有当具体子模块的forward被调用时才加载到加速器；它依赖accelerate >= 0.14.0的cpu_offloadAPI。

组合到 2GB 以内：offload + attention slicing

把 sequential offload 与注意力切片叠加，可以把内存需求再压低到2GB 以下：

import torch from diffusers import StableDiffusionPipeline pipe = StableDiffusionPipeline.from_pretrained( "stable-diffusion-v1-5/stable-diffusion-v1-5", dtype=torch.float16, ) prompt = "a photo of an astronaut riding a horse on mars" pipe.enable_sequential_cpu_offload() pipe.enable_attention_slicing(1) image = pipe(prompt).images[0]

一个关键禁忌：使用enable_sequential_cpu_offload()时，不要提前把管线.to("cuda")——先整管线上 GPU 再 offload，会让显存节省的收益几乎全部失效。

速度与显存的折中：model offload

Sequential offload 省显存最多但最慢。模型卸载（model offload）是另一个方向的折中：它不按子模块搬运，而是以整个模型为单位在 CPU 与 GPU 之间切换。任一时刻 GPU 上只驻留管线的主要组件之一（通常是 text encoder、unet、vae 中的一个），其余在 CPU 上等待；像 UNet 这种需要多步迭代执行的组件，会在不再需要之前一直留在 GPU 上，避免反复搬运。

因此它对推理时延的影响远小于 sequential offload（与直接把管线放到cuda相比只多一点点开销），同时仍省下若干 GB 显存：

import torch from diffusers import StableDiffusionPipeline pipe = StableDiffusionPipeline.from_pretrained( "stable-diffusion-v1-5/stable-diffusion-v1-5", dtype=torch.float16, ) prompt = "a photo of an astronaut riding a horse on mars" pipe.enable_model_cpu_offload() image = pipe(prompt).images[0]

同样可与注意力切片叠加：

pipe.enable_model_cpu_offload() pipe.enable_attention_slicing(1) image = pipe(prompt).images[0]

环境要求：该功能依赖accelerate0.17.0 及以上版本。

源码视角：offload 的实现依赖model_cpu_offload_seq

查看enable_model_cpu_offload的实现可以发现其工作细节：

• 它要求管线类上定义了model_cpu_offload_seq类属性（否则抛出 ValueError），按该序列字符串（以->分隔）依次对每个组件挂载cpu_offload_with_hook；不在链上的组件会被 offload 到 CPU，直到maybe_free_model_hooks被调用。这与"GPU 上同时只有一个主要组件"的行为一致；
• 它校验了版本下限：is_accelerate_version(">=", "0.17.0.dev0")，低于该版本直接ImportError，与文档中"需要 accelerate 0.17.0+"的说明完全吻合；
• 若管线已经启用了 device mapping 策略，调用会报错，需要先reset_device_map()。

两种 offload 的选择建议可以归纳为：显存极度紧张（消费级小显存卡、边缘设备）→ sequential offload；显存中等紧张且在意生成速度 → model offload。

Channels Last 内存格式：重排内存布局而非减少计算

Channels Last 是保持维度顺序不变、但改变内存中 NCHW 张量排布的方式：通道维成为最内层（最密集）的维度，即俗称的"按每像素存储图像"。并非所有算子都支持该格式，个别算子甚至可能因此变慢，所以建议启用后实测确认。

对管线中的 UNet 启用 Channels Last 只需一行 in-place 操作，可以用卷积权重张量的 stride 变化来验证布局确实改变了：

print(pipe.unet.conv_out.state_dict()["weight"].stride()) # (2880, 9, 3, 1) pipe.unet.to(memory_format=torch.channels_last) # in-place 操作 # 变成 (2880, 1, 960, 320)，通道维（第 2 维）stride 为 1，证明布局已改变 print(pipe.unet.conv_out.state_dict()["weight"].stride())

在总览表中，Channels Last 使延迟从 3.61s 降到 3.30s；它与 fp16、tracing、torch.compile 都是正交且可叠加的优化。

UNet 追踪（tracing）：用 torch.jit 固化执行图

追踪（tracing）指让一个示例输入张量穿过模型，捕获各层被调用的操作序列，产出一个可被 just-in-time 优化的执行图。对 Stable Diffusion 的 UNet 做 tracing 的完整流程分为三步：准备输入并做预热、torch.jit.trace固图、基准测试与保存。

import time import torch from diffusers import StableDiffusionPipeline import functools # 关闭梯度 torch.set_grad_enabled(False) n_experiments = 2 unet_runs_per_experiment = 50 def generate_inputs(): sample = torch.randn((2, 4, 64, 64), device="cuda", dtype=torch.float16) timestep = torch.rand(1, device="cuda", dtype=torch.float16) * 999 encoder_hidden_states = torch.randn((2, 77, 768), device="cuda", dtype=torch.float16) return sample, timestep, encoder_hidden_states pipe = StableDiffusionPipeline.from_pretrained( "stable-diffusion-v1-5/stable-diffusion-v1-5", dtype=torch.float16, ).to("cuda") unet = pipe.unet unet.eval() unet.to(memory_format=torch.channels_last) # 叠加 Channels Last 格式 unet.forward = functools.partial(unet.forward, return_dict=False) # 固定 return_dict=False # 预热 for _ in range(3): with torch.inference_mode(): inputs = generate_inputs() orig_output = unet(*inputs) # 追踪 print("tracing..") unet_traced = torch.jit.trace(unet, inputs) unet_traced.eval() print("done tracing") # 预热与图优化 for _ in range(5): with torch.inference_mode(): inputs = generate_inputs() orig_output = unet_traced(*inputs) # 基准测试：对比 traced 与 eager 两种执行的耗时 with torch.inference_mode(): for _ in range(n_experiments): torch.cuda.synchronize() start_time = time.time() for _ in range(unet_runs_per_experiment): orig_output = unet_traced(*inputs) torch.cuda.synchronize() print(f"unet traced inference took {time.time() - start_time:.2f} seconds") for _ in range(n_experiments): torch.cuda.synchronize() start_time = time.time() for _ in range(unet_runs_per_experiment): orig_output = unet(*inputs) torch.cuda.synchronize() print(f"unet inference took {time.time() - start_time:.2f} seconds") # 保存追踪产物 unet_traced.save("unet_traced.pt")

几个细节值得注意：

• 输入形状(2, 4, 64, 64)对应 SD 1.5 的 latent 空间（batch=2、4 通道、512px 下采样 8 倍后的 64x64）；(2, 77, 768)对应 CLIP 文本编码（77 个 token、768 维）；
• functools.partial(unet.forward, return_dict=False)把return_dict=False固化为默认行为，避免追踪到的图包含 dataclass 分支；
• 保存下来的unet_traced.pt可以跨进程加载复用。

把 traced UNet 接回管线

由于torch.jit.load得到的模块不是 diffusers 的 UNet 类型，需要用一个薄包装类适配输出格式（管线内部会访问UNet2DConditionOutput.sample）：

from diffusers import StableDiffusionPipeline import torch from dataclasses import dataclass @dataclass class UNet2DConditionOutput: sample: torch.Tensor pipe = StableDiffusionPipeline.from_pretrained( "stable-diffusion-v1-5/stable-diffusion-v1-5", dtype=torch.float16, ).to("cuda") # 加载 traced unet unet_traced = torch.jit.load("unet_traced.pt") class TracedUNet(torch.nn.Module): def __init__(self): super().__init__() self.in_channels = pipe.unet.config.in_channels self.device = pipe.unet.device def forward(self, latent_model_input, t, encoder_hidden_states): sample = unet_traced(latent_model_input, t, encoder_hidden_states)[0] return UNet2DConditionOutput(sample=sample) pipe.unet = TracedUNet() with torch.inference_mode(): image = pipe([prompt] * 1, num_inference_steps=50).images[0]

在总览表中 traced UNet 使延迟从 3.30s 降到 3.21s——它是在已启用 fp16 + Channels Last 之后的进一步压缩。现代替代方案（torch.compile）在英文版文档中已作为更推荐的加速手段展开介绍，可参阅 docs/source/en/optimization/fp16.md 的 "torch.compile" 一节；tracing 的价值在于产物可序列化、可跨环境复用。

Memory-efficient attention：最大的单点加速

注意力计算是 UNet 中的带宽热点。FlashAttention 等工作通过优化注意力块的重算/分块策略，大幅降低了 GPU 显存占用并提升了吞吐。官方文档给出的实测加速（batch size 1、512x512 推理）：

各档硬件普遍拿到 1.5 倍左右的提升。使用前提：

• PyTorch > 1.12
• 可用 CUDA
• 已安装 xformers 库（安装方式见 docs/source/ko/optimization/xformers.md）

from diffusers import StableDiffusionPipeline import torch pipe = StableDiffusionPipeline.from_pretrained( "stable-diffusion-v1-5/stable-diffusion-v1-5", dtype=torch.float16, ).to("cuda") pipe.enable_xformers_memory_efficient_attention() with torch.inference_mode(): sample = pipe("a small cat") # 可选：用下面这行来关闭 # pipe.disable_xformers_memory_efficient_attention()

源码视角：递归下发 xFormers 注意力后端

在实现层面，enable_xformers_memory_efficient_attention会递归遍历管线的所有组件，凡实现了set_use_memory_efficient_attention_xformers方法的模块（即各 Transformer/UNet 注意力层）都会被切换；模型侧的入口在 src/diffusers/models/modeling_utils.py，管线侧的分发逻辑位于 src/diffusers/pipelines/pipeline_utils.py 的fn_recursive_set_mem_eff递归函数中——它把开关沿模块树下发到叶子层，这正是"一行调用、全模型生效"的原因。

结合基类 docstring 中的警告（与前述注意力切片一节呼应）：当使用 SDPA/xFormers 这类高效注意力后端时，不要再叠加enable_attention_slicing()，否则会带来严重减速。

组合策略速查

把文档中全部技术按"显存收益 / 速度代价"维度整理成决策参考：

一个典型的小显存组合（3.2GB 级显存即可跑 SD 1.5）：dtype=torch.float16+enable_attention_slicing()+enable_xformers_memory_efficient_attention()；而显存只有 1–2GB 的极限场景则应切换到enable_sequential_cpu_offload()+enable_attention_slicing(1)，以可接受的时延换下运行可行性。

【免费下载链接】diffusers🤗 Diffusers: State-of-the-art diffusion models for image, video, and audio generation in PyTorch.项目地址: https://gitcode.com/GitHub_Trending/di/diffusers