Publique um MDG com TPUs multi-anfitrião no GKE com o Saxml

Aviso: a prioridade do Saxml no GKE vai ser reduzida a partir de 24 de abril de 2025. Isto significa que o projeto não vai receber mais atualizações. As implementações do Saxml existentes vão continuar a funcionar sem interrupções. Recomendamos vivamente que migre para o JetStream, o framework de inferência de código aberto atualizado da Google para a publicação de LLMs de alto desempenho em TPUs e GPUs. O JetStream oferece processamento em lote e quantização contínuos para um melhor débito e eficiência de memória. Para ver um exemplo de migração, consulte o artigo Apresente o Gemma com TPUs no GKE com o JetStream.

Este tutorial mostra como implementar e publicar um modelo de linguagem (conteúdo extenso) (MDI/CE) usando um conjunto de nós de fatia de TPU com vários anfitriões no Google Kubernetes Engine (GKE) com o Saxml para uma arquitetura escalável eficiente.

Contexto

O Saxml é um sistema experimental que serve frameworks Paxml, JAX e PyTorch. Pode usar as TPUs para acelerar o processamento de dados com estas frameworks. Para demonstrar a implementação de TPUs no GKE, este tutorial disponibiliza o modelo de teste LmCloudSpmd175B32Test de 175 mil milhões de parâmetros. O GKE implementa este modelo de teste em dois conjuntos de nós de fatias de TPU v5e com topologia 4x8, respetivamente.

Para implementar corretamente o modelo de teste, a topologia da TPU foi definida com base no tamanho do modelo. Dado que o modelo de N mil milhões de bits requer aproximadamente 2 vezes (2xN) GB de memória, o modelo LmCloudSpmd175B32Test de 175 mil milhões requer cerca de 350 GB de memória. O chip de TPU único v5e tem 16 GB. Para suportar 350 GB, o GKE precisa de 21 chips de TPU v5e (350/16= 21). Com base no mapeamento da configuração da TPU, a configuração da TPU adequada para este tutorial é:

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

Selecionar a topologia de TPU certa para publicar um modelo é importante quando implementa TPUs no GKE. Para saber mais, consulte o artigo Planeie a configuração da TPU.

Prepare o ambiente

1. Na Google Cloud consola, inicie uma instância do Cloud Shell:
   Abrir Cloud Shell

2. Defina as variáveis de ambiente predefinidas:

     gcloud config set project PROJECT_ID
     export PROJECT_ID=$(gcloud config get project)
     export CONTROL_PLANE_LOCATION=CONTROL_PLANE_LOCATION
     export BUCKET_NAME=PROJECT_ID-gke-bucket
   

   Substitua os seguintes valores:

   • PROJECT_ID: o seu Google Cloud ID do projeto.
   • CONTROL_PLANE_LOCATION: a zona do Compute Engine do plano de controlo do seu cluster. Selecione a zona onde o ct5lp-hightpu-4t está disponível.

   Neste comando, BUCKET_NAME especifica o nome do Google Cloud contentor de armazenamento para armazenar as configurações do servidor de administrador do Saxml.

Crie um cluster padrão do GKE

Use o Cloud Shell para fazer o seguinte:

1. Crie um cluster Standard que use a Workload Identity Federation para o GKE:

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

   Substitua VERSION pelo número da versão do GKE. O GKE suporta a TPU v5e na versão 1.27.2-gke.2100 e posteriores. Para mais informações, consulte o artigo Disponibilidade de TPUs no GKE.

   A criação do cluster pode demorar vários minutos.

2. Crie o primeiro node pool com o nome tpu1:

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

   O valor da flag --num-nodes é calculado dividindo a topologia da TPU pelo número de chips de TPU por fatia de TPU. Neste caso: (4 * 8) / 4.

3. Crie o segundo node pool com o nome tpu2:

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

   O valor da flag --num-nodes é calculado dividindo a topologia da TPU pelo número de chips de TPU por fatia de TPU. Neste caso: (4 * 8) / 4.

