此页面由 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。如需了解 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 优化的 Cloud GPU 服务容器。对于受支持的模型，当您使用一键式部署或 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 容器部署模型，并将预测请求发送到端点。部署的代码段如下所示：

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（布尔值）：通过一次解码一个瓦片来减少 GPU 内存用量。此参数对于较大的图片或视频非常有用，并且有助于防止内存不足错误。默认值为 false。
ENABLE_SLICING（布尔值）：通过将输入张量拆分为用于 VAE 解码的 slice，减少 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 类。