mini-infer系统实战-21-Expert Sharding：真正的收益不只是“把权重切一半”

做完 Phase 17 之后，mini-infer 已经有了一版能跑、能测、也能讲清楚的 synthetic MoE + 2-GPU Expert Parallelism 原型。它已经回答了三个关键问题：

1. dense MoELayer 能不能作为数值 oracle
2. 2 卡 EPMoELayer 的 dispatch / gather / all-to-all / combine 能不能对齐 dense
3. 在一个正式 synthetic workload 上，2 卡 EP 能不能明确快过 1 卡 dense

Phase 17 的正式 benchmark 结果是：

• dense：22622.14 tok/s
• EP：42778.41 tok/s
• EP / dense = 1.891x
• max_abs_diff = 0.000000

看起来已经很像“Phase 18 可做可不做”了。但如果目标不是做一个 demo，而是朝真正的推理系统能力线靠，这一版还有一个很硬的缺口：

• 每个 rank 仍持有完整 expert 权重

这意味着你虽然已经有了 EP 的通信生命周期，但还没有真正得到 expert-sharded EP。
Phase 18 要解决的，就是把这个缺口补成一个正式闭环。

这篇文章讲的不是“MoE 是什么”，而是一个更具体的问题：

当你已经有了一版能工作的 EP 原型，为什么“把 expert 权重真正切到每个 rank 本地”并不只是少一半参数那么简单？

问题定义

Phase 18 想回答的其实是三件很具体的事情：

1. EPMoELayer 能不能从“每个 rank 都有完整 experts”收紧成“每个 rank 只持有本地 expert shard”
2. 收紧权重所有权以后，dense vs EP 的数值还能不能保持严格对齐
3. 真分片以后，2 卡 EP 的 throughput 会不会塌掉

对应的正式验收标准也很明确：

• ep_rank_param_bytes / dense_param_bytes <= 0.52
• EP / dense throughput >= 1.5x
• max_abs_diff < 1e-4

这三个条件里，最容易被低估的是第一个。
因为“每个 rank 参数少一半”听起来像一个静态结构变化，但真正一做就会发现，它会同时碰到三层问题：

1. 权重 key 的所有权和重编号
2. worker 初始化路径
3. benchmark 的真实运行入口

如果这三层没有一起收好，你很可能会得到一种看似“已经分片”的状态：

• 代码里 experts 数组变短了
• 但 worker 启动时还是把完整 shard 打包传进去
• 最后 benchmark 根本跑不起来

这也是 Phase 18 最有价值的部分：它不是又加了一个模块，而是把 权重所有权、初始化方式和 benchmark 可运行性 串成了一件事。

方案设计

这一版方案有三个核心决策。

1. EPMoELayer 必须真正只实例化本地 expert shard

在 Phase 17 里，EP 原型已经把 dispatch / gather / all-to-all 这些主链路做对了，但每个 rank 仍然保留完整 expert 列表。这样做的好处是简单，坏处也很明显：它没有真实的内存收益。

Phase 18 的第一步，就是把 EPMoELayer 从“完整 experts + local mask”收紧成“只实例化本 rank 对应的 local expert shard”。

