Disponibilizar um LLM usando TPUs de vários hosts no GKE com o Saxml

Padrão

Este tutorial mostra como implantar e disponibilizar um modelo de linguagem grande (LLM) usando um pool de nós de fração de TPU com vários hosts em o Google Kubernetes Engine (GKE) com Saxml para uma arquitetura escalonável e eficiente.

Contexto

O Saxml é um sistema experimental que atende aos frameworks Paxml, JAX e PyTorch. É possível usar TPUs para acelerar o processamento de dados com esses frameworks. Para demonstrar a implantação de TPUs no GKE, este tutorial disponibiliza o modelo de teste LmCloudSpmd175B32Test 175B. O GKE implanta esse modelo de teste em dois pools de nós de fração da TPU v5e com a topologia 4x8, respectivamente.

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

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

Ao implantar TPUs no GKE, é importante selecionar a topologia de TPU correta para disponibilizar um modelo. Para saber mais, consulte Planejar a configuração da TPU.

Objetivos

Este tutorial é destinado a engenheiros de MLOps ou DevOps ou administradores de plataforma que querem usar os recursos de orquestração do GKE para disponibilizar modelos de dados.

Este tutorial inclui as etapas a seguir:

Prepare seu ambiente com um cluster do GKE Standard. O cluster tem dois pools de nós de fatia da TPU v5e com a topologia 4x8.
Implante o Saxml. O Saxml precisa de um servidor administrador, um grupo de pods que funcionam como o servidor de modelos, um servidor HTTP pré-criado e um balanceador de carga.
Use o Saxml para disponibilizar o LLM.

O seguinte diagrama mostra a arquitetura que este tutorial implementa:

Arquitetura de uma TPU de vários hosts no GKE. — **Figura**: exemplo de arquitetura de uma TPU de vários hosts no GKE.

Antes de começar

Faça login na sua conta do Google Cloud. Se você começou a usar o Google Cloud agora, crie uma conta para avaliar o desempenho de nossos produtos em situações reais. Clientes novos também recebem US$ 300 em créditos para executar, testar e implantar cargas de trabalho.

No console do Google Cloud, na página do seletor de projetos, selecione ou crie um projeto do Google Cloud.

Acessar o seletor de projetos

Verifique se a cobrança está ativada para o seu projeto do Google Cloud.

Ative a API necessária.

Ative a API

No console do Google Cloud, na página do seletor de projetos, selecione ou crie um projeto do Google Cloud.

Acessar o seletor de projetos

Verifique se a cobrança está ativada para o seu projeto do Google Cloud.

Ative a API necessária.

Ative a API

Verifique se você tem os seguintes papéis no projeto: roles/container.admin, roles/iam.serviceAccountAdmin
Verificar os papéis
1. No console do Google Cloud, abra a página IAM.
  Acessar IAM
2. Selecionar um projeto.
3. Na coluna Principal, encontre a linha que contém seu endereço de e-mail.
  
  Caso seu endereço de e-mail não esteja nessa coluna, isso significa que você não tem papéis.
4. Na coluna Papel da linha com seu endereço de e-mail, verifique se a lista de papéis inclui os papéis necessários.
Conceder os papéis
1. No console do Google Cloud, abra a página IAM.
  Acesse o IAM
2. Selecionar um projeto.
3. Clique em CONCEDER ACESSO.
4. No campo Novos participantes, digite seu endereço de e-mail.
5. Na lista Selecionar um papel, escolha um.
6. Para conceder outros papéis, clique em Adicionar outro papel e adicione cada papel adicional.
7. Clique em Save.

Verifique se o projeto tem cota suficiente para a Cloud TPU no GKE.

prepare o ambiente

No console do Google Cloud, inicie uma instância do Cloud Shell:
Abrir o Cloud Shell
Defina as variáveis de ambiente padrão:
```
  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
```
Substitua os seguintes valores:
- PROJECT_ID: o ID do projeto do Google Cloud.
- COMPUTE_REGION: a região do Compute Engine.
- COMPUTE_ZONE: a zona em que ct5lp-hightpu-4t está disponível.

