处理大型模型
现代扩散模型，例如 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}