DIstributed branch base by 3outeille · Pull Request #46269 · huggingface/transformers

3outeille · 2026-05-28T17:32:45Z

Current state of things

Training has been confirming to work in TRL thanks to @AmineDiro cf EP + Trainer integration on top of DistributedConfig (#45028) #46126

I was validated end-to-end with trl.SFTTrainer on Qwen3-30B-A3B at multiple shapes (16k / 32k / 64k, 2n / 4n / 8n). Healthy training (loss bit-exact match against our prod fork at LR=0) and very competitive MFU 👏 👏
Inference works but depends on Dtensor (which we don't want as it is very slow)
DTensor-based TP
Native FSDP2 support in transformer
Sequence parallelism for activations and norms via per-model _sp_plan entries
MoE expert parallelism + packed-weight sharding for fused gate_up_proj /
grouped_mm kernels
Shard-on-read via DtensorShardOperation
Resume under a different parallelism layout.
Standard safetensors export. model.save_pretrained(dir) (no split files checkpoint when distributed) writes a fully-gathered
Distributed checkpointing for both model and optimizer.

What's missing:

The PR has been reverted FSDP + TP & native save/load distributed #45028 due to:
- missing support of EP
- Inference on plain tensors for Continuous batching

TODO

Critical (blocking)
- Bring EP => ep_router + moe_experts_ep_allreduce ParallelStyles (work began by @AmineDiro at EP + Trainer integration on top of DistributedConfig (#45028) #46126)
- Mix plain tensors & dtensors in Colwise/Rowwise to enable inference for continuous batching
Medium:
- push checkpoint to hf_buckets by default
- MoEExpertsParallel
  - go from module level to param level ("layers.*.mlp.experts.gate_up_proj": "grouped_gemm", "layers.*.mlp.experts.down_proj": "grouped_gemm", "layers.*.mlp.experts": "moe_experts_allreduce",)
Minor:
- handling sparse + dense in the sp_plan for models like cohere2_moe, mellum, qwen3_moe
- default tp_plan + ep_plan for inference / sp_plan + ep_plan training
- Add DistributedMixin to isolate (save_pretrained(distributed_checkpoint=True) from PreTrainedModel
- cleaning FSDP by unifying manual/auto path

PR

Tp param level #46290 (param_level, ep_router, handling sparse + dense)
- sp + ep training / tp + ep inference #46292 (tp + ep_plan / sp + ep_plan)
  - cleaning _accumulate_local_param_grad #46394
  - tp_inference_training
simplify fsdp #46679
DistributedMixin

Split PR

This reverts commit 295cee3.

HuggingFaceDocBuilderDev · 2026-05-28T17:46:34Z

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.

AmineDiro · 2026-05-29T07:45:14Z

@3outeille & @ArthurZucker should I port #46126 or thats something already in work and just close it ?

3outeille · 2026-05-29T17:49:51Z

@AmineDiro you can change the base to this branch !

…to distributed

* simplify fsdp * add verify_fsdp_plan * linting * remove tie_weights() before distribute_model. Now use config.tie_word_embeddings for plan expectations, not pre-shard aliasing.

* [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

github-actions · 2026-06-17T03:10:14Z

View the CircleCI Test Summary for this PR:

https://huggingface.co/spaces/transformers-community/circle-ci-viz?pr=46269&sha=309729

github-actions · 2026-06-24T05:41:12Z

[For maintainers] Suggested jobs to run (before merge)

run-slow: afmoe, apertus, arcee, aria, audioflamingo3, bamba, bitnet, cohere, cohere2, cohere2_moe, csm, cwm, data2vec, dbrx, deepseek_ocr2, deepseek_v2

@ArthurZucker

* add distributed config * Add native FSDP2 module and migrate FSDP imports (Phase A PR-2). Move FSDP2 wrapping and plan verification to distributed/fsdp.py, keep integrations/fsdp.py as a backward-compatible re-export, and update core call sites to import from transformers.distributed.fsdp. * linting * unecessary * copyright edit * revert * add shard on read * jsut shard on read * cleaning * linting * fix * fix * remove redundant test file * Update src/transformers/distributed/fsdp.py naming Co-authored-by: Arthur <48595927+ArthurZucker@users.noreply.github.com> * avoid looping, just look at dict * expand_fsdp returns reshard_targets, no_reshard_targets right away * better _resolve_tied_embed_lm_head_plan * cleaning * ruff * more robust detection of embed and lm_head * cleaning * ruff * typo * cleaner * cleaner * typo * refactor dense path + apply_contiguous_shard * linting * cleaning * refactor _apply_strided_shard * better * refactor _slice_and_cat * better comment * refactor moe dtensor shard ops * better comment * comment * cleaning * linting * Add FSDP orchestration: mesh init, distribute-before-load, and DCP save. Wire distributed_config from_pretrained/save_pretrained alongside the legacy tp_plan path, add distributed/utils.py for mesh orchestration and checkpoint I/O, and extend sharding_utils with DTensor gather/optimizer fusion helpers needed by save/load. * add fsdp plan to 2 models for now * add tests fsdp mixin * linting * refactor test fsdp mixin * test fsdp mixin cleaning * remove fsdp policy in tests + trim down further * test fsdp clean * restore test_modeling_utils * linting * start trim down stuff * fix * breaking: cleaning modeling_utils.py * load path with fsdp (dtensor) and tp (old tp) is linked * linting * add saving * styling * fix tp ci * add fsdp to ci * linting * pick one model only for this PR * restore * trigger fsdp ci * doc cleaning + tp_size remove * fix tp ci for ep * edit doc * move distributed function to utils + guarding * linting * expand_fsdp_plan iterate over modules * comment about tie embedding * add comment tied embedding * add DistributedMixin * some cleaning * cleaning + comment * rename function for clarity * Apply suggestion from @ArthurZucker Co-authored-by: Arthur <48595927+ArthurZucker@users.noreply.github.com> * doc * comment * linting * refactor * abstract to mixin * typo * Add FSDP plans to all models from distributed branch. Port base_model_fsdp_plan and ForCausalLM _fsdp_plan entries from PR huggingface#46269 and expand FSDP distributed test coverage to the pilot model subset. * fsdp plans * linting * linting * Add distributed runtime utils and DistributedMixin (FSDP orchestration 1/3). Introduce distributed/utils.py and DistributedMixin, defer DistributedConfig validation to load time, and refactor PreTrainedModel plan properties without changing the from_pretrained distributed_config API yet. * Wire DistributedConfig through from_pretrained and save_pretrained (FSDP orchestration 2/3). Route distributed loading and saving through DistributedMixin, migrate TP tests and docs off tp_plan="auto", and add FSDP gather/DCP save paths. * Add FSDP CI and end-to-end FSDP tests (FSDP orchestration 3/3). Add FSDPTesterMixin, cohere2_moe base_fsdp_plan, dedicated fsdp_ci job, and pytest markers for distributed FSDP load/save/generation coverage. * addd ep_plan * restore validate module * Wire DistributedConfig through from_pretrained and save_pretrained. Route distributed load/save orchestration through DistributedMixin so TP and FSDP paths share the same entry points. * revert * inline distribute_model * revert * remove saving/loading * leaner mixin * downgrade torch version guarding * remove * linting * revert * revert * post_init() parallel plan move to mixin * revert tp mixin * add save/load * only FSDP save/load for now * revert * refactor * modular * ea * add it to pretrained model * linting * fix the test by moving in init class the fsdp plan instead of post init * edit * revert unwanted changes * fix ci --------- Co-authored-by: Arthur <48595927+ArthurZucker@users.noreply.github.com>

@ArthurZucker

* add distributed config * Add native FSDP2 module and migrate FSDP imports (Phase A PR-2). Move FSDP2 wrapping and plan verification to distributed/fsdp.py, keep integrations/fsdp.py as a backward-compatible re-export, and update core call sites to import from transformers.distributed.fsdp. * linting * unecessary * copyright edit * revert * add shard on read * jsut shard on read * cleaning * linting * fix * fix * remove redundant test file * Update src/transformers/distributed/fsdp.py naming Co-authored-by: Arthur <48595927+ArthurZucker@users.noreply.github.com> * avoid looping, just look at dict * expand_fsdp returns reshard_targets, no_reshard_targets right away * better _resolve_tied_embed_lm_head_plan * cleaning * ruff * more robust detection of embed and lm_head * cleaning * ruff * typo * cleaner * cleaner * typo * refactor dense path + apply_contiguous_shard * linting * cleaning * refactor _apply_strided_shard * better * refactor _slice_and_cat * better comment * refactor moe dtensor shard ops * better comment * comment * cleaning * linting * Add FSDP orchestration: mesh init, distribute-before-load, and DCP save. Wire distributed_config from_pretrained/save_pretrained alongside the legacy tp_plan path, add distributed/utils.py for mesh orchestration and checkpoint I/O, and extend sharding_utils with DTensor gather/optimizer fusion helpers needed by save/load. * add fsdp plan to 2 models for now * add tests fsdp mixin * linting * refactor test fsdp mixin * test fsdp mixin cleaning * remove fsdp policy in tests + trim down further * test fsdp clean * restore test_modeling_utils * linting * start trim down stuff * fix * breaking: cleaning modeling_utils.py * load path with fsdp (dtensor) and tp (old tp) is linked * linting * add saving * styling * fix tp ci * add fsdp to ci * linting * pick one model only for this PR * restore * trigger fsdp ci * doc cleaning + tp_size remove * fix tp ci for ep * edit doc * move distributed function to utils + guarding * linting * expand_fsdp_plan iterate over modules * comment about tie embedding * add comment tied embedding * add DistributedMixin * some cleaning * cleaning + comment * rename function for clarity * Apply suggestion from @ArthurZucker Co-authored-by: Arthur <48595927+ArthurZucker@users.noreply.github.com> * doc * comment * linting * refactor * abstract to mixin * typo * Add FSDP plans to all models from distributed branch. Port base_model_fsdp_plan and ForCausalLM _fsdp_plan entries from PR huggingface#46269 and expand FSDP distributed test coverage to the pilot model subset. * fsdp plans * linting * linting * Add distributed runtime utils and DistributedMixin (FSDP orchestration 1/3). Introduce distributed/utils.py and DistributedMixin, defer DistributedConfig validation to load time, and refactor PreTrainedModel plan properties without changing the from_pretrained distributed_config API yet. * Wire DistributedConfig through from_pretrained and save_pretrained (FSDP orchestration 2/3). Route distributed loading and saving through DistributedMixin, migrate TP tests and docs off tp_plan="auto", and add FSDP gather/DCP save paths. * Add FSDP CI and end-to-end FSDP tests (FSDP orchestration 3/3). Add FSDPTesterMixin, cohere2_moe base_fsdp_plan, dedicated fsdp_ci job, and pytest markers for distributed FSDP load/save/generation coverage. * addd ep_plan * restore validate module * Wire DistributedConfig through from_pretrained and save_pretrained. Route distributed load/save orchestration through DistributedMixin so TP and FSDP paths share the same entry points. * revert * inline distribute_model * revert * remove saving/loading * leaner mixin * downgrade torch version guarding * remove * linting * revert * revert * post_init() parallel plan move to mixin * revert tp mixin * add save/load * only FSDP save/load for now * revert * refactor * modular * ea * add it to pretrained model * linting * fix the test by moving in init class the fsdp plan instead of post init * edit * revert unwanted changes * fix ci --------- Co-authored-by: Arthur <48595927+ArthurZucker@users.noreply.github.com>

@ArthurZucker

* add distributed config * Add native FSDP2 module and migrate FSDP imports (Phase A PR-2). Move FSDP2 wrapping and plan verification to distributed/fsdp.py, keep integrations/fsdp.py as a backward-compatible re-export, and update core call sites to import from transformers.distributed.fsdp. * linting * unecessary * copyright edit * revert * add shard on read * jsut shard on read * cleaning * linting * fix * fix * remove redundant test file * Update src/transformers/distributed/fsdp.py naming Co-authored-by: Arthur <48595927+ArthurZucker@users.noreply.github.com> * avoid looping, just look at dict * expand_fsdp returns reshard_targets, no_reshard_targets right away * better _resolve_tied_embed_lm_head_plan * cleaning * ruff * more robust detection of embed and lm_head * cleaning * ruff * typo * cleaner * cleaner * typo * refactor dense path + apply_contiguous_shard * linting * cleaning * refactor _apply_strided_shard * better * refactor _slice_and_cat * better comment * refactor moe dtensor shard ops * better comment * comment * cleaning * linting * Add FSDP orchestration: mesh init, distribute-before-load, and DCP save. Wire distributed_config from_pretrained/save_pretrained alongside the legacy tp_plan path, add distributed/utils.py for mesh orchestration and checkpoint I/O, and extend sharding_utils with DTensor gather/optimizer fusion helpers needed by save/load. * add fsdp plan to 2 models for now * add tests fsdp mixin * linting * refactor test fsdp mixin * test fsdp mixin cleaning * remove fsdp policy in tests + trim down further * test fsdp clean * restore test_modeling_utils * linting * start trim down stuff * fix * breaking: cleaning modeling_utils.py * load path with fsdp (dtensor) and tp (old tp) is linked * linting * add saving * styling * fix tp ci * add fsdp to ci * linting * pick one model only for this PR * restore * trigger fsdp ci * doc cleaning + tp_size remove * fix tp ci for ep * edit doc * move distributed function to utils + guarding * linting * expand_fsdp_plan iterate over modules * comment about tie embedding * add comment tied embedding * add DistributedMixin * some cleaning * cleaning + comment * rename function for clarity * Apply suggestion from @ArthurZucker Co-authored-by: Arthur <48595927+ArthurZucker@users.noreply.github.com> * doc * comment * linting * refactor * abstract to mixin * typo * Add FSDP plans to all models from distributed branch. Port base_model_fsdp_plan and ForCausalLM _fsdp_plan entries from PR huggingface#46269 and expand FSDP distributed test coverage to the pilot model subset. * fsdp plans * linting * linting * Add distributed runtime utils and DistributedMixin (FSDP orchestration 1/3). Introduce distributed/utils.py and DistributedMixin, defer DistributedConfig validation to load time, and refactor PreTrainedModel plan properties without changing the from_pretrained distributed_config API yet. * Wire DistributedConfig through from_pretrained and save_pretrained (FSDP orchestration 2/3). Route distributed loading and saving through DistributedMixin, migrate TP tests and docs off tp_plan="auto", and add FSDP gather/DCP save paths. * Add FSDP CI and end-to-end FSDP tests (FSDP orchestration 3/3). Add FSDPTesterMixin, cohere2_moe base_fsdp_plan, dedicated fsdp_ci job, and pytest markers for distributed FSDP load/save/generation coverage. * addd ep_plan * restore validate module * Wire DistributedConfig through from_pretrained and save_pretrained. Route distributed load/save orchestration through DistributedMixin so TP and FSDP paths share the same entry points. * revert * inline distribute_model * revert * remove saving/loading * leaner mixin * downgrade torch version guarding * remove * linting * revert * revert * post_init() parallel plan move to mixin * revert tp mixin * add save/load * only FSDP save/load for now * revert * refactor * modular * ea * add it to pretrained model * linting * fix the test by moving in init class the fsdp plan instead of post init * edit * replace everything * revert unwanted changes * fix ci --------- Co-authored-by: Arthur <48595927+ArthurZucker@users.noreply.github.com>

AmineDiro · 2026-08-03T10:21:29Z

I'm flagging a missing path in select_parallel_plan and clarifying a previous recommendation of mine that likely contributed to this issue 🙈 .

Currently, there is no supported setup for training Expert Parallelism (EP) without Sequence Parallelism (SP) !

So using EP alone defaults to _tp_ep_plan (inference TP + EP) and because I run experiment on the trl.SFTTrainer with only FSDP2 + EP , I didnt need to use SP or attention TP. It incorrectly falls back to this inference plan and shards q/k/v column-wise.

This silently caps the EP size at num_key_value_heads for GQA MoE Models. In my case Qwen3-30B-A3B has 32 query heads but only 4 KV heads. At tp_size=8, k_proj outputs 64 features per rank (half a KV head), causing an attention reshape failure:

RuntimeError: shape '[1, 8192, -1, 128]' is invalid for input of size 524288

This means we effectively forces expert_parallel_size <= num_key_value_heads instead of properly bounding it by the 128 experts.

The correct plan for this case already exists: it's base_model_ep_plan is expert-only with replicated attention, but it is never selected. The standalone _ep_plan handling in modeling_utils.py seems to be a dropped path rather than an intentional decision.

Context on my previous advice: I previously told @3outeille that SP+EP scales best. However, I was specifically referring to Ulysses SP (my bad 🙈 ) , not Megatron-style SP. Megatron SP and EP share a single TP mesh dimension, meaning a single sharding size must satisfy both the attention head cap and the expert count. That was my mistake.

So waht I recommend is the following:

1. Add a non-breaking fallback: Introduce DistributedConfig.shard_attention: bool = True.
- If True (default behavior), use _tp_ep_plan.
- If False, use _ep_plan and restore the missing plumbing in modeling_utils.py.
- Result: Testing False on tiny-Qwen3MoeForCausalLM correctly replicates attention. This successfully allows EP=8 on a 4-KV-head model (Qwen3-30B-A3B at 1 node EP=8 reaches 27.8% MFU).
**2. Raise a clear error early when tp_size > num_key_value_heads instead of failing deeper in the stack during the reshape

wdyt @3outeille ?

@ArthurZucker

* add distributed config * Add native FSDP2 module and migrate FSDP imports (Phase A PR-2). Move FSDP2 wrapping and plan verification to distributed/fsdp.py, keep integrations/fsdp.py as a backward-compatible re-export, and update core call sites to import from transformers.distributed.fsdp. * linting * unecessary * copyright edit * revert * add shard on read * jsut shard on read * cleaning * linting * fix * fix * remove redundant test file * Update src/transformers/distributed/fsdp.py naming Co-authored-by: Arthur <48595927+ArthurZucker@users.noreply.github.com> * avoid looping, just look at dict * expand_fsdp returns reshard_targets, no_reshard_targets right away * better _resolve_tied_embed_lm_head_plan * cleaning * ruff * more robust detection of embed and lm_head * cleaning * ruff * typo * cleaner * cleaner * typo * refactor dense path + apply_contiguous_shard * linting * cleaning * refactor _apply_strided_shard * better * refactor _slice_and_cat * better comment * refactor moe dtensor shard ops * better comment * comment * cleaning * linting * Add FSDP orchestration: mesh init, distribute-before-load, and DCP save. Wire distributed_config from_pretrained/save_pretrained alongside the legacy tp_plan path, add distributed/utils.py for mesh orchestration and checkpoint I/O, and extend sharding_utils with DTensor gather/optimizer fusion helpers needed by save/load. * add fsdp plan to 2 models for now * add tests fsdp mixin * linting * refactor test fsdp mixin * test fsdp mixin cleaning * remove fsdp policy in tests + trim down further * test fsdp clean * restore test_modeling_utils * linting * start trim down stuff * fix * breaking: cleaning modeling_utils.py * load path with fsdp (dtensor) and tp (old tp) is linked * linting * add saving * styling * fix tp ci * add fsdp to ci * linting * pick one model only for this PR * restore * trigger fsdp ci * doc cleaning + tp_size remove * fix tp ci for ep * edit doc * move distributed function to utils + guarding * linting * expand_fsdp_plan iterate over modules * comment about tie embedding * add comment tied embedding * add DistributedMixin * some cleaning * cleaning + comment * rename function for clarity * Apply suggestion from @ArthurZucker Co-authored-by: Arthur <48595927+ArthurZucker@users.noreply.github.com> * doc * comment * linting * refactor * abstract to mixin * typo * Add FSDP plans to all models from distributed branch. Port base_model_fsdp_plan and ForCausalLM _fsdp_plan entries from PR huggingface#46269 and expand FSDP distributed test coverage to the pilot model subset. * fsdp plans * linting * linting * Add distributed runtime utils and DistributedMixin (FSDP orchestration 1/3). Introduce distributed/utils.py and DistributedMixin, defer DistributedConfig validation to load time, and refactor PreTrainedModel plan properties without changing the from_pretrained distributed_config API yet. * Wire DistributedConfig through from_pretrained and save_pretrained (FSDP orchestration 2/3). Route distributed loading and saving through DistributedMixin, migrate TP tests and docs off tp_plan="auto", and add FSDP gather/DCP save paths. * Add FSDP CI and end-to-end FSDP tests (FSDP orchestration 3/3). Add FSDPTesterMixin, cohere2_moe base_fsdp_plan, dedicated fsdp_ci job, and pytest markers for distributed FSDP load/save/generation coverage. * addd ep_plan * restore validate module * Wire DistributedConfig through from_pretrained and save_pretrained. Route distributed load/save orchestration through DistributedMixin so TP and FSDP paths share the same entry points. * revert * inline distribute_model * revert * remove saving/loading * leaner mixin * downgrade torch version guarding * remove * linting * revert * revert * post_init() parallel plan move to mixin * revert tp mixin * add save/load * only FSDP save/load for now * revert * refactor * modular * ea * add it to pretrained model * linting * fix the test by moving in init class the fsdp plan instead of post init * edit * revert unwanted changes * fix ci --------- Co-authored-by: Arthur <48595927+ArthurZucker@users.noreply.github.com>

@ArthurZucker

* add distributed config * Add native FSDP2 module and migrate FSDP imports (Phase A PR-2). Move FSDP2 wrapping and plan verification to distributed/fsdp.py, keep integrations/fsdp.py as a backward-compatible re-export, and update core call sites to import from transformers.distributed.fsdp. * linting * unecessary * copyright edit * revert * add shard on read * jsut shard on read * cleaning * linting * fix * fix * remove redundant test file * Update src/transformers/distributed/fsdp.py naming Co-authored-by: Arthur <48595927+ArthurZucker@users.noreply.github.com> * avoid looping, just look at dict * expand_fsdp returns reshard_targets, no_reshard_targets right away * better _resolve_tied_embed_lm_head_plan * cleaning * ruff * more robust detection of embed and lm_head * cleaning * ruff * typo * cleaner * cleaner * typo * refactor dense path + apply_contiguous_shard * linting * cleaning * refactor _apply_strided_shard * better * refactor _slice_and_cat * better comment * refactor moe dtensor shard ops * better comment * comment * cleaning * linting * Add FSDP orchestration: mesh init, distribute-before-load, and DCP save. Wire distributed_config from_pretrained/save_pretrained alongside the legacy tp_plan path, add distributed/utils.py for mesh orchestration and checkpoint I/O, and extend sharding_utils with DTensor gather/optimizer fusion helpers needed by save/load. * add fsdp plan to 2 models for now * add tests fsdp mixin * linting * refactor test fsdp mixin * test fsdp mixin cleaning * remove fsdp policy in tests + trim down further * test fsdp clean * restore test_modeling_utils * linting * start trim down stuff * fix * breaking: cleaning modeling_utils.py * load path with fsdp (dtensor) and tp (old tp) is linked * linting * add saving * styling * fix tp ci * add fsdp to ci * linting * pick one model only for this PR * restore * trigger fsdp ci * doc cleaning + tp_size remove * fix tp ci for ep * edit doc * move distributed function to utils + guarding * linting * expand_fsdp_plan iterate over modules * comment about tie embedding * add comment tied embedding * add DistributedMixin * some cleaning * cleaning + comment * rename function for clarity * Apply suggestion from @ArthurZucker Co-authored-by: Arthur <48595927+ArthurZucker@users.noreply.github.com> * doc * comment * linting * refactor * abstract to mixin * typo * Add FSDP plans to all models from distributed branch. Port base_model_fsdp_plan and ForCausalLM _fsdp_plan entries from PR huggingface#46269 and expand FSDP distributed test coverage to the pilot model subset. * fsdp plans * linting * linting * Add distributed runtime utils and DistributedMixin (FSDP orchestration 1/3). Introduce distributed/utils.py and DistributedMixin, defer DistributedConfig validation to load time, and refactor PreTrainedModel plan properties without changing the from_pretrained distributed_config API yet. * Wire DistributedConfig through from_pretrained and save_pretrained (FSDP orchestration 2/3). Route distributed loading and saving through DistributedMixin, migrate TP tests and docs off tp_plan="auto", and add FSDP gather/DCP save paths. * Add FSDP CI and end-to-end FSDP tests (FSDP orchestration 3/3). Add FSDPTesterMixin, cohere2_moe base_fsdp_plan, dedicated fsdp_ci job, and pytest markers for distributed FSDP load/save/generation coverage. * addd ep_plan * restore validate module * Wire DistributedConfig through from_pretrained and save_pretrained. Route distributed load/save orchestration through DistributedMixin so TP and FSDP paths share the same entry points. * revert * inline distribute_model * revert * remove saving/loading * leaner mixin * downgrade torch version guarding * remove * linting * revert * revert * post_init() parallel plan move to mixin * revert tp mixin * add save/load * only FSDP save/load for now * revert * refactor * modular * ea * add it to pretrained model * linting * fix the test by moving in init class the fsdp plan instead of post init * edit * replace everything * revert unwanted changes * fix ci --------- Co-authored-by: Arthur <48595927+ArthurZucker@users.noreply.github.com>

@ArthurZucker

* add distributed config * Add native FSDP2 module and migrate FSDP imports (Phase A PR-2). Move FSDP2 wrapping and plan verification to distributed/fsdp.py, keep integrations/fsdp.py as a backward-compatible re-export, and update core call sites to import from transformers.distributed.fsdp. * linting * unecessary * copyright edit * revert * add shard on read * jsut shard on read * cleaning * linting * fix * fix * remove redundant test file * Update src/transformers/distributed/fsdp.py naming Co-authored-by: Arthur <48595927+ArthurZucker@users.noreply.github.com> * avoid looping, just look at dict * expand_fsdp returns reshard_targets, no_reshard_targets right away * better _resolve_tied_embed_lm_head_plan * cleaning * ruff * more robust detection of embed and lm_head * cleaning * ruff * typo * cleaner * cleaner * typo * refactor dense path + apply_contiguous_shard * linting * cleaning * refactor _apply_strided_shard * better * refactor _slice_and_cat * better comment * refactor moe dtensor shard ops * better comment * comment * cleaning * linting * Add FSDP orchestration: mesh init, distribute-before-load, and DCP save. Wire distributed_config from_pretrained/save_pretrained alongside the legacy tp_plan path, add distributed/utils.py for mesh orchestration and checkpoint I/O, and extend sharding_utils with DTensor gather/optimizer fusion helpers needed by save/load. * add fsdp plan to 2 models for now * add tests fsdp mixin * linting * refactor test fsdp mixin * test fsdp mixin cleaning * remove fsdp policy in tests + trim down further * test fsdp clean * restore test_modeling_utils * linting * start trim down stuff * fix * breaking: cleaning modeling_utils.py * load path with fsdp (dtensor) and tp (old tp) is linked * linting * add saving * styling * fix tp ci * add fsdp to ci * linting * pick one model only for this PR * restore * trigger fsdp ci * doc cleaning + tp_size remove * fix tp ci for ep * edit doc * move distributed function to utils + guarding * linting * expand_fsdp_plan iterate over modules * comment about tie embedding * add comment tied embedding * add DistributedMixin * some cleaning * cleaning + comment * rename function for clarity * Apply suggestion from @ArthurZucker Co-authored-by: Arthur <48595927+ArthurZucker@users.noreply.github.com> * doc * comment * linting * refactor * abstract to mixin * typo * Add FSDP plans to all models from distributed branch. Port base_model_fsdp_plan and ForCausalLM _fsdp_plan entries from PR huggingface#46269 and expand FSDP distributed test coverage to the pilot model subset. * fsdp plans * linting * linting * Add distributed runtime utils and DistributedMixin (FSDP orchestration 1/3). Introduce distributed/utils.py and DistributedMixin, defer DistributedConfig validation to load time, and refactor PreTrainedModel plan properties without changing the from_pretrained distributed_config API yet. * Wire DistributedConfig through from_pretrained and save_pretrained (FSDP orchestration 2/3). Route distributed loading and saving through DistributedMixin, migrate TP tests and docs off tp_plan="auto", and add FSDP gather/DCP save paths. * Add FSDP CI and end-to-end FSDP tests (FSDP orchestration 3/3). Add FSDPTesterMixin, cohere2_moe base_fsdp_plan, dedicated fsdp_ci job, and pytest markers for distributed FSDP load/save/generation coverage. * addd ep_plan * restore validate module * Wire DistributedConfig through from_pretrained and save_pretrained. Route distributed load/save orchestration through DistributedMixin so TP and FSDP paths share the same entry points. * revert * inline distribute_model * revert * remove saving/loading * leaner mixin * downgrade torch version guarding * remove * linting * revert * revert * post_init() parallel plan move to mixin * revert tp mixin * add save/load * only FSDP save/load for now * revert * refactor * modular * ea * begin migration TP * clean up * add it to pretrained model * linting * fix the test by moving in init class the fsdp plan instead of post init * edit * replace everything * begin migration TP * clean up * revert merge conflict * local params for forward * cleaning * migration from integration.tensor_parallel to distributed.tensor_parallel * revert models * revert * fix * breaking: cleaner way to_local for quantize weight(almost done need to fix backward) * fix all tests * TP dtensor handle natively the to_local() for deepgem + fp8 (huggingface#47634) * Refactor distributed tensor_parallel module and TP mixin tests * better approach to to_local * pass as class method * cleaning PackedColwise * remove the use of to_local in kernels to offload this task to TensorParallel * rowwise bias after redistribute * make it easier to understand * requires grad only for floating point numbers * fix tp meagamoe fp8 with dtensor * linting * revert gitignore * fix run slow quantization ci * fix deepseek v4 ep backward tests * fix deepseek glm4 moe tp backward pass * Refactor tensor parallel loading logic by removing unused tp_plan handling and updating is_dtensor type hint for better type safety. * linting * bench: dtensor vs legacy TP * bench: dtensor vs legacy TP (huggingface#47728) * bench: dtensor vs legacy TP * clean rowise * refactor colwise * comment * comment * refactor MoEExperParallel * ruff * add rowwise input * remove helper script * remove dead code in mxfp4 * remove old TP * cleaning * cleaning * claening * add test_shard_tensor_shape_consistency * cleaning * small fix * typo * comment * rename colwise_gather_output to colwise_rep * keep tp import to avoid breaking chnges * remove async_op=True for redistribute * linting * renaming MoeExpertsParallel class * ep router doc * mlinter * add todo * fix * remove unused keys from global_mapping to avoid BC --------- Co-authored-by: Arthur <48595927+ArthurZucker@users.noreply.github.com>

@ArthurZucker

* add fsdp plan to 2 models for now * add tests fsdp mixin * linting * refactor test fsdp mixin * test fsdp mixin cleaning * remove fsdp policy in tests + trim down further * test fsdp clean * restore test_modeling_utils * linting * start trim down stuff * fix * breaking: cleaning modeling_utils.py * load path with fsdp (dtensor) and tp (old tp) is linked * linting * add saving * styling * fix tp ci * add fsdp to ci * linting * pick one model only for this PR * restore * trigger fsdp ci * doc cleaning + tp_size remove * fix tp ci for ep * edit doc * move distributed function to utils + guarding * linting * expand_fsdp_plan iterate over modules * comment about tie embedding * add comment tied embedding * add DistributedMixin * some cleaning * cleaning + comment * rename function for clarity * Apply suggestion from @ArthurZucker Co-authored-by: Arthur <48595927+ArthurZucker@users.noreply.github.com> * doc * comment * linting * refactor * abstract to mixin * typo * Add FSDP plans to all models from distributed branch. Port base_model_fsdp_plan and ForCausalLM _fsdp_plan entries from PR huggingface#46269 and expand FSDP distributed test coverage to the pilot model subset. * fsdp plans * linting * linting * Add distributed runtime utils and DistributedMixin (FSDP orchestration 1/3). Introduce distributed/utils.py and DistributedMixin, defer DistributedConfig validation to load time, and refactor PreTrainedModel plan properties without changing the from_pretrained distributed_config API yet. * Wire DistributedConfig through from_pretrained and save_pretrained (FSDP orchestration 2/3). Route distributed loading and saving through DistributedMixin, migrate TP tests and docs off tp_plan="auto", and add FSDP gather/DCP save paths. * Add FSDP CI and end-to-end FSDP tests (FSDP orchestration 3/3). Add FSDPTesterMixin, cohere2_moe base_fsdp_plan, dedicated fsdp_ci job, and pytest markers for distributed FSDP load/save/generation coverage. * addd ep_plan * restore validate module * Wire DistributedConfig through from_pretrained and save_pretrained. Route distributed load/save orchestration through DistributedMixin so TP and FSDP paths share the same entry points. * revert * inline distribute_model * revert * remove saving/loading * leaner mixin * downgrade torch version guarding * remove * linting * revert * revert * post_init() parallel plan move to mixin * revert tp mixin * add save/load * adapt distributed config to handle pp_size * edit validate * init pipeline parallelism * avoid loading weights with TP when pp_size>1 * cleaning args in initialize_pipeline_parallelism * add very dumb layers splitting * add test pipeline parallel * draft pp + generate test to be clean * add weight loading logging * cleaning * remove comm_device to make it an attribute * begin to refactor * avoid loading report when tp is enabled * cleaning pp * better cleaning pp * cleaning pp * better cleaning broadcast_from_last * cleaning pp * pp tests no mixing * handle tied embeddings * linting * doc * unecessary if * comment * linting * make sure pp doesnt go into tp shard and distribute * only FSDP save/load for now * revert * refactor * modular * ea * begin migration TP * clean up * add it to pretrained model * linting * fix the test by moving in init class the fsdp plan instead of post init * edit * replace everything * adapt distributed config to handle pp_size * init pipeline parallelism * avoid loading weights with TP when pp_size>1 * cleaning args in initialize_pipeline_parallelism * add very dumb layers splitting * add test pipeline parallel * draft pp + generate test to be clean * add weight loading logging * cleaning * remove comm_device to make it an attribute * begin to refactor * avoid loading report when tp is enabled * cleaning pp * better cleaning pp * cleaning pp * better cleaning broadcast_from_last * cleaning pp * pp tests no mixing * handle tied embeddings * linting * doc * unecessary if * comment * linting * make sure pp doesnt go into tp shard and distribute * begin migration TP * clean up * revert merge conflict * local params for forward * cleaning * migration from integration.tensor_parallel to distributed.tensor_parallel * revert models * revert * fix * breaking: cleaner way to_local for quantize weight(almost done need to fix backward) * fix all tests * TP dtensor handle natively the to_local() for deepgem + fp8 (huggingface#47634) * Refactor distributed tensor_parallel module and TP mixin tests * better approach to to_local * pass as class method * cleaning PackedColwise * remove the use of to_local in kernels to offload this task to TensorParallel * rowwise bias after redistribute * make it easier to understand * requires grad only for floating point numbers * fix tp meagamoe fp8 with dtensor * linting * revert gitignore * fix run slow quantization ci * fix deepseek v4 ep backward tests * fix deepseek glm4 moe tp backward pass * Refactor tensor parallel loading logic by removing unused tp_plan handling and updating is_dtensor type hint for better type safety. * linting * bench: dtensor vs legacy TP * bench: dtensor vs legacy TP (huggingface#47728) * bench: dtensor vs legacy TP * clean rowise * refactor colwise * comment * comment * refactor MoEExperParallel * ruff * add rowwise input * remove helper script * remove dead code in mxfp4 * remove old TP * cleaning * cleaning * claening * add test_shard_tensor_shape_consistency * cleaning * small fix * typo * comment * clean * typo * rename colwise_gather_output to colwise_rep * keep tp import to avoid breaking chnges * remove async_op=True for redistribute * linting * renaming MoeExpertsParallel class * ep router doc * mlinter * add todo * fix * new name * guard import * cache inspect signature * merge init and from_device_mesh * add todo * avoid nesting * pp stages logs only ad INFO level * better doc func * linting * linting * fix --------- Co-authored-by: Arthur <48595927+ArthurZucker@users.noreply.github.com>

@ArthurZucker

* add distributed config * Add native FSDP2 module and migrate FSDP imports (Phase A PR-2). Move FSDP2 wrapping and plan verification to distributed/fsdp.py, keep integrations/fsdp.py as a backward-compatible re-export, and update core call sites to import from transformers.distributed.fsdp. * linting * unecessary * copyright edit * revert * add shard on read * jsut shard on read * cleaning * linting * fix * fix * remove redundant test file * Update src/transformers/distributed/fsdp.py naming Co-authored-by: Arthur <48595927+ArthurZucker@users.noreply.github.com> * avoid looping, just look at dict * expand_fsdp returns reshard_targets, no_reshard_targets right away * better _resolve_tied_embed_lm_head_plan * cleaning * ruff * more robust detection of embed and lm_head * cleaning * ruff * typo * cleaner * cleaner * typo * refactor dense path + apply_contiguous_shard * linting * cleaning * refactor _apply_strided_shard * better * refactor _slice_and_cat * better comment * refactor moe dtensor shard ops * better comment * comment * cleaning * linting * Add FSDP orchestration: mesh init, distribute-before-load, and DCP save. Wire distributed_config from_pretrained/save_pretrained alongside the legacy tp_plan path, add distributed/utils.py for mesh orchestration and checkpoint I/O, and extend sharding_utils with DTensor gather/optimizer fusion helpers needed by save/load. * add fsdp plan to 2 models for now * add tests fsdp mixin * linting * refactor test fsdp mixin * test fsdp mixin cleaning * remove fsdp policy in tests + trim down further * test fsdp clean * restore test_modeling_utils * linting * start trim down stuff * fix * breaking: cleaning modeling_utils.py * load path with fsdp (dtensor) and tp (old tp) is linked * linting * add saving * styling * fix tp ci * add fsdp to ci * linting * pick one model only for this PR * restore * trigger fsdp ci * doc cleaning + tp_size remove * fix tp ci for ep * edit doc * move distributed function to utils + guarding * linting * expand_fsdp_plan iterate over modules * comment about tie embedding * add comment tied embedding * add DistributedMixin * some cleaning * cleaning + comment * rename function for clarity * Apply suggestion from @ArthurZucker Co-authored-by: Arthur <48595927+ArthurZucker@users.noreply.github.com> * doc * comment * linting * refactor * abstract to mixin * typo * Add FSDP plans to all models from distributed branch. Port base_model_fsdp_plan and ForCausalLM _fsdp_plan entries from PR huggingface#46269 and expand FSDP distributed test coverage to the pilot model subset. * fsdp plans * linting * linting * Add distributed runtime utils and DistributedMixin (FSDP orchestration 1/3). Introduce distributed/utils.py and DistributedMixin, defer DistributedConfig validation to load time, and refactor PreTrainedModel plan properties without changing the from_pretrained distributed_config API yet. * Wire DistributedConfig through from_pretrained and save_pretrained (FSDP orchestration 2/3). Route distributed loading and saving through DistributedMixin, migrate TP tests and docs off tp_plan="auto", and add FSDP gather/DCP save paths. * Add FSDP CI and end-to-end FSDP tests (FSDP orchestration 3/3). Add FSDPTesterMixin, cohere2_moe base_fsdp_plan, dedicated fsdp_ci job, and pytest markers for distributed FSDP load/save/generation coverage. * addd ep_plan * restore validate module * Wire DistributedConfig through from_pretrained and save_pretrained. Route distributed load/save orchestration through DistributedMixin so TP and FSDP paths share the same entry points. * revert * inline distribute_model * revert * remove saving/loading * leaner mixin * downgrade torch version guarding * remove * linting * revert * revert * post_init() parallel plan move to mixin * revert tp mixin * add save/load * only FSDP save/load for now * revert * refactor * modular * ea * begin migration TP * clean up * add it to pretrained model * linting * fix the test by moving in init class the fsdp plan instead of post init * edit * replace everything * begin migration TP * clean up * revert merge conflict * local params for forward * cleaning * migration from integration.tensor_parallel to distributed.tensor_parallel * revert models * revert * fix * breaking: cleaner way to_local for quantize weight(almost done need to fix backward) * fix all tests * TP dtensor handle natively the to_local() for deepgem + fp8 (huggingface#47634) * Refactor distributed tensor_parallel module and TP mixin tests * better approach to to_local * pass as class method * cleaning PackedColwise * remove the use of to_local in kernels to offload this task to TensorParallel * rowwise bias after redistribute * make it easier to understand * requires grad only for floating point numbers * fix tp meagamoe fp8 with dtensor * linting * revert gitignore * fix run slow quantization ci * fix deepseek v4 ep backward tests * fix deepseek glm4 moe tp backward pass * Refactor tensor parallel loading logic by removing unused tp_plan handling and updating is_dtensor type hint for better type safety. * linting * bench: dtensor vs legacy TP * bench: dtensor vs legacy TP (huggingface#47728) * bench: dtensor vs legacy TP * clean rowise * refactor colwise * comment * comment * refactor MoEExperParallel * ruff * add rowwise input * remove helper script * remove dead code in mxfp4 * remove old TP * cleaning * cleaning * claening * add test_shard_tensor_shape_consistency * cleaning * small fix * typo * comment * rename colwise_gather_output to colwise_rep * keep tp import to avoid breaking chnges * remove async_op=True for redistribute * linting * renaming MoeExpertsParallel class * ep router doc * mlinter * add todo * fix * remove unused keys from global_mapping to avoid BC --------- Co-authored-by: Arthur <48595927+ArthurZucker@users.noreply.github.com>

@ArthurZucker

* add fsdp plan to 2 models for now * add tests fsdp mixin * linting * refactor test fsdp mixin * test fsdp mixin cleaning * remove fsdp policy in tests + trim down further * test fsdp clean * restore test_modeling_utils * linting * start trim down stuff * fix * breaking: cleaning modeling_utils.py * load path with fsdp (dtensor) and tp (old tp) is linked * linting * add saving * styling * fix tp ci * add fsdp to ci * linting * pick one model only for this PR * restore * trigger fsdp ci * doc cleaning + tp_size remove * fix tp ci for ep * edit doc * move distributed function to utils + guarding * linting * expand_fsdp_plan iterate over modules * comment about tie embedding * add comment tied embedding * add DistributedMixin * some cleaning * cleaning + comment * rename function for clarity * Apply suggestion from @ArthurZucker Co-authored-by: Arthur <48595927+ArthurZucker@users.noreply.github.com> * doc * comment * linting * refactor * abstract to mixin * typo * Add FSDP plans to all models from distributed branch. Port base_model_fsdp_plan and ForCausalLM _fsdp_plan entries from PR huggingface#46269 and expand FSDP distributed test coverage to the pilot model subset. * fsdp plans * linting * linting * Add distributed runtime utils and DistributedMixin (FSDP orchestration 1/3). Introduce distributed/utils.py and DistributedMixin, defer DistributedConfig validation to load time, and refactor PreTrainedModel plan properties without changing the from_pretrained distributed_config API yet. * Wire DistributedConfig through from_pretrained and save_pretrained (FSDP orchestration 2/3). Route distributed loading and saving through DistributedMixin, migrate TP tests and docs off tp_plan="auto", and add FSDP gather/DCP save paths. * Add FSDP CI and end-to-end FSDP tests (FSDP orchestration 3/3). Add FSDPTesterMixin, cohere2_moe base_fsdp_plan, dedicated fsdp_ci job, and pytest markers for distributed FSDP load/save/generation coverage. * addd ep_plan * restore validate module * Wire DistributedConfig through from_pretrained and save_pretrained. Route distributed load/save orchestration through DistributedMixin so TP and FSDP paths share the same entry points. * revert * inline distribute_model * revert * remove saving/loading * leaner mixin * downgrade torch version guarding * remove * linting * revert * revert * post_init() parallel plan move to mixin * revert tp mixin * add save/load * adapt distributed config to handle pp_size * edit validate * init pipeline parallelism * avoid loading weights with TP when pp_size>1 * cleaning args in initialize_pipeline_parallelism * add very dumb layers splitting * add test pipeline parallel * draft pp + generate test to be clean * add weight loading logging * cleaning * remove comm_device to make it an attribute * begin to refactor * avoid loading report when tp is enabled * cleaning pp * better cleaning pp * cleaning pp * better cleaning broadcast_from_last * cleaning pp * pp tests no mixing * handle tied embeddings * linting * doc * unecessary if * comment * linting * make sure pp doesnt go into tp shard and distribute * only FSDP save/load for now * revert * refactor * modular * ea * begin migration TP * clean up * add it to pretrained model * linting * fix the test by moving in init class the fsdp plan instead of post init * edit * replace everything * adapt distributed config to handle pp_size * init pipeline parallelism * avoid loading weights with TP when pp_size>1 * cleaning args in initialize_pipeline_parallelism * add very dumb layers splitting * add test pipeline parallel * draft pp + generate test to be clean * add weight loading logging * cleaning * remove comm_device to make it an attribute * begin to refactor * avoid loading report when tp is enabled * cleaning pp * better cleaning pp * cleaning pp * better cleaning broadcast_from_last * cleaning pp * pp tests no mixing * handle tied embeddings * linting * doc * unecessary if * comment * linting * make sure pp doesnt go into tp shard and distribute * begin migration TP * clean up * revert merge conflict * local params for forward * cleaning * migration from integration.tensor_parallel to distributed.tensor_parallel * revert models * revert * fix * breaking: cleaner way to_local for quantize weight(almost done need to fix backward) * fix all tests * TP dtensor handle natively the to_local() for deepgem + fp8 (huggingface#47634) * Refactor distributed tensor_parallel module and TP mixin tests * better approach to to_local * pass as class method * cleaning PackedColwise * remove the use of to_local in kernels to offload this task to TensorParallel * rowwise bias after redistribute * make it easier to understand * requires grad only for floating point numbers * fix tp meagamoe fp8 with dtensor * linting * revert gitignore * fix run slow quantization ci * fix deepseek v4 ep backward tests * fix deepseek glm4 moe tp backward pass * Refactor tensor parallel loading logic by removing unused tp_plan handling and updating is_dtensor type hint for better type safety. * linting * bench: dtensor vs legacy TP * bench: dtensor vs legacy TP (huggingface#47728) * bench: dtensor vs legacy TP * clean rowise * refactor colwise * comment * comment * refactor MoEExperParallel * ruff * add rowwise input * remove helper script * remove dead code in mxfp4 * remove old TP * cleaning * cleaning * claening * add test_shard_tensor_shape_consistency * cleaning * small fix * typo * comment * clean * typo * rename colwise_gather_output to colwise_rep * keep tp import to avoid breaking chnges * remove async_op=True for redistribute * linting * renaming MoeExpertsParallel class * ep router doc * mlinter * add todo * fix * new name * guard import * cache inspect signature * merge init and from_device_mesh * add todo * avoid nesting * pp stages logs only ad INFO level * better doc func * linting * linting * fix --------- Co-authored-by: Arthur <48595927+ArthurZucker@users.noreply.github.com>

3outeille added 2 commits May 28, 2026 17:31

Revert "[Revert] FSDP+Dtensor refactor related changes (#46246)"

18a712b

This reverts commit 295cee3.

not unifying but still better than before

863e12b

3outeille changed the title ~~DIstributed branch~~ DIstributed branch base May 28, 2026

3outeille added 2 commits May 28, 2026 18:08

mellum now inherits tp and sp plan from qwen3_moe

f914e30

handle dense + sparse mixing for mellum model in SP plan

2d01eb3

3outeille added 2 commits June 16, 2026 02:40

Merge branch 'main' of https://github.com/huggingface/transformers in…

d672665

…to distributed

check-repository-consistency

fb95e26

3outeille force-pushed the distributed branch from 00922bb to fb95e26 Compare June 16, 2026 03:37

3outeille and others added 5 commits June 16, 2026 04:00

check-code-quality

61d3cbd

fix

c3d75af

simplify fsdp (#46679)

a540167

* simplify fsdp * add verify_fsdp_plan * linting * remove tie_weights() before distribute_model. Now use config.tie_word_embeddings for plan expectations, not pre-shard aliasing.

Merge branch 'main' into distributed

5845f3a

This was referenced Jun 17, 2026

Add native FSDP2 module and import migration (Phase A PR-2) 3outeille/transformers#1

Closed

Add native FSDP2 module and import migration #46706

Closed

3outeille added 2 commits June 23, 2026 14:25

Merge branch 'main' into distributed

08911d5

Merge branch 'main' into distributed

245a7f6

3outeille mentioned this pull request Jul 8, 2026

Add FSDP plans to all models #47165

Merged

AmineDiro mentioned this pull request Jul 30, 2026

fix: never scatter the uninitialized tail of the colvec, fix backward grad ds IlyasMoutawwakil/sonic-moe#1

Closed

3outeille closed this Sep 17, 2026

Navigation Menu

DIstributed branch base - #46269

DIstributed branch base#46269
3outeille wants to merge 13 commits into
mainfrom
distributed

3outeille commented May 28, 2026 •

edited

Loading

Uh oh!

HuggingFaceDocBuilderDev commented May 28, 2026

Uh oh!

AmineDiro commented May 29, 2026

Uh oh!

3outeille commented May 29, 2026

Uh oh!

github-actions Bot commented Jun 17, 2026

Uh oh!

github-actions Bot commented Jun 24, 2026

Uh oh!

AmineDiro commented Aug 3, 2026 •

edited

Loading

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

Conversation

3outeille commented May 28, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Current state of things

What's missing:

TODO

PR

Split PR

Uh oh!

HuggingFaceDocBuilderDev commented May 28, 2026

Uh oh!

AmineDiro commented May 29, 2026

Uh oh!

3outeille commented May 29, 2026

Uh oh!

github-actions Bot commented Jun 17, 2026

Uh oh!

github-actions Bot commented Jun 24, 2026

Uh oh!

AmineDiro commented Aug 3, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

3outeille commented May 28, 2026 •

edited

Loading

AmineDiro commented Aug 3, 2026 •

edited

Loading