データ並列性

データ並列性では、同じモデルを各GPUに複製し、各GPUに異なるプロンプトを割り当てます。モデルがすでに1つのGPUに収まる場合、スループットを向上させ、単一プロンプトのレイテンシは変わりません。まだ試していない場合は、まず1つのGPUでバッチ推論を試してください。他の戦略との比較については、分散推論を参照してください。

Accelerate または PyTorch Distributed を直接使用して、データ並列性を設定できます。

Accelerate

Accelerate はプロセスを設定し、それぞれに GPU を割り当てます。

次のコマンドでAccelerateをインストールします。

uv pip install accelerate

run_distributed.py というファイルを作成し、PartialState を作成して分散環境を設定します。DiffusionPipeline を PartialState.device に移動して、各プロセスに GPU を割り当てます。

import torch
from accelerate import PartialState
from diffusers import DiffusionPipeline

pipeline = DiffusionPipeline.from_pretrained(
    "Qwen/Qwen-Image", dtype=torch.bfloat16
)
distributed_state = PartialState()
pipeline.to(distributed_state.device)

同じファイル内で、split_between_processes をコンテキストマネージャーとして使用し、プロンプトをプロセス間で分割します。

with distributed_state.split_between_processes(["a dog", "a cat"]) as prompt:
    result = pipeline(prompt).images[0]
    result.save(f"result_{distributed_state.process_index}.png")

accelerate launchでスクリプトを実行し、--num_processesをGPUの数に設定します。

accelerate launch --num_processes=2 run_distributed.py

完全な例については、この最小限のマルチGPU推論スクリプトを参照してください。詳細については、Accelerateによる分散推論ガイドを参照してください。

PyTorch Distributed

Accelerate を使わずに、torch.distributed を直接使用してデータ並列性を設定することもできます。

分散環境を作成するには、init_process_groupを呼び出します。torchrunは環境変数を通じて各プロセスのランクとプロセスの総数（ワールドサイズ）を設定するため、これらを自分で渡す必要はありません。

各プロセスには、それぞれ専用のGPUとプロンプトが必要です。torchrunは、LOCAL_RANK環境変数をそのノード上のプロセスのGPUインデックスに設定するため、GPUの選択にはこれを使用します。プロンプトの選択には、すべてのノードにまたがってプロセスを数えるget_rankを使用します。

import os
import torch
import torch.distributed as dist
from diffusers import DiffusionPipeline

def main():
    dist.init_process_group("nccl")
    rank = dist.get_rank()
    local_rank = int(os.environ["LOCAL_RANK"])
    device = torch.device(f"cuda:{local_rank}")

    pipeline = DiffusionPipeline.from_pretrained(
        "Qwen/Qwen-Image", dtype=torch.bfloat16,
    ).to(device)

    prompts = ["a dog", "a cat"]
    prompt = prompts[rank]

    image = pipeline(prompt).images[0]
    image.save(f"{prompt.replace(' ', '_')}.png")

    dist.destroy_process_group()


if __name__ == "__main__":
    main()

スクリプトをrun_distributed.pyとして保存し、torchrunで実行します。--nproc-per-nodeにはGPUの数を指定します。promptsリストにはプロセスごとに1つのプロンプトが含まれるため、使用するGPUごとにプロンプトを追加してください。

torchrun --nproc-per-node 2 run_distributed.py

Update on GitHub

Diffusers

データ並列性

Accelerate

PyTorch Distributed