Publique modelos abertos Llama 3 com TPUs na nuvem multi-anfitrião no Vertex AI com o Saxml

Pré-visualização

Esta funcionalidade está sujeita aos "Termos das Ofertas de Pré-DG" na secção Termos Gerais de Serviço dos Termos Específicos do Serviço. As funcionalidades de pré-DG estão disponíveis "tal como estão" e podem ter apoio técnico limitado. Para mais informações, consulte as descrições das fases de lançamento.

O Llama 3 é um grande modelo de linguagem (GML) de código aberto da Meta. Este guia mostra como publicar um LLM Llama 3 usando unidades de processamento de tensores (TPUs) de vários anfitriões no Vertex AI com Saxml.

Neste guia, transfere os pesos e o tokenizador do modelo Llama 3 70B e implementa-os no Vertex AI que executa o Saxml em TPUs.

Nota: está disponível uma versão do Llama 3 apenas com GPU no Model Garden. Para mais informações sobre o Model Garden, consulte o artigo Explore modelos de IA no Model Garden.

Antes de começar

Recomendamos que use uma VM otimizada para memória M2 para transferir o modelo e convertê-lo em Saxml. Isto deve-se ao facto de o processo de conversão do modelo exigir uma quantidade significativa de memória e poder falhar se escolher um tipo de máquina que não tenha memória suficiente.

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.

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.

Note: If you don't plan to keep the resources that you create in this procedure, create a project instead of selecting an existing project. After you finish these steps, you can delete the project, removing all resources associated with the project.

Go to project selector

Verify that billing is enabled for your Google Cloud project.

Enable the Vertex AI and Artifact Registry APIs.

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 APIs

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.

Note: If you don't plan to keep the resources that you create in this procedure, create a project instead of selecting an existing project. After you finish these steps, you can delete the project, removing all resources associated with the project.

Go to project selector

Verify that billing is enabled for your Google Cloud project.

Enable the Vertex AI and Artifact Registry APIs.

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 APIs

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

2. Siga a documentação do Artifact Registry para instalar o Docker.
3. Certifique-se de que tem quotas suficientes para 16 chips de TPU v5e para o Vertex AI.

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

Se estiver a usar uma shell diferente da Cloud Shell para a implementação do modelo, certifique-se de que a versão da CLI Google Cloud é posterior à 475.0.0. Pode atualizar a CLI gcloud executando o comando gcloud components update.

Se estiver a implementar o seu modelo através do SDK do Vertex AI, certifique-se de que tem a versão 1.50.0 ou posterior.

Aceda ao modelo e transfira as ponderações do modelo

Os passos seguintes destinam-se a uma instância do Vertex AI Workbench que tenha uma VM otimizada para memória M2. Para obter informações sobre como alterar o tipo de máquina de uma instância do Vertex AI Workbench, consulte o artigo Alterar o tipo de máquina de uma instância do Vertex AI Workbench.

1. Aceda à página de consentimento do modelo Llama.

2. Selecione o Llama 3, preencha o formulário de consentimento e aceite os termos e condições.

3. Verifique se recebeu um email com um URL assinado na sua caixa de entrada.

4. Transfira o script download.sh do GitHub executando o seguinte comando:

   wget https://raw.githubusercontent.com/meta-llama/llama3/main/download.sh
   chmod +x download.sh
   

5. Para transferir os pesos do modelo, execute o script download.sh que transferiu do GitHub.

6. Quando lhe for pedido, introduza o URL assinado do email que recebeu na secção anterior.

7. Quando lhe for pedido para transferir os modelos, introduza 70B.

Converta as ponderações do modelo para o formato Saxml

1. Execute o seguinte comando para transferir o Saxml:

   git clone https://github.com/google/saxml.git
   

2. Execute os seguintes comandos para configurar um ambiente virtual do Python:

   python -m venv .
   source bin/activate
   

