Disponibilizar modelos de Transformer de difusão usando o contêiner xDiT em GPUs do Cloud

O xDiT é uma biblioteca de código aberto que acelera a inferência de modelos Diffusion Transformer (DiT) usando técnicas de paralelismo e otimização. Essas técnicas permitem uma configuração escalonável de várias GPUs para cargas de trabalho exigentes. Nesta página, mostramos como implantar modelos DiT usando o xDiT e as GPUs do Cloud no Vertex AI.

Para mais informações sobre o xDiT, consulte o projeto xDiT no GitHub (em inglês).

Vantagens

A lista a seguir descreve os principais benefícios de usar o xDiT para veicular modelos DiT na Vertex AI:

  • Geração até três vezes mais rápida: gere imagens e vídeos de alta resolução em uma fração do tempo em comparação com outras soluções de veiculação.
  • Suporte escalonável a várias GPUs: distribua cargas de trabalho de maneira eficiente em várias GPUs para um desempenho ideal.
    • Paralelismo híbrido: o xDiT oferece suporte a várias abordagens de processamento paralelo, como paralelismo de sequência unificado, PipeFusion, paralelismo de CFG e paralelismo de dados. Esses métodos podem ser combinados em uma receita exclusiva para otimizar o desempenho.
  • Desempenho otimizado de GPU única: o xDiT oferece inferência mais rápida, mesmo em uma única GPU.
    • Aceleração de GPU: o xDiT incorpora vários métodos de aceleração de kernel e usa técnicas do DiTFastAttn para acelerar a inferência em uma única GPU.
  • Implantação fácil: comece rapidamente com a implantação em um clique ou notebooks do Colab Enterprise no Model Garden da Vertex AI.

Modelos compatíveis

O xDiT está disponível para algumas arquiteturas de modelo DiT no Model Garden da Vertex AI, como Flux.1 Schnell, CogVideoX-2b e variantes do modelo de texto para vídeo Wan2.1. Para saber se um modelo DiT é compatível com xDiT no Model Garden, consulte o card dele em Model Garden.

Paralelismo híbrido para desempenho de várias GPUs:

O xDiT usa uma combinação de técnicas de paralelismo para maximizar a performance em configurações com várias GPUs. Essas técnicas trabalham juntas para distribuir a carga de trabalho e otimizar a utilização de recursos:

  • Paralelismo de sequência unificado: essa técnica divide os dados de entrada (como dividir uma imagem em patches) em várias GPUs, reduzindo o uso de memória e melhorando a escalonabilidade.
  • PipeFusion: o PipeFusion divide o modelo DiT em estágios e atribui cada um deles a uma GPU diferente, permitindo o processamento paralelo de diferentes partes do modelo.
  • Paralelismo CFG: essa técnica otimiza modelos usando a orientação sem classificador, um método comum para controlar o estilo e o conteúdo das imagens geradas. Ele paraleliza a computação de ramificações condicionais e incondicionais, resultando em uma inferência mais rápida.
  • Paralelismo de dados: esse método replica todo o modelo em cada GPU, e cada GPU processa um lote diferente de dados de entrada, aumentando a taxa de transferência geral do sistema.

Para mais informações sobre melhorias de performance, consulte o relatório da xDiT sobre o Flux.1 Schnell ou o CogVideoX-2b. O Google conseguiu reproduzir esses resultados no Model Garden da Vertex AI.

Aceleração de GPU única

A biblioteca xDiT oferece benefícios para a exibição de GPU única usando torch.compile e onediff para aumentar a velocidade de execução em GPUs. Essas técnicas também podem ser usadas em conjunto com o paralelismo híbrido.

O xDiT também tem uma técnica eficiente de computação de atenção, chamada DiTFastAttn, para resolver o gargalo computacional do DiT. Por enquanto, essa técnica só está disponível para configurações de GPU única ou em conjunto com o paralelismo de dados.

Introdução ao Model Garden

O contêiner de exibição de GPU do Cloud otimizado para xDiT é fornecido no Model Garden da Vertex AI. Para modelos compatíveis, as implantações usam esse contêiner quando você usa implantações com um clique ou os exemplos de notebooks do Colab Enterprise.

Os exemplos a seguir usam o modelo Flux.1-schnell para demonstrar como implantar um modelo DiT em um contêiner xDiT.

Usar a implantação com um clique

É possível implantar um endpoint personalizado da Vertex AI com o contêiner xDiT usando um card de modelo.

  1. Acesse a página do card de modelo e clique em Implantar.

  2. Para a variação de modelo a ser usada, selecione um tipo de máquina para sua implantação.

  3. Clique em Implantar para iniciar o processo. Você recebe duas notificações por e-mail: uma quando o modelo é carregado e outra quando o endpoint está pronto.

