社区如何用 Tunix 和 TPU 训练 Gemma“思考”
英文标题:How the community trained Gemma to "Think" with Tunix and TPUs
Google 在 Kaggle 上举办了 Tunix 黑客松,吸引超过 11,000 名参赛者,使用受限的 TPU 算力将非推理基座模型(Gemma-2-2B、Gemma-3-1B)转变为通用推理模型。获胜团队普遍采用多阶段后训练管线,结合监督微调、偏好优化与强化学习,并利用 LLM 裁判或 TF-IDF 等奖励机制提升逻辑推理质量。该比赛证明社区可在九小时、TPU v5e-8 的预算下获得可观推理效果,并公开了完整的训练配方,对想自训练推理模型的开发者具有重要参考意义。
59.8 KB
64.8 KB
496.8 KB
267.7 KB
62.2 KB
中文
大型语言模型(LLM)在处理复杂任务时,往往得益于“先思考,后开口”。像 Gemini 3 这样的前沿 LLM 和 Gemma 4 等领先的开放权重模型,能在回答用户问题前生成显式的推理轨迹,通常称为思维链。但这一推理能力如何训练得到却很少被公开。尽管网上有很多针对数学或编程等简单可验证任务的推理教程,但面向通用推理的、易于获取且可复现的训练配方(包括数据、训练策略、可运行代码和评测)仍十分稀缺。
这促使我们在 Kaggle 上举办了 Google Tunix 黑客松:训练模型展示其思考过程:我们向开发者发起挑战,要求他们利用 Tunix 和 Kaggle TPU,将非推理基座模型(Gemma-2-2B 和 Gemma-3-1B)转变为通用推理模型。反响极为热烈:超过 11,000 名参赛者和 300+ 份高质量提交证明,即使在非常有限的计算预算(Kaggle TPU v5e-8、9 小时)下,社区也能完成相当出色的推理训练。在这篇文章中,我们将着重介绍获胜者采用的技术,并分享能让模型在关键垂直行业中推理的核心配方,以便你也能训练自己的推理模型。
聚焦获奖方案:关键创新
获奖提交展现了对后训练的深入理解,创造性地融合了监督学习、偏好优化和强化学习。
G-RaR 通过结合监督微调(SFT)与 GRPO,并以一种新颖的基于评分细则的 LLM-as-judge 奖励系统驱动,训练 Gemma 模型生成结构化推理。
- 如何提升推理能力:通过显式训练模型在
<reasoning>标签内“展示其思考过程”,再输出答案,模型的推理能力得到增强。其 GRPO 背后的技术——G-RaR(Rubrics as Rewards),使用更大的裁判模型(Gemma-3-12B)基于任务特定的评分细则来评估这些中间逻辑步骤的质量。该技术将离散的细则分数转换为连续、归一化的奖励信号,为模型的逻辑提供密集且平滑的反馈。这让模型能够持续提升推理能力,而无须完全依赖完全匹配的正确性,因此即使对于开放、不可验证的任务也极为有效。 - 技术方案:团队采用了两阶段后训练管线:
- 第一阶段(SFT):通过 LoRA 在约 33k 样本的数据集上微调 Gemma-2-2B-IT 模型,建立基线。这个“热身启动”教会模型可靠地输出
<reasoning>...</reasoning><answer>...</answer>结构。 - 第二阶段(GRPO):随后使用基于复合奖励函数(格式奖励 + 精确答案奖励 + G-RaR 分数)的 GRPO 对模型进一步优化。为了克服算力限制,团队在单个 Kaggle TPU v5e-8 上采用分片网格架构,将策略/参考模型放在一个网格上,裁判模型放在另一个网格上,实现真正的并行执行。
- 第一阶段(SFT):通过 LoRA 在约 33k 样本的数据集上微调 Gemma-2-2B-IT 模型,建立基线。这个“热身启动”教会模型可靠地输出
将 1B 参数模型演化为结构化推理引擎(“Pinocchio”),依赖高效、9 小时的 TPU 管线(SFT → SimPO → GRPO)
- 如何提升推理能力:模型学会在回答前生成结构化的
<reasoning>轨迹,从基础的模式匹配转向逻辑演绎。这一能力是依次构建的:SFT 灌输基础思维链,SimPO 锁定严格格式(防止冗长 hack),GRPO 则使用 LLM-as-a-Judge 奖励连贯性并严厉惩罚幻觉,由此精炼逻辑。 - 技术方案:该管线包含三个阶段:
- SFT(蒸馏):使用 OSS-120B 教师模型和 Gemini 任务路由器在 70k 提示上进行训练。
- SimPO(对齐):替代内存密集的 DPO,高效强制严格 XML 格式。
- GRPO(精炼):使用 Gemini 2.0 Flash 作为异步裁判,动态奖励准确性、逻辑和格式。
- 定制 Tunix:团队显式扩展了 Tunix 库以支持此工作流:
- 将自定义 SimPO 损失函数(带长度归一化)注入
DPOTrainer。 - 创建高吞吐量、异步评估引擎,实时处理 GRPO 奖励信号。
- 将自定义 SimPO 损失函数(带长度归一化)注入
将结构化的“IDEA-E”伦理推理框架蒸馏到 2B 模型中,采用课程引导的 GRPO 和快速 TF-IDF 奖励系统。
- 为何能提升推理:IDEA-E 框架强制模型在回答前进行逐步逻辑演绎,防止过早猜测。同时,TF-IDF 奖励通过激励推理轨迹中使用上下文相关词汇,避免冗长“废话”。
- 技术方案:管线包含两个阶段:
- SFT:在教师数据上微调以建立 IDEA-E 格式。
- GRPO:使用课程引导和 TF-IDF 奖励替代缓慢的 LLM 裁判进行强化学习。
- 定制 Tunix:团队通过将自定义 TF-IDF 奖励函数集成到 Tunix 的 GRPO 管线中,实现了在 CPU 上快速、非阻塞的奖励计算。
荣誉提名
除了前三名登上领奖台,还有几份提交展现了突出的创造力和技术深度:
-
方案:不依赖静态离线数据集,而是在 Tunix 框架内从零实现在线策略蒸馏方法。他们使用一个更大、能力更强的教师模型(分三个阶段训练过),在训练过程中根据学生模型的生成动态产出推理轨迹,形成了更紧密的反馈环路。
-
方案:由参赛者 TheItCrow 开发,该项目专注于自定义数据集构建和结构化奖励建模。
- 他们整理了 Deep-CoRGI(Cognitive Reasoning Guided Interface)数据集,专门用于教授思维链。
- 他们训练了一个自定义 ThoughtTeacher 奖励模型,不仅评估最终答案的正确性,还评估推理步骤本身的逻辑流程。
我们同样对几份聚焦于特定领域推理训练的提交印象深刻,例如医疗、化学、法律和机器人学。
- 医疗:GRPO 生成结构化、逐步推理轨迹,增强了其复杂临床问题解决输出的可解释性和可靠性。
- 化学:逐步推理轨迹使小语言模型能够解决复杂化学推理任务,让化学用例从中受益。
- 法律:通过 GRPO 进行后训练,强化了结构化、逐步推理,使 Gemma 3 1B 模型能够准确分析复杂法律数据,产生可靠、逻辑合理的解释。
- 机器人学:逐步推理生成让模型能够在单次会话训练的限制下,解决多步机器人规划与决策任务。
准备好构建了吗?
Tunix 黑客松通过产出大量令人印象深刻、现已全部公开的推理训练配方,使训练高性能、结构化推理模型的工作变得民主化。借助 Tunix 和免费的 Kaggle TPU,开发者现在可以在可及的硬件上取得出色成果。
如果你已准备开始对自己的推理模型进行后训练,以下资源可供入门:
- 在 GitHub 上探索 Tunix:查阅 Tunix 官方仓库,获取代码、文档和社区示例。
- 尝试 Colab 教程:在 Google Colab 中启动免费 TPU 实例,试用 Tunix 示例运行你的第一个 SFT 或 RL 循环。
- 进一步了解强化学习:阅读 Tunix 中的 RL 文档,理解如何利用强化学习微调模型。
English / 原文
Large Language Models (LLMs) often benefit from "thinking" before they speak for complex tasks. Frontier LLMs like Gemini 3 and leading open weight models like Gemma 4 can produce explicit reasoning traces, commonly called Chain-of-Thought, before answering user questions. But how this reasoning capability is trained is often not disclosed. While there are many reasoning tutorials available on the Internet to train for simple verifiable tasks such as math or coding, accessible and easy-to-reproduce training recipes (including data, training strategy, runnable code and evaluations) for general reasoning remain scarce.
This motivated us to hold the Google Tunix Hack: Train a model to show its work hackathon on Kaggle: we challenged developers to transform non-reasoning base models (Gemma-2-2B and Gemma-3-1B) into general reasoning models, using Tunix and Kaggle TPUs. The response was overwhelming: over 11,000 entrants and 300+ high-quality submissions proved that decent reasoning training can be done by the community even with a very limited compute budget (Kaggle TPU v5e-8 for 9 hours). In this post, we’ll highlight the techniques used by the winners and share key recipes that allow models to reason across key vertical industries, so you can train your own reasoning models.
Highlighting the Winners: Key Innovations
The winning submissions demonstrated a sophisticated understanding of post-training, combining supervised learning, preference optimization, and reinforcement learning in creative ways.
G-RaR trains Gemma models to produce structured reasoning by combining Supervised Fine-Tuning (SFT) with GRPO, driven by a novel rubric-based LLM-as-judge reward system.
- How It Improves Reasoning The model's reasoning power is improved by explicitly training it to "show its work" inside
<reasoning>tags before outputting an answer. The underlying technique (for GRPO), G-RaR (Rubrics as Rewards), uses a larger judge model (Gemma-3-12B) to evaluate the quality of these intermediate logical steps based on task-specific rubrics. By converting discrete rubric scores into continuous, normalized reward signals, the technique provides dense, smooth feedback on the model's logic. This allows the model to continuously improve its reasoning capabilities without relying solely on exact-match correctness, making it highly effective even for open-ended, non-verifiable tasks. - Technical Solution The team utilized a two-stage post-training pipeline:
- Stage 1 (SFT): The Gemma-2-2B-IT model is fine-tuned via LoRA on a ~33k sample dataset to establish a baseline. This "warm start" teaches the model to reliably output the
<reasoning>...</reasoning><answer>...</answer>structure. - Stage 2 (GRPO): The model is then refined using GRPO-based on a composite reward function (Format Reward + Exact Answer Reward + G-RaR Score). To overcome compute constraints, the team used a split-mesh architecture on a single Kaggle TPU v5e-8, placing the policy/reference models on one mesh and the judge model on the other for true parallel execution.
- Stage 1 (SFT): The Gemma-2-2B-IT model is fine-tuned via LoRA on a ~33k sample dataset to establish a baseline. This "warm start" teaches the model to reliably output the
Evolving a 1B parameter model into a structured reasoning engine ("Pinocchio") via a highly efficient, 9-hour TPU pipeline (SFT → SimPO → GRPO)
- How it Improves Reasoning The model learns to generate a structured
<reasoning>trace before answering, shifting from basic pattern matching to logical deduction. This is built sequentially: SFT instills foundational Chain-of-Thought, SimPO locks in strict formatting (preventing verbosity hacks), and GRPO refines logic by using an LLM-as-a-Judge to reward coherence and heavily penalize hallucinations.. - Technical Solution The pipeline consists of three stages:
- SFT (Distillation): Trained on 70k prompts using an OSS-120B teacher model and a Gemini task-router.
- SimPO (Alignment): Replaced memory-heavy DPO to efficiently enforce strict XML formatting.
- GRPO (Refinement): Used Gemini 2.0 Flash as an asynchronous judge to dynamically reward accuracy, logic, and format.
- Customizing Tunix: The team explicitly extended the Tunix library to support this workflow by:
- Injecting a custom SimPO loss function (with length normalization) into the
DPOTrainer. - Creating a high-throughput, asynchronous evaluation engine to process GRPO reward signals on the fly.
- Injecting a custom SimPO loss function (with length normalization) into the
Distilling the structured "IDEA-E" ethical reasoning framework into a 2B model using curriculum-guided GRPO and a fast TF-IDF reward system.
- Why it Improves Reasoning The IDEA-E scaffold forces the model through a step-by-step logical deduction before answering, preventing premature guessing. Simultaneously, the TF-IDF reward prevents verbose "yapping" by incentivizing the use of context-relevant vocabulary in the reasoning trace.
- Technical Solution The pipeline features two stages:
- SFT: Fine-tuning on teacher data to establish the IDEA-E format.
- GRPO: Reinforcement learning using curriculum guidance and a TF-IDF reward instead of a slow LLM judge.
- Customizing Tunix: The team extended Tunix by integrating their custom TF-IDF reward function into the Tunix GRPO pipeline, allowing for rapid, non-blocking reward calculations on the CPU.
Honorable Mentions
While the top three spots took the podium, several other submissions showcased strong creativity and technical depth:
-
The Approach: Instead of relying solely on static offline datasets, they implemented an on-policy distillation method from scratch within the Tunix framework. They used a larger, highly capable teacher model (trained in 3 phases) to generate reasoning traces dynamically in response to the student model's generations during training, creating a tighter feedback loop.
-
The Approach: Developed by participant TheItCrow , this project focused on custom dataset curation and structured reward modeling.
- They curated the Deep-CoRGI (Cognitive Reasoning Guided Interface) dataset, specifically designed to teach Chain of Thought.
- They trained a custom ThoughtTeacher reward model to evaluate not just the correctness of the final answer, but the logical flow of the reasoning steps themselves.
We are also very impressed with several submissions that focus on reasoning training in specific domains, such as medical, chemistry, legal and robotics.
- Medical: GRPO generates structured, step-by-step reasoning traces enhancing the interpretability and reliability of its complex clinical problem-solving outputs
- Chemistry: step-by-step reasoning traces benefited the chemistry use case by enabling a small language model to solve complex chemistry reasoning tasks.
- Legal: Post-training via GRPO reinforces structured, step-by-step reasoning, enabling the Gemma 3 1B model to accurately analyze complex legal data and produce reliable, logically sound interpretations.
- Robotics: step-by-step reasoning generation allows the model to solve multi-step robotics planning and decision-making tasks under single-session training constraints.
Ready to build?
The Tunix Hackathon democratizes training highly capable, structured reasoning models by producing so many impressive reasoning training recipes that are now all publicly available. With Tunix and free Kaggle TPUs, developers can now achieve strong results on accessible hardware.
If you're ready to start post-training your own reasoning models, here are some resources to get started:
- Explore Tunix on GitHub: Check out the official Tunix repository to access the code, documentation, and community examples.
- Try a Colab Tutorial: Spin up a free TPU instance in Google Colab and try the Tunix examples to run your first SFT or RL loop.
- Learn More About Reinforcement Learning: Read the RL documentation in Tunix to understand how to leverage reinforcement learning to finetune your model.