Menyajikan model terbuka Gemma menggunakan TPU di Vertex AI Prediction dengan Saxml

Pratinjau

Fitur ini tunduk pada "Persyaratan Penawaran Pra-GA" di bagian Persyaratan Layanan Umum dalam Persyaratan Khusus Layanan. Fitur pra-GA tersedia "sebagaimana adanya" dan mungkin memiliki dukungan terbatas. Untuk mengetahui informasi selengkapnya, lihat deskripsi tahap peluncuran.

Panduan ini menunjukkan cara menyalurkan model bahasa besar (LLM) model terbuka Gemma menggunakan Tensor Processing Unit (TPU) di Vertex AI Prediction dengan Saxml. Dalam panduan ini, Anda akan mendownload petunjuk parameter 2B dan 7B yang menyesuaikan model Gemma ke Cloud Storage dan men-deploy-nya di Vertex AI Prediction yang menjalankan Saxml di TPU.

Latar belakang

Dengan menyajikan Gemma menggunakan TPU di Vertex AI Prediction dengan Saxml. Anda dapat memanfaatkan solusi AI terkelola yang menangani infrastruktur tingkat rendah dan menawarkan cara hemat biaya untuk melayani LLM. Bagian ini menjelaskan teknologi utama yang digunakan dalam tutorial ini.

Batu Permata

Gemma adalah serangkaian model kecerdasan buatan (AI) yang tersedia secara terbuka, ringan, dan generatif dan dirilis dengan lisensi terbuka. Model AI ini tersedia untuk berjalan di aplikasi, hardware, perangkat seluler, atau layanan yang dihosting. Anda dapat menggunakan model Gemma untuk pembuatan teks, tetapi Anda juga dapat menyesuaikan model ini untuk tugas khusus.

Untuk mempelajari lebih lanjut, lihat dokumentasi Gemma.

Saxml

Saxml adalah sistem eksperimental yang menyalurkan model Paxml, JAX, dan PyTorch untuk inferensi. Untuk tutorial ini, kami akan membahas cara menyajikan Gemma di TPU yang lebih hemat biaya untuk Saxml. Penyiapan untuk GPU serupa. Saxml menawarkan skrip untuk mem-build container untuk Vertex AI Prediction yang akan kita gunakan dalam tutorial ini.

TPU

TPU adalah sirkuit terintegrasi khusus aplikasi (ASIC) Google yang dikembangkan secara kustom dan digunakan untuk mempercepat framework pemrosesan data seperti TensorFlow, PyTorch, dan JAX.

Tutorial ini menampilkan model Gemma 2B dan Gemma 7B. Vertex AI Prediction menghosting model ini di kumpulan node TPU v5e host tunggal berikut:

• Gemma 2B: Dihosting di kumpulan node TPU v5e dengan topologi 1x1 yang mewakili satu chip TPU. Jenis mesin untuk node adalah ct5lp-hightpu-1t.
• Gemma 7B: Dihosting di kumpulan node TPU v5e dengan topologi 2x2 yang mewakili empat chip TPU. Jenis mesin untuk node adalah ct5lp-hightpu-4t.

Sebelum memulai

1. Login ke akun Google Cloud Anda. Jika Anda baru menggunakan Google Cloud, buat akun untuk mengevaluasi performa produk kami dalam skenario dunia nyata. Pelanggan baru juga mendapatkan kredit gratis senilai $300 untuk menjalankan, menguji, dan men-deploy workload.

2. Di konsol Google Cloud, pada halaman pemilih project, pilih atau buat project Google Cloud.

   Catatan: Jika tidak berencana untuk menyimpan resource yang Anda buat dalam prosedur ini, buatlah project, bukan memilih project yang ada. Setelah menyelesaikan langkah ini, Anda dapat menghapus project, yang menghapus semua resource yang terkait dengan project ini.

   Buka pemilih project

3. Pastikan penagihan telah diaktifkan untuk project Google Cloud Anda.

4. Aktifkan API Vertex AI API and Artifact Registry API.

   Mengaktifkan API

5. Di konsol Google Cloud, pada halaman pemilih project, pilih atau buat project Google Cloud.

   Catatan: Jika tidak berencana untuk menyimpan resource yang Anda buat dalam prosedur ini, buatlah project, bukan memilih project yang ada. Setelah menyelesaikan langkah ini, Anda dapat menghapus project, yang menghapus semua resource yang terkait dengan project ini.

   Buka pemilih project

