AMD、Radeon AI PRO R9700S最大4台でxDiT拡散推論を分散する手順を公開

AMDは、Radeon AI PRO R9700S(RDNA 4、1台あたり32GB)を1台・2台・4台に増やして拡散モデル推論をスケールする手順を、オープンソース推論エンジンxDiTとともにAMD ROCm Blogsで公開しました。4台構成では合計128GBのGPUメモリを使い、32GBを超えるモデルを分散して実行できます。記事の著者はPaul dos Santosです。

AMD ROCmは、AMDが提供するGPU計算向けのオープンなソフトウェア群です。

何が新しいか

AMDのDiffusion Model InferenceリポジトリはxDiTをコンテナにまとめ、対象ワークロード向けに調整したRDNA 4用設定を提供します。対象としてFLUX.2、Qwen-Image、Wan 2.2、Stable Diffusion 3.5、Z-Imageを挙げています。調整済み設定は、GPU間のPCIe転送を抑え、ホストメモリのスパイクを起こさずに重みを配置し、ネイティブFP8実行を使い、適したデノイズ処理を再利用し、有効な場合はVAEデコードを分散します。

仕様・数値

項目	値
GPUあたりメモリ	32GB
4 GPU合計メモリ	128GB
クイックスタートのモデル	black-forest-labs/FLUX.2-klein-4B(4B)
解像度	2048×2048
推論ステップ数	4
シード	42
分散VAEタイル	高さ352、幅2048、高さ方向オーバーラップ88、幅方向0

要件と試し方

• GPU:Radeon AI PRO R9700Sを1台・2台・4台
• ドライバー:ROCm 10.0をサポートするAMD GPUドライバー(記事の測定はamdgpu-dkms 6.19.14 (31.40))
• Linux、ROCm向けに設定したDocker(/dev/kfdと/dev/driへのアクセス)
• Hugging Faceアカウント、HF_TOKEN環境変数、必要に応じたモデル規約への同意
• Hugging Faceキャッシュ、出力、初回コンパイル成果物を置くローカルストレージ

イメージはamdsiloai/pytorch-xdit-custom:rdna4-rocm10.0.0-ubuntu24.04-py3.12-torch2.13.0-xdit-ecf0f58-20261016を使います。xdit_runヘルパーがGPUデバイスをコンテナへ渡し、Hugging Faceキャッシュ、~/.cache/xdit、./outputsをマウントします。コンパイル済みカーネルをキャッシュするため、2回目以降の実行は初回より早く生成を始めます。

1 GPUのクイックスタートでは--num_inference_steps 4、--gemm_quantization fp8、--use_fp8_text_encoder、--attention_backend aiter_flydsl_fp8、--use_torch_compile、--enable_tilingを指定します。--use_fp8_text_encoderは--gemm_quantization fp8が前提です。

2台・4台ではULYSSES_DEGREEを2または4に設定し、--memory_efficient_replicated_load、--ulysses_degree、--use_fp8_comms、--use_parallel_vae、--enable_tilingと上記タイル設定(--vae_tile_size_height 352など)を追加します。

チューニングの指針

AMDは、重みを収める→GPU間の分担を決める→トランスフォーマー実行(精度、コンパイル、キャッシュ)を調整する→目的解像度でVAEを調整する、の順を推奨しています。

• FP8はBF16比で重み容量を半減し、RDNA 4はFP8をネイティブ実行します。出力が変わり得るため、代表的なプロンプトとシードで比較するよう記事は求めています。
• 各GPUに収まらない場合はFSDP(--fully_shard_degree Dで各GPUに約1/D)を使い、収まる最小の分割度を選びます。--fully_shard_components text_encoderでテキストエンコーダーだけを分割できます。
• 1 GPUで32GBを超えるモデルには--enable_model_cpu_offloadか--enable_group_cpu_offloadを使えますが、転送のオーバーヘッドが加わります。2 GPUで大きく高速化した場合、CPUオフロードが不要になった効果も含まれ得ます。
• 並列化手法は3種類です。CFG並列(--use_cfg_parallel)は通信量が少ないもののCFGを使うモデル限定です。PipeFusion(--pipefusion_parallel_degree)はUlyssesより通信が少ない傾向がありますが、ステージ間の均衡が必要です。Ulysses(--ulysses_degree)は対応モデルが多い一方、通信量が最大です。手法は組み合わせられ、--use_cfg_parallel --ulysses_degree 2で4 GPUを使います。

単一リクエストの低遅延や1 GPUに収まらないモデルには分割が有効です。一方、各GPUに収まるワークロードで時間あたりの生成数を増やすなら、GPUごとに独立したリクエストを処理する方がPCIe通信を避けられるため、スループットは通常高くなるとAMDは説明しています。