device_map を使用してモデルのコンポーネントを複数の GPU に分散

device_map 引数は、モデルのコンポーネントを複数の GPU に配置します。モデルが 1 つの GPU に収まらず、かつ GPU 間に高速なインターコネクトがない場合に使用してください。他の戦略との比較については分散推論を、device_map のその他のオプションについてはデバイスの配置を参照してください。

以下の例では、各コンポーネントを必要とするステージでのみ読み込み、device_mapを使用して大きなものを2つのGPUに分散しています。

プロンプトをエンコード

device_map="balanced" を設定すると、テキストエンコーダーが別々の GPU に配置されます。各 GPU が使用するメモリ量を上限値で制限するには max_memory を使用します。テキストエンコーダーのみを読み込むようにするには、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
    )

テキスト埋め込みを取得したら、transformerをGPUに載せるために、テキストエンコーダーをGPUから解放します。

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()

sharded transformerでデノイズ

device_map="auto" を設定すると、Accelerate が transformer の配置を決定します。Accelerate は最初の GPU にメモリを充填してから次の GPU に移動し、GPU のメモリが不足した場合にのみ CPU またはディスクを使用します。2 つの小型 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 で読み込まれたパイプラインにも、各コンポーネントがどこに配置されているかを示す 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

潜在変数をデコードする

パイプラインとトランスフォーマーをメモリから解放し、VAE を読み込んで潜在変数をデコードします。VAE は 1 つの 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 に分散

プロンプトをエンコード

sharded transformerでデノイズ

潜在変数をデコードする