Publique modelos Diffusion Transformer com o contentor xDiT em GPUs na nuvem

A xDiT é uma biblioteca de código aberto que acelera a inferência para modelos de transformadores de difusão (DiT) através de técnicas de paralelismo e otimização. Estas técnicas permitem uma configuração com várias GPUs escalável para cargas de trabalho exigentes. Esta página demonstra como implementar modelos DiT usando o xDiT e GPUs na nuvem na Vertex AI.

Para mais informações sobre o xDiT, consulte o projeto do GitHub do xDiT.

Vantagens

A lista seguinte descreve as principais vantagens da utilização do xDiT para publicar modelos DiT no Vertex AI:

  • Geração até três vezes mais rápida: gere imagens e vídeos de alta resolução em muito menos tempo do que com outras soluções de publicação.
  • Suporte multi-GPU escalável: distribua eficientemente as cargas de trabalho por várias GPUs para um desempenho ideal.
    • Paralelismo híbrido: o xDiT suporta várias abordagens de processamento paralelo, como o paralelismo de sequência unificado, o PipeFusion, o paralelismo de CFG e o paralelismo de dados. Estes métodos podem ser combinados numa receita única para otimizar o desempenho.
  • Desempenho otimizado de uma única GPU: o xDiT oferece uma inferência mais rápida, mesmo numa única GPU.
    • Aceleração da GPU: o xDiT incorpora vários métodos de aceleração do kernel e usa técnicas do DiTFastAttn para acelerar a inferência numa única GPU.
  • Implementação fácil: comece rapidamente com a implementação com um clique ou os blocos de notas do Colab Enterprise no Model Garden do Vertex AI.

Modelos suportados

O xDiT está disponível para determinadas arquiteturas de modelos DiT no Vertex AI Model Garden, como o Flux.1 Schnell, o CogVideoX-2b e as variantes do modelo de texto para vídeo Wan2.1. Para ver se um modelo DiT suporta o xDiT no Model Garden, consulte o respetivo cartão de modelo no Model Garden.

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

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

  • Paralelismo de sequências unificado: esta técnica divide os dados de entrada (como dividir uma imagem em patches) em várias GPUs, reduzindo a utilização de memória e melhorando a escalabilidade.
  • PipeFusion: o PipeFusion divide o modelo DiT em fases e atribui cada fase a uma GPU diferente, o que permite o processamento paralelo de diferentes partes do modelo.
  • Paralelismo CFG: esta técnica otimiza especificamente os modelos através de orientações sem classificador, um método comum para controlar o estilo e o conteúdo das imagens geradas. Paraleliza a computação de ramificações condicionais e incondicionais, o que leva a uma inferência mais rápida.
  • Paralelismo de dados: este método replica o modelo completo em cada GPU, com cada GPU a processar um lote diferente de dados de entrada, o que aumenta o débito geral do sistema.

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

Aceleração de GPU única

A biblioteca xDiT oferece vantagens para a publicação com uma única GPU através da utilização de torch.compile e onediff para melhorar a velocidade de execução nas GPUs. Estas técnicas também podem ser usadas em conjunto com o paralelismo híbrido.

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

Comece a usar o Model Garden

O contentor de publicação de GPUs na nuvem otimizado para xDiT é fornecido no Vertex AI Model Garden. Para modelos suportados, as implementações usam este contentor quando usa implementações com um clique ou os exemplos de blocos de notas do Colab Enterprise.

Os exemplos seguintes usam o modelo Flux.1-schnell para demonstrar como implementar um modelo DiT num contentor xDiT.

Use a implementação com um clique

Pode implementar um ponto final do Vertex AI personalizado com o contentor xDiT através de um cartão de modelo.

  1. Navegue para a página do cartão do modelo e clique em Implementar.

  2. Para a variação do modelo a usar, selecione um tipo de máquina a usar para a sua implementação.

  3. Clique em Implementar para iniciar o processo de implementação. Recebe duas notificações por email: uma quando o modelo é carregado e outra quando o ponto final está pronto.

