Disponibilizar modelos abertos do Gemma usando TPUs na Gemini Enterprise Agent Platform com Saxml

Neste guia, mostramos como disponibilizar um modelo de linguagem grande (LLM) de modelos abertos do Gemma usando unidades de processamento de tensor (TPUs) na Gemini Enterprise Agent Platform com Saxml. Neste guia, você vai fazer o download dos modelos Gemma ajustados por instruções de parâmetros 2B e 7B para o Cloud Storage e implantá-los na Gemini Enterprise Agent Platform que executa Saxml em TPUs.

Contexto

Disponibilizando o Gemma usando TPUs na Gemini Enterprise Agent Platform com Saxml. É possível usar uma solução de IA gerenciada que cuida da infraestrutura de baixo nível e oferece uma maneira econômica de disponibilizar LLMs. Esta seção descreve as principais tecnologias usadas neste tutorial.

Gemma

O Gemma é um conjunto de modelos de inteligência artificial (IA) generativa, leve e abertamente lançados sob 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.

Saxml

O Saxml é um sistema experimental que exibe modelos Paxml, JAX e PyTorch para inferência. Para este tutorial, abordaremos como disponibilizar o Gemma em TPUs que são mais econômicas para o Saxml. A configuração das GPUs é semelhante. O Saxml oferece scripts para criar contêineres para a Gemini Enterprise Agent Platform que vamos usar neste tutorial.

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.

Este tutorial serve os modelos Gemma 2B e Gemma 7B. A Gemini Enterprise Agent Platform hospeda esses modelos nos seguintes pools de nós da TPU v5e de host único:

  • Gemma 2B: hospedado em um pool de nós da TPU v5e com a topologia 1x1, que representa um chip de TPU. O tipo de máquina dos nós é ct5lp-hightpu-1t.
  • Gemma 7B: hospedado em um pool de nós da TPU v5e com a topologia 2x2, que representa quatro chips de TPU. O tipo de máquina dos nós é ct5lp-hightpu-4t.

Antes de começar

  1. In the Cloud de Confiance console, on the project selector page, select or create a Cloud de Confiance 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

  2. Verify that billing is enabled for your Cloud de Confiance project.

  3. Enable the Agent Platform API, if it is not already enabled.

    Roles required to enable APIs

    To enable APIs, you need the serviceusage.services.enable permission. If you created the project, then you likely already have this permission through the Owner role (roles/owner). Otherwise, you can get this permission through the Service Usage Admin role (roles/serviceusage.serviceUsageAdmin). Learn how to grant roles.

    Enable the API

  4. No console do Cloud de Confiance , ative o Cloud Shell.

    Ativar o Cloud Shell

    Na parte de baixo do console Cloud de Confiance , uma sessão do Cloud Shell é iniciada e exibe um prompt de linha de comando. O Cloud Shell é um ambiente shell com a CLI do Google Cloud já instalada e com valores já definidos para o projeto atual. A inicialização da sessão pode levar alguns segundos.

Neste tutorial, presumimos que você esteja usando o Cloud Shell para interagir com Cloud de Confiance by S3NS. Se você quiser usar um shell diferente em vez do Cloud Shell, execute a seguinte configuração extra:

  1. Instale a CLI do Google Cloud.

  2. Configure a CLI gcloud para usar sua identidade federada.

    Para mais informações, consulte Fazer login na CLI gcloud com sua identidade federada.

  3. Para inicializar a CLI gcloud, execute o seguinte comando:

    gcloud init
  4. Verifique se você tem cota suficiente para chips TPU v5e para a plataforma de agentes do Gemini Enterprise. Por padrão, essa cota é 0. Para uma topologia 1x1, ele precisa ser 1. Para o 2x2, deve ser 4. Para executar ambas as topologias, ele precisa ser 5.
  5. Crie uma conta da Kaggle, caso ainda não tenha uma.

Receber acesso ao modelo

Observe que o Cloud Shell pode não ter recursos suficientes para fazer download dos pesos do modelo. Nesse caso, você pode criar uma instância do Gemini Enterprise Agent Platform Workbench para executar essa tarefa.

