Megatron-DeepSpeed

497aa1bf - Sorry, last change was meant to a PR. This reverts commit d0fcf4170def7205426117016d4622c745f33883.