3. Execute os seguintes comandos para instalar as dependências:

   pip install --upgrade pip
   
   pip install paxml
   
   pip install praxis
   
   pip install torch
   

4. Para converter os pesos do modelo para o formato Saxml, execute o seguinte comando:

   python3 saxml/saxml/tools/convert_llama_ckpt.py \
       --base PATH_TO_META_LLAMA3 \
       --pax PATH_TO_PAX_LLAMA3 \
       --model-size llama3_70b
   

   Substitua o seguinte:

   • PATH_TO_META_LLAMA3: o caminho para o diretório que contém as ponderações do modelo transferidas
   • PATH_TO_PAX_LLAMA3: o caminho para o diretório no qual armazenar os pesos do modelo convertido
   Nota: pode usar este comando para qualquer modelo Llama 2 ou Llama 3.

   Os modelos convertidos são colocados na pasta $PATH_TO_PAX_LLAMA3/checkpoint_00000000.

5. Copie o ficheiro do tokenizador do diretório original para uma subpasta com o nome vocabs da seguinte forma:

   cp $PATH_TO_META_LLAMA3/tokenizer.model $PATH_TO_PAX_LLAMA3/vocabs/tokenizer.model
   

6. Adicione um ficheiro commit_success.txt vazio na pasta $PATH_TO_PAX_LLAMA3 e nas subpastas metadata e state nessa pasta da seguinte forma:

   touch $PATH_TO_PAX_LLAMA3/checkpoint_00000000/commit_success.txt
   touch $PATH_TO_PAX_LLAMA3/checkpoint_00000000/metadata/commit_success.txt
   touch $PATH_TO_PAX_LLAMA3/checkpoint_00000000/state/commit_success.txt
   

   A pasta $PATH_TO_PAX_LLAMA3 contém agora as seguintes pastas e ficheiros:

   $PATH_TO_PAX_LLAMA3/checkpoint_00000000/commit_success.txt
   $PATH_TO_PAX_LLAMA3/checkpoint_00000000/metadata/
   $PATH_TO_PAX_LLAMA3/checkpoint_00000000/metadata/commit_success.txt
   $PATH_TO_PAX_LLAMA3/checkpoint_00000000/state/
   $PATH_TO_PAX_LLAMA3/checkpoint_00000000/state/commit_success.txt
   $PATH_TO_PAX_LLAMA3/vocabs/tokenizer.model
   

Crie um contentor do Cloud Storage

Crie um contentor do Cloud Storage para armazenar os pesos do modelo convertido.

1. No Cloud Shell, execute os seguintes comandos, substituindo PROJECT_ID pelo ID do seu projeto:

   projectid=PROJECT_ID
   gcloud config set project ${projectid}
   

2. Para criar o contentor, execute o seguinte comando:

   gcloud storage buckets create gs://WEIGHTS_BUCKET_NAME
   

   Substitua WEIGHTS_BUCKET_NAME pelo nome que quer usar para o contentor.

Copie os pesos do modelo para o contentor do Cloud Storage

Para copiar as ponderações do modelo para o seu contentor, execute o seguinte comando:

gcloud storage cp PATH_TO_PAX_LLAMA3/* gs://WEIGHTS_BUCKET_NAME/llama3/pax_70b/ --recursive

Carregue o modelo

Está disponível um contentor Saxml pré-criado em us-docker.pkg.dev/vertex-ai/prediction/sax-tpu:latest.

Para carregar um recurso Model para o Vertex AI usando o contentor Saxml pré-criado, execute o comando gcloud ai models upload da seguinte forma:

gcloud ai models upload \
    --region=LOCATION \
    --display-name=MODEL_DISPLAY_NAME \
    --container-image-uri=us-docker.pkg.dev/vertex-ai/prediction/sax-tpu:latest \
    --artifact-uri='gs://WEIGHTS_BUCKET_NAME/llama3/pax_70b' \
    --container-args='--model_path=saxml.server.pax.lm.params.lm_cloud.LLaMA3_70BFP16x16' \
    --container-args='--platform_chip=tpuv5e' \
    --container-args='--platform_topology=4x4' \
    --container-args='--ckpt_path_suffix=checkpoint_00000000' \
    --container-deployment-timeout-seconds=2700 \
    --container-ports=8502 \
    --project=PROJECT_ID

Faça as seguintes substituições:

• LOCATION: 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.
• MODEL_DISPLAY_NAME: o nome a apresentar que quer para o seu modelo
• PROJECT_ID: o ID do seu Google Cloud projeto
Nota: pode usar este comando para qualquer modelo Llama que esteja listado no seguinte link do GitHub: https://github.com/google/saxml/blob/main/saxml/server/pax/lm/params/lm_cloud.py

Crie um ponto final de inferência online

Para criar o ponto final, execute o seguinte comando:

gcloud ai endpoints create \
    --region=LOCATION \
    --display-name=ENDPOINT_DISPLAY_NAME \
    --project=PROJECT_ID

Substitua ENDPOINT_DISPLAY_NAME pelo nome a apresentar que quer para o seu ponto final.

Implemente o modelo no ponto final

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

Neste tutorial, implementa um modelo Llama 3 70B fragmentado para 16 chips Cloud TPU v5e com uma topologia 4x4. No entanto, pode especificar qualquer uma das seguintes topologias de Cloud TPU com vários anfitriões suportadas:

Tipo de máquina	Topologia	Número de chips de TPU	Número de anfitriões
ct5lp-hightpu-4t	4x4	16	2
ct5lp-hightpu-4t	4x8	32	4
ct5lp-hightpu-4t	8x8	64	8
ct5lp-hightpu-4t	8x16	128	16
ct5lp-hightpu-4t	16x16	256	32

Se estiver a implementar um modelo Llama diferente definido no repositório do GitHub do Saxml, certifique-se de que está particionado para corresponder ao número de dispositivos que está a segmentar e que o Cloud TPU tem memória suficiente para carregar o modelo.

Para obter informações sobre a implementação de um modelo em Cloud TPUs de anfitrião único, consulte o artigo Implemente um modelo.

Para ver uma lista completa dos tipos e regiões de TPU na nuvem suportados, consulte as localizações do Vertex AI.

1. Obtenha o ID do ponto final do ponto final de inferência online:

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

2. Obtenha o ID do modelo:

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

3. Implemente o modelo no ponto final:

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

   Substitua DEPLOYED_MODEL_NAME por um nome para o elemento implementado. Pode ser igual ao nome a apresentar do modelo (MODEL_DISPLAY_NAME).

   A operação de implementação pode exceder o tempo limite.

   O comando deploy-model devolve um ID da operação que pode ser usado para verificar quando a operação estiver concluída. Pode sondar o estado da operação até que a resposta inclua "done": true. Use o seguinte comando para consultar o estado:

   gcloud ai operations describe \
   --region=LOCATION \
   OPERATION_ID
   

   Substitua OPERATION_ID pelo ID da operação devolvido pelo comando anterior.

Obtenha inferências online do modelo implementado

Para obter inferências online a partir do ponto final do Vertex AI, execute o comando gcloud ai endpoints predict.

1. Execute o seguinte comando para criar um ficheiro request.json que contenha um pedido de inferência de exemplo:

   cat << EOF > request.json
   {"instances": [{"text_batch": "the distance between Earth and Moon is "}]}
   EOF
   

2. Para enviar o pedido de inferência online para o ponto final, execute o seguinte comando:

   gcloud ai endpoints predict $ENDPOINT_ID \
       --project=PROJECT_ID \
       --region=LOCATION \
       --json-request=request.json
   

Limpar

Para evitar incorrer em mais custos do Vertex AI, 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 os seguintes comandos:

   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
   

2. Para eliminar o modelo, execute os seguintes comandos:

   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