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 变体，通过对每个提示生成多个响应并在组内计算相对优势，消除了对单独价值函数模型的需求。这降低了硬件占用，使得高级 RL 能够在单个 TPU 主机上运行。
• 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 是为规模而构建的。这些工作流可以无缝过渡到多机配置，以训练更大的模型并利用海量数据集。敬请期待此方向的后续更新。