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

Questa guida mostra come pubblicare 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 pubblicare 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 aumentati 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 in 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 di 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 l'erogazione di PyTorch/XLA e JAX TPU per ottimizzare le prestazioni.

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 question answering. GKE offre l'infrastruttura necessaria per supportare le esigenze di addestramento e serving distribuito 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.
  • IFRT proxy server riceve le richieste da IFRT proxy client e le inoltra a 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

  • Accedi al tuo account Google Cloud . Se non conosci Google Cloud, crea un account per valutare le prestazioni dei nostri prodotti in scenari reali. I nuovi clienti ricevono anche 300 $di crediti senza costi per l'esecuzione, il test e il deployment dei workload.

  • 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 role (roles/resourcemanager.projectCreator), which contains the resourcemanager.projects.create permission. Learn how to grant roles.

    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 role (roles/resourcemanager.projectCreator), which contains the resourcemanager.projects.create permission. Learn how to grant roles.

    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

    •

  • Assicurati di disporre dei seguenti ruoli nel progetto: roles/container.admin, roles/iam.serviceAccountAdmin, roles/resourcemanager.projectIamAdmin

    Controlla i ruoli

    1. Nella console Google Cloud vai alla pagina IAM.

      Vai a IAM
    2. Seleziona il progetto.

    3. Nella colonna Entità, trova tutte le righe che identificano te o un gruppo di cui fai parte. Per scoprire a quali gruppi appartieni, contatta il tuo amministratore.

    4. Per tutte le righe che ti specificano o ti includono, controlla la colonna Ruolo per verificare se l'elenco dei ruoli include i ruoli richiesti.

    Concedi i ruoli

    1. Nella console Google Cloud vai alla pagina IAM.

      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. Fai clic su Seleziona un ruolo, quindi cerca il ruolo.
    6. Per concedere altri ruoli, fai clic su Aggiungi un altro ruolo e aggiungi ogni ruolo successivo.
    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.

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 della CLI llama. 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 Icona di attivazione di Cloud Shell Attiva Cloud Shell nella Google Cloud console. 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 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:

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 node pool 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

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à memorizzato il checkpoint:

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

gcloud projects add-iam-policy-binding ${PROJECT_ID} \
  --member "serviceAccount:jetstream-pathways@${PROJECT_ID}.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_ID}.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_NAME}" --project "${PROJECT_ID}"  \
    --location "${CONTROL_PLANE_LOCATION}"

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 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 per i percorsi che hai creato in precedenza.

    perl -pi -e "s|CHECKPOINT_PATH|gs://${BUCKET_NAME}/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

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ù a lungo 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 eseguito 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.

  • Prefill: un passaggio in avanti sul prompt di input per inizializzare la cache della coppia 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_ID} \
  --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 per i percorsi che hai creato in precedenza.

    perl -pi -e "s|CHECKPOINT_PATH|gs://${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 idoneità verificando l'avvio del server del modello ed erogare il modello configurando l'inoltro delle porte in modo da poter interagire con il modello.

Hai eseguito 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 il modello e interagito con esso.

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