扩展用于混合专家推理模型的Puzzle框架及其在GPT-OSS加速中的应用
机器学习
2026-03-30 v2
摘要
以推理为中心的LLM通过生成更长的推理轨迹来提高答案质量，但额外的token会显著增加服务成本，从而激发了推理优化。我们将后训练神经架构搜索(NAS)框架Puzzle进行扩展并应用于gpt-oss-120B，产生了gpt-oss-puzzle-88B，一个部署优化的衍生模型。我们的方法结合了异构MoE专家剪枝、将全上下文注意力选择性替换为窗口注意力、具有校准尺度的FP8 KV缓存量化，以及用于恢复精度的后训练强化学习，同时保持较低的生成长度。在每token速度方面，在8XH100节点上，我们在长上下文和短上下文设置中分别实现了1.63倍和1.22倍的吞吐量加速。gpt-oss-puzzle-88B在单个NVIDIA H100 GPU上也提供了2.82倍的吞吐量加速。然而，由于token数量可能随推理努力和模型变体而变化，每token吞吐量(tok/s)和延迟(ms/token)并不一定能带来端到端的加速：如果推理轨迹增长2倍，2倍的吞吐量增益就会被抵消。相反，吞吐量增益可以用于更多的推理token以提高准确性；因此，我们提倡按请求级别的效率指标，该指标通过生成的token对吞吐量进行归一化，并描绘出跨推理努力的准确率-速度前沿。我们展示了gpt-oss-puzzle-88B在整个前沿上优于gpt-oss-120B，提供了高达1.29倍的请求级效率提升。在各种基准测试中，gpt-oss-puzzle-88B在跨推理努力的平均套件准确率上匹配或略微超过其父模型，保留率从100.8%(高)到108.2%(低)不等，表明后训练架构搜索可以在不牺牲质量的情况下大幅降低推理成本。
引用
@article{arxiv.2602.11937,
title = {Extending Puzzle for Mixture-of-Experts Reasoning Models with Application to GPT-OSS Acceleration},
author = {Akhiad Bercovich and Nir Ailon and Vladimir Anisimov and Tomer Asida and Nave Assaf and Mohammad Dabbah and Ido Galil and Amnon Geifman and Yonatan Geifman and Izhak Golan and Roi Koren and Itay Levy and Zach Moshe and Pavlo Molchanov and Najeeb Nabwani and Mostofa Patwary and Omri Puny and Tomer Ronen and Itamar Schen and Elad Segal and Ido Shahaf and Oren Tropp and Ran Zilberstein and Ran El-Yaniv},
journal= {arXiv preprint arXiv:2602.11937},
year = {2026}
}