refactor deepspeed dataloader prepare logic (#2238)
* refactor deepspeed dataloader prepare logic
Co-Authored-By: Stas Bekman <stas00@users.noreply.github.com>
* address comments and fix issues
Co-Authored-By: Stas Bekman <stas00@users.noreply.github.com>
* further refactor
* add test
* rename test
---------
Co-authored-by: Stas Bekman <stas00@users.noreply.github.com>
