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 RTX50 步 DDIM提示词 a photo of an astronaut riding a horse on mars生成 512x512 单张图像设置延迟相对加速无额外设置9.50sx1cuDNN auto-tuner9.37sx1.01fp163.61sx2.63Channels Last 内存格式3.30sx2.88traced UNet3.21sx2.96memory-efficient attention2.63sx3.61从这张表可以读出清晰的优先级半精度权重fp16是性价比最高的单点优化memory-efficient attention 收益最大而 cuDNN autotuner 几乎不改变延迟但属于零成本设置。文档还建议在优先使用 memory-efficient attention 的前提下参考 xFormers 安装指南 完成依赖安装。零成本设置cuDNN autotuner 与 TF32启用 cuDNN auto-tunerNVIDIA cuDNN 针对卷积提供了多种候选算法。Autotuner 会对候选内核跑一次简短基准测试然后为给定输入尺寸和硬件选择最优内核。由于 Stable Diffusion 类模型以卷积网络为主体其他网络类型当前不在该设置的适用范围内在推理前加入以下两行即可启用import torch torch.backends.cudnn.benchmark True注意该设置仅在输入尺寸固定时才稳定发挥收益若每次推理的输入形状都不同autotuner 会反复重新选核反而引入开销。在 Ampere 及之后的 GPU 上用 TF32 替代 fp32从源码结构看这一节的适用前提是硬件为 NVIDIA AmpereRTX 30 系、A 系数据中心卡等及之后的 CUDA 设备。在这些设备上矩阵乘和卷积可以在 TensorFloat32TF32模式下运行速度更快但精度略低。PyTorch 的默认行为是卷积默认开启 TF32而矩阵乘默认不启用。官方文档的建议是除非网络确实需要完整 float32 精度否则应当把矩阵乘也切到 TF32因为精度损失通常可忽略而计算速度提升明显。只需在推理前追加import torch torch.backends.cuda.matmul.allow_tf32 TrueTF32 与 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, dtypetorch.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 均支持也可将dtypetorch.float16换成dtypetorch.bfloat16bfloat16 具有更大的动态范围数值上更稳健这一点在英文版同名文档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, dtypetorch.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_size1时按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_attentionSDPA或 xFormers就不要开启注意力切片——这些后端本身已经很省显存叠加切片反而会造成显著减速。大批量生成的 sliced VAE 解码在受限 VRAM 下解码大批量图像尤其是 32 张及以上的批次时瓶颈往往出现在 VAE 解码阶段。sliced VAE 解码的策略是把批次中的 latent 图像逐个送入解码器而不是一次性全部解码从而把峰值显存压下来。调用方式是推理前对pipe.vae调用enable_slicing()对应AutoencoderKL上的 [~AutoencoderKL.enable_slicing] APIimport torch from diffusers import StableDiffusionPipeline pipe StableDiffusionPipeline.from_pretrained( stable-diffusion-v1-5/stable-diffusion-v1-5, dtypetorch.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, dtypetorch.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, dtypetorch.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 offloadSequential 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, dtypetorch.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_formattorch.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), devicecuda, dtypetorch.float16) timestep torch.rand(1, devicecuda, dtypetorch.float16) * 999 encoder_hidden_states torch.randn((2, 77, 768), devicecuda, dtypetorch.float16) return sample, timestep, encoder_hidden_states pipe StableDiffusionPipeline.from_pretrained( stable-diffusion-v1-5/stable-diffusion-v1-5, dtypetorch.float16, ).to(cuda) unet pipe.unet unet.eval() unet.to(memory_formattorch.channels_last) # 叠加 Channels Last 格式 unet.forward functools.partial(unet.forward, return_dictFalse) # 固定 return_dictFalse # 预热 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(funet 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(funet inference took {time.time() - start_time:.2f} seconds) # 保存追踪产物 unet_traced.save(unet_traced.pt)几个细节值得注意输入形状(2, 4, 64, 64)对应 SD 1.5 的 latent 空间batch2、4 通道、512px 下采样 8 倍后的 64x64(2, 77, 768)对应 CLIP 文本编码77 个 token、768 维functools.partial(unet.forward, return_dictFalse)把return_dictFalse固化为默认行为避免追踪到的图包含 dataclass 分支保存下来的unet_traced.pt可以跨进程加载复用。把 traced UNet 接回管线由于torch.jit.load得到的模块不是 diffusers 的 UNet 类型需要用一个薄包装类适配输出格式管线内部会访问UNet2DConditionOutput.samplefrom 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, dtypetorch.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(samplesample) pipe.unet TracedUNet() with torch.inference_mode(): image pipe([prompt] * 1, num_inference_steps50).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 推理GPU基准注意力 FP16memory-efficient attention FP16NVIDIA Tesla T43.5 it/s5.5 it/sNVIDIA 3060 RTX4.6 it/s7.8 it/sNVIDIA A10G8.88 it/s15.6 it/sNVIDIA RTX A600011.7 it/s21.09 it/sNVIDIA TITAN RTX12.51 it/s18.22 it/sA100-SXM4-40GB18.6 it/s29 it/sA100-SXM-80GB18.7 it/s29.5 it/s各档硬件普遍拿到 1.5 倍左右的提升。使用前提PyTorch 1.12可用 CUDA已安装 xformers 库安装方式见 docs/source/ko/optimization/xformers.mdfrom diffusers import StableDiffusionPipeline import torch pipe StableDiffusionPipeline.from_pretrained( stable-diffusion-v1-5/stable-diffusion-v1-5, dtypetorch.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()否则会带来严重减速。组合策略速查把文档中全部技术按显存收益 / 速度代价维度整理成决策参考技术显存收益速度影响关键调用适用场景cuDNN autotuner无约持平x1.01torch.backends.cudnn.benchmark True固定输入尺寸的卷积网络TF32无显著提升Amperetorch.backends.cuda.matmul.allow_tf32 Truefloat32 精度非必需fp16 权重约减半x2.63from_pretrained(..., dtypetorch.float16)几乎所有 CUDA 场景注意力切片显著可低至 3.2GB约慢 10%pipe.enable_attention_slicing()小显存勿与 SDPA/xFormers 同开sliced VAE解码峰值显著降低多批次略慢单张无感pipe.vae.enable_slicing()大批量≥32 张解码sequential CPU offload最强3GB叠加切片可 2GB显著变慢pipe.enable_sequential_cpu_offload()显存极度受限勿先.to(cuda)model CPU offload中等接近全量上卡pipe.enable_model_cpu_offload()显存中等受限且在意速度需 accelerate ≥0.17.0Channels Last无直接收益提升x2.88 累积unet.to(memory_formattorch.channels_last)需实测算子兼容性UNet tracing无提升x2.96 累积torch.jit.trace(...) 包装类固定输入尺寸、可序列化部署memory-efficient attention最大x3.61 累积pipe.enable_xformers_memory_efficient_attention()PyTorch1.12 CUDA xformers一个典型的小显存组合3.2GB 级显存即可跑 SD 1.5dtypetorch.float16enable_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创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考