处理大型模型

现代扩散模型，例如 Stable Diffusion XL (SDXL)，不仅仅是一个单一模型，而是一组多个模型的集合。SDXL 具有四个不同的模型级组件

一个变分自动编码器 (VAE)
两个文本编码器
一个用于去噪的 UNet

通常，文本编码器和去噪器与 VAE 相比要大得多。

随着模型越来越大、越来越好，您的模型可能变得非常大，以至于单个副本都无法放入内存中。但这并不意味着它无法加载。如果您有多个 GPU，则有更多内存可用于存储模型。在这种情况下，最好将模型检查点拆分为多个较小的检查点分片。

当文本编码器检查点有多个分片时，例如 SD3 的 T5-xxl，它会由 Transformers 库自动处理，因为当使用 StableDiffusion3Pipeline 时，它是 Diffusers 的必需依赖项。更具体地说，Transformers 将自动处理在请求的模型类中加载多个分片，并使其准备好执行推理。

去噪器检查点也可以有多个分片，并通过 Accelerate 库支持推理。

有关处理难以放入内存的大型模型的一般指南，请参阅推理大型模型指南。

例如，让我们为 SDXL UNet 保存一个分片检查点

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 检查点的 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()

通常，我们建议在检查点大于 5GB（以 fp32 为单位）时进行分片。

设备放置

在分布式设置中，您可以使用 Accelerate 在多个 GPU 上运行推理。

此功能处于实验阶段，其 API 可能会在将来发生更改。

使用 Accelerate，您可以使用device_map来确定如何在多个设备上分配管道的模型。这在您有多个 GPU 的情况下很有用。

例如，如果您有两个 8GB 的 GPU，则使用 enable_model_cpu_offload() 可能效果不佳，因为

它仅适用于单个 GPU
单个模型可能无法放入单个 GPU 中（enable_sequential_cpu_offload() 可能有效，但速度会非常慢，并且也仅限于单个 GPU）

要利用两个 GPU，您可以使用“平衡”设备放置策略，该策略将模型拆分到所有可用的 GPU 上。

目前仅支持“平衡”策略，我们计划在将来支持其他映射策略。

from diffusers import DiffusionPipeline
import torch

pipeline = DiffusionPipeline.from_pretrained(
-    "runwayml/stable-diffusion-v1-5", torch_dtype=torch.float16, use_safetensors=True,
+    "runwayml/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(
    "runwayml/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() 以重置管道的device_map。如果您想在已进行设备映射的管道上使用诸如to()、enable_sequential_cpu_offload() 和 enable_model_cpu_offload() 之类的方法，这也十分必要。

pipeline.reset_device_map()

管道完成设备映射后，您还可以通过hf_device_map访问其设备映射

print(pipeline.hf_device_map)

设备映射示例如下所示

{'unet': 1, 'vae': 1, 'safety_checker': 0, 'text_encoder': 0}

< > 在 GitHub 上更新