DeepSpeed refactoring (#2313)
* DeepSpeed refactoring
Co-Authored-By: Stas Bekman <stas00@users.noreply.github.com>
* add tests
* Update test_deepspeed.py
* Update test_deepspeed.py
---------
Co-authored-by: Stas Bekman <stas00@users.noreply.github.com>