对应代码在 [moe_layer.py](https://github.com/psmarter/mini-infer/blob/main/mini_infer/moe_layer.py) 里，核心语义是：

• experts_per_rank = num_experts // ep_size
• local_expert_offset = rank * experts_per_rank
• 每个 rank 只构造 experts_per_rank 个 MoEFFNExpert

这一步的意义不是节省几行代码，而是把“expert 属于谁”这件事变成模型对象本身的真实语义。
从这一步开始，Phase 18 的 EP 才算真正进入 expert-sharded 状态。

2. state_dict 切分必须发生在 parent 进程，而不是 worker 内临时推导

一旦每个 rank 只保留本地 experts，马上就会遇到一个问题：

dense MoELayer.state_dict() 里的 expert key 还是全局编号，worker 怎么知道自己该加载哪部分？

这就是 [moe_layer.py](https://github.com/psmarter/mini-infer/blob/main/mini_infer/moe_layer.py) 里 shard_moe_state_dict() 的意义。

这段逻辑做了两件事：

1. 把 experts.<global_id>.* 重新映射到每个 rank 的局部 experts.<local_id>.*
2. 把 router 等非 expert 权重复制到每个 shard

也就是说，Phase 18 不是简单把 tensor 按 rank 切掉，而是连 权重 key 空间 一起切成了 per-rank local view。

这件事必须在 parent 进程里显式完成，而不是让 worker 启动后再去做复杂推导。因为一旦把“切片规则 + 初始化逻辑 + 分布式 rank 上下文”混到 worker 内部，排错会立刻变得很痛苦。

3. benchmark 不能只测“模型数学对不对”，还必须测“worker 能不能起来”

这轮最意外、但也最工程化的坑，不在 MoE 数学路径，而在 mp.spawn。

一开始我已经把 Phase 18 的 shard 所有权和参数统计都做完了，相关测试也能过。但一跑正式 benchmark，就会在 worker 启动前直接炸掉：

ValueError: bad value(s) in fds_to_keep

这个错误的根因不是 all-to-all，也不是 expert combine，而是：

• EPEngine 把整包 rank_state_dicts 直接作为 mp.spawn 参数传给 worker
• 到正式 benchmark 规模时，这条启动路径在当前执行环境里触发了 fds_to_keep 问题

这件事很重要，因为它说明一个经常被忽略的事实：

对 EP 来说，worker 初始化路径本身就是系统设计的一部分，不是“启动细节”。

如果 benchmark 根本起不来，那前面所有“参数收缩”“数值等价”的结论都还没真正进入可交付状态。

实现细节

1. local shard 不是只改 experts 数组长度

Phase 18 真正的收紧点在于：EPMoELayer 的 expert 语义、state_dict 的 key 语义、以及 dispatch 后 _apply_experts() 的 expert id 解释方式，三者必须一起对齐。

如果只把 experts 数组长度减半，但 _apply_experts() 仍按全局 id 访问，就会立刻出错。
所以这一版的关键不是“少构造几个模块”，而是让全局 expert id 到局部 expert id 的映射在每条路径上都一致。

对应代码里，最核心的概念是：

• local_expert_offset
• local_expert_ids(rank)
• valid_expert_ids

这几项决定了 distributed forward 时，每个 rank 最终到底会执行哪些 expert。

2. EPEngine 的真正修复点是 worker-init 下发方式

最开始的 Phase 18 版本里，EPEngine 还是直接把 rank_state_dicts 当作 mp.spawn 参数传进去。这在小配置测试下问题不明显，但到了正式 benchmark workload，就会在启动阶段爆掉。

最终稳定下来的修复不是去碰 mp.spawn 的 start method，而是把下发方式改成：

1. parent 进程先把每个 rank 的 local shard torch.save() 到临时目录
2. worker 只接收 rank_state_dict_dir
3. worker 启动后按自己的 rank 读取对应文件
4. 整个 run 结束后统一清理临时目录

这条路径在 [ep_engine.py](https://github.com/psmarter/mini-infer/blob/main/mini_infer/ep_engine.py) 里很直接：

temp_dir = tempfile.mkdtemp(prefix="mini_infer_ep_")
result_file = os.path.join(temp_dir, "result.pt")
rendezvous_file = os.path.join(temp_dir, "rendezvous")
rank_state_dict_dir = os.path.join(temp_dir, "rank_state_dicts")

_dump_rank_state_dicts(self.rank_state_dicts, rank_state_dict_dir)
mp.spawn(..., args=(..., rank_state_dict_dir, ...))

worker 侧对应的是：

rank_state_dict = torch.load(
    _rank_state_dict_path(rank_state_dict_dir, rank),
    map_location="cpu",
)

从功能上看，这只是把“直接传 tensor”换成了“先写文件再读文件”。
但从工程上看，它解决的是更关键的问题：正式 benchmark 入口终于稳定了。

3. benchmark 里必须把“参数收益”和“通信现实”同时说清楚

Phase 18 的 benchmark 文件是 [benchmark_moe.py](https://github.com/psmarter/mini-infer/blob/main/benchmarks/benchmark_moe.py)，这一版里最重要的不是 throughput 本身，而是它把两类结论分开了：

1. 权重参数量
2. 通信字节量

参数量侧，正式输出有四项：

• dense_param_bytes
• ep_rank_param_bytes
• expert_param_bytes
• shard_ratio

通信侧，仍然保留：

• tp_bytes_per_layer
• ep_ideal_bytes_per_layer
• ep_prototype_bytes_per_layer

这两个维度必须同时存在。否则你很容易写出一种误导性的结论：

• “每个 rank 参数量已经减半”
• 但完全不提当前通信仍然是 padded prototype

这会把内存闭环和通信优化混成一件事。Phase 18 最终稳定下来的口径，是把这两件事明确拆开：

• 本阶段已经完成：权重所有权 + 参数量收缩
• 本阶段还没完成：逼近 ideal EP 通信口径

实验结果

下文使用的数据来自该阶段的正式 benchmark 记录。

环境：

• Ubuntu 24.04
• torch 2.1.2+cu121
• RTX 4090 × 2
• ai-infra Python 环境

正式 workload：

• batch_size=4
• seq_len=16
• hidden_size=512
• intermediate_size=1024
• num_experts=8
• top_k=2
• dtype=float16
• warmup=2
• runs=5
• src_rank=1

正式命令是：

python benchmarks/benchmark_moe.py \
  --compare \
  --batch-size 4 \
  --seq-len 16 \
  --hidden-size 512 \
  --intermediate-size 1024 \
  --num-experts 8 \
  --top-k 2 \
  --dtype float16 \
  --warmup 2 \
  --runs 5 \
  --src-rank 1

同配置复跑 2 次的结果是：

指标	Run 1	Run 2	平均
dense throughput	22005.26 tok/s	22920.32 tok/s	22462.79 tok/s
EP throughput	42730.35 tok/s	43341.68 tok/s	43036.02 tok/s
EP / dense	1.942x	1.891x	1.916x
max abs diff	0.000000	0.000000	0.000000
dense param bytes	50348032	50348032	50348032
ep rank param bytes	25182208	25182208	25182208
shard ratio	0.5002	0.5002	0.5002

通信和负载相关的辅助输出是：

• ep_send_counts = [69, 59]
• dense_expert_loads = [20, 12, 23, 14, 19, 15, 14, 11]
• ep_expert_loads = [20, 12, 23, 14, 19, 15, 14, 11]
• tp_bytes_per_layer = 131072
• ep_ideal_bytes_per_layer = 262144
• ep_prototype_bytes_per_layer = 524288

这组结果说明了三件事：

1. 参数收缩目标通过了
   ep_rank_param_bytes / dense_param_bytes = 0.5002 <= 0.52

2. 吞吐 gate 通过了
   EP / dense = 1.916x >= 1.5x

3. 数值一致性仍然非常干净
   max_abs_diff = 0.000000

也就是说，Phase 18 不是只拿到了“参数少一半”的静态结论，而是拿到了 参数收缩、吞吐保持、数值对齐 三条同时成立的正式 benchmark 结果。

这阶段真正踩到的坑

坑 1：真正的 benchmark blocker 不一定在算子里，可能在 worker 初始化路径

这是整个 Phase 18 最值得记录的地方。

一开始我以为这轮最大的风险会是：

• shard 以后 dense / EP 不再数值对齐
• 或者分片后吞吐掉到 gate 以下

结果真正把阶段卡住的，是 benchmark 根本起不来。
正式 workload 下直接报：

ValueError: bad value(s) in fds_to_keep

这件事的教训很直接：

在分布式推理系统里，worker 初始化路径本身就是主链路的一部分。

如果你只盯着 forward 数学、all-to-all 和 throughput，很容易忽略掉“正式 benchmark 配置下，进程到底能不能被稳定拉起来”这件事。

坑 2：参数收缩和通信优化必须分开讲

Phase 18 最容易写错的一种总结是：

我已经做了 true expert sharding，所以 EP 实现更接近生产级了。

这句话只对一半。

正确的说法应该是：

• 权重所有权层面：已经是真正的 per-rank local shard
• 通信实现层面：仍然是 padded prototype，不是 ideal EP

如果不把这两层拆开，你很容易高估当前实现的“生产接近度”。

坑 3：shard_ratio 只是参数量统计，不是完整显存峰值

这一阶段的 benchmark 里：

• throughput 是真实测量值
• TTFT / TPOT / peak memory 都是 N/A

所以 shard_ratio=0.5002 的含义必须说清楚：它表示的是 参数量层面的收缩，不是整个系统运行期显存峰值已经精确减少到一半。

对技术博客来说，这种边界一定要写出来。否则读者会把“参数字节数减半”误读成“整条服务链路的 GPU 内存峰值也已经减半”。

总结

Phase 18 的价值，不在于它又多实现了一个 MoE 模块，而在于它把一个很容易停留在“概念上已经做了”的东西，真正做成了工程上可验证的闭环。

这轮最终交付的是：

• EPMoELayer 的真实 local expert shard 语义
• EPEngine 的 rank-local shard 下发与 worker 初始化修复
• 参数量统计、通信口径和正式 benchmark 的统一输出
• 一组通过硬 gate 的正式结果：
  • shard_ratio = 0.5002
  • EP / dense = 1.916x
  • max_abs_diff = 0.000000

如果说 Phase 17 解决的是：

“MoE 的 dispatch / gather / all-to-all 能不能做成一个正确、可测的 EP 原型？”

那么 Phase 18 解决的就是：

“这个 EP 原型能不能真正进入 expert-sharded 状态，并且不把 benchmark 和 worker 初始化一起弄坏？”

对推理系统来说，后一个问题其实更接近真实工程。

---

系列导航

• 导读：mini-infer系统实战-00-导读：从最小推理链路到 MoE Expert Parallel 的项目路线
• 上一篇：mini-infer系统实战-20-MoE 推理主线：关键不在 Router，而在 Dispatch 与 All-to-All
• 下一篇：mini-infer系统实战-22-Packed Communication：难点不是 bytes 公式，而是实现和 benchmark 口径一起做对