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.

Experiência

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

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 as APIs Vertex AI API and Artifact Registry API.

Ative as APIs

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 as APIs Vertex AI API and Artifact Registry API.

Ative as APIs

No Console do Google Cloud, ative o Cloud Shell.

Ativar o Cloud Shell

Na parte inferior do Console do Google Cloud, 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 o Google Cloud. Se você quiser usar um shell diferente em vez do Cloud Shell, execute a seguinte configuração extra:

Instale a CLI do Google Cloud.
Para inicializar a CLI gcloud, execute o seguinte comando:
```
gcloud init
```

Siga a documentação do Artifact Registry para instalar o Docker.
Verifique se você tem cotas suficientes para cinco chips TPU v5e para a Previsão da Vertex AI.
Crie uma conta da Kaggle, caso ainda não tenha uma.

Receber acesso ao modelo

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/

Crie um repositório do Artifact Registry

Crie um repositório do Artifact Registry para armazenar a imagem do contêiner que você criará na próxima seção.

Ative o serviço da API Artifact Registry para o projeto.

gcloud services enable artifactregistry.googleapis.com

Execute o seguinte comando no shell para criar o repositório do Artifact Registry:

gcloud artifacts repositories create saxml \
 --repository-format=docker \
 --location=LOCATION \
 --description="Saxml Docker repository"

Substitua LOCATION pela região em que o Artifact Registry armazena a imagem do contêiner. Mais tarde, será preciso criar um recurso de modelo da Vertex AI em um endpoint regional que corresponda a essa região. Portanto, escolha uma região em que a Vertex AI tenha um endpoint regional: como us-west1 para TPUs.

Envie a imagem do contêiner para o Artifact Registry

O contêiner Saxml pré-criado está disponível em us-docker.pkg.dev/vertex-ai/prediction/sax-tpu:latest. Copie o arquivo para o Artifact Registry. Configure o Docker para acessar o Artifact Registry. Em seguida, envie a imagem do contêiner para o repositório do Artifact Registry.

Para conceder à instalação local do Docker permissão para transmitir ao Artifact Registry na região escolhida, execute o seguinte comando no shell:
```
gcloud auth configure-docker LOCATION-docker.pkg.dev
```
- Substitua LOCATION pela região em que você criou o repositório.

Para copiar a imagem de contêiner recém-criada ao Artifact Registry, execute o seguinte comando no shell:

docker tag us-docker.pkg.dev/vertex-ai/prediction/sax-tpu:latest LOCATION-docker.pkg.dev/PROJECT_ID/saxml/saxml-tpu:latest

Para enviar a imagem de contêiner recém-criada ao Artifact Registry, execute o seguinte comando no shell:
```
docker push LOCATION-docker.pkg.dev/PROJECT_ID/saxml/saxml-tpu:latest
```
Faça o seguinte, como na seção anterior:
- LOCATION: a região do repositório do Artifact Registry
- PROJECT_ID: o ID do seu projeto do Google Cloud.

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=LOCATION-docker.pkg.dev/PROJECT_ID/saxml/saxml-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=LOCATION-docker.pkg.dev/PROJECT_ID/saxml/saxml-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>"},{...}]}

Limpar

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.

Para excluir o repositório do Artifact Registry e a imagem do contêiner nele, execute o seguinte comando no shell:
```
gcloud artifacts repositories delete saxml \
  --location=LOCATION \
  --quiet
```
Substitua LOCATION pela região em que você criou o repositório do Artifact Registry 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

Experiência

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

Crie um repositório do Artifact Registry

Envie a imagem do contêiner para o Artifact Registry

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

Limpar

Limitações

A seguir