Criar um cluster do GKE Standard

Use o Cloud Shell para fazer o seguinte:

Crie um cluster padrão que use a federação de identidade da carga de trabalho para o GKE:
```
gcloud container clusters create saxml \
    --zone=${ZONE} \
    --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 dá suporte à TPU v5e versão 1.27.2-gke.2100 e a versões mais recentes. Para mais informações, consulte Disponibilidade da TPU no GKE.

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

Crie o primeiro pool de nós chamado tpu1:

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

Crie o segundo pool de nós chamado tpu2:

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

Você criou os seguintes recursos:

Um cluster Standard com quatro nós de CPU.
Dois pools de nós de fatia da TPU v5e com a topologia 4x8. Cada pool de nós representa oito Nós de fração de TPU com 4 chips de TPU cada.

O modelo 175B precisa ser disponibilizado em uma fração da TPU v5e de vários hosts com no mínimo a fração de topologia 4x8 (32 chips da TPU v5e).

Crie um bucket do Cloud Storage

Crie um bucket do Cloud Storage para armazenar as configurações do servidor administrador do Saxml. Um servidor administrador em execução salva periodicamente o próprio estado e os detalhes dos modelos publicados.

No Cloud Shell, execute o seguinte:

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

Configurar o acesso das cargas de trabalho usando a federação de identidade da carga de trabalho para o GKE

Atribua uma conta de serviço do Kubernetes ao aplicativo e configure-a para atuar como uma conta de serviço do IAM.

Configure kubectl para se comunicar com o cluster:

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

Crie uma conta de serviço do Kubernetes para o aplicativo usar:
```
kubectl create serviceaccount sax-sa --namespace default
```
Crie uma conta de serviço do IAM para seu aplicativo:
```
gcloud iam service-accounts create sax-iam-sa
```

Adicione uma vinculação de política do IAM para que sua conta de serviço do IAM leia e grave 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

Permita que a conta de serviço do Kubernetes atue como a conta de serviço do IAM adicionando uma vinculação de política do IAM entre as duas. Essa vinculação permite que a conta de serviço do Kubernetes atue como a conta de serviço do IAM, para poder fazer leituras e gravações 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]"
```
Anote a conta de serviço do Kubernetes com o endereço de e-mail da conta de serviço do IAM Isso permite que seu aplicativo de amostra identifique qual conta de serviço usar para acessar os serviços do Google Cloud. Portanto, quando o aplicativo usa qualquer biblioteca de cliente padrão da API do Google para acessar os serviços do Google Cloud, ele usa essa conta de serviço do IAM.
```
kubectl annotate serviceaccount sax-sa \
  iam.gke.io/gcp-service-account=sax-iam-sa@${PROJECT_ID}.iam.gserviceaccount.com
```

Implantar o Saxml

Nesta seção, você vai implantar o servidor administrador do Saxml e o servidor de modelos do Saxml.

Implantar o servidor administrador do Saxml

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

Substitua BUCKET_NAME pelo nome do bucket do Cloud Storage.

Aplique o manifesto:
```
kubectl apply -f sax-admin-server.yaml
```

Verifique se o pod do servidor administrador está em execução:

kubectl get deployment

O resultado será assim:

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

Implantar o servidor de modelos do Saxml

As cargas de trabalho executadas em frações da TPU de vários hosts exigem um identificador de rede estável para que cada pod descubra pares na mesma fração de TPU. Para definir esses identificadores, use IndexedJob, StatefulSet com um serviço headless ou JobSet, que cria automaticamente um serviço headless para todos os jobs que pertencem ao JobSet. Na seção a seguir, você vai aprender a gerenciar vários grupos de pods do servidor de modelos com o JobSet.

Instale o JobSet v0.2.3 ou mais recente.

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.

Valide o controlador de JobSet em execução no namespace jobset-system:

kubectl get pod -n jobset-system

O resultado será assim:

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

