[ZeroRedundancyOptimizer] Pytorch compliant state (#52960)
Summary:
Same as https://github.com/pytorch/pytorch/issues/52760 which I could not get to land. I just could not live with ghstack/ghimport/randomly broken things, I break enough of them myself, so this is a fresh copy without ghstack shenanigans. I'm hopeful that this can land relatively bug free, and am sorry for the duplications..
What this does:
- call the common_utils test runner instead of unittest, because it seems that it's how it should be done
- change the returned state from ZeroRedundancyOptimizer to be PyTorch compliant, which has the added benefit of being elastic (world size independent)
Pull Request resolved: https://github.com/pytorch/pytorch/pull/52960
Reviewed By: mrshenli
Differential Revision: D26710932
Pulled By: blefaudeux
fbshipit-source-id: 1d914bc9221442ba1bb2b48f5df10c313e674ece
