SparkLM 模型架构文档

SparkLM 模型架构文档

SparkLM-31M · Dense + MoE + MTP + MLA/GQA/SAHA + ECKVC + CRS 全组件架构总览

---

目录

1. 图 1 · 整体架构总览
2. 图 2 · MLA (Multi-head Latent Attention) 内部结构
3. 图 3 · HybridMoE Layer + CRS + ECKVC 协同机制
4. 图 4 · MTP (Multi-Token Prediction) 递归结构
5. 图 5 · SAHA 分层注意力选择策略
6. 图 6 · Transformer Block 内部结构
7. SparkLM 五大原创创新模块
8. SparkLM-31M 核心配置参数

---

图 1 · 整体架构总览

SparkLM 整体采用 Pre-Norm Transformer 架构，核心前向数据流为：

Input Token IDs → Token Embedding → RoPE → Transformer Backbone (×12 层) → Final RMSNorm → [ECKVC] → LM Head → Logits → CrossEntropy Loss

其中 Transformer Block 内部支持 Dense FFN 与 HybridMoE 两种前馈网络切换；MTP 模块作为可选的训练辅助分支，从主模型隐藏状态出发递归预测未来 token；MoE 的 Aux Loss 与 MTP Loss 共同加入总损失。

Total Loss = CrossEntropy + α · MTP_Loss + Aux_Loss

图例说明：

颜色	模块
🔵 蓝色	输入 / Embedding
🟡 黄色	注意力模块 (MLA / GQA / SAHA / CSA / HCA)
🟢 绿色	Dense FFN (SwiGLU)
🟣 紫色	MoE Layer (HybridMoE V4)
🟪 浅紫	MTP 模块 (多 token 预测)
🩷 粉色	输出 / Loss
⬜ 灰色	RMSNorm
🩷粉色虚线	创新模块 (ECKVC / CRS / MTP)

---

图 2 · MLA (Multi-head Latent Attention) 内部结构

MLA (Multi-head Latent Attention) 是 DeepSeek V4 的核心注意力机制，通过低秩压缩大幅减少 KV Cache 的内存占用。

核心数据流

Q 路径： hidden_states → q_down_proj (768→384) → RMSNorm (q_norm) → 分裂为 q_nope + q_rope → q_up_proj / q_rope_proj → RoPE(q_rope)

KV 路径： hidden_states → kv_down_proj (768→192+64) → RMSNorm (kv_norm) → kv_up_proj (192→12×128) + k_rope(64) → RoPE(k_rope)

注意力计算： Concat(q_nope | q_rope_emb) · Concat(k_nope | k_rope_emb)^T / √128 → SDPA / Flash Attention 2 → o_proj (12×64→768)

KV Cache 压缩效果

方案	KV Cache 维度/层	压缩率
标准 MHA	12×(64+64) = 1536	基准
MLA	192 + 64 = 256	83.3% ↓

ECKVC 增强（创新 #4）

在 MLA 的 KV 压缩隐空间上，ECKVC 利用 MoE 路由权重进行条件化调制：

c_kv' = c_kv × (1 + γ(w_e)) + β(w_e)

其中 γ = scale_proj(w_e)、β = shift_proj(w_e)，零初始化渐进学习。

---

图 3 · HybridMoE Layer + CRS + ECKVC 协同机制

HybridMoELayer 是 SparkLM 的混合专家前馈层，融合了 DeepSeek V4 的 Shared Expert 架构与多项原创创新。

架构组成

1. Shared Expert (SwiGLU FFN)：始终激活，处理公共语义。gate_up_proj: 768→2×2048 → SiLU(gate)×up → down_proj: 2048→768
2. MoEGateV4 (Router)：gate: 768→8 → SqrtSoftplus / Sigmoid → Softmax → scores (B·S, 8)
3. CRS (创新 #5)：因果 EMA 平滑路由分数，smoothed = (1-α)·logit + α·EMA(logit[0:t])，α=0.1
4. Top-K 选择：Top-K (k=2/8) + Normalize
5. Aux-Loss-Free 偏置更新：speed=0.001，无需辅助损失即可平衡专家负载
6. MoEFusedExperts (8 × SwiGLU)：sort_moe_routed_tokens → forward_grouped，每 token 仅激活 Top-2 专家
7. ECKVC (创新 #4)：γ = 1 + scale_proj(w_e)、β = shift_proj(w_e)、h' = h × γ + β，零初始化

混合路由策略

层位置	路由方式
layer_idx < 3	Hash Routing（确定性哈希）
layer_idx ≥ 3	Learned Routing（可学习路由 + CRS）

输出聚合

output = shared_expert(h) + Σ(router_weights_i × routed_expert_i(h))

---

图 4 · MTP (Multi-Token Prediction) 递归结构

MTP 是 DeepSeek V4 风格的多 token 预测模块，作为训练辅助分支递归预测未来 token，推理时可丢弃。

递归结构

每一层 MTP 接收上一层的隐藏状态 h_{i-1} 和下一个 token 的 embedding，经过独立的 Transformer Block 后输出 h_i 并预测 token_{t+i+1}：

h_i = Block(RMSNorm(h_{i-1}) + Embed(token_{t+i}))
logits_i = LM_Head(RMSNorm(h_i))

序列长度逐步缩短：S → S-1 → S-2 → ... → S-N

损失计算

# compute_mtp_loss()
logits.float()  # fp32 防 fp16 溢出
depth_weights = [w_0, w_1, ..., w_{N-1}]  # cosine decay
total_mtp_loss = Σ w_i · CE_i / Σ w_i
# exclude_depth_0: skip overlap with main head

关键设计要点

#	设计决策
1	每层独立 Transformer Block
2	递归传播: h_i → h_
3	序列逐步缩短: S → S-1 → S-2
4	无 zero-padding, 自然截断
5	logits.float() 防 fp16 溢出（Bug 修复）
6	MTP block 不参与 hybrid attn
7	深度权重: cosine decay
8	loss_weight=0.3（主 loss）
9	num_mtp_tokens=1 (31M 配置)
10	推理时可丢弃 (训练辅助)

Embedding 共享

• embed_tokens 与主模型共享
• mtp_share_lm_head: 可选共享 LM Head（31M 默认不共享）

---

图 5 · SAHA 分层注意力选择策略

SAHA (Scale-Aware Hybrid Attention) 是 SparkLM 的原创创新 #3，根据模型规模自动为每层选择最优注意力机制。

三档自适应策略

模型规模	注意力配置	KV Cache/层	设计理由
Small (< 100M)	全部 GQA	4×64 = 256 维	信息损失最小，31M 最优
Medium (100M–1B)	前 60% GQA + 后 40% SharedK	渐变	SharedK 无低秩瓶颈，KV 压缩 75%
Large (> 1B)	前 40% GQA + 后 60% MLA	256 维	大模型可承受低秩分解，KV 效率最优

各注意力机制详解

GQA (Grouped Query Attention)

• Q: 12 heads, K/V: 4 groups
• 每 group 共享 K/V → repeat_kv
• head_dim=64, scale=1/√64
• KV cache: 4×64 = 256 维/层
• 信息损失最小

SharedK Attention（本文提出）

• Q: 12 heads（独立）
• K: 1 head（全局共享！）
• V: 12 heads（独立）
• KV cache: 1×64 = 64 维/层
• 强制全局一致注意力模式
• 无低秩瓶颈，KV 压缩 75%

MLA (Multi-head Latent Attention)

• kv_down: 768 → 192+64（压缩）
• kv_up: 192 → 12×128
• KV cache: 256 维/层（83%↓）
• 低秩分解，大模型可承受
• q_nope + q_rope 分离设计

SparkLM-31M 实验结论

• SAHA=small → 全 GQA 最优（eval=1.5277）
• MLA 在 31M 有信息瓶颈（低秩压缩在极小模型上损失过多信息）
• SharedK 在中等模型有潜力

---

图 6 · Transformer Block 内部结构（Pre-Norm 残差）

SparkLMBlock 是 SparkLM 的基础构建单元，采用 Pre-Norm 残差结构，统一支持 Dense/MoE 切换。

前向流程

hidden_states
  ├──────────────────────────────────────────→ (residual 1)
  ↓
RMSNorm (input_layernorm)
  ↓
[ECKVC Layer] (optional, MoE only)
  ↓
Attention Module (MLA / GQA / SharedK / CSA / HCA)
  ↓
⊕ ←────────────── (residual 1 add)
  ├──────────────────────────────────────────→ (residual 2)
  ↓
RMSNorm (post_attention)
  ↓
SwiGLU FFN (Dense)  OR  HybridMoELayer (MoE V4 + CRS)
  ↓
⊕ ←────────────── (residual 2 add)
  ↓
output (B, S, 768)

Gradient Checkpoint

特性	效果
use_gradient_checkpointing	checkpoint(block, h, cos, sin)
显存节省	~40%
吞吐损失	<5%

注意力模块可切换

SparkLMBlock 内部的 Attention Module 支持 4 种注意力变体：

变体	特点
MLA	低秩 KV 压缩，大模型最优
GQA	分组 KV，小模型信息损失最小
SharedK	全局共享 K，中等模型潜力
CSA/HCA	混合注意力

---

SparkLM 五大原创创新模块

#	模块名称	类型	核心机制	参数开销	31M 效果
#1	SparkLMBlock	Core	Pre-Norm 残差结构，统一支持 Dense/MoE FFN 切换；Gradient Checkpoint 可选	—	✓ 基础架构
#2	MoE V4	MoE	Shared Expert + Top-K Routed Experts；SqrtSoftplus 评分 + Aux-Loss-Free 偏置更新；Hash Routing 前几层	+8× SwiGLU experts	✗ 过拟合
#3	SAHA	Attention	根据模型规模自动选择每层注意力：Small→GQA, Medium→GQA+SharedK, Large→GQA+MLA	0 (配置选择)	✓ 参数效率最优
#4	ECKVC	MoE×Attn	用 MoE 路由权重调制 MLA KV 压缩隐空间：c_kv × (1+γ(w_e)) + β(w_e)，零初始化渐进学习	2×E×rank ≈ 3K	— 未单独消融
#5	CRS	MoE	因果 EMA 平滑 MoE 路由分数，减少相邻 token 路由跳变；可学习 α + cosine warmup + 并行 cumsum	0 (无参数)	✗ eval 无提升

创新模块详细说明

#1 SparkLMBlock — 统一 Transformer Block

• Pre-Norm 残差架构（input_layernorm + post_attention_layernorm）
• Dense/MoE 前馈网络运行时切换（use_moe 标志）
• 可选 Gradient Checkpointing（节省 ~40% 显存，吞吐损失 <5%）
• 支持 4 种注意力变体运行时注入

#2 MoE V4 — DeepSeek V4 完整工程化

• Shared Expert：始终激活的 SwiGLU FFN，处理公共语义
• Routed Experts：8 个 SwiGLU 专家，每 token 激活 Top-2
• MoEGateV4：gate → SqrtSoftplus → Softmax → Top-K + Normalize
• Aux-Loss-Free：偏置更新机制（speed=0.001），无需辅助平衡损失
• Hash Routing：前 3 层使用确定性哈希路由，后续层使用可学习路由
• FusedExperts：sort_moe_routed_tokens → forward_grouped 批量执行

#3 SAHA — 尺度感知混合注意力

• 三档自适应策略：Small / Medium / Large
• 小模型全 GQA → 信息损失最小
• 中模型 GQA + SharedK → 无低秩瓶颈
• 大模型 GQA + MLA → KV 效率最优
• 零参数开销（纯配置选择）

#4 ECKVC — 专家条件化 KV 压缩

• 利用 MoE 路由权重 w_e 调制 MLA 的 KV 压缩隐空间 c_kv
• 仿射变换：c_kv' = c_kv × (1 + γ(w_e)) + β(w_e)
• γ = scale_proj(w_e)、β = shift_proj(w_e)
• 零初始化：训练初期等效于无 ECKVC，渐进学习
• 参数开销极小：2×E×rank ≈ 3K

#5 CRS — 因果路由平滑

• 因果 EMA 平滑：smoothed = (1-α)·logit + α·EMA(logit[0:t])
• 减少相邻 token 的路由跳变，提升路由一致性
• 可学习 α + cosine warmup 调度
• 并行 cumsum 实现因果 EMA
• 零参数（纯算法创新）

---

SparkLM-31M 核心配置参数

基础配置

参数	值
hidden_size	768
num_hidden_layers	12
num_attention_heads	12
head_dim	64
intermediate_size	2048
vocab_size	6400
max_seq_len	1024
total params	~31M

注意力配置

参数	值
attention_type	gqa (SAHA auto)
num_kv_heads (GQA)	4
MLAkv_lora_rank	192
MLAq_lora_rank	384
use_qk_norm	True
use_flash_attn_2	True
RoPE base	10000
SharedK K groups	1

MoE 配置

参数	值
use_moe	False (default)
num_experts	8
num_experts_per_tok	2
moe_intermediate_size	1024
hash_routing_layers	3
router_score_function	sqrt_softplus
CRS alpha	0.1
ECKVC bottleneck	0.25×hidden

MTP 配置

参数	值
num_mtp_tokens	1
mtp_loss_weight	0.3
mtp_share_lm_head	False (default)
mtp_block_type	Dense (no MoE)
depth_weights	[1.0]
exclude_depth_0	True
logits upcast	fp32 (防 NaN)

---

31M 消融实验关键结论

发现	结论
Dense vs MoE	Dense 全面优于 MoE（31M 规模 MoE 严重过拟合）
MTP	拖累性能（额外计算开销未带来 eval 收益）
SAHA	参数效率最高（全 GQA 在小模型最优）
MLA	低秩压缩在 31M 造成信息瓶颈
Gradient Checkpoint	近乎免费（显存节省 ~40%，吞吐损失 <5%）
最优配置	train_adamw（eval=1.5277）

---

本文档基于 SparkLM-31M 模型源码（~49K 行）和 35 配置消融实验数据生成。