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 https://huggingface.co/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:
- Faça login na sua Conta do Google. Se ainda não tiver uma, crie uma nova conta.
- No console do Google Cloud, selecione ou crie um projeto Google Cloud na página do seletor de projetos.
- Ative o faturamento do projeto Google Cloud . O faturamento é obrigatório para todo o uso do Google Cloud .
- Instale os componentes da gcloud alfa.
Execute o comando a seguir para instalar a versão mais recente dos componentes
gcloud
.gcloud components update
Ative a API TPU usando o comando
gcloud
a seguir no Cloud Shell. Também é possível ativá-la no console do Google Cloud.gcloud services enable tpu.googleapis.com
Crie uma identidade de serviço para a VM da TPU.
gcloud alpha compute tpus tpu-vm service-identity create --zone=ZONE
Crie uma conta de serviço do TPU e conceda acesso aos serviços Google Cloud .
As contas de serviço permitem que o serviço Google Cloud TPU acesse outros serviços Google Cloud. Recomendamos usar uma conta de serviço gerenciado pelo usuário. Siga estes guias para criar e conceder funções. Os seguintes papéis são necessários:
- Administrador de TPU: necessário para criar uma TPU
- Administrador de armazenamento: necessário para acessar o Cloud Storage
- Logs Writer: necessário para gravar registros com a API Logging.
- Gravador de métricas do Monitoring: necessário para gravar métricas no Cloud Monitoring
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
Entre em contato com a equipe de vendas ou de contas do Cloud TPU para solicitar a cota de TPU e fazer perguntas sobre a capacidade.
Provisionar o ambiente do Cloud TPU
É possível provisionar TPUs v6e 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 . - Este tutorial foi testado com a seguinte configuração:
- Python
3.10 or later
- Versões noturnas do software:
0.4.32.dev20240912
noturno JAX- LibTPU noturno
0.1.dev20240912+nightly
- Versões estáveis do software:
- JAX + JAX Lib de
v0.4.35
- JAX + JAX Lib de
- Python
- Verifique se o projeto tem cota suficiente de TPU para:
- Cota da VM de TPU
- Quota de endereço IP
- Cota 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
Descrições de sinalizações de comando
Variável | Descrição |
NODE_ID | O ID atribuído pelo usuário do TPU que é criado quando a solicitação de recurso em fila é alocada. |
PROJECT_ID | Google Cloud nome do projeto. Use um projeto existente ou crie um novo.> |
ZONA | Consulte o documento Regiões e zonas de TPU para conferir as zonas compatíveis. |
ACCELERATOR_TYPE | Consulte a documentação Tipos de acelerador para saber quais são os tipos de aceleradores compatíveis. |
RUNTIME_VERSION | v2-alpha-tpuv6e
|
SERVICE_ACCOUNT | Esse é o endereço de e-mail da sua conta de serviço, que pode ser encontrado em
console do Google Cloud -> IAM -> Contas de serviço
Por exemplo: tpu-service-account@<your_project_ID>.iam.gserviceaccount.com.com |
Use os comandos list
ou describe
para consultar o status do recurso em fila.
gcloud alpha compute tpus queued-resources describe ${QUEUED_RESOURCE_ID} \
--project ${PROJECT_ID} --zone ${ZONE}
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
Crie um diretório para o Miniconda:
mkdir -p ~/miniconda3
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
Instale o Miniconda:
bash ~/miniconda3/miniconda.sh -b -u -p ~/miniconda3
Remova o script de instalação do Miniconda:
rm -rf ~/miniconda3/miniconda.sh
Adicione o Miniconda à variável
PATH
:export PATH="$HOME/miniconda3/bin:$PATH"
Atualize
~/.bashrc
para aplicar as mudanças à variávelPATH
:source ~/.bashrc
Crie um ambiente da Conda:
conda create -n vllm python=3.10 -y conda activate vllm
Clone o repositório vLLM e navegue até o diretório vLLM:
git clone https://github.com/vllm-project/vllm.git && cd vllm
Limpe os pacotes torch e torch-xla:
pip uninstall torch torch-xla -y
Instale outras dependências do build:
pip install -r requirements-tpu.txt VLLM_TARGET_DEVICE="tpu" python setup.py develop 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:
- Clique em Seu perfil > Configurações > Tokens de acesso.
- Selecione Novo token.
- Especifique um nome e um papel com pelo menos
Read
permissões. - Selecione Gerar um token.
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
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 de vLLM.
cd ~/vllm
vllm serve "meta-llama/Meta-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/Meta-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