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

Padrão

Este tutorial mostra como implementar e publicar um modelo de linguagem (conteúdo extenso) (MDI/CE) usando o conjunto de nós de segmentação de TPU de 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 node pools 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. Uma vez que o modelo de N mil milhões de bits de 16 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.

Objetivos

Este tutorial destina-se a engenheiros de MLOps ou DevOps, ou administradores de plataformas, que querem usar as capacidades de orquestração do GKE para publicar modelos de dados.

Este tutorial abrange os seguintes passos:

Prepare o seu ambiente com um cluster padrão. O cluster tem dois pools de nós de fatias de TPU v5e com topologia 4x8.
Implemente o Saxml. O Saxml precisa de um servidor de 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 publicar o LLM.

O diagrama seguinte mostra a arquitetura que o seguinte tutorial implementa:

Arquitetura de uma TPU multi-host no GKE. — **Figura:** exemplo de arquitetura de uma TPU com vários anfitriões 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.

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

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

In the Google Cloud console, go to the IAM page.
Go to IAM
Select the project.
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.
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

In the Google Cloud console, go to the IAM page.
Aceder ao IAM
Selecione o projeto.
Clique em Conceder acesso.
No campo Novos responsáveis, introduza o identificador do utilizador. Normalmente, este é o endereço de email de uma Conta Google.
Clique em Selecionar uma função e, de seguida, pesquise a função.
Para conceder funções adicionais, clique em Adicionar outra função e adicione cada função adicional.
Clique em Guardar.

Certifique-se de que o seu projeto tem quota suficiente para o Cloud TPU no GKE.

Prepare o ambiente

Na Google Cloud consola, inicie uma instância do Cloud Shell:
Abrir Cloud Shell
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

Use o Cloud Shell para fazer o seguinte:

Crie um cluster Standard que use a Federação de identidades de cargas 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 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.

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.

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

Configure kubectl para comunicar com o seu cluster:

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

Crie uma ServiceAccount do Kubernetes para a sua aplicação usar:
```
kubectl create serviceaccount sax-sa --namespace default
```
Crie uma conta de serviço do IAM para a sua aplicação:
```
gcloud iam service-accounts create sax-iam-sa
```

Adicione uma associação de políticas de IAM para a sua conta de serviço do IAM 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

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]"
```
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 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

Crie o seguinte sax-admin-server.yaml manifesto:

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 armazenamento na nuvem criado anteriormente:

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

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

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 cada Pod para descobrir 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.

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.

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

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

Implemente dois servidores de modelos em dois conjuntos de nós de fatias de TPUs. 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

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 agrupamentos.
- 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
- O comando name: MEGASCALE_NUM_SLICES anula a definição da variável de ambiente porque o GKE não está a executar o treino com vários fragmentos.

Aplique o manifesto:

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

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

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

Substitua BUCKET_NAME pelo armazenamento na nuvem criado anteriormente:
```
perl -pi -e 's|BUCKET_NAME|BUCKET_NAME|g' sax-http.yaml
```
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.

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

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

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

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
}

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.

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

Limpar

Para evitar incorrer em custos na sua conta do Google Cloud pelos recursos usados neste tutorial, elimine o projeto que contém os recursos ou mantenha o projeto e elimine os recursos individuais.

Elimine os recursos implementados

Elimine o cluster que criou para este tutorial:

gcloud container clusters delete saxml --location ${CONTROL_PLANE_LOCATION}

Elimine a conta de serviço:

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

Elimine o contentor do Cloud Storage:

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

O que se segue?

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

Publique um MDG com TPUs multi-anfitrião no GKE com o Saxml Mantenha tudo organizado com as coleções Salve e categorize o conteúdo com base nas suas preferências.

Contexto

Objetivos

Antes de começar

Check for the roles

Grant the roles

Prepare o ambiente

Crie um cluster

Crie um contentor do Cloud Storage

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

Implemente o Saxml

Implemente o servidor de administrador do SAXML

Implemente o servidor de modelos Saxml

Implemente o servidor HTTP Saxml e o balanceador de carga

Use o Saxml

Carregue o modelo

Publique o modelo

Anule a publicação do modelo

Limpar

Elimine os recursos implementados

O que se segue?

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