Questa guida mostra come gestire un modello linguistico di grandi dimensioni (LLM) utilizzando le Tensor Processing Unit (TPU) su Google Kubernetes Engine (GKE) con JetStream tramite PyTorch. In questa guida scaricherai i pesi del modello in Cloud Storage ed eseguirai il deployment su un cluster GKE Autopilot o Standard utilizzando un container che esegue JetStream.
Se hai bisogno della scalabilità, della resilienza e dell'efficacia in termini di costi offerti dalle funzionalità di Kubernetes durante il deployment del modello su JetStream, questa guida è un buon punto di partenza.
Questa guida è rivolta ai clienti dell'AI generativa che utilizzano PyTorch, agli utenti nuovi o esistenti di GKE, agli ingegneri ML, agli ingegneri MLOps (DevOps) o agli amministratori della piattaforma interessati a utilizzare le funzionalità di orchestrazione dei container Kubernetes per la gestione degli LLM.
Sfondo
Se gestisci un LLM utilizzando le TPU su GKE con JetStream, puoi creare una soluzione di gestione solida e pronta per la produzione con tutti i vantaggi di Kubernetes gestito, tra cui efficienza dei costi, scalabilità e maggiore disponibilità. Questa sezione descrive le tecnologie chiave utilizzate in questo tutorial.
Informazioni sulle TPU
Le TPU sono circuiti integrati specifici per le applicazioni (ASIC) sviluppati da Google e utilizzati per accelerare i modelli di machine learning e AI creati utilizzando framework come TensorFlow, PyTorch e JAX.
Prima di utilizzare le TPU in GKE, ti consigliamo di completare il seguente percorso di apprendimento:
- Scopri di più sulla disponibilità della versione attuale della TPU con l'architettura di sistema di Cloud TPU.
- Scopri di più sulle TPU in GKE.
Questo tutorial illustra l'erogazione di vari modelli LLM. GKE esegue il deployment del modello su nodi TPUv5e a singolo host con topologie TPU configurate in base ai requisiti del modello per la gestione dei prompt con bassa latenza.
Informazioni su JetStream
JetStream è un framework open source per la pubblicazione di inferenze sviluppato da Google. JetStream consente un'inferenza ad alte prestazioni, ad alta velocità effettiva e ottimizzato per la memoria su TPU e GPU. JetStream offre ottimizzazioni avanzate delle prestazioni, tra cui batching continuo, ottimizzazioni della cache KV e tecniche di quantizzazione, per facilitare l'implementazione degli LLM. JetStream consente l'erogazione di PyTorch/XLA e JAX TPU per ottenere prestazioni ottimali.
Raggruppamento continuo
Il batch continuo è una tecnica che raggruppa dinamicamente le richieste di inferenza in entrata in batch, riducendo la latenza e aumentando la velocità effettiva.
Quantizzazione della cache KV
La quantizzazione della cache KV prevede la compressione della cache chiave-valore utilizzata nei meccanismi di attenzione, riducendo i requisiti di memoria.
Quantizzazione dei pesi Int8
La quantizzazione dei pesi Int8 riduce la precisione dei pesi del modello da virgola mobile a 32 bit a numeri interi a 8 bit, il che comporta un calcolo più rapido e un utilizzo della memoria ridotto.
Per scoprire di più su queste ottimizzazioni, consulta i repository dei progetti JetStream PyTorch e JetStream MaxText.
Informazioni su PyTorch
PyTorch è un framework di machine learning open source sviluppato da Meta e ora parte della Linux Foundation. PyTorch fornisce funzionalità di alto livello come il calcolo dei tensori e le reti neurali profonde.