xDiT es una biblioteca de código abierto que acelera la inferencia de los modelos de Diffusion Transformer (DiT) a través de técnicas de paralelismo y optimización. Estas técnicas permiten configurar de forma escalable varias GPU para cargas de trabajo exigentes. En esta página, se muestra cómo implementar modelos de DiT con xDiT y GPUs de Cloud en Vertex AI.
Para obtener más información sobre xDiT, consulta el proyecto de GitHub de xDiT.
Beneficios
En la siguiente lista, se describen los principales beneficios de usar xDiT para entregar modelos de DiT en Vertex AI:
- Generación hasta tres veces más rápida: Genera imágenes y videos en alta resolución en una fracción del tiempo en comparación con otras soluciones de publicación.
- Compatibilidad escalable con varias GPUs: Distribuye las cargas de trabajo de manera eficiente en varias GPUs para obtener un rendimiento óptimo.
- Paralelismo híbrido: xDiT admite varios enfoques de procesamiento paralelo, como el paralelismo de secuencia unificado, PipeFusion, el paralelismo de CFG y el paralelismo de datos. Estos métodos se pueden combinar en una receta única para optimizar el rendimiento.
- Rendimiento optimizado de una sola GPU: xDiT proporciona una inferencia más rápida incluso en una sola GPU.
- Aceleración de GPU: xDiT incorpora varios métodos de aceleración del kernel y usa técnicas de DiTFastAttn para acelerar la inferencia en una sola GPU.
- Implementación sencilla: Comienza rápidamente con la implementación en un solo clic o los notebooks de Colab Enterprise en Vertex AI Model Garden.
Modelos compatibles
xDiT está disponible para ciertas arquitecturas de modelos DiT en Vertex AI Model Garden, como Flux.1 Schnell, CogVideoX-2b y las variantes del modelo de texto a video Wan2.1. Para ver si un modelo de DiT admite xDiT en Model Garden, consulta su tarjeta de modelo en Model Garden.
Paralelismo híbrido para el rendimiento de varias GPUs:
xDiT usa una combinación de técnicas de paralelismo para maximizar el rendimiento en configuraciones de varias GPU. Estas técnicas funcionan en conjunto para distribuir la carga de trabajo y optimizar el uso de recursos:
- Paralelismo de secuencia unificado: Esta técnica divide los datos de entrada (por ejemplo, divide una imagen en parches) en varias GPUs, lo que reduce el uso de memoria y mejora la escalabilidad.
- PipeFusion: PipeFusion divide el modelo DiT en etapas y asigna cada etapa a una GPU diferente, lo que permite el procesamiento paralelo de diferentes partes del modelo.
- Paralelismo de CFG: Esta técnica optimiza específicamente los modelos con la guía sin clasificador, un método común para controlar el estilo y el contenido de las imágenes generadas. Paraleliza el cálculo de las ramas condicionales y no condicionales, lo que genera una inferencia más rápida.
- Paralelismo de datos: Este método replica todo el modelo en cada GPU, y cada GPU procesa un lote diferente de datos de entrada, lo que aumenta el rendimiento general del sistema.
Para obtener más información sobre las mejoras de rendimiento, consulta el informe de xDiT sobre Flux.1 Schnell o CogVideoX-2b. Google pudo reproducir estos resultados en Vertex AI Model Garden.
Aceleración de GPU única
La biblioteca xDiT proporciona beneficios para la publicación con una sola GPU, ya que usa torch.compile y onediff para mejorar la velocidad del tiempo de ejecución en las GPUs. Estas técnicas también se pueden usar junto con el paralelismo híbrido.
xDiT también tiene una técnica de cálculo de atención eficiente, llamada DiTFastAttn, para abordar el cuello de botella computacional de DiT. Por el momento, esta técnica solo está disponible para configuraciones de una sola GPU o en conjunto con el paralelismo de datos.
Primeros pasos con Model Garden
El contenedor de entrega de GPU de Cloud optimizado para xDiT se proporciona en Vertex AI Model Garden. En el caso de los modelos compatibles, las implementaciones usan este contenedor cuando usas las implementaciones con un clic o los ejemplos de notebooks de Colab Enterprise.
En los siguientes ejemplos, se usa el modelo Flux.1-schnell para demostrar cómo implementar un modelo DiT en un contenedor xDiT.
Usa la implementación con un solo clic
Puedes implementar un extremo personalizado de Vertex AI con el contenedor de xDiT a través de una tarjeta de modelo.
Navega a la página de la tarjeta de modelo y haz clic en Implementar.
Para la variación del modelo que se usará, selecciona un tipo de máquina para tu implementación.
Haz clic en Implementar para comenzar el proceso de implementación. Recibirás dos notificaciones por correo electrónico: una cuando se suba el modelo y otra cuando el extremo esté listo.
Usa el notebook de Colab Enterprise
Si deseas obtener flexibilidad y personalización, usa los ejemplos de notebook de Colab Enterprise para implementar un extremo de Vertex AI con el contenedor de xDiT a través del SDK de Vertex AI para Python.
Navega a la página de la tarjeta de modelo y haz clic en Abrir notebook.
Selecciona el notebook de Vertex Serving. Se abre el notebook en Colab Enterprise.
Ejecuta el notebook para implementar un modelo con el contenedor xDiT y enviar solicitudes de predicción al extremo. El fragmento de código para la implementación es el siguiente:
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()
Argumentos de xDiT
xDiT ofrece una variedad de argumentos del servidor que se pueden configurar para optimizar el rendimiento en casos de uso específicos. Estos argumentos se configuran como variables de entorno durante la implementación. En la siguiente lista, se incluyen los argumentos clave que tal vez debas configurar:
Configuración del modelo
MODEL_ID(cadena): Especifica el identificador del modelo que se cargará. Debe coincidir con el nombre del modelo en tu registro o ruta de acceso.
Argumentos de optimización del tiempo de ejecución
N_GPUS(número entero): Especifica la cantidad de GPUs que se usarán para la inferencia. El valor predeterminado es 1.WARMUP_STEPS(número entero): Es la cantidad de pasos de preparación necesarios antes de que comience la inferencia. Esto es especialmente importante cuando PipeFusion está habilitado para garantizar un rendimiento estable. El valor predeterminado es 1.USE_PARALLEL_VAE(booleano): Habilita el procesamiento eficiente de imágenes de alta resolución (más de 2,048 píxeles) paralelizando el componente del VAE en los dispositivos. Esto evita problemas de OOM con imágenes grandes. El valor predeterminado es falso.USE_TORCH_COMPILE(booleano): Habilita la aceleración de una sola GPU a través de torch.compile, lo que proporciona optimizaciones a nivel del kernel para mejorar el rendimiento. El valor predeterminado es falso.USE_ONEDIFF(booleano): Habilita la tecnología de aceleración de la compilación de OneDiff para optimizar la velocidad de ejecución del kernel de la GPU. El valor predeterminado es falso.
Argumentos de Data Parallel
DATA_PARALLEL_DEGREE(número entero): Establece el grado de paralelismo de los datos. Déjalo vacío para inhabilitar o establecer el grado de paralelismo seleccionado.USE_CFG_PARALLEL(booleano): Habilita la computación paralela para la orientación sin clasificador (CFG), también conocida como Split Batch. Cuando está habilitado, el grado de paralelismo constante es 2. Se establece como verdadero cuando se usa CFG para controlar el estilo y el contenido del resultado. El valor predeterminado es falso.
Argumentos de paralelismo de secuencia (USP, paralelismo de secuencia unificado)
ULYSSES_DEGREE(número entero): Establece el grado de Ulysses para el enfoque paralelo de la secuencia unificada, que combina DeepSpeed-Ulysses y Ring-Attention. Este parámetro controla el patrón de comunicación de todos con todos. Déjalo vacío para usar el valor predeterminado.RING_DEGREE(número entero): Establece el grado de anillo para la comunicación peer-to-peer en el paralelismo de secuencia. Funciona en conjunto con ULYSSES_DEGREE para formar la malla de proceso 2D. Déjalo vacío para usar el valor predeterminado.
Argumentos de paralelismo de tensores
TENSOR_PARALLEL_DEGREE(número entero): Establece el grado de paralelismo de tensor, que divide los parámetros del modelo entre los dispositivos a lo largo de las dimensiones de las características para reducir los costos de memoria por dispositivo. Déjalo vacío para inhabilitar la opción.SPLIT_SCHEME(cadena): Define cómo dividir los tensores del modelo entre los dispositivos (p.ej., por encabezados de atención o dimensiones ocultas). Déjalo vacío para usar el esquema de división predeterminado.
Argumentos distribuidos de Ray
USE_RAY(booleano): Habilita el framework de ejecución distribuida de Ray para escalar los cálculos en varios nodos. El valor predeterminado es falso.RAY_WORLD_SIZE(número entero): Es la cantidad total de procesos en el clúster de Ray. El valor predeterminado es 1.VAE_PARALLEL_SIZE(número entero): Cantidad de procesos dedicados al procesamiento paralelo del VAE cuando se usa Ray. El valor predeterminado es 0.DIT_PARALLEL_SIZE(número entero): Cantidad de procesos dedicados al procesamiento paralelo de la red troncal de DiT cuando se usa Ray. El valor predeterminado es 0.
Argumentos paralelos de PipeFusion
PIPEFUSION_PARALLEL_DEGREE(número entero): Establece el grado de paralelismo para PipeFusion, un paralelismo de canalización a nivel de la secuencia que aprovecha las características de redundancia temporal de entrada de los modelos de difusión. Los valores más altos aumentan el paralelismo, pero requieren más memoria. El valor predeterminado es 1.NUM_PIPELINE_PATCH(número entero): Cantidad de parches en los que se dividirá la secuencia para el procesamiento de la canalización. Déjalo vacío para que se determine automáticamente.ATTN_LAYER_NUM_FOR_PP(cadena): Especifica qué capas de atención se usarán para el paralelismo de la canalización. Se pueden separar con comas (p.ej., "10,9") o separados por espacios (p.ej., "10 9"). Déjalo vacío para usar todas las capas.
Argumentos de optimización de memoria
ENABLE_MODEL_CPU_OFFLOAD(booleano): Descarga los pesos del modelo en la memoria de la CPU cuando no están en uso, lo que reduce el uso de la memoria de la GPU a costa de una mayor latencia. El valor predeterminado es falso.ENABLE_SEQUENTIAL_CPU_OFFLOAD(booleano): Descarga secuencialmente las capas del modelo en la CPU durante el pase hacia adelante, lo que permite la inferencia de modelos más grandes que la memoria de la GPU. El valor predeterminado es falso.ENABLE_TILING(booleano): Reduce el uso de memoria de la GPU decodificando el componente del VAE una segmentación a la vez. Este argumento es útil para imágenes o videos más grandes y para evitar errores de memoria insuficiente. El valor predeterminado es falso.ENABLE_SLICING(booleano): Reduce el uso de memoria de la GPU dividiendo el tensor de entrada en segmentos para la decodificación del VAE. El valor predeterminado es falso.
Argumentos de DiTFastAttn (optimización de la atención)
USE_FAST_ATTN(booleano): Habilita la aceleración de DiTFastAttn para la inferencia con una sola GPU, utilizando la reducción temporal de entrada para reducir la complejidad computacional. El valor predeterminado es falso.N_CALIB(número entero): Es la cantidad de muestras de calibración para la optimización de DiTFastAttn. El valor predeterminado es 8.THRESHOLD(float): Umbral de similitud para la reducción de similitud temporal en DiTFastAttn. El valor predeterminado es 0.5.WINDOW_SIZE(número entero): Tamaño de la ventana para la atención de ventana con almacenamiento en caché residual para reducir la redundancia espacial. El valor predeterminado es 64.COCO_PATH(cadena): Es la ruta de acceso al conjunto de datos de COCO para la calibración de DiTFastAttn. Se requiere cuando USE_FAST_ATTN es verdadero. Déjalo vacío si no lo usas.
Argumentos de optimización de caché
USE_CACHE(booleano): Habilita mecanismos generales de almacenamiento en caché para reducir los cálculos redundantes. El valor predeterminado es falso.USE_TEACACHE(booleano): Habilita el método de optimización de TeaCache para almacenar en caché los resultados intermedios. El valor predeterminado es falso.USE_FBCACHE(booleano): Habilita el método de optimización de First-Block-Cache. El valor predeterminado es falso.
Argumentos de optimización de precisión
USE_FP8_T5_ENCODER(booleano): Habilita la precisión FP8 (punto flotante de 8 bits) para el codificador de texto T5, lo que reduce el uso de memoria y puede mejorar el rendimiento con un impacto mínimo en la calidad. El valor predeterminado es falso.
Personalizaciones de entrega
Model Garden proporciona configuraciones de paralelización xDiT predeterminadas para los modelos admitidos. Puedes inspeccionar estos parámetros de configuración predeterminados con el SDK de Vertex AI para Python.
Para ver la configuración de implementación predeterminada de un modelo, como "black-forest-labs/FLUX.1-schnell", puedes ejecutar el siguiente fragmento de código:
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"
# }
# ..........]
El método list_deploy_options() devuelve las especificaciones del contenedor, incluidas las variables de entorno (env) que definen la configuración de xDiT.
Para personalizar la estrategia de paralelismo, puedes anular estas variables de entorno cuando implementes el modelo. En el siguiente ejemplo, se muestra cómo modificar RING_DEGREE y ULYSSES_DEGREE para una configuración de 2 GPUs, lo que cambia el enfoque de paralelismo:
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
)
Recuerda consultar la sección "Cómo comprender los argumentos específicos de xDiT" para obtener detalles sobre cada variable de entorno. Asegúrate de que el producto de los grados de paralelismo (p.ej., PIPEFUSION_PARALLEL_DEGREE, ULYSSES_DEGREE, RING_DEGREE y USE_CFG_PARALLEL) es igual a la cantidad total de GPUs (N_GPUS).
Para obtener más ejemplos de recetas y configuraciones de diferentes modelos, consulta la documentación oficial de xDiT. Para obtener más información sobre el SDK de Model Garden, consulta la documentación.