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

Standard

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.

Importante: per accedere a Percorsi, contatta il rappresentante dell'account Google Cloud.

Sfondo

I modelli linguistici di grandi dimensioni sono aumentati di dimensioni 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 il batching continuo, l'ottimizzazione 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 pubblicare il modello.

Il seguente diagramma mostra i componenti di questa architettura:

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 OSS Interim Framework Runtime (IFRT) 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.

Prima di iniziare

Sign in to your Google Cloud account. If you're new to Google Cloud, create an account to evaluate how our products perform in real-world scenarios. New customers also get $300 in free credits to run, test, and deploy workloads.

In the Google Cloud console, on the project selector page, select or create a Google Cloud project.

Roles required to select or create a project

• Select a project: Selecting a project doesn't require a specific IAM role—you can select any project that you've been granted a role on.
• Create a project: To create a project, you need the Project Creator (roles/resourcemanager.projectCreator), which contains the resourcemanager.projects.create permission. Learn how to grant roles.

Note: If you don't plan to keep the resources that you create in this procedure, create a project instead of selecting an existing project. After you finish these steps, you can delete the project, removing all resources associated with the project.

Go to project selector

Verify that billing is enabled for your Google Cloud project.

Enable the required API.

Roles required to enable APIs

To enable APIs, you need the Service Usage Admin IAM role (roles/serviceusage.serviceUsageAdmin), which contains the serviceusage.services.enable permission. Learn how to grant roles.

Enable the API

In the Google Cloud console, on the project selector page, select or create a Google Cloud project.

Roles required to select or create a project

• Select a project: Selecting a project doesn't require a specific IAM role—you can select any project that you've been granted a role on.
• Create a project: To create a project, you need the Project Creator (roles/resourcemanager.projectCreator), which contains the resourcemanager.projects.create permission. Learn how to grant roles.

Note: If you don't plan to keep the resources that you create in this procedure, create a project instead of selecting an existing project. After you finish these steps, you can delete the project, removing all resources associated with the project.

Go to project selector

Verify that billing is enabled for your Google Cloud project.

Enable the required API.

Roles required to enable APIs

To enable APIs, you need the Service Usage Admin IAM role (roles/serviceusage.serviceUsageAdmin), which contains the serviceusage.services.enable permission. Learn how to grant roles.

Enable the API

