Initial FSDP2 support #3394
Feat: initial conversion tool draft
2cfe2db7
Feat: add value mapping to conversion tool
d055eab2
Refactor: move from os to pathlib
23ea666f
Feat: add first tests
7deaaaad
Feat: more tests
d506b951
Feat: minor fixes + dataclass conversions
3eab2262
Feat: more remapping
269c2173
Fix: namespace has no attribute version + style
ba5372cc
Fix: offload params behavior
b16131bd
Feat: add option to only rename keys in the config file to
b0cc66b3
Fix: wrong attr name
8bc5cb34
Merge branch 'main' into dev/fsdp2
c31cc550
Fix: partially resolve comments
bc70ec2b
Feat: work on config command + minor fixes to reflect changes
00dafc45
Refactor: style + quality
7a92dac0
Feat: fsdp2 initial work
b724f9cd
Feat: some cleanups and first running fsdp2
7cd25879
Fix: version checks + mixed precision policy
d920e94c
Refactor: style + quality
3cc7c205
Remove obsolete todos
432b4ff0
Feat: grad norm clipping
bb839851
Fix: tests + rename attrs
7759d799
Refactor: style + quality
f8bfa04d
Fix: None object is not iterable
47a12b97
Fix: default cpu_offload for fsdp2
030fa5df
Fix: cpu offload now behaves correctly
c5526aaa
Feat: apply_activation_checkpointing
4984c8ad
Fix: append to models
914e55f5
Feat: start on concept guide
1304284c
Merge branch 'main' into dev/fsdp2
f9061f6d
wip: concept guide
c596c127
Fix: toctree
d55bab7a
cleanup of the concept guide
8494f7bc
Fix: minor fixes + mp
f298b5eb
Fix: quality + | to union
4e4b6643
Feat: backwards compatibility + args cleanup
30961c91
Fix: style + quality
314df772
Feat: enable dropping refs when getting named params
5238cb27
Merge branch 'main' into dev/fsdp2
e104c100
Fix: memory footprint with fsdp2
8347b917
Feat: cpu ram efficient loading
70f48825
Fix: mp
c5975983
Fix: not warn about sync_modules if fsdp version is 1
fa732959
Refactor: minor changes
ee1267a1
Small fixes + refactors
021e153e
Feat: docs + cleanup
04a5ffbd
Feat: saving works (not sure about optim)
d9fba399
More loading/saving work
0696c961
Feat: disable local_state_dict for fsdp2
0d0abde7
Fix: fsdp2 convergence
1fe97f62
Merge branch 'main' into dev/fsdp2
ab4e58dd
Feat: working comparison script
d6edd8cd
Feat: memory tracking fsdp2
8a480385
Feat: memory visualizer
a563493b
Feat: more work on benchmark
a8d127bd
Fix: raise error if model+optimizer arent prepared together
55605558
Minor fixes
ab5a6cf6
Style
31ce351a
More warnings
b619589c
Fix: reshard_after_forward vs sharding_strategy conflict
4dab060d
Refactor: clean up accelerator
c7ef53da
Feat: more testing in fsdp2 benchmark
8fadfbf9
Fix: memory visualizer
79d5cf71
Untested: support load/save_state
d633c125
Feat: concept guide improvements
8a4de4d4
Refactor: concept guide
4881bbf3
Feat: benchmark works
9fd8fd51
Feat: more work on fsdp2 benchmark
0d53cadd
Fix: note syntax
9251dcd6
Fix: small fixes + make original tests work
25f15c78
Fix: grad scaling
4dfb12c2
muellerzr
approved these changes
on 2025-03-25
Feat: reshard after forward tests
d0244e2c
Feat: backward prefetch tests
d2891a2d
Feat: tests for fsdp2
f99e980d
Refactor: minor fixes
5619feaa
Feat: fsdp_utils docstrings
3fba7e21
Feat: autodoc fsdp.md
6df64832
Docs: get_module_children_bottom_up
d24aa9d4
Fix: remove unused images
050488d1
Refactor: benchmark cleanup
e9845716
Fix: docs
ccf3d426
Feat: final doc changes
124a56fd
Fix: torch.distributed has no attribute tensor
5bb4f306
Fix: style
0f71c8c6
Feat: tests include version in failures
0a2791ba
S1ro1
changed the title WIP: Initial FSDP2 support Initial FSDP2 support 1 year ago
S1ro1
marked this pull request as ready for review 1 year ago
Fix: benchmark force model to load in fp32
9f37abd6
Fix: rename runs
a7333d1b
Feat: last minor fixes
4d4f829b
Feat: new benchmark images
7ad58a15
muellerzr
merged
d7c741a6
into main 1 year ago
Assignees
No one assigned
Login to write a write a comment.
Login via GitHub