Memperkenalkan Tunix: Library Native JAX untuk Pasca Pelatihan LLM

Bagi developer dan peneliti dalam ekosistem JAX, jalur dari model yang telah dilatih sebelumnya ke LLM siap produksi yang sepenuhnya selaras menjadi jauh lebih sederhana.

Hari ini, kami sangat senang bisa memperkenalkan Tunix, library native JAX open source baru yang dibuat khusus untuk pasca pelatihan LLM. Tunix menutup celah kritis dengan menyediakan toolkit yang komprehensif dan ramah developer untuk menyelaraskan model dalam skala besar.

Dibangun untuk performa pada TPU, terutama ketika dikombinasikan dengan MaxText, Tunix menawarkan:

• Rangkaian Algoritme Lengkap: Dapatkan pelatih siap produksi untuk Supervised Fine-Tuning (SFT), penyesuaian preferensi, distilasi pengetahuan, dan metode Reinforcement Learning lanjutan seperti PPO, GRPO, dan GSPO dalam satu library terpadu.
• Desain "White-Box": Pegang kontrol penuh atas pekerjaan Anda. Tunix dirancang untuk mengoptimalkan pengalaman developer, sehingga Anda bisa dengan mudah menyesuaikan loop pelatihan dan kode pasca-pelatihan lainnya tanpa harus melalui lapisan abstraksi.
• Integrasi JAX yang mulus: Sebagai library native JAX, Tunix adalah solusi yang kuat dan mudah digunakan untuk menyelaraskan model open source yang sudah Anda gunakan.

Fitur yang tersedia dalam rilis awal

Rilis awal ini menyediakan API yang modular dan mudah digunakan untuk alur kerja pasca-pelatihan yang paling umum, terintegrasi secara mulus dengan ekosistem JAX:

• Supervised Fine-Tuning (SFT): PeftTrainer adalah model-agnostik dan mendukung metode penyempurnaan bobot penuh dan penyesuaian parameter yang efisien serta populer seperti LoRA dan QLoRA (melalui integrasi kami dengan library qwix).
• Penyesuaian Preferensi: DPOTrainer menyederhanakan penyelarasan dengan menerapkan Direct Preference Optimization (DPO). Teknik canggih ini menggunakan set data sederhana dari respons yang disukai dan ditolak, tanpa perlu melatih dan mengelola model reward secara terpisah.
• Reinforcement Learning (RL): Tunix menyediakan rangkaian pelatih RL untuk menyelaraskan perilaku model dengan preferensi dan petunjuk manusia:
  • PPOLearner: Menyediakan metode kritik-aktor dengan standar emas untuk RLHF dengan mengimplementasikan Proximal Policy Optimization (PPO). Hal ini sangat penting untuk model pelatihan pada tugas berurutan yang kompleks, terutama untuk alur kerja agentik baru yang melibatkan penggunaan alat.
  • GRPOLearner: Menawarkan algoritme RL yang sangat efisien dan bebas kritik. Algoritme ini mengimplementasikan Group Relative Policy Optimization(GRPO), yang menormalkan reward di kelompok respons yang dihasilkan untuk memandu model tanpa kompleksitas dan biaya dari model kritik terpisah.
  • Group Sequence Policy Optimization (GSPO-token): Menawarkan varian algoritme GRPO yang memberikan fleksibilitas yang lebih baik untuk menyesuaikan komputasi keuntungan tingkat token, dan bisa meningkatkan stabilitas untuk pelatihan RL multi-giliran.
• Distilasi pengetahuan: DistillationTrainer memungkinkan kompresi model dengan melatih model ‘siswa’ yang lebih kecil dan efisien untuk mereplikasi output dari model ‘guru’ yang lebih besar. Ini adalah teknik penting untuk menerapkan model berkinerja tinggi di lingkungan produksi dengan latensi ketat atau batasan biaya. Tunix menyediakan algoritme distilasi berikut secara langsung:
  • Distilasi Berbasis Logit: memanfaatkan probabilitas output akhir model guru sebagai “target lunak” untuk memandu model siswa.
  • Transfer perhatian: memanfaatkan fitur perhatian model guru untuk memandu model siswa.
• Paket PyPI: Tunix tersedia sebagai paket yang dirilis di PyPI -
  • pip install google-tunix
• Contoh: Contoh untuk semua algoritme yang didukung dengan implementasi kanonik dari beberapa model open source terkemuka di repo Tunix.
• AI Agentik: Tunix memungkinkan pelatihan Agen yang dapat menalar menggunakan LLM dan berinteraksi dengan lingkungan eksternal.

Hasil kuantitatif

Kami telah membuat beberapa notebook Python untuk membantu pengguna menggunakan Tunix. Hasil di bawah ini menunjukkan keefektifan implementasi GRPO Tunix. Pada tolok ukur penalaran matematika GSM8K, menyempurnakan model Gemma 2 2B-IT dengan Tunix menghasilkan ~12% peningkatan relatif dalam akurasi jawaban pass@1. Kami mengamati peningkatan yang menjanjikan di semua metrik, yang menunjukkan kemampuan library untuk menyelaraskan perilaku model dengan cepat dan efektif.

Untuk mengakomodasi sifat stochastic pembuatan teks, kami mengevaluasi performa menggunakan pass@1 (penelusuran greedy) dan pass@5 (pengambilan sampel dengan keragaman) untuk mengukur ketepatan dalam satu atau lima kali percobaan. Evaluasi kami berfokus pada tiga metrik utama:

• Akurasi jawaban: Persentase prediksi dengan jawaban numerik akhir yang benar.
• Akurasi jawaban (parsial): Metrik yang lebih fleksibel dengan jawaban model berada dalam 10% dari jawaban yang benar (rasio antara 0,9 dan 1,1).
• Akurasi format: Persentase sampel saat model menggunakan alasan dan token jawaban yang diperlukan dengan benar.

Untuk validasi, akurasi pass@1 baseline kami sebesar ~52% selaras dengan ~51% yang dilaporkan oleh Eleuther LM Eval Harness untuk model dasar, yang mengonfirmasi validitas pengaturan kami. Meskipun akurasi absolut sensitif terhadap pemformatan prompt (mis. menggunakan <start_answer> vs. <answer>), peningkatan performa yang signifikan dari pasca pelatihan tetap konsisten dalam berbagai setelan.

Dipercaya oleh peneliti dan inovator

Dari lab akademis terkemuka hingga startup AI, Tunix telah mendukung gelombang pengembangan ML yang berikutnya. Kami mengembangkan Tunix bekerja sama dengan mitra kami untuk memecahkan tantangan dunia nyata dalam penyelarasan model dan AI agentik. Inilah yang mereka katakan:

"Penelitian saya berfokus pada pembelajaran yang berpusat pada data, yang melibatkan penyiapan data berkualitas tinggi untuk meningkatkan performa model, terutama pada fase pasca-pelatihan model bahasa besar (LLM). Tantangan utamanya adalah melakukan iterasi dengan cepat pada sampel data untuk mengidentifikasi mana yang bermanfaat dan mana yang tidak. Untuk hal ini, Tunix adalah library yang sempurna. Desain ‘white-box’-nya memberikan tim saya kontrol penuh atas loop pelatihan, sehingga kami dapat dengan mudah memodifikasi dan mengadaptasi kode untuk kebutuhan penelitian kami yang spesifik. Fleksibilitas ini merupakan keunggulan yang signifikan dibandingkan framework lain dan sangat penting untuk mempercepat analisis data iteratif kami."

— Hongfu Liu, Assistant Professor of Computer Science, Brandeis University; Senior Area Chair for NeurIPS; Area Chair for ICLR

"Bottleneck utama dalam reinforcement learning pasca pelatihan adalah kelangkaan lingkungan dengan reward yang dapat diverifikasi. Game menyediakan lingkungan multi-giliran yang sempurna untuk mengatasi hal ini, dan Tunix adalah framework yang ideal untuk penelitian ini. Ini memungkinkan kami untuk membangun langsung di atas JAX, memanfaatkan TPU dan paralelisasi yang mudah. Dibandingkan dengan alternatif lain, Tunix adalah library yang ringan dengan codebase yang bersih dan mudah dikelola. Ia menawarkan kustomisasi model tingkat tinggi dan hiperparameter tanpa lapisan abstraksi yang berlebihan seperti framework lainnya. Pendekatan yang efisien ini sangat penting untuk pekerjaan kami, dan kami menemukan bahwa kurva pembelajarannya cukup landai, karena Anda tidak perlu menjadi ahli JAX untuk bisa menggunakannya secara efektif."

— Hao Zhang, Assistant Professor, UC San Diego, Co-creator vLLM, Chatbot Arena (LMSys), dan Inventor Disaggregated Serving

Precur AI adalah startup yang membangun Agent Compiler yang mentransformasi alur kerja latar belakang menjadi agen berbasis kode yang andal dan efisien. Hanjun Dai, Co-Founder dan CTO, berkata:

"Perusahaan kami berfokus pada agen latar belakang yang berjalan 24/7 tanpa pengawasan. Tujuan utamanya adalah ketahanan agen, jadi kami melakukan “kernel agen” pasca-pelatihan - model yang dioptimalkan untuk tugas jangka panjang tetapi repetitif. Luasnya desain Tunix, yang mencakup SFT, RL, dan distilasi, memungkinkan kami menjaga seluruh agent development stack tetap terpadu. Integrasi aslinya dengan ekosistem JAX dan TPU merupakan keuntungan yang signifikan. Kemudahan kustomisasi dengan Flax untuk pengembangan dan Qwix untuk penyajian dalam bentuk kuantisasi menjadikannya framework yang bersih dan kuat yang sangat cocok dengan alur kerja kami."

— Hanjun Dai, Cofounder dan CTO, PreCur AI

Komunitas dan Kolaborasi - Mari Bergabung

Kami membangun Tunix secara terbuka dan mengundang Anda untuk bergabung, mencoba langsung, dan ikut berkontribusi dalam komunitas kami.

• Berkontribusi untuk Tunix: Kami secara aktif mencari kolaborator dan siap mendukung kontribusi Anda. Jika Anda tertarik untuk mengembangkan fitur atau lingkungan agentik baru, meningkatkan algoritme, atau membentuk kemitraan penelitian, silakan gunakan formulir ini untuk memberi tahu kami bagaimana Anda ingin berkontribusi.
• Repositori dan Dokumentasi GitHub: Anda bisa menemukan kode sumber, issue tracker, dokumentasi lengkap, dan forum diskusi di repositori GitHub kami dan tunix.readthedocs.io.
• Contoh Praktis: Cara terbaik untuk memulainya adalah dengan menjalankan kodenya. Kami telah menyiapkan serangkaian notebook Python yang bisa Anda jalankan untuk memulai dengan setiap pelatih inti kami.
• MaxText: MaxText adalah library LLM open source dengan performa tinggi yang sangat skalabel dan implementasi referensi yang ditulis sepenuhnya dalam bahasa Python/JAX serta menargetkan Google Cloud TPU dan GPU untuk pelatihan.

Kami sangat senang bisa berbagi Tunix dengan komunitas JAX dan tak sabar ingin segera melihat kreasi Anda.