Usar o bloco do Colab Enterprise

Para flexibilidade e personalização, use os exemplos de notebook do Colab Enterprise para implantar um endpoint da Vertex AI com o contêiner xDiT usando o SDK da Vertex AI para Python.

  1. Acesse a página do card de modelo e clique em Abrir notebook.

  2. Selecione o notebook da Vertex Serving. O notebook está aberto no Colab Enterprise

  3. Execute o notebook para implantar um modelo usando o contêiner xDiT e envie solicitações de previsão ao endpoint. O snippet de código da implantação é este:

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 do xDiT

O xDiT oferece uma variedade de argumentos de servidor que podem ser configurados para otimizar a performance em casos de uso específicos. Esses argumentos são definidos como variáveis de ambiente durante a implantação. Confira abaixo os principais argumentos que você pode precisar configurar:

Configuração do modelo
  • MODEL_ID (string): especifica o identificador do modelo a ser carregado. Ele precisa corresponder ao nome do modelo no seu registro ou caminho.
Argumentos de otimização de ambiente de execução
  • N_GPUS (número inteiro): especifica o número de GPUs a serem usadas para inferência. O valor padrão é 1.
  • WARMUP_STEPS (inteiro): número de etapas de pré-aquecimento necessárias antes do início da inferência. Isso é especialmente importante quando o PipeFusion está ativado para garantir uma performance estável. O valor padrão é 1.
  • USE_PARALLEL_VAE (booleano): permite o processamento eficiente de imagens de alta resolução (mais de 2.048 pixels) ao paralelizar o componente VAE em dispositivos. Isso evita problemas de falta de memória para imagens grandes. O valor padrão é falso.
  • USE_TORCH_COMPILE (booleano): ativa a aceleração de GPU única usando torch.compile, oferecendo otimizações no nível do kernel para melhorar o desempenho. O valor padrão é falso.
  • USE_ONEDIFF (booleano): ativa a tecnologia de aceleração de compilação do OneDiff para otimizar a velocidade de execução do kernel da GPU. O valor padrão é falso.
Argumentos de paralelismo de dados
  • DATA_PARALLEL_DEGREE (número inteiro): define o grau de paralelismo de dados. Deixe em branco para desativar ou defina o grau de paralelismo selecionado.
  • USE_CFG_PARALLEL (booleano): ativa a computação paralela para a orientação sem classificador (CFG), também conhecida como Split Batch. Quando ativado, o grau de paralelismo constante é 2. Defina como "true" ao usar CFG para controlar o estilo e o conteúdo da saída. O valor padrão é falso.
Argumentos de paralelismo de sequência (USP: Unified Sequence Parallelism)
  • ULYSSES_DEGREE (inteiro): define o grau de Ulysses para a abordagem paralela de sequência unificada, que combina DeepSpeed-Ulysses e Ring-Attention. Isso controla o padrão de comunicação de todos para todos. Deixe em branco para usar o padrão.
  • RING_DEGREE (número inteiro): define o grau do anel para comunicação ponto a ponto no paralelismo de sequência. Funciona em conjunto com ULYSSES_DEGREE para formar a malha de processo 2D. Deixe em branco para usar o padrão.
Argumentos de paralelismo de tensor
  • TENSOR_PARALLEL_DEGREE (inteiro): define o grau de paralelismo de tensor, que divide os parâmetros do modelo entre dispositivos ao longo das dimensões de recursos para reduzir os custos de memória por dispositivo. Deixe em branco para desativar.
  • SPLIT_SCHEME (string): define como dividir os tensores do modelo entre dispositivos (por exemplo, por cabeças de atenção, dimensões ocultas). Deixe em branco para usar o esquema de divisão padrão.
Argumentos distribuídos do Ray
  • USE_RAY (booleano): ativa o framework de execução distribuída do Ray para escalonar computações em vários nós. O valor padrão é falso.
  • RAY_WORLD_SIZE (número inteiro): número total de processos no cluster do Ray. O valor padrão é 1.
  • VAE_PARALLEL_SIZE (inteiro): número de processos dedicados ao processamento paralelo de VAE ao usar o Ray. O valor padrão é 0.
  • DIT_PARALLEL_SIZE (inteiro): número de processos dedicados ao processamento paralelo de backbones do DiT ao usar o Ray. O valor padrão é 0.
