Cara Men-scale Model Anda: Perspektif Sistem terhadap LLM di TPU

 (jax-ml.github.io)
7 poin oleh GN⁺ 2025-02-07 | 1 komentar | Bagikan ke WhatsApp
  • Mengoptimalkan performa deep learning dalam skala besar sering tampak seperti “alkimia”, tetapi pada praktiknya efisiensi model dapat ditingkatkan dengan prinsip-prinsip sederhana yang bisa dipahami
  • Dari satu akselerator hingga puluhan ribu akselerator, prinsip yang relatif sederhana berlaku di mana-mana, dan dengan memahaminya kita bisa melakukan hal-hal berguna seperti berikut:
    • Memperkirakan secara kasar seberapa dekat tiap bagian model dengan nilai optimal teoretis
    • Menyusun dasar untuk memilih berbagai teknik paralelisasi pada beragam skala
    • Memperkirakan biaya dan waktu yang diperlukan untuk melatih dan menjalankan model Transformer besar
    • Merancang algoritme yang memanfaatkan karakteristik perangkat keras tertentu
    • Merancang perangkat keras dengan memahami secara jelas batas performa algoritme saat ini
  • Pengetahuan latar yang dibutuhkan
    • Perlu memahami konsep dasar LLM dan arsitektur Transformer
    • Pemahaman tentang cara operasi berskala besar bukan keharusan
    • Akan lebih baik jika memiliki pengetahuan dasar pelatihan LLM dan pengalaman menggunakan JAX
    • Disarankan merujuk ke posting blog tentang arsitektur Transformer dan slide tentang scaling LLM di JAX
  • Tujuan
    • Mengembangkan kemampuan untuk memperkirakan cara terbaik memparalelkan model pada perangkat keras yang tersedia
    • Mengembangkan kemampuan untuk menghitung secara kasar waktu dan biaya pelatihan maupun inferensi

Mengapa ini penting

  • Bahkan 3~4 tahun lalu, sebagian besar peneliti ML tidak perlu memahami optimasi skala besar seperti ini secara mendalam
    • Kini, bahkan model yang “kecil” pun berjalan mendekati batas perangkat keras, sehingga memahami cara kerja efisien pada skala besar menjadi hal esensial
    • Sejarah ML dapat dilihat sebagai alur perkembangan silang antara inovasi sistem dan peningkatan perangkat lunak
    • Karena model Transformer belakangan ini menggunakan perangkat keras hingga batasnya, tanpa memahami efisiensi model, arsitektur atau riset baru berisiko gagal saat diterapkan di dunia nyata
    • Bahkan jika benchmark menunjukkan peningkatan performa 20%, bila efisiensi perangkat keras turun 20%, pada akhirnya kegunaannya menjadi rendah
  • Tujuan inti scaling model adalah membuat throughput meningkat secara linear ketika jumlah chip (akselerator) ditambah
    • Ini disebut "strong scaling"
    • Menambah chip mengurangi waktu komputasi tetapi menimbulkan biaya komunikasi antar-chip
    • Jika komunikasi memakan waktu lebih lama daripada komputasi, sistem menjadi berada dalam kondisi "communication bound" sehingga strong scaling tidak mungkin dicapai
    • Jika kita cukup memahami perangkat keras untuk memprediksi di mana bottleneck seperti ini akan muncul, kita bisa merancang atau menyusun ulang model untuk mencegahnya
  • Tujuan buku ini adalah menjelaskan cara kerja perangkat keras TPU (dan GPU) serta bagaimana arsitektur Transformer berkembang agar berjalan baik di perangkat keras saat ini
    • Diharapkan ini bermanfaat baik bagi peneliti yang merancang arsitektur baru maupun engineer yang berusaha menjalankan LLM generasi sekarang dengan cepat

Gambaran umum

  • Tulisan ini disusun sebagai berikut
  • Bagian 1 menjelaskan faktor-faktor yang menentukan batas performa model (komunikasi, komputasi, memori) melalui analisis roofline
  • Bagian 2, Bagian 3 membahas struktur internal TPU dan GPU serta cara koneksi antar-chip
    • Melalui ini, tulisan menjawab pertanyaan seperti berikut
      • Seberapa cepat perkalian matriks berukuran tertentu secara teoretis dapat dijalankan
      • Pada titik mana komputasi menjadi dibatasi oleh bandwidth memori atau bandwidth komunikasi
      • Dengan struktur seperti apa klaster TPU dihubungkan, dan kira-kira berapa lama waktu yang dibutuhkan untuk memindahkan data dari satu chip ke chip lain
      • Bagaimana mengalikan matriks terdistribusi secara efisien
  • Bagian 4 membahas secara rinci rumus-rumus arsitektur Transformer (ukuran matriks, jumlah parameter, FLOPs)
  • Bagian 5 dan Bagian 7 adalah bagian inti, yang memperkenalkan berbagai cara memparalelkan model ke banyak chip
    • Data parallel, Tensor parallel, Pipeline parallel, Expert parallel
    • Juga membahas teknik penghematan memori seperti ZeRO, Rematerialisation, Host offload, Gradient accumulation
  • Bagian 6, Bagian 8 menggunakan contoh pelatihan dan inferensi model LLaMA-3 di TPU untuk menunjukkan biaya, waktu, dan konfigurasi nyata
  • Terakhir, Bagian 9, Bagian 10 membahas cara praktis melakukan profiling, debugging, dan menerapkan pemrosesan paralel pada model di JAX

Rincian lebih lanjut: ringkasan bagian-bagian utama buku

Bacaan terkait

1 komentar

 
GN⁺ 2025-02-07
Komentar Hacker News
  • Ada harapan bahwa JAX akan menggantikan pytorch/cuda dalam beberapa tahun ke depan. Masalah PTX dengan tim Deepseek menunjukkan nilai dari berinvestasi pada pendekatan level lebih rendah untuk memaksimalkan pemanfaatan performa hardware
    • Ini digunakan di internal Google sebagai panduan untuk pekerjaan performa. Mengejutkan bahwa ini dipublikasikan, tetapi tampaknya detail terkait Gemini telah dihapus
    • Panduan ini bagus karena berkat JAX/XLA bisa langsung beralih ke GPU
    • Ada pendapat yang bertanya-tanya mengapa JAX menggunakan tracing alih-alih AST
    • Tautan ke thread tweet penulis dibagikan
    • Ada seseorang yang sedang mencari cara mengonversi situs Jekyll menjadi PDF
    • Ada pujian bahwa ini adalah tulisan yang luar biasa serta ungkapan terima kasih
    • Ada pendapat yang penasaran bagaimana membuat animasi yang keren seperti itu