Inferencia de vLLM en TPUs v6e

En este tutorial se muestra cómo ejecutar la inferencia de vLLM en TPUs v6e. También se explica cómo ejecutar la secuencia de comandos de la prueba comparativa del modelo Meta Llama-3.1-8B.

Para empezar a usar vLLM en las TPUs v6e, consulta la guía de inicio rápido de vLLM.

Si usas GKE, consulta también el tutorial de GKE.

Antes de empezar

Debes firmar el contrato de consentimiento para usar la familia de modelos Llama 3 en el repositorio de Hugging Face. Ve a meta-llama/Llama-3.1-8B, rellena el acuerdo de consentimiento y espera a que se te apruebe.

Prepara el aprovisionamiento de una TPU v6e con 4 chips:

  1. Sigue la guía Configurar el entorno de TPU de Cloud para configurar un proyecto de Google Cloud , configurar la CLI de Google Cloud, habilitar la API Cloud TPU y asegurarte de que tienes acceso para usar las TPU de Cloud.

  2. Autentícate con Google Cloud y configura el proyecto y la zona predeterminados de la CLI de Google Cloud.

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

Capacidad segura

Cuando quieras proteger la capacidad de las TPUs, consulta Cuotas de TPUs de Cloud para obtener más información sobre las cuotas de TPUs de Cloud. Si tienes más preguntas sobre cómo proteger la capacidad, ponte en contacto con el equipo de Ventas o de Cuentas de Cloud TPU.

Aprovisionar el entorno de TPU de Cloud

Puedes aprovisionar VMs de TPU con GKE, con GKE y XPK o como recursos en cola.

Requisitos previos

  • Comprueba que tu proyecto tenga suficiente cuota de TPUS_PER_TPU_FAMILY, que especifica el número máximo de chips a los que puedes acceder en tuGoogle Cloud proyecto.
  • Comprueba que tu proyecto tenga suficiente cuota de TPU para lo siguiente:
    • Cuota de máquinas virtuales de TPU
    • Cuota de direcciones IP
    • Cuota de Hyperdisk Balanced
  • Permisos de proyecto de usuario

Aprovisionar 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

Descripciones de marcas de comandos

Variable Descripción
NODE_ID Es el ID asignado por el usuario de la TPU que se crea cuando se asigna la solicitud de recurso en cola.
PROJECT_ID El Google Cloud nombre del proyecto. Usa un proyecto que ya tengas o crea uno nuevo.
ZONE Consulta el documento Regiones y zonas de TPU para ver las zonas admitidas.
ACCELERATOR_TYPE Consulta la documentación sobre los tipos de aceleradores para ver los tipos de aceleradores admitidos.
RUNTIME_VERSION v2-alpha-tpuv6e
SERVICE_ACCOUNT Esta es la dirección de correo de tu cuenta de servicio, que puedes encontrar en Google Cloud consola > Gestión de identidades y accesos > Cuentas de servicio.

Por ejemplo: tpu-service-account@<your_project_ID>.iam.gserviceaccount.com

Usa los comandos list o describe para consultar el estado del recurso en cola.

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

Para ver una lista completa de los estados de las solicitudes de recursos en cola, consulta la documentación sobre recursos en cola.

Conectarse a la TPU mediante SSH

  gcloud compute tpus tpu-vm ssh TPU_NAME

Instalar dependencias

  1. Crea un directorio para Miniconda:

    mkdir -p ~/miniconda3

  2. Descarga la secuencia de comandos del instalador de Miniconda:

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

  3. Instala Miniconda:

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

  4. Elimina la secuencia de comandos del instalador de Miniconda:

    rm -rf ~/miniconda3/miniconda.sh

  5. Añade Miniconda a tu variable PATH:

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

  6. Vuelve a cargar ~/.bashrc para aplicar los cambios a la variable PATH:

    source ~/.bashrc

  7. Crea un entorno de Conda:

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

  8. Clona el repositorio de vLLM y ve al directorio vllm:

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

  9. Limpia los paquetes torch y torch-xla:

    pip uninstall torch torch-xla -y

  10. Instala otras dependencias de compilación:

    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

Acceder al modelo

Genera un nuevo token de Hugging Face si aún no tienes uno:

  1. Ve a Tu perfil > Configuración > Tokens de acceso.

  2. Selecciona Crear token.

  3. Especifica el nombre que quieras y un rol con al menos Read permisos.

  4. Selecciona Generar un token.

  5. Copia el token generado en el portapapeles, configúralo como variable de entorno y autentícate con huggingface-cli:

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

Iniciar el servidor vLLM

El siguiente comando descarga los pesos del modelo desde Hugging Face Model Hub a la VM de TPU, en el directorio /tmp, precompila una serie de formas de entrada y escribe la compilación del modelo en ~/.cache/vllm/xla_cache.

Para obtener más información, consulta la documentación de 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 &

Ejecutar comparativas de vLLM

Ejecuta la secuencia de comandos de comparativas de 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

Limpieza

Elimina la TPU:

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