xDiT es una biblioteca de código abierto que acelera la inferencia de los modelos de Diffusion Transformer (DiT) mediante el uso de técnicas de paralelismo y optimización. Estas técnicas permiten una configuración de varias GPUs escalable para cargas de trabajo exigentes. En esta página se muestra cómo desplegar modelos de DiT mediante xDiT y GPUs de Cloud en Vertex AI.
Para obtener más información sobre xDiT, consulta el proyecto de GitHub de xDiT.
Ventajas
En la siguiente lista se describen las principales ventajas de usar xDiT para publicar modelos de DiT en Vertex AI:
- Generación hasta tres veces más rápida: genera imágenes y vídeos de alta resolución en una fracción del tiempo en comparación con otras soluciones de publicación.
- Compatibilidad con varias GPUs escalable: distribuye las cargas de trabajo de forma eficiente entre varias GPUs para obtener un rendimiento óptimo.
- Paralelismo híbrido: xDiT admite varios enfoques de procesamiento paralelo, como el paralelismo de secuencias 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 por 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: empieza rápidamente con la implementación con un solo clic o con los cuadernos de Colab Enterprise en Vertex AI Model Garden.
Modelos admitidos
xDiT está disponible para determinadas arquitecturas de modelos de DiT en Model Garden de Vertex AI, como Flux.1 Schnell, CogVideoX-2b y las variantes del modelo de texto a vídeo 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 mejorar el rendimiento de varias GPUs:
xDiT usa una combinación de técnicas de paralelismo para maximizar el rendimiento en configuraciones de varias GPUs. Estas técnicas funcionan conjuntamente para distribuir la carga de trabajo y optimizar el uso de los recursos:
- Paralelismo de secuencia unificado: esta técnica divide los datos de entrada (por ejemplo, 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 fases y asigna cada fase a una GPU diferente, lo que permite procesar en paralelo diferentes partes del modelo.
- Paralelismo de CFG: esta técnica optimiza los modelos específicamente mediante la guía sin clasificador, un método habitual para controlar el estilo y el contenido de las imágenes generadas. Paraleliza el cálculo de las ramas condicionales e incondicionales, lo que acelera la inferencia.
- Paralelismo de datos: este método replica todo el modelo en cada GPU. 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 ha podido reproducir estos resultados en Vertex AI Model Garden.
Aceleración de una sola GPU
La biblioteca xDiT ofrece ventajas para el servicio de una sola GPU mediante el uso de torch.compile y onediff para mejorar la velocidad 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 eficiente de cálculo de la atención, llamada DiTFastAttn, para abordar el cuello de botella computacional de DiT. Por ahora, esta técnica solo se puede usar en configuraciones de una sola GPU o junto con el paralelismo de datos.
Empezar a usar Model Garden
El contenedor de servicio de GPU de Cloud optimizado para xDiT se proporciona en Vertex AI Model Garden. En los modelos admitidos, las implementaciones usan este contenedor cuando se utilizan implementaciones con un solo clic o los ejemplos de cuadernos de Colab Enterprise.
En los siguientes ejemplos se usa el modelo Flux.1-schnell para mostrar cómo implementar un modelo DiT en un contenedor xDiT.
Usar el despliegue con un solo clic
Puedes desplegar un endpoint de Vertex AI personalizado con el contenedor xDiT mediante una tarjeta de modelo.
Ve a la página de la tarjeta de modelo y haz clic en Desplegar.
En la variación del modelo que quieras usar, selecciona un tipo de máquina para tu despliegue.
Haz clic en Implementar para iniciar el proceso de implementación. Recibirás dos notificaciones por correo electrónico: una cuando se suba el modelo y otra cuando el endpoint esté listo.
Usar el cuaderno de Colab Enterprise
Para disfrutar de flexibilidad y personalización, usa los ejemplos de cuadernos de Colab Enterprise para desplegar un endpoint de Vertex AI con el contenedor xDiT mediante el SDK de Vertex AI para Python.
Ve a la página de la tarjeta de modelo y haz clic en Abrir cuaderno.
Selecciona el cuaderno Vertex Serving. El cuaderno se abrirá en Colab Enterprise.
Ejecuta el cuaderno para desplegar un modelo mediante el contenedor xDiT y envía solicitudes de predicción al endpoint. El fragmento de código de 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 serie de argumentos de servidor que se pueden configurar para optimizar el rendimiento en casos prácticos específicos. Estos argumentos se definen como variables de entorno durante la implementación. A continuación se indican los argumentos clave que puede que tengas que configurar:
Configuración del modelo
MODEL_ID(cadena): especifica el identificador del modelo que se va a cargar. Debe coincidir con el nombre del modelo en tu registro o ruta.
Argumentos de optimización del tiempo de ejecución
N_GPUS(entero): especifica el número de GPUs que se van a usar para la inferencia. El valor predeterminado es 1.WARMUP_STEPS(entero): número de pasos de calentamiento necesarios antes de que empiece la inferencia. Esto es especialmente importante cuando PipeFusion está habilitado para asegurar 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 2048 píxeles) paralelizando el componente VAE en los dispositivos. De esta forma, se evitan problemas de falta de memoria en imágenes grandes. El valor predeterminado es false.USE_TORCH_COMPILE(booleano): habilita la aceleración de una sola GPU mediante torch.compile, lo que proporciona optimizaciones a nivel de kernel para mejorar el rendimiento. El valor predeterminado es false.USE_ONEDIFF(booleano): habilita la tecnología de aceleración de compilación de OneDiff para optimizar la velocidad de ejecución del kernel de la GPU. El valor predeterminado es false.
Argumentos de paralelismo de datos
DATA_PARALLEL_DEGREE(entero): define el grado de paralelismo de los datos. Déjelo en blanco para inhabilitar esta opción o defina el grado de paralelismo seleccionado.USE_CFG_PARALLEL(booleano): habilita la computación paralela para la guía sin clasificador (CFG), también conocida como Split Batch. Si se habilita, el grado de paralelismo constante es 2. Se define como true cuando se usa CFG para controlar el estilo y el contenido de la salida. El valor predeterminado es false.
Argumentos de paralelismo de secuencia (USP - Paralelismo de secuencia unificado)
ULYSSES_DEGREE(entero): define el grado de Ulysses para el enfoque paralelo de secuencia unificada, que combina DeepSpeed-Ulysses y Ring-Attention. De esta forma, se controla el patrón de comunicación entre todos. Si quieres aplicar la opción predeterminada, no indiques ningún valor.RING_DEGREE(entero): define el grado de anillo para la comunicación entre iguales en el paralelismo de secuencias. Funciona junto con ULYSSES_DEGREE para formar la malla de proceso 2D. Si quieres aplicar la opción predeterminada, no indiques ningún valor.
Argumentos de paralelismo de tensor
TENSOR_PARALLEL_DEGREE(entero): define el grado de paralelismo de los tensores, que divide los parámetros del modelo entre los dispositivos a lo largo de las dimensiones de las características para reducir los costes de memoria por dispositivo. Si quieres inhabilitar esta opción, no indiques ningún valor.SPLIT_SCHEME(cadena): define cómo dividir los tensores del modelo entre los dispositivos (por ejemplo, por encabezados de atención o dimensiones ocultas). Si quieres usar el esquema de división predeterminado, no indiques ningún valor.
Argumentos distribuidos de Ray
USE_RAY(booleano): habilita el framework de ejecución distribuida de Ray para escalar cálculos en varios nodos. El valor predeterminado es false.RAY_WORLD_SIZE(entero): número total de procesos del clúster de Ray. El valor predeterminado es 1.VAE_PARALLEL_SIZE(entero): número de procesos dedicados al procesamiento paralelo de VAE cuando se usa Ray. El valor predeterminado es 0.DIT_PARALLEL_SIZE(entero): número de procesos dedicados al procesamiento en paralelo de la red troncal de DiT cuando se usa Ray. El valor predeterminado es 0.
Argumentos paralelos de PipeFusion
PIPEFUSION_PARALLEL_DEGREE(entero): define el grado de paralelismo de PipeFusion, un paralelismo de la canalización a nivel de 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(entero): número de parches en los que se dividirá la secuencia para el procesamiento de la canalización. Si quieres que se determine automáticamente, deja el campo en blanco.ATTN_LAYER_NUM_FOR_PP(cadena): especifica qué capas de atención se deben usar para el paralelismo de la canalización. Se pueden separar por comas (por ejemplo, "10,9") o separados por espacios (por ejemplo, "10 9"). Déjalo en blanco 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 se usa, lo que reduce el uso de la memoria de la GPU a costa de aumentar la latencia. El valor predeterminado es false.ENABLE_SEQUENTIAL_CPU_OFFLOAD(booleano): descarga secuencialmente las capas del modelo a la CPU durante el pase hacia delante, lo que permite la inferencia de modelos más grandes que la memoria de la GPU. El valor predeterminado es false.ENABLE_TILING(booleano): reduce el uso de memoria de la GPU decodificando el componente VAE por partes. Este argumento es útil para imágenes o vídeos más grandes y para evitar errores de falta de memoria. El valor predeterminado es false.ENABLE_SLICING(booleano): reduce el uso de memoria de la GPU dividiendo el tensor de entrada en segmentos para la decodificación de VAE. El valor predeterminado es false.
Argumentos de DiTFastAttn (optimización de la atención)
USE_FAST_ATTN(booleano): habilita la aceleración de DiTFastAttn para la inferencia de una sola GPU. Utiliza la reducción temporal de entrada para reducir la complejidad computacional. El valor predeterminado es false.N_CALIB(entero): número 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(entero): tamaño de la ventana de Window Attention with Residual Caching para reducir la redundancia espacial. El valor predeterminado es 64.COCO_PATH(cadena): ruta al conjunto de datos COCO para la calibración de DiTFastAttn. Obligatorio cuando USE_FAST_ATTN es true. Déjalo en blanco si no lo usas.
Argumentos de optimización de la caché
USE_CACHE(booleano): habilita mecanismos generales de almacenamiento en caché para reducir los cálculos redundantes. El valor predeterminado es false.USE_TEACACHE(booleano): habilita el método de optimización TeaCache para almacenar en caché resultados intermedios. El valor predeterminado es false.USE_FBCACHE(booleano): habilita el método de optimización de la caché del primer bloque. El valor predeterminado es false.
Argumentos de optimización de precisión
USE_FP8_T5_ENCODER(booleano): habilita la precisión FP8 (coma 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 false.
Personalizaciones de publicación
Model Garden proporciona configuraciones de paralelización xDiT predeterminadas para los modelos admitidos. Puedes consultar estos ajustes 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 al implementar el modelo. En el siguiente ejemplo se muestra cómo modificar RING_DEGREE y ULYSSES_DEGREE en una configuración de 2 GPUs para cambiar el enfoque del 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
)
Consulta la sección "Understanding xDiT specific arguments" (Información sobre los argumentos específicos de xDiT) para obtener más información sobre cada variable de entorno. Asegúrate de que el producto de los grados de paralelismo (por ejemplo, PIPEFUSION_PARALLEL_DEGREE, ULYSSES_DEGREE, RING_DEGREE y USE_CFG_PARALLEL) es igual al número total de GPUs (N_GPUS).
Para ver más ejemplos de cómo publicar recetas y configuraciones para 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.