Entrega un LLM con TPU de varios hosts en GKE con Saxml

En este instructivo, se muestra cómo implementar y entregar un modelo de lenguaje grande (LLM) mediante un grupo de nodos de porción de TPU de varios hosts en Google Kubernetes Engine (GKE) con Saxml para una arquitectura escalable eficiente.

Segundo plano

Saxml es un sistema experimental que entrega frameworks Paxml, JAX y PyTorch. Puedes usar TPU para acelerar el procesamiento de datos con estos frameworks. Para demostrar la implementación de las TPU en GKE, en este instructivo se entrega el modelo de prueba LmCloudSpmd175B32Test 175B. GKE implementa este modelo de prueba en dos grupos de nodos de porción de TPU v5e con la topología 4x8, respectivamente.

Para implementar correctamente el modelo de prueba, la topología de TPU se definió según el tamaño del modelo. Dado que el modelo de N mil millones de 16 bits requiere aproximadamente 2 veces la memoria (2xN) GB, el modelo LmCloudSpmd175B32Test de 175,000 millones requiere alrededor de 350 GB de memoria. El chip único de TPU v5e tiene 16 GB. Para admitir 350 GB, GKE necesita 21 chips de TPU v5e (350/16= 21). Según la asignación de la configuración de TPU, la configuración de TPU adecuada para este instructivo es la siguiente:

  • Tipo de máquina: ct5lp-hightpu-4t
  • Topología: 4x8 (32 chips TPU)

Seleccionar la topología de TPU correcta para la entrega de un modelo es importante cuando se implementan TPU en GKE. Para obtener más información, consulta Planifica tu configuración de TPU.

Objetivos

Este instructivo está dirigido a administradores de plataformas o ingenieros de MLOps o DevOps que deseen usar las funciones de organización de GKE para entregar modelos de datos.

En este instructivo, se abarcan los siguientes pasos:

  1. Prepara el entorno con un clúster de GKE Standard. El clúster tiene dos grupos de nodos de porción de TPU v5e con la topología 4x8.
  2. Implementa Saxml. Saxml necesita un servidor de administrador, un grupo de Pods que funcionen como el servidor de modelos, un servidor HTTP compilado previamente y un balanceador de cargas.
  3. Usa Saxml para entregar el LLM.

En el siguiente diagrama, se muestra la arquitectura que implementa el siguiente instructivo:

Arquitectura de una TPU de varios hosts en GKE.
Figura: Arquitectura de ejemplo de una TPU de varios hosts en GKE.

Antes de comenzar

  • 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.

    Go to project selector

  • Make sure that billing is enabled for your Google Cloud project.

  • Enable the required API.

    Enable the API

    •

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

    Go to project selector

  • Make sure that billing is enabled for your Google Cloud project.

  • Enable the required API.

    Enable the API

    •

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

    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 colunn 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.

      Ir a IAM
    2. Selecciona el proyecto.
    3. Haz clic en Grant access.

    4. En el campo Principales nuevas, ingresa tu identificador de usuario. Esta suele ser la dirección de correo electrónico de una Cuenta de Google.

    5. En la lista Seleccionar un rol, elige un rol.
    6. Para otorgar funciones adicionales, haz clic en Agregar otro rol y agrega cada rol adicional.
    7. Haz clic en Guardar.

Prepara el entorno

  1. En la consola de Google Cloud, inicia una instancia de Cloud Shell:
    Abrir Cloud Shell

  2. Configura las variables de entorno predeterminadas:

      gcloud config set project PROJECT_ID
  export PROJECT_ID=$(gcloud config get project)
  export REGION=COMPUTE_REGION
  export ZONE=COMPUTE_ZONE
  export GSBUCKET=PROJECT_ID-gke-bucket

    Reemplaza los siguientes valores:

Crea un clúster de GKE Standard