6. Pastikan penagihan telah diaktifkan untuk project Google Cloud Anda.

7. Aktifkan API Vertex AI API and Artifact Registry API.

   Mengaktifkan API

8. Di konsol Google Cloud, aktifkan Cloud Shell.

   Aktifkan Cloud Shell

   Di bagian bawah Google Cloud Console, Cloud Shell sesi akan terbuka dan menampilkan perintah command line. Cloud Shell adalah lingkungan shell dengan Google Cloud CLI yang sudah terinstal, dan dengan nilai yang sudah ditetapkan untuk project Anda saat ini. Diperlukan waktu beberapa detik untuk melakukan inisialisasi sesi.

Tutorial ini mengasumsikan bahwa Anda menggunakan Cloud Shell untuk berinteraksi dengan Google Cloud. Jika Anda ingin menggunakan shell lain, bukan Cloud Shell, lakukan konfigurasi tambahan berikut:

1. Menginstal Google Cloud CLI.

2. Untuk initialize gcloud CLI, jalankan perintah berikut:

   gcloud init

3. Ikuti dokumentasi Artifact Registry untuk Menginstal Docker.
4. Pastikan Anda memiliki kuota yang cukup untuk 5 chip TPU v5e untuk Vertex AI Prediction.
5. Buat akun Kaggle, jika Anda belum memilikinya.

Mendapatkan akses ke model

Guna mendapatkan akses ke model Gemma untuk deployment ke Vertex AI Prediction, Anda harus login ke platform Kaggle, menandatangani perjanjian persetujuan lisensi, dan mendapatkan token API Kaggle. Dalam tutorial ini, Anda akan menggunakan Secret Kubernetes untuk kredensial Kaggle.

Tanda tangani perjanjian izin lisensi

Anda harus menandatangani perjanjian izin untuk menggunakan Gemma. Ikuti petunjuk berikut:

1. Akses halaman izin model di Kaggle.com.
2. Masuk ke Kaggle jika Anda belum melakukannya.
3. Klik Minta Akses.
4. Di bagian Choose Account for Consent, pilih Verify via Kaggle Account untuk menggunakan akun Kaggle Anda untuk mendapatkan izin.
5. Setujui Persyaratan dan Ketentuan model.

Membuat token akses

Untuk mengakses model melalui Kaggle, Anda memerlukan token API Kaggle.

Ikuti langkah-langkah berikut untuk membuat token baru jika Anda belum memilikinya:

1. Di browser, buka setelan Kaggle.

2. Di bagian API, klik Create New Token.

   File bernama kaggle.json telah didownload.

Mengupload token akses ke Cloud Shell

Di Cloud Shell, Anda dapat mengupload token API Kaggle ke project Google Cloud:

1. Di Cloud Shell, klik More > Upload.
2. Pilih File dan klik Choose Files.
3. Buka file kaggle.json.
4. Klik Upload.

Membuat bucket Cloud Storage

Membuat bucket Cloud Storage untuk menyimpan checkpoint model.

Jalankan perintah berikut di Cloud Shell:

gcloud storage buckets create gs://CHECKPOINTS_BUCKET_NAME

Ganti CHECKPOINTS_BUCKET_NAME dengan nama bucket Cloud Storage yang menyimpan checkpoint model.

Salin model ke bucket Cloud Storage

Jalankan perintah berikut di Cloud Shell:

pip install kaggle --break-system-packages

