Google Tunix：TPU上的自主LLM后训练框架

想象一下：你写了一份 Markdown 规格文档后去睡觉，第二天醒来时发现 AI Agent 已经代表你跑了数十个 LLM 微调实验——找到了最优的 LoRA rank、优化了学习率调度器、调整了 batch size，并将每个经过验证的改进提交到了 Git。

这不再是幻想。今年早些时候，autoresearch 项目展示了自主 LLM Agent 如何在自包含的循环中迭代探索预训练。受这一范式启发，我们创建了 autofinetune：将自主研究循环应用于 LLM post-training（Supervised Fine-Tuning 和通过 GRPO 进行的 Reinforcement Learning），使用 Google 完整的 AI 堆栈——Tunix、Gemma 和 Cloud TPU，由 Antigravity CLI 和 Gemini Flash 3.7 编排。

在本文中，我们将探讨自主研究循环如何应用于 post-training，并通过几个真实的 LLM 微调案例研究进行讲解。

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

传统的 post-training 涉及一个重复的手动循环：

1. 提出假设（例如，"在 LoRA target modules 中添加 attn_vec_einsum 会提高准确率吗？"或者"如果在 GRPO 期间改变 rollout temperature 会发生什么？"）。
2. 编辑训练脚本和超参数。
3. 在加速器上启动任务。
4. 监控 loss 曲线和基准评估。
5. 手动回滚失败记录或在电子表格中记录成功。

正如 autoresearch 中所演示的，我们现在可以用 AI Agent 的力量自动化整个过程：

• 设计 Arena（program.md）：人类定义循环、边界条件、评估标准和约束。
• 提供执行代码（run.py）：一个单一的、干净的、自包含的微调脚本。
• 让 Agent 迭代：Agent 遵循 program.md 中的指令：修改 run.py、运行训练任务、测量目标指标、保留成功的提交或回滚退步的提交，并在 results.tsv 中记录结果。

案例研究 1：在 FunctionGemma 上进行 Supervised Fine-Tuning（SFT）

在 autofinetune 的第一个实验中，我们采用了之前博客中相同的 SFT 设置，并通过创建 autoresearch 循环来扩展它，以优化 google/functiongemma-270m-it 在 google/mobile-actions 数据集上的表现。

• 硬件：Cloud TPU v5e-1
• 迭代速度：每次运行几分钟
• 总运行次数：在几个小时内完成 20 次自动化实验，用于快速实验
• 目标指标：函数调用生成准确率（accuracy）

Agent 在 program.md 中被给予了边界：

• 允许：LoRA rank/alpha、target projection layers、学习率、warmup/decay 调度器、优化器（例如 AdamW/Muon、gradient clipping）、batch size 和种子。
• 禁止：更改数据集、epoch 数量或模型架构。

下面是 sample_runs/SFT_results.tsv 中的一个样本轨迹，展示了 Agent 如何进行爬山优化。

如图所示，Agent 能够自动调整 LoRA rank/alpha、优化器、学习率等，以不断提高模型在生成正确函数调用方面的准确率。

案例研究 2：在 Gemma 上进行数学推理的 Reinforcement Learning（GRPO）

监督微调只是一个简单的测试。在我们的第二个案例研究中，我们采用了 Tunix 仓库中的官方 GRPO 示例（使用 GSM8K 训练 Gemma 3 1B 进行数学推理；训练后的模型在其答案中具有更好的数值准确率和格式准确率），并将其设置为自主 RL 微调。强化学习容易受到超参数敏感性、不稳定性和更长的执行时间的影响——这使得这项任务更具挑战性和耗时。

• 硬件：Cloud TPU v6e-1
• 迭代速度：每次运行几小时
• 运行周期：在 2-3 天内完成 40 次实验
• 目标指标：为了简化 Agent 的外层优化循环，我们选择一个单一的人工评估指标 Post_RL_metric，即 numerical_accuracy + format_accuracy（当然你也可以使用其他指标，例如使用不同的权重）。

下面是记录在 sample_runs/RL_results.tsv 中的一个样本轨迹，展示了 Agent 的进展。Agent 能够识别更好的 LoRA 配置、rollout temperature、KL penalty、system prompt 等，将总 reward 提高了约 10%。

我们希望这个项目能向你展示 AI Agent 在 LLM post-training 领域的威力，并启发你思考如何利用它们使用 Tunix on TPUs 自动化你的 LLM 微调工作流。请查看 autofinetune GitHub 仓库中的代码、样本运行和 program.md 模板，探索 Tunix 库，今天就开始构建你自己的自主 post-training 实验室吧！