Repository navigation
[distributed] Default MoE ep_plans to token dispatch - #49160
Conversation
3outeille
changed the title
ep_plans to token dispatchdistributed] Default MoE ep_plans to token dispatch
|
run-slow: afmoe, axk1, axk2, cohere2_moe, deepseek_ocr2, deepseek_v2, deepseek_v3, deepseek_v32, deepseek_v4, dots1, ernie4_5_moe, ernie4_5_vl_moe, exaone_moe, flex_olmo, glm4_moe, glm4_moe_lite |
|
The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update. |
87c3cb3 to
35d66ce
Compare
55e0430 to
93bc375
Compare
51b74d5 to
040541a
Compare
66105e6 to
1cf28ef
Compare
|
run-slow: afmoe, axk1, axk2, cohere2_moe, deepseek_ocr2, deepseek_v2, deepseek_v3, deepseek_v32, deepseek_v4, dots1, ernie4_5_moe, ernie4_5_vl_moe, exaone_moe, flex_olmo, glm4_moe, glm4_moe_lite |
|
Workflow Run ⚙️💔 This comment contains |
| "layers": (["hidden_states", "attention_mask"], ["hidden_states"]), | ||
| "norm": (["hidden_states"], ["hidden_states"]), | ||
| } | ||
| base_model_ep_plan = { |
There was a problem hiding this comment.
this is a dense model, should not have ep_plan
eacebd8 to
8a8b28b
Compare
dc8ec2c to
e3da413
Compare
ArthurZucker
left a comment
ArthurZucker
left a comment
There was a problem hiding this comment.
SGTM, keeping experts in tp plan would IMO mean we are doing TP over EP. Which we are not. I think it woul dbe nice to have that written somewhere?
IDK if we plan to support it, maybenot? but that would clearup potential missunderstanding + the checks in PR stack 2 I think
e3da413 to
f5d7569
Compare
We are not doing ETP (expert tensor parallel) by keeping the experts in As for ETP, I dont think it's a priority to support it for now. |
f7a128c to
5c975a6
Compare
87bc190 to
0161264
Compare
0161264 to
0122045
Compare
`initialize_distributed_mesh` now builds two named views of the same ranks: `(pp, fsdp, tp)` for dense layers and `(pp, efsdp, ep)` for experts, both keeping size-one axes so callers select dimensions by name. `MeshManager` routes `ep`/`efsdp` lookups to the expert view and everything else to the dense view. `DistributedConfig` gains `ep_size` (defaults to `tp_size` when `enable_expert_parallel=True`) and `efsdp_size`, with size validation. Model execution is unchanged: expert sharding and FSDP still use the `tp` and `fsdp` axes, and loading rejects `ep_size != tp_size` until the all-to-all dispatcher lands.
Co-authored-by: Arthur <48595927+ArthurZucker@users.noreply.github.com>
…formers into ep-mesh-infrastructure
|
[For maintainers] Suggested jobs to run (before merge) run-slow: afmoe, axk1, axk2, cohere2_moe, deepseek_ocr2, deepseek_v2, deepseek_v3, deepseek_v32, deepseek_v4, dots1, ernie4_5_moe, ernie4_5_vl_moe, exaone_moe, flex_olmo, glm4_moe, glm4_moe_lite |
|
run-slow: afmoe, axk1, axk2, cohere2_moe, deepseek_ocr2, deepseek_v2, deepseek_v3, deepseek_v32, deepseek_v4, dots1, ernie4_5_moe, ernie4_5_vl_moe, exaone_moe, flex_olmo, glm4_moe, glm4_moe_lite |
CI recapDashboard: View test results in Grafana |
AMD CIThis comment contains models: ["models/afmoe", "models/axk1", "models/axk2", "models/cohere2_moe", "models/deepseek_ocr2", "models/deepseek_v2", "models/deepseek_v3", "models/deepseek_v32", "models/deepseek_v4", "models/dots1", "models/ernie4_5_moe", "models/ernie4_5_vl_moe", "models/exaone_moe", "models/flex_olmo", "models/glm4_moe", "models/glm4_moe_lite"] |
Nvidia CIThis comment contains models: ["models/afmoe", "models/axk1", "models/axk2", "models/cohere2_moe", "models/deepseek_ocr2", "models/deepseek_v2", "models/deepseek_v3", "models/deepseek_v32", "models/deepseek_v4", "models/dots1", "models/ernie4_5_moe", "models/ernie4_5_vl_moe", "models/exaone_moe", "models/flex_olmo", "models/glm4_moe", "models/glm4_moe_lite"] |
3 participants
Converted (44 models, 45 config classes):
afmoe, axk1, axk2, cohere2_moe, deepseek_ocr2, deepseek_v2, deepseek_v3, deepseek_v32, deepseek_v4, dots1, ernie4_5_moe, ernie4_5_vl_moe, exaone_moe, flex_olmo, glm4_moe, glm4_moe_lite, glm4v_moe, glm5_next, glm_moe_dsa, gpt_oss, hunyuan_v1_moe, hy_v3, kimi_linear, laguna, lfm2_moe, mellum, minimax, minimax_m2, minimax_m3_vl, mistral4, mixtral, olmoe, openai_privacy_filter, phimoe, qwen2_moe, qwen3_5_moe, qwen3_moe, qwen3_next, qwen3_omni_moe, qwen3_vl_moe, qwen4_exp, solar_open, step3p7, zaya.Still on masking (6 models):
gemma4, diffusion_gemma, granitemoe_swa, hy_v4, inkling, mimo_v2_flash. To be fix as they were failing already before the stack of PR