xDiT è una libreria open source che accelera l'inferenza per i modelli Diffusion Transformer (DiT) mediante tecniche di parallelismo e ottimizzazione. Queste tecniche consentono una configurazione scalabile su più GPU per workload impegnativi. Questa pagina mostra come eseguire il deployment dei modelli DiT utilizzando xDiT e le GPU Cloud su Vertex AI.
Per saperne di più su xDiT, consulta il progetto GitHub di xDiT.
Vantaggi
Il seguente elenco descrive i principali vantaggi dell'utilizzo di xDiT per pubblicare modelli DiT su Vertex AI:
- Generazione fino a tre volte più veloce: genera immagini e video ad alta risoluzione in una frazione di 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 sequenza 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 per una singola GPU: xDiT offre 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 semplice: inizia rapidamente con il deployment con un solo clic o con i blocchi note 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, CogVideoX-2b e le varianti del modello da testo a video Wan2.1. Per verificare se un modello DiT supporta xDiT in Model Garden, visualizza la scheda del modello in Model Garden.
Parallelismo ibrido per prestazioni multi-GPU:
xDiT utilizza una combinazione di tecniche di parallelismo per massimizzare le prestazioni nelle configurazioni multi-GPU. Queste tecniche agiscono insieme per distribuire il carico di lavoro e ottimizzare l'utilizzo delle risorse:
- Parallelismo della sequenza unificato: questa tecnica suddivide i dati di input (ad esempio, suddivide un'immagine in patch) su più GPU, riducendo l'utilizzo della memoria e migliorando la scalabilità.
- PipeFusion: PipeFusion divide il modello DiT in fasi e assegna ogni fase a una GPU diversa, consentendo l'elaborazione parallela di diverse parti del modello.
- Parallelismo CFG: questa tecnica ottimizza in modo specifico i modelli utilizzando la guida senza classificatore, un metodo comune per controllare lo stile e il contenuto delle immagini generate. Parallelizza il calcolo dei rami condizionali e incondizionati, il che porta a un'inferenza più rapida.
- Parallelismo dei dati: questo metodo replica l'intero modello su ogni GPU, con ogni GPU che elabora un batch diverso di dati di input, aumentando la velocità effettiva complessiva del sistema.
Per ulteriori informazioni sui miglioramenti delle prestazioni, consulta il report di xDiT su Flux.1 Schnell o CogVideoX-2b. Google è riuscita a riprodurre questi risultati in Vertex AI Model Garden.
Accelerazione di una singola GPU
La libreria xDiT offre vantaggi per la gestione di 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 di DiT. Per il momento, questa tecnica è disponibile solo per configurazioni con una sola GPU o in combinazione con il parallelismo dei dati.
Inizia a utilizzare Model Garden
Il container di erogazione delle GPU Cloud ottimizzato per xDiT è fornito all'interno di Vertex AI Model Garden. Per i modelli supportati, i deployment utilizzano questo container quando utilizzi i deployment con un clic o gli esempi di notebook Colab Enterprise.
Gli esempi riportati di seguito utilizzano il modello Flux.1-schnell per mostrare come eseguire il deployment di un modello DiT su un container 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.
Vai alla pagina della scheda del modello e fai clic su Esegui il deployment.
Per la variante del modello da utilizzare, seleziona un tipo di macchina da utilizzare per il deployment.
Fai clic su Esegui il deployment per iniziare la procedura di deployment. Riceverai due notifiche via 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 Colab Enterprise per eseguire il deployment di un endpoint Vertex AI con il contenitore xDiT utilizzando l'SDK Vertex AI per Python.
Vai alla pagina della scheda del modello e fai clic su Apri blocco note.
Seleziona il notebook Vertex Serving. Il notebook viene aperto in Colab Enterprise.
Esegui il notebook per eseguire il deployment di un modello utilizzando il contenitore xDiT e inviare richieste di previsione all'endpoint. Lo snippet di codice per il deployment è il seguente:
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()
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. Il seguente elenco riporta gli argomenti chiave che potresti dover configurare:
Configurazione del modello
MODEL_ID(stringa): specifica l'identificatore del modello da caricare. Deve corrispondere al nome del modello nel registro o nel percorso.
Argomenti di ottimizzazione del runtime
N_GPUS(numero intero): specifica il numero di GPU da utilizzare per l'inferenza. Il valore predefinito è 1.WARMUP_STEPS(numero intero): numero di passaggi di riscaldamento necessari prima dell'inizio dell'inferenza. Ciò è particolarmente importante quando PipeFusion è abilitato per garantire un rendimento stabile. Il valore predefinito è 1.USE_PARALLEL_VAE(booleano): consente l'elaborazione efficiente di immagini ad alta risoluzione (superiori a 2048 pixel) parallelizzando il componente VAE su più dispositivi. In questo modo si evitano problemi di esaurimento della memoria per le immagini di grandi dimensioni. Il valore predefinito è false.USE_TORCH_COMPILE(booleano): attiva l'accelerazione di una singola GPU tramite torch.compile, fornendo ottimizzazioni a livello di kernel per migliorare le prestazioni. Il valore predefinito è false.USE_ONEDIFF(booleano): attiva la tecnologia di accelerazione della compilazione OneDiff per ottimizzare la velocità di esecuzione del kernel della GPU. Il valore predefinito è false.
Argomenti di parallelismo dei dati
DATA_PARALLEL_DEGREE(numero intero): imposta il grado di parallelismo dei dati. Lascia vuoto per disattivare o impostare il grado di parallelismo selezionato.USE_CFG_PARALLEL(booleano): attiva il calcolo parallelo per la guida senza classificatore (CFG), nota anche come Split Batch. Se abilitato, il grado di parallelismo costante è 2. Imposta su true quando utilizzi CFG per controllare lo stile e i contenuti dell'output. Il valore predefinito è false.
Argomenti di parallelismo della sequenza (USP - Unified Sequence Parallelism)
ULYSSES_DEGREE(numero intero): imposta il grado di Ulysses per l'approccio parallelo alla sequenza unificata, che combina DeepSpeed-Ulysses e Ring-Attention. Questo controlla il pattern di comunicazione all-to-all. Lascia vuoto per utilizzare il valore predefinito.RING_DEGREE(numero intero): imposta il grado di anello per la comunicazione peer-to-peer nel parallelismo di sequenza. Funziona in combinazione con ULYSSES_DEGREE per formare la mesh di processo 2D. Lascia vuoto per utilizzare il valore predefinito.
Argomenti di parallelismo dei tensori
TENSOR_PARALLEL_DEGREE(numero intero): imposta il grado di parallelismo dei tensori, che divide i parametri del modello tra i dispositivi lungo le dimensioni delle funzionalità per ridurre i costi di memoria per dispositivo. Lascia il campo vuoto per disattivare la funzionalità.SPLIT_SCHEME(stringa): definisce come dividere i tensori del modello tra i dispositivi (ad es. in base alle teste di attenzione, alle dimensioni nascoste). Lascia vuoto per lo schema di suddivisione predefinito.
Argomenti distribuiti di Ray
USE_RAY(booleano): abilita il framework di esecuzione distribuita Ray per scalare i calcoli su più nodi. Il valore predefinito è false.RAY_WORLD_SIZE(intero): numero totale di processi nel cluster Ray. Il valore predefinito è 1.VAE_PARALLEL_SIZE(numero intero): numero di processi dedicati all'elaborazione parallela VAE quando utilizzi Ray. Il valore predefinito è 0.DIT_PARALLEL_SIZE(numero intero): numero di processi dedicati all'elaborazione parallela del backbone DiT quando utilizzi Ray. Il valore predefinito è 0.
Argomenti paralleli di PipeFusion
PIPEFUSION_PARALLEL_DEGREE(numero intero): imposta il grado di parallelismo per PipeFusion, un parallelismo della pipeline a livello di sequenza che sfrutta le caratteristiche di ridondanza temporale dell'input dei modelli di diffusione. Valori più alti aumentano il parallelismo, ma richiedono più memoria. Il valore predefinito è 1.NUM_PIPELINE_PATCH(integer): Number of patches to split the sequence into for pipeline processing. Lascia il campo vuoto per la determinazione automatica.ATTN_LAYER_NUM_FOR_PP(stringa): specifica i livelli di attenzione da utilizzare per il parallelismo della pipeline. Possono essere separati da virgole (ad es. "10,9") o separati da spazi (ad es. "10 9"). Lascia il campo vuoto per utilizzare tutti i livelli.
Argomenti di ottimizzazione della memoria
ENABLE_MODEL_CPU_OFFLOAD(booleano): scarica i pesi del modello nella memoria della CPU quando non sono in uso, riducendo l'utilizzo della memoria GPU a costo di una maggiore latenza. Il valore predefinito è false.ENABLE_SEQUENTIAL_CPU_OFFLOAD(booleano): esegue l'offload sequenziale dei livelli del modello alla CPU durante il forward pass, consentendo l'inferenza di modelli più grandi della memoria della GPU. Il valore predefinito è false.ENABLE_TILING(booleano): riduce l'utilizzo della memoria GPU decodificando il componente VAE una tessera alla volta. Questo argomento è utile per immagini o video più grandi e per evitare errori di memoria insufficiente. 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.
Argomenti di DiTFastAttn (ottimizzazione dell'attenzione)
USE_FAST_ATTN(booleano): attiva l'accelerazione DiTFastAttn per l'inferenza su una singola GPU, utilizzando la riduzione temporale dell'input per ridurre la complessità computazionale. Il valore predefinito è false.N_CALIB(integer): numero di campioni di calibrazione per l'ottimizzazione di DiTFastAttn. Il valore predefinito è 8.THRESHOLD(float): soglia di similarità per la riduzione della similarità temporale in DiTFastAttn. Il valore predefinito è 0,5.WINDOW_SIZE(numero intero): dimensione della finestra per l'attenzione alla finestra con memorizzazione nella cache dei residui per ridurre la ridondanza spaziale. Il valore predefinito è 64.COCO_PATH(stringa): percorso del set di dati COCO per la calibrazione di DiTFastAttn. Obbligatorio quando USE_FAST_ATTN è true. Lascia vuoto se non lo utilizzi.
Argomenti di ottimizzazione della cache
USE_CACHE(booleano): attiva i meccanismi di memorizzazione nella cache generali per ridurre i calcoli ridondanti. Il valore predefinito è false.USE_TEACACHE(booleano): attiva il metodo di ottimizzazione TeaCache per la memorizzazione nella cache dei risultati intermedi. Il valore predefinito è false.USE_FBCACHE(booleano): attiva il metodo di ottimizzazione della cache del primo blocco. Il valore predefinito è false.
Argomenti di ottimizzazione di precisione
USE_FP8_T5_ENCODER(booleano): attiva la precisione FP8 (virgola mobile a 8 bit) per il codificatore di testo T5, riducendo l'utilizzo di memoria e potenzialmente migliorando il throughput con un impatto minimo sulla qualità. Il valore predefinito è false.
Personalizzazioni della pubblicazione
Model Garden fornisce configurazioni di parallelizzazione xDiT predefinite per i modelli supportati. Puoi esaminare queste impostazioni predefinite utilizzando l'SDK Vertex AI per Python.
Per visualizzare la configurazione di deployment predefinita per un modello, ad esempio "black-forest-labs/FLUX.1-schnell", puoi eseguire il seguente snippet di codice:
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"
# }
# ..........]
Il metodo list_deploy_options() restituisce le specifiche del container, incluse le variabili di ambiente (env) che definiscono la configurazione xDiT.
Per personalizzare la strategia di parallelismo, puoi eseguire l'override di queste variabili di ambiente durante il deployment del modello. L'esempio seguente mostra come modificare RING_DEGREE e ULYSSES_DEGREE per una configurazione a 2 GPU, modificando l'approccio al parallelismo:
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
)
Ricorda di consultare la sezione "Informazioni sugli argomenti specifici di xDiT" per i dettagli su ogni variabile di ambiente. Assicurati che il prodotto dei gradi di parallelismo (ad es. PIPEFUSION_PARALLEL_DEGREE,ULYSSES_DEGREE, RING_DEGREE e USE_CFG_PARALLEL) è uguale al numero totale di GPU (N_GPUS).
Per altri esempi di serving di ricette e configurazioni per modelli diversi, consulta la documentazione ufficiale di xDiT. Per ulteriori informazioni sull'SDK Model Garden, consulta la documentazione.