Diffusers 文档
处理大型模型
并获得增强的文档体验
开始使用
处理大型模型
像 Stable Diffusion XL (SDXL) 这样的现代扩散模型，不仅仅是一个单一模型，而是多个模型的集合。SDXL 具有四个不同的模型级别组件
- 变分自编码器 (VAE)
- 两个文本编码器
- 用于去噪的 UNet
通常，文本编码器和去噪器比 VAE 大得多。
随着模型变得更大更好，您的模型可能变得非常大，以至于即使是单个副本也无法装入内存。但这并不意味着它无法加载。如果您有多个 GPU，则有更多内存可用于存储您的模型。在这种情况下，最好将您的模型 checkpoint 拆分为几个较小的checkpoint 分片。
当文本编码器 checkpoint 有多个分片时，例如 SD3 的 T5-xxl，Transformers 库会自动处理它，因为当使用 StableDiffusion3Pipeline 时，它是 Diffusers 的必需依赖项。更具体地说，Transformers 将自动处理所请求模型类中多个分片的加载，并使其准备就绪，以便可以执行推理。
去噪器 checkpoint 也可以有多个分片，并且由于 Accelerate 库的支持，它支持推理。
有关处理难以装入内存的大型模型的一般指导，请参阅 处理大型模型以进行推理 指南。
例如，让我们为 SDXL UNet 保存分片 checkpoint
from diffusers import UNet2DConditionModel
unet = UNet2DConditionModel.from_pretrained(
"stabilityai/stable-diffusion-xl-base-1.0", subfolder="unet"
)
unet.save_pretrained("sdxl-unet-sharded", max_shard_size="5GB")SDXL UNet checkpoint 的 fp32 变体大小约为 10.4GB。将 max_shard_size 参数设置为 5GB 以创建 3 个分片。保存后，您可以在 StableDiffusionXLPipeline 中加载它们
from diffusers import UNet2DConditionModel, StableDiffusionXLPipeline
import torch
unet = UNet2DConditionModel.from_pretrained(
"sayakpaul/sdxl-unet-sharded", torch_dtype=torch.float16
)
pipeline = StableDiffusionXLPipeline.from_pretrained(
"stabilityai/stable-diffusion-xl-base-1.0", unet=unet, torch_dtype=torch.float16
).to("cuda")
image = pipeline("a cute dog running on the grass", num_inference_steps=30).images[0]
image.save("dog.png")如果一次将所有模型级别组件都放在 GPU 上不可行，请使用 enable_model_cpu_offload() 来帮助您
- pipeline.to("cuda")
+ pipeline.enable_model_cpu_offload()一般来说，我们建议当 checkpoint 大于 5GB（fp32）时进行分片。
设备放置
在分布式设置中，您可以使用 Accelerate 在多个 GPU 上运行推理。
此功能是实验性的，其 API 在未来可能会发生变化。
使用 Accelerate，您可以使用 device_map 来确定如何在多个设备之间分配 pipeline 的模型。这在您有多个 GPU 的情况下非常有用。
例如，如果您有两个 8GB GPU，那么使用 enable_model_cpu_offload() 可能效果不佳，因为
- 它仅适用于单个 GPU
- 单个模型可能无法装入单个 GPU（enable_sequential_cpu_offload() 可能会起作用，但速度会非常慢，并且也仅限于单个 GPU）
为了利用两个 GPU，您可以使用“balanced”设备放置策略，该策略将模型分布在所有可用的 GPU 上。
目前仅支持“balanced”策略，我们计划在未来支持其他映射策略。
from diffusers import DiffusionPipeline
import torch
pipeline = DiffusionPipeline.from_pretrained(
- "stable-diffusion-v1-5/stable-diffusion-v1-5", torch_dtype=torch.float16, use_safetensors=True,
+ "stable-diffusion-v1-5/stable-diffusion-v1-5", torch_dtype=torch.float16, use_safetensors=True, device_map="balanced"
)
image = pipeline("a dog").images[0]
image您还可以传递一个字典来强制限制每个设备上可以使用的最大 GPU 内存
from diffusers import DiffusionPipeline
import torch
max_memory = {0:"1GB", 1:"1GB"}
pipeline = DiffusionPipeline.from_pretrained(
"stable-diffusion-v1-5/stable-diffusion-v1-5",
torch_dtype=torch.float16,
use_safetensors=True,
device_map="balanced",
+ max_memory=max_memory
)
image = pipeline("a dog").images[0]
image如果 max_memory 中不存在某个设备，则该设备将被完全忽略，并且不会参与设备放置。
默认情况下，Diffusers 使用所有设备的最大内存。如果模型无法装入 GPU，则会将它们卸载到 CPU。如果 CPU 没有足够的内存，您可能会看到错误。在这种情况下，您可以考虑使用 enable_sequential_cpu_offload() 和 enable_model_cpu_offload()。
调用 reset_device_map() 以重置 pipeline 的 device_map。如果您想在已设备映射的 pipeline 上使用 to()、enable_sequential_cpu_offload() 和 enable_model_cpu_offload() 等方法，这也是必要的。
pipeline.reset_device_map()
一旦 pipeline 被设备映射，您还可以通过 hf_device_map 访问其设备映射
print(pipeline.hf_device_map)设备映射示例可能如下所示
{'unet': 1, 'vae': 1, 'safety_checker': 0, 'text_encoder': 0}