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:
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.
-
In the Google Cloud console, on the project selector page, select or create a Google Cloud project.
-
Asegúrate de que la facturación esté habilitada para tu proyecto de Google Cloud.
-
Habilita la API necesaria.
-
In the Google Cloud console, on the project selector page, select or create a Google Cloud project.
-
Asegúrate de que la facturación esté habilitada para tu proyecto de Google Cloud.
-
Habilita la API necesaria.
-
Asegúrate de tener los siguientes roles en el proyecto: roles/container.admin, roles/iam.serviceAccountAdmin
Verifica los roles
-
En la consola de Google Cloud, ve a la página IAM.
Ir a IAM - Selecciona el proyecto.
-
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.
- 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
-
En la consola de Google Cloud, ve a la página IAM.
Ir a IAM - Selecciona el proyecto.
- Haz clic en Grant access.
- En el campo Principales nuevas, ingresa tu dirección de correo electrónico.
- En la lista Seleccionar un rol, elige un rol.
- Para otorgar funciones adicionales, haz clic en Agregar otro rol y agrega cada rol adicional.
- 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 ShellConfigura 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
: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
: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
ycompletions: 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
:Reemplaza
BUCKET_NAME
por el nombre de tu bucket de Cloud Storage.Aplica el manifiesto
sax-http.yaml
:kubectl apply -f sax-http.yaml
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.