Inferensi MaxText JetStream pada VM TPU v6e

Tutorial ini menunjukkan cara menggunakan JetStream untuk menayangkan model MaxText 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. Login ke Akun Google Anda. Jika Anda belum melakukannya, daftar untuk membuat akun baru.
  2. Di konsol Google Cloud, pilih atau buat project Google Cloud dari halaman pemilih project.
  3. Aktifkan penagihan untuk project Google Cloud Anda. Penagihan diperlukan untuk semua penggunaan Google Cloud .
  4. Instal komponen gcloud alpha.

  5. Jalankan perintah berikut untuk menginstal versi terbaru komponen gcloud.

    gcloud components update

  6. Aktifkan TPU API melalui perintah gcloud berikut menggunakan Cloud Shell. Anda juga dapat mengaktifkannya dari konsol .

    gcloud services enable tpu.googleapis.com

  7. Buat identitas layanan untuk VM TPU.

    gcloud alpha compute tpus tpu-vm service-identity create --zone=ZONE

  8. Buat akun layanan TPU dan berikan akses ke layanan Google Cloud .

    Akun layanan memungkinkan layanan TPU Google Cloud untuk mengakses layanan Google Cloud lainnya. Akun layanan yang dikelola pengguna direkomendasikan. Ikuti panduan ini untuk membuat dan memberikan peran. Peran berikut diperlukan:

    • TPU Admin: Diperlukan untuk membuat TPU
    • Storage Admin: Diperlukan untuk mengakses Cloud Storage
    • Logs Writer: Diperlukan untuk menulis log dengan Logging API
    • Monitoring Metric Writer: Diperlukan untuk menulis metrik ke Cloud Monitoring

  9. 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

Hubungi tim penjualan atau akun Cloud TPU Anda untuk meminta kuota TPU dan mengajukan pertanyaan tentang kapasitas.

Menyediakan lingkungan Cloud TPU

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

Prasyarat

  • Pastikan project Anda memiliki cukup kuota TPUS_PER_TPU_FAMILY, yang menentukan jumlah maksimum chip yang dapat Anda akses dalam projectGoogle Cloud .
  • Tutorial ini diuji dengan konfigurasi berikut:
    • Python 3.10 or later
    • Versi software harian:
      • JAX harian 0.4.32.dev20240912
      • LibTPU harian 0.1.dev20240912+nightly
    • Versi software stabil:
      • JAX + JAX Lib dari v0.4.35
  • 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 NODE_ID=TPU_NODE_ID # TPU name
export PROJECT_ID=PROJECT_ID
export ACCELERATOR_TYPE=v6e-4
export ZONE=us-central2-b
export RUNTIME_VERSION=v2-alpha-tpuv6e
export SERVICE_ACCOUNT=YOUR_SERVICE_ACCOUNT
export QUEUED_RESOURCE_ID=QUEUED_RESOURCE_ID
export VALID_DURATION=VALID_DURATION

# Additional environment variable needed for Multislice:
export NUM_SLICES=NUM_SLICES

# Use a custom network for better performance as well as to avoid having the
# default network becoming overloaded.
export NETWORK_NAME=${PROJECT_ID}-mtu9k
export NETWORK_FW_NAME=${NETWORK_NAME}-fw

Deskripsi flag perintah

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

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

NUM_SLICES Jumlah slice yang akan dibuat (hanya diperlukan untuk Multislice)
QUEUED_RESOURCE_ID ID teks yang ditetapkan pengguna untuk permintaan resource yang diantrekan.
VALID_DURATION Durasi validitas permintaan resource yang diantrekan.
NETWORK_NAME Nama jaringan sekunder yang akan digunakan.
NETWORK_FW_NAME Nama firewall jaringan sekunder yang akan digunakan.

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

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

Setelah terhubung ke TPU, Anda dapat menjalankan benchmark inferensi.

Menjalankan benchmark inferensi Llama2-7B

Untuk menyiapkan JetStream dan MaxText, mengonversi titik pemeriksaan 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