AI News
官方更新 Google Developers Blog 正文 ok 717 words

在 TPU 上基于 Tunix 实现 LLM 自主后训练：autofinetune 详解

英文标题：Autonomous LLM post-training with Tunix on TPUs

受开源项目 autoresearch 启发，开发者推出了 autofinetune 项目，通过将自主研究循环引入大语言模型（LLM）的后训练阶段，基于 Google AI 技术栈（Tunix、Gemma、Cloud TPU 与 Antigravity CLI/Gemini Flash 3.7）实现了完全自动化的监督微调（SFT）和强化学习（GRPO）。在 FunctionGemma SFT 和 Gemma 3 数学推理强化学习任务中，AI Agent 在限定边界内自主完成数十次微调迭代与调优实验，成功推动模型指标持续攀升并实现总奖励提升约 10%。该项目展示了 AI Agent 将算法工程师从重复试错中解放出来的巨大潜力，标志着大模型后训练正迈向无人值守的自动化实验室新范式。
打开原文相关 6 · 机会 3
media11.7 KB media88.4 KB media172.5 KB media324.3 KB media153.3 KB media950.9 KB media80.2 KB media59.6 KB media72.8 KB media318.3 KB media92.7 KB

中文

想象一下，只需写下一份 Markdown 规范文档便可安然入睡，第二天醒来发现 AI Agent 已经替你连夜运行了数十次 LLM 微调实验——不仅找到了最优的 LoRA rank、优化了学习率调度、微调了批次大小（batch size），还将每次经过验证的改进自动提交到了 Git。

这已不再是天方夜谭。今年早些时候，autoresearch 项目展示了自主 LLM Agent 如何在一个自包含的闭环中迭代探索预训练流程。受此范式启发，我们构建了 autofinetune：将自主研究循环应用于大语言模型后训练（监督微调 SFT 以及基于 GRPO 的强化学习），并全面采用 Google 的 AI 技术栈——通过 Antigravity CLI 和 Gemini Flash 3.7 编排 Tunix、Gemma 以及 Cloud TPU。

在本文中，我们将深入探讨自主研究循环如何赋能后训练，并详细剖析两个真实的 LLM 微调实战案例。

范式转变：从手动调优到自主闭环

传统的后训练包含着繁琐且重复的人工循环：

  1. 提出假设（例如：“在 LoRA 目标模块中加入 attn_vec_einsum 是否能提升准确率？” 或 “如果在 GRPO 期间调整采样温度（rollout temperature）会发生什么？”）。
  2. 手动修改训练脚本与超参数。
  3. 在加速器（accelerators）上启动任务。
  4. 监控损失曲线与基准评测结果。
  5. 手动回滚失败的尝试，或将成功的指标记录在电子表格中。

正如 autoresearch 所验证的，现在我们完全可以借助 AI Agent 将整个流程自动化：

  • 设计试验场（program.md）：由人类研究员定义循环规则、边界条件、评估标准以及约束项。
  • 提供执行代码（run.py）：一份单文件、结构清晰且自包含的微调脚本。
  • 让 Agent 自主迭代：Agent 严格遵循 program.md 中的规范：自主修改 run.py、运行训练任务、评估目标指标、保留效果更优的提交或回滚退步代码，并在 results.tsv 中记录全流程结果。

案例研究 1：在 FunctionGemma 上进行监督微调（SFT）

在 autofinetune 的第一个实验中，我们采用了此前博客中的 SFT 实验配置，并通过构建自动研究闭环进行扩展，在 google/mobile-actions 数据集上对 google/functiongemma-270m-it 进行优化。

配置环境

  • 硬件：Cloud TPU v5e-1
  • 迭代速度：每次运行仅需数分钟
  • 总运行次数：在数小时内快速完成 20 次自动化实验
  • 目标指标：后训练在函数调用生成任务上的评估准确率（accuracy）

我们在 program.md 中为 Agent 明确划定了探索边界：

  • 允许调整：LoRA rank/alpha、目标投影层、学习率、预热/衰减调度（warmup/decay schedules）、优化器选择（例如 AdamW/Muon、梯度裁剪）、批次大小（batch size）以及随机种子。
  • 禁止修改：数据集、训练轮数（epochs）以及底层模型架构。

样本轨迹

以下是来自 sample_runs/SFT_results.tsv 的轨迹记录样本，展示了 Agent 如何逐步实现指标爬升（hill climbing）。

正如结果所示，Agent 能够自主调节 LoRA rank/alpha、优化器和学习率等各项参数，持续提升模型生成正确函数调用的准确率。

案例研究 2：在 Gemma 上针对数学推理进行强化学习（GRPO）

监督微调仅仅是一个基础验证。在第二个案例研究中，我们选用了 Tunix 官方仓库中的 GRPO 示例（该示例基于 GSM8K 训练 Gemma 3 1B 进行数学推理；训练后的模型在回答的数值准确率和格式准确率上均有改善），并将其配置为自主强化学习微调闭环。由于强化学习通常伴随着极高的超参数敏感性、训练不稳定性以及更长的单次耗时，这项任务更具挑战性也更耗费时间。

试验场设置

  • 硬件：Cloud TPU v6e-1
  • 模型：Gemma 3 1B
  • 数据集：GSM8K
  • 迭代速度：单次运行耗时数小时
  • 运行时长：跨越 2 至 3 天，共完成 40 次实验
  • 目标指标：为了简化 Agent 的外层优化循环，我们设定了一个复合评估指标 Post_RL_metric，即 numerical_accuracy + format_accuracy（你当然也可以根据需求采用其他指标或赋予不同的加权权重）。

