微软的 Agent Lightning 于 2026 年 11 月 3 日发布新版本,包含使用低精度推理的强化学习示例和内存优化。智能体训练开发者可以试验在低精度生成答案时稳定训练的校正方法,也可以在使用多个专家的模型中降低重放训练路径时的主机内存负担。

1. 尝试稳定低精度推理下的奖励学习

本次发布为 Qwen3.5-9B 数学训练示例加入分数中心化校正,减去生成答案的模型与训练模型之间概率差异造成的平均更新方向偏差。低精度生成模型与正在训练的模型,选择下一个词元的概率可能不同。分数中心化针对的正是这种差异导致的训练更新方向偏移。

示例在每个生成位置保留采样词元时概率最高的 128 个候选项的概率,并与当前训练模型的概率对照,计算校正值。未记录的其余候选项概率,则近似为与当前模型的概率成比例。

根据公司公布的结果,分数中心化将训练后期的奖励维持在约 0.6。对照基线使用通过比较一组响应的奖励来学习的常规 GRPO,但未加校正,其奖励在约第 45 步降至 0 后持续不变。

这一结果来自单张 B200 上使用 Qwen/Qwen3.5-9B 的数学示例训练曲线,反映了性能随训练的变化。答案生成使用 INT8 权重和激活值以及 INT4 KV 缓存,训练计算采用 BF16。奖励由 math_verify 判断最终答案,正确得 1,否则得 0,生成被截断也得 0 分。

2. 运行示例区分推理精度与训练精度

本次发布还包含模型准备和精度验证工具,以及将训练权重同步到 INT8 推理模型的路径。量化即减少数值表示位数的转换,其对象是将输入转换为加权和的线性层。将词元转换为数值向量的嵌入层和语言模型输出层则保留 BF16。相关设置放在推理部分,如下所示。

actor_rollout_ref:
  rollout:
    dtype: bfloat16
    quantization: int8
    quantization_config_file: ${oc.env:SC_INT8_CONFIG,outputs/precision_validation/Qwen3.5-9B-W8A8-INT8/config.json}
    engine_kwargs:
      vllm:
        linear_backend: triton
        attention_backend: TRITON_ATTN
        kv_cache_dtype: int4_per_token_head

dtype 决定周边计算的精度,显式指定的量化和缓存设置则选择 INT8 与 INT4 路径。示例的完整训练配置为 32 个提示、每个提示 4 个响应,共更新 200 次;快速检查则使用 2 个提示、每个提示 4 个响应,运行 3 步。

3. MoE 路由重放逐个处理临时数组

MoE 是一种从多个专家中选取部分专家处理词元的模型架构。R3 路由重放将生成答案时选择的专家信息传入训练过程,以重放相同路径。此前会先解码整批路由数组,再复制到最终张量;此次修改则每次只解码一个数组,复制后即释放。

改进提案者公布的本地 CPU 测量结果来自对 64 个合成生成样本的处理。如下表所示,用于展开路由信息的临时空间和进程峰值常驻内存(RSS)均有所减少。

表:R3 路由重放的主机内存

条件:本地 CPU 基准测试,合成 rollout,即模型生成的响应批次,共 64 个

单位:MB

MoE 路由重放逐个处理临时数组
项目修改前修改后
临时解码分配峰值126.74.7
进程峰值 RSS660.3543.4
最终张量128128

节省幅度最大的是临时解码缓冲区,最终训练张量在修改前后均为 128 MB。这项改进在保持结果张量格式以及填充、截断行为不变的同时,减少了中间处理空间。

4. 节省的代价与适用范围

低精度推理旨在提升内存与计算效率,但分数中心化额外需要记录并传递候选概率,以及计算训练模型的全词表概率。由于对未记录分布采用近似处理,校正值仍可能存在误差。该示例的 GRPO 与分数中心化都采用基于组内标准差的奖励归一化,即按组内波动幅度调整奖励值。它并非对论文中仅使用均值的设置及梯度下降法的同条件复现;梯度下降法是沿梯度降低损失的训练方法。

本次发布还改进了奖励小数值的保留和归一化精度。它修复了先将奖励转为 BF16,导致 1.0 与 1.001 被当作相同值的问题,也防止了归一化分母溢出,即数值超出表示范围的错误,以及词元数量的舍入问题。这些修改减少了强化学习更新输入值在低精度转换或大数计算中丢失的情况。

本版本分别针对低精度推理与 MoE 路由重放的瓶颈,内存节省来自中间缓冲区的处理,而非最终张量。