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

本文档介绍了如何在 Vertex AI 中使用 Cloud GPU 上的 xDiT 容器来部署 Diffusion Transformer (DiT) 模型。本文档介绍了以下主题：

• 优势：了解使用 xDiT 在性能和可伸缩性方面的优势。
• 支持的模型：了解 Model Garden 中哪些模型与 xDiT 兼容。
• 并行处理和优化技术：了解 xDiT 用于加快推理速度的混合并行处理和单 GPU 加速方法。
• 在 Model Garden 中开始使用：按照说明使用一键式部署或 Colab Enterprise 笔记本部署模型。
• 自定义部署：了解如何替换默认设置，并找到可配置实参的完整列表，以实现最佳性能。

下图总结了整个工作流程：

xDiT 是一个开源库，可使用并行处理和优化技术加快 Diffusion Transformer (DiT) 模型的推理速度。这些技术可针对要求苛刻的工作负载提供可伸缩的多 GPU 设置。本文档介绍了如何在 Vertex AI 上使用 xDiT 和 Cloud GPU 部署 DiT 模型。

如需详细了解 xDiT，请参阅 xDiT GitHub 项目。

优势

使用 xDiT 在 Vertex AI 上部署 DiT 模型具有以下优势：

• 生成速度最多可提高 3 倍：与其他服务解决方案相比，生成高分辨率图片和视频所需的时间更短。
• 可扩缩的多 GPU 支持：高效地将工作负载分配到多个 GPU，以获得最佳性能。xDiT 支持各种并行处理方法，例如统一序列并行处理、PipeFusion、CFG 并行处理和数据并行处理。您可以结合使用这些方法来优化效果。
• 优化了单 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 设置中最大限度地提高性能。这些技术协同工作，以分配工作负载并优化资源利用率。

如需详细了解性能改进，请参阅 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 优化的服务容器。对于受支持的模型，当您使用一键式部署或 Colab Enterprise 笔记本示例时，部署会使用此容器。

下表比较了这两种部署方法。

以下示例使用 Flux.1-schnell 模型来演示如何在 xDiT 容器上部署 DiT 模型。

使用一键式部署

如需使用模型卡片部署提供 xDiT 容器的自定义 Vertex AI 端点，请按以下步骤操作：

1. 前往模型卡片页面，然后点击部署。
2. 为部署选择模型变体和机器类型。
3. 点击部署。部署流程开始。您会在模型上传完成时收到电子邮件通知，并在端点准备就绪时收到另一封电子邮件通知。

使用 Colab Enterprise 笔记本

为了实现更灵活的自定义内容，您可以使用 Colab Enterprise 笔记本示例通过 Vertex AI SDK for Python 部署提供 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()

如需了解如何通过替换默认设置来自定义部署，请参阅自定义部署。

自定义部署

Model Garden 为支持的模型提供默认 xDiT 并行化配置。您可以查看这些默认设置，并根据需要替换它们。

覆盖默认设置

如需查看模型的默认部署配置（例如 black-forest-labs/FLUX.1-schnell），请使用 Vertex AI SDK for Python，如以下代码所示：

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 官方文档。如需详细了解 Model Garden SDK，请参阅文档。

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 - 统一序列并行处理)

• ULYSSES_DEGREE（整数）：为统一序列并行方法设置 Ulysses 度，该方法结合了 DeepSpeed-Ulysses 和 Ring-Attention。此设置用于控制全对全通信模式。如需使用默认值，请将此值留空。
• RING_DEGREE（整数）：在序列并行处理中设置点对点通信的环度。与 ULYSSES_DEGREE 结合使用，以形成二维进程网格。如需使用默认值，请将此值留空。

张量并行参数

• 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 的并行度，PipeFusion 是一种序列级流水线并行处理，可利用 diffusion 模型的输入时间冗余特征。值越大，并行度越高，但需要的内存也越多。默认值为 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 时，必须设置此参数。如果您未使用 DiTFastAttn，请将此值留空。

缓存优化实参

• USE_CACHE（布尔值）：启用常规缓存机制以减少冗余计算。默认值为 false。
• USE_TEACACHE（布尔值）：启用 TeaCache 优化方法来缓存中间结果。默认值为 false。
• USE_FBCACHE（布尔值）：启用“首块缓存”(FBCache) 优化方法。默认值为 false。

精确率优化参数

• USE_FP8_T5_ENCODER（布尔值）：为 T5 文本编码器启用 FP8（8 位浮点）精度。这样可以减少内存用量，并可在对质量影响极小的情况下提高吞吐量。默认值为 false。