使用 device_map 将模型组件拆分到多个 GPU

device_map 参数用于将模型的各个组件放置在不同的 GPU 上。当模型无法装入单个 GPU，且你的 GPU 之间没有高速互连时，请使用该参数。请参阅分布式推理以将其与其他策略进行比较，并参阅设备放置了解 device_map 的更多选项。

以下示例仅在需要该组件的阶段加载它，并使用 device_map 将较大的组件分布在两块 GPU 上。

编码提示

将 device_map="balanced" 设置为将文本编码器放置在单独的 GPU 上。使用 max_memory 来限制每个 GPU 使用的内存量。将 transformer 和 vae 设置为 None，以便仅加载文本编码器。

from diffusers import FluxPipeline
import torch

prompt = """
cinematic film still of a cat sipping a margarita in a pool in Palm Springs, California
highly detailed, high budget hollywood movie, cinemascope, moody, epic, gorgeous, film grain
"""

pipeline = FluxPipeline.from_pretrained(
    "black-forest-labs/FLUX.1-dev",
    transformer=None,
    vae=None,
    device_map="balanced",
    max_memory={0: "16GB", 1: "16GB"},
    dtype=torch.bfloat16
)
with torch.no_grad():
    print("Encoding prompts.")
    prompt_embeds, pooled_prompt_embeds, text_ids = pipeline.encode_prompt(
        prompt=prompt, prompt_2=None, max_sequence_length=512
    )

获取文本嵌入后，将文本编码器从 GPU 上移除，以为 transformer 腾出空间。

import gc

def flush():
    gc.collect()
    torch.cuda.empty_cache()
    torch.cuda.reset_max_memory_allocated()
    torch.cuda.reset_peak_memory_stats()

del pipeline.text_encoder
del pipeline.text_encoder_2
del pipeline.tokenizer
del pipeline.tokenizer_2
del pipeline

flush()

使用分片 Transformer 进行去噪

设置 device_map="auto" 以让 Accelerate 放置 transformer。它会先填满第一块 GPU，然后再移动到下一块，并且仅在 GPU 内存耗尽时才使用 CPU 或磁盘。在两块较小的 GPU 上，transformer 最终会被拆分到两块 GPU 上。位于 CPU 或磁盘上的权重会拖慢推理速度。

from diffusers import AutoModel
import torch

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

检查 transformer.hf_device_map 以查看 transformer 在设备上的分布情况。使用 device_map 加载的 pipeline 也具有 hf_device_map 属性，用于显示每个组件的放置位置。

加载仅包含 transformer 的 pipeline，并设置 output_type="latent"，使其返回潜变量而非图像。

pipeline = FluxPipeline.from_pretrained(
    "black-forest-labs/FLUX.1-dev",
    text_encoder=None,
    text_encoder_2=None,
    tokenizer=None,
    tokenizer_2=None,
    vae=None,
    transformer=transformer,
    dtype=torch.bfloat16
)

print("Running denoising.")
height, width = 768, 1360
latents = pipeline(
    prompt_embeds=prompt_embeds,
    pooled_prompt_embeds=pooled_prompt_embeds,
    num_inference_steps=50,
    guidance_scale=3.5,
    height=height,
    width=width,
    output_type="latent",
).images

解码潜在变量

从内存中移除 pipeline 和 transformer，并加载 VAE 以解码潜变量。VAE 足够小，可以放入单个 GPU 中。

import torch
from diffusers import AutoencoderKL
from diffusers.image_processor import VaeImageProcessor

del pipeline.transformer
del pipeline
del transformer

flush()

vae = AutoencoderKL.from_pretrained(
    "black-forest-labs/FLUX.1-dev", subfolder="vae", dtype=torch.bfloat16
).to("cuda")  # or "mps", "xpu", "cpu"
vae_scale_factor = 2 ** (len(vae.config.block_out_channels) - 1)
image_processor = VaeImageProcessor(vae_scale_factor=vae_scale_factor)

with torch.no_grad():
    print("Running decoding.")
    latents = FluxPipeline._unpack_latents(latents, height, width, vae_scale_factor)
    latents = (latents / vae.config.scaling_factor) + vae.config.shift_factor

    image = vae.decode(latents, return_dict=False)[0]
    image = image_processor.postprocess(image, output_type="pil")
    image[0].save("split_transformer.png")

通过选择性地加载和卸载在特定阶段所需的模型，并将最大的模型分片到多个 GPU 上，您可以在消费级 GPU 上运行大模型的推理。

Update on GitHub

Diffusers

使用 device_map 将模型组件拆分到多个 GPU

编码提示

使用分片 Transformer 进行去噪

解码潜在变量