十年匠心定制 · 商业建站与技术教学双线并行 咨询热线:400-886-1026 service@lmnt.cn
ARTICLE DETAIL

资讯详情

深耕网站建设与运营推广的一线实战洞察。

Diffusers 中 xFormers 内存高效注意力:安装、启用与原理详解

Diffusers 中 xFormers 内存高效注意力:安装、启用与原理详解 Diffusers 中 xFormers 内存高效注意力安装、启用与原理详解【免费下载链接】diffusers Diffusers: State-of-the-art diffusion models for image, video, and audio generation in PyTorch.项目地址: https://gitcode.com/GitHub_Trending/di/diffusersxFormers 是 Diffusers 官方推荐的注意力优化库通过替换 Transformer/UNet 中的注意力计算内核在推理与训练中同时获得更快的速度和更低的内存占用。本文以 xformers.md 为骨架完整讲解 xFormers 的安装方式、在 Diffusers 中启用它的两条 API 路径并结合 attention_dispatch.py 等源码剖析其底层工作原理与已知坑点帮助你在 Stable Diffusion 等管线中安全落地这一优化。xFormers 为什么值得安装推理与训练双收益扩散模型生成过程是一个多步迭代去噪过程每一步都要执行大量注意力Attention计算而注意力块恰恰是 UNet / Transformer 中计算最密集、显存占用最高的部分。xFormers 的核心价值在于对注意力块内部的矩阵运算做了内核级优化——包括 FlashAttention 式分块tiling与重计算从而推理提速减少注意力计算中的访存次数吞吐更高显存降低避免把完整的注意力权重矩阵物化到显存中可以支撑更大的批量或更高分辨率训练友好官方在文档中明确表示推理和训练都推荐使用 xFormers。这一点在 Diffusers 源码中也有直接印证src/diffusers/models/attention_dispatch.py在模块加载时即检测is_xformers_available()与版本并将 xFormers 注册为与 FlashAttention、SageAttention 并列的一等注意力后端AttentionBackendName.XFORMERS见 attention_dispatch.py。安装 xFormers方式一pip 安装推荐v0.0.16 起自 2023 年 1 月发布的 xFormers0.0.16版本起官方开始提供预编译的 pip wheel安装非常简单pip install xformers安装完成后可用以下命令确认版本并验证是否可导入python -c import xformers; print(xformers.__version__)版本与 PyTorch 的匹配要求[!TIP] xFormers 的 pip 包要求较新版本的 PyTorch例如 xFormers 0.0.16 对应 PyTorch 1.13.1。如果你的环境必须使用更旧的 PyTorch建议按照 xFormers 官方安装指引 从源码自行编译安装。这一点在当前仓库的代码中体现为硬性的最低版本门槛src/diffusers/models/attention_dispatch.py中定义了_REQUIRED_XFORMERS_VERSION 0.0.29见 attention_dispatch.py并且只有满足is_xformers_available() and is_xformers_version(, _REQUIRED_XFORMERS_VERSION)时才真正启用 xFormers 后端attention_dispatch.py。如果包缺失或版本过旧Diffusers 会提示Xformers Attention backend xformers is not usable because of missing package or the version is too old. Please install xformers0.0.29.也就是说当前版本的 Diffusers 建议安装 xFormers 0.0.29而不是文档早期提到的 0.0.16。因此安装前请先检查你的 Diffusers 与 PyTorch 版本尽量使用较新的 xFormers 发布版。环境自检清单PyTorch 版本足够新pip 版 xFormers 对 PyTorch 有版本下限要求CUDA 可用xFormers 的内存高效注意力目前仅支持 GPU源码中会在启用时校验torch.cuda.is_available()见下文xFormers 版本 ≥ 0.0.29针对当前 Diffusers 源码。在 Diffusers 中启用 xFormers两种 API安装完成后在 Diffusers 中有两条启用路径。路径一经典的enable_xformers_memory_efficient_attention()这是韩文版 xformers.md 中提到的传统接口在英文版中该接口对应 fp16 文档的 Memory-efficient attention 一节。它定义在 pipeline_utils.py用法如下import torch from diffusers import DiffusionPipeline from xformers.ops import MemoryEfficientAttentionFlashAttentionOp pipe DiffusionPipeline.from_pretrained( stabilityai/stable-diffusion-2-1, dtypetorch.float16, ).to(cuda) # 启用 xFormers 内存高效注意力可指定 attention_op 覆盖默认算子 pipe.enable_xformers_memory_efficient_attention(attention_opMemoryEfficientAttentionFlashAttentionOp) # Flash Attention 对 VAE 的形状支持有限VAE 用默认 None 算子即可 pipe.vae.enable_xformers_memory_efficient_attention(attention_opNone)对应地通过disable_xformers_memory_efficient_attention()可以关闭该功能pipeline_utils.py。该方法内部会递归遍历管线内的所有子模块把set_use_memory_efficient_attention_xformers(True, attention_op)传递给每个暴露该方法的注意力模块pipeline_utils.py。值得注意的兼容性约束见 attention_processor.py仅限 GPU启用时会校验torch.cuda.is_available()为 False 时直接抛ValueError启用前会先做一次冒烟测试随机张量跑xformers.ops.memory_efficient_attention确保算子可用与Custom Diffusion等特定注意力处理器存在组合限制会抛NotImplementedError与注意力切片attention slicing同时开启时xFormers 优先级更高。路径二现代的set_attention_backend(xformers)英文版 xformers.md 已将该接口迁移到统一的Attention backends注意力调度器机制。当前 Diffusers 通过ModelMixin.set_attention_backend在模型如 UNet、Transformer级别一键切换注意力后端modeling_utils.pyimport torch from diffusers import StableDiffusionXLPipeline pipeline StableDiffusionXLPipeline.from_pretrained( stabilityai/stable-diffusion-xl-base-1.0, dtypetorch.bfloat16, ).to(cuda) # 或 mps、xpu、cpu # 在 transformer/UNet 上启用 xFormers 内存高效注意力 pipeline.unet.set_attention_backend(xformers) prompt Astronaut in a jungle, cold color palette, muted colors, detailed, 8k image pipeline(prompt, num_inference_steps30).images[0]该方法的执行流程源码见 modeling_utils.py将传入的后端名小写化并与AttentionBackendName中的合法值比对非法时抛ValueError若模型启用了上下文并行context parallelism而该后端不支持则报错并列出兼容后端调用_check_attention_backend_requirements校验依赖与版本遍历模型全部子模块为Attention、MochiAttention、AttentionModuleMixin类模块的处理器设置_attention_backend。恢复默认PyTorch 原生 SDPA只需调用pipeline.unet.reset_attention_backend()见 modeling_utils.py。此外还可以用attention_backend上下文管理器做临时切换便于在同一管线内对比不同后端效果参考 attention_backends.md 中的with attention_backend(xformers):用法。两条路径如何选择维度enable_xformers_memory_efficient_attention()set_attention_backend(xformers)层级管线级递归设置所有子模块模型级针对 UNet / Transformer 单独设置时代早期 API仍在 pipeline_utils.py 保留当前推荐的新式 API见 attention_backends.md适用范围兼容旧脚本与社区示例与 FlashAttention、SageAttention 等统一管理恢复方式disable_xformers_memory_efficient_attention()reset_attention_backend()底层原理xFormers 后端在 Diffusers 中如何工作注册与可用性判定在 attention_dispatch.py 中Diffusers 维护了一个注意力后端注册表_AttentionBackendRegistry并在导入阶段通过环境探测决定_CAN_USE_XFORMERS_ATTN_CAN_USE_XFORMERS_ATTN is_xformers_available() and is_xformers_version(, _REQUIRED_XFORMERS_VERSION)只有当该标志为真时才尝试import xformers.ops as xops导入失败如 ABI 不匹配会记录 warning 并优雅回退到原生注意力attention_dispatch.py。is_xformers_available/is_xformers_version的实现位于 import_utils.py本质上是对包是否存在与版本比较的封装。核心算子_xformers_attention注册表将AttentionBackendName.XFORMERS映射到_xformers_attention实现attention_dispatch.py其核心逻辑attention_dispatch.py掩码转换因果注意力转换为xops.LowerTriangularMask()2D 布尔/加性掩码会被转成 xFormers 要求的 4D 加性掩码[batch, heads, seq_q, seq_k]并按 8 的倍数做内存对齐填充超出部分置-inf仅接受 2D 与 4D 掩码其他维度直接抛ValueErrorGQA 支持当enable_gqaTrue时校验 query 头数能被 key/value 头数整除随后 unflatten expand 完成多头分组计算后再 flatten 还原最终调用out xops.memory_efficient_attention(query, key, value, attn_mask, dropout_p, scale)把整个注意力计算交给 xFormers 优化内核不支持return_lse请求返回 log-sum-exp 时会明确报错。相关后端与默认行为在当前的注意力调度器体系里PyTorch ≥ 2.0 时默认后端是原生scaled_dot_product_attentionSDPA它内部本身就封装了 FlashAttention、xFormers 与原生 C 三种实现并自动选择最优参见 fp16.md 的 Scaled dot product attention 一节。因此如果你使用 PyTorch ≥ 2.0默认就已经获得高效注意力不需要额外代码如果你想显式指定 xFormers例如与 SDPA 行为做对比、或需要更可控的算子选择再用set_attention_backend(xformers)。完整后端清单见 attention_backends.md其中xformers后端的定位是 Memory-efficient attention支持多种注意力内核。已知问题与注意事项[!WARNING] 根据 diffusers issue #2234 的反馈xFormersv0.0.16在部分 GPU 上无法用于训练微调 fine-tune 或 DreamBooth 训练。如果你遇到该问题请按照 issue 评论中的指引安装 development开发版版本的 xFormers。除上述训练兼容性问题外还有几点实践提醒保持版本新鲜当前 Diffusers 要求xformers0.0.29务必避开旧版0.0.16带来的训练与兼容性缺陷训练提速不保证enable_xformers_memory_efficient_attention的 docstring 明确指出训练时提速不保证主要收益在显存与推理侧pipeline_utils.py仅限 GPUxFormers 内存高效注意力要求 CUDA 可用CPU/MPS 环境请使用 SDPA 等其他后端VAE 的 Flash 算子限制使用MemoryEfficientAttentionFlashAttentionOp时 VAE 可能因形状限制报错实践中常对 VAE 保留默认attention_opNone见上文示例源自 pipeline_utils.py可与 torch.compile 叠加绝大多数注意力后端含 xFormers支持在无 graph break 的情况下配合torch.compile进一步提速attention_backends.md。结合其他优化手段进一步加速xFormers 只是 Diffusers 推理加速工具箱中的一环官方建议组合使用多项技术以获得叠加收益详见 fp16.md低精度权重将管线以dtypetorch.bfloat16或float16加载并在 NVIDIA Ampere GPU 上开启torch.backends.cuda.matmul.allow_tf32 Truetorch.compile对 UNet/VAE 使用torch.compile(..., modemax-autotune, fullgraphTrue)并配合channels_last内存布局注意力后端调度在 xFormers、FlashAttention、SageAttention 之间用set_attention_backend统一切换、按硬件择优。这些技术与 xFormers 内存高效注意力彼此正交可在保持代码简洁的前提下显著压缩端到端延迟。总结xFormers 是 Diffusers 生态中经过验证的注意力优化方案一条pip install xformers命令即可获得推理加速与显存节省再通过enable_xformers_memory_efficient_attention()经典接口或set_attention_backend(xformers)当前推荐接口一行代码接入管线。其底层由 attention_dispatch.py 中的注册表统一调度包含版本门槛≥ 0.0.29、GPU 校验、掩码对齐与 GQA 展开等细节。动手前只需注意避开 v0.0.16 的训练缺陷、确保 CUDA 环境、保持版本新鲜即可在 Stable Diffusion 等管线中安全享受 xFormers 带来的性能收益。【免费下载链接】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),仅供参考
返回列表