Use o bloco de notas do Colab Enterprise

Para flexibilidade e personalização, use os exemplos de blocos de notas do Colab Enterprise para implementar um ponto final da Vertex AI com o contentor xDiT através do SDK da Vertex AI para Python.

  1. Navegue para a página do cartão do modelo e clique em Abrir bloco de notas.

  2. Selecione o notebook Vertex Serving. O bloco de notas é aberto no Colab Enterprise.

  3. Execute o bloco de notas para implementar um modelo através do contentor xDiT e envie pedidos de previsão para o ponto final. O fragmento do código para a implementação é o seguinte:

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 xDiT

A xDiT oferece uma variedade de argumentos de servidor que podem ser configurados para otimizar o desempenho para exemplos de utilização específicos. Estes argumentos são definidos como variáveis de ambiente durante a implementação. Segue-se uma lista dos principais argumentos que pode ter de configurar:

Configuração do modelo
  • MODEL_ID (string): especifica o identificador do modelo a carregar. Deve corresponder ao nome do modelo no seu registo ou caminho.
Argumentos de otimização do tempo de execução
  • N_GPUS (número inteiro): especifica o número de GPUs a usar para a inferência. O valor predefinido é 1.
  • WARMUP_STEPS (integer): número de passos de preparação necessários antes do início da inferência. Isto é particularmente importante quando o PipeFusion está ativado para garantir um desempenho estável. O valor predefinido é 1.
  • USE_PARALLEL_VAE (booleano): permite o processamento eficiente de imagens de alta resolução (superiores a 2048 píxeis) paralelizando o componente VAE em vários dispositivos. Isto evita problemas de falta de memória para imagens grandes. O valor predefinido é false.
  • USE_TORCH_COMPILE (booleano): ativa a aceleração de GPU única através do torch.compile, oferecendo otimizações ao nível do kernel para um desempenho melhorado. O valor predefinido é false.
  • USE_ONEDIFF (booleano): ativa a tecnologia de aceleração de compilação OneDiff para otimizar a velocidade de execução do kernel da GPU. O valor predefinido é false.
Argumentos paralelos de dados
  • DATA_PARALLEL_DEGREE (integer): define o grau de paralelismo de dados. Deixe vazio 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 divisão de lotes. Quando está ativado, o grau de paralelismo constante é 2. Definido como verdadeiro quando usa CFG para controlar o estilo e o conteúdo de saída. O valor predefinido é false.
Argumentos paralelos de sequência (USP – Paralelismo de sequência unificado)
  • ULYSSES_DEGREE (integer): define o grau de Ulysses para a abordagem paralela da sequência unificada, que combina o DeepSpeed-Ulysses e o Ring-Attention. Isto controla o padrão de comunicação de todos para todos. Deixe este campo vazio para usar a predefinição.
  • RING_DEGREE (integer): define o grau de anel para a comunicação ponto a ponto no paralelismo de sequência. Funciona em conjunto com ULYSSES_DEGREE para formar a malha de processamento 2D. Deixe este campo vazio para usar a predefinição.
Argumentos de paralelismo de tensores
  • TENSOR_PARALLEL_DEGREE (integer): define o grau de paralelismo de tensores, que divide os parâmetros do modelo entre dispositivos ao longo das dimensões das caraterísticas 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çalhos de atenção, dimensões ocultas). Deixe vazio para o esquema de divisão predefinido.
Argumentos distribuídos do Ray
  • USE_RAY (booleano): ativa a framework de execução distribuída do Ray para dimensionar os cálculos em vários nós. O valor predefinido é false.
  • RAY_WORLD_SIZE (integer): número total de processos no cluster do Ray. O valor predefinido é 1.
  • VAE_PARALLEL_SIZE (integer): número de processos dedicados ao processamento paralelo de VAE quando usa o Ray. O valor predefinido é 0.
  • DIT_PARALLEL_SIZE (integer): Number of processes dedicated to DiT backbone parallel processing when using Ray. O valor predefinido é 0.
