Merge 6b2f4c939a into cdf0a93e02

model/model_minimind.py

@@ -159,15 +159,30 @@ class MOEFeedForward(nn.Module):
         scores = F.softmax(self.gate(x_flat), dim=-1)
         topk_weight, topk_idx = torch.topk(scores, k=self.config.num_experts_per_tok, dim=-1, sorted=False)
         if self.config.norm_topk_prob: topk_weight = topk_weight / (topk_weight.sum(dim=-1, keepdim=True) + 1e-20)
-        y = torch.zeros_like(x_flat)
+        # 先把 token 按专家排成连续段，每个专家取一段连续切片。
+        # 直觉上更自然的写法是每个专家做一次 mask：`if mask.any(): x_flat[mask.any(-1).nonzero()]`，
+        # 但 mask.any() 的取值、nonzero() 的输出长度、布尔索引的结果长度都必须先回到主机才知道，
+        # 于是每个专家都要 3 次 device→host 同步（每层 3E 次，默认配置下每次前向 96 次）。
+        # 同步会清空 CUDA 流：GPU 停下来等 CPU 读值再重新下发，而这些开销全花在路由的簿记上。
+        # 排序之后每段的边界一次就能算完，整层只剩 ends.tolist() 一次同步，且与专家数无关。
+        # stable=True 保证段内 token 仍按原顺序，因此结果与逐专家 mask 的写法逐位相同
+        #（掩码版本作为参考实现保留在 scripts/test_moe_sorted.py，每次跑测试都会断言两者一致）。
+        e, k = self.config.num_experts, self.config.num_experts_per_tok
+        flat = topk_idx.reshape(-1)
+        idx_sorted, order = torch.sort(flat, stable=True)
+        ends = torch.searchsorted(idx_sorted, torch.arange(e, device=x_flat.device, dtype=flat.dtype), right=True)
+        token_idx = order.div(k, rounding_mode='floor')  # 每个槽位对应的原 token
+        x_sorted = x_flat.index_select(0, token_idx)
+        weight = topk_weight.reshape(-1).index_select(0, order).unsqueeze(1)
+        bounds = [0] + ends.tolist()  # 整层唯一一次 device→host 同步
+        y, outs = torch.zeros_like(x_flat), []
         for i, expert in enumerate(self.experts):
-            mask = (topk_idx == i)
+            lo, hi = bounds[i], bounds[i + 1]
-            if mask.any():
+            if hi > lo:
-                token_idx = mask.any(dim=-1).nonzero().flatten()
+                outs.append(expert(x_sorted[lo:hi]))  # 主机侧的 python int 比较，不会触发同步
-                weight = topk_weight[mask].view(-1, 1)
-                y.index_add_(0, token_idx, (expert(x_flat[token_idx]) * weight).to(y.dtype))
             elif self.training:
                 y[0, 0] += 0 * sum(p.sum() for p in expert.parameters())
+        if outs: y.index_add_(0, token_idx, (torch.cat(outs) * weight).to(y.dtype))
         if self.training and self.config.router_aux_loss_coef > 0:
             load = F.one_hot(topk_idx, self.config.num_experts).float().mean(0)
             self.aux_loss = (load * scores.mean(0)).sum() * self.config.num_experts * self.config.router_aux_loss_coef

scripts/bench_moe.py

@@ -0,0 +1,115 @@
+"""MoE 分发的性能基准 / MoE dispatch benchmark. 自包含，不引入任何新依赖。
+
+  --mode sync-probe   测本机一次 device→host 同步的往返开销
+  --mode layer        单个 MoE 层：当前实现 vs 改动前的掩码写法
+
+为什么要先测同步开销：掩码写法每层要做 3E 次主机往返，一次往返的代价由主机 CPU/PCIe
+延迟决定，而不是 GPU。我实测过从裸机主机的几微秒到 WSL2 笔记本的约 120 微秒，
+所以同一份基准在同型号 GPU、不同主机上可以相差 2-3 倍。先看自己机器的往返开销，
+再看下面的加速比，才知道该期待哪一端。
+
+计时协议：≥10 次预热，计时区间前后 torch.cuda.synchronize()，取中位数并报告 min/max。
+"""
+import argparse
+import os
+import statistics
+import sys
+import time
+
+import torch
+
+sys.path.append(os.path.join(os.path.dirname(os.path.abspath(__file__)), '..'))
+from model.model_minimind import MiniMindConfig, MOEFeedForward
+from scripts.test_moe_sorted import reference_masked_forward
+
+
+def timed(fn, warmup=10, reps=20):
+    for _ in range(warmup): fn()
+    torch.cuda.synchronize()
+    samples = []
+    for _ in range(reps):
+        torch.cuda.synchronize()
+        t0 = time.perf_counter()
+        fn()
+        torch.cuda.synchronize()
+        samples.append(time.perf_counter() - t0)
+    return {'median': statistics.median(samples), 'min': min(samples), 'max': max(samples)}
+
+
+def mode_sync_probe(args):
+    flag = torch.zeros(1, dtype=torch.bool, device='cuda')
+    x = torch.randn(1024, 1024, device='cuda')
+
+    def probe(fn, n):
+        for _ in range(20): fn()
+        torch.cuda.synchronize()
+        s = []
+        for _ in range(args.reps):
+            torch.cuda.synchronize()
+            t0 = time.perf_counter()
+            for _ in range(n): fn()
+            torch.cuda.synchronize()
+            s.append((time.perf_counter() - t0) / n)
+        return statistics.median(s)
+
+    launch = probe(lambda: x.add_(0.0), 200)
+    rt = probe(lambda: flag.any().item(), 100)
+    print(f'  kernel launch (无同步) : {launch * 1e6:8.2f} us')
+    print(f'  device->host 往返      : {rt * 1e6:8.2f} us   <- 掩码写法每层要做 3E 次')
+    e, l = args.experts[0], args.num_hidden_layers
+    print(f'\n  按 3 x {e} experts x {l} layers 估算，仅路由簿记的主机开销下限约 '
+          f'{3 * e * l * rt * 1e3:.2f} ms/前向')
+
+
+def mode_layer(args):
+    print(f'{"E":>4}  {"masked (改动前)":>18}  {"sorted (当前)":>18}  {"加速":>7}  {"峰值显存":>16}')
+    for e in args.experts:
+        torch.manual_seed(0)
+        mod = MOEFeedForward(MiniMindConfig(
+            hidden_size=args.hidden_size, use_moe=True, num_experts=e,
+            num_experts_per_tok=args.top_k, moe_intermediate_size=args.hidden_size * 2)).cuda().train()
+        x = torch.randn(args.batch_size, args.seq_len, args.hidden_size, device='cuda')
+
+        def run(fn):
+            def step():
+                with torch.autocast('cuda', dtype=torch.bfloat16):
+                    fn(x).sum().backward()
+                mod.zero_grad(set_to_none=True)
+            torch.cuda.reset_peak_memory_stats()
+            r = timed(step, args.warmup, args.reps)
+            r['peak'] = torch.cuda.max_memory_allocated() / 1e6
+            return r
+
+        try:
+            a = run(lambda t: reference_masked_forward(mod, t))
+            b = run(mod)
+        except torch.cuda.OutOfMemoryError:
+            print(f'{e:>4}  OOM'); continue
+        print(f'{e:>4}  {a["median"] * 1e3:8.2f} ms [{a["min"] * 1e3:6.2f},{a["max"] * 1e3:6.2f}]  '
+              f'{b["median"] * 1e3:8.2f} ms [{b["min"] * 1e3:6.2f},{b["max"] * 1e3:6.2f}]  '
+              f'{a["median"] / b["median"]:6.2f}x  {a["peak"]:7.0f} -> {b["peak"]:5.0f} MB')
+        del mod; torch.cuda.empty_cache()
+
+
+def main():
+    p = argparse.ArgumentParser()
+    p.add_argument('--mode', default='sync-probe', choices=['sync-probe', 'layer'])
+    p.add_argument('--hidden_size', default=768, type=int)
+    p.add_argument('--num_hidden_layers', default=8, type=int)
+    p.add_argument('--batch_size', default=8, type=int)
+    p.add_argument('--seq_len', default=512, type=int)
+    p.add_argument('--top_k', default=1, type=int)
+    p.add_argument('--experts', default='2,4,8,16', type=lambda s: [int(x) for x in s.split(',')])
+    p.add_argument('--warmup', default=10, type=int)
+    p.add_argument('--reps', default=20, type=int)
+    args = p.parse_args()
+
+    if not torch.cuda.is_available():
+        print('本脚本需要 CUDA。/ This benchmark requires CUDA.'); sys.exit(1)
+    print(f'torch {torch.__version__}  {torch.cuda.get_device_name(0)}  '
+          f'sm_{"".join(map(str, torch.cuda.get_device_capability(0)))}\n')
+    {'sync-probe': mode_sync_probe, 'layer': mode_layer}[args.mode](args)
+
+
+if __name__ == '__main__':
+    main()