Argumentos paralelos do PipeFusion
  • PIPEFUSION_PARALLEL_DEGREE (inteiro): define o grau de paralelismo para o PipeFusion, um paralelismo de pipeline no nível da sequência que aproveita as características de redundância temporal de entrada dos modelos de difusão. Valores mais altos aumentam o paralelismo, mas exigem mais memória. O valor padrão é 1.
  • NUM_PIPELINE_PATCH (número inteiro): número de patches para dividir a sequência no processamento do pipeline. Deixe em branco para determinação automática.
  • ATTN_LAYER_NUM_FOR_PP (string): especifica quais camadas de atenção usar para o paralelismo de pipeline. Podem ser separados por vírgulas (por exemplo, "10,9") ou separados por espaços (por exemplo, "10 9"). Deixe em branco para usar todas as camadas.
Argumentos de otimização de memória
  • ENABLE_MODEL_CPU_OFFLOAD (booleano): descarrega os pesos do modelo para a memória da CPU quando não estão em uso, reduzindo o uso da memória da GPU ao custo de maior latência. O valor padrão é "false".
  • ENABLE_SEQUENTIAL_CPU_OFFLOAD (booleano): descarrega sequencialmente as camadas do modelo para a CPU durante a transmissão direta, permitindo a inferência de modelos maiores que a memória da GPU. O valor padrão é falso.
  • ENABLE_TILING (booleano): reduz o uso de memória da GPU decodificando o componente VAE um bloco por vez. Esse argumento é útil para imagens ou vídeos maiores e para evitar erros de falta de memória. O valor padrão é falso.
  • ENABLE_SLICING (booleano): reduz o uso da memória da GPU dividindo o tensor de entrada em slices para decodificação de VAE. O valor padrão é falso.
Argumentos DiTFastAttn (otimização de atenção)
  • USE_FAST_ATTN (booleano): ativa a aceleração DiTFastAttn para inferência de GPU única, usando a redução temporal de entrada para reduzir a complexidade computacional. O valor padrão é falso.
  • N_CALIB (número inteiro): número de amostras de calibragem para otimização do DiTFastAttn. O valor padrão é 8.
  • THRESHOLD (ponto flutuante): limite de similaridade para redução de similaridade temporal em DiTFastAttn. O valor padrão é 0,5.
  • WINDOW_SIZE (inteiro): tamanho da janela para atenção da janela com cache residual para reduzir a redundância espacial. O valor padrão é 64.
  • COCO_PATH (string): caminho para o conjunto de dados COCO para calibragem do DiTFastAttn. Obrigatório quando USE_FAST_ATTN é verdadeiro. Deixe em branco se não estiver usando.
Argumentos de otimização de cache
  • USE_CACHE (booleano): ativa mecanismos gerais de cache para reduzir cálculos redundantes. O valor padrão é falso.
  • USE_TEACACHE (booleano): ativa o método de otimização do TeaCache para armazenar em cache resultados intermediários. O valor padrão é falso.
  • USE_FBCACHE (booleano): ativa o método de otimização First-Block-Cache. O valor padrão é "false".
Argumentos de otimização de precisão
  • USE_FP8_T5_ENCODER (booleano): ativa a precisão FP8 (ponto flutuante de 8 bits) para o codificador de texto T5, reduzindo o uso de memória e possivelmente melhorando a capacidade de processamento com impacto mínimo na qualidade. O valor padrão é falso.

Personalizações de veiculação

O Model Garden oferece configurações padrão de paralelização xDiT para modelos compatíveis. É possível inspecionar essas configurações padrão usando o SDK da Vertex AI para Python.

Para conferir a configuração de implantação padrão de um modelo, como "black-forest-labs/FLUX.1-schnell", execute o seguinte snippet 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"
#  }
# ..........]

O método list_deploy_options() retorna as especificações do contêiner, incluindo as variáveis de ambiente (env) que definem a configuração do xDiT.

Para personalizar a estratégia de paralelismo, substitua essas variáveis de ambiente ao implantar o modelo. O exemplo a seguir demonstra como modificar RING_DEGREE e ULYSSES_DEGREE para uma configuração de duas GPUs, mudando a abordagem 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
)

Consulte a seção "Entender argumentos específicos do xDiT" para saber mais sobre cada variável de ambiente. Verifique se o produto dos graus de paralelismo (por exemplo, PIPEFUSION_PARALLEL_DEGREE,ULYSSES_DEGREE, RING_DEGREE e USE_CFG_PARALLEL) é igual ao número total de GPUs (N_GPUS).

Para mais exemplos de veiculação de receitas e configurações para diferentes modelos, consulte a documentação oficial do xDiT. Para mais informações sobre o SDK do Model Garden, consulte a documentação.