Skip to content
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

Add FSDP strategy #1553

Merged
merged 37 commits into from
Sep 29, 2022
Merged

Add FSDP strategy #1553

merged 37 commits into from
Sep 29, 2022

Conversation

abhi-mosaic
Copy link
Contributor

@abhi-mosaic abhi-mosaic commented Sep 22, 2022 •
edited
Loading

This PR adds alpha support for FSDP with Composer.

@abhi-mosaic abhi-mosaic requested a review from hanlint September 22, 2022 01:25
@abhi-mosaic abhi-mosaic self-assigned this Sep 22, 2022
@abhi-mosaic
Copy link
Contributor Author

Branch is currently failing in torch 1.11 due to FSDP imports. I think it is necessary to gate all FSDP features behind torch 1.12

hanlint
hanlint reviewed Sep 22, 2022
View reviewed changes
composer/trainer/strategy.py Outdated Show resolved Hide resolved
composer/trainer/trainer.py Outdated Show resolved Hide resolved
@hanlint hanlint requested a review from mvpatel2000 September 23, 2022 18:30
mvpatel2000
mvpatel2000 reviewed Sep 23, 2022
View reviewed changes
Copy link
Contributor

@mvpatel2000 mvpatel2000 left a comment

Choose a reason for hiding this comment

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

Mostly minor nits. The design seems reasonable to me. Can do testing and a deeper pass once its ready for final review

composer/core/state.py Outdated Show resolved Hide resolved
composer/core/state.py Outdated Show resolved Hide resolved
composer/core/state.py Outdated Show resolved Hide resolved
composer/core/state.py Outdated Show resolved Hide resolved
composer/trainer/strategy.py Outdated Show resolved Hide resolved
composer/trainer/trainer.py Show resolved Hide resolved
@hanlint hanlint marked this pull request as ready for review September 28, 2022 23:50
@hanlint hanlint requested a review from a team as a code owner September 28, 2022 23:50
mvpatel2000
mvpatel2000 reviewed Sep 29, 2022
View reviewed changes
composer/utils/misc.py Outdated Show resolved Hide resolved
mvpatel2000
mvpatel2000 reviewed Sep 29, 2022
View reviewed changes
Copy link
Contributor

@mvpatel2000 mvpatel2000 left a comment

Choose a reason for hiding this comment

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

minor style nits to be consistent in codebase. its a little bit of a pain lol im fine doing a followup cleanup PR if this takes a while.

No comments on functionality, will approve once tests pass!

composer/optim/__init__.py Show resolved Hide resolved
composer/trainer/dist_strategy.py Outdated Show resolved Hide resolved
composer/core/state.py Outdated Show resolved Hide resolved
composer/core/state.py Outdated Show resolved Hide resolved
composer/core/state.py Outdated Show resolved Hide resolved
composer/trainer/dist_strategy.py Outdated Show resolved Hide resolved
composer/trainer/dist_strategy.py Show resolved Hide resolved
composer/trainer/dist_strategy.py Outdated Show resolved Hide resolved
abhi-mosaic and others added 4 commits September 28, 2022 18:15
@abhi-mosaic @mvpatel2000
Update composer/core/state.py
64638c8 
Co-authored-by: Mihir Patel <mihir.v.patel7@gmail.com>
@abhi-mosaic @mvpatel2000
Update composer/core/state.py
59e95a7 
Co-authored-by: Mihir Patel <mihir.v.patel7@gmail.com>
@abhi-mosaic @mvpatel2000
Update composer/core/state.py
a6a66de 
Co-authored-by: Mihir Patel <mihir.v.patel7@gmail.com>
@abhi-mosaic @mvpatel2000
Update composer/trainer/dist_strategy.py
8118b1e 
Co-authored-by: Mihir Patel <mihir.v.patel7@gmail.com>
@mvpatel2000
Copy link
Contributor

Also Id just disable the linters in the files from pytorch, eg pydocstyle

@abhi-mosaic abhi-mosaic requested a review from a team as a code owner September 29, 2022 01:41
mvpatel2000
mvpatel2000 approved these changes Sep 29, 2022
View reviewed changes
Copy link
Contributor

@mvpatel2000 mvpatel2000 left a comment

Choose a reason for hiding this comment

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

Approving to unblock. Slack me for specific followups, or otherwise will be offline

composer/core/state.py Outdated Show resolved Hide resolved
composer/trainer/dist_strategy.py Show resolved Hide resolved
@abhi-mosaic abhi-mosaic merged commit aa19e24 into mosaicml:dev Sep 29, 2022
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Reviewers

@hanlint hanlint hanlint left review comments

@mvpatel2000 mvpatel2000 mvpatel2000 approved these changes

Assignees

@abhi-mosaic abhi-mosaic

Labels
None yet
Projects
None yet
Milestone
No milestone
Development

Successfully merging this pull request may close these issues.
3 participants
@abhi-mosaic @mvpatel2000 @hanlint