Disponibilizar modelos abertos do Gemma usando TPUs no GKE com Saxml

Padrão

Neste guia, mostramos como exibir um modelo de linguagem grande (LLM) de modelos abertos do Gemma usando Unidades de Processamento de Tensor (TPUs) no Google Kubernetes Engine (GKE) com o Saxml. Neste guia, faça o download dos modelos Gemma ajustados por instruções de parâmetros 2B e 7B para o Cloud Storage e os implante em um cluster do GKE Standard usando contêineres que executam o Saxml.

Este guia é um bom ponto de partida se você precisar da escalonabilidade, resiliência e economia oferecidos pelos recursos do Kubernetes ao implantar o modelo no Saxml. Se você precisa de uma plataforma de IA gerenciada unificada para criar e disponibilizar modelos de ML rapidamente e de maneira econômica, recomendamos testar nossa solução de implantação da Vertex AI.

Experiência

Ao disponibilizar o Gemma usando TPUs no GKE com o Saxml, é possível implementar uma solução de veiculação de inferência robusta e pronta para produção com todos os benefícios do Kubernetes., incluindo escalonabilidade eficiente e maior disponibilidade. Esta seção descreve as principais tecnologias usadas neste tutorial.

Gemma

O Gemma é um conjunto de modelos de IA generativa leve e abertamente lançados com licença aberta. Esses modelos de IA estão disponíveis para execução em aplicativos, hardware, dispositivos móveis ou serviços hospedados. É possível usar os modelos Gemma para geração de texto, mas também é possível ajustá-los para tarefas especializadas.

Para saber mais, consulte a documentação do Gemma.

TPUs

TPUs são circuitos integrados de aplicação específica (ASICs, na sigla em inglês) desenvolvidos especialmente pelo Google. Eles são usados para acelerar frameworks de processamento de dados, como TensorFlow, PyTorch e JAX.

Antes de usar TPUs no GKE, recomendamos que você conclua o seguinte programa de aprendizado:

Saiba mais sobre a disponibilidade atual da versão da TPU com a arquitetura do sistema do Cloud TPU.
Saiba mais sobre TPUs no GKE.

Este tutorial serve os modelos Gemma 2B e Gemma 7B. O GKE hospeda esses modelos nos seguintes pools de nós da TPU v5e de host único:

Gemma 2B: modelo ajustado por instruções hospedado em um pool de nós da TPU v5e com topologia 1x1 que representa um chip de TPU. O tipo de máquina dos nós é ct5lp-hightpu-1t.
Gemma 7B: modelo ajustado por instruções hospedado em um pool de nós da TPU v5e com topologia 2x2 que representa quatro chips de TPU. O tipo de máquina dos nós é ct5lp-hightpu-4t.

Saxml

O Saxml é um sistema experimental que exibe modelos Paxml, JAX e PyTorch para inferência. O sistema Saxml é composto pelos seguintes componentes:

Célula Saxml ou cluster Sax: composto por um servidor administrador e um grupo de servidores de modelo. O servidor de administrador monitora os servidores de modelos, atribui modelos publicados a servidores de modelo para exibição e ajuda os clientes a localizar servidores de modelos que atendem a modelos publicados específicos.
Cliente Saxml: a interface de programação voltada ao usuário para o sistema Saxml. O cliente Saxml inclui uma ferramenta de linha de comando (saxutil) e um conjunto de bibliotecas de cliente em Python, C++ e Go.

Neste tutorial, você também vai usar o servidor HTTP Saxml. O servidor HTTP Saxml é um servidor HTTP personalizado que encapsula a biblioteca de cliente Python do Saxml e expõe APIs REST para interagir com o sistema Saxml, incluindo endpoints para publicar, listar, cancelar a publicação de modelos e gerar previsões.

Objetivos

Este tutorial é destinado a clientes de IA generativa que usam o JAX, usuários novos ou atuais do GKE, engenheiros de ML, engenheiros de MLOps (DevOps) ou administradores de plataformas interessados em usar os recursos de orquestração de contêineres do Kubernetes para veiculação Gemma.

Este tutorial abrange as etapas a seguir:

Prepare um cluster do GKE Standard com a topologia de TPU recomendada com base nas características do modelo.
Implantar componentes Saxml no GKE.
Conseguir e publicar o modelo de parâmetro Gemma 2B ou Gemma 7B.
Disponibilizar e interagir com os modelos publicados.

Arquitetura

Nesta seção, descrevemos a arquitetura do GKE usada neste tutorial. A arquitetura compreende um cluster do GKE Standard que provisiona TPUs e hospeda componentes Saxml para implantar e disponibilizar modelos Gemma 2B ou 7B. O diagrama a seguir mostra os componentes dessa arquitetura:

