MaxText 扩展后训练能力：在单主机 TPU 上引入 SFT 和 RL

Source: Google Developers Blog

APRIL 16, 2026

在快速演进的大型语言模型（LLM）领域，预训练仅是第一步。要将基础模型转化为专用助理或高性能推理引擎，后训练是必不可少的。今天我们宣布 MaxText 的新功能，简化这一过程：监督微调（SFT） 和 强化学习（RL） 现已在单机 TPU 配置（例如 v5p‑8 和 v6e‑8）上可用。

通过利用 JAX 的强大能力以及 Tunix 库的高效性，MaxText 为开发者提供了一条高性能、可扩展的路径，以使用最新的后训练技术细化模型。查看 SFT 和 RL 的完整文档，立即在 TPU 上开启您的后训练之旅。

监督微调 (SFT)：简易的精确调优

监督微调是将预训练模型适配为遵循特定指令或在细分任务上表现出色的主要方法。借助全新的单机 SFT 支持，用户可以在已有的 MaxText 或 Hugging Face 检查点上，使用标注数据集进行微调，且只需最少的配置工作。

关键亮点

• 无缝集成： 原生支持 Hugging Face 数据集（例如 ultrachat_200k）。
• 灵活的检查点： 可直接使用现有的 MaxText 检查点，或在生态系统内转换 Hugging Face 模型（如 Gemma 3）。
• 优化执行： 由 Tunix 提供动力，这是一款基于 JAX 的库，专为后训练效率而设计。

强化学习 (RL)：提升推理能力

对于需要复杂逻辑和推理的任务——例如数学或编码——强化学习是一个改变游戏规则的技术。MaxText 现在在单机 TPU 上支持多种最先进的 RL 算法，并在训练循环中使用 vLLM 进行高吞吐量推理。

• Group Relative Policy Optimization (GRPO) – 一种内存高效的 PPO（近端策略优化）变体。它消除了对单独价值函数模型的需求，而是对每个提示生成多个响应，并在组内计算相对优势。此方式降低了硬件占用，使得在单个 TPU 主机上也能使用先进的 RL。
• Group Sequence Policy Optimization (GSPO) – 关注序列级别的重要性比率和裁剪。通过在序列层面奖励模型行为，它提升了训练的稳定性和效率，在 GSM8K 等基准测试中表现尤为出色。

入门

确保已安装最新的后训练依赖：

uv pip install maxtext[tpu-post-train]==0.2.1 --resolution=lowest
install_maxtext_tpu_post_train_extra_deps

运行 SFT

使用 train_sft 模块启动 SFT 运行，指定模型、数据集和输出目录：

python3 -m maxtext.trainers.post_train.sft.train_sft \
   model_name=${MODEL?} \
   load_parameters_path=${MAXTEXT_CKPT_PATH?} \
   run_name=${RUN_NAME?} \
   base_output_directory=${BASE_OUTPUT_DIRECTORY?}

运行 RL（GRPO/GSPO）

对于 RL，train_rl 模块负责加载策略和参考模型，执行训练，并在推理基准上提供自动评估：

python3 -m maxtext.trainers.post_train.rl.train_rl \
  model_name=${MODEL?} \
  load_parameters_path=${MAXTEXT_CKPT_PATH?} \
  run_name=${RUN_NAME?} \
  base_output_directory=${BASE_OUTPUT_DIRECTORY?} \
  loss_algo=gspo-token \
  chips_per_vm=${CHIPS_PER_VM?}

接下来是什么？

虽然单主机支持为许多开发者提供了强大的入门点，MaxText 是为规模而构建的。这些工作流可以无缝过渡到多主机配置，以训练更大的模型并利用海量数据集。敬请期待后续更新。