このページは Cloud Translation API によって翻訳されました。

Cloud GPU で xDiT コンテナを使用して Diffusion Transformer モデルを提供する

xDiT は、並列処理と最適化の手法を使用して、Diffusion Transformer（DiT）モデルの推論を高速化するオープンソースライブラリです。これらの手法により、要求の厳しいワークロードに対応できるスケーラブルなマルチ GPU の設定が可能になります。このページでは、xDiT と Vertex AI 上の Cloud GPU を使用して DiT モデルをデプロイする方法について例示的に説明しています。

xDiT の詳細については、xDiT GitHub プロジェクトをご覧ください。

利点

次のリストは、xDiT を使用して Vertex AI で DiT モデルをサービングする主なメリットをまとめたものです。

最大 3 倍の高速生成: 他のサービングソリューションと比較して、高解像度の画像や動画を短時間で生成します。
スケーラブルなマルチ GPU サポート: 複数の GPU にワークロードを効率的に分散し、最適なパフォーマンスを実現します。
- ハイブリッド並列処理: xDiT は、統合シーケンス並列処理、PipeFusion、CFG 並列処理、データ並列処理など、さまざまな並列処理アプローチをサポートしています。これらの方法は、パフォーマンスを最適化するために独自のレシピで組み合わせることができます。
単一 GPU のパフォーマンスの最適化: xDiT は、単一 GPU でも推論を高速化します。
- GPU アクセラレーション: xDiT には、複数のカーネルアクセラレーション手法が組み込まれており、DiTFastAttn の手法を使用して、単一の GPU での推論を高速化します。
簡単なデプロイ: Vertex AI Model Garden のワンクリックデプロイまたは Colab Enterprise ノートブックを使用して、すぐに開始できます。

サポートされているモデル

xDiT は、Vertex AI Model Garden の特定の DiT モデルアーキテクチャ（Flux.1 Schnell、CogVideoX-2b、Wan2.1 テキスト動画変換モデルバリアントなど）で使用できます。Model Garden で DiT モデルが xDiT をサポートしているかどうかを確認するには、Model Garden でモデルカードを表示します。

マルチ GPU パフォーマンスのためのハイブリッド並列処理:

xDiT は、並列処理技術を組み合わせて使用し、マルチ GPU 設定でパフォーマンスを最大化します。これらの手法は連携して、ワークロードを分散し、リソース使用率を最適化します。

統合シーケンス並列処理: この手法では、入力データ（画像をパッチに分割するなど）を複数の GPU に分割し、メモリ使用量を削減してスケーラビリティを向上させます。
PipeFusion: PipeFusion は DiT モデルをステージに分割し、各ステージを異なる GPU に割り当てて、モデルの異なる部分の並列処理を可能にします。
CFG 並列処理: この手法は、生成された画像のスタイルとコンテンツを制御する一般的な方法である分類器なしガイダンスを使用して、モデルを最適化します。条件分岐と無条件分岐の計算を並列化し、推論を高速化します。
データ並列処理: この方法では、各 GPU にモデル全体を複製し、各 GPU が異なる入力データバッチを処理することで、システムの全体的なスループットを向上させます。

パフォーマンスの改善について詳しくは、xDiT の Flux.1 Schnell または CogVideoX-2b に関するレポートをご覧ください。Google は、Vertex AI Model Garden でこれらの結果を再現できました。

単一 GPU アクセラレーション

xDiT ライブラリは、torch.compile と onediff を使用して GPU のランタイム速度を向上させることで、単一 GPU サービングにメリットをもたらします。これらの手法は、ハイブリッド並列処理と組み合わせて使用することもできます。

また、xDiT には、DiT の計算ボトルネックに対処するための DiTFastAttn という効率的なアテンション計算手法もあります。現時点では、この手法は単一の GPU 設定でのみ、またはデータ並列処理と組み合わせて使用できます。

Model Garden を使ってみる

xDiT 最適化 Cloud GPU サービングコンテナは、Vertex AI Model Garden 内で提供されます。サポートされているモデルの場合、ワンクリックデプロイまたは Colab Enterprise ノートブックの例を使用すると、デプロイでこのコンテナが使用されます。

