模型自蒸馏
搜索文档
仅需15%全量Attention!「RTPurbo」阿里Qwen3长文本推理5倍压缩方案来了
机器之心· 2025-12-23 12:15
行业核心问题:长序列推理的成本与效率瓶颈 - 在当前主流的Full Attention机制下,计算开销会随着输入长度平方增长,导致处理长序列时“又贵又慢”,成为大模型应用中最昂贵的奢侈品 [1] - 尽管各大模型厂商竞相推出更长的上下文窗口(如128K),但在计费上长文本显著更贵,且实际落地产品常需将用户输入变短,这凸显了长序列处理在成本与性能上的核心矛盾 [1] 现有技术路线及其局限性 - 业界主流的Attention压缩方案分为Linear Attention和Sparse Attention两种 [8] - Linear Attention(以Qwen-Next和Kimi-K2为代表)可将存储代价压缩到O(n),计算代价压缩到O(n),但长序列下的信息召回能力显著弱于Full Attention,常需混合使用,限制了性能与加速收益 [10] - Sparse Attention通过稀疏化优化计算,实践中能达到接近**90%**以上的稀疏度,已被DeepSeek-V3.2采用,但其算子与调度设计复杂,影响工程生态的通用性与可维护性 [8][10] - Sliding Window Attention(SWA)因其简单和可规模化复用,被gpt-oss和MiMo等项目采用,并非仅是权宜之计 [10] RTPurbo解决方案:核心原理与假设 - 阿里RTP-LLM团队提出非侵入式的后训练压缩方案RTPurbo,在不损失模型效果的前提下,实现了Attention计算**5倍**压缩 [1][4] - 其核心假设源于对LLM注意力模式的观察:绝大多数Attention Head天然偏好局部信息,只有极少数“长程头”能在长文本中稳定定位并召回关键关联信息 [10] - 类比人类阅读写作,LLM处理长文本时,先由少量Head从长上下文中“召回”相关信息,再在更局部的上下文范围内完成推理输出 [10][11] - 基于此,RTPurbo采用Headwise级别的混合压缩算法,仅保留关键长程Attention Head的全局信息(使用Full Attention),对剩余冗余Head直接丢弃远程Tokens(使用SWA) [4][7] 技术验证与性能表现 - 对比试验表明,仅**15%**的长程头使用Full Attention配合**85%**的Head使用SWA(方案一),其长文能力显著优于反向配置的方案二(后者保留**85%**的KV cache) [12][13] - 在2Wikimqa任务上,原模型(100% FULL)得分为**35.75%**,方案一为**35.40%**,方案二降至**32.58%**;在Passage_retrieval_en任务上,原模型与方案一均为**100%**,方案二降至**62.25%** [13][14] - 未经微调,方案一在Ruler-32k和Ruler-64k任务上表现几乎无损,分别为**96.06%**和**94.44%**,与原模型(**95.15%**和**94.57%**)相当 [14] - 但在某些特定任务(如Hotpotqa、Musique)上,未经微调的压缩模型会出现性能退化,表明需要额外训练来消化注意力模式切换带来的扰动 [14][15] 训练策略:自蒸馏解决数据与能力保留难题 - 直接使用长文SFT/预训练语料对RL后模型进行续训,会导致过拟合、灾难性遗忘,并损伤短文本任务能力,例如在MMLU-Pro上分数从**72.54%**降至**67.28%** [17][18] - RTPurbo采用“模型自蒸馏”作为关键训练策略,让压缩后模型对齐原模型输出,仅需约**1万条32k**长度的预训练语料(训练时间小时级),即可使长文任务表现与原模型持平 [19] - 该策略仅对模型自身输出进行对齐,避免依赖特定领域数据,从而确保短文本下游指标基本无损 [24] 压缩效果与通用能力保持 - 在长文本测试集Longbench上,压缩后的Qwen3-Coder-30B-A3B-Instruct平均得分达**52.02%**,略高于原模型的**51.20%**;Qwen3-Coder-Plus压缩后平均得分为**50.03%**,与原模型**49.82%**相当 [22] - 在Ruler测试集上,压缩模型在32k和64k长度下的平均得分分别为**92.31%**和**88.57%**,与原模型持平 [22] - 在多项短文本通用Benchmark上,采用自蒸馏训练的模型未出现性能衰减:例如Qwen3-Coder-30B-A3B在MMLU PRO上得分从**70.08%**微升至**70.55%**,在HumanEval上从**95.10%**升至**96.34%** [23] 工程优化与极致性能 - 为解决Headwise混合计算模式导致的负载不均衡问题,RTP-LLM在算子层与框架层做了针对性优化 [30][31] - 优化包括:对Full Attention Head进行PTX指令级优化、采用稀疏度感知的动态负载均衡调度、实现高效的SWA算子以避免冗余访存、以及用Context Parallel替代Tensor Parallel以提高GPU利用率 [34][35] - 通过综合优化,RTP-LLM能将Attention稀疏的理论收益稳定转化为端到端加速,在**256k**长序列下实现单算子最高**9倍**加速(图中仅使用**15%**的Full Attention) [35] 项目影响与开源 - RTPurbo被定位为一套具有良好可迁移性和通用性的长序列加速方案,可为更大规模、更多架构的LLM提供高性价比的推理加速路径 [25] - RTP-LLM是阿里巴巴智能引擎团队自研的高性能大模型推理引擎,支持了淘宝、天猫、高德、饿了么等核心业务的大模型推理需求 [37] - 项目模型与推理代码已发布至Huggingface和ModelScope平台,项目已在GitHub开源 [5][9][38]