Criou os seguintes recursos:

• Um cluster padrão com quatro nós de CPU.
• Dois conjuntos de nós de fatias de TPU v5e com topologia 4x8. Cada conjunto de nós representa oito nós de fatias de TPU com 4 chips de TPU cada.

O modelo de 175 mil milhões de parâmetros tem de ser publicado numa fatia de TPU v5e com vários anfitriões com, no mínimo, uma 4x8 fatia de topologia (32 chips de TPU v5e).

Crie um contentor do Cloud Storage

Crie um contentor do Cloud Storage para armazenar as configurações do servidor de administrador do Saxml. Um servidor de administrador em execução guarda periodicamente o respetivo estado e os detalhes dos modelos publicados.

No Cloud Shell, execute o seguinte:

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

Configure o acesso das suas cargas de trabalho através da federação de identidades de cargas de trabalho para o GKE

Atribua uma conta de serviço do Kubernetes à aplicação e configure essa conta de serviço do Kubernetes para funcionar como uma conta de serviço do IAM.

1. Configure kubectl para comunicar com o cluster:

   gcloud container clusters get-credentials saxml --location=${CONTROL_PLANE_LOCATION}
   

2. Crie uma ServiceAccount do Kubernetes para a sua aplicação usar:

   kubectl create serviceaccount sax-sa --namespace default
   

3. Crie uma conta de serviço do IAM para a sua aplicação:

   gcloud iam service-accounts create sax-iam-sa
   

4. Adicione uma associação de políticas de IAM para a sua conta de serviço de IAM poder ler e escrever no 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. Permita que a conta de serviço do Kubernetes use a identidade da conta de serviço da IAM adicionando uma associação de políticas da IAM entre as duas contas de serviço. Esta associação permite que a conta de serviço do Kubernetes atue como a conta de serviço do IAM, para que a conta de serviço do Kubernetes possa ler e escrever no 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. Anotar a conta de serviço do Kubernetes com o endereço de email da conta de serviço do IAM. Isto permite que a sua app de exemplo saiba que conta de serviço usar para aceder aos serviços Google Cloud . Assim, quando a app usa quaisquer bibliotecas cliente de APIs Google padrão para aceder a serviços, usa essa conta de serviço da IAM. Google Cloud

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

Implemente o Saxml

Nesta secção, implementa o servidor de administrador do Saxml e o servidor do modelo Saxml.

Implemente o servidor de administrador do SAXML

1. Crie o seguinte manifesto 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

2. Substitua BUCKET_NAME pelo armazenamento na nuvem criado anteriormente:

   perl -pi -e 's|BUCKET_NAME|BUCKET_NAME|g' sax-admin-server.yaml
   

3. Aplique o manifesto:

   kubectl apply -f sax-admin-server.yaml
   

4. Verifique se o pod do servidor de administrador está em funcionamento:

   kubectl get deployment
   

   O resultado é semelhante ao seguinte:

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

Implemente o servidor de modelos Saxml

As cargas de trabalho executadas em fatias de TPU com vários anfitriões requerem um identificador de rede estável para que cada Pod descubra pares na mesma fatia de TPU. Para definir estes identificadores, use IndexedJob, StatefulSet com um serviço sem interface ou JobSet, que cria automaticamente um serviço sem interface para todos os trabalhos que pertencem ao JobSet. Um Jobset é uma API de carga de trabalho que lhe permite gerir um grupo de tarefas do Kubernetes como uma unidade. O exemplo de utilização mais comum de um JobSet é o treino distribuído, mas também o pode usar para executar cargas de trabalho em lote.

A secção seguinte mostra como gerir vários grupos de pods do servidor de modelos com o JobSet.

1. Instale o JobSet v0.2.3 ou posterior.

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

   Substitua JOBSET_VERSION pela versão do JobSet. Por exemplo, v0.2.3.

