Menjalankan inferensi di Cloud TPU

Penyajian adalah proses men-deploy model machine learning terlatih ke lingkungan produksi untuk menjalankan inferensi. TPU v5e dan versi yang lebih baru mendukung inferensi. Memenuhi tujuan tingkat layanan (SLO) latensi adalah prioritas untuk menjalankan workload.

Dokumen ini membahas penyajian model di TPU host tunggal. Slice TPU dengan delapan chip atau kurang memiliki satu VM atau host TPU dan disebut TPU host tunggal. Untuk mengetahui informasi tentang inferensi multi-host dengan PyTorch, lihat vLLM-TorchTPU dan SGLang-TorchTPU. Untuk mengetahui informasi tentang inferensi multi-host dengan JAX, lihat Melakukan inferensi multi-host menggunakan Pathways.

Menyiapkan TPU

Apa pun stack yang Anda pilih, Anda memerlukan project Google Cloud dengan TPU host tunggal sebelum dapat menyajikan model. Langkah-langkah berikut berlaku untuk setiap kombinasi library dan framework:

  1. Buat project. Anda memerlukan akun dan project Google Cloud untuk menggunakan Cloud TPU. Untuk mengetahui informasi selengkapnya, lihat Menyiapkan projectGoogle Cloud untuk TPU.

  2. Periksa kuota Anda. Pastikan Anda memiliki kuota yang cukup untuk jumlah core TPU yang ingin Anda gunakan untuk inferensi. TPU v5e menggunakan kuota terpisah untuk pelatihan dan inferensi. Di Cloud TPU API, kuota khusus inferensi untuk TPU v5e adalah sebagai berikut:

    • Resource v5e on-demand: TPU v5 lite pod cores for serving per project per zone
    • Resource v5e preemptible: Preemptible TPU v5 lite pod cores for serving per project per zone

    Untuk versi TPU lainnya, workload pelatihan dan inferensi menggunakan kuota yang sama. Untuk mengetahui informasi selengkapnya, lihat Kuota Cloud TPU.

  3. Membuat VM TPU. Pilih versi dan zona TPU yang memiliki kapasitas inferensi, lalu buat VM TPU host tunggal. Untuk mengetahui informasi selengkapnya, lihat Region dan zona TPU serta Mengelola resource TPU.

Setelah VM TPU berjalan, ikuti panduan penyiapan untuk library dan framework Anda di Memilih stack inferensi.

Memilih stack inferensi

Cloud TPU mendukung inferensi model bahasa besar (LLM) melalui dua library penyajian open source: vLLM dan SGLang. Kedua library berjalan di Cloud TPU dengan backend JAX atau PyTorch, sehingga Anda dapat memilih library yang paling sesuai dengan workload Anda, terlepas dari framework yang digunakan untuk menulis model Anda.

Pilih library terlebih dahulu, lalu pilih backend yang cocok dengan framework model Anda, seperti yang ditunjukkan dalam tabel berikut:

Library inferensi Backend JAX Backend PyTorch
vLLM vllm-jax vllm-torchtpu
SGLang sglang-jax sglang-torchtpu

Inferensi dengan vLLM

vLLM adalah library open source untuk inferensi dan penyajian LLM yang cepat, yang dibangun berdasarkan pengelompokan berkelanjutan dan PagedAttention. Cloud TPU berfungsi dengan vLLM melalui plugin vllm-jax untuk model JAX dan plugin vllm-torchtpu untuk model PyTorch.

Untuk contoh tambahan penggunaan vLLM untuk menyajikan model di TPU, lihat referensi berikut:

vLLM dengan JAX

Dokumentasi TPU vLLM membahas penginstalan, panduan memulai, model yang didukung, dan opsi penyesuaian. Untuk mengetahui informasi selengkapnya, lihat Menyajikan model JAX dengan vLLM di TPU.

vLLM dengan PyTorch

Plugin vllm-torchtpu menambahkan dukungan runtime dan kernel TorchTPU ke vLLM. Untuk mengetahui informasi selengkapnya, lihat Menyajikan model PyTorch dengan vLLM di TPU.

Inferensi dengan SGLang

SGLang adalah framework penyajian open source untuk model bahasa besar dan multimodal yang menggabungkan runtime berperforma tinggi dengan frontend pembuatan terstruktur. Di Cloud TPU, SGLang berjalan melalui backend SGLang-JAX (sglang-jax) untuk model JAX dan backend SGLang-TorchTPU (sglang-torchtpu) untuk model PyTorch.

SGLang dengan PyTorch

Untuk mengetahui informasi tentang penginstalan, peluncuran server, dan konfigurasi yang didukung, lihat Menayangkan model PyTorch dengan runtime SGLang-TorchTPU di TPU.

SGLang dengan JAX

Untuk mengetahui informasi tentang penginstalan, peluncuran server, dan konfigurasi yang didukung, lihat Menyajikan model JAX dengan runtime SGLang-JAX di TPU.

Membuat profil performa inferensi

Setelah menyiapkan inferensi, Anda dapat menggunakan profiler untuk menganalisis performa dan pemanfaatan TPU. Untuk mengetahui informasi selengkapnya tentang pembuatan profil, lihat referensi berikut: