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

  • Faça login na sua conta do Google Cloud . Se você começou a usar o Google Cloud, 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.

  • In the Google Cloud console, on the project selector page, select or create a Google Cloud project.

    Roles required to select or create a project

    • Select a project: Selecting a project doesn't require a specific IAM role—you can select any project that you've been granted a role on.
    • Create a project: To create a project, you need the Project Creator role (roles/resourcemanager.projectCreator), which contains the resourcemanager.projects.create permission. Learn how to grant roles.

    Go to project selector

  • Verify that billing is enabled for your Google Cloud project.

  • Enable the required API.

    Roles required to enable APIs

    To enable APIs, you need the Service Usage Admin IAM role (roles/serviceusage.serviceUsageAdmin), which contains the serviceusage.services.enable permission. Learn how to grant roles.

    Enable the API

  • In the Google Cloud console, on the project selector page, select or create a Google Cloud project.

    Roles required to select or create a project

    • Select a project: Selecting a project doesn't require a specific IAM role—you can select any project that you've been granted a role on.
    • Create a project: To create a project, you need the Project Creator role (roles/resourcemanager.projectCreator), which contains the resourcemanager.projects.create permission. Learn how to grant roles.

    Go to project selector

  • Verify that billing is enabled for your Google Cloud project.

  • Enable the required API.

    Roles required to enable APIs

    To enable APIs, you need the Service Usage Admin IAM role (roles/serviceusage.serviceUsageAdmin), which contains the serviceusage.services.enable permission. Learn how to grant roles.

    Enable the API

    •

  • Verifique se você tem os seguintes papéis no projeto: roles/container.admin, roles/iam.serviceAccountAdmin, roles/iam.policyAdmin

    Verificar os papéis

    1. No console do Google Cloud , acesse a página IAM.

      Acessar IAM
    2. Selecione o projeto.

    3. Na coluna Principal, encontre todas as linhas que identificam você ou um grupo no qual você está incluído. Para saber em quais grupos você está incluído, entre em contato com o administrador.

    4. Em todas as linhas que especificam ou incluem você, verifique a coluna Papel para ver se a lista de papéis inclui os papéis necessários.

    Conceder os papéis

    1. No console do Google Cloud , acesse a página IAM.

      Acessar 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. Clique em Selecionar um papel e pesquise o papel.
    6. Para conceder outros papéis, adicione-os clicando em Adicionar outro papel.
    7. Clique em Salvar.

Prepare o ambiente

  1. No console do Google Cloud , 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 CONTROL_PLANE_LOCATION=CONTROL_PLANE_LOCATION
  export BUCKET_NAME=PROJECT_ID-gke-bucket

    Substitua os seguintes valores:

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

Criar um cluster

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 \
    --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 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 \
    --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 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 \
    --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 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 --location=${CONTROL_PLANE_LOCATION}

  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 app de exemplo 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 --location ${CONTROL_PLANE_LOCATION}

  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