6 基建与训练稳定性优化 / 让 1.6T 模型真的能跑起来

前几节讲了"算法层"如何省 FLOPs、压 KV、稳训练。但 1.6T 参数 × 32T+ token 的训练， 真正的拦路虎是工程层的细节——通信能不能藏进计算、低精度会不会崩、loss spike 怎么救、 KV cache 怎么落盘。本节把 V4 论文 §3 / §4.2.3 / §5.2 里 8 个"工程亮点"一次性讲完， 每一个单独看都不大，合起来就是"能不能跑得动"的差别。

本节速览 8 个工程亮点分成三组：①推理/训练速度（Fine-Grained EP、FP4 QAT、Hash Routing、MTP）· ②训练稳定性（Anticipatory Routing、SwiGLU Clamping、Attention Sink）· ③显存与服务（异构 KV Cache + On-Disk Cache）。下面用 tab 切换查看。

6.1 Fine-Grained EP Scheme（MegaMoE Kernel）

MoE 推理在多卡间的标准流程是三步串行：

dispatch (token 分发到对应 expert 所在的卡) → compute (各卡跑自己的 expert) → combine (结果收回原卡)

传统实现里这三步严格串行——通信时 GPU 在等，计算时网卡在等。 MegaMoE 把三步融合成一个 CUDA kernel，按"波次"（wave）调度专家： 每波处理一小部分专家，第 k 波在 combine 时，第 k+1 波已经在 dispatch、第 k-1 波的 compute 还在跑。

通信-计算重叠时序对比

Naive (串行) Comet (双流) MegaMoE (波次调度) dispatch compute combine 1.00× t=0 end dispatch-1 dispatch-2 compute-1 compute-2 comb-1 comb-2 ~1.3× ... ... 1.50–1.73× (rollout: 1.96×) dispatch compute

在同样的 token 总量下，MegaMoE 通过波次调度把通信完全藏进上一波的计算里。

能 hide 多少计算？计算-通信比公式

C / B ≤ 2 d = 2 × 3072 = 6144 FLOPs / Byte

这条不等式来自 V4 论文 §3.1：要让通信完全被计算覆盖，每 1 GB/s 的可用带宽就够"喂" 约 6.1 TFLOP/s 的计算。换句话说，只要专家计算量足够大，bandwidth 就不再是瓶颈。

已开源 MegaMoE kernel 已通过 DeepGEMM PR #304 开源，社区可直接复用到自己的 MoE 推理栈。

6.2 FP4 QAT（Quantization-Aware Training）

主流的 FP4 部署都是 PTQ（post-training quantization）——训完用 BF16 / FP8，再 round 到 FP4。 V4 反过来做：训练时就让模型"知道"自己最终会跑在 FP4 上，把量化误差当成训练信号来 fit。

用在哪两处？

  • MoE 专家权重——参数量最大的一块（~90% 的总参数）
  • Lightning Indexer 的 Q/K 投影——CSA 索引器的算路径

两处恰好都是"反复读、读得多"的部分——显存和带宽收益最大。

为什么不是 Attention 主体？

主 Attention 的 Q/K/V/O 投影对精度敏感（CSA 的 score 一旦失真，topk 就错）。 V4 把 FP4 用在对噪声不敏感、反向梯度光滑的位置。

训练时的数据流：FP32 主权重 → FP4 → FP8 计算

优化器 FP32 master W 量化 FP32 → FP4 (E2M1) 反量化（无损） FP4 → FP8 前向计算 FP8 × FP8 GEMM 反向 STE：梯度跳过量化层，直通 FP32 主权重

关键数学：FP4 → FP8 反量化为何无损？

FP8（E4M3 / E5M2）有 4 或 5 位指数，能表示的动态范围远超 FP4（E2M1，仅 ±6 的可表示值）。 因此 FP4 的 16 个离散值乘上每个 block 的 scale factor 后，结果都能被 FP8 精确容纳， 量化只发生在 FP32→FP4 这一步，FP4→FP8 是数学上的恒等映射。

Blackwell 硬件红利

B200/GB200 原生支持 FP4 × FP8 mixed-precision matmul。

  • 现阶段 V4 的 FP4 主要省 显存 / 带宽（参数体积砍掉一半）
  • 未来开启 FP4×FP8 kernel 后，理论上比 FP8×FP8 再快 约 1/3

6.3 Hash Routing（前 3 层 MoE）

config 里有一行不起眼的：

num_hash_layers = 3

意思是：前 3 层 MoE 不走可学习的路由网络，而是用一张固定的 hash 表把 token ID 直接映射到 expert ID。

普通路由

  1. 过 router MLP：scores = W·h
  2. top-k 选 expert
  3. 计算负载均衡 loss

每次前向都要算 → 有计算开销 + 训练不稳定时 routing 会抖。

Hash 路由（前 3 层）

  1. expert_id = hash(token_id) % n_experts
  2. O(1) 查表，无 router 参数
  3. 不需要 balance loss

来自 Roller et al. 2021 (Hash Layers)——简单粗暴但极快。

为什么只用在前 3 层？ 底层主要做"token 级表征"（subword 拼合、词性、局部句法），路由模式相对稳定， 固定 hash 既省 router 计算又避开训练早期负载均衡问题。 上层是"语义/任务级表征"，路由必须动态，仍用 sigmoid + balance loss。
易混淆术语 Hash Routing 和 V3/V4 论文里的 Engram（哈希记忆表）不是一回事—— Engram 是 KV cache 的内容寻址机制，只是名字里都带"哈希"。

6.4 Anticipatory Routing（防 loss spike 的黑科技）

MoE 训练最经典的崩法：

  1. router 参数在某 step 突然变了一点
  2. 大量 token 涌入之前没怎么训过的冷门 expert
  3. 该 expert 输出与残差流分布不匹配 → loss 瞬间飙升 10 倍
  4. 梯度爆炸 → 整个训练直接玩完

核心思路：路由和计算"解耦时间步"

t-Δt ... t (current) 路由决策 用 θ_{t-Δt}（旧） 专家计算 用 θ_t（最新） 路由"慢半拍"，专家不变 检测 loss spike → 启用 ; 稳定后 → 自动关闭

这样做的好处：router 不会因为本步的剧烈更新就把 token 倒进冷门 expert—— 因为路由决策来自几步之前已经被验证"稳定"的 θ。

收益

几乎完全消除 loss spike 引发的训练崩溃，1.6T 模型不用每次 spike 都 rollback 回 checkpoint。

代价

spike 期间额外开销 ~20%（要保留旧 router 参数副本 + 多一次路由计算）。

自适应

只在监测到 loss spike 时启用，日常训练 0 开销——平均吞吐几乎不受影响。

6.5 SwiGLU Clamping

SwiGLU 是现代 LLM 的标配（LLaMA、Gemma、Qwen、DeepSeek 全在用）：

SwiGLU(x) = SiLU(Wg x) ⊙ (Wv x)

问题是输出可以任意大——两个线性分量逐元素相乘，没有上界。 超大模型训到几万亿 token 时，个别 token 的激活值能飙到几百，FP4 直接溢出 （E2M1 只有 ±6 的可表示范围），梯度也跟着爆。

V4 的做法：硬裁剪两个分量

SwiGLU(x) = SiLU(clamp(Wg x, max=10)) ⊙ clamp(Wv x, -10, 10)

视觉对比：clamp 前后

x y 原始 SwiGLU（无界） → FP4 溢出 x y Clamped SwiGLU（V4） y ≤ 10（FP4 安全区）
一举两得：①梯度爆炸没了 ②数值落进 FP4 安全范围，QAT 训练不会因为长尾激活而崩。 代价：极少数 token 会丢一点精度，但全局 loss 反而更稳。

6.6 Attention Sink

标准 softmax 的死结：所有权重必须加和为 1。 即便所有 token 都"无关"，softmax 还是被迫把这 1 单位注意力分配出去—— 结果就是模型必须"看"一些不该看的东西。

标准 softmax

wi = exp(si) / Σj exp(sj)

权重必和为 1 → 长上下文中无关 token 被迫分到注意力 → 信号被噪声稀释。

+ Sink Logit（V4）

wi = exp(si) / (Σj exp(sj) + exp(ssink))

多一个可学习的 sink logit——无效注意力可以"弃权"倒进水槽。

直觉：弃权票箱

原本 100% 注意力必须分给 N 个 token；现在有了 sink，模型可以决定"只给前 5 个 token 分 60%，剩下 40% 全部弃权"。 在百万级 token 上下文里，真正相关的可能只有几百个，sink 让模型有说"我不在乎"的权利。

OpenAI GPT-5 系列也采用了类似设计，这已经成为长上下文模型的标配。

6.7 MTP（Multi-Token Prediction）

config 行：

num_nextn_predict_layers = 1

含义：在主预测头之外，额外加 1 个 MTP head，同时预测 下下一个 token（t+2 位置）。

训练时

  • 主 head 预测 t+1（标准 NTP loss）
  • MTP head 预测 t+2（辅助 loss）
  • 强迫模型在 t 步的 hidden state 里已经准备好 t+1 和 t+2的语义

推理时

  • 可以丢掉 MTP head（不影响主预测）
  • 也可以保留——作为推测解码（speculative decoding）的 draft 模型，大幅加速
为什么 depth=1？ V3 的 ablation 显示 depth=1 收益/开销比最优；depth=2/3 的额外收益边际递减， 但训练显存和计算成倍上升。V4 沿用 V3 结论。

6.8 异构 KV Cache 管理 + On-Disk KV Cache

V4 的 KV 来源很复杂：SWA（局部最近 128 token）、CSA（压缩 KV）、HCA（更深层压缩）。 一刀切的 PagedAttention 显然不合适——V4 用了分类型异构存储。

State Cache

  • 存 SWA 最近 128 token 的原始 KV
  • 外加待压缩的 tail buffer
  • 每请求固定大小块——好分配、好回收

KV Cache

  • 存 CSA / HCA 压缩后的 KV entries
  • 按 lcm(m, m') 对齐分块（适配多种压缩粒度）
  • 跨请求可共享（前缀复用）

On-Disk KV Cache：把"压缩前缀"扔到磁盘

服务化场景里大量请求共享同一段系统 prompt / 工具描述 / few-shot 示例。 V4 把这些共享前缀的压缩 KV持久化到磁盘，新请求来时直接读盘，免去重复 prefill。

SWA KV 的三种磁盘策略

策略磁盘占用命中延迟适用场景
Full Caching 最大（每层 SWA 都存） 最低（直接读） 磁盘充裕、延迟敏感
Periodic Checkpointing 中等（每 N 层存一次） 中等（部分需重算） 权衡折中
Zero SWA Caching 最小（只存压缩 KV） 最高（SWA 全部重算） 磁盘紧张、长前缀

SWA 部分按 token 数线性增长，是落盘成本的大头——三种策略本质都在权衡存储 vs prefill 计算。

6.9 总览：8 个工程亮点一表速查

# 名称 解决什么问题 关键收益
6.1 Fine-Grained EP
(MegaMoE)		 MoE 推理 dispatch/compute/combine 串行，GPU 等通信 1.50–1.73× 加速（rollout 1.96×）
6.2 FP4 QAT PTQ 量化丢精度，FP8 还是太贵 MoE 权重砍半，未来 FP4×FP8 再快 ~33%
6.3 Hash Routing 底层 MoE 路由开销大、训练早期负载抖 前 3 层 O(1) 路由，省 router 计算
6.4 Anticipatory Routing router 突变 → loss spike → 训练崩溃 spike 自动救火，~0% 日常开销
6.5 SwiGLU Clamping 激活值飙升 → 梯度爆炸 + FP4 溢出 梯度稳定，QAT 友好
6.6 Attention Sink softmax 必和为 1，被迫"看"无关 token 长上下文降噪，百万 token 可用
6.7 MTP (depth=1) 表征不够 forward-looking 主任务收益 + 推测解码 draft 模型
6.8 异构 KV + On-Disk SWA/CSA/HCA 混合 KV 难统一管理；共享前缀重复 prefill 显存利用率↑、TTFT↓（命中时）
合起来看 任意一条单独拿出来都不算"震撼"，但放在 1.6T × 32T+ token 的训练上： 没有 6.1 推理就崩；没有 6.2 显存装不下；没有 6.4 / 6.5 训练会反复 spike； 没有 6.6 长上下文跑不动；没有 6.8 服务化撑不起来。工程的本质，就是把这些"小事"全部做对。