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。如要查看 DiT 模型是否支援 Model Garden 中的 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 最佳化的雲端 GPU 服務容器。對於支援的模型，如果您使用一鍵部署或 Colab Enterprise 筆記本範例，部署作業會使用這個容器。
以下範例使用 Flux.1-schnell 模型，示範如何在 xDiT 容器上部署 DiT 模型。
使用一鍵部署
您可以使用模型資訊卡，透過 xDiT 容器部署自訂 Vertex AI 端點。
前往模型資訊卡頁面，然後按一下「部署」。
針對要使用的模型變化版本，選取要用於部署的機器類型。
按一下「部署」即可開始部署程序。您會收到兩封電子郵件通知，一封是在模型上傳時收到，另一封是在端點準備就緒時收到。
使用 Colab Enterprise 筆記本
為求靈活性和客製化，請使用 Colab Enterprise 筆記本範例，透過 Python 適用的 Vertex AI SDK 部署含有 xDiT 容器的 Vertex AI 端點。
前往模型資訊卡頁面，然後按一下「開啟 Notebook」。
選取 Vertex Serving 筆記本。筆記本會在 Colab Enterprise 中開啟。
透過筆記本執行，使用 xDiT 容器部署模型，並將預測要求傳送至端點。部署作業的程式碼片段如下所示：
XDIT_DOCKER_URI=us-docker.pkg.dev/deeplearning-platform-release/vertex-model-garden/xdit-serve.cu125.0-1.ubuntu2204.py310
serving_env = {
"MODEL_ID": "black-forest-labs/FLUX.1-schnell",
"TASK": "text-to-image",
"DEPLOY_SOURCE": "notebook",
"N_GPUS": "2",
"ULYSSES_DEGREE": "1",
"RING_DEGREE": "2",
"PIPEFUSION_PARALLEL_DEGREE": "1",
"USE_TORCH_COMPILE": "true",
}
model = aiplatform.Model.upload(
display_name=model_name,
serving_container_image_uri=XDIT_DOCKER_URI,
serving_container_ports=[7080],
serving_container_predict_route="/predict",
serving_container_health_route="/health",
serving_container_environment_variables=serving_env,
)
model.deploy(
endpoint=endpoint,
machine_type="a3-highgpu-2g",
accelerator_type="NVIDIA_H100_80GB",
accelerator_count=2,
deploy_request_timeout=1800,
service_account=SERVICE_ACCOUNT,
)
環境變數
MODEL_ID：指定要部署的 DiT 模型 ID，例如「black-forest-labs/FLUX.1-schnell」。TASK：定義模型執行的工作，例如「text-to-image-flux-xdit」。N_GPUS：設定用於推論的 GPU 數量。ULYSSES_DEGREE、RING_DEGREE、PIPEFUSION_PARALLEL_DEGREE：控制 xDiT 使用的平行處理技術。如要進一步瞭解各個引數，請參閱「xDiT 引數」。USE_TORCH_COMPILE：使用 torch.compile 啟用單一 GPU 加速功能。
xDiT 引數
xDiT 提供各種伺服器引數，可設定為針對特定用途進行最佳化。這些引數會在部署期間設為環境變數。以下是您可能需要設定的主要引數：
N_GPUS(整數)：指定用於推論的 GPU 數量。預設值為1。ENABLE_TILING(布林值)：一次解碼一個 VAE 元件，藉此減少 GPU 記憶體用量。這個引數適用於較大的圖片或影片，可避免記憶體不足錯誤。預設值為false。ENABLE_SLICING(布林值)：透過將輸入張量切割成 VAE 解碼用的切片，減少 GPU 記憶體用量。預設值為false。USE_TORCH_COMPILE(布林值)：透過 torch.compile 啟用單一 GPU 加速功能，提升編譯速度。預設值為false。PIPEFUSION_PARALLEL_DEGREE(整數)：設定 PipeFusion 的平行處理程度。值越高，平行處理的項目就越多，但可能需要更多記憶體。預設值為1。WARMUP_STEPS(整數)：如果已啟用 PipeFusion，這個引數會指定在推論開始前所需的暖機步驟數量。預設值為0。ULYSSES_DEGREE(整數)：設定 Ulysses 的角度。預設值為1。RING_DEGREE(整數)：設定環形圖形的角度。預設值為1。USE_CFG_PARALLEL(布林值)：為無分類器指引 (CFG) 啟用平行運算功能，這是用來控制 DiT 模型輸出的一種技術。啟用後，常數平行度會設為 2。使用 CFG 時設為「true」。預設值為false。USE_PARALLEL_VAE(布林值)：透過並行處理 VAE 元件，有效處理高解析度圖片 (超過 2048 像素)。預設值為false。
如需完整的引數清單，請參閱 xDiT GitHub 專案中的 xFuserArgs 類別。