2. Valide se o controlador JobSet está a ser executado no jobset-system espaço de nomes:

   kubectl get pod -n jobset-system
   

   O resultado é semelhante ao seguinte:

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

3. Implemente dois servidores de modelos em dois conjuntos de nós de fatias de TPU. Guarde o seguinte manifesto 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

4. Substitua BUCKET_NAME pelo armazenamento na nuvem criado anteriormente:

   perl -pi -e 's|BUCKET_NAME|BUCKET_NAME|g' sax-model-server-set.yaml
   

   Neste manifesto:

   • replicas: 2 é o número de réplicas de tarefas. Cada tarefa representa um servidor de modelos. Por isso, um grupo de 8 Pods.
   • parallelism: 8 e completions: 8 são iguais ao número de nós em cada conjunto de nós.
   • backoffLimit: 0 tem de ser zero para marcar a tarefa como falhada se algum pod falhar.
   • A porta ports.containerPort: 8471 é a porta predefinida para a comunicação das VMs
   • name: MEGASCALE_NUM_SLICES anula a definição da variável de ambiente porque o GKE não está a executar o treino de vários fragmentos.

5. Aplique o manifesto:

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

6. Verifique o estado dos pods do servidor de administração e do servidor de modelos do Saxml:

   kubectl get pods
   

   O resultado é semelhante ao seguinte:

   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
   

Neste exemplo, existem 16 contentores de servidores de modelos: sax-model-server-set-sax-model-server-0-0-nj4sm e sax-model-server-set-sax-model-server-1-0-8mlf5 são os dois servidores de modelos principais em cada grupo.

O seu cluster Saxml tem dois servidores de modelos implementados em dois node pools de fatias de TPU v5e com a topologia 4x8, respetivamente.

Implemente o servidor HTTP Saxml e o balanceador de carga

1. Use a seguinte imagem do servidor HTTP de imagem pré-criada. Guarde o seguinte manifesto 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

2. Substitua BUCKET_NAME pelo armazenamento na nuvem criado anteriormente:

   perl -pi -e 's|BUCKET_NAME|BUCKET_NAME|g' sax-http.yaml
   

3. Aplique o manifesto sax-http.yaml:

   kubectl apply -f sax-http.yaml
   

   Nota: pode usar o seu próprio servidor HTTP criado para o Saxml. Para saber mais, consulte o artigo Inferência com o Saxml e um servidor HTTP.

4. Aguarde pela conclusão da criação do contentor do servidor HTTP:

   kubectl get pods
   

   O resultado é semelhante ao seguinte:

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

5. Aguarde até que o serviço tenha um endereço IP externo atribuído:

   kubectl get svc
   

   O resultado é semelhante ao seguinte:

   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
   

Use o Saxml

Carregue, implemente e publique o modelo no Saxml na fatia de vários anfitriões da TPU v5e:

Carregue o modelo

1. Obtenha o endereço IP do balanceador de carga para o Saxml.

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

2. Carregue o modelo de teste LmCloudSpmd175B em dois conjuntos de nós de fatias 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
   }'
   

   O modelo de teste não tem um ponto de verificação otimizado. Os pesos são gerados aleatoriamente. O carregamento do modelo pode demorar até 10 minutos.

   O resultado é semelhante ao seguinte:

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

3. Verifique a prontidão do modelo:

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

   O resultado é semelhante ao seguinte:

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

   O modelo está totalmente carregado.

4. Obtenha informações sobre o modelo:

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

   O resultado é semelhante ao seguinte:

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

Publique o modelo

Apresentar um pedido de comando:

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?"
}'

O resultado mostra um exemplo da resposta do modelo. Esta resposta pode não ser significativa porque o modelo de teste tem ponderações aleatórias.

Êxito: publicou com êxito um MDG usando TPUs multi-host no GKE.

Anule a publicação do modelo

Execute o seguinte comando para anular a publicação do modelo:

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

O resultado é semelhante ao seguinte:

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