"""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()