Add FSDP plans to all models - #47165
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>
Comment on lines
618
to
+619
|
|
||
| _fsdp_plan = {"lm_head": "keep_full_weight"} |
Collaborator
There was a problem hiding this comment.
Suggested change
| _fsdp_plan = {"lm_head": "keep_full_weight"} | |
| _fsdp_plan = {"lm_head": "keep_full_weight"} |
Comment on lines
+63
to
+67
| base_model_fsdp_plan = { | ||
| "embed_tokens": "free_full_weight", | ||
| "layers.*": "free_full_weight", | ||
| "norm": "keep_full_weight", | ||
| } |
Collaborator
There was a problem hiding this comment.
if literally all models have the same default, its a good default -> parent distributed config can use it imo its fairly safe.
ArthurZucker
approved these changes
Jul 27, 2026
ArthurZucker
left a comment
Collaborator
There was a problem hiding this comment.
just some extra line here and there
3outeille
enabled auto-merge
July 27, 2026 12:20
…ace/transformers into split/a-pr-5-fsdp-plans
Contributor
|
[For maintainers] Suggested jobs to run (before merge) run-slow: afmoe, apertus, arcee, aria, axk1, axk2, bamba, bitnet, cohere, cohere2, cohere2_moe, csm, cwm, dbrx, deepseek_v2, deepseek_v3 |
Contributor
CI recapDashboard: View test results in Grafana |
stevhliu
pushed a commit
to stevhliu/transformers
that referenced
this pull request
Jul 30, 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 * 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>
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
base_model_fsdp_planand_fsdp_plan = {"lm_head": "keep_full_weight"}to 94 ForCausalLM classes, extracted from PR #46269 (origin/distributed).FSDP_DISTRIBUTED_TEST_MODEL_TYPESintests/test_fsdp_mixin.pyto run distributed FSDP tests on subset:cohere2_moe,llama,mistral,mixtral,qwen3,qwen3_moe