Menyajikan model terbuka Llama 3 menggunakan Cloud TPU multi-host di Vertex AI dengan Saxml

Llama 3 adalah model bahasa besar (LLM) open source dari Meta. Panduan ini menunjukkan cara menyajikan LLM Llama 3 menggunakan Tensor Processing Unit (TPU) multi-host di Vertex AI dengan Saxml.

Dalam panduan ini, Anda akan mendownload bobot dan tokenizer model Llama 3 70B, lalu men-deploy-nya di Vertex AI yang menjalankan Saxml di TPU.

Sebelum memulai

Sebaiknya gunakan VM M2 yang dioptimalkan untuk memori untuk mendownload model dan mengonversinya ke Saxml. Hal ini karena proses konversi model memerlukan memori yang besar dan dapat gagal jika Anda memilih jenis mesin yang tidak memiliki memori yang cukup.

Sign in to your Google Cloud account. If you're new to Google Cloud, create an account to evaluate how our products perform in real-world scenarios. New customers also get $300 in free credits to run, test, and deploy workloads.

In the Google Cloud console, on the project selector page, select or create a Google Cloud project.

Go to project selector

Make sure that billing is enabled for your Google Cloud project.

Enable the Vertex AI and Artifact Registry APIs.

Enable the APIs

In the Google Cloud console, on the project selector page, select or create a Google Cloud project.

Go to project selector

Make sure that billing is enabled for your Google Cloud project.

Enable the Vertex AI and Artifact Registry APIs.

Enable the APIs

1. In the Google Cloud console, activate Cloud Shell.

   Activate Cloud Shell

   At the bottom of the Google Cloud console, a Cloud Shell session starts and displays a command-line prompt. Cloud Shell is a shell environment with the Google Cloud CLI already installed and with values already set for your current project. It can take a few seconds for the session to initialize.

2. Ikuti dokumentasi Artifact Registry untuk Menginstal Docker.
3. Pastikan Anda memiliki kuota yang cukup untuk 16 chip TPU v5e untuk Vertex AI.

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

1. Install the Google Cloud CLI.

2. If you're using an external identity provider (IdP), you must first sign in to the gcloud CLI with your federated identity.

3. To initialize the gcloud CLI, run the following command:

   gcloud init

Jika Anda menggunakan shell lain, bukan Cloud Shell, untuk deployment model, pastikan versi Google Cloud CLI lebih baru dari 475.0.0. Anda dapat mengupdate Google Cloud CLI dengan menjalankan perintah gcloud components update.

Jika Anda men-deploy model menggunakan Vertex AI SDK, pastikan Anda memiliki versi 1.50.0 atau yang lebih baru.

Dapatkan akses ke model dan download bobot model

Langkah-langkah berikut ditujukan untuk instance Vertex AI Workbench yang memiliki VM yang dioptimalkan untuk memori M2. Untuk mengetahui informasi tentang cara mengubah jenis mesin instance Vertex AI Workbench, lihat Mengubah jenis mesin instance Vertex AI Workbench.

1. Buka halaman izin model Llama.

2. Pilih Llama 3, isi formulir izin, dan setujui persyaratan dan ketentuan.

3. Periksa kotak masuk Anda untuk menemukan email yang berisi URL bertanda tangan.

4. Download skrip download.sh dari GitHub dengan menjalankan perintah berikut:

   wget https://raw.githubusercontent.com/meta-llama/llama3/main/download.sh
   chmod +x download.sh
   

5. Untuk mendownload bobot model, jalankan skrip download.sh yang Anda download dari GitHub.

6. Saat diminta, masukkan URL bertanda tangan dari email yang Anda terima di bagian sebelumnya.

7. Saat diminta memilih model yang akan didownload, masukkan 70B.

Mengonversi bobot model ke format Saxml

1. Jalankan perintah berikut untuk mendownload Saxml:

   git clone https://github.com/google/saxml.git
   

