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) relativi alla 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.
Configura la TPU
Indipendentemente dallo stack scelto, per poter pubblicare un modello devi avere un progetto Google Cloud con una TPU single-host. 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 l'erogazione. Nell'API Cloud TPU, le quote specifiche per la distribuzione per 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 altre versioni di TPU, i workload di addestramento e di distribuzione utilizzano la stessa quota. Per saperne di più, 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 l'avvio della VM TPU, 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 gestione 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, quindi 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'erogazione 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 della TPU. Per saperne di più sulla profilazione, consulta le seguenti risorse: