Pubblica modelli Diffusion Transformer utilizzando il contenitore xDiT su GPU Cloud

xDiT è una libreria open source che accelera l'inferenza per i modelli Diffusion Transformer (DiT) utilizzando tecniche di parallelismo e ottimizzazione. Queste tecniche consentono una configurazione multi-GPU scalabile per i carichi di lavoro impegnativi. Questa pagina dimostra come eseguire il deployment di modelli DiT utilizzando xDiT e GPU Cloud su Vertex AI.

Per ulteriori informazioni su xDiT, consulta il progetto GitHub di xDiT.

Vantaggi

Il seguente elenco descrive i vantaggi principali dell'utilizzo di xDiT per pubblicare modelli DiT su Vertex AI:

  • Generazione fino a tre volte più rapida: genera immagini e video ad alta risoluzione in una frazione del tempo rispetto ad altre soluzioni di pubblicazione.
  • Supporto multi-GPU scalabile: distribuisci in modo efficiente i carichi di lavoro su più GPU per prestazioni ottimali.
    • Parallelismo ibrido: xDiT supporta vari approcci di elaborazione parallela, come il parallelismo di sequenze unificato, PipeFusion, il parallelismo CFG e il parallelismo dei dati. Questi metodi possono essere combinati in una ricetta unica per ottimizzare il rendimento.
  • Prestazioni ottimizzate con una singola GPU: xDiT fornisce un'inferenza più rapida anche su una singola GPU.
    • Accelerazione GPU: xDiT incorpora diversi metodi di accelerazione del kernel e utilizza tecniche di DiTFastAttn per velocizzare l'inferenza su una singola GPU.
  • Deployment facile: inizia rapidamente con il deployment con un solo clic o con i notebook Colab Enterprise in Vertex AI Model Garden.

Modelli supportati

xDiT è disponibile per alcune architetture di modelli DiT in Vertex AI Model Garden, come Flux.1 Schnell e CogVideoX-2b. Per sapere se un modello DiT supporta xDiT in Model Garden, visualizza la relativa scheda in Model Garden.

Parallelismo ibrido per le prestazioni multi-GPU:

xDiT utilizza una combinazione di tecniche di parallelismo per massimizzare le prestazioni su configurazioni multi-GPU. Queste tecniche agiscono congiuntamente per distribuire il carico di lavoro e ottimizzare l'utilizzo delle risorse:

  • Parallelismo di sequenza unificato: questa tecnica suddivide i dati di input (ad esempio la suddivisione di un'immagine in patch) su più GPU, riducendo l'utilizzo della memoria e migliorando la scalabilità.
  • PipeFusion: PipeFusion suddivide il modello DiT in fasi e assegna ogni fase a una GPU diversa, consentendo l'elaborazione in parallelo di parti diverse del modello.
  • Parallello CFG: questa tecnica ottimizza in modo specifico i modelli utilizzando indicazioni senza classificatori, un metodo comune per controllare lo stile e i contenuti delle immagini generate. Esegue in parallelo il calcolo dei rami condizionali e incondizionali, il che consente di eseguire l'inferenza più rapidamente.
  • Parallismo di dati: questo metodo replica l'intero modello su ogni GPU, con ogni GPU che elabora un batch diverso di dati di input, aumentando il throughput complessivo del sistema.

Per ulteriori informazioni sui miglioramenti del rendimento, consulta il report di xDiT su Flux.1 Schnell o CogVideoX-2b. Google è riuscita a riprodurre questi risultati in Model Garden di Vertex AI.

Accelerazione con una singola GPU

La libreria xDiT offre vantaggi per il servizio con una singola GPU utilizzando torch.compile e onediff per migliorare la velocità di runtime sulle GPU. Queste tecniche possono essere utilizzate anche in combinazione con il parallelismo ibrido.

xDiT dispone anche di una tecnica di calcolo dell'attenzione efficiente, chiamata DiTFastAttn, per risolvere il collo di bottiglia computazionale del DiT. Per il momento, questa tecnica è disponibile solo per le configurazioni con una singola GPU o in combinazione con il parallelismo dei dati.

Iniziare a utilizzare Model Garden

Il contenitore di pubblicazione GPU Cloud ottimizzato per xDiT è fornito all'interno di Model Garden di Vertex AI. Per i modelli supportati, i deployment utilizzano questo contenitore quando utilizzi i deployment con un solo clic o gli esempi di notebook di Colab Enterprise.

Gli esempi riportati di seguito utilizzano il modello Flux.1-schnell per dimostrare come eseguire il deployment di un modello DiT in un contenitore xDiT.

Utilizzare il deployment con un solo clic

Puoi eseguire il deployment di un endpoint Vertex AI personalizzato con il container xDiT utilizzando una scheda del modello.

  1. Vai alla pagina della scheda del modello e fai clic su Esegui il deployment.

  2. Per la variante del modello da utilizzare, seleziona un tipo di macchina da utilizzare per il deployment.

  3. Fai clic su Esegui il deployment per avviare la procedura di deployment. Riceverai due email: una quando il modello viene caricato e un'altra quando l'endpoint è pronto.

Utilizzare il notebook Colab Enterprise

Per flessibilità e personalizzazione, utilizza gli esempi di notebook di Colab Enterprise per eseguire il deployment di un endpoint Vertex AI con il contenitore xDiT utilizzando l'SDK Vertex AI per Python.

  1. Vai alla pagina della scheda del modello e fai clic su Apri notebook.

  2. Seleziona il notebook Vertex Serving. Il notebook viene aperto in Colab Enterprise.

  3. Esegui il blocco note per eseguire il deployment di un modello utilizzando il contenitore xDiT e invia richieste di previsione all'endpoint. Lo snippet di codice per il deployment è il seguente:

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

Variabili di ambiente

  • MODEL_ID: specifica l'ID del modello DiT da implementare, ad esempio 'black-forest-labs/FLUX.1-schnell'.
  • TASK: definisce l'attività svolta dal modello, ad esempio "text-to-image-flux-xdit".
  • N_GPUS: imposta il numero di GPU da utilizzare per l'inferenza.
  • ULYSSES_DEGREE, RING_DEGREE, PIPEFUSION_PARALLEL_DEGREE: controlla le tecniche di parallelismo utilizzate da xDiT. Per informazioni dettagliate su ciascun argomento, consulta Argomenti xDiT.
  • USE_TORCH_COMPILE: attiva l'accelerazione su una singola GPU utilizzando torch.compile.

Argomenti xDiT

xDiT offre una serie di argomenti del server che possono essere configurati per ottimizzare le prestazioni per casi d'uso specifici. Questi argomenti vengono impostati come variabili di ambiente durante il deployment. Di seguito sono riportati gli argomenti chiave che potresti dover configurare:

  • N_GPUS (numero intero): specifica il numero di GPU da utilizzare per l'inferenza. Il valore predefinito è 1.
  • ENABLE_TILING (booleano): riduce l'utilizzo della memoria GPU decodificando il componente VAE uno riquadro alla volta. Questo argomento è utile per immagini o video di dimensioni maggiori e per evitare errori di esaurimento della memoria. Il valore predefinito è false.
  • ENABLE_SLICING (booleano): riduce l'utilizzo della memoria GPU dividendo il tensore di input in sezioni per la decodifica VAE. Il valore predefinito è false.
  • USE_TORCH_COMPILE (booleano): attiva l'accelerazione GPU singola tramite torch.compile, migliorando la velocità di compilazione. Il valore predefinito è false.
  • PIPEFUSION_PARALLEL_DEGREE (numero intero): imposta il grado di parallelismo per PipeFusion. Valori più elevati aumentano il parallelismo, ma potrebbero richiedere più memoria. Il valore predefinito è 1.
  • WARMUP_STEPS (intero): se PipeFusion è abilitato, questo argomento specifica il numero di passaggi di riscaldamento necessari prima dell'inizio dell'inferenza. Il valore predefinito è 0.
  • ULYSSES_DEGREE (intero): imposta il grado di Ulysses. Il valore predefinito è 1.
  • RING_DEGREE (intero): imposta il grado dell'anello. Il valore predefinito è 1.
  • USE_CFG_PARALLEL (booleano): attiva il calcolo parallelo per le indicazioni senza classificatore (CFG), una tecnica utilizzata per controllare l'output dei modelli DiT. Se è attivato, il grado di parallelismo costante è 2. Imposta su "true" quando utilizzi CFG. Il valore predefinito è false.
  • USE_PARALLEL_VAE (booleano): consente un'elaborazione efficiente di immagini ad alta risoluzione (maggiori di 2048 pixel) parallelizzando il componente VAE. Il valore predefinito è false.

Per un elenco completo degli argomenti, consulta la classe xFuserArgs nel progetto GitHub xDiT.