2. Jalankan perintah berikut untuk mengonfigurasi lingkungan virtual Python:

   python -m venv .
   source bin/activate
   

3. Jalankan perintah berikut untuk menginstal dependensi:

   pip install --upgrade pip
   
   pip install paxml
   
   pip install praxis
   
   pip install torch
   

4. Untuk mengonversi bobot model ke format Saxml, jalankan perintah berikut:

   python3 saxml/saxml/tools/convert_llama_ckpt.py \
       --base PATH_TO_META_LLAMA3 \
       --pax PATH_TO_PAX_LLAMA3 \
       --model-size llama3_70b
   

   Ganti kode berikut:

   • PATH_TO_META_LLAMA3: jalur ke direktori yang berisi bobot model yang didownload
   • PATH_TO_PAX_LLAMA3: jalur ke direktori tempat bobot model yang dikonversi akan disimpan

   Model yang dikonversi akan dimasukkan ke folder $PATH_TO_PAX_LLAMA3/checkpoint_00000000.

5. Salin file tokenizer dari direktori asli ke dalam subfolder bernama vocabs sebagai berikut:

   cp $PATH_TO_META_LLAMA3/tokenizer.model $PATH_TO_PAX_LLAMA3/vocabs/tokenizer.model
   

6. Tambahkan file commit_success.txt kosong di folder $PATH_TO_PAX_LLAMA3 dan subfolder metadata dan state di folder tersebut sebagai berikut:

   touch $PATH_TO_PAX_LLAMA3/checkpoint_00000000/commit_success.txt
   touch $PATH_TO_PAX_LLAMA3/checkpoint_00000000/metadata/commit_success.txt
   touch $PATH_TO_PAX_LLAMA3/checkpoint_00000000/state/commit_success.txt
   

   Folder $PATH_TO_PAX_LLAMA3 kini berisi folder dan file berikut:

   $PATH_TO_PAX_LLAMA3/checkpoint_00000000/commit_success.txt
   $PATH_TO_PAX_LLAMA3/checkpoint_00000000/metadata/
   $PATH_TO_PAX_LLAMA3/checkpoint_00000000/metadata/commit_success.txt
   $PATH_TO_PAX_LLAMA3/checkpoint_00000000/state/
   $PATH_TO_PAX_LLAMA3/checkpoint_00000000/state/commit_success.txt
   $PATH_TO_PAX_LLAMA3/vocabs/tokenizer.model
   

Membuat bucket Cloud Storage

Buat bucket Cloud Storage untuk menyimpan bobot model yang dikonversi.

1. Di Cloud Shell, jalankan perintah berikut, ganti PROJECT_ID dengan project ID Anda:

   projectid=PROJECT_ID
   gcloud config set project ${projectid}
   

2. Untuk membuat bucket, jalankan perintah berikut:

   gcloud storage buckets create gs://WEIGHTS_BUCKET_NAME
   

   Ganti WEIGHTS_BUCKET_NAME dengan nama yang ingin Anda gunakan untuk bucket.

Menyalin bobot model ke bucket Cloud Storage

Untuk menyalin bobot model ke bucket Anda, jalankan perintah berikut:

