Menyajikan LLM menggunakan TPU multi-host di GKE dengan JetStream dan Pathways

Panduan ini menunjukkan cara menayangkan model bahasa besar (LLM) canggih seperti Llama 3.1 405B di Google Kubernetes Engine (GKE) menggunakan unit pemrosesan tensor (TPU) di beberapa node.

Panduan ini menunjukkan cara menggunakan teknologi open source portabel—Kubernetes, JetStream, Pathways on Cloud, dan API LeaderWorkerSet (LWS)—untuk men-deploy dan menayangkan workload AI/ML di GKE, dengan memanfaatkan kontrol terperinci, skalabilitas, ketahanan, portabilitas, dan efektivitas biaya GKE.

Latar belakang

Model bahasa besar telah bertambah ukurannya dan tidak lagi sesuai dengan slice TPU host tunggal. Untuk inferensi ML, Anda dapat menggunakan Pathways on Cloud untuk menjalankan inferensi multi-host berskala besar di GKE di beberapa node TPU yang saling terhubung. Dalam panduan ini, Anda akan mempelajari cara menyediakan cluster GKE dengan slice TPU multi-host dan menggunakan biner Pathways on Cloud, meluncurkan Server JetStream dengan framework MaxText, dan membuat permintaan inferensi multi-host.

Dengan menyajikan LLM menggunakan TPU di GKE dengan JetStream, MaxText, dan Pathways, Anda dapat membuat solusi penayangan 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 dibangun 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 menyajikan model Llama 3.1-405B. GKE men-deploy model pada node TPU v6e multi-host dengan topologi TPU yang dikonfigurasi berdasarkan persyaratan model untuk menyajikan perintah dengan latensi rendah.

Pathways on Cloud

Pathways adalah lapisan orkestrasi berskala besar untuk akselerator. Pathways dirancang secara eksplisit untuk memungkinkan eksplorasi sistem baru dan ide penelitian ML, sekaligus mempertahankan performa canggih untuk model saat ini. Pathways memungkinkan satu proses klien JAX mengoordinasikan komputasi di satu atau beberapa slice TPU besar, sehingga menyederhanakan komputasi ML yang mencakup ratusan atau ribuan chip TPU.

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 mengoptimalkan performa.

MaxText

MaxText adalah implementasi LLM JAX yang berperforma tinggi, skalabel, dan dapat disesuaikan, yang dibangun di atas library JAX open source seperti Flax, Orbax, dan Optax. Implementasi LLM khusus dekoder MaxText ditulis dalam Python. Hal ini memanfaatkan compiler XLA secara intensif untuk mencapai performa tinggi tanpa perlu membuat kernel kustom.

Untuk mengetahui informasi selengkapnya tentang model dan ukuran parameter terbaru yang didukung MaxText, lihat repositori project MaxText.

Llama 3.1 405B

Llama 3.1 405B adalah model bahasa besar dari Meta yang dirancang untuk berbagai tugas natural language processing, termasuk pembuatan teks, terjemahan, dan question answering. GKE menawarkan infrastruktur yang diperlukan untuk mendukung kebutuhan pelatihan dan penyajian terdistribusi model dalam skala ini.

Untuk mengetahui informasi selengkapnya, lihat dokumentasi Llama.

Arsitektur

Bagian ini menjelaskan arsitektur GKE yang digunakan dalam tutorial ini. Arsitektur ini mencakup cluster Standar GKE yang menyediakan TPU dan menghosting komponen JetStream dan Pathways untuk men-deploy dan menayangkan model.

Diagram berikut menunjukkan komponen arsitektur ini:

Arsitektur cluster GKE dengan node pool TPU multi-host yang berisi komponen JetStream dan Pathways.

Arsitektur ini mencakup komponen berikut:

  • Cluster regional GKE Standard.
  • Node pool slice TPU multi-host yang menghosting komponen JetStream deployment dan Pathways.
  • Pathways resource manager mengelola resource akselerator dan mengoordinasikan alokasi akselerator untuk tugas pengguna.
  • Pathways client berkoordinasi dengan Pathways resource manager untuk menentukan tempat program yang dikompilasi ditempatkan untuk dieksekusi.
  • Pathways worker berjalan dan melakukan komputasi di mesin akselerator, serta mengirimkan data kembali ke workload Anda melalui server proxy IFRT.
  • IFRT proxy client mengimplementasikan API Interim Framework Runtime (IFRT) OSS dan bertindak sebagai jembatan komunikasi antara workload dan komponen Pathways Anda.
  • IFRT proxy server menerima permintaan dari IFRT proxy client dan meneruskannya ke Pathways client, sehingga mendistribusikan tugas.
  • Penampung JetStream-Pathways menyediakan server inferensi berbasis JAX yang menerima permintaan inferensi dan mendelegasikan proses eksekusinya ke Pathways workers
  • Komponen Service menyebarkan traffic masuk ke semua replika JetStream HTTP.
  • JetStream HTTP adalah server HTTP yang menerima permintaan sebagai wrapper ke format yang diperlukan JetStream dan mengirimkannya ke klien GRPC JetStream.