Disponibilizar modelos abertos do Gemma usando TPUs na Vertex AI Prediction com Saxml

Neste guia, mostramos como exibir um modelo de linguagem grande (LLM) de modelos abertos do Gemma usando unidades de processamento de tensor (TPUs) na Vertex AI Prediction 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 previsão da Vertex AI que executa Saxml em TPUs.

Contexto

Disponibilizando o Gemma usando TPUs na previsão da Vertex AI 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 Previsão da Vertex AI 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 Previsão da Vertex AI 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

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.

Go to project selector

Make sure that billing is enabled for your Google Cloud project.

Enable the Vertex AI API.

Enable the API

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

Go to project selector

Make sure that billing is enabled for your Google Cloud project.

Enable the Vertex AI API.

Enable the API

In the Google Cloud console, activate Cloud Shell.

Activate Cloud Shell

At the bottom of the Google Cloud console, a Cloud Shell session starts and displays a command-line prompt. Cloud Shell is a shell environment with the Google Cloud CLI already installed and with values already set for your current project. It can take a few seconds for the session to initialize.

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

Install the Google Cloud CLI.
To initialize the gcloud CLI, run the following command:
```
gcloud init
```
Verifique se você tem cota suficiente para chips TPU v5e para Previsão da Vertex AI. 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.
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 Vertex AI Workbench para executar essa tarefa.

Para ter acesso aos modelos Gemma para implantação no Vertex AI Prediction, 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, caso ainda não o tenha feito.
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 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
gsutil -m cp -R /data/gemma_2b-it/* gs://CHECKPOINTS_BUCKET_NAME/gemma_2b-it/

# 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
gsutil -m cp -R /data/gemma_7b-it/* gs://CHECKPOINTS_BUCKET_NAME/gemma_7b-it/

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 projeto do Google Cloud.
LOCATION_ID: a região em que você está usando a Vertex AI. 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.

Criar um endpoint

Implante o modelo em um endpoint para que ele possa ser usado para exibir previsões 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 a Vertex AI.
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 a Vertex AI.
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 previsões on-line com base no modelo implantado

Para invocar o modelo por meio do endpoint de previsão da Vertex AI, formate a solicitação de previsão 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 a Vertex AI.
ENDPOINT_NAME: o nome de exibição do endpoint.
instances.json tem o seguinte formato: {"instances": [{"text_batch": "<your prompt>"},{...}]}

Limpeza

Para evitar outras cobranças da Vertex AI e cobranças do Artifact Registry, exclua os recursos do Google Cloud criados durante este tutorial:

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.

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 Previsão da Vertex AI, 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.

Disponibilizar modelos abertos do Gemma usando TPUs na Vertex AI Prediction com Saxml

Contexto

Gemma

Saxml

TPUs

Antes de começar

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 o bucket do Cloud Storage

Copiar o modelo para o bucket do Cloud Storage

Implantar o modelo.

Fazer upload de um modelo

Gemma 2B-it

Gemma 7B-it

Criar um endpoint

Implantar o modelo no endpoint

Como receber previsões on-line com base no modelo implantado

Limpeza

Limitações

A seguir