Skip to content

[FSDP2] Init FSDP2 based checkpointing #3824

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom
Merged
merged 29 commits into from
Apr 17, 2025
Merged

[FSDP2] Init FSDP2 based checkpointing #3824

merged 29 commits into from
Apr 17, 2025

Conversation

bowenyang008
Copy link
Contributor

@bowenyang008 bowenyang008 commented Apr 17, 2025 •
edited
Loading

What does this PR do?

Kick start support of FSDP2 checkpointing along with Trainer and State:

  1. Support same interface in State.fsdp_config between FSDP2Config and FSDPConfig so both config can be passed to Trainer/State. Currently FSDP2Config supports a minimum default interface methods/properties to make State functional
  2. Support FSDP2 based checkpointing by State
  3. Support loading checkpoints from FSDP1 artifacts to FSDP2

Tests

  1. Two unit tests to test FSDP2 checkpoint saving/loading and FSDP1 saving -> FSDP2 loading
  2. One unit test to test GradScaler

bowenyang008 and others added 29 commits April 11, 2025 00:17
@bowenyang008
simplify fsdp2 config to reduce version complexity first
1a59cc0
@bowenyang008
update fsdp2 config and test
f6a47ec
@bowenyang008
update fields
7878678
@bowenyang008
guardrails against FSDP1 and grad scaler handling
12a9d8e
@bowenyang008
more property
722221b
@bowenyang008
update test
c2dc262
@bowenyang008
assert equal
5911217
@bowenyang008
refactor test
fb3bd0d
@bowenyang008
dummy func
6264587
@bowenyang008
too many weird reference in root module
578844b
@bowenyang008
test pass
0cacc1b
@bowenyang008
simplify
bb38ef6
@dakinggg @bowenyang008
Remove some more calls to HF CI (#3817)
5fa8288
@dependabot @bowenyang008
Update psutil requirement from <7,>=5.8.0 to >=5.8.0,<8 (#3818)
ef407cb 
Signed-off-by: dependabot[bot] <[email protected]>
Co-authored-by: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com>
@bowenyang008
Revert "Update psutil requirement from <7,>=5.8.0 to >=5.8.0,<8 (#3818)"
e8687f8 
This reverts commit ef407cb.
@bowenyang008
Revert "Remove some more calls to HF CI (#3817)"
e667fd1 
This reverts commit 5fa8288.
@bowenyang008
Merge remote-tracking branch 'origin/main' into boweny/fsdp2/state-ch…
1f62617 
…eckpoint
@bowenyang008
cleanup
13ae07a
@bowenyang008
centralize test context
976c6a0
@bowenyang008
new test
230b3c4
@bowenyang008
format
f052cfb
@bowenyang008
format again
a89b731
@bowenyang008
fix import
4ff48df
@bowenyang008
more
bcd9bcf
@bowenyang008
type fix
7ebcd90
@bowenyang008
doc
792e306
@bowenyang008
format
a34f835
@bowenyang008
dummy change
f06707e
@bowenyang008
avoid destroy?
6803f71
@bowenyang008 bowenyang008 changed the title Boweny/fsdp2/state checkpoint [FSDP2] Init FSDP2 based checkpointing Apr 17, 2025
@bowenyang008 bowenyang008 marked this pull request as ready for review April 17, 2025 06:32
dakinggg
dakinggg approved these changes Apr 17, 2025
View reviewed changes
Copy link
Contributor

@dakinggg dakinggg left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM!

@bowenyang008 bowenyang008 merged commit 03ac462 into main Apr 17, 2025
14 checks passed
@bowenyang008 bowenyang008 deleted the boweny/fsdp2/state-checkpoint branch April 17, 2025 17:39
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Reviewers

@rithwik-db rithwik-db Awaiting requested review from rithwik-db

@ethantang-db ethantang-db Awaiting requested review from ethantang-db

1 more reviewer

@dakinggg dakinggg dakinggg approved these changes

Reviewers whose approvals may not affect merge requirements
Assignees
No one assigned
Labels
None yet
Projects
None yet
Milestone
No milestone
Development

Successfully merging this pull request may close these issues.
2 participants
@bowenyang008 @dakinggg