样本轨迹

以下是记录在 sample_runs/RL_results.tsv 中的样本轨迹，清晰展现了 Agent 的探索过程。Agent 成功挖掘出了更优的 LoRA 配置、采样温度（rollout temperature）、KL 散度惩罚（KL penalty）以及系统提示词（system prompt），最终将整体奖励（total reward）提升了约 10%。

未来展望

我们希望本项目能展现出 AI Agent 在 LLM 后训练领域中的巨大能量，并启发大家思考如何结合 Tunix 与 TPU 构建完全自动化的模型微调工作流。欢迎前往 autofinetune GitHub 仓库 查看完整代码、运行日志与 program.md 模板，探索 Tunix 库，立即开启属于你自己的自主后训练实验室！

English / 原文

Imagine going to sleep after writing a single Markdown specification and waking up to find that an AI agent ran dozens of LLM fine-tuning experiments overnight on your behalf - discovering optimal LoRA ranks, refining learning rate schedules, tuning batch sizes and committing each verified improvement to Git.

This is no longer a fantasy. Earlier this year, the autoresearch project showcased how autonomous LLM agents can iteratively explore pre-training in a self-contained loop. Taking inspiration from this paradigm, we created autofinetune: applying autonomous research loops to LLM post-training (Supervised Fine-Tuning and Reinforcement Learning via GRPO), using Google’s full AI stack—Tunix, Gemma, and Cloud TPUs orchestrated with Antigravity CLI and Gemini Flash 3.7.

In this post, we’ll explore how the autonomous research loop works for post-training and walk through a couple of real-world LLM finetuning case studies.

The Paradigm Shift: From Manual Tuning to Autonomous Loops

Traditional post-training involves a repetitive, manual cycle:

  1. Formulate a hypothesis (e.g., "Will adding _attn_vec_einsum_ to the LoRA target modules improve accuracy?" or "What happens if we change rollout temperature during GRPO?").
  2. Edit training scripts and hyperparameters.
  3. Launch a job on accelerators.
  4. Monitor loss curves and benchmark evaluations.
  5. Manually revert failures or record successes in a spreadsheet.

As demonstrated in autoresearch, we can now automate this whole process with the power of AI agents:

  • Design the Arena (program.md) : human defines the loop, boundary conditions, evaluation criteria, and constraints.
  • Provide the Execution Code (run.py) : A single, clean, self-contained finetuning script.
  • Let the Agent Iterate : An agent follows instructions in program.md: it modifies run.py, runs the training job, measures the target metric, retains winning commits or reverts regressions, and logs results in results.tsv.

Case Study 1: Supervised Fine-Tuning (SFT) on FunctionGemma

In the first experiment in autofinetune, we took the same SFT setup in our previous blog and extended it by creating the autoresearch loop to optimize google/functiongemma-270m-it on the google/mobile-actions dataset.

The Setup

  • Hardware : Cloud TPU v5e-1
  • Iteration Speed : A few minutes per run
  • Total Runs : 20 automated experiments in a couple of hours for quick experimentation
  • Objective Metric : Post-training evaluation accuracy on function call generation (accuracy)

The agent was given boundaries in program.md:

  • Allowed : LoRA rank/alpha, target projection layers, learning rates, warmup/decay schedules, optimizers (e.g., AdamW/Muon, gradient clipping), batch size, and seeds.
  • Disallowed : Changing the dataset, number of epochs, or model architecture.

Sample Trajectory

Here is a sample trajectory from sample_runs/SFT_results.tsv demonstrating how the agent hill climbed.

As you can see, the agent is able to automatically adjust LoRA rank/alpha, optimizer, learning rate, etc. to keep improving the model’s accuracy in terms of generating correct function calls.

Case Study 2: Reinforcement Learning (GRPO) on Gemma for Math Reasoning

Supervised fine-tuning is only a simple test. For our second case study, we took the official GRPO example from the Tunix repository (which trains Gemma 3 1B for math reasoning using GSM8K; the trained model has better numerical accuracy and format accuracy in its answers) and set it up for autonomous RL finetuning. Reinforcement learning is subject to hyperparameter sensitivity, instability, and longer execution times - making this task more challenging and time-consuming.

The Arena Setup

  • Hardware : Cloud TPU v6e-1
  • Model : Gemma 3 1B
  • Dataset : GSM8K
  • Iteration Speed : A couple of hours per run
  • Run Duration : 40 experiments over 2–3 days
  • Objective Metric : to simplify the outer optimization loop for the agent, we chose a single artificial evaluation metric Post_RL_metric, which is simply numerical_accuracy + format_accuracy (you can of course use other metrics, i.e., using different weights).

Sample trajectory

Below is a sample trajectory logged in sample_runs/RL_results.tsv, showing the agent’s progress. The agent was able to identify better LoRA configurations, rollout temperature, KL penalty, system prompt, etc. to improve the total reward by ~10%.

What’s Next?

We hope this project shows you the power of AI agents in the domain of LLM post-training and inspires you to think about how to leverage them to automate your LLM finetuning workflows using Tunix on TPUs. Please check out the code, sample runs, and program.md templates in the autofinetune GitHub repository, explore the Tunix library, and start building your own autonomous post-training lab today!