Entrega modelos de Transformer de difusión con el contenedor xDiT en GPUs de Cloud

xDiT es una biblioteca de código abierto que acelera la inferencia de los modelos de transformador de difusión (DiT) mediante el uso de técnicas de paralelismo y optimización. Estas técnicas permiten una configuración escalable de varias GPUs 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 beneficios clave de usar xDiT para publicar modelos de DiT en Vertex AI:

  • Generación hasta tres veces más rápida: Genera imágenes y videos de alta resolución en una fracción del tiempo en comparación con otras soluciones de publicación.
  • Compatibilidad escalable con varias GPUs: Distribuye cargas de trabajo de manera eficiente en varias GPUs para obtener un rendimiento óptimo.
    • Paralelismo híbrido: xDiT admite varios enfoques de procesamiento en 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 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 de un clic o los notebooks de Colab Enterprise en Model Garden de Vertex AI.

Modelos compatibles

xDiT está disponible para ciertas arquitecturas de modelos de DiT en Vertex AI Model Garden, como Flux.1 Schnell y CogVideoX-2b. 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 GPU:

xDiT usa una combinación de técnicas de paralelismo para maximizar el rendimiento en configuraciones de varias GPUs. Estas técnicas funcionan en conjunto para distribuir la carga de trabajo y optimizar el uso de recursos:

  • Paralelismo de secuencias unificado: Esta técnica divide los datos de entrada (como dividir una imagen en parches) en varias GPUs, lo que reduce el uso de la memoria y mejora la escalabilidad.
  • PipeFusion: PipeFusion divide el modelo de DiT en etapas y asigna cada etapa a una GPU diferente, lo que permite el procesamiento en paralelo de diferentes partes del modelo.
  • Paralelismo de CFG: Esta técnica optimiza los modelos de forma específica mediante la guía sin clasificadores, un método común para controlar el estilo y el contenido de las imágenes generadas. Paraleliza el cálculo de 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 Vertex AI Model Garden.

Aceleración de una sola GPU

La biblioteca de xDiT proporciona beneficios para la publicación de una sola GPU con torch.compile y onediff para mejorar la velocidad del entorno 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 procesamiento de atención eficiente, llamada DiTFastAttn, para abordar el cuello de botella computacional de DiT. Por ahora, 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 de xDiT se proporciona en Model Garden de Vertex AI. En el caso de los modelos compatibles, las implementaciones usan este contenedor cuando usas 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 de DiT en un contenedor de xDiT.

Usa la implementación con un solo clic

Puedes implementar un extremo personalizado de Vertex AI con el contenedor de xDiT mediante una tarjeta de modelo.

  1. Navega a la página de la tarjeta de modelo y haz clic en Implementar.

  2. Para la variación del modelo que deseas usar, selecciona un tipo de máquina que se usará para la implementación.

  3. Haz clic en Implementar para comenzar el proceso de implementación. Recibirás dos notificaciones por correo electrónico; una cuando se sube 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 notebooks de Colab Enterprise para implementar un extremo de Vertex AI con el contenedor xDiT mediante el SDK de Vertex AI para Python.

  1. Navega a la página de la tarjeta de modelo y haz clic en Abrir notebook.

  2. Selecciona el notebook de Vertex Serving. Se abre el notebook en Colab Enterprise.

  3. Ejecuta el notebook para implementar un modelo con el contenedor xDiT y envía solicitudes de predicción al extremo. El fragmento de código para la implementación es el siguiente:

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,
)

Variables de entorno

  • MODEL_ID: Especifica el ID del modelo de DiT que se implementará, como "black-forest-labs/FLUX.1-schnell".
  • TASK: Define la tarea que realiza el modelo, como "text-to-image-flux-xdit".
  • N_GPUS: Establece la cantidad de GPUs que se usarán para la inferencia.
  • ULYSSES_DEGREE, RING_DEGREE, PIPEFUSION_PARALLEL_DEGREE: Controlan las técnicas de paralelismo que usa xDiT. Para obtener detalles sobre cada argumento, consulta Argumentos de xDiT.
  • USE_TORCH_COMPILE: Habilita la aceleración de una sola GPU con torch.compile.

Argumentos de xDiT

xDiT ofrece una variedad de argumentos de servidor que se pueden configurar para optimizar el rendimiento en casos de uso específicos. Estos argumentos se establecen como variables de entorno durante la implementación. En la siguiente lista, se muestran los argumentos clave que podrías necesitar configurar:

  • N_GPUS (número entero): Especifica la cantidad de GPUs que se usarán para la inferencia. El valor predeterminado es 1.
  • ENABLE_TILING (booleano): Reduce el uso de memoria de la GPU decodificando el componente de VAE una tarjeta a la vez. Este argumento es útil para imágenes o videos 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 rebanadas para la decodificación de VAE. El valor predeterminado es false.
  • USE_TORCH_COMPILE (booleano): Habilita la aceleración de una sola GPU a través de torch.compile, lo que mejora la velocidad de compilación. El valor predeterminado es false.
  • PIPEFUSION_PARALLEL_DEGREE (número entero): Establece el grado de paralelismo para PipeFusion. Los valores más altos aumentan el paralelismo, pero pueden requerir más memoria. El valor predeterminado es 1.
  • WARMUP_STEPS (número entero): Si PipeFusion está habilitado, este argumento especifica la cantidad de pasos de preparación necesarios antes de que comience la inferencia. El valor predeterminado es 0.
  • ULYSSES_DEGREE (número entero): Establece el grado de Ulysses. El valor predeterminado es 1.
  • RING_DEGREE (número entero): Establece el grado de anillo. El valor predeterminado es 1.
  • USE_CFG_PARALLEL (booleano): Habilita el procesamiento en paralelo para la guía sin clasificador (CFG), una técnica que se usa para controlar el resultado de los modelos de DiT. Cuando se habilita, el grado de paralelismo constante es 2. Configúralo como “true” cuando uses CFG. El valor predeterminado es false.
  • USE_PARALLEL_VAE (booleano): Habilita el procesamiento eficiente de imágenes de alta resolución (más de 2,048 píxeles) mediante la paralelización del componente VAE. El valor predeterminado es false.

Para obtener una lista completa de los argumentos, consulta la clase xFuserArgs en el proyecto de GitHub de xDiT.