Um diagrama da arquitetura implantada neste tutorial

Essa arquitetura inclui os seguintes componentes:

Um cluster zonal do GKE Standard.
Um pool de nós de fração de TPU de host único que depende do modelo Gemma que você quer disponibilizar:
- Gemma 2B: configurado com uma TPU v5e com uma topologia 1x1. Uma instância do servidor do modelo Saxml está configurada para usar esse pool de nós.
- Gemma 7B: configurado com uma TPU v5e com uma topologia 2x2. Uma instância do servidor do modelo Saxml está configurada para usar esse pool de nós.
Um pool de nós de CPU padrão em que o servidor de administração Saxml e o servidor HTTP Saxml são implantados.
Dois buckets do Cloud Storage:
- Um bucket do Cloud Storage armazena o estado gerenciado por um servidor administrador.
- Um bucket do Cloud Storage armazena os checkpoints do modelo Gemma.

Essa arquitetura tem as seguintes características:

Um Artifact Registry público gerencia as imagens de contêineres para os componentes do Saxml.
O cluster do GKE usa a federação de identidade da carga de trabalho do GKE. Todos os componentes do Saxml usam uma federação de identidade da carga de trabalho que integra uma conta de serviço do IAM para acessar serviços externos, como buckets do Cloud Storage.
Os registros gerados pelos componentes do Saxml são integrados ao Cloud Logging.
Use o Cloud Monitoring para analisar as métricas de desempenho dos pools de nós do GKE criados por este tutorial.

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 você tem cotas suficientes para cinco chips TPU v5e. Neste tutorial, você usará as instâncias sob demanda.
Crie uma conta da Kaggle, caso ainda não tenha uma.

Preparar o ambiente para o Gemma

Iniciar o Cloud Shell

Neste tutorial, você usará o Cloud Shell para gerenciar recursos hospedados no Google Cloud. O Cloud Shell vem pré-instalado com o software necessário para este tutorial, incluindo kubectl e gcloud CLI.

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 LOCATION=LOCATION
export CLUSTER_NAME=saxml-tpu
```
Substitua os seguintes valores:
- PROJECT_ID: o ID do projeto do Google Cloud.
- LOCATION: o nome da zona do Compute Engine em que os tipos de máquina TPU v5e estão disponíveis.

Criar um cluster do GKE Standard

Nesta seção, você vai criar o cluster do GKE e o pool de nós.

Gemma 2B-it

Use o Cloud Shell para fazer o seguinte:

Crie um cluster padrão que use a federação de identidade da carga de trabalho do GKE:

gcloud container clusters create ${CLUSTER_NAME} \
    --enable-ip-alias \
    --machine-type=e2-standard-4 \
    --num-nodes=2 \
    --release-channel=rapid \
    --workload-pool=${PROJECT_ID}.svc.id.goog \
    --location=${LOCATION}

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

Crie um pool de nós de TPU v5e com uma topologia 1x1 e um nó:

gcloud container node-pools create tpu-v5e-1x1 \
    --cluster=${CLUSTER_NAME} \
    --machine-type=ct5lp-hightpu-1t \
    --num-nodes=1 \
    --location=${LOCATION}

Você exibe o modelo Gemma 2B neste pool de nós.

Gemma 7B-it

Use o Cloud Shell para fazer o seguinte:

Crie um cluster padrão que use a federação de identidade da carga de trabalho do GKE:

gcloud container clusters create ${CLUSTER_NAME} \
    --enable-ip-alias \
    --machine-type=e2-standard-4 \
    --num-nodes=2 \
    --release-channel=rapid \
    --workload-pool=${PROJECT_ID}.svc.id.goog \
    --location=${LOCATION}

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

Crie um pool de nós de TPU v5e com uma topologia 2x2 e um nó:

gcloud container node-pools create tpu-v5e-2x2 \
    --cluster=${CLUSTER_NAME} \
    --machine-type=ct5lp-hightpu-4t \
    --num-nodes=1 \
    --location=${LOCATION}

Você exibe o modelo Gemma 7B neste pool de nós.

Criar os buckets do Cloud Storage

Crie dois buckets do Cloud Storage para gerenciar o estado do servidor Saxml Admin e os checkpoints do modelo.

No Cloud Shell, execute o seguinte:

Crie um bucket do Cloud Storage para armazenar as configurações do servidor administrador do Saxml.
```
gcloud storage buckets create gs://ADMIN_BUCKET_NAME
```
Substitua ADMIN_BUCKET_NAME pelo nome do bucket do Cloud Storage que armazena o servidor Administrador do Saxml.
Crie um bucket do Cloud Storage para armazenar checkpoints do modelo:
```
gcloud storage buckets create gs://CHECKPOINTS_BUCKET_NAME
```
Substitua CHECKPOINTS_BUCKET_NAME pelo nome do bucket do Cloud Storage que armazena os checkpoints do modelo.

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

Atribua uma Kubernetes ServiceAccount 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 ${CLUSTER_NAME} --location=${LOCATION}

Crie uma conta de serviço do Kubernetes para o aplicativo usar:
```
gcloud iam service-accounts create wi-sax
```

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:wi-sax@${PROJECT_ID}.iam.gserviceaccount.com" \
    --role roles/storage.objectUser

gcloud projects add-iam-policy-binding ${PROJECT_ID} \
    --member "serviceAccount:wi-sax@${PROJECT_ID}.iam.gserviceaccount.com" \
    --role roles/storage.insightsCollectorService

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 Kubernetes ServiceAccount atue como a conta de serviço do Google.
```
gcloud iam service-accounts add-iam-policy-binding wi-sax@${PROJECT_ID}.iam.gserviceaccount.com \
    --role roles/iam.workloadIdentityUser \
    --member "serviceAccount:${PROJECT_ID}.svc.id.goog[default/default]"
