🚨 TP dtensor API inference + training - #47579
Merged
Merged
Conversation
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.
…uggingface/transformers into split/a-pr-3-dual-path-loading
naming Co-authored-by: Arthur <48595927+ArthurZucker@users.noreply.github.com>
Member
Author
|
run-slow: doge, llama4, fbgemm_fp8, finegrained_fp8 |
Contributor
|
This comment contains models: ["models/doge", "models/llama4"] |
Contributor
CI ResultsCommit Info
The test failure analysis could not be completed. Please check the workflow run for details. |
ArthurZucker
approved these changes
Aug 21, 2026
ArthurZucker
left a comment
ArthurZucker
left a comment
Collaborator
There was a problem hiding this comment.
Happy to merge as long as the mapping between bc names is there! let's make this as non breaking as possible!
| self.use_local_output = use_local_output | ||
|
|
||
| def should_use_local_tensors(self, module): | ||
| use_local_quantized_path = getattr(module, "_hf_quantized_needs_local_tp", False) |
Collaborator
There was a problem hiding this comment.
yeah, but training is opt in / not with quantization
| num_local_experts = num_experts // ep_size | ||
|
|
||
| router_logits, router_scores, router_indices = output | ||
| router_logits, router_scores, router_indices, *extra_outputs = output |
Collaborator
There was a problem hiding this comment.
are there actually any?
Member
Author
There was a problem hiding this comment.
yeah InklingTopKRouter returns 4 parameters return routed_logits, topk_weights, topk_indices, shared_gammas
Contributor
|
[For maintainers] Suggested jobs to run (before merge) run-slow: doge, hunyuan_vl, llama4, fbgemm_fp8, finegrained_fp8 |
Contributor
CI recapDashboard: View test results in Grafana |
pull Bot
pushed a commit
to AmirulAndalib/DeepSpeed
that referenced
this pull request
Aug 23, 2026
…eepspeedai#8294) ## Problem huggingface/transformers#47579 (on `main` since `861f4c41`, 2026-08-21) makes `PretrainedConfig.__init__` inject `"embed_tokens": "embedding_rowwise"` into `base_model_tp_plan` whenever `tie_word_embeddings` is true. `SUPPORTED_STYLES` is a strict allowlist and `convert()` raises on any style outside it, rejecting the whole plan, so AutoTP plan conversion now fails for every tied-embedding model (Qwen2/Qwen3, Llama, Gemma) built against transformers `main`. ## Fix Recognize `embedding_rowwise` and convert it to a `SKIP` spec, which is option 2 in deepspeedai#8290: the entry is understood and the embedding is deliberately left replicated. What follows is the behaviour DeepSpeed already implements rather than a new policy. `lm_head` still converts to a gathered column spec, and `_configure_gathered_column_tie_fallbacks` then sees that `lm_head.weight is embed_tokens.weight` and leaves both modules replicated, logging that coupled vocabulary-parallel embedding is not supported yet. A tied model is therefore left in the shape it has on a transformers release without the injection, with both modules replicated and the tie intact. The entry maps to `SKIP` with `grad_allreduce` left false, unlike `replicated_with_grad_allreduce`. The parameter is never split, so `register_replicated_grad_hooks` must not register an all-reduce for it. Styles that are still unknown continue to reject the whole plan. `test_unsupported_style_rejects_whole_plan` is unchanged and still passes. ## Verification Run on CPU in a container at `edaa7221`, against transformers `main` (5.16.0.dev0) and torch 2.13.0+cpu. - The two added tests fail on master with the reported `ValueError` and pass with this change. - `tests/unit/module_inject/` and `tests/unit/runtime/test_tp_plan_extraction.py`: 47 passed, on Python 3.11 and on 3.12. - `pre-commit run --files` on the three changed files passes yapf, check-torchdist, check-license and codespell; flake8 5.0.4 exits 0 on them under Python 3.11. - Not verified here: `test_qwen2_tied_lm_head_falls_back_to_replicated`, which needs 2 GPUs. That is the test deepspeedai#8290 reports as failing and the one this change is meant to restore. ## Two things worth deciding separately Scoping this to `embedding_rowwise` leaves the next transformers-side style to fail the same way, since the injection is unconditional and the allowlist is deny-by-default against a vocabulary DeepSpeed does not own. A general rule for unknown styles looks like a maintainer call rather than something to settle here. Related to that, the `convert()` docstring says entries with an unsupported style become SKIP specs instead of invalidating the plan, but no code path does that, and none does after this change either: an unsupported style still raises before the loop is reached. The docstring and the raise arrived together in deepspeedai#8204, so I have left both alone. Happy to follow up once you have picked the policy. Refs deepspeedai#8290. This covers the conversion failure only, and does not implement vocabulary-parallel tied embeddings (option 1 or 3 in that issue), so I have not used a closing keyword. --------- Signed-off-by: Ehsan Barkhordar <realbarkhordar@gmail.com> Co-authored-by: Ma, Guokai <guokai.ma@gmail.com>
SunMarc
added a commit
that referenced
this pull request
Aug 26, 2026
…rides `_impl_tp_layer_overrides` only declares `deepgemm_megamoe`, so `.get(impl)` returns `None` for every other experts implementation (including the default `None`). #47579 dropped the `if layer_overrides:` guard around the plan rewrite in order to run the new `_scale_inv` sharding loop unconditionally, which left `layer_overrides.get(v, v)` dereferencing `None` for any model whose `base_model_tp_plan` is non-empty -- i.e. all FP8 quantization of e.g. Llama: AttributeError: 'NoneType' object has no attribute 'get' Default to `{}` so the rewrite is a no-op while the `_scale_inv` loop still runs.
SunMarc
added a commit
that referenced
this pull request
Aug 26, 2026
…rides `_impl_tp_layer_overrides` only declares `deepgemm_megamoe`, so `.get(impl)` returns `None` for every other experts implementation (including the default `None`). #47579 dropped the `if layer_overrides:` guard around the plan rewrite in order to run the new `_scale_inv` sharding loop unconditionally, which left `layer_overrides.get(v, v)` dereferencing `None` for any model whose `base_model_tp_plan` is non-empty -- i.e. all FP8 quantization of e.g. Llama: AttributeError: 'NoneType' object has no attribute 'get' Default to `{}` so the rewrite is a no-op while the `_scale_inv` loop still runs.
37 tasks
BenjaminBossan
pushed a commit
to huggingface/peft
that referenced
this pull request
Sep 16, 2026
Make the TP integration in PEFT work with the new Transformers approach using DTensors: huggingface/transformers#47579 The legacy TP integration is still supported.
sbucaille
pushed a commit
to sbucaille/transformers
that referenced
this pull request
Sep 16, 2026
* 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>
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.
Summary
This PR replaces the legacy Transformers tensor-parallel implementation with a DTensor-native backend while preserving efficient inference for kernels that operate on local tensors + loading time. We outsource all the placement and parameters sharding logic to Dtensor which will enable us to combine parallelism easily in the future.
We made sure the TP dtensor api is compatible with kernels (deepgem, sonicmoe) and quantization (mxfp4, fp8, torchao) and other libraries (vllm, peft in progress)
Performance
Benchmark using 2 GPUS with batch size 1 and 256 generated tokens shows performance and peak memory similar to the legacy TP.
Benchmark of model loading time between both APIs