FSDP x TP Saving / Loading (#48802)
* wip: add script to track potential changes
* feat: workaround for _StridedShard
* feat: handle optimizer loading / saving
* wip: add clip_grad_norm_
* doc: add docstring to clip_grad_norm_
* feat: uniformize the clip_grad_norm_ implem
* refactor: use the proper function for clip_grad_norm_
* feat: save shards / consolidate utility functions
* feat: save / load sharded / consolidated optimizer states
* feat: add checkpoint state check
* feat: add distributed checkpoint loading
* test: add saving, loading and clipping tests
* wip: remove .bin support
* doc: update docstring
* chore: remove testing script
* align: match planner branch loader surface
* fix: remove method that was directly calling a util
* doc: rename word in docstring
* wip: add guards in utility functions
* doc: add docstring explaining why `_StridedShard` is broken
* wip: main function to load checkpoints in DistributedMixin
* test: add test for loading paths
* wip: add strict for sharded checkpoints as well
* doc: added tree_map comments
* doc: explain why options are required
* refactor: move functions to distributed/checkpoint.py
* feat: remove clip_grad_norm from current PR
* fix: remove consolidate option
* wip: refactor distributed tests
* fix: tiny fixes
* feat: add loading path in from_pretrained
* tests: tests for checkpointing
* chore: restore fsdp mixin
* fix: broken imports in tests
* fix: remove the support for full optimizer loading
* fix: remove consolidation support for optimizer saving
* refactor: rename to load_model_checkpoint_distributed
* refactor: group two if statements into a for loop
* refactor: separate the logic of finding the files and loading
* refactor: restore optimizer changes for stacking
* style: ruff
* refactor: use custom error message for guard function
* fix: use is_torch_distributed_available
* feat: consolidate checkpoint feature
* doc: add context for DTensor hooks
* refactor: add get_distributed_checkpoint_dir helper
* refactor: _load_and_finalize_distributed_pretrained_model in mixin
* refactor: make helper private
* style: ruff
* doc: improve docstring
* style: ruff
* feat: remove loading paths that cannot be reached
* refactor: refactor save_distributed_checkpoint
* feat: remove consolidated dir if it exists
* feat: implement reverse logic for `_prepare_state_dict_for_dcp`
* refactor: rename helper function
* chore: add parameter name
* docs: add comment
* refactor: move distributed utils to distributed/
* fix: broken test because of refactor
---------
Co-authored-by: Ferdinand Mom <47445085+3outeille@users.noreply.github.com>
