開放原始碼程式庫 xDiT 採用平行處理和最佳化技術,加快 Diffusion Transformer (DiT) 模型的推論速度。有了這些技術,您就能採用可擴充的多 GPU 設定,處理要求嚴苛的工作負載。這個頁面展示如何在 Vertex AI 使用 xDiT 和 Cloud GPU 部署 DiT 模型。
如要進一步瞭解 xDiT,請參閱 xDiT GitHub 專案。
優點
以下列出在 Vertex AI 上使用 xDiT 服務 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 端點。
前往模型資訊卡頁面,然後按一下「部署」。
如要使用模型變體,請選取要用於部署作業的機器類型。
按一下「部署」即可開始部署程序。您會收到兩封電子郵件通知,一封是在模型上傳時收到,另一封則是在端點準備就緒時收到。
使用 Colab Enterprise 筆記本
如要彈性自訂,請使用 Colab Enterprise 筆記本範例,透過 Python 適用的 Vertex AI SDK,使用 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
(字串):指定要載入的模型 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,請參閱說明文件。