Publique modelos abertos Gemma com TPUs na Vertex AI com o Saxml

Este guia mostra como apresentar um modelo de linguagem (conteúdo extenso) (MDI/CE) de modelos abertos Gemma através de unidades de processamento tensor (TPUs) na Vertex AI com o Saxml. Neste guia, transfere os modelos Gemma com ajuste fino de instruções de 2B e 7B parâmetros para o Cloud Storage e implementa-os no Vertex AI que executa o Saxml em TPUs.

Contexto

Publicando o Gemma através de TPUs no Vertex AI com o Saxml. Pode tirar partido de uma solução de IA gerida que cuida da infraestrutura de baixo nível e oferece uma forma económica de publicar MDIs. Esta secção descreve as principais tecnologias usadas neste tutorial.

Gemma

O Gemma é um conjunto de modelos de inteligência artificial (IA) generativa, leve e disponível abertamente, lançado ao abrigo de uma licença aberta. Estes modelos de IA estão disponíveis para execução nas suas aplicações, hardware, dispositivos móveis ou serviços alojados. Pode usar os modelos Gemma para a geração de texto. No entanto, também pode otimizar estes modelos para tarefas especializadas.

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

Saxml

O Saxml é um sistema experimental que publica modelos Paxml, JAX e PyTorch para inferência. Para este tutorial, vamos abordar como publicar o Gemma em TPUs mais rentáveis para o Saxml. A configuração para GPUs é semelhante. O Saxml oferece scripts para criar contentores para o Vertex AI que vamos usar neste tutorial.

TPUs

As TPUs são circuitos integrados específicos da aplicação (ASICs) desenvolvidos pela Google e usados para acelerar frameworks de processamento de dados, como o TensorFlow, o PyTorch e o JAX.

Este tutorial disponibiliza os modelos Gemma 2B e Gemma 7B. O Vertex AI aloja estes modelos nos seguintes conjuntos de nós de TPU v5e de anfitrião único:

  • Gemma 2B: alojado num conjunto de nós da TPU v5e com uma 1x1topologia que representa um chip da TPU. O tipo de máquina dos nós é ct5lp-hightpu-1t.
  • Gemma 7B: alojado num conjunto de nós da TPU v5e com uma 2x2topologia que representa quatro chips da TPU. O tipo de máquina dos nós é ct5lp-hightpu-4t.

