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

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:

  1. 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.
  2. 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.
  3. 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

  • 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

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

  • Verify 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, roles/iam.policyAdmin

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

      Acessar o IAM
    2. Selecione o projeto.
    3. Clique em Conceder acesso.

    4. No campo Novos principais, digite seu identificador de usuário. Normalmente, é o endereço de e-mail de uma Conta do Google.

    5. Na lista Selecionar papel, escolha um.
    6. Para conceder outros papéis, clique em Adicionar outro papel e adicione cada papel adicional.
    7. Clique em Salvar.

      8. Prepare o ambiente

      1. No Google Cloud console, inicie uma instância do Cloud Shell:
        Abrir o Cloud Shell

      2. Defina as variáveis de ambiente padrão:

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

        Substitua os seguintes valores:

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

      Criar um cluster do GKE Standard

      Use o Cloud Shell para fazer o seguinte:

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

      2. Crie o primeiro pool de nós chamado tpu1:

        gcloud container node-pools create tpu1 \
    --zone=${ZONE} \
    --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 fração de TPU. Nesse caso: (4 * 8) / 4.

      3. Crie o segundo pool de nós chamado tpu2:

        gcloud container node-pools create tpu2 \
    --zone=${ZONE} \
    --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 fração de TPU. Nesse caso: (4 * 8) / 4.

      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://${BUCKET_NAME}

      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.

      1. Configure kubectl para se comunicar com o cluster:

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

      2. Crie uma conta de serviço do Kubernetes para o aplicativo usar:

        kubectl create serviceaccount sax-sa --namespace default

      3. Crie uma conta de serviço do IAM para seu aplicativo:

        gcloud iam service-accounts create sax-iam-sa

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

      5. 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]"

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

      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 bucket do Cloud Storage 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 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. Um Jobset é uma API de carga de trabalho que permite gerenciar um grupo de jobs do Kubernetes como uma unidade. O caso de uso mais comum para um JobSet é o treinamento distribuído, mas também é possível usá-lo para executar cargas de trabalho em lote.

      Na seção a seguir, você vai aprender a gerenciar vários grupos de pods do servidor de modelos com o JobSet.

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

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

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

      4. Substitua BUCKET_NAME pelo bucket do Cloud Storage criado anteriormente:

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

        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.

      5. Aplique o manifesto:

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

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

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

      2. Substitua BUCKET_NAME pelo bucket do Cloud Storage 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

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

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

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

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

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

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

      1. Exclua o cluster criado para este tutorial:

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

      2. Exclua a conta de serviço:

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

      3. Exclua o bucket do Cloud Storage:

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

      A seguir