Commit Graph

20 Commits (310057773e942ad0318f4b4c738b44c88f6f2e8e)

Author SHA1 Message Date
Justin Guese 6b2f4c939a [update] MoE 分发:每层主机同步从 3E 次降到 1 次
专家循环里 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 加载。
2026-09-03 15:05:25 +07:00
jingyaogong dddedc6881 [fix] repetition_penalty 2026-05-07 19:08:52 +08:00
jingyaogong 1718e9a44d [fix] transformers-5.x 2026-04-19 23:48:54 +08:00
jingyaogong 5704766352 [update] tie embedding 2026-04-19 21:57:28 +08:00
jingyaogong 2ab6455d9d [update] open causal 2026-04-02 15:28:58 +08:00
jingyaogong 101d7df2da [update] minimind-3 2026-03-25 23:57:45 +08:00
jingyaogong c090b69c4d [update] align loss 2026-01-15 00:56:32 +08:00
jingyaogong 1279a61681 [update] prompt prefill 2026-01-13 17:46:54 +08:00
jingyaogong 9d898576ac [update] aux loss 2026-01-01 22:41:46 +08:00
jingyaogong c65335b56f [fix] experts unused 2025-12-31 21:47:04 +08:00
jingyaogong 5129f0e2a2 [fix] dtype & lr 2025-12-09 13:01:38 +08:00
jingyaogong ecd1ae1563 [fix] reduce aux_loss_alpha 2025-12-05 23:08:29 +08:00
jingyaogong 151fdf7e76 [feat] update yarn 2025-12-01 16:15:05 +08:00
jingyaogong 6b86ea399a [feat] release memory 2025-11-27 19:39:49 +08:00
jingyaogong f5374dc87f [fix] model attn_mask 2025-11-19 22:26:53 +08:00
jingyaogong a044578d73 [fix] update model 2025-11-18 13:07:20 +08:00
yuyu5333 7d02ce673c fix: attn_forwad when is_causal=True assert attn_mask is None 2025-11-18 03:17:17 +00:00
jingyaogong 4e35fb9da8 [fix] update model 2025-10-17 00:09:32 +08:00
jingyaogong 29454c31af fix bugs 2025-04-27 09:56:49 +08:00
jingyaogong a62faf34bd 250426 2025-04-26 10:05:47 +08:00