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

Estándar

En este instructivo, se muestra cómo entregar un modelo grande de lenguaje (LLM) mediante unidades de procesamiento tensorial (TPU) en Google Kubernetes Engine (GKE) con Saxml.

Antecedentes

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 de 175B. GKE implementa este modelo de prueba en dos grupos de nodos TPU v5e con la topología 4x8, respectivamente.

Para implementar correctamente el modelo de prueba, se definió la topología de TPU 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 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 entregar un modelo es importante cuando se implementan TPU en GKE. Para obtener más información, consulta Planifica la configuración de TPU.

Objetivos

Este instructivo está dirigido a ingenieros 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:

Prepara el entorno con un clúster de GKE Standard. El clúster tiene dos grupos de nodos TPU v5e con topología 4x8.
Implementar 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.
Usa el 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:** Ejemplo de arquitectura de una TPU de varios hosts en GKE.

Antes de comenzar

Accede a tu cuenta de Google Cloud. Si eres nuevo en Google Cloud, crea una cuenta para evaluar el rendimiento de nuestros productos en situaciones reales. Los clientes nuevos también obtienen $300 en créditos gratuitos para ejecutar, probar y, además, implementar cargas de trabajo.

En la página del selector de proyectos de la consola de Google Cloud, selecciona o crea un proyecto de Google Cloud.

Ir al selector de proyectos

Asegúrate de que la facturación esté habilitada para tu proyecto de Google Cloud.

Habilita la API necesaria.

Habilita la API

En la página del selector de proyectos de la consola de Google Cloud, selecciona o crea un proyecto de Google Cloud.

Ir al selector de proyectos

Asegúrate de que la facturación esté habilitada para tu proyecto de Google Cloud.

Habilita la API necesaria.

Habilita la API

Asegúrate de tener los siguientes roles en el proyecto: roles/container.admin, roles/iam.serviceAccountAdmin
Verifica los roles
1. En la consola de Google Cloud, ve a la página IAM.
  Ir a IAM
2. Selecciona el proyecto.
3. En la columna Principal, busca la fila que tiene tu dirección de correo electrónico.
  
  Si tu dirección de correo electrónico no está en esa columna, no tienes ningún rol.
4. En la columna Función de la fila con la dirección de correo electrónico, verifica si la lista de roles incluye los roles necesarios.
Otorga los roles
1. En la consola de Google Cloud, ve a la página IAM.
  Ir a IAM
2. Selecciona el proyecto.
3. Haz clic en Grant access.
4. En el campo Principales nuevas, ingresa tu dirección de correo electrónico.
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.

Asegúrate de que tu proyecto tenga la cuota suficiente para Cloud TPU en GKE.

Prepara el entorno

En la consola de Google Cloud, inicia una instancia de Cloud Shell:
Abrir Cloud Shell
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:
- PROJECT_ID: El ID del proyecto de Google Cloud.
- COMPUTE_REGION: La región de Compute Engine
- COMPUTE_ZONE: La zona en la que ct5lp-hightpu-4t está disponible.

Crea un clúster de GKE Standard

Usa Cloud Shell para realizar las siguientes acciones:

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.

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

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 estándar con cuatro nodos de CPU.
Dos grupos de nodos TPU v5e con topología 4x8. Cada grupo de nodos representa ocho nodos TPU con 4 chips 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 administración de Saxml. Un servidor de administración en ejecución guarda de forma periódica su estado y los detalles de los modelos publicados.

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.

Configura kubectl para comunicarse con tu clúster:

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

Crea una cuenta de servicio de Kubernetes para que tu aplicación use:
```
kubectl create serviceaccount sax-sa --namespace default
```
Crea una cuenta de servicio de IAM para tu aplicación:
```
gcloud iam service-accounts create sax-iam-sa
```

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

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]"
```
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 administración de Saxml y el servidor de modelos de Saxml.

Implementa el servidor de administración de Saxml

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 de tu bucket de Cloud Storage.

Aplica el manifiesto
```
kubectl apply -f sax-admin-server.yaml
```

Verifica que el Pod del servidor de administración esté en funcionamiento:

kubectl get deployment

El resultado es similar al siguiente:

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

Implementa un servidor de modelos 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.

Instala JobSet v0.2.3 o 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

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

kubectl get pod -n jobset-system

El resultado es similar al siguiente:

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

Implementa dos servidores de modelos en dos grupos de nodos 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 de tu bucket de Cloud Storage.

En el manifiesto se muestra lo siguiente:

replicas: 2 es la cantidad de réplicas de 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 de TPU.
name: MEGASCALE_NUM_SLICES anula la configuración de la variable de entorno porque GKE no ejecuta el entrenamiento de varias porciones.

Aplica el manifiesto

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

Verifica el estado del Pod del administrador de Saxml y de los Pods del servidor de modelos:

kubectl get pods

El resultado es similar al siguiente:

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 servidor de modelo: 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 TPU v5e con la topología 4x8 respectivamente.

Implementa el servidor HTTP de Saxml y el balanceador de cargas

Usa la siguiente imagen de servidor HTTP de imagen compilada con anterioridad. 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 de tu bucket de Cloud Storage.

Aplica el manifiesto sax-http.yaml:
```
kubectl apply -f sax-http.yaml
```
Nota: Puedes usar tu propio servidor HTTP compilado para Saxml. Para obtener más información, consulta Infiere mediante Saxml y un servidor HTTP.

Espera a que el contenedor del servidor HTTP termine de crearse:

kubectl get pods

El resultado es similar al siguiente:

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

Espera a que el Service tenga asignada una dirección IP externa:

kubectl get svc

El resultado es similar al siguiente:

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 el Saxml en la porción de varios hosts de TPU v5e:

Carga el modelo

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"

Carga el modelo de prueba LmCloudSpmd175B en dos grupos de nodos 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 al siguiente:

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

Verifica la preparación del modelo:

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

El resultado es similar al siguiente:

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

El modelo está completamente cargado.

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 al siguiente:

{
"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

Entrega una solicitud de mensaje:

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. Esta respuesta puede no ser 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 al siguiente:

{
  "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

Borra el clúster que creaste para este instructivo:

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

Borra la cuenta de servicio:

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

Borra el bucket de Cloud Storage:
```
gcloud storage rm -r gs://${GSBUCKET}
```

¿Qué sigue?

Obtén información sobre las versiones actuales de TPU con la arquitectura del sistema de Cloud TPU.
Obtén más información sobre las TPU en GKE.

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

Antecedentes

Objetivos

Antes de comenzar

Verifica los roles

Otorga los roles

Prepara el entorno

Crea un clúster de GKE Standard

Crear un bucket de Cloud Storage

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

Implementa Saxml

Implementa el servidor de administración de Saxml

Implementa un servidor de modelos Saxml

Implementa el servidor HTTP de Saxml y el balanceador de cargas

Usa Saxml

Carga el modelo

Entrega el modelo

Anula la publicación del modelo

Limpia

Borra los recursos implementados

¿Qué sigue?