昇腾SWA混合注意力训练实践，加速模型训练迭代，长序列推理性能跃升3.5倍

大模型长序列化面临算力与内存双重挑战，模型注意力计算量和 KV cache 占用急剧增长，推理性能挑战大。昇腾提出一种高效的SWA（Sliding Window Attention）模型局部稀疏化方案，以Qwen2.5-1.5B模型为例将模型部分Full Attention层替换为 Sliding Window Attention，通过两阶段Off-Policy蒸馏后训练，仅需12B Token数据即可恢复原始模型能力，构建类SWA-Hybrid模型，在不重新进行大规模预训练的前提下，实现长序列场景下推理吞吐3.5倍提升，提供高效落地的模型改造路径参考实践。

昇腾局部稀疏化与后训练方案

以Qwen2.5-1.5B-Base作为教师模型（即用于提供输出分布监督信号的原始模型）和初始化基座，在保持主体网络结构和参数继承的基础上，将部分Global Attention（GA，全局注意力）层替换为Sliding Window Attention（SWA，滑动窗口注意力）层，形成SWA-GA-Hybrid混合注意力架构（即由全局注意力和滑动窗口注意力组成的混合结构）。相比从零预训练，该方法主要依赖基于教师模型监督的蒸馏优化完成结构迁移，可在较低数据和算力成本下恢复模型通用能力，并提升长序列推理效率。整体流程分为两个训练阶段（Stage 1和Stage 2）以及阶段间的结构选择，如下图1所示。

图1. 局部稀疏化与适应性后训练方案架构图

超网络蒸馏，一次训练获得多种注意力的模型

首先，构建包含GA/SWA双路径的超网络（Supernet Distillation，即每层包含多个注意力分支的网络），实现一次训练获得多种注意力层组合配置的模型。学生模型直接复用教师模型的Q（Query）、K（Key）、V（Value）投影以及输出投影等核心权重，仅扩展Attention模块的计算路径。为避免训练初期破坏原模型已具备的语义能力，该阶段仅更新Attention层相关参数，冻结MLP、Embedding、LayerNorm等非注意力模块。训练过程中，每个Transformer层在前向时随机选择GA或SWA分支，通过结构扰动增强学生模型对不同注意力路径的适应能力。监督信号由两部分组成：一是学生最终logits与教师logits之间的KL loss，用于约束整网输出分布；二是逐层hidden states的MSE loss，用于对齐中间层表征。该阶段的目标不是直接得到最终模型，而是让学生模型具备在多种注意力结构下稳定拟合教师模型的能力，为后续固定Hybrid拓扑提供更好的初始化。

基于Attention-weights统计方法确定最优网络拓扑

在完成Stage1训练后，需要确定适合替换为SWA的层，从而确定最优网络拓扑。本文采用注意力头维度的Attention-weights统计选层方法，具体做法是在保持原始GA结构不变的条件下，对MMLU、GSM8K、BBH等不同类型任务进行抽样推理，统计各层各注意力头在固定512 Token窗口内的Attention-weights。若某一层的大部分注意力头的Attention-weights都集中在此窗口内，说明该层对全局信息依赖较弱，可以替换为SWA；若该层注意力分布更分散，则说明其保留较强的长程依赖，应保留GA。基于此原理，具体实现中采用双重判定规则：层内所有注意力头窗口内Attention-weights均值不低于70%，且窗口内weights占比≥70%的注意力头数占比不低于70%，同时满足两项条件的层被判定为局部主导层并替换为SWA。

除Attention-weights统计方法外，本研究还对比了均匀放置、种群搜索和训练中自适应选层三种方案。具体来讲，自适应选层在每层同时保留GA和SWA两条路径，并引入可学习标量α∈(0,1) 控制两类注意力输出的加权比例，通过冻结其他参数、仅训练α的方式估计每层对全局注意力的依赖程度。种群搜索方法则通过采样不同SWA放置组合，并在下游Benchmark上进行评估，选择较优拓扑。实验结果表明，基于Attention-Weights的统计方法不依赖大规模搜索，选层逻辑可解释性更强，在多数任务上的结果更稳定。

固定拓扑全参数蒸馏，对齐模型最终能力

