DeepSpeed

[engine] train should be able to get `mode` arg

#571

Merged