```

Anote a conta de serviço do Kubernetes com o endereço de e-mail da conta de serviço do IAM

kubectl annotate serviceaccount default \
    iam.gke.io/gcp-service-account=wi-sax@${PROJECT_ID}.iam.gserviceaccount.com

Receber acesso ao modelo

Para ter acesso aos modelos Gemma para implantação no GKE, faça login na plataforma Kaggle, assine o contrato de consentimento de licença e receba um token de API da Kaggle. Neste tutorial, você usa um Secret do Kubernetes para as credenciais do Kaggle.

É necessário assinar o contrato de consentimento para usar o Gemma. Siga estas instruções:

Acesse a página de consentimento do modelo em Kaggle.com.
Faça login no Kaggle se ainda não tiver feito isso.
Clique em Solicitar acesso.
Na seção Escolher conta para consentimento, selecione Verificar via conta do Kaggle para usar sua conta do Kaggle para consentimento.
Aceite os Termos e Condições do modelo.

Gerar um token de acesso

Para acessar o modelo pelo Kaggle, você precisa de um token da API Kaggle.

Siga as etapas abaixo para gerar um novo token, caso ainda não tenha um:

No navegador, acesse as configurações da Kaggle.
Na seção API, clique em Criar novo token.

Um arquivo chamado kaggle.json é transferido por download.

Fazer upload do token de acesso no Cloud Shell

No Cloud Shell, é possível fazer upload do token da API Kaggle para seu projeto do Google Cloud:

No Cloud Shell, clique em Mais > Fazer upload.
Selecione "Arquivo" e clique em Escolher arquivos.
Abra o arquivo kaggle.json.
Clique em Fazer upload.

Criar um secret do Kubernetes para credenciais do Kaggle

No Cloud Shell, faça o seguinte:

Configure kubectl para se comunicar com o cluster:

gcloud container clusters get-credentials ${CLUSTER_NAME} --location=${LOCATION}

Crie um Secret para armazenar as credenciais do Kaggle:

kubectl create secret generic kaggle-secret \
    --from-file=kaggle.json

Implantar o Saxml

Nesta seção, você vai implantar o servidor de administrador Saxml, os servidores de modelo e o servidor HTTP.

Implantar o servidor de administrador Saxml

Crie o seguinte manifesto saxml-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
      containers:
      - name: sax-admin-server
        image: us-docker.pkg.dev/cloud-tpu-images/inference/sax-admin-server:v1.2.0
        securityContext:
          privileged: true
        ports:
        - containerPort: 10000
        env:
        - name: GSBUCKET
          value: ADMIN_BUCKET_NAME

Substitua ADMIN_BUCKET_NAME pelo nome do bucket criado na seção Criar buckets do Cloud Storage. Não inclua o prefixo gs://.

Aplique o manifesto:

kubectl apply -f saxml-admin-server.yaml

Verifique a implantação do servidor de administrador:

kubectl get deployment

A saída será semelhante a esta:

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

Implantar o servidor de modelos do Saxml

Siga estas instruções para implantar o servidor para os modelos Gemma 2B ou Gemma 7B.

Gemma 2B-it

Crie o seguinte manifesto saxml-model-server-1x1.yaml:

apiVersion: apps/v1
kind: Deployment
metadata:
  name: sax-model-server-v5e-1x1
spec:
  replicas: 1
  selector:
    matchLabels:
      app: gemma-server
  strategy:
    type: Recreate
  template:
    metadata:
      labels:
        app: gemma-server
        ai.gke.io/model: gemma-2b-it
        ai.gke.io/inference-server: saxml
        examples.ai.gke.io/source: user-guide
    spec:
      nodeSelector:
        cloud.google.com/gke-tpu-topology: 1x1
        cloud.google.com/gke-tpu-accelerator: tpu-v5-lite-podslice
      hostNetwork: false
      restartPolicy: Always
      containers:
      - name: inference-server
        image: us-docker.pkg.dev/cloud-tpu-images/inference/sax-model-server:v1.2.0
        args:
        - "--jax_platforms=tpu"
        - "--platform_chip=tpuv5e"
        - "--platform_topology=1x1"
        - "--port=10001"
        - "--sax_cell=/sax/test"
        ports:
        - containerPort: 10001
        securityContext:
          privileged: true
        env:
        - name: SAX_ROOT
          value: "gs://ADMIN_BUCKET_NAME/sax-root"
        resources:
          requests:
            google.com/tpu: 1
          limits:
            google.com/tpu: 1

Substitua ADMIN_BUCKET_NAME pelo nome do bucket criado na seção Criar buckets do Cloud Storage. Não inclua o prefixo gs://.

Aplique o manifesto:

kubectl apply -f saxml-model-server-1x1.yaml

Verifique o status da implantação do servidor de modelo:

kubectl get deployment

A saída será semelhante a esta:

NAME                                              READY   STATUS    RESTARTS   AGE
sax-admin-server                                  1/1     Running   0          ##m
sax-model-server-v5e-1x1                          1/1     Running   0          ##s

Gemma 7B-it

Crie o seguinte manifesto saxml-model-server-2x2.yaml:

apiVersion: apps/v1
kind: Deployment
metadata:
  name: sax-model-server-v5e-2x2
spec:
  replicas: 1
  selector:
    matchLabels:
      app: gemma-server
  strategy:
    type: Recreate
  template:
    metadata:
      labels:
        app: gemma-server
        ai.gke.io/model: gemma-7b-it
        ai.gke.io/inference-server: saxml
        examples.ai.gke.io/source: user-guide
    spec:
      nodeSelector:
        cloud.google.com/gke-tpu-topology: 2x2
        cloud.google.com/gke-tpu-accelerator: tpu-v5-lite-podslice
      hostNetwork: false
      restartPolicy: Always
      containers:
      - name: inference-server
        image: us-docker.pkg.dev/cloud-tpu-images/inference/sax-model-server:v1.2.0
        args:
        - "--jax_platforms=tpu"
        - "--platform_chip=tpuv5e"
        - "--platform_topology=2x2"
        - "--port=10001"
        - "--sax_cell=/sax/test"
        ports:
        - containerPort: 10001
        securityContext:
          privileged: true
        env:
        - name: SAX_ROOT
          value: "gs://ADMIN_BUCKET_NAME/sax-root"
        resources:
          requests:
            google.com/tpu: 4
          limits:
            google.com/tpu: 4

Substitua ADMIN_BUCKET_NAME pelo nome do bucket criado na seção Criar buckets do Cloud Storage. Não inclua o prefixo gs://.

Aplique o manifesto:

kubectl apply -f saxml-model-server-2x2.yaml

Verifique o status da implantação do servidor de modelo:

kubectl get deployment

A saída será semelhante a esta:

NAME                                              READY   STATUS    RESTARTS   AGE
sax-admin-server                                  1/1     Running   0          ##m
sax-model-server-v5e-2x2                          1/1     Running   0          ##s

Implantar o servidor HTTP Saxml

Nesta seção, você implanta o servidor HTTP Saxml e cria um serviço IP de cluster que você usa para acessar o servidor.

Crie o seguinte manifesto saxml-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
      containers:
      - name: sax-http
        image: us-docker.pkg.dev/cloud-tpu-images/inference/sax-http:v1.2.0
        imagePullPolicy: Always
        ports:
        - containerPort: 8888
        env:
        - name: SAX_ROOT
          value: "gs://ADMIN_BUCKET_NAME/sax-root"
---
apiVersion: v1
kind: Service
metadata:
  name: sax-http-svc
spec:
  selector:
    app: sax-http
  ports:
  - protocol: TCP
    port: 8888
    targetPort: 8888
  type: ClusterIP

Substitua ADMIN_BUCKET_NAME pelo nome do bucket do Cloud Storage que armazena o servidor Administrador do Saxml.

Aplique o manifesto:
```
kubectl apply -f saxml-http.yaml
```
Verifique o status da implantação do servidor HTTP Saxml:
```
kubectl get deployment
```

Gemma 2B-it

A saída será semelhante a esta:

NAME                                              READY   STATUS    RESTARTS   AGE
sax-admin-server                                  1/1     Running   0          ##m
sax-model-server-v5e-1x1                          1/1     Running   0          ##m
sax-http                                          1/1     Running   0          ##s

Gemma 7B-it

A saída será semelhante a esta:

NAME                                              READY   STATUS    RESTARTS   AGE
sax-admin-server                                  1/1     Running   0          ##m
sax-model-server-v5e-2x2                          1/1     Running   0          ##m
sax-http                                          1/1     Running   0          ##s

Fazer o download do checkpoint do modelo

Nesta seção, você executará um job do Kubernetes que busca, faz o download e armazena o checkpoint do modelo. Siga as etapas com base no modelo do Gemma que você quer usar:

Gemma 2B-it

Crie o seguinte manifesto job-2b.yaml:

apiVersion: v1
kind: ConfigMap
metadata:
  name: fetch-model-scripts
data:
  fetch_model.sh: |-
    #!/usr/bin/bash -x
    pip install kaggle --break-system-packages && \

    MODEL_NAME=$(echo ${MODEL_PATH} | awk -F'/' '{print $2}') && \
    VARIATION_NAME=$(echo ${MODEL_PATH} | awk -F'/' '{print $4}') && \

    mkdir -p /data/${MODEL_NAME}_${VARIATION_NAME} &&\
    kaggle models instances versions download ${MODEL_PATH} --untar -p /data/${MODEL_NAME}_${VARIATION_NAME} && \
    echo -e "\nCompleted extraction to /data/${MODEL_NAME}_${VARIATION_NAME}" && \

    gcloud storage rsync --recursive --no-clobber /data/${MODEL_NAME}_${VARIATION_NAME} gs://${BUCKET_NAME}/${MODEL_NAME}_${VARIATION_NAME} && \
    echo -e "\nCompleted copy of data to gs://${BUCKET_NAME}/${MODEL_NAME}_${VARIATION_NAME}"
---
apiVersion: batch/v1
kind: Job
metadata:
  name: data-loader-2b
  labels:
    app: data-loader-2b
spec:
  ttlSecondsAfterFinished: 120
  template:
    metadata:
      labels:
        app: data-loader-2b
    spec:
      restartPolicy: OnFailure
      containers:
      - name: gcloud
        image: gcr.io/google.com/cloudsdktool/google-cloud-cli:slim
        command:
        - /scripts/fetch_model.sh
        env:
        - name: BUCKET_NAME
          value: CHECKPOINTS_BUCKET_NAME
        - name: KAGGLE_CONFIG_DIR
          value: /kaggle
        - name: MODEL_PATH
          value: "google/gemma/pax/2b-it/2"
        volumeMounts:
        - mountPath: "/kaggle/"
          name: kaggle-credentials
          readOnly: true
        - mountPath: "/scripts/"
          name: scripts-volume
          readOnly: true
      volumes:
      - name: kaggle-credentials
        secret:
          defaultMode: 0400
          secretName: kaggle-secret
      - name: scripts-volume
        configMap:
          defaultMode: 0700
          name: fetch-model-scripts

Substitua CHECKPOINTS_BUCKET_NAME pelo nome do bucket criado na seção Criar buckets do Cloud Storage. Não inclua o prefixo gs://.

Aplique o manifesto:
```
kubectl apply -f job-2b.yaml
```

Aguarde a conclusão do job:

kubectl wait --for=condition=complete --timeout=180s job/data-loader-2b

O resultado será assim:

job.batch/data-loader-2b condition met

Verifique se o job foi concluído:

kubectl get job/data-loader-2b

O resultado será assim:

NAME             COMPLETIONS   DURATION   AGE
data-loader-2b   1/1           ##s        #m##s

Visualize os registros do job:

kubectl logs --follow job/data-loader-2b

O upload dos checkpoints é feito para gs://CHECKPOINTS_BUCKET_NAME/gemma_2b-it/checkpoint_00000000.

Gemma 7B-it

Crie o seguinte manifesto job-7b.yaml:

apiVersion: v1
kind: ConfigMap
metadata:
  name: fetch-model-scripts
data:
  fetch_model.sh: |-
    #!/usr/bin/bash -x
    pip install kaggle --break-system-packages && \

    MODEL_NAME=$(echo ${MODEL_PATH} | awk -F'/' '{print $2}') && \
    VARIATION_NAME=$(echo ${MODEL_PATH} | awk -F'/' '{print $4}') && \

    mkdir -p /data/${MODEL_NAME}_${VARIATION_NAME} &&\
    kaggle models instances versions download ${MODEL_PATH} --untar -p /data/${MODEL_NAME}_${VARIATION_NAME} && \
    echo -e "\nCompleted extraction to /data/${MODEL_NAME}_${VARIATION_NAME}" && \

    gcloud storage rsync --recursive --no-clobber /data/${MODEL_NAME}_${VARIATION_NAME} gs://${BUCKET_NAME}/${MODEL_NAME}_${VARIATION_NAME} && \
    echo -e "\nCompleted copy of data to gs://${BUCKET_NAME}/${MODEL_NAME}_${VARIATION_NAME}"
---
apiVersion: batch/v1
kind: Job
metadata:
  name: data-loader-7b
  labels:
    app: data-loader-7b
spec:
  ttlSecondsAfterFinished: 120
  template:
    metadata:
      labels:
        app: data-loader-7b
    spec:
      restartPolicy: OnFailure
      containers:
      - name: gcloud
        image: gcr.io/google.com/cloudsdktool/google-cloud-cli:slim
        command:
        - /scripts/fetch_model.sh
        env:
        - name: BUCKET_NAME
          value: CHECKPOINTS_BUCKET_NAME
        - name: KAGGLE_CONFIG_DIR
          value: /kaggle
        - name: MODEL_PATH
          value: "google/gemma/pax/7b-it/2"
        volumeMounts:
        - mountPath: "/kaggle/"
          name: kaggle-credentials
          readOnly: true
        - mountPath: "/scripts/"
          name: scripts-volume
          readOnly: true
      volumes:
      - name: kaggle-credentials
        secret:
          defaultMode: 0400
          secretName: kaggle-secret
      - name: scripts-volume
        configMap:
          defaultMode: 0700
          name: fetch-model-scripts

Substitua CHECKPOINTS_BUCKET_NAME pelo nome do bucket criado na seção Criar buckets do Cloud Storage. Inclua o prefixo gs://.

Aplique o manifesto:
```
kubectl apply -f job-7b.yaml
```

Aguarde a conclusão do job:

kubectl wait --for=condition=complete --timeout=360s job/data-loader-7b

O resultado será assim:

job.batch/data-loader-7b condition met

Verifique se o job foi concluído:

kubectl get job/data-loader-7b

O resultado será assim:

NAME             COMPLETIONS   DURATION   AGE
data-loader-7b   1/1           ##s        #m##s

Visualize os registros do job:

kubectl logs --follow job/data-loader-7b

O upload dos checkpoints é feito para gs://CHECKPOINTS_BUCKET_NAME/gemma_7b_it/checkpoint_00000000.

Expor o servidor HTTP Saxml

É possível acessar o servidor HTTP Saxml usando um Serviço ClusterIP que você criou na etapa anterior. Os serviços ClusterIP só podem ser acessados dentro do cluster. Portanto, para acessar o serviço de fora do cluster, conclua as seguintes etapas:

Estabeleça uma sessão de encaminhamento de portas:

kubectl port-forward service/sax-http-svc 8888:8888

Para verificar se é possível acessar o servidor HTTP Saxml, abra um novo terminal e execute o seguinte comando:
```
curl -s localhost:8888
```
O resultado será assim:
```
{
    "Message": "HTTP Server for SAX Client"
}
```

No servidor HTTP Saxml, a interface do cliente é encapsulada no sistema Saxml e ela é exposta por um conjunto de APIs REST. Você usará essas APIs para publicar, gerenciar e interagir com os modelos Gemma 2B e Gemma 7B.

Publicar o modelo Gemma

Depois de configurar os componentes Saxml, é possível publicar o modelo Gemma em um servidor de modelo que é executado em um pool de nós de fração de TPU. Use a API publish do servidor HTTP Saxml para publicar um modelo. Para saber mais sobre a API do servidor HTTP Saxml, consulte APIs HTTP do Saxml. Siga estas etapas para publicar o modelo de parâmetro Gemma 2B ou 7B.

Gemma 2B-it

Verifique se a sessão de encaminhamento de portas ainda está ativa:
```
curl -s localhost:8888
```

Publique o parâmetro Gemma 2B:

curl --request POST \
--header "Content-type: application/json" \
-s \
localhost:8888/publish \
--data \
'{
    "model": "/sax/test/gemma2bfp16",
    "model_path": "saxml.server.pax.lm.params.gemma.Gemma2BFP16",
    "checkpoint": "gs://CHECKPOINTS_BUCKET_NAME/gemma_2b-it/checkpoint_00000000",
    "replicas": "1"
}'

O resultado será assim:

{
    "model": "/sax/test/gemma2bfp16",
    "model_path": "saxml.server.pax.lm.params.gemma.Gemma2BFP16",
    "checkpoint": "gs://CHECKPOINTS_BUCKET_NAME/gemma_2b-it/checkpoint_00000000",
    "replicas": 1
}

A implantação do modelo pode levar alguns minutos.

Monitore o progresso observando os registros em um pod de servidor de modelo da implantação sax-model-server-v5e-1x1.

kubectl logs --follow deployment/sax-model-server-v5e-1x1

Essa implantação pode levar até cinco minutos para ser concluída. Aguarde até uma mensagem semelhante a esta aparecer:

I0125 15:34:31.685555 139063071708736 servable_model.py:699] loading completed.
I0125 15:34:31.686286 139063071708736 model_service_base.py:532] Successfully loaded model for key: /sax/test/gemma2bfp16

Verifique se você pode acessar o modelo exibindo as informações dele:

curl --request GET \
--header "Content-type: application/json" \
-s \
localhost:8888/listcell \
--data \
'{
    "model": "/sax/test/gemma2bfp16"
}'

O resultado será assim:

{
    "model": "/sax/test/gemma2bfp16",
    "model_path": "saxml.server.pax.lm.params.gemma.Gemma2BFP16",
    "checkpoint": "gs://CHECKPOINTS_BUCKET_NAME/gemma_2b-it/checkpoint_00000000",
    "max_replicas": 1,
    "active_replicas": 1
}

Gemma 7B-it

Verifique se a sessão de encaminhamento de portas ainda está ativa:
```
curl -s localhost:8888
```

Publique o parâmetro Gemma 7B:

curl --request POST \
--header "Content-type: application/json" \
-s \
localhost:8888/publish \
--data \
'{
    "model": "/sax/test/gemma7bfp16",
    "model_path": "saxml.server.pax.lm.params.gemma.Gemma7BFP16",
    "checkpoint": "gs://CHECKPOINTS_BUCKET_NAME/gemma_7b-it/checkpoint_00000000",
    "replicas": "1"
}'

O resultado será assim:

{
    "model": "/sax/test/gemma7bfp16",
    "model_path": "saxml.server.pax.lm.params.gemma.Gemma7BFP16",
    "checkpoint": "gs://CHECKPOINTS_BUCKET_NAME/gemma_7b-it/checkpoint_00000000",
    "replicas": 1
}

A implantação do modelo pode levar alguns minutos.

Monitore o progresso observando os registros em um pod de servidor de modelo da implantação sax-model-server-v5e-2x2.

kubectl logs --follow deployment/sax-model-server-v5e-2x2

Aguarde até uma mensagem semelhante a esta aparecer:

I0125 15:34:31.685555 139063071708736 servable_model.py:699] loading completed.
I0125 15:34:31.686286 139063071708736 model_service_base.py:532] Successfully loaded model for key: /sax/test/gemma7bfp16

Verifique se o modelo foi publicado exibindo as informações dele:

curl --request GET \
--header "Content-type: application/json" \
-s \
localhost:8888/listcell \
--data \
'{
    "model": "/sax/test/gemma7bfp16"
}'

O resultado será assim:

{
    "model": "/sax/test/gemma7bfp16",
    "model_path": "saxml.server.pax.lm.params.gemma.Gemma7BFP16",
    "checkpoint": "gs://CHECKPOINTS_BUCKET_NAME/gemma_7b-it/checkpoint_00000000",
    "max_replicas": 1,
    "active_replicas": 1
}

use o modelo

É possível interagir com os modelos Gemma 2B ou 7B. Use a API generate do servidor HTTP Saxml para enviar um prompt ao modelo.

Gemma 2B-it

Disponibilize uma solicitação de prompt usando o endpoint generate do servidor HTTP Saxml:

curl --request POST \
--header "Content-type: application/json" \
-s \
localhost:8888/generate \
--data \
'{
  "model": "/sax/test/gemma2bfp16",
  "query": "What are the top 5 most popular programming languages?"
}'

Confira a seguir um exemplo de resposta do modelo. A saída pode variar em cada comando exibido:

[
    [
        "\n\n1. **Python**\n2. **JavaScript**\n3. **Java**\n4. **C++**\n5. **Go**",
        -3.0704939365386963
    ]
]

É possível executar o comando com diferentes parâmetros query. Também é possível modificar parâmetros extras, como temperature, top_k e topc_p, usando a API generate. Para saber mais sobre a API do servidor HTTP Saxml, consulte APIs HTTP do Saxml.

Gemma 7B-it

Disponibilize uma solicitação de prompt usando o endpoint generate do servidor HTTP Saxml:

curl --request POST \
--header "Content-type: application/json" \
-s \
localhost:8888/generate \
--data \
'{
  "model": "/sax/test/gemma7bfp16",
  "query": "What are the top 5 most popular programming languages?"
}'

Confira a seguir um exemplo de resposta do modelo. A saída pode variar em cada comando exibido:

[
    [
        "\n\n**1. JavaScript**\n\n* Most widely used language on the web.\n* Used for front-end development, such as websites and mobile apps.\n* Extensive libraries and frameworks available.\n\n**2. Python**\n\n* Known for its simplicity and readability.\n* Versatile, used for various tasks, including data science, machine learning, and web development.\n* Large and active community.\n\n**3. Java**\n\n* Object-oriented language widely used in enterprise applications.\n* Used for web applications, mobile apps, and enterprise software.\n* Strong ecosystem and support.\n\n**4. Go**\n\n",
        -16.806324005126953
    ]
]

Cancelar a publicação do modelo

Siga estas etapas para cancelar a publicação do modelo:

Gemma 2B-it

Para cancelar a publicação do modelo Gemma 2B-it, execute o seguinte comando:

curl --request POST \
--header "Content-type: application/json" \
-s \
localhost:8888/unpublish \
--data \
'{
    "model": "/sax/test/gemma2bfp16"
}'

O resultado será mais ou menos assim:

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

É possível executar o comando com diferentes solicitações transmitidas no parâmetro query.

Gemma 7B-it

Para cancelar a publicação do modelo Gemma 7B-it, execute o seguinte comando:

curl --request POST \
--header "Content-type: application/json" \
-s \
localhost:8888/unpublish \
--data \
'{
    "model": "/sax/test/gemma7bfp16"
}'

O resultado será mais ou menos assim:

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

É possível executar o comando com diferentes solicitações transmitidas no parâmetro query.

Resolver problemas

Se você receber a mensagem Empty reply from server, é possível que o contêiner não tenha concluído o download dos dados do modelo. Verifique os registros do pod novamente para ver a mensagem Connected, que indica que o modelo está pronto para ser disponibilizado.
Se aparecer Connection refused, verifique se o encaminhamento de portas está ativo.

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

Para evitar cobranças na sua conta do Google Cloud pelos recursos criados neste guia, execute o seguinte comando:

gcloud container clusters delete ${CLUSTER_NAME} --location=${LOCATION}
gcloud iam service-accounts delete --quiet wi-sax@${PROJECT_ID}.iam.gserviceaccount.com
gcloud storage rm --recursive gs://ADMIN_BUCKET_NAME
gcloud storage rm --recursive gs://CHECKPOINTS_BUCKET_NAME

Substitua:

ADMIN_BUCKET_NAME: o nome do bucket do Cloud Storage que armazena o servidor Administrador do Saxml.
CHECKPOINTS_BUCKET_NAME: o nome do bucket do Cloud Storage que armazena os checkpoints do modelo.

A seguir

Saiba mais sobre TPUs no GKE.
Conheça o repositório do GitHub do Saxml.
Conheça o Model Garden da Vertex AI.
Descubra como executar cargas de trabalho de IA/ML otimizadas com os recursos de orquestração da plataforma GKE.
Confira arquiteturas de referência, diagramas, tutoriais e práticas recomendadas do Google Cloud. Confira o Centro de arquitetura do Cloud.

Disponibilizar modelos abertos do Gemma usando TPUs no GKE com Saxml

Experiência

Gemma

TPUs

Saxml

Objetivos

Arquitetura

Antes de começar

Verificar os papéis

Conceder os papéis

Preparar o ambiente para o Gemma

Iniciar o Cloud Shell

Criar um cluster do GKE Standard

Gemma 2B-it

Gemma 7B-it

Criar os buckets do Cloud Storage

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

Receber acesso ao modelo

Assinar o contrato de consentimento de licença

Gerar um token de acesso

Fazer upload do token de acesso no Cloud Shell

Criar um secret do Kubernetes para credenciais do Kaggle

Implantar o Saxml

Implantar o servidor de administrador Saxml

Implantar o servidor de modelos do Saxml

Gemma 2B-it

Gemma 7B-it

Implantar o servidor HTTP Saxml

Gemma 2B-it

Gemma 7B-it

Fazer o download do checkpoint do modelo

Gemma 2B-it

Gemma 7B-it

Expor o servidor HTTP Saxml

Publicar o modelo Gemma

Gemma 2B-it

Gemma 7B-it

use o modelo

Gemma 2B-it

Gemma 7B-it

Cancelar a publicação do modelo

Gemma 2B-it

Gemma 7B-it

Resolver problemas

Limpar

Excluir os recursos implantados

A seguir