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.

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

enhancements and fixes for FSDP and DeepSpeed #532

Merged

Conversation

pacman100
Copy link
Contributor

@pacman100 pacman100 commented Jul 19, 2022

What does this PR do?

  1. Enables StateDictType functionality of FSDP for checkpointing Full and Sharded model and optimizer states.
  2. Enabled DeepSpeed checkpointing functionality to save and load Full and Sharded model, optimizer and scheduler states.
  3. Fixes Issues with saving model/optimizer and loading them back #285 related to deepspeed and the optimizer issue of Don't unwrap in save_state() #489 when using deepspeed
  4. Renaming all FSDP arguments with fsdp prefix for clarity and readability
  5. Fixing accelerator.wait_for_everyone to account for FSDP

One can now use accelerator.save_state(output_dir) and accelerator.load_state(input_dir) without having to deal with complexities of DeepSpeed and FSDP.

ToDo

  • Write tests
  • Run experiments to verify
  • update the docs related to this PR

@HuggingFaceDocBuilderDev
Copy link

HuggingFaceDocBuilderDev commented Jul 19, 2022

The documentation is not available anymore as the PR was closed or merged.

Copy link
Collaborator

@sgugger sgugger left a comment

Choose a reason for hiding this comment

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

Thanks for your PR! I left a few comments, but mainly take care of not breaking backward compatibility by changing public functions.

src/accelerate/checkpointing.py Outdated Show resolved Hide resolved
src/accelerate/checkpointing.py Outdated Show resolved Hide resolved
src/accelerate/checkpointing.py Outdated Show resolved Hide resolved
src/accelerate/commands/config/cluster.py Show resolved Hide resolved
src/accelerate/commands/config/cluster.py Outdated Show resolved Hide resolved
src/accelerate/commands/launch.py Show resolved Hide resolved
src/accelerate/utils/dataclasses.py Outdated Show resolved Hide resolved
1. Adding deprecation args and warnings in launcher for FSDP
2. Handling old configs to work with new launcher args wrt FSDP.
3. Reverting changes to public methods in `checkpointing.py` and handling it in `Accelerator`
4. Explicitly writing the defaults of various FSDP options in `dataclasses` for readability.
1. FSDP wrapped model being added to the `_models`.
2. Not passing the env variables when args are None.
@pacman100 pacman100 changed the title checkpointing enhancements and fixes for FSDP and DeepSpeed enhancements and fixes for FSDP and DeepSpeed Jul 20, 2022
Copy link
Collaborator

@sgugger sgugger left a comment

Choose a reason for hiding this comment

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

Thanks for iterating!

1. Removes mrpc datafiles and directly relies on HF datasets as it was throwing `file not found` error when running from within `tests` folder. Updating `moke_dataloaders` as a result.
2. adding `test_performance.py`, `test_memory.py` and `test_checkpointing.py` for multi-gpu FSDP and DeepSpeed tests
@pacman100 pacman100 marked this pull request as ready for review July 22, 2022 15:14
@pacman100
Copy link
Contributor Author

make test_deepspeed results below (~6 minutes) :

Screenshot 2022-07-26 at 1 11 22 PM

make test_fsdp results below (~4 minutes):

Screenshot 2022-07-26 at 1 25 05 PM

Copy link
Collaborator

@sgugger sgugger left a comment

Choose a reason for hiding this comment

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

LGTM, thanks for working on this!

@pacman100 pacman100 merged commit 0c6bdc2 into huggingface:main Jul 26, 2022
@pacman100 pacman100 deleted the smangrul/saving-and-loading-utilities branch July 27, 2022 04:20
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

Issues with saving model/optimizer and loading them back
3 participants