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 GPU de Cloud en la plataforma de agentes de Gemini Enterprise.
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 publicar modelos de DiT en Agent Platform:
- 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 con un solo clic o los notebooks de Colab Enterprise en Model Garden de la Agent Platform de Gemini Enterprise.
Modelos compatibles
xDiT está disponible para ciertas arquitecturas de modelos DiT en Model Garden de Agent Platform de Gemini Enterprise, 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 utilizando 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 la capacidad de procesamiento 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 Model Garden de Gemini Enterprise Agent Platform.
Aceleración de GPU única
La biblioteca xDiT proporciona beneficios para la entrega con una sola GPU, ya que usa 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 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 combinación con el paralelismo de datos.
Primeros pasos con Model Garden
El contenedor de entrega de GPU de Cloud optimizado para xDiT se proporciona en Model Garden de la plataforma de agentes de Gemini Enterprise. 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 Agent Platform 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
Para obtener flexibilidad y personalización, usa los ejemplos de notebook de Colab Enterprise para implementar un extremo de Agent Platform con el contenedor xDiT a través del SDK de Agent Platform 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 todos 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 paralelismo de datos
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 se habilita, el grado de paralelismo constante es 2. Se establece en 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 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 de igual a igual 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 tensor
TENSOR_PARALLEL_DEGREE(número entero): Establece el grado de paralelismo de tensor, que divide los parámetros del modelo en los dispositivos a lo largo de las dimensiones de las características para reducir los costos de memoria por dispositivo. Deja el campo 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 neuronal 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): Es la cantidad de parches en los que se divide 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 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 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 la 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 predeterminadas de paralelización de xDiT para los modelos admitidos. Puedes inspeccionar estos parámetros de configuración predeterminados con el SDK de Agent Platform 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) sea 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.