arXiv:2602.11937cs.LG2026-02被引 2

通过神经架构搜索优化大模型推理,实现加速不降质。

Extending Puzzle for Mixture-of-Experts Reasoning Models with Application to GPT-OSS Acceleration

  • 结合专家剪枝、窗口注意力等技术,压缩模型体积。
  • 长/短上下文吞吐提升1.63倍和1.22倍,单卡提速2.82倍。
  • 以请求级效率为核心指标,兼顾速度与推理质量。

聚焦推理的LLM通过生成更长的推理链提升答案质量,但额外的token显著增加服务成本,亟需推理优化。本文将Puzzle框架扩展应用于gpt-oss-120B,构建出部署优化版gpt-oss-puzzle-88B。方法包括异构MoE专家剪枝、全上下文注意力替换为窗口注意力、FP8 KV缓存量化(带校准比例),并采用后训练强化学习恢复精度,同时保持生成长度较低。在8XH100节点上,长/短上下文每令牌吞吐分别提升1.63倍和1.22倍;单张NVIDIA H100 GPU上实现2.82倍吞吐提升。由于推理长度随努力程度变化,仅看吞吐或延迟无法反映端到端效率——若推理链增长2倍,2倍吞吐优势将被抵消。因此我们倡导以请求级效率(标准化吞吐/生成令牌数)为指标,绘制准确率-速度前沿曲线。结果表明,gpt-oss-puzzle-88B在整个前沿均优于原模型,最高达1.29倍请求级效率提升。跨多个基准测试,其平均准确率与原模型持平或略优,准确率保留率达100.8%至108.2%,证明后训练架构搜索可大幅降低推理成本而无质量损失。

原文摘要 · Abstract (English)

Reasoning-focused LLMs improve answer quality by generating longer reasoning traces, but the additional tokens dramatically increase serving cost, motivating inference optimization. We extend and apply Puzzle, a post-training neural architecture search (NAS) framework, to gpt-oss-120B to produce gpt-oss-puzzle-88B, a deployment-optimized derivative. Our approach combines heterogeneous MoE expert pruning, selective replacement of full-context attention with window attention, FP8 KV-cache quantization with calibrated scales, and post-training reinforcement learning to recover accuracy, while maintaining low generation length. In terms of per-token speeds, on an 8XH100 node we achieve 1.63X and 1.22X throughput speedups in long-context and short-context settings, respectively. gpt-oss-puzzle-88B also delivers throughput speedups of 2.82X on a single NVIDIA H100 GPU. However, because token counts can change with reasoning effort and model variants, per-token throughput (tok/s) and latency (ms/token) do not necessarily lead to end-to-end speedups: a 2X throughput gain is erased if traces grow 2X. Conversely, throughput gains can be spent on more reasoning tokens to improve accuracy; we therefore advocate request-level efficiency metrics that normalize throughput by tokens generated and trace an accuracy--speed frontier across reasoning efforts. We show that gpt-oss-puzzle-88B improves over gpt-oss-120B along the entire frontier, delivering up to 1.29X higher request-level efficiency. Across various benchmarks, gpt-oss-puzzle-88B matches or slightly exceeds the parent on suite-average accuracy across reasoning efforts, with retention ranging from 100.8% (high) to 108.2% (low), showing that post-training architecture search can substantially reduce inference costs without sacrificing quality.

推理优化MoE大模型压缩效率评估

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。