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. Ikuti panduan Menyiapkan lingkungan Cloud TPU untuk memastikan Anda memiliki akses yang sesuai untuk menggunakan Cloud TPU.

2. Buat identitas layanan untuk VM TPU.

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

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

   Akun layanan memungkinkan layanan Google Cloud TPU mengakses layanan Google Cloudlainnya. Akun layanan yang dikelola pengguna direkomendasikan. Anda dapat membuat akun layanan dari Konsol Google Cloud atau melalui perintah gcloud.

   Buat akun layanan menggunakan alat command line gcloud:

   gcloud iam service-accounts create your-service-account-name \
   --description="your-sa-description" \
   --display-name="your-sa-display-name"
   export SERVICE_ACCOUNT_NAME=your-service-account-name

   Buat akun layanan dari konsol Google Cloud:

   1. Buka halaman Akun Layanan di konsol Google Cloud.
   2. Klik Create service account.
   3. Masukkan nama akun layanan.
   4. (Opsional) Masukkan deskripsi akun layanan.
   5. Klik Buat dan lanjutkan.
   6. Pilih peran yang ingin Anda berikan ke akun layanan.
   7. Klik Lanjutkan.
   8. (Opsional) Tentukan pengguna atau grup yang dapat mengelola akun layanan.
   9. Klik Selesai untuk menyelesaikan pembuatan akun layanan.

   Setelah membuat akun layanan, ikuti langkah-langkah berikut untuk memberikan peran akun layanan.

   Peran berikut diperlukan:

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

   Administrator harus memberi Anda roles/resourcemanager.projectIamAdmin agar Anda dapat menetapkan peran IAM kepada pengguna. Pengguna dengan peran Project IAM Admin roles/resourcemanager.projectIamAdmin juga dapat memberikan peran ini.

   Gunakan perintah gcloud berikut untuk menambahkan peran akun layanan:

   gcloud projects add-iam-policy-binding ${PROJECT_ID} \
      --member serviceAccount:${SERVICE_ACCOUNT_NAME}@${PROJECT_ID}.iam.gserviceaccount.com \
      --role roles/tpu.admin
   gcloud projects add-iam-policy-binding ${PROJECT_ID} \
      --member serviceAccount:${SERVICE_ACCOUNT_NAME}@${PROJECT_ID}.iam.gserviceaccount.com \
      --role roles/storage.admin
   gcloud projects add-iam-policy-binding ${PROJECT_ID} \
      --member serviceAccount:${SERVICE_ACCOUNT_NAME}@${PROJECT_ID}.iam.gserviceaccount.com \
      --role roles/logging.logWriter
   gcloud projects add-iam-policy-binding ${PROJECT_ID} \
      --member serviceAccount:${SERVICE_ACCOUNT_NAME}@${PROJECT_ID}.iam.gserviceaccount.com \
      --role roles/monitoring.metricWriter

   Anda juga dapat menetapkan peran menggunakan konsol Google Cloud.

   Dari konsol Google Cloud, pilih peran berikut:

   1. Pilih akun layanan Anda, lalu klik Tambahkan Akun Utama.
   2. Di kolom New Principals, masukkan alamat email akun layanan Anda.
   3. Di drop-down Select a role, telusuri peran (misalnya, Storage Admin) dan pilih peran tersebut.
   4. Klik Simpan.

4. 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 sudah siap untuk mendapatkan kapasitas TPU, tinjau halaman kuota untuk mempelajari sistem Kuota Cloud. Jika ada pertanyaan tambahan tentang cara mendapatkan kapasitas, hubungi tim penjualan atau akun Cloud TPU Anda.

Menyediakan lingkungan Cloud TPU

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

Prasyarat

• Tutorial ini telah diuji dengan Python 3.10 atau yang lebih baru.
• 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
  • Jika Anda menggunakan GKE dengan XPK, lihat Izin Konsol Cloud di akun pengguna atau layanan untuk mengetahui izin yang diperlukan untuk menjalankan XPK.

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

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