大模型做强化学习,像学生先批量做题,再按得分订正:模型要反复生成大量回答,这个“回放”(rollout)阶段很吃显存和算力。TRACE瞄准的正是这一步。它让混合专家模型(MoE——每次只调用部分“专家”模块)用FP4——大致以4位浮点数表示数值——完成回放,希望用更小的“账本”换取更快生成。
难点不只是低精度会产生误差。论文指出,训练和回放即使各自都接近高精度结果,也可能因舍入方向不同而彼此错位,进而让强化学习不稳定。TRACE因此记录回放侧的FP4量化结果,用它指导训练侧如何舍入;为控制额外开销,只缓存部分较深层的尾数与缩放信息。这相当于不再分别校准两把尺子,而是让它们直接对齐。
作者在四个大型MoE模型及推理、编程和长程任务上测试,称TRACE可同时使用FP4权重、激活与KV缓存,强化学习表现接近BF16回放,回放最高提速5.4倍。结果仍来自论文作者自述,但它把低精度优化从“各算各的误差”推进到直接约束训练与回放的一致性,这一点值得继续观察。