Hugging Face Transformers 中 GraniteMoeSWA：融合共享专家稀疏 MoE、滑窗注意力与可学习注意力 Sink 的因果语言模型

GraniteMoeSWA 是 Hugging Face Transformers 中一个组合式混合专家（MoE）因果语言模型实现：它把 GraniteMoeShared 的稀疏 MoE 结构（含可选的共享专家）与 GraniteSWA 的分层滑窗注意力（sliding-window attention）及可学习注意力 Sink 结合起来。本文将从模型架构要点、注意力后端选择、配置类逐参数解析、生成示例到并行切分方案，给出可直接复现与二次开发所需的完整指引。

GraniteMoeSWA 是什么：三种机制的叠加

GraniteMoeSWA 于 2026-07-29 被贡献进 Hugging Face Transformers，模型类型标识为 granitemoe_swa，其设计目标是让一个 MoE 解码器同时具备三方面特性：

• 混合专家（Mixture of Experts）：每个解码层都包含一个稀疏专家层（block_sparse_moe），对每个 token 在 num_local_experts 个局部专家中路由到其中 num_experts_per_tok 个。同时支持可选的共享专家（shared experts），但默认被关闭（shared_intermediate_size=0）；只要将该参数设置为正值即可启用。
• 分层滑窗注意力：每一层的注意力类型由 layer_types 决定，取值为 "full_attention"（全注意力）或 "sliding_attention"（滑窗注意力）。默认策略是每 4 层保留 1 层全注意力（i % 4 == 0），其余各层只关注最近 sliding_window 个 token。
• 可学习的逐头注意力 Sink：每个注意力头学习一个标量 sink，对注意力输出按 sigmoid(logsumexp(attn_logits) - sink) 进行重缩放。数学上这等价于在 softmax 分母中追加一个可学习的额外 logit（即 GPT-OSS 采用过的 attention-sink 机制）。

代码结构上，GraniteMoeSWA 的配置类是 GraniteMoeSharedConfig 的子类，模型主体（GraniteMoeSWAModel、GraniteMoeSWAForCausalLM、各层与模块）分别继承自 GraniteMoeSharedModel、GraniteMoeSharedForCausalLM 系列，而注意力部分则复用了 GraniteSWAAttention 的实现。因此可以将其理解为“以 GraniteMoeShared 为骨架、以 GraniteSWA 的注意力替换原自注意力”。该组合过程在模块化源文件 modular_granitemoe_swa.py 中有清晰的注解说明。

每个解码层内部的长什么样

从 modeling_granitemoe_swa.py 与其父类 modeling_granitemoeshared.py 的结构看，每一层 GraniteMoeSWADecoderLayer 由三大部分组成（见父类 GraniteMoeSharedDecoderLayer 的构造函数）：

1. input_layernorm / post_attention_layernorm：RMSNorm 归一化；
2. self_attn：GraniteMoeSWA 专用注意力模块（等价于 GraniteSWA 的注意力）；
3. block_sparse_moe：稀疏 MoE 模块，外加一个可选的 shared_mlp（当 shared_intermediate_size == 0 时为 None，不创建）。

在前向过程中，MoE 的输出与共享 MLP 的输出会相加：

• 若 shared_mlp 为 None：hidden_states = moe_hidden_states；
• 否则：hidden_states = moe_hidden_states + self.shared_mlp(hidden_states)。

共享专家是一个小型的稠密 FFN（同样由 gate_proj/up_proj/down_proj 组成），对全体 token 无差别地计算，用于兜底那些路由分配不够充分的特征表达。

注意力 Sink 的实现细节与公式

GraniteMoeSWA 的注意力 Sink 实现被定义在 GraniteSWAAttention 中：

• 每个头的 sink 是一个零初始化的可学习参数：self.sinks = nn.Parameter(torch.zeros(config.num_attention_heads))（modeling_granite_swa.py）；
• 该层是否为滑窗注意力取决于 self.layer_type = config.layer_types[layer_idx]，若为 "sliding_attention" 则 self.sliding_window = config.sliding_window，否则为 None；
• 前向时将 sliding_window 与 s_aux=self.sinks 一并传入实际的注意力实现。

在 eager 路径上（eager_attention_forward），sink 的计算分成三步，且全程刻意保持 fp32 精度以保证数值稳定性：

# 1. 对未归一化的注意力权重做 logsumexp
lse = torch.logsumexp(attn_weights, dim=-1)                 # (batch, num_heads, q_len)
# 2. 用 sink 计算重缩放因子
sink_scale = (lse - module.sinks.view(1, -1, 1)).to(torch.float32).sigmoid()
# 3. softmax 之后把输出乘上缩放因子
attn_output = attn_output * sink_scale.unsqueeze(-1).to(attn_output.dtype)

该设计把一个额外可学习 logit 放进 softmax 分母，使得哪怕 query 与 KV 缓存中任意 token 的相关性都很低，注意力分布也不会退化成无法训练的形式——这正是 attention-sink 提升长文本外推与训练稳定性的关键。

注意力后端选型：SDPA 不可用，五选其四

由于 sink 无法通过 torch.nn.functional.scaled_dot_product_attention（SDPA）表达，GraniteMoeSWA 不支持 SDPA，其模型基类显式声明了 _supports_sdpa = False。官方支持的后端矩阵如下：

后端	阶段	说明
"eager"	训练 + 推理	显式实现 sigmoid 重缩放，默认后端
"flex_attention"	训练 + 推理	官方推荐用于训练的后端
"flash_attention_3"	推理（也可训练）	基于 vLLM FA3 hub kernel；当本机未装 FA3 但装有 kernels 时，也会回退到该实现
"flash_attention_4"	推理	由 _compatible_flash_implementations 声明兼容

在模块化实现中，模型声明了兼容的 flash 实现列表 ["kernels-community/vllm-flash-attn3", "flash_attention_4"]。共享注意力分发（attention dispatch）通过把 s_aux=self.sinks 传给 FlexAttention、FlashAttention-3、FlashAttention-4 后端，从而在那些优化内核上以相同的语义应用 sink，保证不同后端之间数学行为一致。

[!NOTE] 在测试套件 test_modeling_granitemoe_swa.py 中，有一个专门的跳过项说明“GraniteMoeSWA 滑窗注意力层与 QuantizedCache 不兼容”，因此 test_generate_with_quant_cache 被跳过。这意味着如果你打算结合量化 KV 缓存（如 QuantizedCache）使用，需要注意兼容性限制。

快速上手：Pipeline 与 AutoModel 两种生成方式

官方推荐的参考检查点为 ibm-granite/granite-swash-3b-a600m。以下两段示例完整保留了原文的用法，可直接运行。

方式一：Pipeline

from transformers import pipeline


pipe = pipeline(
    task="text-generation",
    model="ibm-granite/granite-swash-3b-a600m",
)
pipe("Explain quantum computing in simple terms", max_new_tokens=50)

方式二：AutoModelForCausalLM

from transformers import AutoModelForCausalLM, AutoTokenizer


tokenizer = AutoTokenizer.from_pretrained("ibm-granite/granite-swash-3b-a600m")
model = AutoModelForCausalLM.from_pretrained(
    "ibm-granite/granite-swash-3b-a600m",
    device_map="auto",
    # eager 是默认后端，也支持 "flex_attention"、"flash_attention_3"、"flash_attention_4"
    attn_implementation="eager",
)

inputs = tokenizer("Explain quantum computing in simple terms", return_tensors="pt").to(model.device)
outputs = model.generate(**inputs, max_new_tokens=50)
print(tokenizer.decode(outputs[0], skip_special_tokens=True))

在 CUDA 环境中，慢速集成测试会以 bf16 + attn_implementation="eager" 的方式加载该检查点并逐 token 校验 logits 均值与切片值（见 test_modeling_granitemoe_swa.py 中的 GraniteMoeSWAIntegrationTest），可作为端到端正确性的回归基准。

GraniteMoeSWAConfig 配置全解析

GraniteMoeSWAConfig 定义于 configuration_granitemoe_swa.py。除继承通用 PreTrainedConfig 与 GraniteMoeSharedConfig 的全部字段外，它特有的核心参数如下。

MoE 路由相关

参数	默认值	含义
num_local_experts	8	每层局部专家的总数
num_experts_per_tok	2	每个 token 被路由到的专家数（top-2）
shared_intermediate_size	0	共享专家 FFN 的中间维度；0 表示禁用共享专家，设为正值即可启用
output_router_logits	False	是否额外输出路由器的原始 logits（供辅助损失分析用）
router_aux_loss_coef	0.001	路由器辅助损失的加权系数

关于共享专家的默认值需要特别留意：GraniteMoeSWA 与 GraniteMoeShared 不同——后者默认开启共享专家，而前者默认关闭。从源码可见，GraniteMoeSharedConfig 文档中注释的示例默认共享中间层较大，而两个配置类最终默认字段都落在 shared_intermediate_size = 0；GraniteMoeSWA 的模块化源码 modular_granitemoe_swa.py 更是显式注释“shared experts … disabled by default via shared_intermediate_size=0”。

路由的辅助损失（aux loss）在损失计算处按 loss += self.router_aux_loss_coef * aux_loss 累加（见 modeling_granitemoeshared.py），用于鼓励专家间的负载均衡。

注意力与 RoPE 相关

参数	默认值	含义
sliding_window	128	滑窗注意力层关注的最远历史 token 数
layer_types	None	逐层注意力类型列表，每项为 "full_attention" 或 "sliding_attention"；为 None 时自动生成：第 i 层若 i % 4 == 0 则为 "full_attention"，否则为 "sliding_attention"
layer_rope_theta	None	逐层 RoPE 基频（theta）列表；某项为 0 表示该层采用 NoPE（不施加位置编码）；为 None 时回退为全局 rope_parameters["rope_theta"]
rope_parameters	None	全局 RoPE 参数（包含 rope_theta、rope_scaling 等）
attention_multiplier	1.0	注意力 logits 的缩放系数
attention_bias	False	是否在 QKV/O 投影中使用偏置
attention_dropout	0.0	注意力 dropout 概率

通用结构参数

参数	默认值	含义
vocab_size	32000	词表大小
hidden_size	4096	隐藏维度
intermediate_size	11008	MoE 专家 FFN 的中间维度
num_hidden_layers	32	解码层总数
num_attention_heads	32	注意力头数（num_key_value_heads 为 None 时默认为 32，即不做 GQA 分组压缩）
hidden_act	"silu"	激活函数
max_position_embeddings	2048	最大位置编码长度
initializer_range	0.02	参数初始化范围
rms_norm_eps	1e-6	RMSNorm 的 epsilon
embedding_multiplier / logits_scaling / residual_multiplier	1.0	嵌入乘子、logits 缩放除数、残差乘子（IBM 训练技巧）
tie_word_embeddings	False	是否共享输入/输出词嵌入
use_cache	True	是否使用 KV 缓存加速自回归生成

layer_types 与逐层 RoPE 的运行时行为

在 __post_init__ 中，若未显式提供 layer_types，配置会按 num_hidden_layers 自动铺开默认模式；若 num_key_value_heads 为空则补齐等于 num_attention_heads；若 layer_rope_theta 为空则整层统一为全局 rope_theta。

测试基类 GraniteMoeSWAModelTester 的注释印证了这一点：默认 num_hidden_layers=2 时，layer_types 会被解析成 ["full_attention", "sliding_attention"]，从而在单测中同时覆盖两条注意力路径。

模型主体对这两种逐层特性做了专门的工程处理（见 modular_granitemoe_swa.py 中的 GraniteMoeSWAModel）：

• 掩码按类型分组复用：前向时只按“全注意力 / 滑窗注意力”两类各自创建一次因果掩码（create_causal_mask 与 create_sliding_window_causal_mask），再按各层 layer_types[i] 索引复用，避免每层重复建掩码的开销；
• RoPE 按去重后的 theta 只算一次：self.rotary_embs 只对 layer_rope_theta 中出现的每个非零 theta 各建一份旋转嵌入；前向时按 theta 缓存 (cos, sin)，theta 为 0（NoPE）的层则拿到 None 并跳过旋转位置编码。

配置对象快速初始化

>>> from transformers import GraniteMoeSWAModel, GraniteMoeSWAConfig

>>> # 初始化一份 GraniteMoeSWA 配置
>>> configuration = GraniteMoeSWAConfig()

>>> # 由配置初始化模型
>>> model = GraniteMoeSWAModel(configuration)

>>> # 读取模型配置
>>> configuration = model.config

如需定制，例如启用共享专家并把滑窗加宽，可以这样写：

from transformers import GraniteMoeSWAConfig

config = GraniteMoeSWAConfig(
    num_hidden_layers=16,
    num_local_experts=8,
    num_experts_per_tok=2,
    shared_intermediate_size=2048,   # 默认 0（禁用），设为正值启用共享专家
    sliding_window=256,              # 滑窗注意力层关注最近 256 个 token
    # layer_types 未提供时默认：第 0,4,8,12 层为全注意力，其余为滑窗注意力
)

可用类 API 一览

GraniteMoeSWAConfig、GraniteMoeSWAModel、GraniteMoeSWAForCausalLM 均已通过 granitemoe_swa 包 对外导出，也可直接 from transformers import ... 导入：

• GraniteMoeSWAConfig：配置类，model_type = "granitemoe_swa"，其 forward 相关的输出对象为 MoeModelOutputWithPast（MoE 变体的 BaseModelOutputWithPast）；
• GraniteMoeSWAModel：裸解码器主干，输出 last_hidden_state 与 past_key_values；
• GraniteMoeSWAForCausalLM：带语言建模头（lm_head）与 GenerationMixin 的生成接口。

模型基类 GraniteMoeSWAPreTrainedModel 中值得注意的实现细节：

• _no_split_modules = ["GraniteMoeSWADecoderLayer"]，供 device_map="auto" 做层级切分；
• _supports_sdpa = False，即前文所述的 SDPA 不支持；
• 权重初始化时会把注意力头的 sinks 参数显式置零（init.zeros_(module.sinks)），保证训练起点等价于标准的 softmax 注意力。

分布式并行：Tensor Parallel 与 Expert Parallel 的切分蓝图

配置类内置了两套并行切分方案（对应 Transformers 的 tensor parallelism / expert parallelism 基础设施）：

张量并行 base_model_tp_plan（与 Granite 一致地切分注意力，逐头 sinks 随列并行走，路由专家走 TP 切分）：

• q_proj / k_proj / v_proj 与 sinks：colwise；
• o_proj：rowwise；
• 路由专家 gate_up_proj：packed_colwise（gate/up 打包按列切分）；
• down_proj：rowwise；
• block_sparse_moe.experts：moe_tp_experts；
• 路由器（router）保持复制；可选的共享专家（shared_mlp，默认关闭）也保持复制——因为它体量小，且其完整输出与 all-reduce 后的 MoE 输出能一致地相加。

专家并行 base_model_ep_plan：把被路由的专家切分到不同 rank 上，每个 rank 持有专家切片，由路由器驱动 dispatch：

• 路由器：ep_router；
• gate_up_proj / down_proj：grouped_gemm；
• 专家集：moe_tp_experts；共享专家仍保持复制。

GraniteMoeSWA 的 GraniteMoeSWATopKRouter 与 GraniteMoeSWAMoE 相对于旧版 GraniteMoe 唯一的差异就是返回值顺序（router_logits, router_scores, router_indices），这是为了启用 EP（专家并行）而做的调整；源码中的 TODO 注释也提到未来会把旧 granitemoe 系列一并重构以支持 EP。

小结

GraniteMoeSWA 用三行“配方”概括：8 个局部专家的 top-2 稀疏路由 + 每 4 层一次的滑窗注意力 + 每头一个可学习注意力 sink。它把稀疏专家带来的参数效率、滑窗带来的长序列开销控制与 attention-sink 带来的长程稳定性合并在一个解码器里。使用时务必记住两点：SDPA 不可用，请从 eager / flex_attention / flash_attention_3 / flash_attention_4 中选择与硬件匹配的后端；共享专家默认关闭，需要时通过 shared_intermediate_size 显式开启。若想深入了解其两个“父架构”，可继续阅读 GraniteMoeShared 文档 与 GraniteSWA 文档，以及两者的建模源码 modeling_granitemoeshared.py 和 modeling_granite_swa.py。