DeepSpeed MoE and Model Compression

导言

MoE 和模型压缩看似方向相反：前者扩大总参数，后者缩小部署对象。它们实际都在重写“每个 token 访问哪些参数”。DeepSpeed-MoE/MoE Inference 管理稀疏路由和专家放置；Model Compression/MoQ 管理层、权重和精度的删减。


自绘示意图：MoE 改变 token 的动态参数访问，压缩改变实际存在或部署的对象。

DeepSpeed-MoE：容量增长不等于每 token 计算增长

MoE 用 router 为 token 选择少量专家。若有 (E) 个专家、top-(k)，每个 token 只执行 (k) 个 FFN：

[
p(x)=\operatorname{softmax}(W_rx),\qquad
y=\sum_{e\in\operatorname{TopK}(p(x),k)}p_e(x)E_e(x).
]

DeepSpeed-MoE 把专家分到 expert parallel group，通过 All-to-All 派发 token、在本地专家计算，再 All-to-All 返回。[^moe-training]

scores = router(hidden)
expert_ids, gates = topk(scores, k)
dispatch = pack_by_expert(hidden, expert_ids, capacity)
remote_tokens = all_to_all(dispatch, ep_group)
remote_out = local_experts(remote_tokens)
returned = all_to_all(remote_out, ep_group)
output = combine(returned, gates)
loss += load_balance_loss(scores, expert_ids)

三个容易漏掉的对象

• capacity buffer：每个专家接收 token 的上限，太小会 drop，太大浪费显存；
• router auxiliary loss：抑制专家拥塞，但会改变训练目标；
• expert optimizer state：总参数增大后，EP/ZeRO 如何组合决定真正的状态容量。

因此“总参数很大、激活参数很小”不能单独证明训练便宜。All-to-All、负载不均和小专家 GEMM 可能成为关键路径。

MoE Inference：目标从吞吐扩展变成延迟与放置

训练时可以用较大 batch 聚合专家 token；在线推理尤其 decode 时，每步 token 少，专家 GEMM 更碎，All-to-All 更难摊薄。DeepSpeed MoE Inference 将 expert parallel、tensor parallel 和推理 kernel 组合，重点是专家放置与低延迟执行。[^moe-inference]

request tokens
  -> replicated/local router
  -> EP dispatch
  -> TP inside selected expert if configured
  -> EP combine
  -> next transformer sublayer

需要分别测 prefill 和 decode：

• prefill token 多，专家 batching 较容易；
• decode 每序列每步通常一个 token，路由离散和跨 Rank 尾延迟更明显；
• TP 会降低单专家权重容量，但增加专家内部 collective；
• EP 会分散专家容量，但增加 token dispatch。

训练配置可加载不代表它是最优推理布局。

Model Compression：四类变换作用于不同对象

DeepSpeed Model Compression 把 layer reduction、knowledge distillation、sparse pruning 和 quantization 组织成流水。[^compression]

Layer reduction

减少 Transformer 层数，直接缩短串行深度。学生层可从教师层映射或初始化，但层数减少后的表达能力要由蒸馏弥补。

Knowledge distillation

学生拟合标签、教师 logits 或 hidden state：

[
\mathcal L=\lambda_{\mathrm{task}}\mathcal L_{\mathrm{task}}
+\lambda_{\mathrm{KD}}T^2\operatorname{KL}
\lambda_{\mathrm{hidden}}\mathcal L_{\mathrm{hidden}}.
]

教师前向增加训练成本；tokenizer、hidden shape 和层映射是接口约束。

Sparse pruning

结构化稀疏只有在 kernel/硬件真的跳过零块时才产生速度；非结构化零值若仍走 dense GEMM，主要收益可能只是可压缩存储。

Quantization

把权重/激活映射为低 bit：

[
q=\operatorname{clip}\left(\operatorname{round}(x/s)+z,q_{\min},q_{\max}\right).
]

部署收益取决于真实低精度 kernel、scale 粒度和反量化位置。fake quant 的训练成功不等于线上一定加速。

MoQ：逐级降 bit，而不是一步跳到目标精度

Mixture-of-Quantization（MoQ）在 fine-tuning 中从 start_bits 逐步降到 target_bits。可选的 eigenvalue schedule 用二阶敏感度让脆弱层更慢降 bit。[^moq]

for step in training:
    for layer in quantized_layers:
        sensitivity = cached_or_recomputed_eigenvalue(layer)
        bits = schedule(
            step=step,
            start_bits=8,
            target_bits=4,
            period=period,
            sensitivity=sensitivity,
        )
        layer.weight = fake_quantize(layer.weight, bits, groups)
    optimizer_step()

官方 GLUE 教程展示了其设置中的质量结果，但边界同样重要：

• 二阶信息计算会显著增加训练时间；
• eigenvalue 可能出现 NaN/Inf；
• group size、offset、period 和目标 bit 都是模型相关超参；
• QAT 质量不保证目标硬件存在相应 kernel。


自绘示意图：从 router、EP dispatch、expert kernel 到压缩部署的完整链路。

选择与组合

目标	方法	必测指标
增大模型容量、控制每 token FLOPs	DeepSpeed-MoE	expert load、drop rate、A2A、任务质量
部署已有 MoE	MoE Inference	prefill/decode 延迟、EP/TP 放置、尾延迟
形成更小的 dense/sparse 学生	Model Compression	各组件独立消融、真实 kernel、最终质量
让模型逐步适应低 bit	MoQ	bit schedule、敏感度开销、目标硬件端到端性能

MoE 也能继续量化或蒸馏，但必须区分 router、共享层和专家权重的校准分布。把所有专家混用一套 scale，可能掩盖专家间离群值；逐专家量化又会增加 metadata 和 kernel 复杂度。

结论

MoE 改变 token 到参数的动态连接，压缩改变模型中实际存在或实际访问的层、权重与 bit。工程决策应从 token → router → expert → kernel → memory 的整条路径衡量，而不是只比较总参数或权重文件大小。

[^moe-training]: DeepSpeed-MoE training tutorial，教程文件最早提交于 2021-08-17。
[^moe-inference]: DeepSpeed-MoE inference tutorial，教程文件最早提交于 2022-01-18。
[^compression]: DeepSpeed Model Compression tutorial，教程文件最早提交于 2022-07-19；论文：arXiv:2206.01859。
[^moq]: DeepSpeed MoQ tutorial，教程文件最早提交于 2021-05-24。