Add FSDP plans to all models (#47165)
* 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 #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>