Usa Cloud Shell para realizar las siguientes acciones:

  1. Crea un clúster estándar que use la federación de identidades para cargas de trabajo para GKE:

    gcloud container clusters create saxml \
    --zone=${ZONE} \
    --workload-pool=${PROJECT_ID}.svc.id.goog \
    --cluster-version=VERSION \
    --num-nodes=4

    Reemplaza VERSION por el número de versión de GKE. GKE es compatible con TPU v5e en la versión 1.27.2-gke.2100 y posteriores. Para obtener más información, consulta la disponibilidad de TPU en GKE.

    La creación del clúster puede tomar varios minutos.

  2. Crea el primer grupo de nodos llamado tpu1:

    gcloud container node-pools create tpu1 \
    --zone=${ZONE} \
    --num-nodes=8 \
    --machine-type=ct5lp-hightpu-4t \
    --tpu-topology=4x8 \
    --cluster=saxml

  3. Crea el segundo grupo de nodos llamado tpu2:

    gcloud container node-pools create tpu2 \
    --zone=${ZONE} \
    --num-nodes=8 \
    --machine-type=ct5lp-hightpu-4t \
    --tpu-topology=4x8 \
    --cluster=saxml

Creaste los siguientes recursos:

  • Un clúster de Standard con cuatro nodos de CPU.
  • Dos grupos de nodos de porción de TPU v5e con topología 4x8. Cada grupo de nodos representa ocho nodos de porción de TPU con 4 chips TPU cada uno.

El modelo 175B debe entregarse en una porción de TPU v5e de varios hosts con una porción de topología 4x8 (32 chips TPU v5e) como mínimo.

Crear un bucket de Cloud Storage

Crea un bucket de Cloud Storage para almacenar las opciones de configuración del servidor de administrador de Saxml. Un servidor administrador en ejecución guarda su estado y los detalles de los modelos publicados de forma periódica.

En Cloud Shell, ejecuta lo siguiente:

gcloud storage buckets create gs://${GSBUCKET}

Configura el acceso a tus cargas de trabajo mediante la federación de identidades para cargas de trabajo para GKE

Asigna una Cuenta de servicio de Kubernetes a la aplicación y configúrala para que actúe como una cuenta de servicio de IAM.

  1. Configura kubectl para comunicarse con tu clúster:

    gcloud container clusters get-credentials saxml --zone=${ZONE}

  2. Crea una cuenta de servicio de Kubernetes para que tu aplicación use:

    kubectl create serviceaccount sax-sa --namespace default

  3. Crea una cuenta de servicio de IAM para tu aplicación:

    gcloud iam service-accounts create sax-iam-sa

  4. Agrega una vinculación de política de IAM para que tu cuenta de servicio de IAM lea y escriba en Cloud Storage:

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

  5. Para permitir que la cuenta de servicio de Kubernetes actúe en nombre de la cuenta de servicio de IAM, agrega una vinculación de política de IAM entre las dos. Esta vinculación permite que la cuenta de servicio de Kubernetes actúe como la cuenta de servicio de IAM para que la cuenta de servicio de Kubernetes pueda leer y escribir en Cloud Storage.

    gcloud iam service-accounts add-iam-policy-binding sax-iam-sa@${PROJECT_ID}.iam.gserviceaccount.com \
  --role roles/iam.workloadIdentityUser \
  --member "serviceAccount:${PROJECT_ID}.svc.id.goog[default/sax-sa]"

  6. Anota la cuenta de servicio de Kubernetes con la dirección de correo electrónico de la cuenta de servicio de IAM. Esto permite que tu app de ejemplo sepa qué cuenta de servicio usar para acceder a los servicios de Google Cloud. Por lo tanto, cuando la app usa cualquier biblioteca cliente estándar de la API de Google para acceder a los servicios de Google Cloud, usa esa cuenta de servicio de IAM.

    kubectl annotate serviceaccount sax-sa \
  iam.gke.io/gcp-service-account=sax-iam-sa@${PROJECT_ID}.iam.gserviceaccount.com

Implementa Saxml

En esta sección, implementarás el servidor de administrador de Saxml y el servidor de modelos de Saxml.

Implementa el servidor administrador de Saxml

  1. Crea el siguiente manifiesto sax-admin-server.yaml:

    apiVersion: apps/v1
kind: Deployment
metadata:
  name: sax-admin-server