gcloud storage cp PATH_TO_PAX_LLAMA3/* gs://WEIGHTS_BUCKET_NAME/llama3/pax_70b/ --recursive

Mengupload model

Container Saxml bawaan tersedia di us-docker.pkg.dev/vertex-ai/prediction/sax-tpu:latest.

Untuk mengupload resource Model ke Vertex AI menggunakan container Saxml bawaan, jalankan perintah gcloud ai models upload sebagai berikut:

gcloud ai models upload \
    --region=LOCATION \
    --display-name=MODEL_DISPLAY_NAME \
    --container-image-uri=us-docker.pkg.dev/vertex-ai/prediction/sax-tpu:latest \
    --artifact-uri='gs://WEIGHTS_BUCKET_NAME/llama3/pax_70b' \
    --container-args='--model_path=saxml.server.pax.lm.params.lm_cloud.LLaMA3_70BFP16x16' \
    --container-args='--platform_chip=tpuv5e' \
    --container-args='--platform_topology=4x4' \
    --container-args='--ckpt_path_suffix=checkpoint_00000000' \
    --container-deployment-timeout-seconds=2700 \
    --container-ports=8502 \
    --project=PROJECT_ID

Lakukan penggantian berikut:

• LOCATION: region tempat Anda menggunakan Vertex AI. Perhatikan bahwa TPU hanya tersedia di us-west1.
• MODEL_DISPLAY_NAME: nama tampilan yang Anda inginkan untuk model Anda
• PROJECT_ID: ID Google Cloud project Anda

Membuat endpoint inferensi online

Untuk membuat endpoint, jalankan perintah berikut:

gcloud ai endpoints create \
    --region=LOCATION \
    --display-name=ENDPOINT_DISPLAY_NAME \
    --project=PROJECT_ID

Ganti ENDPOINT_DISPLAY_NAME dengan nama tampilan yang Anda inginkan untuk endpoint Anda.

Men-deploy model ke endpoint

Setelah endpoint siap, deploy model ke endpoint.

Dalam tutorial ini, Anda akan men-deploy model Llama 3 70B yang di-shard untuk 16 chip Cloud TPU v5e menggunakan topologi 4x4. Namun, Anda dapat menentukan salah satu topologi Cloud TPU multi-host yang didukung berikut:

Jenis Mesin	Topologi	Jumlah TPU chip	Jumlah Host
ct5lp-hightpu-4t	4x4	16	2
ct5lp-hightpu-4t	4x8	32	4
ct5lp-hightpu-4t	8x8	64	8
ct5lp-hightpu-4t	8x16	128	16
ct5lp-hightpu-4t	16x16	256	32

Jika Anda men-deploy model Llama lain yang ditentukan di repositori GitHub Saxml, pastikan model tersebut dipartisi agar sesuai dengan jumlah perangkat yang Anda targetkan dan Cloud TPU memiliki memori yang cukup untuk memuat model.

Untuk mengetahui informasi tentang men-deploy model di Cloud TPU host tunggal, lihat Men-deploy model.

Untuk mengetahui informasi selengkapnya tentang jenis Cloud TPU v5e, lihat TPU v5e.

1. Dapatkan ID endpoint untuk endpoint inferensi online:

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

2. Dapatkan ID model untuk model Anda:

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

3. Deploy model ke endpoint:

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

   Ganti DEPLOYED_MODEL_NAME dengan nama untuk yang di-deploy. Nama ini bisa sama dengan nama tampilan model (MODEL_DISPLAY_NAME).

   Operasi deployment mungkin kehabisan waktu.

   Perintah deploy-model menampilkan ID operasi yang dapat digunakan untuk memeriksa kapan operasi selesai. Anda dapat melakukan polling untuk status operasi hingga responsnya menyertakan "done": true. Gunakan perintah berikut untuk melakukan polling status:

   gcloud ai operations describe \
   --region=LOCATION \
   OPERATION_ID
   

   Ganti OPERATION_ID dengan ID operasi yang ditampilkan oleh perintah sebelumnya.

Mendapatkan inferensi online dari model yang di-deploy

Untuk mendapatkan inferensi online dari endpoint Vertex AI, jalankan perintah gcloud ai endpoints predict.

1. Jalankan perintah berikut untuk membuat file request.json yang berisi contoh permintaan inferensi:

   cat << EOF > request.json
   {"instances": [{"text_batch": "the distance between Earth and Moon is "}]}
   EOF
   

2. Untuk mengirim permintaan inferensi online ke endpoint, jalankan perintah berikut:

   gcloud ai endpoints predict $ENDPOINT_ID \
       --project=PROJECT_ID \
       --region=LOCATION \
       --json-request=request.json
   

Pembersihan

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

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

   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
   

2. Untuk menghapus model, jalankan perintah berikut:

   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