Inferensi vLLM di TPU v6e

Tutorial ini menunjukkan cara menjalankan inferensi vLLM di TPU v6e. Bagian ini juga menunjukkan cara menjalankan skrip benchmark untuk model Meta Llama-3.1 8B.

Untuk mulai menggunakan vLLM di TPU v6e, lihat panduan memulai vLLM.

Jika Anda menggunakan GKE, lihat juga tutorial GKE.

Sebelum memulai

Anda harus menandatangani perjanjian izin untuk menggunakan keluarga model Llama3 di repo HuggingFace. Buka meta-llama/Llama-3.1-8B, isi perjanjian izin, dan tunggu hingga Anda disetujui.

Bersiaplah untuk menyediakan TPU v6e dengan 4 chip:

  1. Ikuti panduan Menyiapkan lingkungan Cloud TPU untuk menyiapkan Google Cloud project, mengonfigurasi Google Cloud CLI, mengaktifkan Cloud TPU API, dan memastikan Anda memiliki akses untuk menggunakan Cloud TPU.

  2. Lakukan autentikasi dengan Google Cloud dan konfigurasikan project dan zona default untuk Google Cloud CLI.

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

Kapasitas aman

Jika Anda siap untuk mengamankan kapasitas TPU, lihat Kuota Cloud TPU untuk mengetahui informasi selengkapnya tentang kuota Cloud TPU. Jika ada pertanyaan tambahan tentang cara mengamankan kapasitas, hubungi tim akun atau tim penjualan Cloud TPU Anda.

Menyediakan lingkungan Cloud TPU

Anda dapat menyediakan VM TPU dengan GKE, dengan GKE dan XPK, atau sebagai resource dalam antrean.

Prasyarat

  • Pastikan project Anda memiliki kuota TPUS_PER_TPU_FAMILY yang cukup, yang menentukan jumlah maksimum chip yang dapat Anda akses dalam projectGoogle Cloud .
  • Pastikan project Anda memiliki cukup kuota TPU untuk:
    • Kuota VM TPU
    • Kuota alamat IP
    • Kuota Hyperdisk Balanced
  • Izin project pengguna

Menyediakan 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

Deskripsi flag perintah

Variabel Deskripsi
NODE_ID ID TPU yang ditetapkan pengguna yang dibuat saat permintaan resource dalam antrean dialokasikan.
PROJECT_ID Google Cloud nama project. Gunakan project yang ada atau buat project baru.
ZONA Lihat dokumen Region dan zona TPU untuk zona yang didukung.
ACCELERATOR_TYPE Lihat dokumentasi Jenis Akselerator untuk jenis akselerator yang didukung.
RUNTIME_VERSION v2-alpha-tpuv6e
SERVICE_ACCOUNT Ini adalah alamat email untuk akun layanan yang dapat Anda temukan di Google Cloud console -> IAM -> Service Accounts.

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

Gunakan perintah list atau describe untuk membuat kueri status resource yang diantrekan.

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

Untuk mengetahui daftar lengkap status permintaan resource yang diantrekan, lihat dokumentasi Resource yang Diantrekan.

Menghubungkan ke TPU menggunakan SSH

  gcloud compute tpus tpu-vm ssh TPU_NAME

Menginstal dependensi

  1. Buat direktori untuk Miniconda:

    mkdir -p ~/miniconda3
  2. Download skrip penginstal Miniconda:

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

    bash ~/miniconda3/miniconda.sh -b -u -p ~/miniconda3
  4. Hapus skrip penginstal Miniconda:

    rm -rf ~/miniconda3/miniconda.sh
  5. Tambahkan Miniconda ke variabel PATH Anda:

    export PATH="$HOME/miniconda3/bin:$PATH"
  6. Muat ulang ~/.bashrc untuk menerapkan perubahan pada variabel PATH:

    source ~/.bashrc
  7. Buat lingkungan Conda:

    conda create -n vllm python=3.11 -y
    conda activate vllm
  8. Clone repositori vLLM dan buka direktori vllm:

    git clone https://github.com/vllm-project/vllm.git && cd vllm
    
  9. Bersihkan paket torch dan torch-xla yang ada:

    pip uninstall torch torch-xla -y
    
  10. Instal dependensi build lainnya:

    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
    

Mendapatkan akses ke model

Buat token Hugging Face baru jika Anda belum memilikinya:

  1. Klik Profil Anda > Setelan > Token Akses.

  2. Pilih New Token.

  3. Tentukan Nama pilihan Anda dan Peran dengan setidaknya Read izin.

  4. Pilih Buat token.

  5. Salin token yang dihasilkan ke papan klip Anda, tetapkan sebagai variabel lingkungan, dan autentikasi dengan huggingface-cli:

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

Meluncurkan server vLLM

Perintah berikut mendownload bobot model dari Hugging Face Model Hub ke direktori /tmp VM TPU, mengompilasi berbagai bentuk input terlebih dahulu, dan menulis kompilasi model ke ~/.cache/vllm/xla_cache.

Untuk mengetahui detail selengkapnya, lihat dokumen 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 &

Menjalankan benchmark vLLM

Jalankan skrip benchmark vLLM:

export MODEL = <var>your-model</var>
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

Pembersihan

Hapus TPU:

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