spec:
  replicas: 1
  selector:
    matchLabels:
      app: sax-admin-server
  template:
    metadata:
      labels:
        app: sax-admin-server
    spec:
      hostNetwork: false
      serviceAccountName: sax-sa
      containers:
      - name: sax-admin-server
        image: us-docker.pkg.dev/cloud-tpu-images/inference/sax-admin-server:v1.1.0
        securityContext:
          privileged: true
        ports:
        - containerPort: 10000
        env:
        - name: GSBUCKET
          value: BUCKET_NAME

    Reemplaza BUCKET_NAME por el nombre del bucket de Cloud Storage.

  2. Aplica el manifiesto

    kubectl apply -f sax-admin-server.yaml

  3. Verifica que el Pod del servidor administrador esté en funcionamiento:

    kubectl get deployment

    El resultado es similar a este:

    NAME               READY   UP-TO-DATE   AVAILABLE   AGE
sax-admin-server   1/1     1            1           52s

Implementa el servidor de modelos de Saxml

Las cargas de trabajo que se ejecutan en porciones de TPU de varios hosts requieren un identificador de red estable para que cada Pod descubra pares en la misma porción de TPU. Para definir estos identificadores, usa IndexedJob, StatefulSet con un Service sin interfaz gráfica o JobSet que crea automáticamente un Service sin interfaz gráfica para todos los trabajos que pertenecen al JobSet. En la siguiente sección, se muestra cómo administrar varios grupos de Pods del servidor de modelos con JobSet.

  1. Instala JobSet v0.2.3 o una versión posterior.

    kubectl apply --server-side -f https://github.com/kubernetes-sigs/jobset/releases/download/JOBSET_VERSION/manifests.yaml

    Reemplaza JOBSET_VERSION por la versión de JobSet. Por ejemplo, v0.2.3

  2. Valida que el controlador JobSet se ejecute en el espacio de nombres jobset-system:

    kubectl get pod -n jobset-system

    El resultado es similar a este:

    NAME                                        READY   STATUS    RESTARTS   AGE
jobset-controller-manager-69449d86bc-hp5r6   2/2     Running   0          2m15s

  3. Implementa dos servidores de modelos en dos grupos de nodos de porción de TPU. Guarda el siguiente manifiesto sax-model-server-set:

    apiVersion: jobset.x-k8s.io/v1alpha2
kind: JobSet
metadata:
  name: sax-model-server-set
  annotations:
    alpha.jobset.sigs.k8s.io/exclusive-topology: cloud.google.com/gke-nodepool
spec:
  failurePolicy:
    maxRestarts: 4
  replicatedJobs:
    - name: sax-model-server
      replicas: 2
      template:
        spec:
          parallelism: 8
          completions: 8
          backoffLimit: 0
          template:
            spec:
              serviceAccountName: sax-sa
              hostNetwork: true
              dnsPolicy: ClusterFirstWithHostNet
              nodeSelector:
                cloud.google.com/gke-tpu-accelerator: tpu-v5-lite-podslice
                cloud.google.com/gke-tpu-topology: 4x8
              containers:
              - name: sax-model-server
                image: us-docker.pkg.dev/cloud-tpu-images/inference/sax-model-server:v1.1.0
                args: ["--port=10001","--sax_cell=/sax/test", "--platform_chip=tpuv5e"]
                ports:
                - containerPort: 10001
                - containerPort: 8471
                securityContext:
                  privileged: true
                env:
                - name: SAX_ROOT
                  value: "gs://BUCKET_NAME/sax-root"
                - name: MEGASCALE_NUM_SLICES
                  value: ""
                resources:
                  requests:
                    google.com/tpu: 4
                  limits:
                    google.com/tpu: 4

    Reemplaza BUCKET_NAME por el nombre del bucket de Cloud Storage.

    En el manifiesto se muestra lo siguiente:

    • replicas: 2 es la cantidad de réplicas del trabajo. Cada trabajo representa un servidor de modelos. Por lo tanto, un grupo de 8 Pods.
    • parallelism: 8 y completions: 8 son iguales a la cantidad de nodos en cada grupo de nodos.
    • backoffLimit: 0 debe ser cero para marcar el trabajo como con errores si falla cualquier Pod.
    • ports.containerPort: 8471 es el puerto predeterminado para la comunicación de las VMs.
    • name: MEGASCALE_NUM_SLICES desactiva la variable de entorno porque GKE no ejecuta el entrenamiento de Multislice.

  4. Aplica el manifiesto

    kubectl apply -f sax-model-server-set.yaml

  5. Verifica el estado del servidor de administración de Saxml y los Pods del servidor de modelos:

    kubectl get pods

    El resultado es similar a este:

    NAME                                              READY   STATUS    RESTARTS   AGE