# For Gemma 2B
mkdir -p /data/gemma_2b-it
kaggle models instances versions download google/gemma/pax/2b-it/1 --untar -p /data/gemma_2b-it
gsutil -m cp -R /data/gemma_2b-it/* gs://CHECKPOINTS_BUCKET_NAME/gemma_2b-it/

# For Gemma 7B
mkdir -p /data/gemma_7b-it
kaggle models instances versions download google/gemma/pax/7b-it/1 --untar -p /data/gemma_7b-it
gsutil -m cp -R /data/gemma_7b-it/* gs://CHECKPOINTS_BUCKET_NAME/gemma_7b-it/

Membuat repositori Artifact Registry

Buat repositori Artifact Registry untuk menyimpan image container yang akan Anda buat di bagian berikutnya.

Aktifkan layanan Artifact Registry API untuk project Anda.

gcloud services enable artifactregistry.googleapis.com

Jalankan perintah berikut di shell Anda untuk membuat repositori Artifact Registry:

gcloud artifacts repositories create saxml \
 --repository-format=docker \
 --location=LOCATION \
 --description="Saxml Docker repository"

Ganti LOCATION dengan region tempat Artifact Registry menyimpan image container Anda. Kemudian, Anda harus membuat resource model Vertex AI di endpoint regional yang cocok dengan region ini, jadi pilih region tempat Vertex AI memiliki endpoint regional, seperti us-west1 untuk TPU.

Mengirim container ke Artifact Registry

Container Saxml bawaan tersedia di us-docker.pkg.dev/vertex-ai/prediction/sax-tpu:latest. Salin ke Artifact Registry Anda. Mengonfigurasi Docker untuk mengakses Artifact Registry. Kemudian, kirim image container Anda ke repositori Artifact Registry.

1. Untuk memberikan izin penginstalan Docker lokal kepada Artifact Registry di region yang dipilih, jalankan perintah berikut di shell Anda:

   gcloud auth configure-docker LOCATION-docker.pkg.dev
   

   • Ganti LOCATION dengan region tempat Anda membuat repositori.

2. Untuk menyalin image container yang baru Anda buat ke Artifact Registry, jalankan perintah berikut di shell Anda:

   docker tag us-docker.pkg.dev/vertex-ai/prediction/sax-tpu:latest LOCATION-docker.pkg.dev/PROJECT_ID/saxml/saxml-tpu:latest
   

3. Untuk mengirim image container yang baru saja Anda buat ke Artifact Registry, jalankan perintah berikut di shell Anda:

   docker push LOCATION-docker.pkg.dev/PROJECT_ID/saxml/saxml-tpu:latest
   

   Ganti baris berikut, seperti yang Anda lakukan di bagian sebelumnya:

   • LOCATION: region repositori Artifact Registry Anda.
   • PROJECT_ID: ID project Google Cloud Anda

Men-deploy model

Mengupload model

Untuk mengupload resource Model yang menggunakan container Saxml, jalankan perintah gcloud ai models upload berikut:

Gemma 2B-it

gcloud ai models upload \
  --region=LOCATION \
  --display-name=DEPLOYED_MODEL_NAME \
  --container-image-uri=LOCATION-docker.pkg.dev/PROJECT_ID/saxml/saxml-tpu:latest \
  --artifact-uri='gs://CHECKPOINTS_BUCKET_NAME/gemma_2b-it/' \
  --container-args='--model_path=saxml.server.pax.lm.params.gemma.Gemma2BFP16' \
  --container-args='--platform_chip=tpuv5e' \
  --container-args='--platform_topology=2x2' \
  --container-args='--ckpt_path_suffix=checkpoint_00000000' \
  --container-ports=8502

Gemma 7B-it

gcloud ai models upload \
  --region=LOCATION \
  --display-name=DEPLOYED_MODEL_NAME \
  --container-image-uri=LOCATION-docker.pkg.dev/PROJECT_ID/saxml/saxml-tpu:latest \
  --artifact-uri='gs://CHECKPOINTS_BUCKET_NAME/gemma_7b-it/' \
  --container-args='--model_path=saxml.server.pax.lm.params.gemma.Gemma7BFP16' \
  --container-args='--platform_chip=tpuv5e' \
  --container-args='--platform_topology=2x2' \
  --container-args='--ckpt_path_suffix=checkpoint_00000000' \
  --container-ports=8502

Ganti kode berikut:

• PROJECT_ID: ID project Google Cloud Anda
• LOCATION_ID: Region tempat Anda menggunakan Vertex AI. Perhatikan bahwa TPU hanya tersedia di us-west1.
• DEPLOYED_MODEL_NAME: Nama untuk DeployedModel. Anda juga dapat menggunakan nama tampilan Model untuk DeployedModel.

Membuat endpoint

Anda harus men-deploy model ke endpoint sebelum model ini dapat digunakan untuk menayangkan prediksi online. Jika Anda men-deploy model ke endpoint yang sudah ada, Anda dapat melewati langkah ini. Contoh berikut menggunakan perintah gcloud ai endpoints create:

gcloud ai endpoints create \
  --region=LOCATION \
  --display-name=ENDPOINT_NAME

Ganti kode berikut:

• LOCATION_ID: Region tempat Anda menggunakan Vertex AI.
• ENDPOINT_NAME: Nama tampilan endpoint.

Alat Google Cloud CLI mungkin memerlukan waktu beberapa detik untuk membuat endpoint.

Men-deploy model ke endpoint

Setelah endpoint siap, deploy model ke endpoint.

ENDPOINT_ID=$(gcloud ai endpoints list \
   --region=LOCATION \
   --filter=display_name=ENDPOINT_NAME \
   --format="value(name)")

MODEL_ID=$(gcloud ai models list \
   --region=LOCATION \
   --filter=display_name=DEPLOYED_MODEL_NAME \
   --format="value(name)")

gcloud ai endpoints deploy-model $ENDPOINT_ID \
  --region=LOCATION \
  --model=$MODEL_ID \
  --display-name=DEPLOYED_MODEL_NAME \
  --machine-type=ct5lp-hightpu-4t \
  --traffic-split=0=100

Ganti kode berikut:

• LOCATION_ID: Region tempat Anda menggunakan Vertex AI.
• ENDPOINT_NAME: Nama tampilan endpoint.
• DEPLOYED_MODEL_NAME: Nama untuk DeployedModel. Anda juga dapat menggunakan nama tampilan Model untuk DeployedModel.

Gemma 2B dapat di-deploy pada mesin ct5lp-hightpu-1t yang lebih kecil. Dalam hal ini, Anda harus menentukan --platform_topology=1x1 saat mengupload model.

Alat Google Cloud CLI mungkin memerlukan waktu beberapa menit untuk men-deploy model ke endpoint. Setelah model berhasil di-deploy, perintah ini akan mencetak output berikut:

  Deployed a model to the endpoint xxxxx. Id of the deployed model: xxxxx.

Mendapatkan prediksi online dari model yang di-deploy

Untuk memanggil model melalui endpoint Prediksi Vertex AI, format permintaan prediksi menggunakan Objek JSON Permintaan Inferensi standar .

Contoh berikut menggunakan perintah gcloud ai endpoints predict:

ENDPOINT_ID=$(gcloud ai endpoints list \
   --region=LOCATION \
   --filter=display_name=ENDPOINT_NAME \
   --format="value(name)")

gcloud ai endpoints predict $ENDPOINT_ID \
  --region=LOCATION \
  --http-headers=Content-Type=application/json \
  --json-request instances.json

Ganti kode berikut:

• LOCATION_ID: Region tempat Anda menggunakan Vertex AI.
• ENDPOINT_NAME: Nama tampilan endpoint.
• instances.json memiliki format berikut: {"instances": [{"text_batch": "<your prompt>"},{...}]}

Pembersihan

Agar tidak menimbulkan biaya Vertex AI dan biaya Artifact Registry lebih lanjut, hapus resource Google Cloud yang Anda buat selama tutorial ini:

1. Untuk membatalkan deployment model dari endpoint dan menghapus endpoint, jalankan perintah berikut di shell Anda:

   ENDPOINT_ID=$(gcloud ai endpoints list \
      --region=LOCATION \
      --filter=display_name=ENDPOINT_NAME \
      --format="value(name)")
   
   DEPLOYED_MODEL_ID=$(gcloud ai endpoints describe $ENDPOINT_ID \
      --region=LOCATION \
      --format="value(deployedModels.id)")
   
   gcloud ai endpoints undeploy-model $ENDPOINT_ID \
     --region=LOCATION \
     --deployed-model-id=$DEPLOYED_MODEL_ID
   
   gcloud ai endpoints delete $ENDPOINT_ID \
      --region=LOCATION \
      --quiet
   

   Ganti LOCATION dengan region tempat Anda membuat model di bagian sebelumnya.

2. Untuk menghapus model, jalankan perintah berikut di shell Anda:

   MODEL_ID=$(gcloud ai models list \
      --region=LOCATION \
      --filter=display_name=DEPLOYED_MODEL_NAME \
      --format="value(name)")
   
   gcloud ai models delete $MODEL_ID \
      --region=LOCATION \
      --quiet
   

   Ganti LOCATION dengan region tempat Anda membuat model di bagian sebelumnya.

3. Untuk menghapus repositori Artifact Registry dan image container di dalamnya, jalankan perintah berikut di shell Anda:

   gcloud artifacts repositories delete saxml \
     --location=LOCATION \
     --quiet
   

   Ganti LOCATION dengan region tempat Anda membuat repositori Artifact Registry di bagian sebelumnya.

Batasan

• Di Vertex AI Prediction, Cloud TPU hanya didukung di us-west1. Untuk mengetahui informasi selengkapnya, lihat lokasi.

Langkah selanjutnya

• Pelajari cara men-deploy model SAxml lainnya seperti Llama2 dan GPT-J.