Implante dois servidores de modelos em dois pools de nós da fração TPU. Salve 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

Substitua BUCKET_NAME pelo nome do bucket do Cloud Storage.

Nesse manifesto:

replicas: 2 é o número de réplicas de jobs. Cada job representa um servidor de modelos. Portanto, um grupo de oito pods.
parallelism: 8 e completions: 8 são iguais ao número de nós em cada pool de nós.
backoffLimit: 0 precisa ser zero para que o job seja marcado com uma indicação de falha em caso de falha em um pod.
ports.containerPort: 8471 é a porta padrão para a comunicação de VMs.
name: MEGASCALE_NUM_SLICES cancela a configuração da variável de ambiente porque o GKE não está executando o treinamento multislice.

Aplique o manifesto:

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

Verifique o status do servidor administrador do Saxml e dos pods do servidor de modelos:

kubectl get pods

O resultado será assim:

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, há 16 contêineres 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 primários em cada grupo.

O cluster do Saxml tem dois servidores de modelos implantados em dois pools de nós de fração de TPU v5e com a topologia 4x8, respectivamente.

Implantar o servidor HTTP e o balanceador de carga do Saxml

Use a imagem pré-criada do servidor HTTP a seguir. Salve 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

Substitua BUCKET_NAME pelo nome do bucket do Cloud Storage.

Aplique o manifesto sax-http.yaml:
```
kubectl apply -f sax-http.yaml
```
Observação: é possível usar seu próprio servidor HTTP criado para Saxml. Para saber mais, consulte Inferência com o Saxml e um servidor HTTP.

Aguarde a conclusão da criação do contêiner do servidor HTTP:

kubectl get pods

O resultado será assim:

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

Aguarde até que um endereço IP externo seja atribuído:

kubectl get svc

O resultado será assim:

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

Usar o Saxml

Carregue, implante e disponibilize o modelo no Saxml na fração de vários hosts da TPU v5e:

Carregar o modelo

Recupere o endereço IP do balanceador de carga para Saxml.

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

Carregue o modelo de teste LmCloudSpmd175B em dois pools de nós da fração 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 checkpoint ajustado. Os pesos são gerados aleatoriamente. O carregamento do modelo pode levar até 10 minutos.

O resultado será assim:

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

Verifique a prontidão do modelo:

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

O resultado será assim:

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

O modelo está totalmente carregado.

Receba 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 será assim:

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

Disponibilizar o modelo

Realize uma solicitação 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?"
}'

A saída mostra um exemplo da resposta do modelo. Essa resposta pode não ser significativa porque o modelo de teste tem pesos aleatórios.

Cancelar a publicação do modelo

Execute o seguinte comando para cancelar 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 será assim:

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

Limpar

Para evitar cobranças na sua conta do Google Cloud pelos recursos usados no tutorial, exclua o projeto que os contém ou mantenha o projeto e exclua os recursos individuais.

Excluir os recursos implantados

Exclua o cluster criado para este tutorial:

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

Exclua a conta de serviço:

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

Exclua o bucket do Cloud Storage:
```
gcloud storage rm -r gs://${GSBUCKET}
```

A seguir

Saiba mais sobre as versões atuais da TPU com a arquitetura do sistema da Cloud TPU.
Saiba mais sobre TPUs no GKE.

Disponibilizar um LLM usando TPUs de vários hosts no GKE com o Saxml

Contexto

Objetivos

Antes de começar

Verificar os papéis

Conceder os papéis

prepare o ambiente

Criar um cluster do GKE Standard

Crie um bucket do Cloud Storage

Configurar o acesso das cargas de trabalho usando a federação de identidade da carga de trabalho para o GKE

Implantar o Saxml

Implantar o servidor administrador do Saxml

Implantar o servidor de modelos do Saxml

Implantar o servidor HTTP e o balanceador de carga do Saxml

Usar o Saxml

Carregar o modelo

Disponibilizar o modelo

Cancelar a publicação do modelo

Limpar

Excluir os recursos implantados

A seguir