次の例では、Flux.1-schnell モデルを使用して、xDiT コンテナに DiT モデルをデプロイする方法を示します。

ワンクリックデプロイを使用する

モデルカードを使用して、xDiT コンテナを備えたカスタム Vertex AI エンドポイントをデプロイできます。

モデルカードページに移動し、[デプロイ] をクリックします。
使用するモデルのバリエーションに対して、デプロイに使用するマシンタイプを選択します。
[デプロイ] をクリックして、デプロイプロセスを開始します。2 通のメール通知が届きます。1 通はモデルがアップロードされたとき、もう 1 通はエンドポイントの準備が整ったときです。

Colab Enterprise ノートブックを使用する

柔軟性とカスタマイズのために、Colab Enterprise ノートブックの例を使用して、Vertex AI SDK for Python を使用して xDiT コンテナで Vertex AI エンドポイントをデプロイします。

モデルカードのページに移動し、[ノートブックを開く] をクリックします。
Vertex Serving ノートブックを選択します。ノートブックが Colab Enterprise で開きます。
ノートブックを実行して、xDiT コンテナを使用してモデルをデプロイし、予測リクエストをエンドポイントに送信します。デプロイのコードスニペットは次のとおりです。

import vertexai
from vertexai import model_garden

vertexai.init(project=<YOUR_PROJECT_ID>, location=<REGION>)

model = model_garden.OpenModel("black-forest-labs/FLUX.1-schnell")
endpoint = model.deploy()

xDiT 引数

xDiT には、特定のユースケースのパフォーマンスを最適化するために構成できるさまざまなサーバー引数が用意されています。これらの引数は、デプロイ時に環境変数として設定されます。構成する必要がある主な引数は次のとおりです。

モデルの構成

MODEL_ID（文字列）: 読み込むモデル識別子を指定します。これは、レジストリまたはパスのモデル名と一致する必要があります。

ランタイム最適化引数

N_GPUS（整数）: 推論に使用する GPU の数を指定します。デフォルト値は 1 です。
WARMUP_STEPS（整数）: 推論を開始する前に必要なウォームアップステップの数。これは、PipeFusion が有効になっている場合に、安定したパフォーマンスを確保するために特に重要です。デフォルト値は 1 です。
USE_PARALLEL_VAE（ブール値）: デバイス間で VAE コンポーネントを並列化することで、高解像度画像（2,048 ピクセル超）の効率的な処理を可能にします。これにより、大きな画像で OOM の問題が発生するのを防ぎます。デフォルト値は false です。
USE_TORCH_COMPILE（ブール値）: torch.compile を使用して単一 GPU アクセラレーションを有効にし、パフォーマンスを向上させるカーネルレベルの最適化を提供します。デフォルト値は false です。
USE_ONEDIFF（ブール値）: OneDiff コンパイルアクセラレーションテクノロジーを有効にして、GPU カーネルの実行速度を最適化します。デフォルト値は false です。

データ並列処理の引数

DATA_PARALLEL_DEGREE（整数）: データ並列処理の度合いを設定します。無効にする場合は空欄にします。並列処理の度合いを設定する場合は、その値を指定します。
USE_CFG_PARALLEL（ブール値）: 分類なしガイダンス（CFG）（スプリットバッチとも呼ばれます）の並列計算を有効にします。有効にすると、並列処理の定数は 2 になります。CFG を使用して出力スタイルとコンテンツを制御する場合は、true に設定します。デフォルト値は false です。

シーケンス並列引数（USP - Unified Sequence Parallelism）

ULYSSES_DEGREE（整数）: DeepSpeed-Ulysses と Ring-Attention を組み合わせた統合シーケンス並列アプローチの Ulysses 度を設定します。これは、全対全通信パターンを制御します。デフォルトを使用する場合は空欄にします。
RING_DEGREE（整数）: シーケンス並列処理でピアツーピア通信のリング次数を設定します。ULYSSES_DEGREE と連携して 2D プロセスメッシュを形成します。デフォルトを使用する場合は空欄にします。

テンソル並列処理の引数

TENSOR_PARALLEL_DEGREE（整数）: テンソル並列処理の度合いを設定します。これにより、特徴ディメンションに沿ってデバイス間でモデルパラメータが分割され、デバイスあたりのメモリ費用が削減されます。無効にする場合は空欄にします。
SPLIT_SCHEME（文字列）: モデルテンソルをデバイス間で分割する方法（アテンションヘッド、隠れ次元など）を定義します。デフォルトの分割スキームを使用する場合は空欄にします。

Ray 分散引数

USE_RAY（ブール値）: 複数のノードにわたって計算をスケーリングするための Ray 分散実行フレームワークを有効にします。デフォルト値は false です。
RAY_WORLD_SIZE（整数）: Ray クラスタ内のプロセスの合計数。デフォルト値は 1 です。
VAE_PARALLEL_SIZE（整数）: Ray を使用する場合に VAE 並列処理専用のプロセス数。デフォルト値は 0 です。
DIT_PARALLEL_SIZE（整数）: Ray を使用する場合に DiT バックボーンの並列処理専用のプロセス数。デフォルト値は 0 です。

PipeFusion 並列引数

PIPEFUSION_PARALLEL_DEGREE（整数）: 拡散モデルの入力時間冗長性の特性を利用するシーケンスレベルのパイプライン並列処理である PipeFusion の並列処理の程度を設定します。値を大きくすると並列処理は増加しますが、より多くのメモリが必要になります。デフォルト値は 1 です。
NUM_PIPELINE_PATCH（整数）: パイプライン処理のためにシーケンスを分割するパッチの数。自動判定する場合は空欄にします。
ATTN_LAYER_NUM_FOR_PP（文字列）: パイプライン並列処理に使用する注意レイヤを指定します。カンマ区切りで指定できます（例: "10,9"）またはスペース区切り（例: "10 9"）。すべてのレイヤを使用する場合は、空白のままにします。

メモリ最適化引数

ENABLE_MODEL_CPU_OFFLOAD（ブール値）: 使用されていないときにモデルの重みを CPU メモリにオフロードし、レイテンシの増加と引き換えに GPU メモリの使用量を削減します。デフォルト値は false です。
ENABLE_SEQUENTIAL_CPU_OFFLOAD（ブール値）: 順方向パス中にモデルレイヤを CPU に順次オフロードし、GPU メモリよりも大きなモデルの推論を可能にします。デフォルト値は false です。
ENABLE_TILING（ブール値）: VAE コンポーネントを一度に 1 つのタイルでデコードして、GPU メモリの使用量を減らします。この引数は、大きな画像や動画に便利で、メモリ不足エラーを防ぐことができます。デフォルト値は false です。
ENABLE_SLICING（ブール値）: 入力テンソルをスライスに分割して VAE デコードを行うことで、GPU メモリ使用量を削減します。デフォルト値は false です。

DiTFastAttn 引数（注意最適化）

USE_FAST_ATTN（ブール値）: 単一 GPU 推論で DiTFastAttn アクセラレーションを有効にし、入力時間削減を利用して計算の複雑さを軽減します。デフォルト値は false です。
N_CALIB（整数）: DiTFastAttn 最適化のキャリブレーションサンプル数。デフォルト値は 8 です。
THRESHOLD（浮動小数点）: DiTFastAttn の時間的類似性削減の類似性しきい値。デフォルト値は 0.5 です。
WINDOW_SIZE（整数）: 空間冗長性を削減する残差キャッシュを使用した Window Attention のウィンドウサイズ。デフォルト値は 64 です。
COCO_PATH（文字列）: DiTFastAttn キャリブレーション用の COCO データセットのパス。USE_FAST_ATTN が true の場合は必須です。使用しない場合は空のままにします。

キャッシュ最適化の引数

USE_CACHE（ブール値）: 一般的なキャッシュメカニズムを有効にして、冗長な計算を減らします。デフォルト値は false です。
USE_TEACACHE（ブール値）: 中間結果のキャッシュ保存に TeaCache 最適化メソッドを有効にします。デフォルト値は false です。
USE_FBCACHE（ブール値）: First-Block-Cache 最適化メソッドを有効にします。デフォルト値は false です。