• Make sure that you have the following role or roles on the project: roles/container.admin, roles/iam.serviceAccountAdmin, roles/resourcemanager.projectIamAdmin

  Check for the roles

  1. In the Google Cloud console, go to the IAM page.

     Go to IAM
  2. Select the project.

  3. In the Principal column, find all rows that identify you or a group that you're included in. To learn which groups you're included in, contact your administrator.

  4. For all rows that specify or include you, check the Role column to see whether the list of roles includes the required roles.

  Grant the roles

  1. In the Google Cloud console, go to the IAM page.

     Vai a IAM
  2. Seleziona il progetto.
  3. Fai clic su Concedi l'accesso.

  4. Nel campo Nuove entità, inserisci il tuo identificatore dell'utente. In genere si tratta dell'indirizzo email di un Account Google.

  5. Nell'elenco Seleziona un ruolo, seleziona un ruolo.
  6. Per concedere altri ruoli, fai clic su Aggiungi un altro ruolo e aggiungi ogni ruolo aggiuntivo.
  7. Fai clic su Salva.
  • Assicurati di disporre di una quota sufficiente per 16 chip TPU v6e PodSlice Lite. In questo tutorial utilizzi le istanze on demand.
  • Assicurati che il tuo progetto Google Cloud sia incluso nella lista consentita per Pathways.

  Ottenere l'accesso al modello

  Per accedere al checkpoint Meta Llama 3.1-405B per il deployment su GKE, segui questi passaggi:

  1. Firma il contratto di consenso alla licenza.
  2. Accedi alla pagina di download di Meta Llama.
  3. Leggi e accetta i Termini e condizioni modello e ottieni l'URL necessario per scaricare il modello.
  4. Per scaricare il checkpoint del modello, trova l'ID modello appropriato. Per un elenco dei modelli supportati e dei relativi ID, consulta la documentazione di llama CLI. Ad esempio, utilizza Llama 3.1-405B-Instruct:bf16-mp16 per il modello Llama 3.1-405B.

  Prepara l'ambiente

  In questo tutorial utilizzerai Cloud Shell per gestire le risorse ospitate su Google Cloud. Cloud Shell include il software necessario per questo tutorial, tra cui kubectl e gcloud CLI.

  Per configurare l'ambiente con Cloud Shell:

  1. Nella console Google Cloud , avvia una sessione di Cloud Shell facendo clic su Attiva Cloud Shell nella consoleGoogle Cloud . Viene avviata una sessione nel riquadro inferiore della console Google Cloud .

  2. Imposta le variabili di ambiente predefinite:

     gcloud config set project PROJECT_ID
     gcloud config set billing/quota_project PROJECT_ID
     export PROJECT_ID=$(gcloud config get project)
     export CLUSTER_NAME=CLUSTER_NAME
     export BUCKET_NAME=BUCKET_NAME
     export CONTROL_PLANE_LOCATION=CONTROL_PLANE_LOCATION
     export NODE_LOCATION=NODE_LOCATION
     export CLUSTER_VERSION=CLUSTER_VERSION
     export MACHINE_TYPE=ct6e-standard-4t
     export TPU_TYPE=v6e
     export TOPOLOGY=4x4
     export WORKERS_PER_SLICE=4
     

     Sostituisci i seguenti valori:

     • PROJECT_ID: il tuo Google Cloud ID progetto.
     • CLUSTER_NAME: il nome del tuo cluster GKE.
     • BUCKET_NAME: il nome del tuo bucket Cloud Storage. Non è necessario specificare il prefisso gs://.
     • CONTROL_PLANE_LOCATION: la regione Compute Engine in cui si trovano il cluster GKE, il bucket Cloud Storage e i nodi TPU. La regione contiene zone in cui sono disponibili i tipi di macchine TPU v6e (ad esempio us-east1, us-east5, europe-west4, asia-northeast1 o us-south1).
     • NODE_LOCATION: la zona in cui sono disponibili le risorse TPU (ad esempio, us-east1-d).
     • CLUSTER_VERSION: la versione di GKE, che deve supportare il tipo di macchina che vuoi utilizzare. Tieni presente che la versione GKE predefinita potrebbe non essere disponibile per la TPU di destinazione. Per un elenco delle versioni GKE minime disponibili per tipo di macchina TPU, consulta Disponibilità delle TPU in GKE.
     • MACHINE_TYPE: il tipo di macchina v6e.
     • TPU_TYPE: un prefisso utilizzato per denominare i node pool (v6e).
     • TOPOLOGY: la topologia TPU v6e.
     • WORKERS_PER_SLICE: il numero di nodi per pool di nodi o slice TPU.

  Creare e configurare risorse Google Cloud

  Per creare le risorse richieste, segui queste istruzioni:

  Nota: per utilizzare alcuni acceleratori, potrebbe essere necessario creare una prenotazione di capacità. Per saperne di più su come prenotare e utilizzare le risorse prenotate, consulta Utilizzo di risorse di zona prenotate.

  Crea un cluster GKE

  1. Crea un cluster GKE Standard regionale:

     gcloud container clusters create CLUSTER_NAME \
         --project=PROJECT_ID \
         --cluster-version=CLUSTER_VERSION \
         --location=CONTROL_PLANE_LOCATION \
         --scopes=cloud-platform \
         --machine-type=n2-standard-32
     

     La creazione del cluster potrebbe richiedere diversi minuti.

     Sostituisci CLUSTER_VERSION con la versione del cluster appropriata.

  2. Crea un pool di nodi TPU v6e con una topologia 4x4 e quattro nodi ciascuno:

     gcloud container node-pools create multihost-np \
     --project=PROJECT_ID \
     --location=CONTROL_PLANE_LOCATION \
     --node-locations=NODE_LOCATION \
     --cluster=CLUSTER_NAME \
     --machine-type=MACHINE_TYPE \
     --num-nodes=WORKERS_PER_SLICE \
     --tpu-topology=TOPOLOGY \
     --scopes cloud-platform \
     --placement-type=COMPACT \
     --workload-metadata=GCE_METADATA
     

     Nota: il simbolo --workload-metadata=GCE_METADATA indica che la federazione delle identità per i carichi di lavoro per GKE è disabilitata per il pool di nodi. I node pool che eseguono Pathways con JetStream e MaxText devono disabilitare la federazione delle identità per i carichi di lavoro per GKE.

  Configura un account di servizio per l'accesso agli oggetti Storage

  Configura un account di servizio Kubernetes in modo che funga da account di servizio IAM.

  1. Crea un account di servizio IAM per la tua applicazione:

     gcloud iam service-accounts create jetstream-pathways
     

  2. Aggiungi un'associazione di policy IAM per il tuo account di servizio IAM per gestire Cloud Storage. In questo modo, il account di servizio IAM potrà accedere al bucket di archiviazione in cui verrà archiviato il checkpoint:

     gcloud projects add-iam-policy-binding ${PROJECT} \
       --member "serviceAccount:jetstream-pathways@${PROJECT}.iam.gserviceaccount.com" \
       --role roles/storage.objectUser
     
     gcloud projects add-iam-policy-binding ${PROJECT} \
       --member "serviceAccount:jetstream-pathways@${PROJECT}.iam.gserviceaccount.com" \
       --role roles/storage.insightsCollectorService
     

  3. Annota il account di servizio Kubernetes con l'indirizzo email del account di servizio IAM.

     kubectl annotate serviceaccount default \
     iam.gke.io/gcp-service-account=jetstream-pathways@${PROJECT}.iam.gserviceaccount.com
     

  Configura Docker per l'autenticazione in Artifact Registry

  Configura Docker per l'autenticazione in Artifact Registry in modo che possa eseguire il pull delle immagini di Pathways consentite:

  gcloud auth login
  gcloud auth configure-docker
  

  Conversione del checkpoint

  Per convertire un checkpoint Meta Llama 3.1-405B in un checkpoint di inferenza int8 compatibile con MaxText, completa i passaggi descritti in Conversione del checkpoint con Llama3.1-405B. Il deployment utilizza il checkpoint con il flag load_parameters_path.

  Crea un bucket Cloud Storage per archiviare i file temporanei di Pathways

  Crea un bucket Cloud Storage per archiviare i file temporanei di Pathways, come la cache di compilazione:

  export PATHWAYS_BUCKET=PATHWAYS_BUCKET
  gcloud storage buckets create gs://$PATHWAYS_BUCKET
  

  Eseguire il deployment di JetStream-MaxText e Pathways

  Esegui il deployment del server di modelli JetStream-MaxText e Pathways.

  Connettiti al cluster GKE

  gcloud container clusters get-credentials "${CLUSTER}" --project "${PROJECT}" --location "${ZONE}"
  

  Esegui il deployment dell'API LeaderWorkerSet (LWS)

  LWS è una risorsa personalizzata progettata per il deployment e la gestione di applicazioni distribuite con stato, in particolare quelle con un'architettura leader-worker. È particolarmente adatto per i workload di AI/ML in cui un modello di grandi dimensioni viene suddiviso e pubblicato su più dispositivi su più nodi.

  VERSION=v0.6.1
  kubectl apply --server-side -f https://github.com/kubernetes-sigs/lws/releases/download/$VERSION/manifests.yaml
  

  Attendi che il controller LeaderWorkerSet sia completamente disponibile:

  kubectl wait deploy/lws-controller-manager -n lws-system --for=condition=available --timeout=5m
  

  L'output dovrebbe essere simile al seguente:

  deployment.apps/lws-controller-manager condition met
  

  Verifica che il controller LeaderWorkerSet sia in esecuzione nello spazio dei nomi lws-system:

  kubectl get pod -n lws-system
  

  L'output dovrebbe essere simile al seguente:

  NAME                          READY   STATUS    RESTARTS    AGE
  lws-controller-manager-abcd   1/1     Running   0           40s
  lws-controller-manager-efgh   1/1     Running   0           40s
  

  Esegui il deployment del manifest del workload

  1. Salva il seguente manifest come jetstream-pathways-llama-3-1-405b-4x4.yaml:

     apiVersion: leaderworkerset.x-k8s.io/v1
     kind: LeaderWorkerSet
     metadata:
       name: jetstream-pathways
       annotations:
         leaderworkerset.sigs.k8s.io/exclusive-topology: cloud.google.com/gke-nodepool
     spec:
       replicas: 1
       leaderWorkerTemplate:
         leaderTemplate:
           metadata:
             labels:
               app: jetstream-pathways
           spec:
             nodeSelector:
               cloud.google.com/gke-tpu-accelerator: tpu-v6e-slice
               cloud.google.com/gke-tpu-topology: 4x4
             tolerations:
             - key: "google.com/tpu"
               operator: "Exists"
               effect: "NoSchedule"
             containers:
             - name: pathways-proxy
               image: us-docker.pkg.dev/cloud-tpu-v2-images/pathways/proxy_server:jax-0.5.3
               args:
               imagePullPolicy: Always
               ports:
               - containerPort: 38681
             - name: pathways-rm
               env:
               - name: HOST_ADDRESS
                 value: "$(LWS_LEADER_ADDRESS)"
               - name: TPU_SKIP_MDS_QUERY
                 value: "true"
               image: us-docker.pkg.dev/cloud-tpu-v2-images/pathways/server:jax-0.5.3
               args:
               - --server_port=38677
               - --gcs_scratch_location=PATHWAYS_BUCKET
               - --node_type=resource_manager
               - --instance_count=1
               - --instance_type=tpuv6e:4x4
               imagePullPolicy: Always
               ports:
               - containerPort: 38677
             - name: jax-tpu
               image: us-docker.pkg.dev/cloud-tpu-images/inference/jetstream-pathways:v0.2.0
               env:
               - name: LOG_LEVEL
                 value: "INFO"
               args:
               - MaxText/configs/v5e/inference/llama3_405b_v5e-64.yml
               - model_name=llama3.1-405b
               - load_parameters_path=CHECKPOINT_PATH
               - max_prefill_predict_length=1024
               - max_target_length=2048
               - async_checkpointing=false
               - steps=1
               - ici_fsdp_parallelism=1
               - ici_autoregressive_parallelism=2
               - ici_tensor_parallelism=8
               - scan_layers=false
               - weight_dtype=bfloat16
               - per_device_batch_size=6
               - enable_single_controller=true
               - quantization=int8
               - quantize_kvcache=true
               - checkpoint_is_quantized=true
               - enable_model_warmup=true
               imagePullPolicy: Always
               ports:
               - containerPort: 9000
               startupProbe:
                 httpGet:
                   path: /healthcheck
                   port: 8000
                   scheme: HTTP
                 periodSeconds: 1
                 initialDelaySeconds: 600
                 failureThreshold: 10000
               livenessProbe:
                 httpGet:
                   path: /healthcheck
                   port: 8000
                   scheme: HTTP
                 periodSeconds: 60
                 failureThreshold: 10
               readinessProbe:
                 httpGet:
                   path: /healthcheck
                   port: 8000
                   scheme: HTTP
                 periodSeconds: 60
                 failureThreshold: 10
             - name: jetstream-http
               image: us-docker.pkg.dev/cloud-tpu-images/inference/jetstream-http:v0.2.3
               imagePullPolicy: Always
               ports:
               - containerPort: 8000
         size: 5
         workerTemplate:
           spec:
             nodeSelector:
               cloud.google.com/gke-tpu-accelerator: tpu-v6e-slice
               cloud.google.com/gke-tpu-topology: 4x4
             tolerations:
             - key: "google.com/tpu"
               operator: "Exists"
               effect: "NoSchedule"
             containers:
             - name: worker
               args:
               - --server_port=38679
               - --resource_manager_address=$(LWS_LEADER_ADDRESS):38677
               - --gcs_scratch_location=PATHWAYS_BUCKET
               image: us-docker.pkg.dev/cloud-tpu-v2-images/pathways/server:jax-0.5.3
               imagePullPolicy: Always
               ports:
               - containerPort: 38679
               resources:
                 limits:
                   google.com/tpu: "4"
     --- 
     apiVersion: v1
     kind: Service
     metadata:
       name: jetstream-svc
     spec:
       selector:
         app: jetstream-pathways
       ports:
       - protocol: TCP
         name: jetstream-http
         port: 8000
         targetPort: 8000

  2. Imposta il valore del campo load_parameters_path sul percorso del checkpoint prodotto nel processo di conversione del checkpoint.

     • Per un checkpoint bf16, il percorso dovrebbe essere simile a gs://OUTPUT_BUCKET_DIRECTORY/bf16/unscanned/checkpoints/0/items.
     • Per un checkpoint int8, dovrebbe essere simile a gs://OUTPUT_BUCKET_DIRECTORY/int8.

     Imposta il valore del campo gcs_scratch_location sul bucket Pathways che hai creato in precedenza.

     perl -pi -e 's|CHECKPOINT_PATH|gs://OUTPUT_BUCKET_DIRECTORY/int8|g' jetstream-pathways-llama-3-1-405b-4x4.yaml
     perl -pi -e 's|PATHWAYS_BUCKET|gs://PATHWAYS_BUCKET|g' jetstream-pathways-llama-3-1-405b-4x4.yaml
     

     Descrizione per variabili e impostazioni.

     Le seguenti descrizioni spiegano le variabili e le impostazioni pertinenti per JetStream, Pathways e LeaderWorkerSet.

     • Jetstream
       • Questo container è creato dal progetto OSS JetStream. Inoltre, incapsula la libreria client basata su gRPC del proxy IFRT, che fornisce un backend JAX in grado di delegare l'esecuzione di programmi TPU al server proxy Pathways.
       • Per ulteriori informazioni sui flag JetStream-MaxText, consulta Descrizioni dei flag del server JetStream-MaxText.
       • Per utilizzare un checkpoint quantizzato, puoi impostare i flag quantization, quantize_kvcache o checkpoint_is_quantized. Se il checkpoint è bf16, puoi omettere questi flag.
     • Percorsi: per una descrizione dei componenti pertinenti dei percorsi, consulta Componenti dei percorsi.
     • LeaderWorkerSet: per ulteriori informazioni sulle impostazioni di LeaderWorkerSet, consulta la definizione dell'API.

  Applica il manifest del deployment

  Applica il manifest per eseguire il deployment del server:

  kubectl apply -f jetstream-pathways-llama-3-1-405b-4x4.yaml
  

  Il server del modello dovrebbe avviarsi.

  Verificare l'avvio del server del modello

  Un modello 405B potrebbe richiedere circa 10-20 minuti per ripristinare il checkpoint. Potresti anche dover attendere più tempo durante il riscaldamento del modello se hai attivato il flag enable_model_warmup.

  kubectl logs -f jetstream-pathways-0 -c jax-tpu
  

  L'output è simile al seguente:

  2025-03-02 02:15:07,682 - JetstreamLogger - INFO - Initializing the driver with 1 prefill engines and 1 generate engines in interleaved mode
  2025-03-02 02:15:07,683 - JetstreamLogger - INFO - Spinning up prefill thread 0.
  2025-03-02 02:15:07,683 - JetstreamLogger - INFO - Spinning up transfer thread 0.
  2025-03-02 02:15:07,684 - JetstreamLogger - INFO - Spinning up generate thread 0.
  2025-03-02 02:15:07,684 - JetstreamLogger - INFO - Spinning up detokenize thread 0.
  2025-03-02 02:15:07,685 - JetstreamLogger - INFO - Driver initialized.
  ...
  ...
  ...
  INFO:     Started server process [7]
  INFO:     Waiting for application startup.
  INFO:     Application startup complete.
  INFO:     Uvicorn running on http://0.0.0.0:9999 (Press CTRL+C to quit)
  

  Distribuisci Llama 3.1-405b

  Per pubblicare il modello Llama 3.1-405b, configura il port forwarding:

  kubectl port-forward svc/jetstream-svc 8000:8000
  

  Il port forwarding ti consente di accedere al servizio dall'esterno del cluster. Puoi accedere al deployment di JetStream-Pathways tramite il servizio ClusterIP di GKE. I servizi ClusterIP sono raggiungibili solo dall'interno del cluster.

  Interagire con il modello

  In un nuovo terminale, esegui questo comando:

  curl --request POST \
  --header "Content-type: application/json" \
  -s \
  localhost:8000/generate \
  --data \
  '{
      "prompt": "What are the top 5 programming languages",
      "max_tokens": 200
  }'
  

  Il completamento della richiesta iniziale può richiedere diversi secondi a causa del riscaldamento del modello. L'output dovrebbe essere simile al seguente:

  {
      "response": " for web development?\nThe top 5 programming languages for web development are:\n1. **JavaScript**: JavaScript is the most popular language for web development, used by over 90% of websites for client-side scripting. It's also popular for server-side programming with technologies like Node.js.\n2. **HTML/CSS**: HTML (Hypertext Markup Language) and CSS (Cascading Style Sheets) are not programming languages, but are essential for building websites. HTML is used for structuring content, while CSS is used for styling and layout.\n3. **Python**: Python is a popular language for web development, especially with frameworks like Django and Flask. It's known for its simplicity, flexibility, and large community of developers.\n4. **Java**: Java is a popular language for building enterprise-level web applications, especially with frameworks like Spring and Hibernate. It's known for its platform independence, strong security features, and large community of developers.\n5. **PHP**: PHP is a mature language for web"
  }
  

  Hai completato correttamente le seguenti operazioni:

  1. È stato eseguito il deployment del server di modelli JetStream con MaxText e Pathways su GKE utilizzando le TPU.
  2. È stato creato un checkpoint Llama 3.1-405B int8 alle ore gs://BUCKET_NAME.
  3. Ha pubblicato e interagito con il modello.

  Pubblicazione disaggregata

  La pubblicazione disaggregata è una tecnica per pubblicare LLM dividendo le fasi di precompilazione e decodifica in host diversi. Questo approccio ottimizza l'utilizzo delle risorse, il che può portare a un miglioramento della velocità effettiva e della latenza.

  • Precompilazione: un passaggio in avanti sul prompt di input per inizializzare la cache delle coppie chiave-valore.

  • Decodifica: una procedura che genera token di output in modo incrementale, un token per passaggio e un valore della cache KV per iterazione.

  1. Imposta le variabili di ambiente predefinite:

     export NODE_POOL_NAME=dis-v6e-8
     export NODE_POOL_SIZE=2
     export MACHINE_TYPE=ct6e-standard-4t
     export TOPOLOGY=2x4
     export WORKERS_PER_SLICE=2
     

  2. Crea due node pool che utilizzano nodi v6e-8:

     for i in $(seq 1 NODE_POOL_SIZE); do
       gcloud container node-pools create NODE_POOL_NAME-${i}-np \
       --project=PROJECT \
       --cluster=CLUSTER_NAME \
       --location=CONTROL_PLANE_LOCATION \
       --node-locations=NODE_LOCATION \
       --machine-type=MACHINE_TYPE \
       --num-nodes=WORKERS_PER_SLICE \
       --tpu-topology=TOPOLOGY \
       --scopes=cloud-platform \
       --workload-metadata=GCE_METADATA
     done
     

  Conversione del checkpoint

  Per convertire un checkpoint Meta Llama 2-70B in un checkpoint di inferenza int8 compatibile con MaxText, completa i passaggi descritti in Conversione del checkpoint con Llama2-70B. Seleziona Llama2-70B come modello quando accetti i Termini e condizioni di Meta. Il deployment utilizza il checkpoint con il flag load_parameters_path.

  Sostituisci i seguenti parametri nel file checkpoint-job.yaml:

  - --meta_url=META_URL
  - --model_name=llama-2
  - --model_path=Llama-2-70b-chat
  - --output_directory=gs://BUCKET_NAME/maxtext/llama-2-70b
  

  Il checkpoint verrà utilizzato nella tua implementazione con il flag load_parameters_path.

  Esegui il deployment di JetStream Pathways con la pubblicazione disaggregata

  1. Salva il seguente manifest come jetstream-pathways-disagg-llama-2-70b-2-2x4.yaml:

     apiVersion: leaderworkerset.x-k8s.io/v1
     kind: LeaderWorkerSet
     metadata:
       name: jetstream-pathways
       annotations:
         leaderworkerset.sigs.k8s.io/subgroup-exclusive-topology: cloud.google.com/gke-nodepool
     spec:
       replicas: 1
       leaderWorkerTemplate:
         subGroupPolicy:
           subGroupSize: 2
         leaderTemplate:
           metadata:
             labels:
               app: jetstream-pathways
           spec:
             nodeSelector:
               cloud.google.com/gke-tpu-accelerator: tpu-v6e-slice
               cloud.google.com/gke-tpu-topology: 2x4
             tolerations:
             - key: "google.com/tpu"
               operator: "Exists"
               effect: "NoSchedule"
             containers:
             - name: pathways-proxy
               image: us-docker.pkg.dev/cloud-tpu-v2-images/pathways/proxy_server:jax-0.5.3
               args:
               - --resource_manager_address=$(LWS_LEADER_ADDRESS):38677
               - --server_port=38681
               - --gcs_scratch_location=gs://cloud-pathways-staging/tmp
               - --xla_jf_auto_cross_replica_sharding=false
               - --xla_tpu_enable_windowed_einsum_for_reduce_scatter=false
               - --xla_tpu_enable_windowed_einsum_for_all_gather=false
               - --xla_tpu_prefer_latch_optimized_rhs_layouts=true
               - --xla_tpu_enable_experimental_fusion_cost_model=false
               - --xla_tpu_dot_dot_fusion_duplicated=false
               - --xla_tpu_dot_dot_fusion=true
               - --xla_jf_conv_input_fusion=true
               - --xla_jf_conv_output_fusion=true
               - --xla_tpu_rwb_fusion=false
               - --xla_tpu_copy_fusion_pad_unpad_ratio=0
               - --xla_tpu_licm_size_inflation_ratio=1
               - --xla_tpu_copy_elision_analysis_allowance=150000
               - --xla_tpu_copy_insertion_use_region_analysis_limit=10000
               - --xla_tpu_order_dot_after_layout=true
               - --xla_jf_rematerialization_percent_shared_memory_limit=100
               - --xla_tpu_use_repeated_instance_for_preferred_prefetch_time=true
               - --xla_tpu_enforce_prefetch_fifo_order=false
               - --xla_tpu_prefetch_interval_picker_size_override=6000000
               - --xla_tpu_async_copy_bandwidth_scaling_factor=1
               - --xla_tpu_nd_short_transfer_max_chunks=-1
               - --xla_tpu_enable_aggressive_broadcast_priority_update=true
               - --xla_tpu_alternate_memory_benefit_scaling_factor_for_large_buffers=SQRT
               - --xla_tpu_memory_bound_loop_optimizer_options=enabled:true
               - --xla_tpu_enable_copy_fusion=true
               - --xla_tpu_enable_cross_program_prefetch_freeing=false
               - --xla_tpu_enable_dot_strength_reduction=true
               - --xla_tpu_layout_use_dot_grouping=false
               - --xla_tpu_msa_inefficient_use_to_copy_ratio=0.5
               - --xla_tpu_reduce_loop_fusion_dup_with_unfusable_user=false
               - --xla_tpu_vector_load_fusion_window=1024
               - --xla_tpu_vector_store_fusion_window=256
               - --xla_jf_conv_reshape_fusion=false
               - --xla_tpu_input_conv_multi_users=false
               - --xla_tpu_enable_multi_level_input_dot_dot_fusion=false
               - --xla_tpu_enable_multi_level_output_dot_dot_fusion=false
               - --xla_tpu_dot_dot_fusion_separable_convs_only=false
               - --xla_tpu_enable_multi_level_nested_loop_fusion=true
               - --xla_tpu_nested_dot_fusion=true
               - --xla_tpu_enable_multi_level_nested_dot_fusion=false
               - --xla_jf_enable_multi_output_fusion=true
               - --xla_tpu_use_lp_llo_scheduler_for_dot_dot_fusions=false
               - --xla_tpu_enable_flash_attention=true
               imagePullPolicy: Always
               ports:
               - containerPort: 38681
             - name: pathways-rm
               env:       
               - name: HOST_ADDRESS
                 value: "$(LWS_LEADER_ADDRESS)"
               - name: TPU_SKIP_MDS_QUERY
                 value: "true"
               image: us-docker.pkg.dev/cloud-tpu-v2-images/pathways/server:jax-0.5.3
               args:
               - --server_port=38677
               - --gcs_scratch_location=PATHWAYS_BUCKET
               - --node_type=resource_manager
               - --instance_count=2
               - --instance_type=tpuv6e:2x4
               imagePullPolicy: Always
               ports:
               - containerPort: 38677
             - name: jax-tpu
               image: us-docker.pkg.dev/cloud-tpu-images/inference/jetstream-pathways:v0.2.0
               args:
               - MaxText/configs/base.yml
               - tokenizer_path=assets/tokenizer.llama2
               - load_parameters_path=CHECKPOINT_PATH
               - max_prefill_predict_length=1024
               - max_target_length=2048
               - model_name=llama2-70b
               - ici_fsdp_parallelism=1
               - ici_autoregressive_parallelism=1
               - ici_tensor_parallelism=-1
               - scan_layers=false
               - weight_dtype=bfloat16
               - per_device_batch_size=27
               - checkpoint_is_quantized=true 
               - quantization=int8
               - quantize_kvcache=true
               - compute_axis_order=0,2,1,3
               - ar_cache_axis_order=0,2,1,3
               - stack_prefill_result_cache=True
               - inference_server=ExperimentalMaxtextDisaggregatedServer_8
               - inference_benchmark_test=True
               - enable_model_warmup=True
               env:
               - name: LOG_LEVEL
                 value: "INFO"
               imagePullPolicy: Always
               securityContext:
                 capabilities:
                   add: ["SYS_PTRACE", "NET_ADMIN", "SYS_TIME"]
               ports: 
               - containerPort: 9000
               startupProbe:
                 httpGet:
                   path: /healthcheck
                   port: 8000
                   scheme: HTTP
                 periodSeconds: 1
                 initialDelaySeconds: 240
                 failureThreshold: 10000
               livenessProbe:
                 httpGet:
                   path: /healthcheck
                   port: 8000
                   scheme: HTTP
                 periodSeconds: 60
                 failureThreshold: 100
               readinessProbe:
                 httpGet:
                   path: /healthcheck
                   port: 8000
                   scheme: HTTP
                 periodSeconds: 60
                 failureThreshold: 100
             - name: jetstream-http
               image: us-docker.pkg.dev/cloud-tpu-images/inference/jetstream-http:v0.2.3
               imagePullPolicy: Always
               ports:
               - containerPort: 8000
         size: 5
         workerTemplate:
           spec:
             nodeSelector:
               cloud.google.com/gke-tpu-accelerator: tpu-v6e-slice
               cloud.google.com/gke-tpu-topology: 2x4
             containers:
             - name: worker
               args:
               - --server_port=38679
               - --resource_manager_address=$(LWS_LEADER_ADDRESS):38677
               - --gcs_scratch_location=PATHWAYS_BUCKET
               image: us-docker.pkg.dev/cloud-tpu-v2-images/pathways/server:jax-0.5.3
               imagePullPolicy: Always
               ports:
               - containerPort: 38679
               resources:
                 limits:
                   google.com/tpu: "4"
     --- 
     apiVersion: v1
     kind: Service
     metadata:
       name: jetstream-svc
     spec:
       selector:
         app: jetstream-pathways
       ports:
       - protocol: TCP
         name: jetstream-http
         port: 8000
         targetPort: 8000

  2. Imposta il valore del campo load_parameters_path sul percorso del checkpoint prodotto nel processo di conversione del checkpoint.

     • Per un checkpoint bf16, il percorso dovrebbe essere simile a gs://OUTPUT_BUCKET_DIRECTORY/bf16/unscanned/checkpoints/0/items.
     • Per un checkpoint int8, dovrebbe essere simile a gs://OUTPUT_BUCKET_DIRECTORY/int8.

     Imposta il valore del campo gcs_scratch_location sul bucket Pathways che hai creato in precedenza.

     perl -pi -e 's|CHECKPOINT_PATH|BUCKET_NAME/maxtext/llama-2-70b/int8|g' jetstream-pathways-disagg-llama-2-70b-2-2x4.yaml
     perl -pi -e 's|PATHWAYS_BUCKET|gs://PATHWAYS_BUCKET|g' jetstream-pathways-disagg-llama-2-70b-2-2x4.yaml
     

  3. Applica il manifest:

     kubectl apply -f jetstream-pathways-disagg-llama-2-70b-2-2x4.yaml
     

     Il server del modello potrebbe richiedere del tempo per ripristinare il checkpoint, a seconda delle dimensioni del checkpoint. Il ripristino del checkpoint di un modello da 70 miliardi di parametri potrebbe richiedere circa 8 minuti, incluso il riscaldamento del modello. Puoi osservare ulteriormente i log per determinare il punto di preparazione verificando l'avvio del server del modello e servire il modello configurando l'inoltro delle porte in modo da poter interagire con il modello.

  Hai completato correttamente le seguenti operazioni:

  1. È stato eseguito il deployment del server di modelli JetStream con MaxText e Pathways su GKE utilizzando le TPU e il servizio disaggregato.
  2. È stato creato un checkpoint Llama 2-70B int8 in gs://BUCKET_NAME.
  3. Ha pubblicato e interagito con il modello.

  Risoluzione dei problemi

  • Se visualizzi il messaggio Empty reply from server, è possibile che il container non abbia completato il download dei dati del modello. Controlla di nuovo i log del pod per il messaggio Connected, che indica che il modello è pronto per la pubblicazione.
  • Se visualizzi il messaggio Connection refused, verifica che l'inoltro delle porte sia attivo.

  Esegui la pulizia

  Per evitare che al tuo account Google Cloud vengano addebitati costi relativi alle risorse utilizzate in questo tutorial, elimina il progetto che contiene le risorse oppure mantieni il progetto ed elimina le singole risorse.

  Elimina le risorse di cui è stato eseguito il deployment

  Per evitare che al tuo account Google Cloud vengano addebitati costi relativi alle risorse che hai creato in questa guida, esegui i seguenti comandi e segui le istruzioni:

  gcloud container clusters delete CLUSTER_NAME --location=CONTROL_PLANE_LOCATION
  
  gcloud iam service-accounts delete jetstream-pathways@PROJECT_ID.iam.gserviceaccount.com
  
  gcloud storage rm --recursive gs://BUCKET_NAME
  

  Passaggi successivi

  • Scopri come eseguire i modelli Gemma su GKE e come eseguire carichi di lavoro di AI/ML ottimizzati con le funzionalità di orchestrazione della piattaforma GKE.
  • Scopri di più sulle TPU in GKE.
  • Esplora il repository GitHub di JetStream.
  • Esplora Vertex AI Model Garden.