Disponibilizar o Llama-3.1-8B com o vLLM em TPUs

Neste tutorial, mostramos como disponibilizar o modelo meta-llama/Llama-3.1-8B usando o framework de disponibilização de TPU vLLM em uma VM de TPU v6e.

Objetivos

1. Prepare o ambiente.
2. Execute o vLLM com o Llama-3.1-8B.
3. Envie uma solicitação de inferência.
4. Execute uma carga de trabalho de comparativo de mercado.
5. Fazer a limpeza.

Custos

Neste tutorial, há componentes faturáveis do Google Cloud, entre eles:

• Cloud TPU
• Compute Engine

Para gerar uma estimativa de custo baseada na projeção de uso deste tutorial, use a calculadora de preços.

Antes de começar

Antes de seguir este tutorial, siga as instruções na página Configurar o ambiente do Cloud TPU. As instruções orientam você pelas etapas necessárias para criar um projeto do Google Cloud e configurá-lo para usar o Cloud TPU. Também é possível usar um projeto doGoogle Cloud . Se preferir, pule a etapa de criação de um projetoGoogle Cloud e comece com Configurar o ambiente para usar o Cloud TPU.

Você precisa de um token de acesso do Hugging Face para usar este tutorial. Você pode se inscrever para uma conta sem custo financeiro no Hugging Face. Depois de ter uma conta, gere um token de acesso:

1. Na página Bem-vindo ao Hugging Face, clique no avatar da sua conta e selecione Tokens de acesso.
2. Na página Tokens de acesso, clique em Criar novo token.
3. Selecione o tipo de token Read e insira um nome para ele.
4. Seu token de acesso vai aparecer. Salve o token em um local seguro.

Observação: antes de continuar com este tutorial, observe que ele usa um modelo restrito no Hugging Face. Você precisa solicitar e receber a aprovação do modelo pelos criadores na página do modelo meta-llama/Llama-3.1-8B.

Configurar o ambiente

Recursos na fila

Crie uma VM do Cloud TPU v6e usando a API de recursos em fila. Para o Llama-3.1-8B, recomendamos usar uma TPU v6e-1.

export PROJECT_ID=<PROJECT>
export TPU_NAME=<TPU_NAME>
export ZONE=<ZONE>
export QR_ID=<QR_ID>
export TPU_TYPE=<TPU_TYPE>

Defina as variáveis:

• PROJECT: o nome do projeto.
• TPU_NAME: nome da máquina da VM de TPU que você vai criar.
• ZONE: a zona da nuvem em que você cria a nova VM.
• TPU_TYPE: o tipo de VM de TPU que você cria. Por exemplo, v6e-1 ou v6e-4.
• QR_ID: nome do recurso em fila que você cria.

Crie a solicitação de recurso em fila:

gcloud alpha compute tpus queued-resources create $QR_ID \
 --node-id $TPU_NAME \
 --project $PROJECT_ID \
 --zone $ZONE \
 --accelerator-type $TPU_TYPE \
 --runtime-version v2-alpha-tpuv6e

Verifique se a VM da TPU está pronta.

gcloud compute tpus queued-resources describe $QR_ID \
  --project $PROJECT_ID \
  --zone $ZONE

Dica: inicialmente, o status da solicitação de recurso em fila será WAITING_FOR_RESOURCES. Quando a VM de TPU for criada, o status vai mudar para ACTIVE. Isso pode levar algum tempo, dependendo da disponibilidade de recursos. Se o tempo de espera for muito longo, tente usar outra zona. Aguarde até que o status seja ACTIVE antes de continuar.

Por exemplo, quando o status é ACTIVE:

name: projects/your-project-id/locations/your-zone/queuedResources/your-queued-resource-id
  state:
  state: ACTIVE
  tpu:
  nodeSpec:
  - node:
      acceleratorType: v6e-1
      bootDisk: {}
      networkConfig:
          enableExternalIps: true
      queuedResource: projects/your-project-number/locations/your-zone/queuedResources/your-queued-resource-id
      runtimeVersion: v2-alpha-tpuv6e
      schedulingConfig: {}
      serviceAccount: {}
      shieldedInstanceConfig: {}
      useTpuVm: true
      nodeId: your-node-id
      parent: projects/your-project-number/locations/your-zone

Reserva

Crie uma VM do Cloud TPU v6e usando uma reserva. Para o Llama-3.1-8B, recomendamos usar uma TPU v6e-1. Comece definindo as variáveis de ambiente:

export PROJECT_ID="YOUR_PROJECT_ID"
export TPU_NAME="TPU_MACHINE_NAME"
export ZONE="ZONE_NAME"
export TPU_TYPE=v6e-8
export HF_TOKEN="YOUR_HF_TOKEN"
export RESERVATION="RESERVATION_NAME"

Defina as variáveis:

• PROJECT: o nome do projeto.
• TPU_NAME: nome da máquina da VM de TPU que você vai criar.
• ZONE: a zona da nuvem em que você cria a nova VM.
• TPU_TYPE: o tipo de VM de TPU que você cria. Por exemplo, v6e-1 ou v6e-4.
• RESERVATION: nome da reserva com suas TPUs.

Crie a VM de TPU usando sua reserva:

gcloud alpha compute tpus tpu-vm create $TPU_NAME \
    --zone=$ZONE \
    --project $PROJECT_ID \
    --accelerator-type=$TPU_TYPE \
    --version=v2-alpha-tpuv6e \
    --provisioning-model=reservation-bound \
    --reservation=$RESERVATION

Conecte-se à VM de TPU.

gcloud compute tpus tpu-vm ssh $TPU_NAME \
  --project $PROJECT_ID \
  --zone $ZONE

Executar o vLLM com o Llama-3.1-8B

1. Defina as variáveis de token e nome do modelo do Hugging Face.

     export HF_TOKEN="YOUR_HF_TOKEN"
     export MODEL_NAME="meta-llama/Llama-3.1-8B"
   

2. Na VM de TPU, execute o contêiner do Docker vLLM no modo desvinculado e inicie o servidor vLLM. Esse comando usa um tamanho de memória compartilhada de 10 GB.

   export DOCKER_URI="vllm/vllm-tpu:v0.18.0"
   export CONTAINER_NAME="${USER}-vllm"
   export MAX_MODEL_LEN=4096
   export TP=1 # number of chips
   
   sudo docker run -d --name "${CONTAINER_NAME}" \
       --privileged --net=host \
       -v /dev/shm:/dev/shm \
       --shm-size 10gb \
       -e "HF_HOME=/dev/shm" \
       -e "HF_TOKEN=${HF_TOKEN}" \
       -p 8000:8000 "${DOCKER_URI}" \
           vllm serve ${MODEL_NAME} \
               --seed 42 \
               --gpu-memory-utilization 0.98 \
               --max-num-batched-tokens 1024 \
               --max-num-seqs 128 \
               --tensor-parallel-size $TP \
               --max-model-len $MAX_MODEL_LEN

3. Verifique os registros do servidor para confirmar se ele está em execução.

   sudo docker logs -f "${CONTAINER_NAME}"

   Quando o servidor vLLM estiver em execução, você vai ver uma saída parecida com esta: Depois que a saída aparecer, pressione CTRL+C para voltar ao terminal.

   (APIServer pid=7) INFO:     Started server process [7]
   (APIServer pid=7) INFO:     Waiting for application startup.
   (APIServer pid=7) INFO:     Application startup complete.
   

Enviar uma solicitação de inferência

Depois que o servidor vLLM estiver em execução, você poderá enviar solicitações para a API. Para mais informações, consulte a documentação de referência da API vLLM.

1. Envie uma solicitação de teste ao servidor usando curl.

   sudo docker exec "${CONTAINER_NAME}" \
     curl http://localhost:8000/v1/completions \
       -H "Content-Type: application/json" \
       -d '{
           "prompt": "The future of AI is",
           "max_tokens": 200,
           "temperature": 0
         }'

A resposta é retornada no formato JSON.

Executar uma carga de trabalho de comparativo

Você pode executar comparativos de mercado no servidor em execução no segundo terminal.

1. Dentro do contêiner, instale a biblioteca datasets.

   sudo docker exec "${CONTAINER_NAME}" pip install datasets

2. Dentro do contêiner, execute o comando vllm bench serve.

   sudo docker exec "${CONTAINER_NAME}" \
       vllm bench serve \
           --backend vllm \
           --dataset-name random \
           --num-prompts 1000 \
           --seed 100

Os resultados do comparativo aparecem da seguinte maneira:

============ Serving Benchmark Result ============
Successful requests:                     1000
Failed requests:                         0
Benchmark duration (s):                  73.97
Total input tokens:                      1024000
Total generated tokens:                  128000
Request throughput (req/s):              13.52
Output token throughput (tok/s):         1730.38
Peak output token throughput (tok/s):    2522.00
Peak concurrent requests:                1000.00
Total Token throughput (tok/s):          15573.42
---------------Time to First Token----------------
Mean TTFT (ms):                          34834.97
Median TTFT (ms):                        34486.19
P99 TTFT (ms):                           70234.40
-----Time per Output Token (excl. 1st token)------
Mean TPOT (ms):                          47.30
Median TPOT (ms):                        48.57
P99 TPOT (ms):                           48.60
---------------Inter-token Latency----------------
Mean ITL (ms):                           47.31
Median ITL (ms):                         53.49
P99 ITL (ms):                            54.58
==================================================

Limpar

Para evitar cobranças na sua conta do Google Cloud pelos recursos usados no tutorial, exclua o projeto que os contém ou mantenha o projeto e exclua os recursos individuais.

1. No terminal, digite exit para se desconectar da VM de TPU.

Excluir os recursos

Você pode excluir o projeto, o que vai remover todos os recursos, ou manter o projeto e excluir os recursos.

Excluir o projeto

Para excluir o projeto Google Cloud e todos os recursos associados, execute:

    gcloud projects delete $PROJECT_ID

Excluir recursos da TPU

Recursos na fila

Exclua os recursos do Cloud TPU. O comando a seguir exclui a solicitação de recurso em fila e a VM de TPU usando o parâmetro --force.

gcloud alpha compute tpus queued-resources delete $QR_ID \
  --project=$PROJECT_ID \
  --zone=$ZONE \
  --force

Reserva

Exclua a VM do Cloud TPU. Use o comando a seguir para encerrar a VM, liberando as TPUs de volta para sua reserva.

gcloud compute tpus tpu-vm delete $TPU_NAME --zone $ZONE --project $PROJECT_ID --quiet

Observação: você ainda vai pagar pelas TPUs enquanto tiver a reserva.

A seguir

• Saiba mais sobre o vLLM no Cloud TPU.
• Saiba mais sobre o Cloud TPU.