[Feat] DeepSpeed single file saving

[SeanNaren]

What does this PR do?

DeepSpeed has released a new version which contains all gathering weights to a single process and saving. As discussed in #6691 it's probably best for now to prioritise one file saving as this is easiest to deploy etc. As a result I've made this the default behaviour when using ZeRO 3.

In addition, if a user uses the configure_sharded_model hook, it is necessary to call this in the load_from_checkpoint function as the test indicates via the model hook. This could potentially be done for the user, however right now I'll make it clear the docs in a separate PR that this needs to be handled manually.

cc @ananthsub @tchaton

Before submitting

• Was this discussed/approved via a GitHub issue? (not for typos and docs)
• Did you read the contributor guideline, Pull Request section?
• Did you make sure your PR does only one thing, instead of bundling different changes together?
• Did you make sure to update the documentation with your changes? (if necessary)
• Did you write any new necessary tests? (not for typos and docs)
• Did you verify new and existing tests pass locally with your changes?
• Did you update the CHANGELOG? (not for typos, docs, test updates, or internal minor changes/refactorings)

PR review

Anyone in the community is free to review the PR once the tests have passed.
Before you start reviewing make sure you have read Review guidelines. In short, see the following bullet-list:

• Is this pull request ready for review? (if not, please submit in draft mode)
• Check that all items from Before submitting are resolved
• Make sure the title is self-explanatory and the description concisely explains the PR
• Add labels and milestones (and optionally projects) to the PR so it can be classified

Did you have fun?

Make sure you had fun coding 🙃

[tchaton]

Overall LGTM ! Small nits

[codecov]

Codecov Report

Merging #6900 (dbff37d) into master (e9c3e02) will decrease coverage by 5%.
The diff coverage is 0%.

@@           Coverage Diff           @@
##           master   #6900    +/-   ##
=======================================
- Coverage      92%     87%    -5%     
=======================================
  Files         194     194            
  Lines       12346   12355     +9     
=======================================
- Hits        11327   10688   -639     
- Misses       1019    1667   +648     

[kaushikb11]

Hi @SeanNaren, this particular test is failing.

[SeanNaren]

I'm also noticing the container is failing when DeepSpeed is added to the extras. I might move to having DeepSpeed installed within the GPU container instead of the extras file.

[ananthsub]

@SeanNaren totally separate from this PR, but have you seen checkpoint saving become a bottleneck at all with deepspeed? I think there's room to optimize the saving logic inside the lightning checkpoint callback, since we are potentially re-running the broadcast/state dict modification 2 times here: https://github.com/PyTorchLightning/pytorch-lightning/blob/fe0d08899eba94d275ff42253f495d9e70d86f89/pytorch_lightning/callbacks/model_checkpoint.py#L278-L286

[SeanNaren]

pytorch-lightning/pytorch_lightning/callbacks/model_checkpoint.py

If multiple of these modes are used, then definitely it would be an expensive operation. Currently I've only tested with save_last to ensure the final model can be saved. There also seems to be a limit in terms of the the all gather required to save, but this might be from Lightning not de-allocating memory (need to check after this PR is merged).

Rudimentary numbers I collected for 8 A100s using minGPT and using time probing when saving the model to disk

13B parameters with ZeRO 3 + offloading
Time taken to consolidate 10.642701387405396s (picked one of the 8 values logged from the GPUs, all were about the same)
Time taken to dump 95.15039467811584s

20B parameters with ZeRO 3 + offloading
Blows up CUDA OOM

20B parameters with ZeRO 3 + offloading + deepspeed checkpointing
Blows up CUDA OOM