DeepSpeed feat: Add support for `NamedTuple` when sharding parameters [#3029]

As discussed in https://github.com/microsoft/DeepSpeed/issues/3029 this MR adds support for NamedTuple outputs when parameters are sharded.

I also took the liberty to add extra documentation and extend the tests to test previously untested outputs (List/Tuple). Please do not hesitate to let me know where I can improve this PR 🙇

Mar 16 '23 11:03 alexandervaneck

@microsoft-github-policy-service agree

Mar 16 '23 11:03 alexandervaneck

@AlexanderVanEck, thanks for this amazing contribution!

Mar 16 '23 15:03 tjruwase

Thank you @tjruwase 🙇

I see that all of the tests are failing, but for reasons unrelated to the PR. I am happy to rebase this once master has been updated to allow the pytorch2.0 release 🎉

Mar 16 '23 15:03 alexandervaneck

@AlexanderVanEck, yes, sorry for the failing tests. They all seem to be due to py 2.0 issues in other libraries. I will push a temporary fix that reverts CI to py1.13*.

Mar 16 '23 16:03 tjruwase

Thanks @tjruwase . I wasn't sure if you need collab permissions on my fork, but just in case I sent it to you! 🙏

Also; no need for apologising about the failing tests. Red builds are part of the job 🤗

Mar 17 '23 08:03 alexandervaneck

@AlexanderVanEck, thanks for sending collab invitation, but I don't actually need it since the CI issues need to be resolved in master branch. A teammate has taken over that responsibility. I don't think any action is needed on your side from this point on.

Apologies for the delay while we resolve our CI issues. I will keep an eye on things to ensure your PR is merged asap.

Mar 17 '23 12:03 tjruwase

Understood, thanks for updating me. 🙇 I look forward to seeing this pass.

👋

Mar 17 '23 12:03 alexandervaneck