减少内存使用

像 Flux 和 Wan 这样的现代扩散模型拥有数十亿参数，在硬件上进行推理时会占用大量内存。由于常见的 GPU 往往内存不足，这带来了挑战。为了克服内存限制，您可以使用多块 GPU（如果可用）、将部分流水线组件卸载到 CPU 上，以及采用其他方法。

本指南将向你展示如何降低内存使用量。

请注意，这些技术可能需要根据模型进行调整。例如，基于 Transformer 的扩散模型从这些内存优化中获得的收益，可能不如基于 UNet 的模型。

多 GPU

如果您拥有多块 GPU，则有几种选项可以高效地将大型模型加载并分布到您的硬件上。这些功能由 Accelerate 库提供支持，因此请确保先安装该库。

pip install -U accelerate

分片检查点

分片加载大型检查点很有用，因为分片是逐个加载的。这样可以保持内存占用较低，只需满足模型大小和最大分片大小所需的内存即可。我们建议当 fp32 检查点大于 5GB 时进行分片。默认的分片大小为 5GB。

使用 max_shard_size 参数在 save_pretrained() 中对检查点进行分片。

from diffusers import AutoModel

unet = AutoModel.from_pretrained(
    "stabilityai/stable-diffusion-xl-base-1.0", subfolder="unet"
)
unet.save_pretrained("sdxl-unet-sharded", max_shard_size="5GB")

现在，您可以使用分片检查点（sharded checkpoint）来替代常规检查点，以节省内存。

import torch
from diffusers import AutoModel, StableDiffusionXLPipeline

unet = AutoModel.from_pretrained(
    "username/sdxl-unet-sharded", dtype=torch.float16
)
pipeline = StableDiffusionXLPipeline.from_pretrained(
    "stabilityai/stable-diffusion-xl-base-1.0",
    unet=unet,
    dtype=torch.float16
).to("cuda")  # or "mps", "xpu", "cpu"

设备放置

设备放置是一项实验性功能，API 可能会发生变化。目前仅支持 balanced 策略。我们计划在未来支持更多的映射策略。

device_map 参数用于控制流水线中的模型组件或单个模型中的层如何在不同设备上进行分配。

流水线级别

模型级别

在设计您自己的 device_map 时，它应该是一个字典，其中包含模型特定的模块名称或层以及设备标识符（GPU 使用整数，CPU 使用 cpu，磁盘使用 disk）。

对模型调用 hf_device_map 以查看模型层的分布情况，然后自行设计。

print(transformer.hf_device_map)
{'pos_embed': 0, 'time_text_embed': 0, 'context_embedder': 0, 'x_embedder': 0, 'transformer_blocks': 0, 'single_transformer_blocks.0': 0, 'single_transformer_blocks.1': 0, 'single_transformer_blocks.2': 0, 'single_transformer_blocks.3': 0, 'single_transformer_blocks.4': 0, 'single_transformer_blocks.5': 0, 'single_transformer_blocks.6': 0, 'single_transformer_blocks.7': 0, 'single_transformer_blocks.8': 0, 'single_transformer_blocks.9': 0, 'single_transformer_blocks.10': 'cpu', 'single_transformer_blocks.11': 'cpu', 'single_transformer_blocks.12': 'cpu', 'single_transformer_blocks.13': 'cpu', 'single_transformer_blocks.14': 'cpu', 'single_transformer_blocks.15': 'cpu', 'single_transformer_blocks.16': 'cpu', 'single_transformer_blocks.17': 'cpu', 'single_transformer_blocks.18': 'cpu', 'single_transformer_blocks.19': 'cpu', 'single_transformer_blocks.20': 'cpu', 'single_transformer_blocks.21': 'cpu', 'single_transformer_blocks.22': 'cpu', 'single_transformer_blocks.23': 'cpu', 'single_transformer_blocks.24': 'cpu', 'single_transformer_blocks.25': 'cpu', 'single_transformer_blocks.26': 'cpu', 'single_transformer_blocks.27': 'cpu', 'single_transformer_blocks.28': 'cpu', 'single_transformer_blocks.29': 'cpu', 'single_transformer_blocks.30': 'cpu', 'single_transformer_blocks.31': 'cpu', 'single_transformer_blocks.32': 'cpu', 'single_transformer_blocks.33': 'cpu', 'single_transformer_blocks.34': 'cpu', 'single_transformer_blocks.35': 'cpu', 'single_transformer_blocks.36': 'cpu', 'single_transformer_blocks.37': 'cpu', 'norm_out': 'cpu', 'proj_out': 'cpu'}

例如，下面的 device_map 将 single_transformer_blocks.10 到 single_transformer_blocks.20 放置在第二块 GPU（1）上。

import torch
from diffusers import AutoModel

device_map = {
    'pos_embed': 0, 'time_text_embed': 0, 'context_embedder': 0, 'x_embedder': 0, 'transformer_blocks': 0, 'single_transformer_blocks.0': 0, 'single_transformer_blocks.1': 0, 'single_transformer_blocks.2': 0, 'single_transformer_blocks.3': 0, 'single_transformer_blocks.4': 0, 'single_transformer_blocks.5': 0, 'single_transformer_blocks.6': 0, 'single_transformer_blocks.7': 0, 'single_transformer_blocks.8': 0, 'single_transformer_blocks.9': 0, 'single_transformer_blocks.10': 1, 'single_transformer_blocks.11': 1, 'single_transformer_blocks.12': 1, 'single_transformer_blocks.13': 1, 'single_transformer_blocks.14': 1, 'single_transformer_blocks.15': 1, 'single_transformer_blocks.16': 1, 'single_transformer_blocks.17': 1, 'single_transformer_blocks.18': 1, 'single_transformer_blocks.19': 1, 'single_transformer_blocks.20': 1, 'single_transformer_blocks.21': 'cpu', 'single_transformer_blocks.22': 'cpu', 'single_transformer_blocks.23': 'cpu', 'single_transformer_blocks.24': 'cpu', 'single_transformer_blocks.25': 'cpu', 'single_transformer_blocks.26': 'cpu', 'single_transformer_blocks.27': 'cpu', 'single_transformer_blocks.28': 'cpu', 'single_transformer_blocks.29': 'cpu', 'single_transformer_blocks.30': 'cpu', 'single_transformer_blocks.31': 'cpu', 'single_transformer_blocks.32': 'cpu', 'single_transformer_blocks.33': 'cpu', 'single_transformer_blocks.34': 'cpu', 'single_transformer_blocks.35': 'cpu', 'single_transformer_blocks.36': 'cpu', 'single_transformer_blocks.37': 'cpu', 'norm_out': 'cpu', 'proj_out': 'cpu'
}

transformer = AutoModel.from_pretrained(
    "black-forest-labs/FLUX.1-dev", 
    subfolder="transformer",
    device_map=device_map,
    dtype=torch.bfloat16
)

传入一个将最大内存使用量映射到每个设备的字典以强制实施限制。如果某个设备不在 max_memory 中，则会被忽略，且流水线组件不会分发到该设备上。

import torch
from diffusers import AutoModel, StableDiffusionXLPipeline

max_memory = {0:"1GB", 1:"1GB"}
pipeline = StableDiffusionXLPipeline.from_pretrained(
    "stabilityai/stable-diffusion-xl-base-1.0",
    dtype=torch.float16,
    device_map="balanced",
    max_memory=max_memory
)

Diffusers 默认使用所有设备的最大内存，但如果模型无法装入 GPU，则需要使用单个 GPU，并通过以下方法将部分计算卸载到 CPU。

enable_model_cpu_offload() 仅适用于单 GPU，但非常大的模型可能无法装入单个 GPU
enable_sequential_cpu_offload() 可能可以运行，但速度极慢，且仅限于单 GPU

使用 reset_device_map() 方法重置 device_map。如果您想在已进行设备映射的 pipeline 上使用诸如 .to()、enable_sequential_cpu_offload() 和 enable_model_cpu_offload() 之类的方法，则必须执行此操作。

pipeline.reset_device_map()

VAE 切片

VAE 切片通过将一个大批次输入拆分为单个数据批次并分别处理来节省内存。此方法在同时生成多张图像时效果最佳。

例如，如果您一次性生成 4 张图像，解码过程会使峰值激活内存增加 4 倍。VAE 切片通过每次仅解码 1 张图像（而非同时解码全部 4 张）来降低这一开销。

调用 enable_slicing() 以启用切片式 VAE。在解码多图像批次时，性能会有小幅提升；对于单图像批次，则不会有性能影响。

import torch
from diffusers import AutoModel, StableDiffusionXLPipeline

pipeline = StableDiffusionXLPipeline.from_pretrained(
    "stabilityai/stable-diffusion-xl-base-1.0",
    dtype=torch.float16,
).to("cuda")  # or "mps", "xpu", "cpu"
pipeline.vae.enable_slicing()
pipeline(["An astronaut riding a horse on Mars"]*32).images[0]
print(f"Max memory reserved: {torch.cuda.max_memory_allocated() / 1024**3:.2f} GB")

AutoencoderKLWan 和 AsymmetricAutoencoderKL 类不支持切片。

VAE 分块

VAE 分块（tiling）通过将图像划分为较小的重叠块，而非一次性处理整张图像，从而节省内存。由于 GPU 每次仅处理一个块，这也有助于降低峰值内存占用。

调用 enable_tiling() 以启用 VAE 分块。由于各分块是独立解码的，生成的图像在分块之间可能会出现一些色调差异，但分块之间不应出现明显的接缝。对于低于预设（但可配置）阈值的分辨率，分块功能将被禁用。例如，StableDiffusionPipeline 中 VAE 的该阈值为 512x512。

import torch
from diffusers import AutoPipelineForImage2Image
from diffusers.utils import load_image

pipeline = AutoPipelineForImage2Image.from_pretrained(
    "stabilityai/stable-diffusion-xl-base-1.0", dtype=torch.float16
).to("cuda")  # or "mps", "xpu", "cpu"
pipeline.vae.enable_tiling()

init_image = load_image("https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/img2img-sdxl-init.png")
prompt = "Astronaut in a jungle, cold color palette, muted colors, detailed, 8k"
pipeline(prompt, image=init_image, strength=0.5).images[0]
print(f"Max memory reserved: {torch.cuda.max_memory_allocated() / 1024**3:.2f} GB")

AutoencoderKLWan 和 AsymmetricAutoencoderKL 不支持分块处理。

卸载

Offloading 策略会将当前未激活的层或模型移至 CPU，以避免增加 GPU 内存占用。这些策略可以与量化和 torch.compile 结合使用，以平衡推理速度和内存使用。

有关更多详细信息，请参阅编译和卸载量化模型指南。

CPU 卸载

CPU offloading 会将权重从 GPU 有选择地移至 CPU。当某个组件被需要时，它会被传输到 GPU；当不再需要时，它会被移回 CPU。该方法针对子模块而非整个模型进行操作。通过避免在 GPU 上存储整个模型，此方法可以节省内存。

CPU offloading 能大幅降低内存占用，但由于子模块需要在设备之间多次来回传递，其速度极其缓慢。因此，在实际应用中往往难以采用。

在调用 enable_sequential_cpu_offload() 之前，请勿将 pipeline 移至 CUDA，否则节省的内存量微乎其微（更多详情请参见此 issue）。这是一个有状态的操作，会在模型上安装钩子。

调用 enable_sequential_cpu_offload() 以在管道上启用该功能。

import torch
from diffusers import DiffusionPipeline

pipeline = DiffusionPipeline.from_pretrained(
    "black-forest-labs/FLUX.1-schnell", dtype=torch.bfloat16
)
pipeline.enable_sequential_cpu_offload()

pipeline(
    prompt="An astronaut riding a horse on Mars",
    guidance_scale=0.,
    height=768,
    width=1360,
    num_inference_steps=4,
    max_sequence_length=256,
).images[0]
print(f"Max memory reserved: {torch.cuda.max_memory_allocated() / 1024**3:.2f} GB")

模型卸载

模型卸载（Model offloading）会将整个模型移至 GPU，而不是选择性地移动部分层或模型组件。其中一个主要的流水线模型（通常是文本编码器、UNet 和 VAE）会被放置在 GPU 上，而其他组件则保留在 CPU 上。像 UNet 这样需要多次运行的组件会一直留在 GPU 上，直到其完全执行完毕且不再需要为止。这消除了CPU 卸载带来的通信开销，使模型卸载成为一种更快的替代方案。其代价是内存节省幅度不会那么大。

请注意，如果在安装钩子之后，模型在流水线之外被复用（有关更多详细信息，请参阅移除钩子），则需要按照预期的顺序运行整个流水线及模型，才能正确卸载它们。这是一个有状态的操作，会在模型上安装钩子。

调用 enable_model_cpu_offload() 以在管道中启用该功能。

import torch
from diffusers import DiffusionPipeline

pipeline = DiffusionPipeline.from_pretrained(
    "black-forest-labs/FLUX.1-schnell", dtype=torch.bfloat16
)
pipeline.enable_model_cpu_offload()

pipeline(
    prompt="An astronaut riding a horse on Mars",
    guidance_scale=0.,
    height=768,
    width=1360,
    num_inference_steps=4,
    max_sequence_length=256,
).images[0]
print(f"Max memory reserved: {torch.cuda.max_memory_allocated() / 1024**3:.2f} GB")

enable_model_cpu_offload() 在单独使用 encode_prompt() 方法生成文本编码器的隐藏状态时也有帮助。

组卸载

Group offloading 将内部层的组（torch.nn.ModuleList 或 torch.nn.Sequential）移至 CPU。与 model offloading 相比，它占用的内存更少；与 CPU offloading 相比，它速度更快，因为它降低了通信开销。

如果前向实现中包含依赖于权重的输入设备转换，则组卸载可能无法在所有模型上正常工作，因为这可能会与组卸载的设备转换机制发生冲突。

通过将 offload_type 参数配置为 block_level 或 leaf_level 来启用组卸载。

block_level 卸载根据 num_blocks_per_group 参数对层进行分组。例如，对于一个具有 40 层的模型，如果 num_blocks_per_group=2，则每次同时加载和卸载 2 层（总共进行 20 次加载/卸载）。这可以大幅降低内存需求。
leaf_level 在最低层级对单个层进行卸载，其效果等同于 CPU 卸载。但如果你使用流（streams），可以在不牺牲推理速度的情况下使其运行得更快。

组卸载（Group offloading）支持应用于整个流水线或单个模型。将组卸载应用于整个流水线是最简单的选择，而将其选择性地应用于单个模型则能赋予用户更大的灵活性，以便为不同的模型使用不同的卸载技术。

pipeline

model

CUDA 流

use_stream 参数可在支持异步数据传输流的 CUDA 设备上启用，与 CPU offloading 相比，可降低整体执行时间。它通过层预取来重叠数据传输与计算：在执行当前层的同时，将下一个待执行的层加载到 GPU 上。这可能会显著增加 CPU 内存占用，因此请确保拥有模型大小 2 倍的内存。

设置 record_stream=True 可以进一步提速，但会略微增加内存占用。请参阅 torch.Tensor.record_stream 文档以了解更多信息。

当在启用 tiling 的 VAE 上设置 use_stream=True 时，请在推理前执行一次前向传播（也可以使用虚拟输入），以避免设备不匹配错误。此方法可能不适用于所有实现，因此如果您遇到任何问题，请随时提交 issue。

如果您在使用 block_level 组卸载且启用了 use_stream，则应将 num_blocks_per_group 参数设置为 1，否则将触发警告。

pipeline.transformer.enable_group_offload(onload_device=onload_device, offload_device=offload_device, offload_type="leaf_level", use_stream=True, record_stream=True)

low_cpu_mem_usage 参数可以设置为 True，以在使用流进行组卸载时降低 CPU 内存占用。该设置最适合 leaf_level 卸载以及 CPU 内存成为瓶颈的场景。其通过动态创建固定张量（pinned tensors）而非预先固定来节省内存。不过，这可能会增加整体执行时间。

卸载到磁盘

根据模型大小，Group offloading 可能会占用大量系统内存。在内存有限的系统上，可以尝试将 group offloading 到磁盘，将其作为二级内存使用。

在 enable_group_offload() 或 apply_group_offloading() 中设置 offload_to_disk_path 参数，即可将模型卸载到磁盘。

pipeline.transformer.enable_group_offload(onload_device=onload_device, offload_device=offload_device, offload_type="leaf_level", offload_to_disk_path="path/to/disk")

apply_group_offloading(pipeline.text_encoder, onload_device=onload_device, offload_type="block_level", num_blocks_per_group=2, offload_to_disk_path="path/to/disk")

请参考这两个表格，以比较速度和内存之间的权衡。

逐层类型转换

将逐层类型转换与组卸载结合使用，可进一步节省内存。

逐层转换（Layerwise casting）将权重存储为更小的数据格式（例如 torch.float8_e4m3fn 和 torch.float8_e5m2）以减少内存占用，并在计算时将权重上转为更高精度，如 torch.float16 或 torch.bfloat16。某些层（与归一化和调制相关的权重）会被跳过，因为以 fp8 存储它们可能会降低生成质量。

如果前向实现中包含对权重的内部类型转换，逐层转换可能无法在所有模型上正常工作。当前逐层转换的实现假设前向传播与权重精度无关，且输入数据类型始终指定为 compute_dtype（参见此处了解一个不兼容的实现）。

在包含 PEFT 层的自定义建模实现中，逐层类型转换也可能失败。虽然存在一些检查机制，但它们并未经过充分测试，也无法保证在所有情况下都能正常工作。

调用 enable_layerwise_casting() 以设置存储和计算的数据类型。

import torch
from diffusers import CogVideoXPipeline, CogVideoXTransformer3DModel
from diffusers.utils import export_to_video

transformer = CogVideoXTransformer3DModel.from_pretrained(
    "THUDM/CogVideoX-5b",
    subfolder="transformer",
    dtype=torch.bfloat16
)
transformer.enable_layerwise_casting(storage_dtype=torch.float8_e4m3fn, compute_dtype=torch.bfloat16)

pipeline = CogVideoXPipeline.from_pretrained("THUDM/CogVideoX-5b",
    transformer=transformer,
    dtype=torch.bfloat16
).to("cuda")  # or "mps", "xpu", "cpu"
prompt = (
    "A panda, dressed in a small, red jacket and a tiny hat, sits on a wooden stool in a serene bamboo forest. "
    "The panda's fluffy paws strum a miniature acoustic guitar, producing soft, melodic tunes. Nearby, a few other "
    "pandas gather, watching curiously and some clapping in rhythm. Sunlight filters through the tall bamboo, "
    "casting a gentle glow on the scene. The panda's face is expressive, showing concentration and joy as it plays. "
    "The background includes a small, flowing stream and vibrant green foliage, enhancing the peaceful and magical "
    "atmosphere of this unique musical performance."
)
video = pipeline(prompt=prompt, guidance_scale=6, num_inference_steps=50).frames[0]
print(f"Max memory reserved: {torch.cuda.max_memory_allocated() / 1024**3:.2f} GB")
export_to_video(video, "output.mp4", fps=8)

如果需要更多的控制和灵活性，也可以使用 apply_layerwise_casting() 方法。通过在特定的内部模块上调用该方法，可以将其部分应用于模型层。使用 skip_modules_pattern 或 skip_modules_classes 参数来指定需要避免的模块，例如归一化层和调制层。

import torch
from diffusers import CogVideoXTransformer3DModel
from diffusers.hooks import apply_layerwise_casting

transformer = CogVideoXTransformer3DModel.from_pretrained(
    "THUDM/CogVideoX-5b",
    subfolder="transformer",
    dtype=torch.bfloat16
)

# skip the normalization layer
apply_layerwise_casting(
    transformer,
    storage_dtype=torch.float8_e4m3fn,
    compute_dtype=torch.bfloat16,
    skip_modules_classes=["norm"],
    non_blocking=True,
)

torch.channels_last

torch.channels_last 将张量的存储方式从 (batch size, channels, height, width) 翻转为 (batch size, height, width, channels)。这使得张量的布局与硬件按顺序访问内存中张量的方式保持一致，从而避免了在内存中来回跳跃以访问像素值。

并非所有算子目前都支持 channels-last 格式，使用它可能会导致性能下降，但仍值得尝试。

print(pipeline.unet.conv_out.state_dict()["weight"].stride())  # (2880, 9, 3, 1)
pipeline.unet.to(memory_format=torch.channels_last)  # in-place operation
print(
    pipeline.unet.conv_out.state_dict()["weight"].stride()
)  # (2880, 1, 960, 320) having a stride of 1 for the 2nd dimension proves that it works

内存高效注意力

Diffusers 通过 set_attention_backend() 支持多种内存高效的注意力后端（FlashAttention、xFormers、SageAttention 等）。请参阅 Attention backends 指南，了解如何在它们之间进行切换。

Update on GitHub