sax-admin-server-557c85f488-lnd5d                 1/1     Running   0          35h
sax-model-server-set-sax-model-server-0-0-nj4sm   1/1     Running   0          24m
sax-model-server-set-sax-model-server-0-1-sl8w4   1/1     Running   0          24m
sax-model-server-set-sax-model-server-0-2-hb4rk   1/1     Running   0          24m
sax-model-server-set-sax-model-server-0-3-qv67g   1/1     Running   0          24m
sax-model-server-set-sax-model-server-0-4-pzqz6   1/1     Running   0          24m
sax-model-server-set-sax-model-server-0-5-nm7mz   1/1     Running   0          24m
sax-model-server-set-sax-model-server-0-6-7br2x   1/1     Running   0          24m
sax-model-server-set-sax-model-server-0-7-4pw6z   1/1     Running   0          24m
sax-model-server-set-sax-model-server-1-0-8mlf5   1/1     Running   0          24m
sax-model-server-set-sax-model-server-1-1-h6z6w   1/1     Running   0          24m
sax-model-server-set-sax-model-server-1-2-jggtv   1/1     Running   0          24m
sax-model-server-set-sax-model-server-1-3-9v8kj   1/1     Running   0          24m
sax-model-server-set-sax-model-server-1-4-6vlb2   1/1     Running   0          24m
sax-model-server-set-sax-model-server-1-5-h689p   1/1     Running   0          24m
sax-model-server-set-sax-model-server-1-6-bgv5k   1/1     Running   0          24m
sax-model-server-set-sax-model-server-1-7-cd6gv   1/1     Running   0          24m

En este ejemplo, hay 16 contenedores de servidores de modelos: sax-model-server-set-sax-model-server-0-0-nj4sm y sax-model-server-set-sax-model-server-1-0-8mlf5 son los dos servidores de modelos principales en cada grupo.

Tu clúster de Saxml tiene dos servidores de modelos implementados en dos grupos de nodos de porción de TPU v5e con la topología 4x8, respectivamente.

Implementa el servidor HTTP de Saxml y el balanceador de cargas

  1. Usa la siguiente imagen del servidor HTTP de imágenes compilada previamente. Guarda el siguiente manifiesto sax-http.yaml:

    apiVersion: apps/v1
kind: Deployment
metadata:
  name: sax-http
spec:
  replicas: 1
  selector:
    matchLabels:
      app: sax-http
  template:
    metadata:
      labels:
        app: sax-http
    spec:
      hostNetwork: false
      serviceAccountName: sax-sa
      containers:
      - name: sax-http
        image: us-docker.pkg.dev/cloud-tpu-images/inference/sax-http:v1.0.0
        ports:
        - containerPort: 8888
        env:
        - name: SAX_ROOT
          value: "gs://BUCKET_NAME/sax-root"
---
apiVersion: v1
kind: Service
metadata:
  name: sax-http-lb
spec:
  selector:
    app: sax-http
  ports:
  - protocol: TCP
    port: 8888
    targetPort: 8888
  type: LoadBalancer

    Reemplaza BUCKET_NAME por el nombre del bucket de Cloud Storage.

  2. Aplica el manifiesto sax-http.yaml:

    kubectl apply -f sax-http.yaml

  3. Espera a que el contenedor del servidor HTTP termine de crear:

    kubectl get pods

    El resultado es similar a este:

    NAME                                              READY   STATUS    RESTARTS   AGE
