Add ep - #39501
Merged
Merged
Conversation
Co-authored-by: Nouamane Tazi <NouamaneTazi@users.noreply.github.com> Co-authored-by: drbh <drbh@users.noreply.github.com>
Merged
Collaborator
Author
|
Does not work for mixtral, tricky because of the sequentiality of the weights. I'll add something that allows for this, but also a plan that allows for one the fly merging the modulelist to the format for megablocks |
25 tasks
Collaborator
|
Will getting deepseek working too be pretty straightforward? |
Collaborator
Author
|
Yes and no! for both I used nn.Modulelist() (deepseek I had no choice and Mixtral I was junior) and so it's a bit more annoying, but yes because next pr will make sur we have a bit of a better interface! |
|
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. |
ArthurZucker
commented
Jul 25, 2025
Contributor
|
[For maintainers] Suggested jobs to run (before merge) run-slow: llama4 |
5 tasks
BenjaminBossan
added a commit
to BenjaminBossan/peft
that referenced
this pull request
Jul 28, 2025
For testing target_parameters, we use a tiny Llama4 model. This model was refactored in huggingface/transformers#39501, resulting in one parameter being accessed an additional time: https://github.com/huggingface/transformers/pull/39501/files#diff-e668ec07f78afdb2cb805d939e47453757f0b9437436cb860fcb7cb2431c9cf5R69 Therefore, a unit test that relied on how often this parameter was accessed started failing. This PR updates the count to the correct number. Additionally debug print statements that were accidentally left over are now removed.
BenjaminBossan
added a commit
to huggingface/peft
that referenced
this pull request
Jul 30, 2025
For testing target_parameters, we use a tiny Llama4 model. This model was refactored in huggingface/transformers#39501, resulting in one parameter being accessed an additional time: https://github.com/huggingface/transformers/pull/39501/files#diff-e668ec07f78afdb2cb805d939e47453757f0b9437436cb860fcb7cb2431c9cf5R69 Therefore, a unit test that relied on how often this parameter was accessed started failing. This PR updates the count to the correct number. Additionally debug print statements that were accidentally left over are now removed.
5 tasks
nvpohanh
added a commit
to nvpohanh/transformers
that referenced
this pull request
Sep 4, 2025
Llama4 accuracy is broken by a bug in huggingface#39501 . It forgot to transpose the router_scores before applying it to routed_in, causing Llama4 to generate garbage output. This PR fixes that issue by adding back the transpose() and adding some comments explaining why the transpose() is needed. Signed-off-by: Po-Han Huang <pohanh@nvidia.com>
Cyrilvallez
added a commit
that referenced
this pull request
Sep 4, 2025
* Fix broken Llama4 accuracy in MoE part Llama4 accuracy is broken by a bug in #39501 . It forgot to transpose the router_scores before applying it to routed_in, causing Llama4 to generate garbage output. This PR fixes that issue by adding back the transpose() and adding some comments explaining why the transpose() is needed. Signed-off-by: Po-Han Huang <pohanh@nvidia.com> * remove comment --------- Signed-off-by: Po-Han Huang <pohanh@nvidia.com> Co-authored-by: Cyril Vallez <cyril.vallez@gmail.com>
Cyrilvallez
added a commit
that referenced
this pull request
Sep 4, 2025
* Fix broken Llama4 accuracy in MoE part Llama4 accuracy is broken by a bug in #39501 . It forgot to transpose the router_scores before applying it to routed_in, causing Llama4 to generate garbage output. This PR fixes that issue by adding back the transpose() and adding some comments explaining why the transpose() is needed. Signed-off-by: Po-Han Huang <pohanh@nvidia.com> * remove comment --------- Signed-off-by: Po-Han Huang <pohanh@nvidia.com> Co-authored-by: Cyril Vallez <cyril.vallez@gmail.com>
zaristei
pushed a commit
to zaristei/transformers
that referenced
this pull request
Sep 9, 2025
* EP + updates Co-authored-by: Nouamane Tazi <NouamaneTazi@users.noreply.github.com> Co-authored-by: drbh <drbh@users.noreply.github.com> * remove unrelated change * not working yet but let's see where it goes! * update the api a bit * udpate * where I am at for now * fix ep * refactor the API * yups * fix * fixup * clean modeling * just support llama4 for now! * properly avoid * fix * nits * Update src/transformers/models/llama4/modeling_llama4.py * Update src/transformers/integrations/tensor_parallel.py * style * ,,,, * update --------- Co-authored-by: Nouamane Tazi <NouamaneTazi@users.noreply.github.com> Co-authored-by: drbh <drbh@users.noreply.github.com>
zaristei
pushed a commit
to zaristei/transformers
that referenced
this pull request
Sep 9, 2025
* EP + updates Co-authored-by: Nouamane Tazi <NouamaneTazi@users.noreply.github.com> Co-authored-by: drbh <drbh@users.noreply.github.com> * remove unrelated change * not working yet but let's see where it goes! * update the api a bit * udpate * where I am at for now * fix ep * refactor the API * yups * fix * fixup * clean modeling * just support llama4 for now! * properly avoid * fix * nits * Update src/transformers/models/llama4/modeling_llama4.py * Update src/transformers/integrations/tensor_parallel.py * style * ,,,, * update --------- Co-authored-by: Nouamane Tazi <NouamaneTazi@users.noreply.github.com> Co-authored-by: drbh <drbh@users.noreply.github.com>
zaristei
pushed a commit
to zaristei/transformers
that referenced
this pull request
Sep 9, 2025
* EP + updates Co-authored-by: Nouamane Tazi <NouamaneTazi@users.noreply.github.com> Co-authored-by: drbh <drbh@users.noreply.github.com> * remove unrelated change * not working yet but let's see where it goes! * update the api a bit * udpate * where I am at for now * fix ep * refactor the API * yups * fix * fixup * clean modeling * just support llama4 for now! * properly avoid * fix * nits * Update src/transformers/models/llama4/modeling_llama4.py * Update src/transformers/integrations/tensor_parallel.py * style * ,,,, * update --------- Co-authored-by: Nouamane Tazi <NouamaneTazi@users.noreply.github.com> Co-authored-by: drbh <drbh@users.noreply.github.com>
zaristei
pushed a commit
to zaristei/transformers
that referenced
this pull request
Sep 9, 2025
* EP + updates Co-authored-by: Nouamane Tazi <NouamaneTazi@users.noreply.github.com> Co-authored-by: drbh <drbh@users.noreply.github.com> * remove unrelated change * not working yet but let's see where it goes! * update the api a bit * udpate * where I am at for now * fix ep * refactor the API * yups * fix * fixup * clean modeling * just support llama4 for now! * properly avoid * fix * nits * Update src/transformers/models/llama4/modeling_llama4.py * Update src/transformers/integrations/tensor_parallel.py * style * ,,,, * update --------- Co-authored-by: Nouamane Tazi <NouamaneTazi@users.noreply.github.com> Co-authored-by: drbh <drbh@users.noreply.github.com>
zaristei
pushed a commit
to zaristei/transformers
that referenced
this pull request
Sep 9, 2025
* EP + updates Co-authored-by: Nouamane Tazi <NouamaneTazi@users.noreply.github.com> Co-authored-by: drbh <drbh@users.noreply.github.com> * remove unrelated change * not working yet but let's see where it goes! * update the api a bit * udpate * where I am at for now * fix ep * refactor the API * yups * fix * fixup * clean modeling * just support llama4 for now! * properly avoid * fix * nits * Update src/transformers/models/llama4/modeling_llama4.py * Update src/transformers/integrations/tensor_parallel.py * style * ,,,, * update --------- Co-authored-by: Nouamane Tazi <NouamaneTazi@users.noreply.github.com> Co-authored-by: drbh <drbh@users.noreply.github.com>
zaristei
pushed a commit
to zaristei/transformers
that referenced
this pull request
Sep 9, 2025
* EP + updates Co-authored-by: Nouamane Tazi <NouamaneTazi@users.noreply.github.com> Co-authored-by: drbh <drbh@users.noreply.github.com> * remove unrelated change * not working yet but let's see where it goes! * update the api a bit * udpate * where I am at for now * fix ep * refactor the API * yups * fix * fixup * clean modeling * just support llama4 for now! * properly avoid * fix * nits * Update src/transformers/models/llama4/modeling_llama4.py * Update src/transformers/integrations/tensor_parallel.py * style * ,,,, * update --------- Co-authored-by: Nouamane Tazi <NouamaneTazi@users.noreply.github.com> Co-authored-by: drbh <drbh@users.noreply.github.com>
zaristei
pushed a commit
to zaristei/transformers
that referenced
this pull request
Sep 9, 2025
* EP + updates Co-authored-by: Nouamane Tazi <NouamaneTazi@users.noreply.github.com> Co-authored-by: drbh <drbh@users.noreply.github.com> * remove unrelated change * not working yet but let's see where it goes! * update the api a bit * udpate * where I am at for now * fix ep * refactor the API * yups * fix * fixup * clean modeling * just support llama4 for now! * properly avoid * fix * nits * Update src/transformers/models/llama4/modeling_llama4.py * Update src/transformers/integrations/tensor_parallel.py * style * ,,,, * update --------- Co-authored-by: Nouamane Tazi <NouamaneTazi@users.noreply.github.com> Co-authored-by: drbh <drbh@users.noreply.github.com>
5 tasks
4 tasks
Guy-Bilitski
pushed a commit
to Guy-Bilitski/UIOrthoLoRA
that referenced
this pull request
Feb 5, 2026
For testing target_parameters, we use a tiny Llama4 model. This model was refactored in huggingface/transformers#39501, resulting in one parameter being accessed an additional time: https://github.com/huggingface/transformers/pull/39501/files#diff-e668ec07f78afdb2cb805d939e47453757f0b9437436cb860fcb7cb2431c9cf5R69 Therefore, a unit test that relied on how often this parameter was accessed started failing. This PR updates the count to the correct number. Additionally debug print statements that were accidentally left over are now removed.
3outeille
added a commit
that referenced
this pull request
May 29, 2026
Decompose MoE tensor/expert parallelism per review feedback: weight sharding is declared per-parameter, while the experts module entry stays forward-comm only. - MoEParamShard: parameter-only style wrapping named expert weights as DTensor placeholders (no forward hook). grouped_gemm shards the expert dim and updates module.num_experts to the per-rank local count. - Register grouped_gemm (Shard(0)), moe_gate_up_colwise (_StridedShard(-2)), moe_gate_up_colwise_alt (_StridedShard(-1)), moe_down_rowwise (Shard(-1)). - EpRouterParallel (ep_router): forward-only slicing of router outputs to local experts, ported from the original RouterParallel (#39501). - moe_experts_allreduce is now forward-comm only: strip the baked shard_plan and drop the now-dead shard_plan ctor arg / _moe_shard_plan / shard_parameters override from MoEExpertsParallel; skip _AllReduceBackward on routing weights under EP. - verify_tp_plan: treat moe_experts_allreduce / ep_router as forward-only.
3outeille
added a commit
that referenced
this pull request
Jun 17, 2026
* [distributed] Add param-level MoE TP/EP styles and ep_router Decompose MoE tensor/expert parallelism per review feedback: weight sharding is declared per-parameter, while the experts module entry stays forward-comm only. - MoEParamShard: parameter-only style wrapping named expert weights as DTensor placeholders (no forward hook). grouped_gemm shards the expert dim and updates module.num_experts to the per-rank local count. - Register grouped_gemm (Shard(0)), moe_gate_up_colwise (_StridedShard(-2)), moe_gate_up_colwise_alt (_StridedShard(-1)), moe_down_rowwise (Shard(-1)). - EpRouterParallel (ep_router): forward-only slicing of router outputs to local experts, ported from the original RouterParallel (#39501). - moe_experts_allreduce is now forward-comm only: strip the baked shard_plan and drop the now-dead shard_plan ctor arg / _moe_shard_plan / shard_parameters override from MoEExpertsParallel; skip _AllReduceBackward on routing weights under EP. - verify_tp_plan: treat moe_experts_allreduce / ep_router as forward-only. * [distributed] Add param-level apply pass to apply_tensor_parallel Run tensor parallelism in two passes: - Pass 1 (param-level): walk named_parameters() and, for styles in PARAM_ONLY_STYLES (grouped_gemm, moe_gate_up_colwise[_alt], moe_down_rowwise), shard the parameter directly via shard_parameters(). No forward hook. - Pass 2 (module-level): the existing named_modules() loop for forward hooks, now skipping PARAM_ONLY_STYLES. Param sharding runs first so module forward hooks (moe_experts_allreduce) see the already-sharded DTensor params. Also wire the EP-plan fallback so enable_expert_parallel uses model._ep_plan when no explicit plan is passed. * [distributed] Add MoE TP/EP plan tests and a two-sided sharding assertion - New tests/distributed/test_moe_tensor_parallel_plan.py: plan resolution, placement expectations for grouped_gemm / moe_gate_up_colwise[_alt] / moe_down_rowwise, gloo distributed integration (EP Shard(0), TP _StridedShard(-2)+Shard(-1), ep_router slicing), and a registry guard that moe_experts_allreduce carries no baked shard plan. - _verify_tp_sharding: add a two-sided check asserting that every parameter whose plan entry is a weight-sharding style actually comes back as a non-replicate DTensor. The prior check only validated params that happened to be sharded, so a style that gracefully degrades to replicated when unsharded (e.g. MoEExpertsParallel) could pass output-equality while silently running unparallelized. * [distributed] Migrate MoE configs to decomposed TP/SP expert plans For every TP/SP plan that sharded experts, declare per-parameter entries: "layers.*.mlp.experts.gate_up_proj": "moe_gate_up_colwise" "layers.*.mlp.experts.down_proj": "moe_down_rowwise" while keeping the forward-only "layers.*.mlp.experts": "moe_experts_allreduce". This matches the now-empty moe_experts_allreduce shard_plan; sharding is declared in config at parameter granularity. EP plans already used "grouped_gemm" and are unchanged. hy_v3 and laguna previously used "packed_colwise" / "rowwise_allreduce" on the 3D expert *parameters*; those styles are module-level and were silently no-ops on params (the bundled shard_plan did the work). They now use the param-level moe_gate_up_colwise / moe_down_rowwise like every other MoE model. Edited modular files where they own the plan literal; generated configs and inherited plans (e.g. from qwen3_moe) propagated via modular conversion. * [distributed] Document decomposed MoE TP/EP plans - expert_parallelism.md: describe the param-level decomposition (grouped_gemm, ep_router, moe_experts_allreduce) instead of the removed GroupedGemmParallel class, and note the TP equivalents (moe_gate_up_colwise / moe_down_rowwise). - weightconverter.md: note that fused expert weights are sharded at parameter granularity by the parallel plan. * [distributed] Rename MoE intra-expert TP styles to moe_tp_* Rename registry and plan entries so TP-on-expert sharding is distinct from EP (grouped_gemm) and dense packed_colwise: moe_gate_up_colwise -> moe_tp_gate_up_colwise, moe_down_rowwise -> moe_tp_down_rowwise. Drop unused moe_tp_gate_up_colwise_alt (GPT-OSS-style layouts stay EP-only). * handle sparse and dense sp plan for qwen3_moe * better tests coverage for sp & ep * linting * uniformize TP Api to avoid confusion with torch native ops * inline tp * rename * cleaning * inline * cleaning * cleaning * linting * fix ci ep_backward * linting * remove flag expert parallel * fix * add tp plan + ep_plan * revert doc * fix install_forward * linting * add moe identity back * no need aymore * update tp_plan for ernie4_5_vl_moe * sp + ep training / tp + ep inference (#46292) * [distributed] Add resolve_parallel_plan merge helper Compose SP/TP dense recipes with an optional EP overlay and strip intra-expert moe_tp_* when expert parallelism is enabled. Add unit tests for training (SP+EP), inference (TP+EP), and TP-only paths. * [distributed] Wire resolve_parallel_plan into apply_tensor_parallel Replace exclusive SP|EP|TP plan selection with merged plans when tp_plan is unset. Add distributed test for TP+EP merged expert sharding. * [distributed] Use merged plan in tp_plan property and load path Expose resolve_parallel_plan via PreTrainedModel.tp_plan and set active_tp_plan during from_pretrained so checkpoint sharding matches the applied layout. * [distributed] Drop intra-expert moe_tp_* from MoE SP plans Expert weight TP under sequence parallelism comes from the EP overlay (grouped_gemm) when enable_expert_parallel is set; keep moe_tp_* only in base_model_tp_plan for TP-only MoE. * [distributed] Document SP+EP and TP+EP flag combinations Update expert_parallelism guide and DistributedConfig docs for merged plans. Export resolve_parallel_plan and extend resolve-plan tests for trimmed SP sources. * refactor merging plans * add test sp_ep and tp_ep * extend verify_tp_plan to verify_tp_sp_ep_plan * add ep_plan to mixtral and olmoe * cleaning _accumulate_local_param_grad (#46394) * remove _accumulate_local_param_grad * comments * linting * fix * clean _accumulate_local_param_grad * linting * cleaning * cleaning * fix mellun test because of bug in parsing sp_ep plan with regex * aea * Add select_parallel_plan and explicit combo plan config fields Introduce base_model_tp_ep_plan / base_model_sp_ep_plan on PreTrainedConfig, select_parallel_plan() with legacy resolve_parallel_plan fallback, and wire apply_tensor_parallel to use the selector. Model post_init tracks _tp_ep_plan and _sp_ep_plan for composite models. * Add base_model_tp_ep_plan and base_model_sp_ep_plan for Mixtral and Qwen3-MoE Define complete inference TP+EP and training SP+EP plans on the pilot MoE configs. Qwen3-MoE expands per-layer entries in _update_parallel_plans. Add plan_utils and golden tests against legacy resolve_parallel_plan merge. * Add explicit tp_ep / sp_ep plans for remaining MoE models Populate combo plans via init_combo_plans() at config init time for MoE configs that still use split tp/sp/ep recipes. Dynamic configs call it after _update_sp_plan(); modular sources updated for generated configuration files. * Remove resolve_parallel_plan and use explicit combo plan selection Delete runtime plan merging; select_parallel_plan now requires a complete combo dict and raises when missing. apply_tensor_parallel uses DistributedConfig flags directly for SP/EP behavior. Drop model._ep_plan aggregation; load-time verification checks the active plan only. Refresh combo plans after MXFP4 quantizer patches. * Sync modular MoE configs and update expert parallelism docs Propagate init_combo_plans from modular sources to generated configuration files and document select_parallel_plan combo lookup in expert_parallelism.md. * Refactor select_parallel_plan flag lookup for readability Use explicit if/elif branches for the SP/EP flag matrix and derive config_attr from plan_attr instead of parallel lookup dicts. * Write explicit combo parallel plans in MoE configs and remove plan_utils Define base_model_tp_ep_plan and base_model_sp_ep_plan directly in each MoE configuration (or via config-time _update_parallel_plans for dynamic models). Delete plan_utils.py and all init_combo_plans / refresh_combo_plans usage. * Add lm_head entries to _tp_ep_plan and _sp_ep_plan on CausalLM classes Explicit combo plan selection no longer merges _sp_plan with _ep_plan, so head-level lm_head rules must live on _tp_ep_plan/_sp_ep_plan directly. Fixes SP+EP training loss shape mismatch under sequence parallelism. * cleaning * cleaning * cleaning * cleaning * linting * add verify tp and fsdp pla aeaea * revert doc * cleaning * check-repository-consistency * linting
iambogeumkim
pushed a commit
to iambogeumkim/peft
that referenced
this pull request
Jun 21, 2026
For testing target_parameters, we use a tiny Llama4 model. This model was refactored in huggingface/transformers#39501, resulting in one parameter being accessed an additional time: https://github.com/huggingface/transformers/pull/39501/files#diff-e668ec07f78afdb2cb805d939e47453757f0b9437436cb860fcb7cb2431c9cf5R69 Therefore, a unit test that relied on how often this parameter was accessed started failing. This PR updates the count to the correct number. Additionally debug print statements that were accidentally left over are now removed.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
3 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What does this PR do?
Add support for expert parallel!