Huấn luyện Gemma-3 giải toán logic với Tunix GRPO, LoRA và bộ dữ liệu GSM8K

Trong hướng dẫn này, chúng ta sẽ xây dựng một quy trình đào tạo GRPO từ đầu đến cuối để dạy Gemma-3 tư duy giải các bài toán GSM8K bằng cách sử dụng Tunix, JAX, LoRA và các hàm phần thưởng (reward functions) tùy chỉnh. Chúng ta bắt đầu bằng việc chuẩn bị môi trường, xác thực với Hugging Face, tải mô hình Gemma-3 và đóng gói các ví dụ GSM8K vào một định dạng prompt yêu cầu cả tư duy có cấu trúc lẫn đáp án số cuối cùng. Sau đó, chúng ta xác định các hàm phần thưởng để đánh giá mức độ tuân thủ định dạng và tính chính xác về mặt toán học, gắn các bộ điều hợp (adapter) LoRA để giữ cho quá trình đào tạo nhẹ nhàng, đánh giá mô hình cơ sở và chạy GRPO để cải thiện chính sách thông qua các thế hệ được lấy mẫu theo nhóm (group-sampled generations). Bài viết cung cấp một hướng dẫn học tăng cường (reinforcement learning), trong đó chúng ta chỉ đào tạo các trọng số của bộ điều hợp trong khi vẫn giữ quy trình đủ gọn nhẹ cho thiết lập một bộ tăng tốc (single-accelerator) duy nhất.

Cài đặt Tunix và Cấu hình Đào tạo GRPO

Chúng ta thiết lập toàn bộ môi trường Colab bằng cách cài đặt Tunix, JAX, Flax, Qwix, TensorFlow, datasets và các tiện ích notebook hỗ trợ cần thiết cho việc đào tạo GRPO. Chúng ta cấu hình xác thực, vô hiệu hóa các đường dẫn ghi nhật ký không cần thiết, ngăn TensorFlow truy cập vào bộ tăng tốc và xác minh rằng JAX có thể nhận diện các thiết bị TPU hoặc GPU khả dụng. Chúng ta cũng xác định các siêu tham số đào tạo cốt lõi, cài đặt LoRA, giới hạn tạo văn bản, đường dẫn checkpoint và lưới thiết bị (device mesh) để kiểm soát cách mô hình đào tạo trên phần cứng hiện có.

Định dạng Prompt GSM8K và các Hàm Phần thưởng

Chúng ta xác định định dạng tư duy có cấu trúc, yêu cầu mô hình đặt phần tư duy vào trong các thẻ tư duy (reasoning tags) và đáp án số cuối cùng vào trong các thẻ đáp án (answer tags). Chúng ta tải GSM8K từ Hugging Face, trích xuất các đáp án đúng (ground-truth) và chuyển đổi từng bài toán toán học sang định dạng prompt mà quy trình triển khai GRPO yêu cầu. Sau đó, chúng ta tạo các hàm phần thưởng để chấm điểm việc khớp định dạng chính xác, sử dụng thẻ gần đúng, độ chính xác của đáp án và trích xuất số dự phòng để mô hình nhận được phản hồi hữu ích từ nhiều tín hiệu khác nhau.

Tải Gemma-3 và Gắn các Bộ điều hợp LoRA

Chúng ta tải xuống checkpoint Gemma-3 đã chọn, tạo mô hình cơ sở từ safetensors, đồng thời chuẩn bị tokenizer và danh sách token EOS để tạo văn bản. Chúng ta gắn các bộ điều hợp LoRA vào các mô-đun chú ý (attention) và dự báo MLP, cho phép chúng ta đào tạo một chính sách nhẹ nhàng mà không cần cập nhật toàn bộ trọng số mô hình. Chúng ta cũng xây dựng một hàm đánh giá dựa trên bộ lấy mẫu (sampler) và chạy thử nghiệm cơ sở trước khi đào tạo GRPO để đo lường độ chính xác ban đầu và mức độ tuân thủ định dạng của mô hình.

Cấu hình Cụm RL Tunix và Đào tạo

Chúng ta tạo lịch trình tốc độ học (learning-rate schedule), trình tối ưu hóa, cắt gradient (gradient clipping) và cấu hình AdamW để hướng dẫn các cập nhật chính sách LoRA trong quá trình đào tạo. Chúng ta xác định cụm RL Tunix với các vai trò actor, reference và rollout, sau đó kết nối cấu hình rollout với tokenizer, giới hạn tạo văn bản, nhiệt độ lấy mẫu (sampling temperature) và các token EOS. Chúng ta khởi tạo trình học GRPO, khởi chạy TensorBoard để theo dõi các chỉ số trực tiếp và bắt đầu vòng lặp đào tạo GRPO trên các lô (batch) GSM8K đã chuẩn bị.

Đánh giá và Xuất Mô hình đã Đào tạo

Chúng ta đánh giá chính sách LoRA đã đào tạo sau GRPO để so sánh độ chính xác toán học và định dạng phản hồi của nó so với kết quả cơ sở. Sau đó, chúng ta thử xuất mô hình Gemma-3 đã hợp nhất LoRA vào một thư mục safetensors của Hugging Face để checkpoint đã đào tạo có thể được tái sử dụng bên ngoài notebook. Chúng ta kết thúc bằng việc báo cáo đường dẫn checkpoint đầu ra, mang đến một quy trình hoàn chỉnh từ thiết lập đào tạo đến đánh giá sau đào tạo và tùy chọn xuất mô hình.

Kết luận

Tóm lại, chúng ta đã hoàn thành một vòng lặp tinh chỉnh GRPO đầy đủ cho tư duy toán học với Gemma-3, từ chuẩn bị tập dữ liệu và thiết kế phần thưởng đến đào tạo chính sách dựa trên LoRA và đánh giá sau đào tạo. Chúng ta đã thấy cách Tunix tổ chức actor, mô hình tham chiếu, công cụ rollout, trình tối ưu hóa, checkpoint và các chỉ số thành một quy trình học tăng cường có thể tái sử dụng. Chúng ta cũng so sánh mô hình trước và sau GRPO để quan sát xem định dạng đáp án và độ chính xác số có cải thiện trong quá trình đào tạo hay không. Cuối cùng, chúng ta đã xuất mô hình LoRA đã hợp nhất, cung cấp một lộ trình hoàn chỉnh từ các ví dụ GSM8K thô đến một checkpoint Gemma-3 đã được đào tạo và định hướng tư duy.

Xem TOÀN BỘ MÃ NGUỒN TẠI ĐÂY. Ngoài ra, hãy thoải mái theo dõi chúng tôi trên Twitter và đừng quên tham gia SubReddit ML 150k+ thành viên của chúng tôi và Đăng ký Bản tin của chúng tôi. Khoan đã! Bạn có dùng Telegram không? Bây giờ bạn cũng có thể tham gia cùng chúng tôi trên Telegram.

Cần hợp tác với chúng tôi để quảng bá GitHub Repo, Trang Hugging Face, Phát hành Sản phẩm hoặc Hội thảo trực tuyến của bạn, v.v.? Hãy kết nối với chúng tôi.

Sana Hassan, một thực tập sinh tư vấn tại Marktechpost và là sinh viên bằng kép tại IIT Madras, rất đam mê việc ứng dụng công nghệ và AI để giải quyết các thách thức trong thế giới thực. Với sự quan tâm sâu sắc đến việc giải quyết các vấn đề thực tiễn, cô mang đến một góc nhìn mới mẻ cho sự giao thoa giữa AI và các giải pháp đời sống.