MoE feed-forward performs ~3×num_experts device→host syncs per layer, per forward

Summary

MOEFeedForward.forward (model/model_minimind.py:163-170) loops over experts and, for each one, executes three operations that each force a device→host synchronization:

for i, expert in enumerate(self.experts):
    mask = (topk_idx == i)
    if mask.any():                                    # 1) .item() on a GPU bool -> D2H round trip
        token_idx = mask.any(dim=-1).nonzero().flatten()   # 2) nonzero() must read the count back
        weight = topk_weight[mask].view(-1, 1)             # 3) boolean-mask indexing, same reason
        y.index_add_(0, token_idx, (expert(x_flat[token_idx]) * weight).to(y.dtype))

Each sync drains the CUDA stream: the GPU idles until the CPU has read the value and enqueued the next work. The cost is paid per expert, per layer, per forward, and it is spent on routing bookkeeping rather than on any expert math.

README.md:566 attributes the MoE slowdown to kernel launch and scheduling overhead (kernel 启停和调度开销), and suggests that optimizing it requires a fused-MoE operator library — Triton custom kernels, DeepSpeed-MoE, Megatron-LM — which the project declines in order to stay native-PyTorch, accepting ~50% slower than dense at 4 experts / top-1.

Host syncs are a separate cost from launch overhead, and this one is removable without leaving native PyTorch. I am not claiming it accounts for the whole 50%; I am pointing out that a measurable part of it is bookkeeping that does not need to be there.

Measurement

Counted with PyTorch's own instrumentation, torch.cuda.set_sync_debug_mode('warn'), on one MOEFeedForward layer (hidden 512, batch 8 × seq 512, top-1):

num_experts	syncs per layer per forward	3E
4 (default)	12	12
8	24	24
16	48	48

Exactly 3E. At the default 4 experts × 8 layers that is 96 syncs per forward; at 16 experts, 384. The same loop runs under eval() — there is no separate inference path — so generation pays it too.

Reproduce:

import warnings, torch
from model.model_minimind import MiniMindConfig, MOEFeedForward

cfg = MiniMindConfig(hidden_size=512, intermediate_size=1408, moe_intermediate_size=1408,
                     num_experts=8, num_experts_per_tok=1, use_moe=True)
m = MOEFeedForward(cfg).cuda().train()
x = torch.randn(8, 512, 512, device='cuda')
m(x); torch.cuda.synchronize()                      # warm up

torch.cuda.set_sync_debug_mode('warn')
with warnings.catch_warnings(record=True) as w:
    warnings.simplefilter('always')
    m(x)
torch.cuda.set_sync_debug_mode('default')
print(sum('sync' in str(r.message).lower() for r in w))   # -> 24

Why it matters more on some machines than others

A sync costs a host round trip, so the price is set by CPU and PCIe latency rather than by the GPU. This makes the MoE slowdown strongly machine-dependent in a way that is easy to misread as measurement noise: on identical GPU models on two different hosts, the same isolated-layer benchmark differed by 2–3×.

Rather than quote a speedup that will not reproduce on your machine, here is the ruler. Measure your own host's sync cost:

import time, torch
flag = torch.zeros(1, dtype=torch.bool, device='cuda')
for _ in range(20): flag.any().item()                 # warm up
torch.cuda.synchronize(); t0 = time.perf_counter()
for _ in range(100): flag.any().item()
torch.cuda.synchronize()
print(f'{(time.perf_counter() - t0) / 100 * 1e6:.1f} us per device->host round trip')

Multiply by 3 × num_experts × num_hidden_layers for a rough floor on what the current loop spends per forward on routing bookkeeping alone. On a WSL2 laptop I measure ~120 µs per round trip, which at the default 4 experts × 8 layers is on the order of 10 ms per forward; a bare-metal Linux host is typically a few µs and the same arithmetic gives well under 1 ms. Both are pure overhead, and both go away.

It is removable in plain PyTorch

The loop needs one piece of host-side information: how many tokens each expert received. That can be computed once per layer instead of 3E times — sort the tokens by expert, take the group boundaries with searchsorted, and read them back in a single .tolist(). The expert loop itself stays, and each expert receives a contiguous slice:

flat = topk_idx.reshape(-1)
idx_s, order = torch.sort(flat, stable=True)      # stable -> tokens keep their original order within a group
ends = torch.searchsorted(idx_s, torch.arange(E, device=dev, dtype=flat.dtype), right=True)
tok = order.div(k, rounding_mode='floor')
xs = x_flat.index_select(0, tok)
bound = [0] + ends.tolist()                        # the only device->host sync in the layer
for i, expert in enumerate(self.experts):
    lo, hi = bound[i], bound[i + 1]
    if hi > lo: outs.append(expert(xs[lo:hi]))     # plain Python ints; no further syncs

