undo the revert (#7536)

Commit

117 days ago

undo the revert (#7536) replay https://github.com/deepspeedai/DeepSpeed/pull/3019 as it got reverted