Esegui l'inferenza su Cloud TPU
La distribuzione è il processo di deployment di un modello di machine learning addestrato in un ambiente di produzione per eseguire l'inferenza. TPU v5e e versioni successive supportano l'inferenza. Il rispetto degli obiettivi del livello di servizio (SLO) di latenza è una priorità per i workload di gestione.
Questo documento illustra l'erogazione di un modello su una TPU a host singolo. Le slice TPU con otto o meno chip hanno una VM TPU o un host e sono chiamate TPU a host singolo. Per informazioni sull'inferenza multi-host con PyTorch, consulta vLLM-TorchTPU e SGLang-TorchTPU. Per informazioni sull'inferenza multi-host con JAX, consulta Esegui l'inferenza multi-host utilizzando Pathways.
Configura la TPU
Indipendentemente dallo stack scelto, prima di poter pubblicare un modello devi avere un progetto Google Cloud con una TPU a host singolo. I seguenti passaggi si applicano a ogni combinazione di libreria e framework:
- Crea un progetto. Per utilizzare Cloud TPU, devi disporre di un account e di un progetto Google Cloud . Per saperne di più, consulta Configura un progettoGoogle Cloud per le TPU.
Controlla la quota. Assicurati di disporre di una quota sufficiente per il numero di core TPU che prevedi di utilizzare per l'inferenza. TPU v5e utilizza quote separate per l'addestramento e la pubblicazione. Nell'API Cloud TPU, le quote specifiche per la gestione di TPU v5e sono le seguenti:
- Risorse v5e on demand:
TPU v5 lite pod cores for serving per project per zone - Risorse v5e preemptible:
Preemptible TPU v5 lite pod cores for serving per project per zone
Per le altre versioni di TPU, i workload di addestramento e di distribuzione utilizzano la stessa quota. Per maggiori informazioni, consulta Quote di Cloud TPU.
- Risorse v5e on demand:
Crea una VM TPU. Scegli una versione e una zona TPU con capacità di gestione, quindi crea una VM TPU single-host. Per saperne di più, consulta Regioni e zone TPU e Gestisci le risorse TPU.
Dopo che la VM TPU è in esecuzione, segui la guida alla configurazione della libreria e del framework in Scegli uno stack di inferenza.
Scegli uno stack di inferenza
Cloud TPU supporta l'inferenza di modelli linguistici di grandi dimensioni (LLM) tramite due librerie di distribuzione open source: vLLM e SGLang. Entrambe le librerie vengono eseguite su Cloud TPU con un backend JAX o PyTorch, quindi puoi scegliere la libreria più adatta al tuo workload indipendentemente dal framework in cui è scritto il modello.
Scegli prima la libreria, poi seleziona il backend che corrisponde al framework del tuo modello, come mostrato nella tabella seguente:
| Libreria di pubblicazione | Backend JAX | Backend PyTorch |
|---|---|---|
| vLLM |
vllm-jax
|
vllm-torchtpu
|
| SGLang |
sglang-jax
|
sglang-torchtpu
|
Inferenza con vLLM
vLLM è una libreria open source per l'inferenza e l'erogazione rapide di LLM, basata su batching continuo e PagedAttention. Cloud TPU funziona con vLLM tramite il plug-in vllm-jax per i modelli JAX e il plug-in vllm-torchtpu per i modelli PyTorch.
Per ulteriori esempi di utilizzo di vLLM per l'erogazione di un modello su TPU, consulta le seguenti risorse:
- Inizia a utilizzare vLLM TPU
- Pubblica un LLM utilizzando TPU Trillium su Google Kubernetes Engine con vLLM
- Ricette per l'utilizzo di vLLM su TPU Trillium (v6e)
vLLM con JAX
La documentazione di vLLM TPU illustra l'installazione, una guida rapida, i modelli supportati e le opzioni di ottimizzazione. Per saperne di più, consulta Pubblica modelli JAX con vLLM su TPU.
vLLM con PyTorch
Il plug-in vllm-torchtpu aggiunge il supporto del kernel e del runtime TorchTPU a vLLM. Per saperne di più, consulta Pubblica modelli PyTorch con vLLM su TPU.
Inferenza con SGLang
SGLang è un framework di distribuzione open source per modelli linguistici di grandi dimensioni e multimodali che combina un runtime ad alte prestazioni con un frontend di generazione strutturato. Su Cloud TPU, SGLang viene eseguito tramite il backend SGLang-JAX (sglang-jax) per i modelli JAX e il backend SGLang-TorchTPU (sglang-torchtpu) per i modelli PyTorch.
SGLang con PyTorch
Per informazioni sull'installazione, sull'avvio del server e sulle configurazioni supportate, consulta Pubblica modelli PyTorch con il runtime SGLang-TorchTPU su TPU.
SGLang con JAX
Per informazioni sull'installazione, sull'avvio del server e sulle configurazioni supportate, consulta Pubblica modelli JAX con il runtime SGLang-JAX su TPU.
Profila le prestazioni di inferenza
Dopo aver configurato l'inferenza, puoi utilizzare i profiler per analizzare le prestazioni e l'utilizzo delle TPU. Per saperne di più sulla profilazione, consulta le seguenti risorse: