【轻量化复刻大模型 / KIMI K3】MoE (Stable LatentMoE) 交互式演示文档
0x01 实现概述
注：
- 项目github.com/RossiLH/Mod… 和系列博客是我的一个尝试，也是我学习的过程。不确定是否满足读者的需求，我自己会反思改进，也希望大家提出宝贵意见。
- 涉及 KIMI K3，DeepSeek V4，PI0.5，PI0.7 等模型。后续可能会跟踪新模型。
- 全 CPU 复刻，不涉及 GPU，也不涉及大规模训练（这是另外一个庞杂深邃的领域）和硬件。
- 模型的组件尽量是原始大模型的缩小版，如果组件没有开源版本，则用近似组件替代。
- 会解释每个组件在解决一个什么问题、这个问题难在哪、为什么只能这么解、解法的原理是什么，以及组件之间如何咬合成一个整体。重点不是记名词，而是理解一条主线。
- 每个主要组件都有可直接运行的 Gradio demo——无需 GPU。另有若干可视化脚本（
plot_*.py），可生成静态图。 - 所有文件行号基于本地源码快照；源码后续变化时，行号可能移动，但函数名和调用关系仍可作为定位入口。
- 为了保证每个文章可以尽量独立成文，所以有些内容会在不同文章中都存在，可能在系列层面上存在冗余，还请大家理解。
1.1 模块范围
本 Gradio 演示覆盖 K3 Stable LatentMoE 模块的全部核心机制，对应论文 §2.3 与附录 B/C/D，共 12 个交互 Tab：
| # | Tab 名称 | 验证对象 | 核心命题 | 代码来源 |
|---|---|---|---|---|
| 1 | 原理总览 | 架构直觉 | width 分离 + shared/routed 双路径 + 三处 Stable 修复 | plot_moe_principle.py |
| 2 | Demo 四联图 | 综合演示 | QB 均衡 / 直方图精度 / 激活有界 / RMSNorm 稳定 | demo_stable_latent_moe.py |
| 3 | 前向契约 | 形状与 Eq.11 | 输出形状、路由契约、独立参考实现逐点一致 | test_moe.py |
| 4 | 潜空间 width 分离 | 瓶颈维度 | 路由专家在 l 维、共享专家在 d 维 | test_moe.py |
| 5 | RMSNorm 稳定性 | §2.3.1 | 放大专家权重不改变路由分支输出尺度 | test_moe.py |
| 6 | 梯度稀疏性 | 稀疏性 | 只有被选中的专家获得梯度 | test_moe.py |
| 7 | Dispatch 正确性 | 分组执行 | grouped 与 dense 逐点一致，且置换不变 | test_moe.py |
| 8 | QB 语义 | §2.3.3 语义 | cutoff / 偏置只导流不改权重 / 权重归一化 | test_router_qb.py |
| 9 | QB 均衡效果 | 收敛性 | 一步改善 / 单调收敛 / 直方图精度 | test_router_qb.py |
| 10 | QB 工程行为 | 工程约定 | 延迟更新 / mean-centring / 推理冻结 / 开关 | test_router_qb.py |
| 11 | 896 专家压力测试 | 论文规模 | QB 在近千专家规模下仍一步见效 | test_router_qb.py |
| 12 | SiTU-GLU | §2.3.2 / 附录 B | 有界 + 光滑可导 + 原点匹配 SwiGLU | test_situ_glu.py |
1.2 实现思路
复用测试层的 measure 函数，零侵入构建可视化层：
-
每个模块的测试文件都导出
measure_*纯函数，返回结构化 dict（数值 + 布尔 + 列表），不做 assert；同时服务 pytest 单元测试和演示层可视化。 -
Gradio UI 层只做"取数 -> 画图 -> 写结论文本"三件事，不重复实现任何数学逻辑，保证演示数据与测试数据同宗同源。
-
左控制 / 右显示的统一布局：每个 Tab 左侧是参数滑块 + 运行按钮，右侧是 gr.Plot + gr.Textbox（结论）。点击"运行"后，滑块参数传入
tab_*函数，调用对应的measure_*或 plot 脚本，返回(fig, text)。 -
before / after 对比：每个创新组件都有标准基线对照——
- LatentMoE vs 全宽 MoE（Tab 1/4）
- 有 RMSNorm vs 无 RMSNorm（Tab 5）
- SiTU-GLU vs SwiGLU（Tab 12）
- QB 前 vs QB 后（Tab 9/11）
- grouped dispatch vs dense 评估（Tab 7）
- histogram 近似 vs exact 分位数（Tab 9）
1.3 运行方式
cd Model-Viz/experiments/models/llm/kimi_k3/modules/moe
# 方式 1：完整交互式演示（12 个 Tab）
python gradio_moe_demo.py
# 方式 2：单出图脚本（独立运行）
python plot_moe_principle.py
python demo_stable_latent_moe.py
所有 Tab 均在 CPU 上实时计算，默认 2~4 线程，无需 GPU。
0x02 核心数学
2.1 总体公式
Stable LatentMoE 的核心定义（论文 Eq. 11）：
# ===== 输入输出维度 =====
# x in R^{b x t x d} # 输入 token 序列 (batch, seq_len, model_dim)
# y in R^{b x t x d} # 输出（与输入同形）
# ===== 路由决策 (QuantileBalancingRouter, Eq.13 / §2.3.3) =====
# s_i = sigmoid(W_r . x_i) # 每个 token 的 per-expert 原始分数
# T_i = topk(s_i + b, k) # Top-k 选中集, 偏置 b 只参与选择
# p_ij = s_ij / sum_{r in T_i} s_ir # 混合权重（不含偏置, Eq.13 的权重）
# a_i = topk(s_i + b, k+1)[k] # cutoff = 第 k+1 个 biased score
# ===== 潜空间路由分支 (width 分离, Eq.11 第一项) =====
# z_i = W_down . x_i # W_down: d -> l, 降维到潜空间
# u_i = Σ_{j in T_i} p_ij . E_j^routed(z_i) # E_j: SiTU-GLU(l, H, l), 在 l 维工作
# u_i = RMSNorm(u_i) # §2.3.1 稳定化
# y_routed_i = W_up . u_i # W_up: l -> d, 升维回模型空间
# ===== 共享分支 (全宽, Eq.11 第二项) =====
# y_shared_i = Σ_j E_j^shared(x_i) # E_j: SiTU-GLU(d, H, d), 在 d 维工作
# ===== 最终输出 (Eq. 11) =====
# y_i = y_shared_i + y_routed_i # 残差式合并: 共享基底 + 路由增强
# ===== SiTU-GLU 激活 (Eq.12 / 附录 B) =====
# softcap(z, β) = β . tanh(z / β) # 光滑软截断, |·| ≤ β
# SiTU-GLU(g, u; β_g, β_u)
# = softcap(g, β_g) . sigmoid(g) . softcap(u, β_u) # 门分支（有界 Swish）× up 分支（有界线性）
# 且 |SiTU-GLU| ≤ β_g . β_u # 全局有界
# ===== Quantile Balancing 更新 (Eq.14 / 附录 C-D) =====
# r_ij = a_i - s_ij # 把专家 j 推到 cutoff 所需偏置
# b_hat_j = quantile(r_·j, k/n) # 第 k/n 分位 = 刚好让 q 个 token 越 cutoff
# b_j = b_hat_j - mean(b_hat) # mean-centring, sum(b) = 0
关键洞察：Stable LatentMoE 的三处「Stable」修复——RMSNorm 去尺度敏感性、SiTU-GLU 去激活离群、Quantile Balancing 去负载失衡——每一处都对应一个独立的数值失稳模式，合起来才让 896 专家 × 16 激活的极端稀疏配置在 2.8T 规模上可训练。少任何一处，病态链都会在某个尺度上爆炸。
”
2.2 子机制分解
2.2.1 Width 分离（Latent Bottleneck）
公式：
routed 参数 = E_routed × (2·l·H + H·l) = 3·E_routed·l·H # 在 l 维工作
shared 参数 = E_shared × (2·d·H + H·d) = 3·E_shared·d·H # 在 d 维工作
投影参数 = 2·d·l # W_down + W_up
每 token FLOPs = E_shared·3dH + top_k·3lH + 2dl # 只激活 top_k 个路由专家
直觉解释：把路由专家塞进更窄的潜空间，让「专家数量多」与「每个专家便宜」解耦——896 个专家听起来很夸张，但每个只在 3584 维而非 4096 维工作，叠加 top_k=16 的稀疏性，实际计算量远低于全宽稠密。
关键性质：
- 参数量节省 ≈ (1 - l/d) × 路由分支占比（K3 约 12.4% 参数、6.8% FLOPs）
- 稀疏度 = n_routed / top_k（K3 = 56）
- shared 路径始终稠密、始终有梯度，是训练的稳定基底
2.2.2 RMSNorm 稳定化（§2.3.1）
公式：
u' = RMSNorm(u) = u / RMS(u) × γ # γ 为可学习增益
y = W_up · u'
直觉解释：路由分支近似串了 4 个矩阵乘（W_down → 专家 FFN 两步 → W_up），这是一条病态链——专家权重的幅度、选了哪些专家、权重分配如何，都会放大成输出尺度的波动。RMSNorm 在聚合之后、上投影之前把尺度归一化，让下游永远接到单位 RMS 的信号。
关键性质：
- 尺度不敏感：experts.down 权重放大 k 倍，RMSNorm 输出不变（Tab 5 验证 ratio ≈ 1.0）
- 非可选：论文指出它本身也改善验证损失
- 与 shared 路径无关：shared 是全宽 SiTU-GLU，本身已经尺度可控
2.2.3 SiTU-GLU 有界激活（§2.3.2 / 附录 B）
公式：
softcap(z, β) = β · tanh(z/β) # 光滑截断, |z|→∞ 时 → ±β
SiTU-GLU(g, u) = softcap(g, β_g) · σ(g) · softcap(u, β_u)
直觉解释：SwiGLU 的 gate × up 两个无界因子相乘，MoE 宽度又大，低精度下偶尔会撞上「两个坐标都很大」的巧合，产生离群点溢出 FP8。SiTU-GLU 给两边都套 softcap，乘积天然不超过 （K3 取 4×25=100，小于 FP8 E4M3 上限 448）。
关键性质：
- 全局有界：|SiTU-GLU| ≤ β_g·β_u（Tab 12 验证，输入再大也不超过上界）
- 原点匹配 SwiGLU：tanh(z) ≈ z - z³/(3β²)，小输入下与 SwiGLU 一阶一致
- 光滑可导：softcap 用 tanh 而非硬 clamp，饱和区仍有 tanh' > 0 的梯度
- 负尾部归零：σ(g) 让负输入趋近 0，保留 Swish 的稀疏性
2.2.4 Quantile Balancing（§2.3.3 / 附录 C-D）
公式：
目标负载 q = m·k / n # m tokens 每 token 选 k 个, 均分到 n 专家
r_ij = a_i - s_ij # 专家 j 达到 token i cutoff 所需偏置
b_hat_j = quantile(r_·j, k/n) # 第 q 个 required 值 = 刚好让 q 个 token 越 cutoff
b_j = b_hat_j - mean(b_hat) # mean-centring, sum(b) = 0
直觉解释：传统的「b ← b + γ·sign(target - load)」像用固定步长爬山，步长是超参，收敛慢。QB 换个角度想：既然 cutoff 固定，「要让专家 j 恰好吃到 q 个 token」等价于找一个偏置值，使得刚好有 q 个 token 的 s_ij + b_j ≥ a_i。这个值就是 required 的第 q 个分位数——一步到位，无需学习率。
关键性质：
- 一步收敛：单次 update_bias 就能把几十倍的失衡压到 2 倍以内（Tab 9 验证 51x → 1.95x）
- 单调收敛：迭代步数越多越均衡（Tab 9 验证）
- 无损失偏置：偏置只进 Top-k 选择、不进权重归一化，不扭曲路由器优化
- 直方图近似：大规模下用分箱直方图估分位数，误差 ≤ 1 个 bin 宽（Tab 9 验证 0.43 bin）
- 延迟更新：本批从不用本批算出的偏置路由自己，避免统计自洽性问题
- mean-centring 不动点：整体平移偏置不改变路由选择（Tab 10 验证）
0x03 每个 Tab 的详细说明
Tab 1：原理总览
目的
展示 Stable LatentMoE 的架构全景——数据流、width 分离节省的参数量/FLOPs、稀疏度与计算量的关系、以及层调度方式。建立整体直觉，为后续 11 个细节 Tab 做铺垫。
验证的是论文 §2.3 开篇的架构描述（Eqs. 11-14 的整体图景），不做数值断言，属于「一图看懂」级别的总览。
UI 组件与参数影响
| 参数 | 范围 | 默认 | 调节效果 |
|---|---|---|---|
| - | - | - | 无参数，点击按钮直接运行 |
此 Tab 无交互参数，固定使用 K3 报告的 2.8T 配置（d=4096, l=3584, H=3072, E=896, k=16, n_shared=2）。
输出图解读
tab-1
2×2 四联图布局。
左上 (a)：架构数据流图
- 图类型：方块箭头示意图
- 含义：展示 token x (d 维) 分流到两条路径——上支路为 shared experts（全宽 SiTU-GLU，n_shared=2），下支路为 W_down 降维 → routed experts（Top-k 路由 + SiTU-GLU，896 专家/激活 16/稀疏度 56） → RMSNorm(latent) → W_up 升维；两路最后在 d 维相加得到 y
- 颜色：蓝色=输入输出，橙色=投影，红色=路由分支，绿色=共享分支，紫色=RMSNorm
- 你该看什么：确认 width 分离的物理位置（W_down 之后才进入 routed experts）、三处 Stable 修复各自对应哪个环节
- 理想情况：图结构与 Eq. 11 逐字对应
- 异常指示：如果缺少 RMSNorm 模块，说明 rmsnorm_before_up 配置被关闭
右上 (b)：width 分离参数量 / 计算量节省
- 图类型：分组柱状图
- X 轴：参数 (B)、每 token FLOPs (M)
- Y 轴：绝对数量（线性）
- 颜色：红色 = LatentMoE，灰色 = 全宽 MoE（基线）
- 你该看什么：对比红色柱与灰色柱的高度差，体会潜空间 + 稀疏的累计节省
- 理想情况：LatentMoE 的两柱都显著低于全宽基线（参数节省 ~12.4%，FLOPs 节省 ~6.8%）
- 异常指示：如果两柱等高，说明 l=d（width 分离未生效）
左下 (c)：稀疏度 vs 计算量占比
- 图类型：双折线图，X 轴对数刻度
- X 轴：稀疏度 (n_routed / top_k)，log scale
- Y 轴：占全宽 MoE 的比例 (%)，线性
- 颜色：红色实线 = FLOPs，紫色虚线 = 参数；橙色竖线 = K3 配置（稀疏度=56）
- 你该看什么：随着稀疏度升高（top_k 减少），FLOPs 和参数量占比如何单调下降；K3 配置（56）在曲线的哪个位置
- 理想情况：两条曲线均单调递减，稀疏度越大越省
- 异常指示：若曲线呈平台或上升，说明稀疏度公式有误
右下 (d)：K3 层调度（dense warm-up + MoE body）
- 图类型：堆叠柱状图
- X 轴：层编号（0~92）
- Y 轴：层类型激活状态（0/1）
- 颜色：蓝色 = DenseFFN（前 2 层 warm-up），红色 = StableLatentMoE（后续 91 层主体）
- 你该看什么：前 n_dense_layers=2 层是全宽稠密 FFN，过渡到 MoE 主体
- 理想情况：第 0-1 层为蓝色（稠密），第 2-92 层为红色（MoE）
- 异常指示：若所有层都是同色，说明层调度配置错误
为什么重要
此 Tab 是整个演示的导航图——读者对 LatentMoE 的整体直觉都从这里建立。后续 11 个 Tab 的细节验证，都需要回溯到这张全景图来定位：width 分离在哪、RMSNorm 在哪、QB 在哪、SiTU-GLU 在哪。没有这张图，读者看到后面的数值验证会缺乏上下文。
Tab 2：Demo 四联图
目的
用最小的配置一次性展示 Stable LatentMoE 的四个标志性效果：QB 负载均衡前后对比、直方图分位数估计精度、SiTU-GLU 有界性 vs SwiGLU 发散、RMSNorm 的尺度不敏感性。这是整个模块的「快速概览」四联图，对应 demo_stable_latent_moe.py 的默认演示。
验证的是论文 §2.3 的四个核心论断（§2.3.1 RMSNorm、§2.3.2 SiTU-GLU、§2.3.3 QB）在同一微型配置下的综合表现。
UI 组件与参数影响
| 参数 | 范围 | 默认 | 调节效果 |
|---|---|---|---|
| n_experts | 4-64 | 8 | 增大 → 专家更多，初始失衡更严重，QB 改善幅度可能更大 |
| top_k | 1-8 | 2 | 增大 → 每 token 选更多专家，稀疏度下降，负载更天然均匀 |
| n_batches | 10-100 | 40 | 增大 → 统计样本更多，直方图估计更精确，误差更小 |
| batch | 100-1000 | 200 | 增大 → 同 n_batches，样本量增加 |
| bins | 128-1024 | 256 | 增大 → 直方图分箱更细，分位数估计误差更小（以内存为代价） |
| hot_experts | 字符串 | "0,1" | 调整热门专家 id（逗号分隔），决定初始失衡的来源 |
hot_experts 对应的专家权重行被放大 6 倍，使其天然获得更多路由；这些专家就是 QB 要降温的对象。
注意：子图 (c)(d)（SiTU-GLU 曲线 / RMSNorm 缩放）使用固定的内部配置，不响应 n_experts / top_k / hot_experts 等路由参数的调节；这两个子图展示的是激活函数与归一化本身的性质。
”
输出图解读
tab-2
2×2 四联图布局。
左上 (a)：Quantile Balancing 前后负载
- 图类型：分组柱状图
- X 轴：expert id（0 到 n_experts-1）
- Y 轴：token 被路由次数
- 颜色：蓝色 = before QB，橙色 = after QB
- 你该看什么：before 中热门专家（如 0、1）的柱子明显高于其他专家；after 后所有柱子高度接近一致
- 理想情况：after 的 max/min ratio ≈ 1.5 左右（默认配置下从 ~3.7 降到 ~1.5）
- 异常指示：after 仍有专家柱子显著高于其他 → QB 未生效或样本不足
右上 (b)：分位数估计 exact vs histogram
- 图类型：双折线对比图
- X 轴：expert id
- Y 轴：偏置值 b
- 颜色/线型：蓝色实线圆点 = exact 分位数（torch.quantile 真值），橙色虚线方块 = histogram 估计
- 你该看什么：两条线几乎重合，说明直方图近似的精度损失很小
- 理想情况：max diff < 0.01（默认 256 bins 下约 0.003），且 diff 以 bin 宽计 < 1.5
- 异常指示：两条线明显分叉 → 直方图实现有 bug 或 bins 太少
左下 (c)：SwiGLU vs SiTU-GLU 激活幅值
- 图类型：双折线图，X 轴对数刻度
- X 轴：输入 scale（0.5 → 8），log scale
- Y 轴：max|activation|（线性）
- 颜色：蓝色 = SwiGLU（无界），橙色 = SiTU-GLU（有界）；红色虚线 = 上界 beta_g·beta_u = 100
- 你该看什么：SwiGLU 随输入 scale 增长而单调上升（发散），SiTU-GLU 很快就趋于平稳（被钳住）
- 理想情况：SiTU-GLU 的曲线在上界附近水平，SwiGLU 持续上升
- 异常指示：SiTU-GLU 也随 scale 增长 → softcap 未生效
右下 (d)：RMSNorm 尺度不敏感性
- 图类型：双折线图
- X 轴：down 权放大倍数（1 → 50）
- Y 轴：输出幅值比 (放大/原始)
- 颜色：蓝色 = 无 latent_norm（随 factor 线性增长），橙色 = 有 latent_norm（保持 ≈1.0）
- 你该看什么：有 RMSNorm 的线几乎水平（在 1.0 附近），无 RMSNorm 的线 45 度上升
- 理想情况：有 norm 的比值 ≈ 1.0（误差 < 1e-3），无 norm 的比值 ≈ factor
- 异常指示：两条线都接近 1 → 权重放大未生效；两条线都增长 → RMSNorm 被关闭
调节建议
- 想看「QB 改善更夸张」→ 调 n_experts 到 32、hot_experts 到 "0,1,2,3"
- 想看「直方图估计更精确」→ 调 bins 到 1024、n_batches 到 100
- 想看「SwiGLU 完全爆炸」→ 此 Tab 内置 scale 上限为 8，想看更大 scale 请去 Tab 12
为什么重要
这是读者接触模块的第一个综合印象——四个子图各代表一处核心创新，看完这 Tab 就知道 LatentMoE 到底解决了哪四个问题、每个问题的效果量级有多大。它是后续深入 Tab 的钩子和索引。
Tab 3：前向契约
目的
验证 StableLatentMoE 层的前向输出契约：形状正确、路由结构存在、与独立 Eq.11 参考实现逐点一致、梯度只流向被选中专家。这是模块实现正确性的主锚点——所有其他性质的验证都建立在前向算对了的基础上。
对应测试 est_forward_shape_and_routing / est_matches_an_independent_eq11_reference / est_only_selected_experts_receive_gradient。
UI 组件与参数影响
| 参数 | 范围 | 默认 | 调节效果 |
|---|---|---|---|
| dim | 8-128 | 32 | 增大 → 模型维度更大，输出尺寸成比例变化 |
| batch | 1-8 | 3 | 增大 → 输出 batch 维变大，总 token 数 = batch×seq |
| seq_len | 2-16 | 6 | 增大 → 序列更长，总 token 数增多 |
| seed | 0-10 | 0 | 随机种子，改变具体权重和输入，但不影响性质结论 |
输出图解读
tab-3
1×3 三联图布局。
左 (a)：前向形状检查
- 图类型：柱状图
- X 轴：hidden_dim、n_indices、n_tokens
- Y 轴：数值
- 颜色：蓝/橙/绿三色柱
- 你该看什么：hidden_dim 应等于 dim（32），n_indices 应等于 batch×seq×top_k（18 = 3×6×2 中的前两个维度乘 top_k），n_tokens 等于 batch×seq
- 理想情况：hidden_shape = (batch, seq, dim)，indices_shape = (batch*seq, top_k)
- 异常指示：hidden_dim ≠ dim → 投影层维度错；indices 形状不对 → top_k 或 flatten 逻辑错
中 (b)：Eq.11 参考实现 vs MoE 对拍
eq-11
- 图类型：柱状图（两个误差指标）
- X 轴：max_abs_diff、mean_abs_diff
- Y 轴：绝对误差
- 颜色：红/橙
- 你该看什么：两个柱子都几乎贴地（接近 0），说明模块输出与独立手写的 Eq.11 参考实现逐点一致
- 理想情况：max_abs_diff < 1e-5，mean_abs_diff < 1e-6
- 异常指示：柱子明显升高（> 1e-3）→ 模块实现与公式有偏差，可能是 SiTU-GLU 参数用错、加权方式反了、或共享专家漏加
右 (c)：梯度稀疏性检查
- 图类型：柱状图
- X 轴：touched_experts（有梯度的专家数）、n_routed（总路由专家数）
- Y 轴：count
- 颜色：绿/蓝
- 你该看什么：touched_experts 应小于 n_routed（因为 top_k=1，只有被选中的那些专家有梯度），且 touched 是 used 的子集
- 理想情况：touched < n_routed，且 touched_subset_of_used = True
- 异常指示：touched == n_routed → 所有专家都有梯度，说明 dispatch 走了 dense 路径而非稀疏路径
为什么重要
这是实现正确性的底线验证。如果前向契约不成立——形状错了、输出与公式对不上、梯度不稀疏——后续所有性质验证都没有意义。此 Tab 是整个演示的信任锚点，建议每次改动代码后先跑它确认基础正确。
Tab 4：潜空间 width 分离
目的
直观展示 LatentMoE 的核心结构特征——路由专家工作在潜空间 l 维，共享专家工作在全宽 d 维（width 分离）。验证的是论文 §2.3 开篇的 width separation 定义（Eq. 11 的 W_down 投影）。
对应测试 est_latent_bottleneck_is_where_routed_experts_live。
背景概念
Width 分离：常规 MoE 中所有专家都是全宽（d 维）的，专家越多参数量和计算量线性放大。LatentMoE 在路由专家前加一个 W_down 投影，把 token 从 d 维降到 l 维（l < d），路由专家全部在 l 维工作，最后 W_up 再升回 d 维。这样路由专家的参数和计算量都从 d 级降到 l 级，K3 中 l=3584、d=4096，约节省 12.4% 的专家侧参数。
代码位置：moe.py:StableLatentMoE.__init__ 中的 down_proj（d→l）、experts（l→H→l）、up_proj（l→d）。
UI 组件与参数影响
| 参数 | 范围 | 默认 | 调节效果 |
|---|---|---|---|
| dim (d) | 16-512 | 64 | 增大 → 全宽维度变大，shared_gate_in_features 随之变大，width 节省比例上升 |
| latent_dim (l) | 8-256 | 16 | 显示值不变（当前实现使用小配置 latent_dim=16 固定值，此滑块暂未接入） |
| n_routed | 2-64 | 8 | 显示值不变（同上，固定为 8） |
| n_shared | 1-8 | 2 | 显示值不变（同上，固定为 2） |
| expert_hidden | 8-128 | 24 | 显示值不变（同上，固定为 24） |
注意：此 Tab 的 latent_dim / n_routed / n_shared / expert_hidden 四个滑块当前仅为占位，底层 measure_latent_bottleneck_dimensions 函数只接收 dim 参数，内部使用固定的 small_cfg() 配置（latent_dim=16, n_routed=8, n_shared=2, expert_hidden=24）。只有 dim 滑块真正影响输出。
”
输出图解读
tab-4
1×2 双联图布局。
左 (a)：维度对比
- 图类型：柱状图
- X 轴：experts_dim（路由专家门控输入维）、down_dim（路由专家降维输出维）、shared_in（共享专家门控输入维）
- Y 轴：特征维度
- 颜色：红/绿/蓝
- 你该看什么：experts_dim 和 down_dim 应该远小于 shared_in——路由专家在窄的潜空间（l=16）工作，共享专家在宽的全空间（d=64）工作
- 理想情况：experts_dim = latent_dim（16），shared_in = dim（64），节省比例 = (1 - 16/64) = 75%
- 异常指示：experts_dim == shared_in → 没有 width 分离（l=d），路由专家与共享专家宽度相同
右 (b)：专家数量
- 图类型：柱状图
- X 轴：n_routed（路由专家数）、n_shared（共享专家数）
- Y 轴：count
- 颜色：红/绿
- 你该看什么：路由专家远多于共享专家（K3 是 896 vs 2；小配置是 8 vs 2）
- 理想情况：n_routed > n_shared
- 异常指示：n_routed == 0 → 没有路由专家，退化为纯共享 FFN
为什么重要
Width 分离是 LatentMoE 的灵魂——没有它，896 个专家的参数和计算量会爆炸到不可接受。此 Tab 把「路由专家在窄带、共享专家在宽带」这件事用柱状图直观展示出来，是理解 K3 为什么能支撑近千专家的第一个关键点。
Tab 5：RMSNorm 稳定性
目的
验证 §2.3.1 的核心论断——在 up-projection 前加 RMSNorm 可以消除路由分支的尺度敏感性：放大专家权重若干倍，有 RMSNorm 的输出几乎不变，无 RMSNorm 的输出随之线性放大。
对应测试 est_rmsnorm_before_up_projection_removes_scale_sensitivity。
背景概念
病态链问题：LatentMoE 的路由分支把 W_down → 专家 FFN（两步 matmul）→ 加权聚合 → W_up 串成近似 4 个连续矩阵乘。在 2.8T 规模下，这条链的条件数很差——专家权重的幅度、选了哪些专家、权重如何分配，都会放大成输出尺度的波动。RMSNorm 在聚合之后、上投影之前做归一化（u / RMS(u)），把输出尺度固定下来，下游永远接到单位 RMS 的信号。
UI 组件与参数影响
| 参数 | 范围 | 默认 | 调节效果 |
|---|---|---|---|
| dim | 8-128 | 32 | 增大 → 维度更大，但比例关系不变（比值应仍接近 1.0 / factor） |
| down 权重放大倍数 | 2-100 | 50 | 增大 → 无 RMSNorm 的输出比值线性增大；有 RMSNorm 的仍保持 ≈1.0 |
| seed | 0-10 | 0 | 随机种子，不影响性质结论 |
输出图解读
tab-5
单图布局（1 个柱状图） 。
中心图：RMSNorm 尺度敏感性
- 图类型：双柱状图
- X 轴：有 RMSNorm (ratio)、无 RMSNorm (ratio)
- Y 轴：output ratio（放大后幅值 / 原始幅值）
- 颜色：绿色 = 有 RMSNorm，红色 = 无 RMSNorm；灰色虚线 = 理想值 1.0
- 你该看什么：绿色柱子几乎贴在 1.0 线上（有 RMSNorm 放大 50 倍后输出几乎不变），红色柱子远超 1.0（接近 50）
- 理想情况：ratio_with_norm ≈ 1.0（误差 < 1e-3），ratio_no_norm ≈ factor（50 左右）
- 异常指示：两根柱子都接近 1.0 → 权重放大未生效；两根柱子都接近 50 → RMSNorm 被关闭或失效
调节建议
- 想看「极端对比」→ 调 factor 到 100，无 norm 柱子会升到 100，有 norm 仍贴地
- 想看「小倍率下的微差」→ 调 factor 到 2，两根柱子分别约为 2.0 和 1.0
为什么重要
这是「Stable LatentMoE」之所以叫 Stable 的第一根支柱。没有 RMSNorm，路由分支的输出尺度会随路由决策剧烈波动，在大模型规模下会引发数值失稳。RMSNorm 不是可选项——论文指出它本身也改善验证损失。在生产中，如果发现训练不稳定，第一个要排查的就是 RMSNorm 是否被意外关闭。
Tab 6：梯度稀疏性
目的
验证稀疏 MoE 的核心性质——只有被路由选中的专家会获得梯度，未被选中的专家梯度为 0。这是 MoE 能节省训练计算量的本质原因：每一步只有 top_k 个专家更新权重。
对应测试 est_only_selected_experts_receive_gradient。
UI 组件与参数影响
| 参数 | 范围 | 默认 | 调节效果 |
|---|---|---|---|
| dim | 4-64 | 16 | 增大 → 维度更大，但稀疏性比例不变 |
| n_routed | 2-32 | 8 | 显示值不变（当前实现使用小配置 n_routed=8 固定值，此滑块暂未接入） |
| top_k | 1-4 | 1 | 增大 → 每 token 选更多专家，touched_experts 数量增加（稀疏度下降） |
| seed | 0-10 | 0 | 增大 → 不同随机种子会选中不同的专家，但稀疏性性质不变 |
注意：n_routed 滑块当前未接入底层函数（measure 内部使用 small_cfg 的 8 个专家），调节它不会改变结果。实际有效的参数是 dim、top_k、seed。
”
输出图解读
tab-6
单图布局（1 个柱状图） 。
中心图：梯度稀疏性
- 图类型：三柱状图
- X 轴：used_experts（路由选中的专家数）、touched_experts（有梯度的专家数）、n_routed（总路由专家数）
- Y 轴：专家数量
- 颜色：蓝/绿/红
- 你该看什么：touched_experts 应该等于 used_experts（有梯度的恰好就是被选中的那些），且都远小于 n_routed
- 理想情况：touched_subset_of_used = True，touched_experts = top_k 左右（seq_len 个 token 可能选中不同专家，所以通常是 top_k~top_k×seq 之间）
- 异常指示：touched_experts == n_routed → 所有专家都有梯度，说明走了 dense 路径，MoE 稀疏性失效
调节建议
- 想看「极端稀疏」→ 调 top_k 到 1，只有 1 个专家有梯度
- 想看「接近稠密」→ 调 top_k 到 4，更多专家获得梯度（但仍 < n_routed）
为什么重要
梯度稀疏性是 MoE 架构的本质收益来源——如果未被选中的专家也有梯度，那 MoE 就退化为了全连接（只是更慢）。此 Tab 验证 dispatch 实现的正确性：index_add 只对被选中的专家加权累加，未选中专家的权重梯度严格为零。在生产中，如果发现训练速度远慢于理论稀疏度，第一个排查点就是梯度稀疏性。
Tab 7：Dispatch 正确性
目的
验证 GroupedExperts 的 orward_indexed（按专家分组 dispatch）与 orward_all（稠密评估）逐点一致，且对 token 顺序置换不变。这是 MoE 专家执行层的数值正确性验证。
对应测试 est_grouped_dispatch_matches_dense_evaluation / est_dispatch_is_permutation_invariant。
背景概念
Dispatch：把 (tokens, dim_in) 的 token 按 expert_ids 分配给各自对应的专家执行。生产中这会被优化为 grouped GEMM（一次矩阵乘法覆盖所有被选中的专家）。CPU 教学实现里使用「按专家排序 → 分段稠密 matmul → 回填原始顺序」的朴素方式，但数值上必须与「所有专家对所有 token 计算再按 id 取行」完全等价。
代码位置：moe.py:GroupedExperts.forward_indexed。
设计思路
背景：Dispatch 是 MoE 的"心脏手术"
MoE 的 forward_indexed 是生产路径——它把 token 按专家分组（torch.argsort(expert_ids) 让同专家的 token 连续），然后每专家只做自己那部分的 GEMM，最后 scatter 回原顺序。这个流程有 3 个容易出错的地方：
- 分组时专家权重寻址错（用了错的
W_e） - scatter 回填顺序错（结果对不上原 token 顺序）
- 排序/反排序不对称（导致结果依赖输入排列）
Tab 7 设计两个图分别对应 2 个独立验证场景：
左图：grouped vs dense（验证"专家寻址正确性"）
场景：同一个 token 批次，用两种方式计算：
- grouped 路径（
forward_indexed）：生产实现，分组 + 分组 GEMM + 回填 - dense 路径（
forward_all）：教科书参考，每个 token 过所有 n_experts 个专家，再按ids[i]选第 i 个 token 对应的专家输出
为什么这么设计：
- dense 路径不可能错——它不用排序、不用分组、不用 scatter，就是简单的
einsum('ne,tnd->ted')然后按索引取 - 如果 grouped 和 dense 不一致，说明 grouped 的"每专家取自己的权重子矩阵"逻辑有 bug
- 数学等价性：
forward_indexed(x, ids)[i] = forward_all(x)[ids[i], i]，恒等式
典型 bug 它能抓到：
argsort后没记录逆置换 → scatter 回填顺序错- 专家权重
W_e在分组时取错（比如第 3 组用了第 4 个专家的权重）
右图：置换不变性（验证"排序/反排序对称性"）
场景：同一批 token，先按原顺序 dispatch 得 direct；把 token 按随机置换 perm 打乱后再 dispatch 得 shuffled；最后检查 direct[perm] == shuffled。
为什么这么设计：
- dispatch 内部用了
argsort（排序），但结果应该与输入顺序无关——token i 的输出只取决于x[i]和ids[i]，不应该因为"它排在第几个"而不同 - 这是强性质检验：即使 grouped vs dense 通过，如果排序/反排序不对称，特定排列下仍可能出错
- 能抓到 grouped vs dense 抓不到的 bug：比如 scatter 时把
inverse_perm写成了perm（正向用了反向）
典型 bug 它能抓到：
scatter_index = perm而不是scatter_index = inverse_perm→ 大部分情况下结果碰巧对（dense 验证通过），但某些排列下错- 分组时按
ids排序，回填时按ids[perm]排序 → 两个排序不一致
两图互补
| 图 | 验证性质 | 抓什么 bug | 抓不到什么 |
|---|---|---|---|
| grouped vs dense | 专家寻址正确 | 权重寻址错、scatter 顺序错（与排列无关的那种） | 排序/反排序不对称 |
| 置换不变性 | 排序/反排序对称 | 排序用错、回填用错、两个排序不一致 | 与排列无关的寻址错 |
一个过、一个不过的情况恰好能定位 bug 类型：
- 左过右不过 → 排序/反排序不对称
- 左不过右过 → 寻址错但碰巧对称（罕见）
- 都不过 → 多重 bug
UI 组件与参数影响
| 参数 | 范围 | 默认 | 调节效果 |
|---|---|---|---|
| n_experts | 2-32 | 6 | 增大 → 更多专家参与，dispatch 的排序和分段更复杂 |
| dim_in | 4-32 | 8 | 增大 → 输入维度更大 |
| hidden | 8-64 | 12 | 增大 → 专家内部隐藏层更宽 |
| dim_out | 4-32 | 8 | 增大 → 输出维度更大 |
| n_tokens | 8-128 | 20 | 增大 → 更多 token 参与 dispatch |
| seed | 0-10 | 0 | 随机种子，不影响性质结论 |
输出图解读
tab-7
1×2 双联图布局。
左 (a)：grouped vs dense
- 图类型：双柱状图，Y 轴对数刻度
- X 轴：max_abs_diff、mean_abs_diff
- Y 轴：abs diff（log scale，1e-12 起）
- 颜色：红/橙；灰色虚线 = 1e-6 容差线
- 你该看什么：两根柱子都远低于 1e-6 容差线（在 1e-8 ~ 1e-9 量级），说明分组 dispatch 与稠密 oracle 的误差是浮点舍入级别，可视为逐点一致
- 理想情况：max_abs_diff ~ 1e-7 ~ 1e-8，mean_abs_diff 比 max 小 1~2 个数量级，两者均 << 1e-6
- 异常指示：任何一根柱子高于 1e-6 容差线 → dispatch 实现有 bug（排序索引错位、回填顺序错、或专家权重寻址错）
- 可视化陷阱：max 与 mean 差 ~40 倍，线性刻度下 mean 会几乎看不见；log 刻度下两者都能看清
右 (b)：置换不变性
- 图类型：双柱状图，Y 轴对数刻度
- X 轴：max_abs_diff、mean_abs_diff
- Y 轴：abs diff（log scale）
- 颜色：蓝/紫；灰色虚线 = 1e-6 容差线
- 你该看什么：两根柱子都远低于 1e-6 容差线，说明打乱 token 顺序后再 dispatch、再还原，结果与原顺序逐点一致
- 理想情况：max_abs_diff ~ 1e-8 或更小（纯索引置换通常误差为 0）
- 异常指示：柱子高于 1e-6 容差线 → dispatch 的排序/回填逻辑不对称（比如 gather 是 argsort 但 scatter 用错了索引）
- 可视化陷阱：若 max_abs_diff 为 0（精确相等），柱子会落到 1e-12 下限（0 值替换为 1e-12 避免 log(0)）
为什么重要
Dispatch 是 MoE 的执行核心。如果 grouped 与 dense 不一致，说明专家权重的寻址或 token 的排列/逆排列有 bug——在大模型下这种误差会被专家数和层数放大，训练出来的模型与公式定义完全不同。置换不变性是更强的验证：如果同一批 token 换个顺序结果就变了，说明实现有状态污染或索引错误。
Tab 8：QB 语义
目的
验证 Quantile Balancing 路由器的三项核心语义：cutoff 等于第 k+1 个 biased score、偏置只导流不改权重、输出权重归一化。这是 QB 算法的语义基础——如果这些性质不成立，后续的均衡效果无从谈起。
对应测试 est_top_k_plus_one_yields_the_cutoff / est_bias_steers_dispatch_but_not_the_mixture_weights / est_weights_sum_to_one。
背景概念
QB 语义三件套：
- cutoff 语义：路由器用 Top-(k+1) 选择，前 k 个是实际路由，第 k+1 个 biased score 就是 cutoff α——一个专家的 biased score 必须超过 α 才会进入 Top-k。cutoff 取自路由本身，省去了单独的 token 侧分位数计算。
- 偏置只导流不改权重：偏置 b 只加在 biased score 上参与 Top-k 选择，但混合权重 p_ij = s_ij / Σs_ir 只用原始分数 s（不含 b）。这就是「无损失」（auxiliary-loss-free）的含义——负载均衡不扭曲路由器的优化目标。
- 权重归一化：每个 token 的 k 个被选专家的权重加起来等于 1，构成一个概率分布。
UI 组件与参数影响
| 参数 | 范围 | 默认 | 调节效果 |
|---|---|---|---|
| n_experts | 4-32 | 8 | 增大 → 专家更多，cutoff 与选中集的差距可能更明显 |
| top_k | 1-8 | 3 | 增大 → 选更多专家，cutoff 是第 k+1 个 |
| dim | 4-64 | 16 | 增大 → 维度更大，但语义性质不变 |
| n_tokens | 16-256 | 64 | 增大 → 更多 token 样本，归一化误差统计更稳定 |
| bias_std | 0.5-4.0 | 2.0 | 增大 → 偏置扰动更大，但 weights 应仍与 scores 归一化一致 |
| seed | 0-10 | 0 | 随机种子，不影响性质结论 |
输出图解读
tab-8
1×3 三联图布局。
左 (a)：cutoff 语义
- 图类型：双柱状图（布尔值）
- X 轴：cutoff_matches_topk、selected_above_cutoff
- Y 轴：bool（1=True，0=False）
- 颜色：绿/蓝
- 你该看什么：两根柱子都应该是 1（True），即 cutoff 恰好等于第 k 个排序后的 biased score，且所有被选中专家的 biased score 都 ≥ cutoff
- 理想情况：两项均为 True，min_gap > 0（选中专家与 cutoff 之间有正间隙）
- 异常指示：任何一项为 False → Top-(k+1) 选择或 cutoff 语义实现错误
中 (b)：bias 只导流不改权重
- 图类型：双柱状图（布尔值）
- X 轴：scores_unchanged、weights_match_scores
- Y 轴：bool（1=True，0=False）
- 颜色：均为绿色（表示应该通过）
- 你该看什么：两根柱子都为 1——加偏置前后 scores 完全不变（scores 与 bias 无关），weights 等于选中集内 scores 归一化
- 理想情况：两项均为 True
- 异常指示：scores_unchanged = False → bias 错误地进入了分数计算；weights_match_scores = False → 权重归一化时混入了 bias
右 (c)：权重归一化误差
- 图类型：双柱状图，Y 轴对数刻度
- X 轴：max_err、mean_err
- Y 轴：|sum(weights) - 1|（log scale，1e-12 起）
- 颜色：红/橙；灰色虚线 = 1e-5 容差线
- 你该看什么：两根柱子都远低于 1e-5 容差线（在 1e-7 ~ 1e-9 量级），说明权重和非常接近 1
- 理想情况：max_err ~ 1e-7 或更小，mean_err 比 max_err 小 1~2 个数量级
- 异常指示：max_err 高于 1e-5 容差线 → 权重归一化有 bug（可能是分母计算错了或数值溢出）
- 可视化陷阱：max 与 mean 差几十倍，线性刻度下 mean 会几乎看不见；log 刻度下两者都能看清
调节建议
- 想看「偏置很大但权重仍正确」→ 调 bias_std 到 4.0，验证权重归一化对偏置幅度不敏感
- 想看「极端小 top_k 的 cutoff 语义」→ 调 top_k 到 1，cutoff 就是第 2 个 biased score
为什么重要
这是 QB 算法的语义地基。如果偏置同时影响了权重（扭曲优化）、或者 cutoff 语义不对（均衡推导的前提不成立）、或者权重不归一（输出尺度不对），那么后续所有关于均衡效果的讨论都没有意义。此 Tab 是 QB 正确性的第一道关口。
Tab 9：QB 均衡效果
目的
展示 Quantile Balancing 的负载均衡效果：一步改善有多大、是否单调收敛、直方图估计与精确分位数的误差有多少。这是 QB 论文 §2.3.3 的核心卖点——「一步到达目标负载，无需学习率」。
对应测试 est_one_update_collapses_severe_imbalance / est_convergence_is_monotone / est_histogram_matches_the_exact_quantile。
UI 组件与参数影响
| 参数 | 范围 | 默认 | 调节效果 |
|---|---|---|---|
| n_experts | 4-64 | 16 | 增大 → 专家更多，初始失衡更严重，一步改善更明显 |
| top_k | 1-8 | 2 | 增大 → 选更多专家，天然负载更均匀，改善空间更小 |
| dim | 8-128 | 32 | 增大 → 维度更大，但均衡效果主要由输入偏斜度决定 |
| bins | 500-4000 | 2000 | 增大 → 直方图分箱更细，分位数估计误差更小（右下子图更接近） |
| tokens | 1024-8192 | 4096 | 增大 → 样本更多，统计更稳定 |
| seed | 0-10 | 0 | 随机种子，影响初始失衡程度 |
输出图解读
tab-9
2×2 四联图布局。
左上 (a)：一步均衡 max/min 负载比
- 图类型：双柱状图
- X 轴：before、after 1 step
- Y 轴：max_min_ratio（最忙专家 / 最闲专家）
- 颜色：红色 = before，绿色 = after
- 你该看什么：before 柱子很高（严重失衡），after 柱子显著降低（接近均衡）
- 理想情况：before/after 比值改善 > 10 倍（默认配置下约 51 → 1.95，改善 26 倍）
- 异常指示：after 仍很高（> 5）→ QB 未生效或样本太少
右上 (b)：单调收敛曲线
- 图类型：折线图
- X 轴：均衡步数（1/2/4/8）
- Y 轴：max_min_ratio
- 颜色/标记：紫色实线圆点
- 你该看什么：曲线单调下降——步数越多，负载比越低，越接近 1.0
- 理想情况：is_monotone_descending = True，总改善 > 1.5 倍
- 异常指示：曲线先升后降或震荡 → 直方图估计噪声过大，或 update 逻辑有 bug
左下 (c)：直方图 vs 精确分位数
- 图类型：双折线对比图
- X 轴：expert id
- Y 轴：bias 偏置值
- 颜色/线型：蓝色实线圆点 = exact，橙色虚线方块 = histogram
- 你该看什么：两条线几乎重合，说明直方图近似的精度损失很小
- 理想情况：max_diff < 0.001（2000 bins 下约 0.0004），diff_in_units_of_bin < 1.5
- 异常指示：两条线明显分叉 → 直方图实现有 bug 或 bins 太少
右下 (d)：负载变异系数 CV
- 图类型：双柱状图
- X 轴：before、after 1 step
- Y 轴：cv（标准差 / 均值）
- 颜色：红色 = before，绿色 = after
- 你该看什么：CV 从较高值降到很低，说明负载分布从偏斜变平坦
- 理想情况：after cv < 0.3（默认配置下约 0.1~0.2 级别）
- 异常指示：CV 下降很少 → 均衡效果差
调节建议
- 想看「最夸张的一步改善」→ 调 n_experts 到 64、tokens 到 8192，before 的 max_min_ratio 会非常高
- 想看「直方图高精度」→ 调 bins 到 4000，exact 与 histogram 几乎重合
- 想看「稀疏度极高时的均衡」→ 调 top_k 到 1，每 token 只选 1 个专家，均衡更有挑战性
为什么重要
这是 QB 的卖点验证——传统固定步长方法需要很多步缓慢收敛，QB 一步就把几十倍的失衡压到 2 倍以内。在 K3 896 专家的规模下，这种一步收敛的特性直接决定了 EP 训练是否能稳定跑起来。单调收敛和直方图精度则分别验证了算法的稳定性和大规模可部署性。
Tab 10：QB 工程行为
目的
验证 Quantile Balancing 的工程约定：延迟一步更新、mean-centring 不动点、推理冻结、可关闭开关。这些是生产部署中的硬约束——任何一项不满足，QB 就不能安全地用在真实训练中。
对应测试 est_bias_update_is_deferred_by_one_step / est_mean_centring_does_not_change_selection / est_bias_is_mean_centred / est_frozen_at_inference / est_disabled_router_leaves_the_bias_at_zero。
背景概念
QB 的 4 项工程约定：
- 延迟一步更新：forward 只统计数据（_accumulate），不修改 bias；bias 只有在调用 update_bias() 后才改变。这保证了「某批算出的偏置不会用来路由同一批」，避免统计自洽性问题。
- mean-centring 不动点：整体平移 bias（b → b + c）不改变路由选择（indices 和 weights 都不变），只会整体平移 cutoff。因此 QB 更新时做 mean-centring（b - mean(b)）是零副作用的。
- 推理冻结：eval 模式下 forward 不 _accumulate，histogram 和 token_count 保持为 0。推理时偏置是训练收敛的死常数。
- 可关闭：enabled=False 时，即使 train 模式下 update_bias() 也不会改变 bias（始终为 0），干净地旁路 QB。
设计思路
背景：QB（Quantile Balancing）是 K3 的"无损失负载均衡"
传统 MoE 用辅助损失（load balancing loss）推动专家均衡，但损失函数会改变主任务梯度。K3 用 QB：只调 bias，不改 loss——通过直方图估计分位数，把"拥挤专家"的 bias 调低、"冷清专家"的 bias 调高。
QB 是有状态的（bias、histogram、token_count），工程上有 5 个关键不变量必须保证，否则训练会污染或不可复现。Tab 10 设计 7 个 case 对应 5 个工程场景：
Case 1-2: 延迟一步更新（measure_deferred_update）
场景：训练循环里，forward 和 update_bias() 的时序关系。
为什么重要：绝不能"用某批算出的 bias 去路由同批" ——否则该批的 Top-k 选择与 cutoff 统计互相污染（第 5 个 token 的路由受前 4 个 token 的影响，但前 4 个又是基于旧 bias 路由的）。
设计：两个检查点：
unchanged_after_forward：forward 后 bias 不变（本批只是累积统计，没动 bias）changed_after_update：update_bias()后 bias 确实变了（更新生效）
抓什么 bug：如果 forward 里就改了 bias，说明时序错了；如果 update_bias() 没改 bias，说明更新逻辑失效。
Case 3-4: 平移不动点（measure_mean_centring_invariance）
场景：update_bias() 里会做 b ← b - mean(b)（mean-centring），控制 bias 整体尺度。
为什么重要：mean-centring 必须零副作用于路由语义——整体平移 bias 不应该改变 Top-k 排序（因为 cutoff 也随动）。
设计：把 bias 整体 +7.25，检查： 3. indices_unchanged：Top-k 索引不变 4. weights_unchanged：路由权重不变（因为 weights 只看原始 scores，不看 bias）
抓什么 bug：如果 mean-centring 改了路由，说明 cutoff 没随 bias 平移，或者 weights 计算混入了 bias。
Case 5: mean-centring 生效（measure_bias_mean_centred）
场景：update_bias() 后 bias 的均值应该接近 0。
为什么重要：mean-centring 防止"所有 bias 同时漂移"（全局正偏），否则多卡训练时各 rank 的 bias 尺度会发散。
设计：5. |mean(bias)| < 1e-5
抓什么 bug：漏了 mean-centring 这一步。
Case 6: 推理冻结（measure_inference_frozen）
场景：eval 模式下，路由器的 histogram 和 token_count 应该保持 0——推理时不累积统计。
为什么重要：推理时 bias 是训练收敛的"死常数"，如果还持续累积统计（或更糟，调用 update_bias()），会把线上流量当成训练数据、污染已收敛的 bias。
设计：6. histogram.sum() == 0（forward 后直方图仍为空）
抓什么 bug：eval 模式下 _accumulate() 没被跳过。
Case 7: 关闭开关彻底（measure_disabled_router）
场景：enabled=False 时，QB 完全旁路——bias 恒为 0，即使训练模式下也无效。
为什么重要：QB 是可选机制，必须能干净关闭（比如先用辅助损失训练的阶段，或做对比实验）。关闭不彻底会让"partial 均衡"悄悄进入模型，测试结果不可复现。
设计：7. bias.all() == 0（即使 update_bias() 后）
抓什么 bug：enabled=False 时 _accumulate 或 update_bias 没短路。
5 个场景覆盖 QB 的完整生命周期
训练 forward ──→ 累积统计（Case 1, 6 的反向）
↓
update_bias() ──→ 改 bias（Case 2）
↓
mean-centring ──→ 零副作用（Case 3, 4, 5）
↓
推理 forward ──→ 冻结统计（Case 6）
↓
enabled=False ──→ 完全旁路（Case 7）
7 个 case 全过 = QB 的工程实现可信，任何一条挂了都指向具体子系统的 bug。
UI 组件与参数影响
| 参数 | 范围 | 默认 | 调节效果 |
|---|---|---|---|
| dim | 4-32 | 16 | 增大 → 维度更大，但工程行为不变 |
| n_experts | 4-32 | 8 | 增大 → 专家更多，但工程行为不变 |
| top_k | 1-4 | 2 | 增大 → 选更多专家，但工程行为不变 |
| n_tokens | 32-256 | 128 | 增大 → 更多 token 样本 |
| 平移量 | 1.0-15.0 | 7.25 | 增大 → 测试 mean-centring 不动点的平移幅度更大，但结果应仍不变 |
| seed | 0-10 | 0 | 随机种子，不影响性质结论 |
输出图解读
tab-10
单图布局（横向条形图，7 项检查） 。
中心图：QB 工程行为检查
-
图类型：横向条形图（7 项 PASS/FAIL 检查）
-
Y 轴：7 项检查名称
-
X 轴：是/否（0 或 1）
-
颜色：绿色 = PASS（1），红色 = FAIL（0）
-
你该看什么：所有 7 条都应该是绿色（PASS）
-
7 项检查内容：
- forward 后 bias 不变（延迟更新）
- update 后 bias 改变（延迟更新）
- 平移不改变 indices（mean-centring 不动点）
- 平移不改变 weights（mean-centring 不动点）
- 更新后 bias 均值 ≈ 0（mean-centring）
- 推理冻结：histogram_sum == 0
- 关闭后 bias 全 0
-
理想情况：7/7 PASS，全部绿色
-
异常指示：任何一项红色（FAIL）→ 对应工程约定被破坏，需要排查
为什么重要
这些是 QB 能安全部署在生产训练中的硬约束。延迟更新保证统计口径自洽；mean-centring 不动点保证偏置去中心化没有副作用；推理冻结保证线上行为确定可复现；可关闭开关支持消融实验与分阶段训练。少任何一项，在多卡分布式训练中都可能出难以排查的问题。
Tab 11：896 专家压力测试
目的
在论文规模（896 专家、top_k=16、稀疏度 56）下验证 Quantile Balancing 仍然有效——一步就有显著改善、几步就收敛到极致均匀。这是 QB 能支撑 K3 极端稀疏配置的直接证据。
对应测试 est_scales_to_the_papers_expert_count。
背景概念
论文规模：K3 每层有 896 个 routed experts，每 token 激活 16 个（稀疏度 56）。在这种近千专家的规模下，传统固定步长均衡方法收敛极慢，容易造成持续的 EP 负载失衡——有的专家忙到爆、有的专家闲到饿死，训练效率和效果都受影响。QB 的设计动机就是在这种规模下仍能一步均衡。
UI 组件与参数影响
| 参数 | 范围 | 默认 | 调节效果 |
|---|---|---|---|
| dim | 16-128 | 64 | 增大 → 维度更大，运行时间增加，但效果趋势不变 |
| n_experts | 128-896 | 896 | 增大 → 专家越多，初始失衡越严重，一步改善越明显 |
| top_k | 4-32 | 16 | 增大 → 激活专家更多，稀疏度下降，天然更均衡 |
| bins | 500-2000 | 1000 | 增大 → 分箱更细，直方图精度更高 |
| tokens | 2048-16384 | 8192 | 增大 → 样本更多，统计更稳定 |
| seed | 0-10 | 0 | 随机种子，影响初始失衡程度 |
输出图解读
tab-11
1×3 三联图布局。
左 (a)：max_min_ratio（对数刻度）
- 图类型：三柱状图
- X 轴：初始失衡、1 步后、5 步后
- Y 轴：max_min_ratio（log scale）
- 颜色：红/橙/绿
- 你该看什么：初始值极高（上百甚至上千），1 步后显著下降，5 步后降到很低
- 理想情况：before > 100，one_step 大幅下降，converged 接近 1（< 2）
- 异常指示：1 步后下降不多 → QB 在大规模下失效
中 (b)：max_over_mean
- 图类型：三柱状图
- X 轴：初始失衡、1 步后、5 步后
- Y 轴：max_over_mean（最忙专家 / 平均负载）
- 颜色：红/橙/绿
- 你该看什么：5 步后 max_over_mean 应接近 1.0（1.1 以下）
- 理想情况：converged max_over_mean < 1.1
- 异常指示：5 步后仍 > 1.3 → 收敛不足
右 (c)：cv（变异系数）
- 图类型：三柱状图
- X 轴：初始失衡、1 步后、5 步后
- Y 轴：cv = std / mean
- 颜色：红/橙/绿
- 你该看什么：cv 从 0.9 左右降到 0.03 以下，说明负载分布从高度偏斜变成几乎平坦
- 理想情况：converged cv < 0.03
- 异常指示：cv 下降很少 → 均衡效果差
调节建议
- 想看「极端初始失衡」→ n_experts 拉到 896、tokens 拉到 16384，初始 max_min_ratio 可能达到几百
- 想看「低稀疏度下的均衡」→ top_k 拉到 32，稀疏度=28，初始失衡会小一些
- 想看「直方图精度影响」→ bins 降到 500，看收敛是否变慢或变差
为什么重要
这是 QB 的「压力测试」——论文中所有关于 QB 的好处（一步收敛、无需学习率）都是在大规模下才有意义。如果在 896 专家规模下 QB 不能一步见效，那整个设计动机就不成立。此 Tab 直接验证了 QB 在 K3 真实配置下的有效性，是整个模块最接近「生产指标」的验证。
Tab 12：SiTU-GLU
目的
全面验证 SiTU-GLU（Sigmoid Tanh Unit GLU）的数学性质：有界性、与 SwiGLU 的曲线对比、原点匹配、饱和区梯度（softcap vs 硬 clamp）、模块级有界性、以及同权重下 SwiGLU 离群 vs SiTU-GLU 守界的对比。这是「Stable」的第二根支柱——用有界激活消除 FP8 溢出风险。
对应测试 test_output_is_bounded / test_matches_swiglu_near_the_origin / test_gradients_survive_saturation / test_module_shapes_and_bound / test_swiglu_baseline_explodes_where_situ_does_not。
背景概念
SwiGLU 的离群点问题：SwiGLU = silu(gate) × up = sigmoid(gate) × gate × up。gate 和 up 两个因子都是无界的，在 MoE 这种极宽网络中，偶尔会有某个神经元的 gate 与 up 同时取到大值，二者乘积产生离群点（outlier）。在低精度（FP8/BF16）下，这些离群点会溢出或严重损失精度。
SiTU-GLU 的解决方案： 给 gate 和 up 各套一个 softcap（β · tanh(z/β)），让两个因子都有界，因此乘积也有界（≤ β_g × β_u）。K3 取 β_g=4、β_u=25，上界恰好为 100，小于 FP8 E4M3 的表示上限 448，保证 FP8 激活不溢出。
幅值差异
SiTU-GLU 是一个有界的 GLU 变体，SwiGLU 无界。相同权重、相同输入下，两者在小输入时几乎一样，幅值差异随输入增大而拉开。
公式对比（Kimi K3 的定义）
SwiGLU(x) = Swish(W_g x) ⊙ (W_u x) = [ W_g x · σ(W_g x) ] ⊙ (W_u x)
SiTU-GLU(x)= [ β₁·tanh(W_g x/β₁) ⊙ σ(W_g x) ] ⊙ [ β₂·tanh(W_u x/β₂) ]
其中 Kimi K3 取 β₁=4（门分支）、β₂=25（up 分支）
也就是说，SiTU-GLU 对 gate 分支和 up 分支各加了一个 scaled tanh 软封顶（softcap）。
三个层面的幅值差异
1. 门分支的幅值 ——SiTU 有上界，SwiGLU 无界
- Swish (g) = g・σ(g)：g 很大时 σ(g)→1，Swish ≈ g，无上界（g=100 → 幅值≈100）。
- SiTU 门 = β₁・tanh (g/β₁)・σ(g)：tanh 把它封在 ±β₁，正大值时趋于 β₁=4。
- 数例：g=4 → Swish≈3.93 vs SiTU 门≈2.99；g=10 → Swish≈10 vs SiTU 门≈3.95（已经封顶在 4 附近）。
2. up 分支的幅值 ——SiTU 有界，SwiGLU 线性
- SwiGLU 的 up = W_u x，线性、无上界。
- SiTU 的 up = β₂・tanh (W_u x/β₂)，封在 ±β₂=±25。
- 数例：u=50 → SwiGLU up=50 vs SiTU up≈24.1（趋近 25）；u=100 → 100 vs ≈24.98。
3. 输出乘积幅值 —— 本质区别
- SwiGLU 输出 = Swish (g)・u，两个因子都无界 → 整体无上界，可爆。
- SiTU-GLU 输出 = [有界门]・[有界 up]， | 输出 | ≤ β₁・β₂ = 4×25 = 100，严格有界。
关键：差异是 "随输入增大而放大"，不是全程不同
scaled tanh 在原点附近近似线性（β・tanh (x/β)≈x），所以：
- 小输入（线性区） ：SiTU-GLU 与 SwiGLU 的激活幅值几乎一致 —— 这正是 Kimi 论文里说的 "保留 SwiGLU 的局部响应"。
- 大输入（饱和区） ：SiTU-GLU 被软封顶（门≤4、up≤±25、输出≤100），SwiGLU 线性放大 ——幅值差异随 | 输入 | 增大而系统性放大。
为什么这么设计
SiTU-GLU 的动机就是防止激活爆炸、控制幅值上界：在大规模 / 深栈（Kimi K3 还配了 Attention Residuals）、低精度量化、以及需要数值稳定的场景下，SwiGLU 无界的大幅值容易成为不稳定源；把两个因子都软封顶后，模型在保留 SwiGLU 小信号行为的同时，保证任意输入下激活幅值不越界。
一句话：有区别，且是有意的—— 局部响应相同，大信号幅值 SiTU-GLU 被封顶（β₁、β₂、β₁β₂ 三个上界），SwiGLU 无界；相同权重同输入下，越大的输入两者幅值差得越多。
示意图（两分支的幅值对比）
门分支 gate（β₁=4）
幅值
7┤ S ─────────── Swish≈g，无上界，越走越高
6┤ ··
5┤ ··
4┤ ·························o──────────────────── ← β₁=4 封顶线（SiTU 门饱和）
3┤ o···::
2┤ o···:
1┤ o···:
0┤ ····o··o··o····o·····:
-0.3┤o·····:
└────┬────┬────┬────┬────┬────┬────┬────┬──→ g
-6 -4 -2 0 2 4 6
· = Swish(g) = g·σ(g) o = SiTU门 = 4·tanh(g/4)·σ(g)
小输入（线性区）两者几乎重合；随 |g| 增大，Swish 无界走高，SiTU 门封顶在 4
up 分支（β₂=25）
幅值
+25┤─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ← 上界 +β₂=25
│ o ── SiTU = 25·tanh(u/25) 封顶
│ o
│ o
0┤──────o────────────────────────────────────── (u=0)
│ o
│ o
-25┤─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ← 下界 -β₂=25
└────┬────┬────┬────┬────┬────┬────┬────┬──→ u
-30 -20 -10 0 10 20 30
SwiGLU up = u（线性，无界） vs SiTU up（封在 ±25）
小输入两者重合；大输入 SiTU 饱和、SwiGLU 线性放大
结论
输出幅值：|SiTU-GLU| ≤ β₁·β₂ = 4×25 = 100（严格有界）
SwiGLU = Swish(g)·u，两个因子都无界 → 整体无界（可爆炸）
小输入（线性区）：两者几乎一致（β·tanh≈x，保留 SwiGLU 局部响应）
大输入（饱和区）：SiTU 门→4、up→±25，幅值差异随 |输入| 增大而放大
UI 组件与参数影响
| 参数 | 范围 | 默认 | 调节效果 |
|---|---|---|---|
| beta_gate | 1.0-10.0 | 4.0 | 增大 → 上界线性增大，原点匹配误差减小（softcap 更接近恒等），饱和区梯度变化更缓 |
| beta_up | 5.0-50.0 | 25.0 | 增大 → 上界线性增大，激活曲线整体更高 |
| x_range | 2.0-30.0 | 10.0 | 增大 → 曲线展示范围更宽，能看到更远处的饱和行为 |
| input_scale | 10.0-100.0 | 40.0 | 增大 → 模块级有界测试中输入更大，SwiGLU 离群更严重（右下图对比更悬殊） |
| seed | 0-10 | 0 | 随机种子，影响模块权重初始化和爆炸测试的具体数值 |
输出图解读
tab-12
2×2 四联图布局。
左上 (a)：激活曲线 SiTU-GLU vs SwiGLU
- 图类型：双折线图
- X 轴：x（输入标量）
- Y 轴：SiTU-GLU(x, x) / SwiGLU(x, x)
- 颜色：绿色实线 = SiTU-GLU（有界），红色实线 = SwiGLU（无界）；灰色虚线 = 上界 ±β_g·β_u
- 你该看什么：在原点附近两条线几乎重合（原点匹配）；x 增大后 SiTU-GLU 趋于平坦（被 softcap 钳住），而 SwiGLU 继续上升
- 理想情况：|x| 大时 SiTU-GLU 接近水平直线（在上界附近），SwiGLU 持续增长；原点附近二者偏差 < 1%
- 异常指示：SiTU-GLU 也随 x 线性增长 → softcap 未生效或 beta 设置过大
右上 (b)：饱和区梯度 softcap vs 硬 clamp
- 图类型：分组柱状图
- X 轴：不同 z 值（如 z=20/50/100）
- Y 轴：梯度值
- 颜色：绿色 = softcap 梯度，红色 = 硬 clamp 梯度
- 你该看什么：softcap 梯度在饱和区仍为正值（虽然很小），硬 clamp 梯度为 0
- 理想情况：softcap_grad > 0 且 softcap_grad > hard_clamp_grad（后者在饱和区为 0）
- 异常指示：softcap 梯度也为 0 → 可能用了 clamp 而非 tanh
左下 (c)：核心性质验证（横向条形图）
- 图类型：横向条形图（4 项性质检查）
- Y 轴：4 项性质（上界可达、原点匹配、模块级有界、SwiGLU/SiTU 比）
- X 轴：0/1（PASS/FAIL）或比值
- 颜色：绿色 = PASS，橙色 = 数值展示项
- 你该看什么：上界可达（rel_err ≈ 0）、原点匹配误差很小（< 0.001）、模块级 all_finite = True、SwiGLU 离群是 SiTU-GLU 的几十倍
- 理想情况：4 项都通过，SwiGLU/SiTU 比 > 10
- 异常指示：任何一项红色 FAIL → 对应性质不成立
右下 (d)：同权重同输入下的激活幅值对比
- 图类型：双柱状图，Y 轴对数刻度
- X 轴：SiTU-GLU max、SwiGLU max
- Y 轴：max|activation|（log scale）
- 颜色：绿色 = SiTU-GLU，红色 = SwiGLU；灰色虚线 = 上界 β_g·β_u
- 你该看什么：SiTU-GLU 柱在上界附近（不超过），SwiGLU 柱远超上界（高几个数量级）
- 理想情况：SiTU max ≤ β_g·β_u，SwiGLU max >> β_g·β_u（比 SiTU 大 10x 以上）
- 异常指示：SiTU-GLU 也远超上界 → softcap 未生效；两根柱子接近 → 输入不够极端或 beta 太大
调节建议
- 想看「上界更高 + 更接近 SwiGLU」→ 调 beta_gate 到 10、beta_up 到 50，上界升至 500，原点匹配误差更小
- 想看「极端离群对比」→ 调 input_scale 到 100，SwiGLU 的 max 会非常夸张，SiTU-GLU 仍守在上界
- 想看「饱和区梯度差异最明显」→ 调 beta_gate 到 1，softcap 更早饱和，与硬 clamp 的梯度差更突出
为什么重要
SiTU-GLU 是「Stable LatentMoE」的第二根支柱。在 K3 的 FP8 推理场景中，如果激活无界，就需要额外的量化缩放和溢出处理，既慢又损失精度。SiTU-GLU 的有界性从数学上保证了激活永远不会超出 FP8 表示范围，是 K3 能高效跑 FP8 推理的关键前提。此 Tab 把有界、光滑、原点匹配三项核心性质一次性验证清楚。
0x04 附录
4.1 Tab 与测试函数的对应关系
| Tab | 测试文件 | measure/test 函数 |
|---|---|---|
| 1 原理总览 | - | plot_moe_principle.py:plot_moe_principle_four |
| 2 Demo 四联图 | - | demo_stable_latent_moe.py:plot_demo_four |
| 3 前向契约 | test_moe.py | measure_forward_shape_and_routing, measure_eq11_reference_consistency, measure_expert_gradient_sparsity |
| 4 潜空间 width 分离 | test_moe.py | measure_latent_bottleneck_dimensions |
| 5 RMSNorm 稳定性 | test_moe.py | measure_rmsnorm_scale_sensitivity |
| 6 梯度稀疏性 | test_moe.py | measure_expert_gradient_sparsity |
| 7 Dispatch 正确性 | test_moe.py | measure_grouped_dispatch_vs_dense, measure_dispatch_permutation_invariance |
| 8 QB 语义 | test_router_qb.py | measure_cutoff_semantics, measure_bias_steering, measure_weights_normalization |
| 9 QB 均衡效果 | test_router_qb.py | measure_one_step_balance, measure_monotonic_convergence, measure_histogram_vs_exact |
| 10 QB 工程行为 | test_router_qb.py | measure_deferred_update, measure_mean_centring_invariance, measure_bias_mean_centred, measure_inference_frozen, measure_disabled_router |
| 11 896 专家压力测试 | test_router_qb.py | measure_paper_scale_performance |
| 12 SiTU-GLU | test_situ_glu.py | measure_situ_vs_swiglu_curve, measure_situ_bound, measure_origin_match, measure_saturation_gradients, measure_module_bound, measure_swiglu_explodes_vs_situ |
4.2 已知参数限制
以下 Tab 中的部分滑块当前仅为占位，未接入底层 measure 函数：
| Tab | 未接入参数 | 说明 |
|---|---|---|
| 4 | latent_dim / n_routed / n_shared / expert_hidden | measure 函数使用固定 small_cfg，仅 dim 有效 |
| 6 | n_routed | 同上，仅 dim / top_k / seed 有效 |
4.3 默认配置下的典型数值
默认参数、seed=0 的实测值（供参考）：
- Tab 5 RMSNorm：有 norm ratio = 1.000, 无 norm ratio = 50.000
- Tab 9 QB 均衡：max_min_ratio 51.16 → 1.95（26x 改善），直方图误差 0.43 bin 宽
- Tab 11 论文规模：初始 358.7x → 1 步 max/mean 1.326 → 5 步 max/mean 1.046, cv 0.0149
- Tab 12 SiTU-GLU：上界 100 (rel_err=0.0)，原点匹配误差 0.02%，SwiGLU/SiTU 比 42.9x
TransFormer-封面
本文使用 markdown.com.cn 排版
