6b2f4c939a
专家循环里 mask.any() / nonzero() / 布尔索引每个专家各触发一次 device→host 同步，实测每层正好 3E 次（默认 4 专家 × 8 层 = 每次前向 96 次）。同步会清空 CUDA 流，而这些开销全部花在路由的簿记上。 改为先按专家把 token 排成连续段，用 searchsorted 取边界、一次 tolist() 读回 主机，每个专家取一段连续切片。专家循环保留，整层只剩 1 次同步，与专家数无关。 纯 PyTorch，无新依赖，对硬件和 dtype 没有要求。 stable=True 保证段内 token 顺序不变，因此前向与全部参数梯度在所有 k 下与原 写法逐位相同；输入梯度在 k<=2 时逐位相同，k>=3 时因浮点加法不满足结合律 差约 1 ulp（相对 1e-7）。原掩码写法作为参考实现保留在 scripts/test_moe_sorted.py， 每次跑测试都会执行并断言一致。 state_dict、aux_loss、router 与模块结构均未改动，旧 checkpoint 仍可 strict 加载。 |
||
|---|---|---|
| .. | ||
| __init__.py | ||
| model_lora.py | ||
| model_minimind.py | ||
| tokenizer.json | ||
| tokenizer_config.json | ||