Inferensi PyTorch JetStream pada VM TPU v6e

Tutorial ini menunjukkan cara menggunakan JetStream untuk menayangkan model PyTorch di TPU v6e. JetStream adalah mesin yang dioptimalkan untuk throughput dan memori untuk inferensi model bahasa besar (LLM) di perangkat XLA (TPU). Dalam tutorial ini, Anda akan menjalankan benchmark inferensi untuk model Llama2-7B.

Sebelum memulai

Bersiap 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 mendapatkan 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

Membuat variabel lingkungan

Di Cloud Shell, buat variabel lingkungan berikut:

export PROJECT_ID=your-project-id
export TPU_NAME=your-tpu-name
export ZONE=us-central2-b
export ACCELERATOR_TYPE=v6e-4
export RUNTIME_VERSION=v2-alpha-tpuv6e
export SERVICE_ACCOUNT=your-service-account
export QUEUED_RESOURCE_ID=your-queued-resource-id

Deskripsi variabel lingkungan

Variabel Deskripsi
PROJECT_ID Project ID Google Cloud Anda. Gunakan project yang ada atau buat project baru.
TPU_NAME Nama TPU.
ZONE Zona tempat VM TPU akan dibuat. Untuk mengetahui informasi selengkapnya tentang zona yang didukung, lihat Region dan zona TPU.
ACCELERATOR_TYPE Jenis akselerator menentukan versi dan ukuran Cloud TPU yang ingin Anda buat. Untuk mengetahui informasi selengkapnya tentang jenis akselerator yang didukung untuk setiap versi TPU, lihat versi TPU.
RUNTIME_VERSION Versi software Cloud TPU.
SERVICE_ACCOUNT Alamat email untuk akun layanan Anda. Anda dapat menemukannya dengan membuka halaman Akun Layanan di konsol Google Cloud .

Misalnya: tpu-service-account@PROJECT_ID.iam.gserviceaccount.com

QUEUED_RESOURCE_ID ID teks yang ditetapkan pengguna untuk permintaan resource yang diantrekan.

Menyediakan TPU v6e

    gcloud alpha compute tpus queued-resources create ${QUEUED_RESOURCE_ID} \
        --node-id ${TPU_NAME} \
        --project ${PROJECT_ID} \
        --zone ${ZONE} \
        --accelerator-type ${ACCELERATOR_TYPE} \
        --runtime-version ${RUNTIME_VERSION} \
        --service-account ${SERVICE_ACCOUNT}

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}

Menjalankan benchmark JetStream PyTorch Llama2-7B

Untuk menyiapkan JetStream-PyTorch, mengonversi checkpoint model, dan menjalankan benchmark inferensi, ikuti petunjuk di repositori GitHub.

Setelah benchmark inferensi selesai, pastikan untuk membersihkan resource TPU.

Pembersihan

Hapus TPU:

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