FSDP x TP Saving / Loading #48802
michaelbenayoun
merged
71 commits into
huggingface:main
from
michaelbenayoun:fsdp_tp_saving
wip: add script to track potential changes
4454e4be
feat: workaround for _StridedShard
ec979c45
feat: handle optimizer loading / saving
6d9a2399
wip: add clip_grad_norm_
6868815b
doc: add docstring to clip_grad_norm_
f20be8d4
Merge branch 'main' into fsdp_tp_saving
fec373e8
michaelbenayoun
commented
on 2026-09-14
feat: uniformize the clip_grad_norm_ implem
f9341a88
refactor: use the proper function for clip_grad_norm_
b09c7eec
feat: save shards / consolidate utility functions
9f677ef2
feat: save / load sharded / consolidated optimizer states
9c742be7
feat: add checkpoint state check
dc192d93
Merge branch 'main' into fsdp_tp_saving
24558d8b
feat: add distributed checkpoint loading
b2719522
test: add saving, loading and clipping tests
6a5075c0
wip: remove .bin support
78d55a4a
doc: update docstring
607d7f9b
chore: remove testing script
55f0814b
Merge branch 'main' into fsdp_tp_saving
b5890fe8
align: match planner branch loader surface
6724f5ae
fix: remove method that was directly calling a util
2818860f
doc: rename word in docstring
61896f01
wip: add guards in utility functions
b58cdbb9
doc: add docstring explaining why `_StridedShard` is broken
7605b920
wip: main function to load checkpoints in DistributedMixin
c99185c3
test: add test for loading paths
1d012c12
wip: add strict for sharded checkpoints as well
afe8bcb8
doc: added tree_map comments
c1329666
doc: explain why options are required
e1e8048f
Merge branch 'main' into fsdp_tp_saving
f5d00940
3outeille
commented
on 2026-09-25
3outeille
commented
on 2026-09-25
Merge branch 'main' into fsdp_tp_saving
1ce39825
refactor: move functions to distributed/checkpoint.py
7384f704
feat: remove clip_grad_norm from current PR
2b6e7116
fix: remove consolidate option
2e81400f
wip: refactor distributed tests
9ba6352a
fix: tiny fixes
c9751154
feat: add loading path in from_pretrained
3d5932d8
tests: tests for checkpointing
151017d9
michaelbenayoun
commented
on 2026-09-29
chore: restore fsdp mixin
676e1b2b
fix: broken imports in tests
7b889135
3outeille
commented
on 2026-09-29
3outeille
commented
on 2026-09-29
3outeille
commented
on 2026-09-29
3outeille
commented
on 2026-09-29
3outeille
commented
on 2026-09-29
fix: remove the support for full optimizer loading
15913b5c
fix: remove consolidation support for optimizer saving
86827426
refactor: rename to load_model_checkpoint_distributed
14cfa9a8
refactor: group two if statements into a for loop
81f34cd8
refactor: separate the logic of finding the files and loading
332b3b3c
refactor: restore optimizer changes for stacking
e4b7ba0e
style: ruff
e8bdcde0
Merge branch 'main' into fsdp_tp_saving
0a46816a
3outeille
commented
on 2026-10-06
3outeille
commented
on 2026-10-06
3outeille
commented
on 2026-10-06
refactor: use custom error message for guard function
d843a519
fix: use is_torch_distributed_available
dec7bdad
feat: consolidate checkpoint feature
3cee2059
doc: add context for DTensor hooks
7d92bbd2
refactor: add get_distributed_checkpoint_dir helper
b5121aee
refactor: _load_and_finalize_distributed_pretrained_model in mixin
07874468
refactor: make helper private
e6c8604f
Merge branch 'main' into fsdp_tp_saving
f14062a2
style: ruff
ba38c786
3outeille
commented
on 2026-10-09
doc: improve docstring
4bee2469
Merge branch 'main' into fsdp_tp_saving
0ea3898f
style: ruff
5707d839
feat: remove loading paths that cannot be reached
9dbd995f
refactor: refactor save_distributed_checkpoint
055acb28
Merge branch 'main' into fsdp_tp_saving
b1aedd81
3outeille
commented
on 2026-10-14
3outeille
commented
on 2026-10-14
3outeille
commented
on 2026-10-14
3outeille
commented
on 2026-10-14
feat: remove consolidated dir if it exists
42e11bc3
feat: implement reverse logic for `_prepare_state_dict_for_dcp`
5bb6ab52
3outeille
commented
on 2026-10-14
refactor: rename helper function
500998e1
3outeille
commented
on 2026-10-15
3outeille
approved these changes
on 2026-10-15
chore: add parameter name
837ae3bd
Merge branch 'main' into fsdp_tp_saving
d4a85ad5
Merge branch 'main' into fsdp_tp_saving
dea96077
docs: add comment
25eb71e3
refactor: move distributed utils to distributed/
321f5388
fix: broken test because of refactor
ec30f1d1
Login to write a write a comment.
Login via GitHub