Repository: microsoft/DeepSpeed

Issue:
[BUG] Local rank mismatch error when training on nodes with different number of GPUs
**Describe the bug**
When training Megatron-DeepSpeed models on nodes with different number of GPUs (e.g., node 0: 2 GPUs, node 1: 4 GPUs), local rank mismatch error (`AssertionError: expected local-rank to be the same as rank % device-count.`) occurs in https://github.com/microsoft/Megatron-DeepSpeed/blob/main/megatron/initialize.py#L207. 

I found that this is caused by local ranks set by DeepSpeed launcher **mismatching** the local ranks set by Megatron-DeepSpeed.
Specifically, while DeepSpeed assigns local ranks corresponding to global ranks sequentially (e.g., global ranks: [2, 3, 4, 5] -> local ranks: [0, 1, 2, 3]), Megatron-DeepSpeed decides each local rank as **global rank % device count of the node** (e.g., global ranks: [2, 3, 4, 5] -> local ranks: [2, 3, 0, 1]). 

**To Reproduce**
1. Clone https://github.com/microsoft/Megatron-DeepSpeed and prepare to run training scripts (e.g., https://github.com/microsoft/Megatron-DeepSpeed/blob/main/examples/MoE/ds_pretrain_gpt_350M_dense.sh) as guided.
2. Prepare a cluster where each node has different number of GPUs (e.g., two nodes each with 2 and 4 GPUs, respectively).
3. Execute a training script.

**Expected behavior**
Training starts normally without an error in such setup. Because DeepSpeed document does not restrict that all nodes must be homogeneous (i.e., exacltly the same number of GPUs), I expected training to work properly for clusters where nodes have different number of GPUs.

**ds_report output**
```bash
--------------------------------------------------
DeepSpeed C++/CUDA extension op report
--------------------------------------------------
NOTE: Ops not installed will be just-in-time (JIT) compiled at
      runtime if needed. Op compatibility means that your system
      meet the required dependencies to JIT install the op.
--------------------------------------------------
JIT compiled ops requires ninja
ninja .................. [OKAY]
--------------------------------------------------
op name ................ installed .. compatible
--------------------------------------------------
 [WARNING]  async_io requires the dev libaio .so object and headers but these were not found.
 [WARNING]  async_io: please install the libaio-dev package with apt
 [WARNING]  If libaio is already installed (perhaps from source), try setting the CFLAGS and LDFLAGS environment variables to where it can be found.
async_io ............... [NO] ....... [NO]
cpu_adagrad ............ [NO] ....... [OKAY]
cpu_adam ............... [NO] ....... [OKAY]
fused_adam ............. [NO] ....... [OKAY]
fused_lamb ............. [NO] ....... [OKAY]
quantizer .............. [NO] ....... [OKAY]
random_ltd ............. [NO] ....... [OKAY]
 [WARNING]  using untested triton version (2.0.0), only 1.0.0 is known to be compatible
sparse_attn ............ [NO] ....... [NO]
spatial_inference ...... [NO] ....... [OKAY]
transformer ............ [NO] ....... [OKAY]
stochastic_transformer . [NO] ....... [OKAY]
transformer_inference .. [NO] ....... [OKAY]
utils .................. [NO] ....... [OKAY]
--------------------------------------------------
DeepSpeed general environment info:
torch install path ............... ['/home/ubuntu/.local/lib/python3.9/site-packages/torch']
torch version .................... 1.10.0+cu111
deepspeed install path ........... ['/home/ubuntu/.local/lib/python3.9/site-packages/deepspeed']
deepspeed info ................... 0.8.3, unknown, unknown
torch cuda version ............... 11.1
torch hip version ................ None
nvcc version ..................... 11.1
deepspeed wheel compiled w. ...... torch 0.0, cuda 0.0
```

**Error log**
```bash
Traceback (most recent call last):
   File "/home/ubuntu/Megatron-DeepSpeed/examples/MoE/../../pretrain_gpt.py", line 345, in <module>
     pretrain(train_valid_test_datasets_provider, model_provider, forward_step,
   File "/home/ubuntu/Megatron-DeepSpeed/megatron/training.py", line 107, in pretrain
     initialize_megatron(extra_args_provider=extra_args_provider,
   File "/home/ubuntu/Megatron-DeepSpeed/megatron/initialize.py", line 79, in initialize_megatron
     finish_mpu_init()
   File "/home/ubuntu/Megatron-DeepSpeed/megatron/initialize.py", line 60, in finish_mpu_init
     _initialize_distributed()
   File "/home/ubuntu/Megatron-DeepSpeed/megatron/initialize.py", line 216, in _initialize_distributed
     assert args.local_rank == device, \
 AssertionError: expected local-rank to be the same as rank % device-count.
```

**System info**
 - OS: Ubuntu 18.04
 - GPU count and types: Two machines, each with 2 T4 GPUs and 4 T4 GPUs, respectively
 - Python version: 3.9.4

**Launcher context**
`deepspeed` launcher

**Docker context**
N/A. Did not use Docker images.


Discussion:

Resolving patch:
diff --git a/deepspeed/launcher/launch.py b/deepspeed/launcher/launch.py
--- a/deepspeed/launcher/launch.py
+++ b/deepspeed/launcher/launch.py
@@ -215,9 +215,10 @@ def main():
                     raise ValueError(f"unable to create directory {args.enable_each_rank_log} for each rank log.")
             log_name_prefix = time.strftime("%Y%m%d%H%M%S", time.localtime())
 
-        for local_rank in range(0, num_local_procs):
+        for local_proc in range(0, num_local_procs):
             # each process's rank
-            dist_rank = global_rank_mapping[local_node][local_rank]
+            dist_rank = global_rank_mapping[local_node][local_proc]
+            local_rank = dist_rank % num_local_procs
             current_env["RANK"] = str(dist_rank)
             current_env["LOCAL_RANK"] = str(local_rank)
 


Test patch: