Menyajikan model terbuka Gemma menggunakan TPU di Gemini Enterprise Agent Platform dengan Saxml

Panduan ini menunjukkan cara menyajikan model bahasa besar (LLM) model terbuka Gemma menggunakan Tensor Processing Unit (TPU) di Gemini Enterprise Agent Platform dengan Saxml. Dalam panduan ini, Anda akan mendownload model Gemma yang di-tuning dengan petunjuk parameter 2B dan 7B ke Cloud Storage dan men-deploy-nya di Gemini Enterprise Agent Platform yang menjalankan Saxml di TPU.

Latar belakang

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

Gemma

Gemma adalah serangkaian model AI generatif yang ringan dan tersedia secara terbuka yang dirilis dengan lisensi terbuka. Model AI ini tersedia untuk dijalankan 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 menyajikan model Paxml, JAX, dan PyTorch untuk inferensi. Untuk tutorial ini, kita akan membahas cara menyajikan Gemma di TPU yang lebih hemat biaya untuk Saxml. Penyiapan untuk GPU serupa. Saxml menawarkan skrip untuk membangun container Gemini Enterprise Agent Platform yang akan kita gunakan dalam tutorial ini.

TPU

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

Tutorial ini menyajikan model Gemma 2B dan Gemma 7B. Gemini Enterprise Agent Platform menghosting model ini di node pool TPU v5e host tunggal berikut:

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

Sebelum memulai

  1. In the Cloud de Confiance console, on the project selector page, select or create a Cloud de Confiance project.

    Roles required to select or create a project

    • Select a project: Selecting a project doesn't require a specific IAM role—you can select any project that you've been granted a role on.
    • Create a project: To create a project, you need the Project Creator role (roles/resourcemanager.projectCreator), which contains the resourcemanager.projects.create permission. Learn how to grant roles.

    Go to project selector

  2. Verify that billing is enabled for your Cloud de Confiance project.

  3. Enable the Agent Platform API, if it is not already enabled.

    Roles required to enable APIs

    To enable APIs, you need the serviceusage.services.enable permission. If you created the project, then you likely already have this permission through the Owner role (roles/owner). Otherwise, you can get this permission through the Service Usage Admin role (roles/serviceusage.serviceUsageAdmin). Learn how to grant roles.

    Enable the API

  4. Di konsol Cloud de Confiance , aktifkan Cloud Shell.

    Aktifkan Cloud Shell

    Di bagian bawah konsol Cloud de Confiance , sesi Cloud Shell akan dimulai 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 pada sesi.

Tutorial ini mengasumsikan bahwa Anda menggunakan Cloud Shell untuk berinteraksi dengan Cloud de Confiance by S3NS. Jika Anda ingin menggunakan shell yang berbeda, bukan Cloud Shell, lakukan konfigurasi tambahan berikut:

  1. Instal Google Cloud CLI.

  2. Konfigurasi gcloud CLI agar menggunakan identitas gabungan Anda.

    Untuk mengetahui informasi selengkapnya, lihat Login ke gcloud CLI dengan identitas gabungan Anda.

  3. Untuk melakukan inisialisasi gcloud CLI, jalankan perintah berikut:

    gcloud init
  4. Pastikan Anda memiliki kuota yang cukup untuk chip TPU v5e untuk Gemini Enterprise Agent Platform. Secara default, kuota ini adalah 0. Untuk topologi 1x1, nilainya harus 1. Untuk 2x2, harus 4. Untuk menjalankan kedua topologi, nilainya harus 5.
  5. Buat akun Kaggle, jika Anda belum memilikinya.

Mendapatkan akses ke model

Perhatikan bahwa Cloud Shell mungkin tidak memiliki resource yang memadai untuk mendownload bobot model. Jika ya, Anda dapat membuat instance Gemini Enterprise Agent Platform Workbench untuk melakukan tugas tersebut.

Untuk mendapatkan akses ke model Gemma untuk deployment ke Gemini Enterprise Agent Platform, Anda harus login ke platform Kaggle, menandatangani perjanjian izin lisensi, dan mendapatkan token API Kaggle. Dalam tutorial ini, Anda akan menggunakan Secret Kubernetes untuk kredensial Kaggle.

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

  1. Akses halaman izin model di Kaggle.com.
  2. Login ke Kaggle jika Anda belum melakukannya.
  3. Klik Minta Akses.
  4. Di bagian Pilih Akun untuk Izin, pilih Verifikasi melalui Akun Kaggle untuk menggunakan akun Kaggle Anda untuk izin.
  5. Setujui Persyaratan dan Ketentuan model.

Membuat token akses

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

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

  1. Di browser Anda, buka setelan Kaggle.
  2. Di bagian API, klik Create New Token.

    File bernama kaggle.json akan didownload.

Mengupload token akses ke Cloud Shell

Di Cloud Shell, Anda dapat mengupload token Kaggle API ke project Cloud de Confiance by S3NS Anda:

  1. Di Cloud Shell, klik Lainnya > Upload.
  2. Pilih File, lalu klik Pilih File.
  3. Buka file kaggle.json.
  4. Klik Upload.

Buat bucket Cloud Storage

Buat bucket Cloud Storage untuk menyimpan checkpoint model.

Di Cloud Shell, jalankan perintah berikut:

gcloud storage buckets create gs://CHECKPOINTS_BUCKET_NAME

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

Menyalin model ke bucket Cloud Storage

Di Cloud Shell, jalankan perintah berikut:

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
gcloud storage cp /data/gemma_2b-it/* gs://CHECKPOINTS_BUCKET_NAME/gemma_2b-it/ --recursive

# 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
gcloud storage cp /data/gemma_7b-it/* gs://CHECKPOINTS_BUCKET_NAME/gemma_7b-it/ --recursive

Men-deploy model

Mengupload model

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

Gemma 2B-it

gcloud ai models upload \
  --region=LOCATION \
  --display-name=DEPLOYED_MODEL_NAME \
  --container-image-uri=us-docker.pkg.dev/vertex-ai/prediction/sax-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=us-docker.pkg.dev/vertex-ai/prediction/sax-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 Cloud de Confianceproject Anda
  • LOCATION_ID: Region tempat Anda menggunakan Agent Platform. 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 tersebut dapat digunakan untuk menyajikan inferensi 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 Agent Platform.
  • 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 Agent Platform.
  • 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 di 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 inferensi online dari model yang di-deploy

Untuk memanggil model melalui endpoint Gemini Enterprise Agent Platform, format permintaan inferensi 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 Agent Platform.
  • ENDPOINT_NAME: Nama tampilan endpoint.
  • instances.json memiliki format berikut: {"instances": [{"text_batch": "<your prompt>"},{...}]}

Pembersihan

Agar tidak menimbulkan biaya Gemini Enterprise Agent Platform lebih lanjut dan biaya Artifact Registry, hapus Cloud de Confiance resource 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.

Batasan

  • Di Gemini Enterprise Agent Platform, TPU Cloud hanya didukung di us-west1. Untuk mengetahui informasi selengkapnya, lihat lokasi.

Langkah berikutnya

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