在确定Hybrid模型拓扑后，进入全参数蒸馏阶段，对模型做最终的能力对齐。训练过程中解冻全部参数，通过KL loss对齐教师模型和学生模型的最终输出分布，使模型在固定混合注意力结构下进一步恢复通用能力、数学推理能力、代码能力和复杂推理能力。相比Stage1，Stage2更强调最终模型能力恢复，是结构改造后能力补偿的关键阶段。

1/4的训练数据量，获得同等模型效果

训练数据累计约12B Token，覆盖通识、数学、代码等多个领域，并对开源数据进行质量筛选，过滤长度小于1k的样本，以增强SWA对较长上下文的适应能力。模型结构保持Qwen2.5-1.5B的基本配置，并最终固定拓扑含20层GA与8层SWA（窗口大小512，结合Head-Wise Gate优化），具体Qwen-SWA模型结构如图2所示。

图2.  Qwen-SWA模型结构

实验在MMLU、BoolQ、GSM8K、HumanEval和BBH等Benchmark上进行评估，分别覆盖综合知识、阅读理解、数学推理、代码生成和复杂推理能力。结果显示，采用8层SWA的Hybrid模型在12Btoken蒸馏后，与原始Qwen2.5-1.5B-Base相比，大部分BenchMark差距控制在约2个点以内，具体数据如下表；与GDN-Hybrid相比，本文方案仅使用其50 Btoken后训练数据量的约24%，模型整体能力恢复效果接近，说明该方法在数据效率和工程成本上具备优势。

模型	训练数据	MMLU	BoolQ	GSM8K	HumanEval	BBH
Qwen2.5-1.5b	/	60.6	73.5	65.4	34.1	43.9
GDN-Hybrid	50B	59.3	71.7	66.6	32.3	43.6
SWA-Hybrid	12B	60.8	71.2	63.2	32.5	43.9

关键方案单点验证

为验证Stage1引入超网络训练的必要性，针对Stage1训练进行了消融实验，结果如下表所示，在相同1B数据、相同8层SWA设置下，经过超网络训练后的模型在MMLU、BoolQ、HumanEval上整体优于直接固定拓扑训练的模型。原因可能在于，固定结构训练的梯度方向更集中，容易使参数过度适配单一路径；而超网络训练通过GA/SWA随机分支引入结构扰动，能够缓解路径依赖，使模型收敛到更平坦、泛化性更好的参数区域。

阶段一	MMLU	BoolQ	HumanEval
超网络训练	60.75	68.7	28.1
固定架构训练	60.31	67.7	28.05

为选择最优的选层方式，对（1）Attention-weights统计选层，（2）均匀放置，（3）自适应选层和（4）种群搜索选层进行了消融实验，结果如图3，Attention-weights统计方法相比均匀放置、自适应选层选层和种群搜索选层，在MMLU、BoolQ、HumanEval、GSM8K等任务上整体表现更优。

图3. 选层实验结果

稀疏化模型性能收益

性能方面，SWA的收益主要体现在长序列推理场景。如图4的推理建模结果所示，随着上下文长度和SWA层替换比例增加，SWA-Hybrid架构可以有效降低长文本推理计算开销和显存压力，当序列长度越长，SWA的局部窗口计算优势越显现，16层SWA配置吞吐加速比约为2.3倍，20层SWA加速比约为3.5倍。

图4. 推理建模结果

在理论建模的基础上，为验证建模结果的可信性，对推理性能进行了实测，如图5所示，在1~1.5k上下文长度下，Qwen-SWA在各榜单上的推理加速比达到1.10~1.14倍，与图3中8层SWA混合模型的理论加速比基本对齐。

图5. 实测性能

后续将继续提升SWA替换比例，并结合Head-Wise的稀疏化方案，探索更高稀疏度/更细粒度稀疏下的模型能力边界；同时结合真实业务长文本任务验证稳定性、吞吐收益和服务部署成本。此外，将进一步引入On-Policy蒸馏，弥补Off-Policy蒸馏在分布对齐上的有限上限，推动Hybrid模型在保持高推理效率的同时，实现对原始教师模型能力的完全追平和超越。

开源链接

当前基于MindSpeed LLM的Qwen-SWA模型后训练代码开源链接：https://gitcode.com/Ascend/MindSpeed-LLM/pull/4736

更多大语言模型训练创新技术，请关注MindSpeed LLM仓库：https://gitcode.com/Ascend/MindSpeed-LLM