Antes de começar

  1. 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.
  2. 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 (roles/resourcemanager.projectCreator), which contains the resourcemanager.projects.create permission. Learn how to grant roles.

    Go to project selector

  3. Verify that billing is enabled for your Google Cloud project.

  4. Enable the Vertex AI 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

  5. 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 (roles/resourcemanager.projectCreator), which contains the resourcemanager.projects.create permission. Learn how to grant roles.

    Go to project selector

  6. Verify that billing is enabled for your Google Cloud project.

  7. Enable the Vertex AI 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

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

  9. Este tutorial pressupõe que está a usar o Cloud Shell para interagir com o Google Cloud. Se quiser usar uma shell diferente em vez do Cloud Shell, faça a seguinte configuração adicional:

    1. Install the Google Cloud CLI.

    2. If you're using an external identity provider (IdP), you must first sign in to the gcloud CLI with your federated identity.

    3. To initialize the gcloud CLI, run the following command:

      gcloud init
    4. Certifique-se de que tem quota suficiente para chips de TPU v5e para o Vertex AI. Por predefinição, esta quota é 0. Para uma topologia 1x1, tem de ser 1. Para 2x2, tem de ser 4. Para executar ambas as topologias, tem de ser 5.
    5. Crie uma conta do Kaggle, se ainda não tiver uma.

    Aceda ao modelo

    Tenha em atenção que o Cloud Shell pode não ter recursos suficientes para transferir os pesos do modelo. Se for o caso, pode criar uma instância do Vertex AI Workbench para realizar essa tarefa.

    Para aceder aos modelos Gemma para implementação no Vertex AI, tem de iniciar sessão na plataforma Kaggle, assinar o contrato de consentimento da licença e obter um token da API Kaggle. Neste tutorial, vai usar um segredo do Kubernetes para as credenciais do Kaggle.

    Tem de assinar o contrato de consentimento para usar o Gemma. Siga estas instruções:

    1. Aceda à página de consentimento do modelo em Kaggle.com.
    2. Inicie sessão no Kaggle, caso ainda não o tenha feito.
    3. Clique em Pedir acesso.
    4. Na secção Escolher conta para consentimento, selecione Validar através da conta do Kaggle para usar a sua conta do Kaggle para consentimento.
    5. Aceite os Termos de Utilização do modelo.

    Gere um token de acesso

    Para aceder ao modelo através do Kaggle, precisa de um token da API Kaggle.

    Siga estes passos para gerar um novo token se ainda não tiver um:

    1. No navegador, aceda às definições do Kaggle.
    2. Na secção API, clique em Criar novo token.

      É transferido um ficheiro com o nome kaggle.json.

    Carregue o token de acesso para o Cloud Shell

    No Cloud Shell, pode carregar o token da API Kaggle para o seu Google Cloud projeto:

    1. No Cloud Shell, clique em Mais > Carregar.
    2. Selecione Ficheiro e clique em Escolher ficheiros.
    3. Abra o ficheiro kaggle.json.
    4. Clique em Carregar.

    Crie o contentor do Cloud Storage

    Crie um contentor do Cloud Storage para armazenar os pontos de verificação do modelo.

    No Cloud Shell, execute o seguinte:

    gcloud storage buckets create gs://CHECKPOINTS_BUCKET_NAME
    

    Substitua CHECKPOINTS_BUCKET_NAME pelo nome do contentor do Cloud Storage que armazena os pontos de verificação do modelo.

    Copie o modelo para o contentor 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
    

    Implementar o modelo

    Carregue um modelo

    Para carregar um recurso Model que use o seu contentor Saxml, execute o seguinte gcloud ai models upload comando:

    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 o seguinte:

    • PROJECT_ID: o ID do seu Google Cloud projeto
    • LOCATION_ID: a região onde está a usar o Vertex AI. Tenha em atenção que as TPUs só estão disponíveis em us-west1.
    • DEPLOYED_MODEL_NAME: um nome para o DeployedModel. Também pode usar o nome a apresentar do Model para o DeployedModel.

    Crie um ponto final

    Tem de implementar o modelo num ponto final antes de o poder usar para publicar inferências online. Se estiver a implementar um modelo num ponto final existente, pode ignorar este passo. O exemplo seguinte usa o comando gcloud ai endpoints create:

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

    Substitua o seguinte:

    • LOCATION_ID: a região onde está a usar o Vertex AI.
    • ENDPOINT_NAME: o nome a apresentar do ponto final.

    A ferramenta Google Cloud CLI pode demorar alguns segundos a criar o ponto final.

    Implemente o modelo no ponto final

    Quando o ponto final estiver pronto, implemente o modelo no ponto final.

    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 o seguinte:

    • LOCATION_ID: a região onde está a usar o Vertex AI.
    • ENDPOINT_NAME: o nome a apresentar do ponto final.
    • DEPLOYED_MODEL_NAME: um nome para o DeployedModel. Também pode usar o nome a apresentar do Model para o DeployedModel.

    Gemma 2B pode ser implementado numa máquina ct5lp-hightpu-1t mais pequena. Nesse caso, deve especificar --platform_topology=1x1 ao carregar o modelo.

    A ferramenta Google Cloud CLI pode demorar alguns minutos a implementar o modelo no ponto final. Quando o modelo é implementado com êxito, este comando imprime o seguinte resultado:

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

    Obter inferências online do modelo implementado

    Para invocar o modelo através do ponto final da Vertex AI, formate o pedido de inferência com um objeto JSON de pedido de inferência padrão .

    O exemplo seguinte 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 o seguinte:

    • LOCATION_ID: a região onde está a usar o Vertex AI.
    • ENDPOINT_NAME: o nome a apresentar do ponto final.
    • instances.json tem o seguinte formato: {"instances": [{"text_batch": "<your prompt>"},{...}]}

    Limpar

    Para evitar incorrer em mais cobranças da Vertex AI e cobranças do Artifact Registry, elimine os Google Cloud recursos que criou durante este tutorial:

    1. Para anular a implementação do modelo a partir do ponto final e eliminar o ponto final, execute o seguinte comando na 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 onde criou o modelo numa secção anterior.

    2. Para eliminar o seu modelo, execute o seguinte comando na 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 onde criou o modelo numa secção anterior.

    Limitações

    • No Vertex AI, os TPUs do Google Cloud só são suportados em us-west1. Para mais informações, consulte as localizações.

    O que se segue?

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