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

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 Cloud de Confiance 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:

    Neste comando, BUCKET_NAME especifica o nome do Cloud de Confiance 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}.s3ns.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}.s3ns.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}.s3ns.iam.gserviceaccount.com \
  --role roles/iam.workloadIdentityUser \
  --member "serviceAccount:${PROJECT_ID}.s3ns.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 Cloud de Confiance . 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. Cloud de Confiance 

    kubectl annotate serviceaccount sax-sa \
  iam.gke.io/gcp-service-account=sax-iam-sa@${PROJECT_ID}.s3ns.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

  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.

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