【轻量化复刻大模型 / KIMI K3】MoE (Stable LatentMoE) 交互式演示文档 - 罗西的思考 - 博客园

# 【轻量化复刻大模型 / KIMI K3】MoE (Stable LatentMoE) 交互式演示文档
本 Gradio 演示覆盖 K3 Stable LatentMoE 模块的全部核心机制，对应论文 §2.3 与附录 B/C/D，共 12 个交互 Tab：
复用测试层的 measure 函数，零侵入构建可视化层：
所有 Tab 均在 CPU 上实时计算，默认 2~4 线程，无需 GPU。
Stable LatentMoE 的核心定义（论文 Eq. 11）：
关键洞察：Stable LatentMoE 的三处「Stable」修复——RMSNorm 去尺度敏感性、SiTU-GLU 去激活离群、Quantile Balancing 去负载失衡——每一处都对应一个独立的数值失稳模式，合起来才让 896 专家 × 16 激活的极端稀疏配置在 2.8T 规模上可训练。少任何一处，病态链都会在某个尺度上爆炸。
直觉解释：把路由专家塞进更窄的潜空间，让「专家数量多」与「每个专家便宜」解耦——896 个专家听起来很夸张，但每个只在 3584 维而非 4096 维工作，叠加 top_k=16 的稀疏性，实际计算量远低于全宽稠密。
直觉解释：路由分支近似串了 4 个矩阵乘（W_down → 专家 FFN 两步 → W_up），这是一条病态链——专家权重的幅度、选了哪些专家、权重分配如何，都会放大成输出尺度的波动。RMSNorm 在聚合之后、上投影之前把尺度归一化，让下游永远接到单位 RMS 的信号。
直觉解释：SwiGLU 的 gate × up 两个无界因子相乘，MoE 宽度又大，低精度下偶尔会撞上「两个坐标都很大」的巧合，产生离群点溢出 FP8。SiTU-GLU 给两边都套 softcap，乘积天然不超过 \(β_g·β_u\)（K3 取 4×25=100，小于 FP8 E4M3 上限 448）。
直觉解释：传统的「b ← b + γ·sign(target - load)」像用固定步长爬山，步长是超参，收敛慢。QB 换个角度想：既然 cutoff 固定，「要让专家 j 恰好吃到 q 个 token」等价于找一个偏置值，使得刚好有 q 个 token 的 s_ij + b_j ≥ a_i。这个值就是 required 的第 q 个分位数——一步到位，无需学习率。
展示 Stable LatentMoE 的架构全景——数据流、width 分离节省的参数量/FLOPs、稀疏度与计算量的关系、以及层调度方式。建立整体直觉，为后续 11 个细节 Tab 做铺垫。
验证的是论文 §2.3 开篇的架构描述（Eqs. 11-14 的整体图景），不做数值断言，属于「一图看懂」级别的总览。
此 Tab 无交互参数，固定使用 K3 报告的 2.8T 配置（d=4096, l=3584, H=3072, E=896, k=16, n_shared=2）。
右上 (b)：width 分离参数量 / 计算量节省
右下 (d)：K3 层调度（dense warm-up + MoE body）
此 Tab 是整个演示的导航图——读者对 LatentMoE 的整体直觉都从这里建立。后续 11 个 Tab 的细节验证，都需要回溯到这张全景图来定位：width 分离在哪、RMSNorm 在哪、QB 在哪、SiTU-GLU 在哪。没有这张图，读者看到后面的数值验证会缺乏上下文。
用最小的配置一次性展示 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）在同一微型配置下的综合表现。
hot_experts 对应的专家权重行被放大 6 倍，使其天然获得更多路由；这些专家就是 QB 要降温的对象。
注意：子图 (c)(d)（SiTU-GLU 曲线 / RMSNorm 缩放）使用固定的内部配置，不响应 n_experts / top_k / hot_experts 等路由参数的调节；这两个子图展示的是激活函数与归一化本身的性质。
左上 (a)：Quantile Balancing 前后负载
右上 (b)：分位数估计 exact vs histogram
左下 (c)：SwiGLU vs SiTU-GLU 激活幅值
右下 (d)：RMSNorm 尺度不敏感性
这是读者接触模块的第一个综合印象——四个子图各代表一处核心创新，看完这 Tab 就知道 LatentMoE 到底解决了哪四个问题、每个问题的效果量级有多大。它是后续深入 Tab 的钩子和索引。
验证 StableLatentMoE 层的前向输出契约：形状正确、路由结构存在、与独立 Eq.11 参考实现逐点一致、梯度只流向被选中专家。这是模块实现正确性的主锚点——所有其他性质的验证都建立在前向算对了的基础上。
对应测试 est_forward_shape_and_routing / est_matches_an_independent_eq11_reference / est_only_selected_experts_receive_gradient。
中 (b)：Eq.11 参考实现 vs MoE 对拍
这是实现正确性的底线验证。如果前向契约不成立——形状错了、输出与公式对不上、梯度不稀疏——后续所有性质验证都没有意义。此 Tab 是整个演示的信任锚点，建议每次改动代码后先跑它确认基础正确。
直观展示 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）。
注意：此 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 滑块真正影响输出。
Width 分离是 LatentMoE 的灵魂——没有它，896 个专家的参数和计算量会爆炸到不可接受。此 Tab 把「路由专家在窄带、共享专家在宽带」这件事用柱状图直观展示出来，是理解 K3 为什么能支撑近千专家的第一个关键点。
验证 §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 的信号。
这是「Stable LatentMoE」之所以叫 Stable 的第一根支柱。没有 RMSNorm，路由分支的输出尺度会随路由决策剧烈波动，在大模型规模下会引发数值失稳。RMSNorm 不是可选项——论文指出它本身也改善验证损失。在生产中，如果发现训练不稳定，第一个要排查的就是 RMSNorm 是否被意外关闭。
验证稀疏 MoE 的核心性质——只有被路由选中的专家会获得梯度，未被选中的专家梯度为 0。这是 MoE 能节省训练计算量的本质原因：每一步只有 top_k 个专家更新权重。