此页面由 Cloud Translation API 翻译。

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

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 性能：即使在单个 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 文生视频模型变体。如需查看 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

xDiT 优化的 Cloud GPU 服务容器在 Vertex AI Model Garden 中提供。对于受支持的模型，当您使用一键式部署或 Colab Enterprise 笔记本示例时，部署会使用此容器。

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

使用一键式部署

您可以使用模型卡片部署提供 xDiT 容器的自定义 Vertex AI 端点。

前往模型卡片页面，然后点击部署。
为要使用的模型变体选择要用于部署的机器类型。
点击部署以开始部署过程。您会收到两条电子邮件通知：在模型上传完成时以及在端点准备就绪时。

使用 Colab Enterprise 笔记本

为了实现灵活的自定义内容，请使用 Colab Enterprise 笔记本示例通过 Vertex AI SDK for Python 部署提供 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（字符串）：指定要加载的模型标识符。此标识符应与注册表或路径中的模型名称匹配。

运行时优化参数

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 时，必须提供此参数。如果不使用，请留空。

缓存优化参数

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

精确率优化参数

USE_FP8_T5_ENCODER（布尔值）：为 T5 文本编码器启用 FP8（8 位浮点）精度，从而减少内存用量并可能提高吞吐量，同时对质量的影响微乎其微。默认值为 false。

部署自定义

Model Garden 为支持的模型提供默认 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，请参阅文档。