[DeepSpeed] Remove partitioning of model in ZeRO 3

[SeanNaren]

What does this PR do?

Partially Fixes #10510

Also adds parametrization to a few tests since it's supported, and handles deprecated args with the latest DeepSpeed.

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 list all the breaking changes introduced by this pull request?
• Did you update the CHANGELOG? (not for typos, docs, test updates, or internal minor changes/refactorings)

PR review

Anyone in the community is welcome to review the PR.
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 🙃

cc @Borda @justusschock @awaelchli @akihironitta @SeanNaren

[SeanNaren]

Seems this is not so simple. If you load the model in inference mode you cannot save (which makes sense as it's the optimizers responsibility to save the ZeRO sharded weights).

Do we think that the case where the user does not call fit before saving should be supported?

trianer = Trainer(...)
trainer.test(model)
trainer.save_checkpoint(...)

If so I think it's better we keep the optimizer code as is and use the configure_optimizers to define the optimizer again.

cc @carmocca @awaelchli @rohitgr7

[rohitgr7]

Seems this is not so simple. If you load the model in inference mode you cannot save (which makes sense as it's the optimizers responsibility to save the ZeRO sharded weights).

Do we think that the case where the user does not call fit before saving should be supported?

trianer = Trainer(...)
trainer.test(model)
trainer.save_checkpoint(...)

If so I think it's better we keep the optimizer code as is and use the configure_optimizers to define the optimizer again.

cc @carmocca @awaelchli @rohitgr7

Even if we do that, we do not force users to define configure_optimizers during inference so possibility of not getting an optimizer still exists. So there are 2 cases now:

1. User calls .fit and then calls .test, which will work right? Since optimizers will be attached during .fit call.
2. User calls .test directly, then I don't think they will do a save_checkpoint because it it's a completely new randomly initialised model then checkpointing is not necessary and if they are using some pretrained model for testing then they will already have the saved checkpoint and calling another won't make sense.

Maybe you can add an error/warning for the case with no optimizer and saving under deepspeed config.

[tchaton]

LGTM !

[SeanNaren]

Pulled out the other issue, will make a separate PR for the deepspeed inference/optimizer fix.

[SeanNaren]

This PR will partially fix the related issue, but there is still the case of partially defined partitioned weights to tackle.