Interruzione vLLM su TPU v6e

Questo tutorial mostra come eseguire l'inferenza vLLM su TPU v6e. Inoltre, mostra come eseguire lo script di benchmark per il modello Meta Llama-3.1-8B.

Per iniziare a utilizzare vLLM sulle TPU v6e, consulta la guida rapida a vLLM.

Se utilizzi GKE, consulta anche il tutorial su GKE.

Prima di iniziare

Devi firmare il contratto di consenso per utilizzare la famiglia di modelli Llama3 nel repo HuggingFace. Vai a meta-llama/Llama-3.1-8B, compila il contratto per il consenso e attendi l'approvazione.

Preparati a eseguire il provisioning di una TPU v6e con 4 chip:

1. Segui la guida Configurare l'ambiente Cloud TPU per configurare un Google Cloud progetto, configurare Google Cloud CLI, attivare l'API Cloud TPU e assicurarti di avere accesso all'utilizzo delle Cloud TPU.

2. Esegui l'autenticazione con Google Cloud e configura il progetto e la zona predefiniti per Google Cloud CLI.

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

Capacità sicura

Quando è tutto pronto per garantire la capacità TPU, consulta Quote Cloud TPU per saperne di più sulle quote Cloud TPU. Se hai altre domande sulla prenotazione della capacità, contatta il team di vendita o di gestione dell'account Cloud TPU.

Esegui il provisioning dell'ambiente Cloud TPU

Puoi eseguire il provisioning delle VM TPU con GKE, con GKE e XPK, o come risorse in coda.

Prerequisiti

• Verifica che il tuo progetto disponga di una quota TPUS_PER_TPU_FAMILY sufficiente, che specifica il numero massimo di chip a cui puoi accedere all'interno del progettoGoogle Cloud .
• Verifica che il tuo progetto disponga di una quota TPU sufficiente per:
  • Quota VM TPU
  • Quota di indirizzi IP
  • Quota Hyperdisk bilanciato
• Autorizzazioni del progetto per gli utenti
  • Se utilizzi GKE con XPK, consulta Autorizzazioni di Cloud Console nell'account utente o di servizio per conoscere le autorizzazioni necessarie per eseguire XPK.

Esegui il provisioning di una 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

Utilizza i comandi list o describe per eseguire query sullo stato della risorsa in coda.

gcloud alpha compute tpus queued-resources describe QUEUED_RESOURCE_ID  \
  --project PROJECT_ID --zone ZONE

Per un elenco completo degli stati delle richieste di risorse in coda, consulta la documentazione relativa alle risorse in coda.

Connettiti alla TPU tramite SSH

  gcloud compute tpus tpu-vm ssh TPU_NAME

Installa le dipendenze

1. Crea una directory per Miniconda:

   mkdir -p ~/miniconda3

2. Scarica lo script di installazione di Miniconda:

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

3. Installa Miniconda:

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

4. Rimuovi lo script di installazione di Miniconda:

   rm -rf ~/miniconda3/miniconda.sh

5. Aggiungi Miniconda alla variabile PATH:

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

6. Ricarica ~/.bashrc per applicare le modifiche alla variabile PATH:

   source ~/.bashrc

7. Crea un ambiente Conda:

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

8. Clona il repository vLLM e vai alla directory vllm:

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

9. Ripulisci i pacchetti torch e torch-xla esistenti:

   pip uninstall torch torch-xla -y
   

10. Installa altre dipendenze di compilazione:

    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
    

Ottieni l'accesso al modello

Genera un nuovo token Abbracciamento se non ne hai già uno:

1. Vai a Il tuo profilo > Impostazioni > Token di accesso.

2. Seleziona Crea nuovo token.

3. Specifica un nome a tua scelta e un ruolo con almeno Read autorizzazioni.

4. Seleziona Genera un token.

5. Copia il token generato negli appunti, impostalo come variabile di ambiente e autenticati con huggingface-cli:

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

Avvia il server vLLM

Il seguente comando scarica i pesi del modello da Hugging Face Model Hub nella directory /tmp della VM TPU, precompila una serie di forme di input e scrive la compilazione del modello in ~/.cache/vllm/xla_cache.

Per maggiori dettagli, consulta la documentazione di vLLM.

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

Eseguire benchmark vLLM

Esegui lo script di benchmarking vLLM:

export MODEL="meta-llama/Llama-3.1-8B"
pip install pandas
pip install datasets
python benchmarks/benchmark_serving.py \
  --backend vllm \
  --model $MODEL  \
  --dataset-name random \
  --random-input-len 1820 \
  --random-output-len 128 \
  --random-prefix-len 0

Esegui la pulizia

Elimina la TPU:

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