rollback #6726 (#7258)
This PR rollback #6726 which caused
https://github.com/deepspeedai/DeepSpeed/issues/7116 .
---------
Signed-off-by: Guokai Ma <guokai.ma@gmail.com>
Co-authored-by: Logan Adams <114770087+loadams@users.noreply.github.com>
Co-authored-by: Olatunji Ruwase <tunji.ruwase@snowflake.com>
