Add a Dist-MoE expert-parallel strategy #49813
3outeille
wants to merge
55 commits into
ep-deepep-strategy
from
ep-dist-moe-strategy
Expert-parallel dispatch runs the experts masked when the EP group sh…
d6f4bf1f
Test the masked path on two real ranks, at the end of the layer tests
0735ccac
Import the distributed test helper relatively, as the other distribut…
cbcff744
Masked EP as its own parallel style, resolved from `ep_dispatch_exper…
89c57aff
Leave experts that route themselves to the router-masked plan
2e30d871
Pick the expert-parallel strategy from `DistributedConfig.ep_strategy`
f2e79153
Expert-parallel strategies as a registry of stateful classes, picked …
bc17ae7f
Write the `ep_all_reduce_max_bytes` default as a literal so the docst…
8c914b38
Write the `ep_all_reduce_max_bytes` default as `32 * 2**20`, as the f…
0edc1935
`ep_all_reduce_max_bytes` defaults to None: 32 MiB, and no limit at `…
d632883e
Keep the expert-parallel docs to what was measured, and say when only…
fd357b9e
"auto" runs the all-reduce at inference and the all-to-all in training
a3248db5
Keep the docs diff to the strategies, and call token dispatch a strat…
5626f97e
`passthrough` expert-parallel strategy, selected for DeepGEMM Mega MoE
4c06b2da
Move the remaining router-masked models onto `ep_dispatch_experts`
eabc2479
Drop the expert-parallel comment from the Qwen3-MoE and Mellum plans
8104384d
Merge remote-tracking branch 'origin/main' into ep-masked-when-shared…
4a5c4c99
Move the all-reduce autograd functions and `_use_local_dtensor_params…
cefdf474
Merge remote-tracking branch 'origin/ep-masked-when-shared-batch' int…
11f62455
Merge remote-tracking branch 'origin/ep-passthrough-strategy' into ep…
ec85eacd
Deprecate the Mega MoE router and experts parallel styles
9fd554ab
Merge remote-tracking branch 'origin/ep-passthrough-strategy' into ep…
224e3539
Remove the Mega MoE router and experts parallel styles
c6d6367e
Merge remote-tracking branch 'origin/ep-passthrough-strategy' into ep…
f9506c53
`deepgemm_megamoe` strategy owns the Mega MoE symmetric buffers, shar…
ae7e2242
Merge remote-tracking branch 'origin/ep-passthrough-strategy' into ep…
559d6627
Allocate the Mega MoE buffer into the cache, then read it back
ce7a0b6b
Merge remote-tracking branch 'origin/ep-passthrough-strategy' into ep…
3eab5814
Keep the Mega MoE forward in its original order
2af6427d
Merge remote-tracking branch 'origin/ep-passthrough-strategy' into ep…
b3a5ed45
Hand each rank's Mega MoE kernel a slice of the shared batch
0f2c56a3
Merge remote-tracking branch 'origin/ep-passthrough-strategy' into ep…
c3d48edc
Note that Mega MoE needs data-parallel inference to run as designed
2feb642d
Merge remote-tracking branch 'origin/ep-passthrough-strategy' into ep…
66c1dd11
MXFP4 GPT-OSS experts take routed tokens, so expert parallelism runs …
1f3a6b1d
Remove the `ep_router` style
53952499
Add the `deepep` expert-parallel strategy
6dba4d93
Prefer DeepEP for token dispatch in `auto` where it runs
7c7693a8
`auto` dispatches with DeepEP at every batch size
f1599342
DeepEP may take more SMs, since nothing overlaps with it
e70b08e0
Expert-parallel layers in `expert_parallel.py`, `ep_implementation`, …
6bfa58b3
Merge branch 'ep-masked-when-shared-batch' into ep-passthrough-strategy
60bf55a7
Merge branch 'ep-passthrough-strategy' into ep-move-router-plans
b4db7983
Merge branch 'ep-move-router-plans' into ep-mxfp4-gpt-oss
7e35df69
Merge branch 'ep-mxfp4-gpt-oss' into ep-remove-router
06034cc9
Merge branch 'ep-remove-router' into ep-deepep-strategy
55d15c0a
Merge remote-tracking branch 'origin/main' into ep-masked-when-shared…
ccfc1335
Merge branch 'ep-masked-when-shared-batch' into ep-passthrough-strategy
a7f22dbf
Merge branch 'ep-passthrough-strategy' into ep-move-router-plans
2b82f8bb
Merge branch 'ep-move-router-plans' into ep-mxfp4-gpt-oss
e67fc052
Merge branch 'ep-mxfp4-gpt-oss' into ep-remove-router
7d13b079
Merge branch 'ep-remove-router' into ep-deepep-strategy
8f57f985
Add a Dist-MoE expert-parallel strategy
58bd94de
Document that dist_moe runs with ep_size == tp_size and ep_size > tp_…
6128a9f6
Add a dist_moe experts implementation
f16cf86d
IlyasMoutawwakil
force-pushed the
ep-deepep-strategy
branch
from
eb8f0495
to
1b8984cb
43 minutes ago
Reviewers
No reviews
Assignees
No one assigned
Labels
None yet
Milestone
No milestone
Login to write a write a comment.
Login via GitHub