This is 3E → 1, independent of expert count, with no new dependency, no custom kernel, and no hardware requirement.

Because stable=True preserves token order inside each group, the result is bit-identical to the current loop, not merely close: I verified torch.equal on the output, the input gradient, and every parameter gradient, across experts ∈ {1,2,4,8} × top-k ∈ {1,2} × train/eval on CPU, and in fp32 and bf16 autocast on CUDA.

That extends to whole training runs. Over 12 steps of a real loop (AdamW, gradient clipping, aux loss, top-k=2), the loss was equal to the last printed digit at every step and all 45 parameter tensors were torch.equal afterwards:

step                loop              sorted       delta
   0      6.256397724152      6.256397724152    0.00e+00
   1      6.267616271973      6.267616271973    0.00e+00
  ...
  11      6.292136669159      6.292136669159    0.00e+00

(Caveat: that run is on CPU. At top-k > 1 the final index_add_ uses atomics on CUDA, so run-to-run bit-exactness there is no better — and no worse — than the current implementation's own.)

Two details worth preserving for anyone attempting this:

• The elif self.training: y[0,0] += 0 * sum(...) branch is load-bearing. DDP is constructed without find_unused_parameters (trainer/train_pretrain.py:154), so experts that receive no tokens still need to appear in the autograd graph. It looks like dead code and is not.
• Deciding emptiness must not reintroduce a sync — compare the host-side ints from bound, not a GPU tensor.

What I would propose

Since the two are bit-identical, I would suggest replacing the masked loop rather than adding a config flag beside it. A flag defaulting to off would add a config key, a branch and a second code path — three things this repo pays for in readability — while delivering the benefit to nobody by default.

There is a reasonable objection that the masked loop is easier to read, and it is your call. I would gently argue the other way: sort-by-expert with contiguous per-expert slices is how MoE dispatch is actually implemented in Megatron-LM, DeepSpeed-MoE and essentially every production stack, whereas the masked loop is a pattern nobody ships. For a repo whose premise is learning the real thing from scratch, the sorted version may be the better teaching artifact — with a comment block explaining the sort, and the masked version kept as an annotated reference implementation in the tests.

Happy to open a PR either way. I would rather agree on the shape first than send an unsolicited diff.

One thing this does not fix

torch.compile (the repo's --use_compile) still breaks the graph on this layer. I checked with fullgraph=True: the masked loop fails on Data-dependent branching, and the sorted version fails on the data-dependent shape coming out of .tolist(). Removing the syncs does not make the layer compile cleanly, and I do not want to imply otherwise. (Note also that torch._dynamo.utils.counters['graph_break'] reports 0 for the masked loop even though fullgraph=True rejects it — the counter is not a reliable way to check this.)

Background

This came out of work on S2-MoE-llm (DOI 10.5281/zenodo.20846758), where the sorted-dispatch pattern was developed and validated. I noticed the same host-sync pattern here while reading minimind's MoE implementation.

Environment

torch 2.11.0+cu128, transformers 5.16.1, RTX 5070 Laptop (sm_120) and RTX 5090 (sm_120). The sync counts above are exact integers and are not hardware-dependent.

---

---

MoE 前馈层每层每次前向会产生约 3×num_experts 次 device→host 同步

摘要

MOEFeedForward.forward（model/model_minimind.py:163-170）在专家循环里，每个专家都会执行三个强制 device→host 同步的操作：

for i, expert in enumerate(self.experts):
    mask = (topk_idx == i)
    if mask.any():                                    # 1) 对 GPU 上的 bool 取值，一次主机往返
        token_idx = mask.any(dim=-1).nonzero().flatten()   # 2) nonzero() 需要把元素个数读回主机
        weight = topk_weight[mask].view(-1, 1)             # 3) 布尔索引，同理
        y.index_add_(0, token_idx, (expert(x_flat[token_idx]) * weight).to(y.dtype))

每次同步都会清空 CUDA 流：GPU 必须等 CPU 把值读回并重新下发任务。这个代价按专家 × 层 × 前向累计，而且花在路由的簿记上，不是任何专家的实际计算。

README.md:566 把 MoE 的性能损失归因于 kernel 启停和调度开销，并指出要优化它得依赖支持 MoE kernel-fused 的算子库（Triton 自定义 kernel、DeepSpeed-MoE、Megatron-LM）；项目为了保留原生 PyTorch 的普适性没有采用，接受了 4 experts / top-1 比 dense 慢约 50% 的现状。

主机同步与 kernel 启停是两回事，而且这一项不需要离开原生 PyTorch 就能去掉。我并不是说它就是那 50% 的全部；我想说明的是，其中有可测量的一部分是本来就不必存在的簿记开销。

实测

用 PyTorch 自带的 torch.cuda.set_sync_debug_mode('warn') 在单个 MOEFeedForward 层上统计（hidden 512，batch 8 × seq 512，top-1）：

num_experts	每层每次前向的同步次数	3E
4（默认）	12	12
8	24	24
16	48	48

正好是 3E。默认的 4 专家 × 8 层即每次前向 96 次同步，16 专家时是 384 次。该循环在 eval() 下完全相同（没有单独的推理路径），所以生成时同样要付这个代价。

复现代码见上方英文部分。

为什么不同机器上差别很大

同步的代价是一次主机往返，取决于 CPU 与 PCIe 延迟，而不是 GPU。我测过的机器上，单次往返（对 1 元素 GPU 张量调用 .item()）相差约一个数量级——数据中心主机几微秒，WSL2 笔记本约 120 µs。

这会让 MoE 的性能损失呈现出很强的机器依赖性，很容易被误当成测量噪声：同型号 GPU、不同主机上，同一份单层基准的结果相差 2–3 倍。

用纯 PyTorch 就能去掉

循环真正需要从主机侧知道的只有一件事：每个专家分到了多少 token。这个信息每层算一次就够，不必算 3E 次——先按专家把 token 排序，用 searchsorted 取出分段边界，再用一次 .tolist() 读回主机。专家循环本身保留，每个专家拿到一段连续切片（代码见上方英文部分）。

这样是 3E → 1，与专家数无关，不引入新依赖、不写自定义 kernel、对硬件没有要求。

由于 stable=True 保证了组内 token 顺序不变，结果与现有循环逐位相同，而非"在容差范围内"：我用 torch.equal 验证了输出、输入梯度和每一个参数梯度，覆盖 experts ∈ {1,2,4,8} × top-k ∈ {1,2} × train/eval（CPU），以及 CUDA 上的 fp32 与 bf16 autocast。

整条训练轨迹同样如此。在真实训练循环下（AdamW + 梯度裁剪 + aux_loss，top-k=2）跑 12 步，每一步的 loss 到最后一位都相同，训练结束后 45 个参数张量全部 torch.equal（输出见上方英文部分）。

（需要说明：这一项是在 CPU 上验证的。top-k > 1 时最后的 index_add_ 在 CUDA 上使用原子累加，因此那里的可复现性与现有实现一样——不会更好，也不会更差。）

有两个细节值得提醒后来者：

• elif self.training: y[0,0] += 0 * sum(...) 这一支是有实际作用的。DDP 构造时没有传 find_unused_parameters（trainer/train_pretrain.py:154），所以没分到 token 的专家仍然必须出现在自动求导图里。它看起来像死代码，其实不是。
• 判断某个专家是否为空时不能把同步加回来——要比较 bound 里的主机侧整数，而不是 GPU 张量。

我的建议

既然两者逐位相同，我倾向于直接替换掩码循环，而不是在旁边加一个配置开关。默认关闭的开关会引入一个配置项、一个分支和第二条代码路径——这三样正是本仓库在可读性上要付出代价的东西——而默认情况下谁也享受不到收益。

"掩码循环更好读"是一个合理的反对意见，最终由您决定。我想温和地提出相反的看法：按专家排序 + 每个专家取连续切片，正是 Megatron-LM、DeepSpeed-MoE 以及几乎所有生产级实现真正采用的 MoE 分发方式，而掩码循环是一种没有人真正部署的写法。对于一个以"从零学真东西"为前提的仓库，排序版本也许是更好的教学素材——配上一段解释排序的注释，并把掩码版本作为带注解的参考实现保留在测试里。

两种方案我都愿意提 PR。比起直接发一个未经沟通的 diff，我更希望先就形态达成一致。

这件事没有解决的问题

torch.compile（即仓库的 --use_compile）在这一层上仍然会断图。我用 fullgraph=True 验证过：掩码循环报 Data-dependent branching，排序版本报 .tolist() 带来的数据依赖形状。去掉同步并不能让这一层干净地编译，我不想给出相反的暗示。（另外，torch._dynamo.utils.counters['graph_break'] 对掩码循环报 0，而 fullgraph=True 明确拒绝它——这个计数器不能用来判断有没有断图。）

来源

这个发现来自 S2-MoE-llm（DOI 10.5281/zenodo.20846758）项目的工作，排序分发的写法是在那里开发和验证的。我在阅读 minimind 的 MoE 实现时注意到了同样的主机同步模式。

环境

torch 2.11.0+cu128，transformers 5.16.1，RTX 5070 Laptop（sm_120）与 RTX 5090（sm_120）。上表的同步次数是精确整数，与硬件无关。