使用 xDiT 容器和 Cloud GPU 提供 Diffusion Transformer 模型

開放原始碼程式庫 xDiT 採用平行處理和最佳化技術，加快 Diffusion Transformer (DiT) 模型的推論速度。有了這些技術，您就能採用可擴充的多 GPU 設定，處理要求嚴苛的工作負載。這個頁面展示如何在 Gemini Enterprise Agent Platform 使用 xDiT 和 Cloud GPU 部署 DiT 模型。

如要進一步瞭解 xDiT，請參閱 xDiT GitHub 專案。

優點

以下列出使用 xDiT 在 Agent Platform 上提供 DiT 模型的主要優點：

  • 生成速度最多提升三倍：與其他放送解決方案相比，生成高解析度圖片和影片所需的時間大幅縮短。
  • 可擴充的多 GPU 支援：在多個 GPU 間有效分配工作負載，達到最佳效能。
    • 混合式平行處理：xDiT 支援各種平行處理方法，例如統一序列平行處理、PipeFusion、CFG 平行處理和資料平行處理。您可以將這些方法組合為獨特的配方，進一步提升成效。
  • 最佳化單一 GPU 效能：即使使用單一 GPU，xDiT 也能提供更快的推論速度。
    • GPU 加速：xDiT 採用多種核心加速方法，並使用 DiTFastAttn 的技術，加快單一 GPU 的推論速度。
  • 輕鬆部署：在 Gemini Enterprise Agent Platform Model Garden 中，只要按一下即可部署，或使用 Colab Enterprise Notebook，快速開始使用。

支援的模型

在 Gemini Enterprise Agent Platform Model Garden 中，xDiT 適用於特定 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 已在 Gemini Enterprise Agent Platform Model Garden 重現這些結果。

單一 GPU 加速

xDiT 程式庫使用 torch.compile 和 onediff，提升 GPU 的執行階段速度，為單一 GPU 服務帶來優勢。這些技術也可以搭配混合式平行處理使用。

xDiT 也採用名為 DiTFastAttn 的高效注意力運算技術，解決 DiT 的運算瓶頸。目前，這項技術僅適用於單一 GPU 設定，或搭配資料平行處理使用。

在 Model Garden 中開始使用

xDiT 最佳化 Cloud GPU 服務容器位於 Gemini Enterprise Agent Platform Model Garden 中。對於支援的模型，當您使用一鍵部署或 Colab Enterprise 筆記本範例時，部署作業會使用這個容器。

下列範例使用 Flux.1-schnell 模型，示範如何在 xDiT 容器上部署 DiT 模型。

使用一鍵部署功能

您可以使用模型資訊卡，透過 xDiT 容器部署自訂 Agent Platform 端點。

  1. 前往模型資訊卡頁面，然後按一下「部署」。

  2. 如要使用模型變體，請選取要用於部署作業的機型。

  3. 按一下「部署」即可開始部署程序。您會收到兩封電子郵件通知，分別是模型上傳完成時，以及端點準備就緒時。

使用 Colab Enterprise 筆記本

如要彈性自訂，請使用 Colab Enterprise 筆記本範例，透過 Agent Platform SDK for Python 部署含有 xDiT 容器的 Agent Platform 端點。

  1. 前往模型資訊卡頁面，然後按一下「開啟筆記本」。

  2. 選取 Vertex Serving 筆記本。筆記本會在 Colab Enterprise 中開啟。

  3. 執行筆記本，使用 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 (字串)：指定要載入的模型 ID。這應與登錄或路徑中的模型名稱相符。
執行階段最佳化引數
  • N_GPUS (整數)：指定用於推論的 GPU 數量。預設值為 1。
  • WARMUP_STEPS (整數)：推論開始前所需的暖機步驟數。啟用 PipeFusion 時，這項設定特別重要，可確保效能穩定。預設值為 1。
  • USE_PARALLEL_VAE (布林值)：在裝置間平行處理 VAE 元件，有效處理高解析度圖片 (大於 2048 像素)。這樣可避免大型圖片發生 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 - 統一序列平行處理)
  • ULYSSES_DEGREE (整數)：為統一序列平行方法設定 Ulysses 程度，結合 DeepSpeed-Ulysses 和 Ring-Attention。這會控管所有對所有通訊模式。留空即可使用預設值。
  • 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 元件，藉此減少 GPU 記憶體用量。這個引數適用於較大的圖片或影片，可避免記憶體不足錯誤。預設值為 false。
  • ENABLE_SLICING (布林值)：將輸入張量分割成切片，以進行 VAE 解碼，藉此減少 GPU 記憶體用量。預設值為 false。
DiTFastAttn 引數 (注意力最佳化)
  • USE_FAST_ATTN (布林值)：為單一 GPU 推論啟用 DiTFastAttn 加速功能，利用輸入時間縮減功能降低運算複雜度。預設值為 false。
  • N_CALIB (整數)：用於 DiTFastAttn 最佳化的校正樣本數量。預設值為 8。
  • THRESHOLD (浮點數)：DiTFastAttn 中時間相似度縮減的類似程度門檻。預設值為 0.5。
  • WINDOW_SIZE (整數)：使用殘餘快取的視窗注意力機制視窗大小，可減少空間冗餘。預設值為 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 專用的 Agent Platform 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，請參閱說明文件。