Argumentos paralelos do PipeFusion
  • PIPEFUSION_PARALLEL_DEGREE (integer): define o grau de paralelismo para o PipeFusion, um paralelismo de pipeline ao nível da sequência que tira partido das características de redundância temporal de entrada dos modelos de difusão. Os valores mais elevados aumentam o paralelismo, mas requerem mais memória. O valor predefinido é 1.
  • NUM_PIPELINE_PATCH (integer): número de patches em que dividir a sequência para o processamento de pipeline. Deixe o campo vazio para uma determinação automática.
  • ATTN_LAYER_NUM_FOR_PP (string): especifica que 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 este campo vazio para usar todas as camadas.
Argumentos de otimização de memória
  • ENABLE_MODEL_CPU_OFFLOAD (booleano): transfere os pesos do modelo para a memória da CPU quando não estão a ser usados, reduzindo a utilização da memória da GPU à custa de um aumento da latência. O valor predefinido é false.
  • ENABLE_SEQUENTIAL_CPU_OFFLOAD (booleano): transfere sequencialmente as camadas do modelo para a CPU durante a passagem direta, permitindo a inferência de modelos maiores do que a memória da GPU. O valor predefinido é false.
  • ENABLE_TILING (booleano): reduz a utilização de memória da GPU descodificando o componente VAE um mosaico de cada vez. Este argumento é útil para imagens ou vídeos maiores e para evitar erros de falta de memória. O valor predefinido é false.
  • ENABLE_SLICING (booleano): reduz a utilização de memória da GPU dividindo o tensor de entrada em fatias para a descodificação de VAE. O valor predefinido é false.
Argumentos DiTFastAttn (otimização da atenção)
  • USE_FAST_ATTN (booleano): ativa a aceleração DiTFastAttn para a inferência de uma única GPU, usando a redução temporal de entrada para reduzir a complexidade computacional. O valor predefinido é false.
  • N_CALIB (integer): número de amostras de calibração para a otimização DiTFastAttn. O valor predefinido é 8.
  • THRESHOLD (float): limite de semelhança para a redução da semelhança temporal em DiTFastAttn. O valor predefinido é 0,5.
  • WINDOW_SIZE (integer): Window size for Window Attention with Residual Caching to reduce spatial redundancy. O valor predefinido é 64.
  • COCO_PATH (string): caminho para o conjunto de dados COCO para a calibração DiTFastAttn. Obrigatório quando USE_FAST_ATTN é verdadeiro. Deixe este campo vazio se não estiver a usar.
Argumentos de otimização da cache
  • USE_CACHE (booleano): ativa mecanismos de colocação em cache gerais para reduzir os cálculos redundantes. O valor predefinido é false.
  • USE_TEACACHE (booleano): ativa o método de otimização TeaCache para colocar em cache os resultados intermédios. O valor predefinido é false.
  • USE_FBCACHE (booleano): ativa o método de otimização First-Block-Cache. O valor predefinido é 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 a utilização de memória e potencialmente melhorando o débito com um impacto mínimo na qualidade. O valor predefinido é false.

Personalizações de publicação

O Model Garden fornece configurações de paralelização xDiT predefinidas para modelos suportados. Pode inspecionar estas predefinições através do SDK Vertex AI para Python.

Para ver a configuração de implementação predefinida de um modelo, como "black-forest-labs/FLUX.1-schnell", pode executar o seguinte 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"
#  }
# ..........]

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

Para personalizar a estratégia de paralelismo, pode substituir estas variáveis de ambiente ao implementar o modelo. O exemplo seguinte demonstra como modificar o RING_DEGREE e o ULYSSES_DEGREE para uma configuração de 2 GPUs, alterando 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
)

Não se esqueça de consultar a secção "Compreender argumentos específicos do xDiT" para ver detalhes sobre cada variável de ambiente. Certifique-se de que o produto dos graus de paralelismo (por exemplo, PIPEFUSION_PARALLEL_DEGREE,ULYSSES_DEGREE, RING_DEGREE e USE_CFG_PARALLEL) for igual ao número total de GPUs (N_GPUS).

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