Menyajikan LLM menggunakan TPU di GKE dengan JetStream dan PyTorch

Panduan ini menunjukkan cara menayangkan model bahasa besar (LLM) menggunakan Unit Pemrosesan Tensor (TPU) di Google Kubernetes Engine (GKE) dengan JetStream melalui PyTorch. Dalam panduan ini, Anda akan mendownload bobot model ke Cloud Storage dan men-deploy-nya di cluster GKE Autopilot atau Standard menggunakan container yang menjalankan JetStream.

Jika Anda memerlukan skalabilitas, ketahanan, dan efektivitas biaya yang ditawarkan oleh fitur Kubernetes saat men-deploy model di JetStream, panduan ini adalah titik awal yang baik.

Panduan ini ditujukan untuk pelanggan AI Generatif yang menggunakan PyTorch, pengguna baru atau lama GKE, Engineer ML, engineer MLOps (DevOps), atau administrator platform yang tertarik menggunakan kemampuan orkestrasi container Kubernetes untuk menayangkan LLM.

Latar belakang

Dengan menyajikan LLM menggunakan TPU di GKE dengan JetStream, Anda dapat membangun solusi penyajian yang tangguh dan siap produksi dengan semua manfaat Kubernetes terkelola, termasuk efisiensi biaya, skalabilitas, dan ketersediaan yang lebih tinggi. Bagian ini menjelaskan teknologi utama yang digunakan dalam tutorial ini.

Tentang TPU

TPU adalah sirkuit terintegrasi khusus aplikasi (ASIC) yang dikembangkan khusus oleh Google dan digunakan untuk mempercepat model machine learning dan AI yang dibuat menggunakan framework seperti TensorFlow, PyTorch, dan JAX.

Sebelum menggunakan TPU di GKE, sebaiknya selesaikan jalur pembelajaran berikut:

  1. Pelajari ketersediaan versi TPU saat ini dengan arsitektur sistem Cloud TPU.
  2. Pelajari TPU di GKE.

Tutorial ini membahas cara menayangkan berbagai model LLM. GKE men-deploy model pada node TPUv5e host tunggal dengan topologi TPU yang dikonfigurasi berdasarkan persyaratan model untuk menyajikan perintah dengan latensi rendah.

Tentang JetStream

JetStream adalah framework penayangan inferensi open source yang dikembangkan oleh Google. JetStream memungkinkan inferensi berperforma tinggi, throughput tinggi, dan dioptimalkan untuk memori di TPU dan GPU. JetStream menyediakan pengoptimalan performa lanjutan, termasuk pengelompokan berkelanjutan, pengoptimalan cache KV, dan teknik kuantisasi, untuk memfasilitasi deployment LLM. JetStream memungkinkan penayangan TPU PyTorch/XLA dan JAX untuk mencapai performa yang optimal.

Pengelompokan Berkelanjutan

Pengelompokan berkelanjutan adalah teknik yang mengelompokkan permintaan inferensi masuk secara dinamis ke dalam batch, sehingga mengurangi latensi dan meningkatkan throughput.

Kuantisasi cache KV

Kuantisasi cache KV melibatkan kompresi cache key-value yang digunakan dalam mekanisme perhatian, sehingga mengurangi persyaratan memori.

Kuantisasi bobot Int8

Kuantisasi bobot Int8 mengurangi presisi bobot model dari floating point 32-bit menjadi bilangan bulat 8-bit, sehingga menghasilkan komputasi yang lebih cepat dan penggunaan memori yang lebih rendah.

Untuk mempelajari lebih lanjut pengoptimalan ini, lihat repositori project JetStream PyTorch dan JetStream MaxText.

Tentang PyTorch

PyTorch adalah framework machine learning open source yang dikembangkan oleh Meta dan kini menjadi bagian dari Linux Foundation. PyTorch menyediakan fitur tingkat tinggi seperti komputasi tensor dan jaringan neural dalam.