scripts/test_moe_sorted.py

@@ -0,0 +1,186 @@
+"""MoE 排序分发的正确性测试 / correctness tests for the sorted MoE dispatch.
+
+本仓库没有 pytest / CI，所以这里不引入任何测试框架，直接 `python scripts/test_moe_sorted.py`
+运行，失败返回非零。
+
+reference_masked_forward 是改动之前逐专家做 mask 的写法，保留在这里作为参考实现：
+它每次运行都会被执行，并断言与当前实现逐位相同，因此既能当作对照读，也不会随代码腐烂。
+"""
+import os
+import sys
+import traceback
+import warnings
+
+import torch
+import torch.nn.functional as F
+
+sys.path.append(os.path.join(os.path.dirname(os.path.abspath(__file__)), '..'))
+from model.model_minimind import MiniMindConfig, MOEFeedForward
+
+
+class Skip(Exception): pass
+
+
+# ============================== 参考实现 / reference ==============================
+
+def reference_masked_forward(mod, x):
+    """改动前的逐专家掩码写法，逐位对照用。
+
+    每个专家都要 3 次 device→host 同步：mask.any() 的取值、nonzero() 的输出长度、
+    布尔索引的结果长度，全都得先回到主机才知道。
+    """
+    batch_size, seq_len, hidden_dim = x.shape
+    x_flat = x.view(-1, hidden_dim)
+    scores = F.softmax(mod.gate(x_flat), dim=-1)
+    topk_weight, topk_idx = torch.topk(scores, k=mod.config.num_experts_per_tok, dim=-1, sorted=False)
+    if mod.config.norm_topk_prob: topk_weight = topk_weight / (topk_weight.sum(dim=-1, keepdim=True) + 1e-20)
+    y = torch.zeros_like(x_flat)
+    for i, expert in enumerate(mod.experts):
+        mask = (topk_idx == i)
+        if mask.any():
+            token_idx = mask.any(dim=-1).nonzero().flatten()
+            weight = topk_weight[mask].view(-1, 1)
+            y.index_add_(0, token_idx, (expert(x_flat[token_idx]) * weight).to(y.dtype))
+        elif mod.training:
+            y[0, 0] += 0 * sum(p.sum() for p in expert.parameters())
+    return y.view(batch_size, seq_len, hidden_dim)
+
+
+# ============================== 工具 / helpers ==============================
+
+def build(seed=0, **kw):
+    base = dict(hidden_size=64, use_moe=True, num_experts=4, num_experts_per_tok=1, moe_intermediate_size=128)
+    base.update(kw)
+    torch.manual_seed(seed)
+    return MOEFeedForward(MiniMindConfig(**base))
+
+
+def force_routing(mod, assign, hidden_size=64):
+    """构造 x 使得 token t 被路由到 assign[t]：gate 取单位阵，x 在目标维上给大值。"""
+    with torch.no_grad():
+        mod.gate.weight.zero_()
+        for i in range(mod.config.num_experts): mod.gate.weight[i, i] = 1.0
+    x = torch.zeros(1, len(assign), hidden_size)
+    for t, a in enumerate(assign): x[0, t, a] = 10.0
+    return x
+
+
+# ============================== 用例 / tests ==============================
+
+def test_bit_identical_to_masked_reference():
+    """与掩码参考实现对照：前向和参数梯度在所有 k 下逐位相同；输入梯度在 k<=2 下逐位相同。
+
+    k>=3 时输入梯度只能做到 1 ulp 以内，这不是 bug：每个 token 要把 k 份专家贡献加起来，
+    浮点加法不满足结合律，(a+b)+c 与 a+(b+c) 可以差 1 ulp；两种写法的求和顺序不同。
+    k<=2 时任何顺序都精确相等，所以那里仍然要求 torch.equal。实测最大相对偏差 ~1.2e-7
+    （fp32 eps ≈ 1.2e-7），这里的阈值 1e-6 只是留一点余量。
+    """
+    for e, k in ((1, 1), (2, 1), (4, 1), (4, 2), (4, 3), (4, 4), (8, 2), (8, 3), (8, 8)):
+        mod = build(seed=3, num_experts=e, num_experts_per_tok=k)
+        torch.manual_seed(21)
+        x = torch.randn(3, 11, 64)
+        for train in (True, False):
+            mod.train(train)
+            xa, xb = x.clone().requires_grad_(), x.clone().requires_grad_()
+            ya = reference_masked_forward(mod, xa)
+            ya.sum().backward()
+            ga = {n: (p.grad.clone() if p.grad is not None else None) for n, p in mod.named_parameters()}
+            mod.zero_grad(set_to_none=True)
+            yb = mod(xb)
+            yb.sum().backward()
+            assert torch.equal(ya, yb), f'E={e} k={k} train={train}: 前向不是逐位相同'
+            if k <= 2:
+                assert torch.equal(xa.grad, xb.grad), f'E={e} k={k} train={train}: 输入梯度不是逐位相同'
+            else:  # 求和顺序不同，只能要求 1 ulp 量级
+                rel = (xa.grad - xb.grad).abs().max() / xa.grad.abs().max()
+                assert rel < 1e-6, f'E={e} k={k} train={train}: 输入梯度相对偏差 {rel:.2e} 超出浮点求和顺序可解释的范围'
+            for n, p in mod.named_parameters():
+                assert (ga[n] is None) == (p.grad is None), f'{n}: 梯度有无不一致'
+                if p.grad is not None:
+                    assert torch.equal(ga[n], p.grad), f'E={e} k={k} train={train} {n}: 参数梯度不是逐位相同'
+            mod.zero_grad(set_to_none=True)
+
+
+def test_edge_cases():
+    """零token专家 / 全部token到同一专家 / batch 小于专家数。"""
+    for name, assign in {'零token专家': [0, 1, 0, 1, 0, 1], '全部到专家0': [0] * 6, 'batch<专家数': [1, 2]}.items():
+        mod = build(seed=5); mod.train()
+        x = force_routing(mod, assign)
+        assert torch.equal(reference_masked_forward(mod, x), mod(x)), f'{name}: 与掩码参考实现不一致'
+
+
+def test_zero_token_expert_still_gets_grad():
+    """DDP 以 find_unused_parameters=False 构造，空专家也必须出现在自动求导图里。"""
+    mod = build(seed=5); mod.train()
+    mod(force_routing(mod, [0] * 6)).sum().backward()
+    missing = [n for n, p in mod.named_parameters() if p.grad is None]
+    assert not missing, f'{missing} 没有梯度 -> DDP(find_unused_parameters=False) 会报错'
+
+
+def test_aux_loss_and_state_dict_unchanged():
+    """aux_loss 与 state_dict 结构不受本次改动影响。"""
+    mod = build(seed=7); mod.train()
+    torch.manual_seed(11)
+    x = torch.randn(2, 16, 64)
+    mod(x); got = mod.aux_loss.clone()
+    scores = F.softmax(mod.gate(x.view(-1, 64)), dim=-1)
+    _, topk_idx = torch.topk(scores, k=1, dim=-1, sorted=False)
+    load = F.one_hot(topk_idx, 4).float().mean(0)
+    want = (load * scores.mean(0)).sum() * 4 * mod.config.router_aux_loss_coef
+    assert torch.equal(got, want), 'aux_loss 变了'
+    expect = {'gate.weight'} | {f'experts.{i}.{p}_proj.weight' for i in range(4) for p in ('gate', 'up', 'down')}
+    assert set(mod.state_dict()) == expect, 'state_dict 键名变了，旧 checkpoint 将无法 strict 加载'
+
+
+def test_host_sync_count():
+    """本次改动的核心：整层同步次数恒为 1，不随专家数增长。"""
+    if not torch.cuda.is_available(): raise Skip('需要 CUDA')
+
+    def syncs(fn, x):
+        fn(x); torch.cuda.synchronize()  # 预热，避免把首次分配算进去
+        torch.cuda.set_sync_debug_mode('warn')
+        with warnings.catch_warnings(record=True) as w:
+            warnings.simplefilter('always'); fn(x)
+        torch.cuda.set_sync_debug_mode('default')
+        return sum('sync' in str(r.message).lower() for r in w)
+
+    for e in (4, 8, 16):
+        mod = build(seed=3, num_experts=e).cuda().train()
+        x = torch.randn(2, 64, 64, device='cuda')
+        ref, now = syncs(lambda t: reference_masked_forward(mod, t), x), syncs(mod, x)
+        assert ref >= 3 * e, f'E={e}: 掩码写法应约 3E={3 * e} 次同步，实测 {ref}'
+        assert now == 1, f'E={e}: 排序写法应恒为 1 次同步，实测 {now}'
+
+
+def test_cuda_bit_identical():
+    """CUDA 上 fp32 与 bf16 autocast 同样逐位相同——排序改变了 kernel 的分块方式。"""
+    if not torch.cuda.is_available(): raise Skip('需要 CUDA')
+    for e, k in ((4, 1), (8, 2)):
+        mod = build(seed=3, num_experts=e, num_experts_per_tok=k).cuda().train()
+        torch.manual_seed(21)
+        x = torch.randn(4, 128, 64, device='cuda')
+        assert torch.equal(reference_masked_forward(mod, x), mod(x)), f'E={e} k={k}: CUDA fp32 不一致'
+        with torch.autocast('cuda', dtype=torch.bfloat16):
+            assert torch.equal(reference_masked_forward(mod, x), mod(x)), f'E={e} k={k}: CUDA bf16 不一致'
+
+
+TESTS = [test_bit_identical_to_masked_reference, test_edge_cases, test_zero_token_expert_still_gets_grad,
+         test_aux_loss_and_state_dict_unchanged, test_host_sync_count, test_cuda_bit_identical]
+
+
+def main():
+    print(f'torch {torch.__version__}  cuda={torch.cuda.is_available()}\n')
+    fails = skips = 0
+    for t in TESTS:
+        try:
+            t(); print(f'  PASS  {t.__name__}')
+        except Skip as s:
+            skips += 1; print(f'  SKIP  {t.__name__} ({s})')
+        except Exception:
+            fails += 1; print(f'  FAIL  {t.__name__}'); traceback.print_exc()
+    print(f'\n{len(TESTS) - fails - skips} passed, {fails} failed, {skips} skipped')
+    sys.exit(1 if fails else 0)
+
+
+if __name__ == '__main__':
+    main()