Questo tutorial mostra come pubblicare un modello linguistico di grandi dimensioni (LLM) Gemma utilizzando le unità di elaborazione tensoriale (TPU) su Google Kubernetes Engine (GKE). Esegui il deployment di un container predefinito con JetStream e MaxText in GKE. Inoltre, configuri GKE per caricare i pesi di Gemma 7B da Cloud Storage in fase di runtime.
Questo tutorial è rivolto a ingegneri di machine learning (ML), amministratori e operatori di piattaforme e a specialisti di dati e AI interessati a utilizzare le funzionalità di orchestrazione dei container Kubernetes per il servizio di LLM. Per scoprire di più sui ruoli comuni e sulle attività di esempio a cui facciamo riferimento nei contenuti diGoogle Cloud , consulta Ruoli e attività comuni degli utenti GKE.
Prima di leggere questa pagina, assicurati di avere familiarità con quanto segue:
- Modalità Autopilot e Standard
- Disponibilità della versione TPU attuale con l'architettura di sistema Cloud TPU
- TPU in GKE
Sfondo
Questa sezione descrive le tecnologie chiave utilizzate in questo tutorial.
Gemma
Gemma è un insieme di modelli di intelligenza artificiale (AI) generativa, leggeri e disponibili pubblicamente, rilasciati con una licenza aperta. Questi modelli di AI sono disponibili per l'esecuzione nelle tue applicazioni, hardware, dispositivi mobili o servizi ospitati. Puoi utilizzare i modelli Gemma per la generazione di testo, ma puoi anche ottimizzarli per attività specializzate.
Per saperne di più, consulta la documentazione di Gemma.
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.
Questo tutorial illustra l'erogazione del modello Gemma 7B. 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.
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. Fornisce ottimizzazioni avanzate delle prestazioni, tra cui tecniche di batching continuo e quantizzazione, per facilitare il deployment degli LLM. JetStream consente a PyTorch/XLA e JAX TPU di ottenere prestazioni ottimali.
Per scoprire di più su queste ottimizzazioni, consulta i repository dei progetti JetStream PyTorch e JetStream MaxText.
MaxText
MaxText è un'implementazione LLM JAX performante, scalabile e adattabile, basata su librerie JAX open source come Flax, Orbax e Optax. L'implementazione LLM solo decodificatore di MaxText è scritta in Python. Sfrutta in modo intensivo il compilatore XLA per ottenere prestazioni elevate senza dover creare kernel personalizzati.
Per scoprire di più sui modelli e sulle dimensioni dei parametri più recenti supportati da MaxText, consulta il repository del progetto MaxText.