精度最適化の引数

USE_FP8_T5_ENCODER（ブール値）: T5 テキストエンコーダで FP8（8 ビット浮動小数点）精度を有効にします。これにより、メモリ使用量が減少し、品質への影響を最小限に抑えながらスループットが向上する可能性があります。デフォルト値は false です。

サービングのカスタマイズ

Model Garden には、サポートされているモデルのデフォルトの xDiT 並列化構成が用意されています。これらのデフォルト設定は、Python 用 Vertex AI SDK を使用して確認できます。

「black-forest-labs/FLUX.1-schnell」などのモデルのデフォルトのデプロイ構成を表示するには、次のコードスニペットを実行します。

import vertexai
from vertexai import model_garden

vertexai.init(project=<YOUR_PROJECT_ID>, location=<REGION>)

model = model_garden.OpenModel("black-forest-labs/FLUX.1-schnell")
deploy_options = model.list_deploy_options()


# Example Response
# ['black-forest-labs/flux1-schnell@flux.1-schnell']
# [model_display_name: "Flux1-schnell"
# container_spec {
#   image_uri: "us-docker.pkg.dev/deeplearning-platform-release/vertex-model-garden/xdit-serve.cu125.0-2.ubuntu2204.py310"
#  env {
#    name: "DEPLOY_SOURCE"
#    value: "UI_NATIVE_MODEL"
#  }
#  env {
#    name: "MODEL_ID"
#    value: "gs://vertex-model-garden-restricted-us/black-forest-labs/FLUX.1-schnell"
#  }
#  env {
#    name: "TASK"
#    value: "text-to-image"
#  }
#  env {
#    name: "N_GPUS"
#    value: "2"
#  }
#  env {
#    name: "USE_TORCH_COMPILE"
#    value: "true"
#  }
#  env {
#    name: "RING_DEGREE"
#    value: "2"
#  }
# ..........]

list_deploy_options() メソッドは、xDiT 構成を定義する環境変数（env）を含むコンテナ仕様を返します。

並列処理戦略をカスタマイズするには、モデルのデプロイ時にこれらの環境変数をオーバーライドします。次の例は、2 つの GPU の設定で RING_DEGREE と ULYSSES_DEGREE を変更し、並列処理のアプローチを変更する方法を示しています。

import vertexai
from vertexai import model_garden

# Replace with your project ID and region
vertexai.init(project="<YOUR_PROJECT_ID>", location="<REGION>")

model = model_garden.OpenModel("black-forest-labs/FLUX.1-schnell")

# Custom environment variables to override default settings
# This example sets N_GPUS as 2, so RING_DEGREE * ULYSSES_DEGREE must equal 2
container_env_vars = {
    "N_GPUS": "2",
    "RING_DEGREE": "1",
    "ULYSSES_DEGREE": "2"
    # Add other environment variables to customize here
}

machine_type = "a3-highgpu-2g"
accelerator_type = "NVIDIA_H100_80GB"
accelerator_count = 2

# Deploy the model with the custom environment variables
endpoint = model.deploy(
    machine_type=machine_type,
    accelerator_type=accelerator_type,
    accelerator_count=accelerator_count,
  container_env_vars=container_env_vars
)

各環境変数の詳細については、「xDiT 固有の引数について」のセクションをご覧ください。並列処理の度合い（PIPEFUSION_PARALLEL_DEGREE、ULYSSES_DEGREE、RING_DEGREE、USE_CFG_PARALLEL の合計が GPU の合計数（N_GPUS）と等しい。

さまざまなモデルのサービングレシピと構成の例については、xDiT の公式ドキュメントをご覧ください。Model Garden SDK の詳細については、ドキュメントをご覧ください。

Cloud GPU で xDiT コンテナを使用して Diffusion Transformer モデルを提供する コレクションでコンテンツを整理 必要に応じて、コンテンツの保存と分類を行います。

利点