Gestisci LLM utilizzando TPU multi-host su GKE con JetStream e Pathways

Questa guida mostra come gestire modelli linguistici di grandi dimensioni (LLM) all'avanguardia, come Llama 3.1 405B, su Google Kubernetes Engine (GKE) utilizzando unità di elaborazione tensoriale (TPU) su più nodi.

Questa guida mostra come utilizzare tecnologie open source portatili, ovvero Kubernetes, JetStream, Pathways on Cloud e l'API LeaderWorkerSet (LWS), per eseguire il deployment e gestire carichi di lavoro di AI/ML su GKE, sfruttando il controllo granulare, la scalabilità, la resilienza, la portabilità e la convenienza di GKE.

Sfondo

I modelli linguistici di grandi dimensioni sono cresciuti e non rientrano più in una singola fetta di TPU host. Per l'inferenza ML, puoi utilizzare Pathways on Cloud per eseguire l'inferenza multi-host su larga scala su GKE su più nodi TPU interconnessi. In questa guida, imparerai a eseguire il provisioning di un cluster GKE con gli slice TPU multi-host e a utilizzare i binari di Pathways on Cloud, avviare il server JetStream con il framework MaxText ed effettuare richieste di inferenza multi-host.

Se gestisci un LLM utilizzando le TPU su GKE con JetStream, MaxText e Pathways, 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:

  1. Scopri di più sulla disponibilità della versione attuale della TPU con l'architettura di sistema di Cloud TPU.
  2. Scopri di più sulle TPU in GKE.

Questo tutorial illustra l'erogazione del modello Llama 3.1-405B. GKE esegue il deployment del modello su nodi TPU v6e multi-host con topologie TPU configurate in base ai requisiti del modello per la gestione dei prompt con bassa latenza.

Pathways on Cloud

Pathways è un livello di orchestrazione su larga scala per gli acceleratori. Pathways è progettato esplicitamente per consentire l'esplorazione di nuovi sistemi e idee di ricerca ML, mantenendo al contempo prestazioni all'avanguardia per i modelli attuali. Pathways consente a un singolo processo client JAX di coordinare il calcolo su una o più sezioni TPU di grandi dimensioni, semplificando i calcoli ML che interessano centinaia o migliaia di chip TPU.

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 di ottimizzare le prestazioni con PyTorch/XLA e JAX TPU.

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 ulteriori informazioni sui modelli e sulle dimensioni dei parametri più recenti supportati da MaxText, consulta il repository del progetto MaxText.

Llama 3.1 405B

Llama 3.1 405B è un modello linguistico di grandi dimensioni di Meta progettato per una serie di attività di elaborazione del linguaggio naturale, tra cui generazione di testo, traduzione e risposta a domande. GKE offre l'infrastruttura necessaria per supportare le esigenze di addestramento e gestione distribuita di modelli di questa scala.

Per saperne di più, consulta la documentazione di Llama.

Architettura

Questa sezione descrive l'architettura GKE utilizzata in questo tutorial. L'architettura include un cluster GKE Standard che esegue il provisioning delle TPU e ospita i componenti JetStream e Pathways per eseguire il deployment e l'hosting del modello.

Il seguente diagramma mostra i componenti di questa architettura:

Architettura del cluster GKE con pool di nodi TPU multi-host contenente i componenti JetStream e Pathways.

Questa architettura include i seguenti componenti:

  • Un cluster regionale GKE Standard.
  • Un node pool di sezioni TPU multi-host che ospita i componenti di deployment di JetStream e Pathways.
  • Pathways resource manager gestisce le risorse dell'acceleratore e coordina l'allocazione degli acceleratori per i job utente.
  • Pathways client si coordina con Pathways resource manager per determinare dove vengono inseriti i programmi compilati per l'esecuzione.
  • Pathways worker viene eseguito ed esegue i calcoli sulle macchine con acceleratore e invia i dati al carico di lavoro tramite il server proxy IFRT.
  • IFRT proxy client implementa l'API Interim Framework Runtime (IFRT) OSS e funge da ponte di comunicazione tra il tuo workload e i componenti di Pathways.
  • Il IFRT proxy server riceve le richieste dal IFRT proxy client e le inoltra al Pathways client, distribuendo il lavoro.
  • Il container JetStream-Pathways fornisce un server di inferenza basato su JAX che riceve le richieste di inferenza e delega i processi di esecuzione a Pathways workers
  • Il componente Service distribuisce il traffico in entrata a tutte le repliche JetStream HTTP.
  • JetStream HTTP è un server HTTP che accetta le richieste come wrapper per il formato richiesto di JetStream e le invia al client GRPC di JetStream.