Nguồn: Google Engineering Blog
Tóm tắt
Google thông báo MaxText, framework training LLM của mình, giờ đây hỗ trợ Supervised Fine-Tuning (SFT) và Reinforcement Learning (RL) trên cấu hình single-host TPU (v5p-8 và v6e-8). Đây là bước tiến quan trọng vì trước đây post-training thường yêu cầu cluster nhiều host, làm tăng chi phí và phức tạp infrastructure đáng kể.
SFT trong MaxText tích hợp native với Hugging Face datasets, hỗ trợ load checkpoint từ cả MaxText lẫn Hugging Face (như Gemma 3), và được tối ưu bằng thư viện Tunix — một JAX-based library thiết kế riêng cho post-training efficiency. Về RL, MaxText hỗ trợ hai thuật toán: GRPO (Group Relative Policy Optimization) và GSPO (Group Sequence Policy Optimization).
GRPO là biến thể memory-efficient của PPO: thay vì dùng separate value function model, GRPO tạo nhiều responses cho mỗi prompt rồi tính relative advantages trong nhóm — giảm đáng kể hardware footprint. GSPO tập trung vào sequence-level importance ratios, cải thiện training stability và đặc biệt hiệu quả cho reasoning benchmarks như GSM8K. Cả hai thuật toán đều dùng vLLM cho high-throughput inference trong training loop.
Việc có thể chạy full RL training trên single-host TPU mở ra khả năng fine-tuning mô hình reasoning mà không cần multi-host cluster. Kết hợp với MaxText v0.2.1 và Tunix, đây là cơ sở hạ tầng post-training production-grade dựa trên JAX cho những ai đang làm việc với TPU trong hệ sinh thái Google Cloud.