DeepSpeed

Fix local rank mismatch error when training on nodes with different number of GPUs

#3409

Merged