sax-admin-server-557c85f488-lnd5d                 1/1     Running   0          35h
sax-http-65d478d987-6q7zd                         1/1     Running   0          24m
sax-model-server-set-sax-model-server-0-0-nj4sm   1/1     Running   0          24m
...

  4. Espera a que se asigne una dirección IP externa al Service:

    kubectl get svc

    El resultado es similar a este:

    NAME           TYPE           CLUSTER-IP    EXTERNAL-IP   PORT(S)          AGE
sax-http-lb    LoadBalancer   10.48.11.80   10.182.0.87   8888:32674/TCP   7m36s

Usa Saxml

Carga, implementa y entrega el modelo en Saxml en la porción de varios hosts de TPU v5e:

Carga el modelo

  1. Recupera la dirección IP del balanceador de cargas para Saxml.

    LB_IP=$(kubectl get svc sax-http-lb -o jsonpath='{.status.loadBalancer.ingress[*].ip}')
PORT="8888"

  2. Carga el modelo de prueba LmCloudSpmd175B en dos grupos de nodos de porción de TPU v5e:

    curl --request POST \
--header "Content-type: application/json" \
-s ${LB_IP}:${PORT}/publish --data \
'{
    "model": "/sax/test/spmd",
    "model_path": "saxml.server.pax.lm.params.lm_cloud.LmCloudSpmd175B32Test",
    "checkpoint": "None",
    "replicas": 2
}'

    El modelo de prueba no tiene un punto de control ajustado, los pesos se generan de forma aleatoria. La carga del modelo puede tardar hasta 10 minutos.

    El resultado es similar a este:

    {
    "model": "/sax/test/spmd",
    "path": "saxml.server.pax.lm.params.lm_cloud.LmCloudSpmd175B32Test",
    "checkpoint": "None",
    "replicas": 2
}

  3. Verifica la preparación del modelo:

    kubectl logs sax-model-server-set-sax-model-server-0-0-nj4sm

    El resultado es similar a este:

    ...
loading completed.
Successfully loaded model for key: /sax/test/spmd

    El modelo está completamente cargado.

  4. Obtén información sobre el modelo:

    curl --request GET \
--header "Content-type: application/json" \
-s ${LB_IP}:${PORT}/listcell --data \
'{
    "model": "/sax/test/spmd"
}'

    El resultado es similar a este:

    {
"model": "/sax/test/spmd",
"model_path": "saxml.server.pax.lm.params.lm_cloud.LmCloudSpmd175B32Test",
"checkpoint": "None",
"max_replicas": 2,
"active_replicas": 2
}

Entrega el modelo

Entregar una solicitud de instrucción:

curl --request POST \
--header "Content-type: application/json" \
-s ${LB_IP}:${PORT}/generate --data \
'{
  "model": "/sax/test/spmd",
  "query": "How many days are in a week?"
}'

El resultado muestra un ejemplo de la respuesta del modelo. Es posible que esta respuesta no sea significativa porque el modelo de prueba tiene pesos aleatorios.

Anula la publicación del modelo

Ejecuta el siguiente comando para anular la publicación del modelo:

curl --request POST \
--header "Content-type: application/json" \
-s ${LB_IP}:${PORT}/unpublish --data \
'{
    "model": "/sax/test/spmd"
}'

El resultado es similar a este:

{
  "model": "/sax/test/spmd"
}

Limpia

Para evitar que se apliquen cargos a tu cuenta de Google Cloud por los recursos usados en este instructivo, borra el proyecto que contiene los recursos o conserva el proyecto y borra los recursos individuales.

Borra los recursos implementados

  1. Borra el clúster que creaste para este instructivo:

    gcloud container clusters delete saxml --zone ${ZONE}

  2. Borra la cuenta de servicio:

    gcloud iam service-accounts delete sax-iam-sa@${PROJECT_ID}.iam.gserviceaccount.com

  3. Borra el bucket de Cloud Storage:

    gcloud storage rm -r gs://${GSBUCKET}

¿Qué sigue?