Para ter acesso aos modelos Gemma para implantação na Gemini Enterprise Agent Platform, 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:

  1. Acesse a página de consentimento do modelo em Kaggle.com.
  2. Faça login no Kaggle, caso ainda não o tenha feito.
  3. Clique em Solicitar acesso.
  4. Na seção Escolher conta para consentimento, selecione Verificar via conta do Kaggle para usar sua conta do Kaggle para consentimento.
  5. 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:

  1. No navegador, acesse as configurações da Kaggle.
  2. 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 Cloud de Confiance by S3NS:

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

Criar o bucket do Cloud Storage

Crie um bucket do Cloud Storage para armazenar os checkpoints do modelo.

No Cloud Shell, execute o seguinte:

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.

Copiar o modelo para o bucket do Cloud Storage

No Cloud Shell, execute o seguinte:

pip install kaggle --break-system-packages

# For Gemma 2B
mkdir -p /data/gemma_2b-it
kaggle models instances versions download google/gemma/pax/2b-it/1 --untar -p /data/gemma_2b-it
gcloud storage cp /data/gemma_2b-it/* gs://CHECKPOINTS_BUCKET_NAME/gemma_2b-it/ --recursive

# For Gemma 7B
mkdir -p /data/gemma_7b-it
kaggle models instances versions download google/gemma/pax/7b-it/1 --untar -p /data/gemma_7b-it
gcloud storage cp /data/gemma_7b-it/* gs://CHECKPOINTS_BUCKET_NAME/gemma_7b-it/ --recursive

Implantar o modelo.

Fazer upload de um modelo

Para fazer upload de um recurso Model que usa o contêiner Saxml, execute o seguinte comando gcloud ai models upload:

Gemma 2B-it

gcloud ai models upload \
  --region=LOCATION \
  --display-name=DEPLOYED_MODEL_NAME \
  --container-image-uri=us-docker.pkg.dev/vertex-ai/prediction/sax-tpu:latest \
  --artifact-uri='gs://CHECKPOINTS_BUCKET_NAME/gemma_2b-it/' \
  --container-args='--model_path=saxml.server.pax.lm.params.gemma.Gemma2BFP16' \
  --container-args='--platform_chip=tpuv5e' \
  --container-args='--platform_topology=2x2' \
  --container-args='--ckpt_path_suffix=checkpoint_00000000' \
  --container-ports=8502

Gemma 7B-it

gcloud ai models upload \
  --region=LOCATION \
  --display-name=DEPLOYED_MODEL_NAME \
  --container-image-uri=us-docker.pkg.dev/vertex-ai/prediction/sax-tpu:latest \
  --artifact-uri='gs://CHECKPOINTS_BUCKET_NAME/gemma_7b-it/' \
  --container-args='--model_path=saxml.server.pax.lm.params.gemma.Gemma7BFP16' \
  --container-args='--platform_chip=tpuv5e' \
  --container-args='--platform_topology=2x2' \
  --container-args='--ckpt_path_suffix=checkpoint_00000000' \
  --container-ports=8502

Substitua:

  • PROJECT_ID: o ID do seu projetoCloud de Confiance
  • LOCATION_ID: A região em que você está usando o Agent Platform. As TPUs só estão disponíveis na região us-west1.
  • DEPLOYED_MODEL_NAME: um nome para DeployedModel. Também é possível usar o nome de exibição do Model para o DeployedModel.

Crie um endpoint

Implante o modelo em um endpoint para que ele possa ser usado para disponibilizar inferências on-line. Pule a etapa abaixo se você estiver implantando um modelo em um endpoint existente. O exemplo a seguir usa o comando gcloud ai endpoints create:

gcloud ai endpoints create \
  --region=LOCATION \
  --display-name=ENDPOINT_NAME

Substitua:

  • LOCATION_ID: A região em que você está usando o Agent Platform.
  • ENDPOINT_NAME: o nome de exibição do endpoint.

A ferramenta CLI do Google Cloud pode levar alguns segundos para criar o endpoint.

Implantar o modelo no endpoint

Quando o endpoint estiver pronto, implante o modelo nele.

ENDPOINT_ID=$(gcloud ai endpoints list \
   --region=LOCATION \
   --filter=display_name=ENDPOINT_NAME \
   --format="value(name)")

MODEL_ID=$(gcloud ai models list \
   --region=LOCATION \
   --filter=display_name=DEPLOYED_MODEL_NAME \
   --format="value(name)")

gcloud ai endpoints deploy-model $ENDPOINT_ID \
  --region=LOCATION \
  --model=$MODEL_ID \
  --display-name=DEPLOYED_MODEL_NAME \
  --machine-type=ct5lp-hightpu-4t \
  --traffic-split=0=100

Substitua:

  • LOCATION_ID: A região em que você está usando o Agent Platform.
  • ENDPOINT_NAME: o nome de exibição do endpoint.
  • DEPLOYED_MODEL_NAME: um nome para DeployedModel. Também é possível usar o nome de exibição do Model para o DeployedModel.

Gemma 2B pode ser implantado em uma máquina ct5lp-hightpu-1t menor. Nesse caso, especifique --platform_topology=1x1 ao fazer upload do modelo.

A ferramenta da CLI do Google Cloud pode levar alguns segundos para implantar o modelo no endpoint. Quando o modelo é implantado com sucesso, esse comando imprime a seguinte saída:

  Deployed a model to the endpoint xxxxx. Id of the deployed model: xxxxx.

Como receber inferências on-line com base no modelo implantado

Para invocar o modelo pelo endpoint da Gemini Enterprise Agent Platform, formate a solicitação de inferência usando um objeto JSON de solicitação de inferência padrão .

O exemplo a seguir usa o comando gcloud ai endpoints predict:

ENDPOINT_ID=$(gcloud ai endpoints list \
   --region=LOCATION \
   --filter=display_name=ENDPOINT_NAME \
   --format="value(name)")

gcloud ai endpoints predict $ENDPOINT_ID \
  --region=LOCATION \
  --http-headers=Content-Type=application/json \
  --json-request instances.json

Substitua:

  • LOCATION_ID: A região em que você está usando o Agent Platform.
  • ENDPOINT_NAME: o nome de exibição do endpoint.
  • instances.json tem o seguinte formato: {"instances": [{"text_batch": "<your prompt>"},{...}]}

Limpar

Para evitar outras cobranças da Gemini Enterprise Agent Platform e do Artifact Registry, exclua os recursos do Cloud de Confiance que você criou durante este tutorial:

  1. Para cancelar a implantação do modelo do endpoint e excluir o endpoint, execute o seguinte comando no shell:

    ENDPOINT_ID=$(gcloud ai endpoints list \
       --region=LOCATION \
       --filter=display_name=ENDPOINT_NAME \
       --format="value(name)")
    
    DEPLOYED_MODEL_ID=$(gcloud ai endpoints describe $ENDPOINT_ID \
       --region=LOCATION \
       --format="value(deployedModels.id)")
    
    gcloud ai endpoints undeploy-model $ENDPOINT_ID \
      --region=LOCATION \
      --deployed-model-id=$DEPLOYED_MODEL_ID
    
    gcloud ai endpoints delete $ENDPOINT_ID \
       --region=LOCATION \
       --quiet
    

    Substitua LOCATION pela região em que você criou o modelo em uma seção anterior.

  2. Para excluir o modelo, execute o seguinte comando no shell:

    MODEL_ID=$(gcloud ai models list \
       --region=LOCATION \
       --filter=display_name=DEPLOYED_MODEL_NAME \
       --format="value(name)")
    
    gcloud ai models delete $MODEL_ID \
       --region=LOCATION \
       --quiet
    

    Substitua LOCATION pela região em que você criou o modelo em uma seção anterior.

Limitações

  • Na Gemini Enterprise Agent Platform, as Cloud TPUs são compatíveis apenas com us-west1. Saiba mais em Locais.

A seguir

  • Saiba como implantar outros modelos Saxml, como Llama2 e GPT-J.