| 发布日期：2026-05-03 15:24 点击次数：101 |
来源：市场资讯
（来源：华为计算）
为突破传统全注意力模型在长文本推理上的瓶颈，业界正加速由全注意力架构向混合架构演进。然而，从零开始预训练新架构模型往往需要TB级别数据，算力与时间成本高昂。本文设计了一套多阶段蒸馏对齐与自动化架构搜索方案实现模型能力的高效迁移，以Qwen模型为例，在昇腾NPU上的验证结果表明，该方案仅需少量数据即可将全注意力模型重塑为混合注意力模型。
核心挑战
核心挑战1：传统的大模型演进通常需要从零开始进行预训练。这不仅意味着需要消耗超大规模高质量数据集，还会产生极其高昂的算力成本和研发周期。如何通过充分利用现有全注意力模型的先验知识，实现跨架构的能力迁移，从而降低成本？
核心挑战2：混合架构中，不同的注意力层该如何排列组合？业界常用的均匀放置虽然简单，但往往不是效果与硬件效率的最优平衡点，难以发挥混合架构的真正潜力。
核心算法方案
本章阐述了混合架构模型从构建到性能恢复的具体流程。方案首先通过Stage1构建包含多注意力路径的超网络，利用双重蒸馏损失（KL散度与MSE）实现学生模型对教师模型先验知识的初步对齐（针对挑战1）；随后引入自动化架构搜索机制，基于遗传算法在海量组合中快速锁定性能最佳拓扑结构（针对挑战2）；最后在Stage2通过对固定拓扑进行全量参数训练更新，完成模型能力的最终对齐与性能恢复。
展开剩余87%图1.训练流程示意图
Stage 1：超网络蒸馏
超网络构建
整体实验研究以Qwen2.5-1.5B-Base作为基准模型。在本阶段，我们参考了Jet-Nemotron的设计，将原始Transformer层的注意力模块拓展为包含多条候选路径的超网络（Supernet）结构，这种设计可以为后续的架构搜索提供灵活的探索空间。超网络具体包含以下三条路径：原模型的注意力路径、目标模型的线性注意力路径（GDN）、目标模型的全注意力路径。
图2.Stage1训练流程示意图
训练对齐策略
参数冻结机制：为加速超网络收敛并避免在训练早期对原模型语义知识破坏，在Stage1仅对目标模型的线性注意力路径和全注意力路径进行训练，其余权重全部保持冻结状态。
路径激活机制：在训练过程中，每个step会随机采样选择每层需要激活的注意力路径。
双重蒸馏损失
1）KL散度蒸馏
计算学生模型与教师模型在输出层的概率分布差异，强制学生模型学习教师模型的软标签分布：
2）MSE特征对齐
对学生模型与教师模型每层注意力的输出特征空间进行均方误差约束：
总损失函数为两项损失的加权和：
自动化架构搜索
在完成Stage1的超网络路径对齐后，接下来的核心挑战在于如何从海量的拓扑组合中，精确锁定性能最优的混合架构。对于拥有N=28层的Qwen2.5-1.5B模型，其搜索空间规模巨大，传统的暴力穷举法在算力成本上难以承受。为此，本方案设计了一套基于任务反馈的遗传算法（Genetic Algorithm, GA）搜索框架，通过模拟进化过程高效探索最优拓扑结构。
层级灵敏度指数与种群初始化
为了规避随机初始化带来的收敛缓慢与搜索低效，我们引入了层级灵敏度指数作为启发式先验，引导初始种群的生成。
层级灵敏度指数：在保持模型其余层均为线性注意力路径的基准下，依次将第i层替换为全注意力路径，并在抽样下游任务（如MMLU、GSM8K等）上计算当前模型表现，用于表示该层对全注意力机制的需求强度，记为得分Si。利用Softmax函数将各层的层级灵敏度指数得分映射为初始概率分布Pi。
种群初始化：根据概率Pi进行路径采样，确保高性能基因在进化初期具有较高的保留概率。
遗传算法建模与演化机制
目标约束：可根据推理性能要求，预设全注意力目标层数K。
适应度函数：以模型在评测集上的综合得分作为打分评价指标。
演化策略：通过精英保留、变异机制、交叉组合策略，经过多轮迭代，搜索框架将在庞大的搜索空间内自动收敛至一组能使下游任务得分最大化的拓扑结构。
Stage2: 全量参数蒸馏
在确定最优拓扑结构后，模型由超网络切换为固定的混合架构。本阶段旨在通过全量参数的蒸馏训练，完成模型能力的最终对齐。
图3.Stage2训练流程示意图
结果分析
整体精度恢复效果分析
如表1所示，实验结果展示了不同训练范式在相同数据量（50B Token）下的表现差异：Qwen2.5 1.5B Continual是基于原模型权重进行持续预训练（Continual Training），其精度在各项核心榜单上相较于原模型均出现了不同程度的下滑。这表明当前所使用的开源数据集在质量或领域覆盖度上可能存在局限，无法直接通过简单的数据吞吐为模型注入有效收益。但在完全相同的数据集上，Stage1+GA+Stage2的蒸馏训练展现了更好的对齐效果，在以下五个关键榜单上，混合架构模型与原模型精度差距已缩减至2个点以内。实验结果表明蒸馏训练方案可降低对训练数据质量的依赖，以较低的数据成本实现了精度恢复。
表1.全榜单上的精度效果
自动化架构搜索效果分析
表2的实验数据表明，相较于传统的均匀放置方案，经由遗传算法搜索确定的混合架构在多项核心评测任务中均展现出显著的性能优势。为了进一步验证搜索算法的鲁棒性，图4展示了在不同全注意力层数配额下，均匀排布与GA搜索方案在MMLU榜单上的表现对比，结果一致表明GA搜索所得的拓扑结构稳定优于均匀排布方案，并且该模型在配置10~12层全注意力时，其MMLU性能已收敛至表现上限。
表2.均匀放置(Uniform)与遗传算法(GA)各榜单效果对比
图4.不同全注意力选层下均匀放置(Uniform)与遗传算法(GA)在MMLU榜单对比
全注意力层位置分布对模型性能的影响分析
图5对各Transformer层的层级灵敏度指数进行了可视化分析。结果显示，不同层对注意力类型的灵敏度呈现出显著的差异性，且关键层主要集中在模型的中后部区域，进一步说明了全局最优拓扑并非简单的线性分布。值得注意的是，基于不同任务识别出的关键层表现出相似性，例如在第15、21层，MMLU与GSM8K两项任务的层级灵敏度得分均处于高位。
图6对比了不同全注意力层位置分布方案在MMLU评测集上的性能差异。结果显示，将全注意层部署于模型中后部时，其性能增益显著优于前置部署方案，这与图5的可视化结果高度吻合。
图5.层级灵敏度指数可视化
图6.不同全注意力层位置分布方案在MMLU榜单对比
Decode阶段推理性能建模
本节通过AI-Workload建模工具针对混合架构模型在昇腾Atlas 800T A2上的推理表现进行了建模，重点分析了其在显存占用与计算效率方面的表现，并量化对比了其相较于传统全注意力模型在长序列场景下的推理增益。
显存占用分析
传统全注意力模型在Decode阶段需要存储每一层的KV Cache，其空间复杂度随序列长度L呈O(L)线性增长。而混合架构会将部分Full Attention层替换为基于递归状态更新的GDN层，其KV状态是固定大小的，空间复杂度降至O(1)常数级别。如图7所示，在相同Batch Size条件下，512k超长序列推理场景中，不同配置下的混合架构模型显存占用仅为原始全注意力模型19%~46%，极大地释放了长序列场景下的显存压力。
图7.Batch=4情况下不同序列长度的理论显存占用对比
推理吞吐量分析
得益于GDN路径的线性计算复杂度，混合架构可降低了单个Token生成时的计算开销。更重要的是，显存占用的降低释放了原本被KV Cache占据的空间，使得系统能够支持更大的Batch Size。图8建模分析表明，在512k超长序列推理场景中，不同配置下的混合架构模型吞吐能达到原始全注意力模型2.24x~6.41x收益。
图8.单卡情况下不同序列长度的理论吞吐情况对比
推理性能约束下的架构搜索
在实际部署场景中，混合架构的设计不仅需要追求精度恢复，更需要满足严格的推理性能约束。为此，我们可以将AI Workload建模工具与自动化架构搜索算法结合，构建一套端到端的硬件感知架构搜索流程。
为确定搜索空间的理论边界，我们需要先预设推理性能阈值Tthreshold（如特定序列长度下的最小吞吐量、最大显存配额等），然后运用建模工具精细评估算子级计算与访存开销，预估不同全注意力层数配置下的推理效率。通过此流程反推全注意力层占比的理论上限Kmax，从而为上述的自动化搜索算法提供核心目标约束：
未来计划
未来，我们将进一步探索On-policy蒸馏策略来优化复杂场景下的生成效果，并尝试验证该方案在更大参数规模下的Scaling Law潜力。针对当前建模中“推理开销仅取决于全注意力层数”的简化假设，随着模型规模增长及并行策略的引入，我们将深度集成AI Workload建模与搜索算法，构建具备实时反馈的闭环优化逻辑，为高性能架构的持续演进提供系统性支撑。
项目开源链接：
发布于：北京市