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 |