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

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

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

優點

以下列出使用 xDiT 在 Vertex AI 上提供 DiT 模型服務的主要優點：

• 生成速度最多可提升三倍：與其他放送解決方案相比，生成高解析度圖片和影片所需的時間大幅縮短。
• 可擴充的多 GPU 支援：在多個 GPU 間有效分配工作負載，以獲得最佳效能。
  • 混合式平行處理：xDiT 支援各種平行處理方法，例如統一序列平行處理、PipeFusion、CFG 平行處理和資料平行處理。您可以將這些方法組合為獨特的配方，進一步提升成效。
• 最佳化單一 GPU 效能：即使使用單一 GPU，xDiT 也能提供更快的推論速度。
  • 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 也採用名為 DiTFastAttn 的高效注意力運算技術，解決 DiT 的運算瓶頸。目前，這項技術僅適用於單一 GPU 設定，或搭配資料平行處理使用。

開始使用 Model Garden

Vertex AI Model Garden 提供經過最佳化的 xDiT Cloud GPU 服務容器。對於支援的模型，當您使用一鍵部署或 Colab Enterprise 筆記本範例時，部署作業會使用這個容器。

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

使用一鍵部署功能

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

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

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

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

使用 Colab Enterprise 筆記本

如要彈性自訂，請使用 Colab Enterprise 筆記本範例，透過 Python 適用的 Vertex AI SDK，使用 xDiT 容器部署 Vertex AI 端點。

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。

放送自訂項目

模型園地會為支援的模型提供預設的 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，請參閱說明文件。