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 model dan tokenizer Llama 3 70B, lalu men-deploy-nya di Vertex AI yang menjalankan Saxml di TPU.

Sebelum memulai

Sebaiknya gunakan VM yang dioptimalkan untuk memori M2 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.

  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. In the Google Cloud console, on the project selector page, select or create a Google Cloud 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

  3. Verify that billing is enabled for your Google Cloud project.

  4. Enable the Vertex AI and Artifact Registry APIs.

    Roles required to enable APIs

    To enable APIs, you need the Service Usage Admin IAM role (roles/serviceusage.serviceUsageAdmin), which contains the serviceusage.services.enable permission. Learn how to grant roles.

    Enable the APIs

  5. In the Google Cloud console, on the project selector page, select or create a Google Cloud 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

  6. Verify that billing is enabled for your Google Cloud project.

  7. Enable the Vertex AI and Artifact Registry APIs.

    Roles required to enable APIs

    To enable APIs, you need the Service Usage Admin IAM role (roles/serviceusage.serviceUsageAdmin), which contains the serviceusage.services.enable permission. Learn how to grant roles.

    Enable the APIs

  8. Di konsol Google Cloud , aktifkan Cloud Shell.

    Aktifkan Cloud Shell

    Di bagian bawah konsol Google Cloud , 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.

  9. Ikuti dokumentasi Artifact Registry untuk Menginstal Docker.
  10. 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. Instal Google Cloud CLI.

  2. Jika Anda menggunakan penyedia identitas (IdP) eksternal, Anda harus login ke gcloud CLI dengan identitas gabungan Anda terlebih dahulu.

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

    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 dalam 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 sekarang 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 daftar lengkap jenis dan region Cloud TPU yang didukung, lihat Lokasi Vertex AI.

  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