Repository: microsoft/DeepSpeed

Issue:
[BUG] Loading model with Zero3 results in NCCL error
I am running this code:
```Python
import os
import deepspeed
import torch
import transformers

local_rank = int(os.getenv('LOCAL_RANK', '0'))
world_size = int(os.getenv('WORLD_SIZE', '1'))

ds_config = {
  'train_batch_size': world_size,
  'fp16': {
    'enabled': True
  },
  "zero_optimization": {
    "stage": 3,
  },
}

from transformers.deepspeed import HfDeepSpeedConfig
dschf = HfDeepSpeedConfig(ds_config)
model = transformers.GPT2LMHeadModel.from_pretrained("gpt2")
```

I run it with `NCCL_DEBUG=INFO deepspeed --num_gpus 2 script.py`.

It fails with
```
RuntimeError: NCCL error in: /opt/conda/conda-bld/pytorch_1634272204863/work/torch/csrc/distributed/c10d/ProcessGroupNCCL.cpp:957, invalid usage, NCCL version 21.0.3
ncclInvalidUsage: This usually reflects invalid usage of NCCL library (such as too many async ops, too many collectives at once, mixing streams in a group, etc).
```

Interestingly, earlier I see these messages:
```
dirkg-zeppelin:4005:4041 [0] init.cc:521 NCCL WARN Duplicate GPU detected : rank 0 and rank 1 both on CUDA device 40
dirkg-zeppelin:4006:4043 [0] init.cc:521 NCCL WARN Duplicate GPU detected : rank 1 and rank 0 both on CUDA device 40
```

I'm using DeepSpeed fc2f378ece3055fddd13cdf309faa9ab66a1b302, but the latest release has the same behavior.

Full stack trace:
```
Traceback (most recent call last):
  File "/home/dirkg/test_inference.py", line 52, in <module>
    model = transformers.GPT2LMHeadModel.from_pretrained("gpt2")
  File "/home/dirkg/miniconda3/envs/megatron/lib/python3.9/site-packages/transformers/modeling_utils.py", line 1441, in from_pretrained
    model, missing_keys, unexpected_keys, mismatched_keys, error_msgs = cls._load_state_dict_into_model(
  File "/home/dirkg/miniconda3/envs/megatron/lib/python3.9/site-packages/transformers/modeling_utils.py", line 1591, in _load_state_dict_into_model
    load(model_to_load, prefix=start_prefix)
  File "/home/dirkg/miniconda3/envs/megatron/lib/python3.9/site-packages/transformers/modeling_utils.py", line 1576, in load
    load(child, prefix + name + ".")
  File "/home/dirkg/miniconda3/envs/megatron/lib/python3.9/site-packages/transformers/modeling_utils.py", line 1568, in load
    with deepspeed.zero.GatheredParameters(list(module.parameters(recurse=False)), modifier_rank=0):
  File "/home/dirkg/deepspeed-microsoft/deepspeed/runtime/zero/partition_parameters.py", line 1310, in __enter__
    self.params[0].all_gather(param_list=self.params)
  File "/home/dirkg/deepspeed-microsoft/deepspeed/runtime/zero/partition_parameters.py", line 595, in all_gather
    return self._all_gather(param_list, async_op=async_op, hierarchy=hierarchy)
  File "/home/dirkg/deepspeed-microsoft/deepspeed/runtime/zero/partition_parameters.py", line 704, in _all_gather
    ret_value = self._allgather_params_coalesced(all_gather_list, hierarchy)
  File "/home/dirkg/deepspeed-microsoft/deepspeed/runtime/zero/partition_parameters.py", line 950, in _allgather_params_coalesced
    h = dist._all_gather_base(allgather_params[param_idx],
  File "/home/dirkg/miniconda3/envs/megatron/lib/python3.9/site-packages/torch/distributed/distributed_c10d.py", line 2070, in _all_gather_base
    work = group._allgather_base(output_tensor, input_tensor)
RuntimeError: NCCL error in: /opt/conda/conda-bld/pytorch_1634272204863/work/torch/csrc/distributed/c10d/ProcessGroupNCCL.cpp:957, invalid usage, NCCL version 21.0.3
ncclInvalidUsage: This usually reflects invalid usage of NCCL library (such as too many async ops, too many collectives at once, mixing streams in a group, etc).
```

Output of `ds_report`:
```
dirkg@dirkg-zeppelin ~> ds_report                                                                                                            
--------------------------------------------------
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
--------------------------------------------------
cpu_adam ............... [NO] ....... [OKAY]
cpu_adagrad ............ [NO] ....... [OKAY]
fused_adam ............. [NO] ....... [OKAY]
fused_lamb ............. [NO] ....... [OKAY]
sparse_attn ............ [NO] ....... [OKAY]
transformer ............ [NO] ....... [OKAY]
stochastic_transformer . [NO] ....... [OKAY]
async_io ............... [NO] ....... [OKAY]
transformer_inference .. [NO] ....... [OKAY]
utils .................. [NO] ....... [OKAY]
quantizer .............. [NO] ....... [OKAY]
--------------------------------------------------
DeepSpeed general environment info:
torch install path ............... ['/home/dirkg/miniconda3/envs/megatron/lib/python3.9/site-packages/torch']
torch version .................... 1.10.0
torch cuda version ............... 11.3
nvcc version ..................... 11.3
deepspeed install path ........... ['/home/dirkg/deepspeed-microsoft/deepspeed']
deepspeed info ................... 0.5.9+fc2f378, fc2f378, master
deepspeed wheel compiled w. ...... torch 1.10, cuda 11.3
```


Discussion:

Resolving patch:
diff --git a/deepspeed/runtime/zero/partition_parameters.py b/deepspeed/runtime/zero/partition_parameters.py
--- a/deepspeed/runtime/zero/partition_parameters.py
+++ b/deepspeed/runtime/zero/partition_parameters.py
@@ -475,9 +475,10 @@ def get_model():
         self.rank = torch.distributed.get_rank(group=self.ds_process_group)
         self.world_size = torch.distributed.get_world_size(group=self.ds_process_group)
 
-        # Local device is the device where the parameters are consumed
+        # Local device is the device where the parameters are consumed, must be default device.
         # It is the device where parameters are fully instantiated using allgather
         self.local_device = torch.device('cuda:{}'.format(os.environ["LOCAL_RANK"]))
+        torch.cuda.set_device(self.local_device)
 
         if _ds_config is not None and _ds_config.zero_config.offload_param is not None:
             remote_device = _ds_config.zero_config.offload_param[OFFLOAD_PARAM_DEVICE]


Test patch: