Hugging Face's logo

Buckets:
hf-doc-build
/
doc

Files
xet
hf-doc-build/doc / diffusers /v0.43.0 /zh /optimization /memory.md
mishig's picture
mishig
preview code
|
download
raw
28.6 kB
# 减少内存使用
像 [Flux](../api/pipelines/flux) 和 [Wan](../api/pipelines/wan) 这样的现代扩散模型拥有数十亿参数，在硬件上进行推理时会占用大量内存。由于常见的 GPU 往往内存不足，这带来了挑战。为了克服内存限制，您可以使用多块 GPU（如果可用）、将部分流水线组件卸载到 CPU 上，以及采用其他方法。
本指南将向你展示如何降低内存使用量。
> [!TIP]
> 请注意，这些技术可能需要根据模型进行调整。例如，基于 Transformer 的扩散模型从这些内存优化中获得的收益，可能不如基于 UNet 的模型。
## 多 GPU
如果您拥有多块 GPU，则有几种选项可以高效地将大型模型加载并分布到您的硬件上。这些功能由 [Accelerate](https://huggingface.co/docs/accelerate/index) 库提供支持，因此请确保先安装该库。
```bash
pip install -U accelerate
```
### 分片检查点
分片加载大型检查点很有用，因为分片是逐个加载的。这样可以保持内存占用较低，只需满足模型大小和最大分片大小所需的内存即可。我们建议当 fp32 检查点大于 5GB 时进行分片。默认的分片大小为 5GB。
使用 `max_shard_size` 参数在 [save_pretrained()](/docs/diffusers/v0.43.0/zh/api/pipelines/overview#diffusers.DiffusionPipeline.save_pretrained) 中对检查点进行分片。
```py
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）来替代常规检查点，以节省内存。
```py
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"
```
### 设备放置
> [!WARNING]
> 设备放置是一项实验性功能，API 可能会发生变化。目前仅支持 `balanced` 策略。我们计划在未来支持更多的映射策略。
`device_map` 参数用于控制流水线中的模型组件或单个模型中的层如何在不同设备上进行分配。
`balanced` 设备放置策略会将流水线均匀地拆分到所有可用设备上。
```py
import torch
from diffusers import AutoModel, StableDiffusionXLPipeline
pipeline = StableDiffusionXLPipeline.from_pretrained(
"stabilityai/stable-diffusion-xl-base-1.0",
dtype=torch.float16,
device_map="balanced"
)
```
您可以通过 `hf_device_map` 检查 pipeline 的设备映射。
```py
print(pipeline.hf_device_map)
{'unet': 1, 'vae': 1, 'safety_checker': 0, 'text_encoder': 0}
```
`device_map` 适用于加载大型模型，例如拥有 12.5B 参数的 Flux 扩散 Transformer。将其设置为 `"auto"` 可自动将模型优先分配到最快的设备上，然后再分配到较慢的设备。有关更多详细信息，请参阅 [分片检查点](#sharded-checkpoints) 部分。
```py
import torch
from diffusers import AutoModel
transformer = AutoModel.from_pretrained(
"black-forest-labs/FLUX.1-dev",
subfolder="transformer",
device_map="auto",
dtype=torch.bfloat16
)
```
您可以通过 `hf_device_map` 检查模型的 device map。
```py
print(transformer.hf_device_map)
```
在设计您自己的 `device_map` 时，它应该是一个字典，其中包含模型特定的模块名称或层以及设备标识符（GPU 使用整数，CPU 使用 `cpu`，磁盘使用 `disk`）。
对模型调用 `hf_device_map` 以查看模型层的分布情况，然后自行设计。
```py
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`）上。
```py
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` 中，则会被忽略，流水线组件也不会被分配到该设备上。
```py
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()](/docs/diffusers/v0.43.0/zh/api/pipelines/overview#diffusers.DiffusionPipeline.enable_model_cpu_offload) 仅适用于单 GPU，但非常大的模型可能无法装入单个 GPU
- [enable_sequential_cpu_offload()](/docs/diffusers/v0.43.0/zh/api/pipelines/overview#diffusers.DiffusionPipeline.enable_sequential_cpu_offload) 可能可以运行，但速度极慢，且仅限于单 GPU
使用 [reset_device_map()](/docs/diffusers/v0.43.0/zh/api/pipelines/overview#diffusers.DiffusionPipeline.reset_device_map) 方法重置 `device_map`。如果您想在已进行设备映射的 pipeline 上使用诸如 `.to()`、[enable_sequential_cpu_offload()](/docs/diffusers/v0.43.0/zh/api/pipelines/overview#diffusers.DiffusionPipeline.enable_sequential_cpu_offload) 和 [enable_model_cpu_offload()](/docs/diffusers/v0.43.0/zh/api/pipelines/overview#diffusers.DiffusionPipeline.enable_model_cpu_offload) 之类的方法，则必须执行此操作。
```py
pipeline.reset_device_map()
```
## VAE 切片
VAE 切片通过将一个大批次输入拆分为单个数据批次并分别处理来节省内存。此方法在同时生成多张图像时效果最佳。
例如，如果您一次生成 4 张图像，解码过程会使峰值激活内存增加 4 倍。VAE 切片通过每次只解码 1 张图像（而不是一次解码全部 4 张）来降低这一开销。
调用 `enable_slicing()` 以启用切片式 VAE。在解码多图像批次时，性能会有小幅提升；对于单图像批次，则不会有性能影响。
```py
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")
```
> [!WARNING]
> [AutoencoderKLWan](/docs/diffusers/v0.43.0/zh/api/models/autoencoder_kl_wan#diffusers.AutoencoderKLWan) 和 [AsymmetricAutoencoderKL](/docs/diffusers/v0.43.0/zh/api/models/asymmetricautoencoderkl#diffusers.AsymmetricAutoencoderKL) 类不支持切片。
## VAE 分块
VAE 分块（tiling）通过将图像划分为较小的重叠块，而非一次性处理整张图像，从而节省内存。由于 GPU 每次仅处理一个块，这也有助于降低峰值内存占用。
调用 `enable_tiling()` 以启用 VAE 分块。由于各分块是独立解码的，生成的图像在分块之间可能会出现一些色调差异，但分块之间不应出现明显的接缝。对于低于预设（但可配置）阈值的分辨率，分块功能将被禁用。例如，[StableDiffusionPipeline](/docs/diffusers/v0.43.0/zh/api/pipelines/stable_diffusion/text2img#diffusers.StableDiffusionPipeline) 中 VAE 的该阈值为 512x512。
```py
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")
```
> [!WARNING]
> [AutoencoderKLWan](/docs/diffusers/v0.43.0/zh/api/models/autoencoder_kl_wan#diffusers.AutoencoderKLWan) 和 [AsymmetricAutoencoderKL](/docs/diffusers/v0.43.0/zh/api/models/asymmetricautoencoderkl#diffusers.AsymmetricAutoencoderKL) 不支持分块处理。
若要完全避免在 GPU 上占用 VAE 内存，请使用 [Remote inference](../hybrid_inference/overview) 在远程端点上执行编码和解码。
## 卸载
Offloading 策略会将当前未激活的层或模型移至 CPU，以避免增加 GPU 内存占用。这些策略可以与量化和 torch.compile 结合使用，以平衡推理速度和内存使用。
有关更多详细信息，请参阅[编译和卸载量化模型](./speed-memory-optims)指南。
### CPU 卸载
CPU offloading 会将权重从 GPU 选择性地移至 CPU。当某个组件被需要时，它会被传输到 GPU；当不再需要时，它会被移回 CPU。此方法针对子模块而非整个模型进行操作。通过避免在 GPU 上存储整个模型，该方法可以节省内存。
CPU offloading 能大幅降低内存占用，但由于子模块需要在设备之间多次来回传递，其速度**极其缓慢**。因此，在实际应用中往往难以采用。
> [!WARNING]
> 在调用 [enable_sequential_cpu_offload()](/docs/diffusers/v0.43.0/zh/api/pipelines/overview#diffusers.DiffusionPipeline.enable_sequential_cpu_offload) 之前，请勿将 pipeline 移至 CUDA，否则节省的内存量将非常有限（更多详情请参见此 [issue](https://github.com/huggingface/diffusers/issues/1934)）。这是一个有状态的操作，会在模型上安装钩子。
调用 [enable_sequential_cpu_offload()](/docs/diffusers/v0.43.0/zh/api/pipelines/overview#diffusers.DiffusionPipeline.enable_sequential_cpu_offload) 以在管道上启用该功能。
```py
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 卸载](#cpu-offloading)带来的通信开销，使模型卸载成为一种更快的替代方案。其代价是内存节省幅度不会那么大。
> [!WARNING]
> 请注意，如果在安装钩子之后，模型在流水线之外被复用（有关更多详细信息，请参阅 [移除钩子](https://huggingface.co/docs/accelerate/en/package_reference/big_modeling#accelerate.hooks.remove_hook_from_module)），则需要按照预期的顺序运行整个流水线及模型，才能正确卸载它们。这是一个有状态的操作，会在模型上安装钩子。
调用 [enable_model_cpu_offload()](/docs/diffusers/v0.43.0/zh/api/pipelines/overview#diffusers.DiffusionPipeline.enable_model_cpu_offload) 以在管道中启用该功能。
```py
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()](/docs/diffusers/v0.43.0/zh/api/pipelines/overview#diffusers.DiffusionPipeline.enable_model_cpu_offload) 在您单独使用 [encode_prompt()](/docs/diffusers/v0.43.0/zh/api/pipelines/stable_diffusion/stable_diffusion_xl#diffusers.StableDiffusionXLPipeline.encode_prompt) 方法生成文本编码器的隐藏状态时也有帮助。
### 组卸载
Group offloading 将内部层的组（[torch.nn.ModuleList](https://pytorch.org/docs/stable/generated/torch.nn.ModuleList.html) 或 [torch.nn.Sequential](https://pytorch.org/docs/stable/generated/torch.nn.Sequential.html)）移至 CPU。与 [model offloading](#model-offloading) 相比，它占用的内存更少；与 [CPU offloading](#cpu-offloading) 相比，它速度更快，因为它降低了通信开销。
> [!WARNING]
> 如果前向实现中包含依赖于权重的输入设备转换，组卸载可能无法在所有模型上正常工作，因为这可能会与组卸载的设备转换机制发生冲突。
通过将 `offload_type` 参数配置为 `block_level` 或 `leaf_level` 来启用组卸载。
- `block_level` 卸载根据 `num_blocks_per_group` 参数对层进行分组。例如，对于一个具有 40 层的模型，如果 `num_blocks_per_group=2`，则每次同时加载和卸载 2 层（总共进行 20 次加载/卸载）。这可以大幅降低内存需求。
- `leaf_level` 在最低层级对单个层进行卸载，其效果等同于 [CPU 卸载](#cpu-offloading)。但如果你使用流（streams），可以在不牺牲推理速度的情况下使其运行得更快。
组卸载（Group offloading）支持应用于整个流水线或单个模型。将组卸载应用于整个流水线是最简单的选择，而将其选择性地应用于单个模型则能赋予用户更大的灵活性，以便为不同的模型使用不同的卸载技术。
在 pipeline 上调用 [enable_group_offload()](/docs/diffusers/v0.43.0/zh/api/pipelines/overview#diffusers.DiffusionPipeline.enable_group_offload)。
```py
import torch
from diffusers import CogVideoXPipeline
from diffusers.hooks import apply_group_offloading
from diffusers.utils import export_to_video
onload_device = torch.device("cuda")
offload_device = torch.device("cpu")
pipeline = CogVideoXPipeline.from_pretrained("THUDM/CogVideoX-5b", dtype=torch.bfloat16)
pipeline.enable_group_offload(
onload_device=onload_device,
offload_device=offload_device,
offload_type="leaf_level",
use_stream=True
)
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)
```
在继承自 [ModelMixin](/docs/diffusers/v0.43.0/zh/api/models/overview#diffusers.ModelMixin) 的标准 Diffusers 模型组件上调用 [enable_group_offload()](/docs/diffusers/v0.43.0/zh/api/models/overview#diffusers.ModelMixin.enable_group_offload)。对于未继承自 [ModelMixin](/docs/diffusers/v0.43.0/zh/api/models/overview#diffusers.ModelMixin) 的其他模型组件（例如通用的 [torch.nn.Module](https://pytorch.org/docs/stable/generated/torch.nn.Module.html)），请改用 [apply_group_offloading()](/docs/diffusers/v0.43.0/zh/api/utilities#diffusers.hooks.apply_group_offloading)。
```py
import torch
from diffusers import CogVideoXPipeline
from diffusers.hooks import apply_group_offloading
from diffusers.utils import export_to_video
onload_device = torch.device("cuda")
offload_device = torch.device("cpu")
pipeline = CogVideoXPipeline.from_pretrained("THUDM/CogVideoX-5b", dtype=torch.bfloat16)
# Use the enable_group_offload method for Diffusers model implementations
pipeline.transformer.enable_group_offload(onload_device=onload_device, offload_device=offload_device, offload_type="leaf_level")
pipeline.vae.enable_group_offload(onload_device=onload_device, offload_type="leaf_level")
# Use the apply_group_offloading method for other model components
apply_group_offloading(pipeline.text_encoder, onload_device=onload_device, offload_type="block_level", num_blocks_per_group=2)
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)
```
#### CUDA 流
`use_stream` 参数可在支持异步数据传输流的 CUDA 设备上启用，与 [CPU offloading](#cpu-offloading) 相比，可降低整体执行时间。它通过层预取来重叠数据传输与计算：在执行当前层的同时，将下一个待执行的层加载到 GPU 上。这可能会显著增加 CPU 内存占用，因此请确保拥有模型大小 2 倍的内存。
设置 `record_stream=True` 可以进一步提速，但会略微增加内存占用。请参阅 [torch.Tensor.record_stream](https://pytorch.org/docs/stable/generated/torch.Tensor.record_stream.html) 文档以了解更多信息。
> [!TIP]
> 当在启用 tiling 的 VAE 上设置 `use_stream=True` 时，请在推理前执行一次前向传播（也可以使用虚拟输入），以避免设备不匹配错误。此方法可能不适用于所有实现，因此如果您遇到任何问题，请随时提交 issue。
如果您在使用 `block_level` 组卸载且启用了 `use_stream`，则应将 `num_blocks_per_group` 参数设置为 `1`，否则将触发警告。
```py
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()](/docs/diffusers/v0.43.0/zh/api/models/overview#diffusers.ModelMixin.enable_group_offload) 或 [apply_group_offloading()](/docs/diffusers/v0.43.0/zh/api/utilities#diffusers.hooks.apply_group_offloading) 中设置 `offload_to_disk_path` 参数，即可将模型卸载到磁盘。
```py
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")
```
请参考这[两个](https://github.com/huggingface/diffusers/pull/11682#issue-3129365363)[表格](https://github.com/huggingface/diffusers/pull/11682#issuecomment-2955715126)，以比较速度和内存之间的权衡。
## 逐层类型转换
> [!TIP]
> 将逐层类型转换与 [组卸载](#group-offloading) 结合使用，可进一步节省内存。
逐层转换（Layerwise casting）将权重存储为更小的数据格式（例如 `torch.float8_e4m3fn` 和 `torch.float8_e5m2`）以减少内存占用，并在计算时将权重上转为更高精度，如 `torch.float16` 或 `torch.bfloat16`。某些层（与归一化和调制相关的权重）会被跳过，因为以 fp8 存储它们可能会降低生成质量。
> [!WARNING]
> 如果前向实现中包含对权重的内部类型转换，逐层转换可能无法在所有模型上正常工作。当前逐层转换的实现假设前向传播与权重精度无关，且输入数据类型始终指定为 `compute_dtype`（参见 [此处](https://github.com/huggingface/transformers/blob/7f5077e53682ca855afc826162b204ebf809f1f9/src/transformers/models/t5/modeling_t5.py#L294-L299) 了解一个不兼容的实现）。
>
> 在包含 [PEFT](https://huggingface.co/docs/peft/index) 层的自定义建模实现中，逐层类型转换也可能失败。虽然存在一些检查机制，但它们并未经过充分测试，也无法保证在所有情况下都能正常工作。
调用 [enable_layerwise_casting()](/docs/diffusers/v0.43.0/zh/api/models/overview#diffusers.ModelMixin.enable_layerwise_casting) 以设置存储和计算的数据类型。
```py
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()](/docs/diffusers/v0.43.0/zh/api/utilities#diffusers.hooks.apply_layerwise_casting) 方法。通过在特定的内部模块上调用该方法，可以将其部分应用于模型层。使用 `skip_modules_pattern` 或 `skip_modules_classes` 参数来指定需要避免的模块，例如归一化层和调制层。
```python
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](https://pytorch.org/tutorials/intermediate/memory_format_tutorial.html) 将张量的存储方式从 `(batch size, channels, height, width)` 翻转为 `(batch size, height, width, channels)`。这使得张量的布局与硬件按顺序访问内存中张量的方式保持一致，从而避免了在内存中来回跳跃以访问像素值。
并非所有算子目前都支持 channels-last 格式，使用它可能会导致性能下降，但仍值得尝试。
```py
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()](/docs/diffusers/v0.43.0/zh/api/models/overview#diffusers.ModelMixin.set_attention_backend) 支持多种内存高效的注意力后端（FlashAttention、xFormers、SageAttention 等）。请参阅 [Attention backends](./attention_backends) 指南，了解如何在它们之间进行切换。

Xet Storage Details

Size:
28.6 kB
·
Xet hash:
19752cfe33ce134a544292a7cb786145d04a5a25d1ab8b871460a48d50acfb37

Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.