Inferência de vLLM em TPUs v6e

Este tutorial mostra como executar a inferência de vLLM em TPUs v6e. Ele também mostra como executar o script de comparação para o modelo Meta Llama 3.1 8B.

Para começar a usar o vLLM em TPUs v6e, consulte o Guia de início rápido do vLLM.

Se você estiver usando o GKE, consulte também o tutorial do GKE.

Antes de começar

É necessário assinar o contrato de consentimento para usar a família de modelos Llama3 no repositório do Hugging Face. Acesse meta-llama/Llama-3.1-8B, preencha o contrato de consentimento e aguarde a aprovação.

Prepare-se para provisionar uma TPU v6e com 4 chips:

1. Siga o guia Configurar o ambiente do Cloud TPU para configurar um projeto Google Cloud , configurar o CLI do Google Cloud, ativar a API Cloud TPU e garantir que você tenha acesso para usar o Cloud TPU.

2. Faça a autenticação com Google Cloud e configure o projeto e a zona padrão para a Google Cloud CLI.

   gcloud auth login
   gcloud config set project PROJECT_ID
   gcloud config set compute/zone ZONE

Capacidade segura

Quando estiver tudo pronto para garantir a capacidade da TPU, consulte Cotas da Cloud TPU para mais informações. Se você tiver outras dúvidas sobre como garantir a capacidade, entre em contato com a equipe de vendas ou de conta do Cloud TPU.

Provisionar o ambiente do Cloud TPU

É possível provisionar VMs do TPU com o GKE, com o GKE e o XPK ou como recursos em fila.

Pré-requisitos

• Verifique se o projeto tem cota de TPUS_PER_TPU_FAMILY suficiente, que especifica o número máximo de chips que você pode acessar no projetoGoogle Cloud .
• Verifique se o projeto tem cota suficiente de TPU para:
  • Cota da VM de TPU
  • Cota de endereços IP
  • Quota do Hyperdisk equilibrado
• Permissões do projeto do usuário
  • Se você estiver usando o GKE com XPK, consulte Permissões do console do Google Cloud na conta de usuário ou de serviço para conferir as permissões necessárias para executar o XPK.

Provisionar um TPU v6e

   gcloud alpha compute tpus queued-resources create QUEUED_RESOURCE_ID \
      --node-id TPU_NAME \
       --project PROJECT_ID \
       --zone ZONE \
       --accelerator-type v6e-4 \
       --runtime-version v2-alpha-tpuv6e \
       --service-account SERVICE_ACCOUNT

Use os comandos list ou describe para consultar o status do recurso em fila.

   gcloud alpha compute tpus queued-resources describe <var>QUEUED_RESOURCE_ID}</var>  \
      --project <var>PROJECT_ID</var> --zone <var>ZONE</var>

Para uma lista completa de status de solicitações de recursos em fila, consulte a documentação de Recursos em fila.

Conectar-se à TPU usando SSH

  gcloud compute tpus tpu-vm ssh TPU_NAME

Instalar dependências

1. Crie um diretório para o Miniconda:

   mkdir -p ~/miniconda3

2. Faça o download do script de instalação do Miniconda:

   wget https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-x86_64.sh -O ~/miniconda3/miniconda.sh

3. Instale o Miniconda:

   bash ~/miniconda3/miniconda.sh -b -u -p ~/miniconda3

4. Remova o script do instalador do Miniconda:

   rm -rf ~/miniconda3/miniconda.sh

5. Adicione o Miniconda à variável PATH:

   export PATH="$HOME/miniconda3/bin:$PATH"

6. Atualize ~/.bashrc para aplicar as mudanças à variável PATH:

   source ~/.bashrc

7. Crie um ambiente da Conda:

   conda create -n vllm python=3.11 -y
   conda activate vllm

8. Clone o repositório vLLM e navegue até o diretório vLLM:

   git clone https://github.com/vllm-project/vllm.git && cd vllm
   

9. Limpe os pacotes torch e torch-xla:

   pip uninstall torch torch-xla -y
   

10. Instale outras dependências do build:

    pip install -r requirements/tpu.txt
    VLLM_TARGET_DEVICE="tpu" python -m pip install --editable .
    sudo apt-get install libopenblas-base libopenmpi-dev libomp-dev
    

Receber acesso ao modelo

Gere um novo token do Huggin' Face, caso ainda não tenha um:

1. Clique em Seu perfil > Configurações > Tokens de acesso.
2. Selecione Novo token.
3. Especifique um nome e um papel com pelo menos Read permissões.
4. Selecione Gerar um token.

5. Copie o token gerado para a área de transferência, defina-o como uma variável de ambiente e faça a autenticação com o huggingface-cli:

   export TOKEN=YOUR_TOKEN
   git config --global credential.helper store
   huggingface-cli login --token $TOKEN

Fazer o download de dados de comparativo de mercado

1. Crie um diretório /data e faça o download do conjunto de dados ShareGPT do Hugging Face.

   mkdir ~/data && cd ~/data
   wget https://huggingface.co/datasets/anon8231489123/ShareGPT_Vicuna_unfiltered/resolve/main/ShareGPT_V3_unfiltered_cleaned_split.json
   

Iniciar o servidor vLLM

O comando a seguir faz o download dos pesos do modelo do Hugging Face Model Hub para o diretório /tmp da VM do TPU, pré-compila uma variedade de formas de entrada e grava a compilação do modelo em ~/.cache/vllm/xla_cache.

Para mais detalhes, consulte os documentos do vLLM.

   cd ~/vllm
   vllm serve "meta-llama/Llama-3.1-8B" --download_dir /tmp --num-scheduler-steps 4 --swap-space 16 --disable-log-requests --tensor_parallel_size=4 --max-model-len=2048 &> serve.log &

Executar comparativos de vLLM

Execute o script de comparação de vLLMs:

   python benchmarks/benchmark_serving.py \
       --backend vllm \
       --model "meta-llama/Llama-3.1-8B"  \
       --dataset-name sharegpt \
       --dataset-path ~/data/ShareGPT_V3_unfiltered_cleaned_split.json  \
       --num-prompts 1000

Limpar

Exclua a TPU:

gcloud compute tpus queued-resources delete QUEUED_RESOURCE_ID \
    --project PROJECT_ID \
    --zone ZONE \
    --force \
    --async