[BUG] ZeRO Stage2 and 3, error while loss backward

In ZeRO stage1, it works.

I used from_pretrained("facebook/bart-base") as backbone in transformers==4.2.1

In ZeRO stage2, backward stop working like infinite loop at some processes.

In ZeRO stage3,

[2022-05-18 14:33:12,828] [WARNING] [stage3.py:106:_apply_to_tensors_only] A module has unknown inputs or outputs type (<class 'torch.device'>) and the tensors embedded in it cannot be detected. The ZeRO-3 hooks designed to trigger before or after backward pass of the module relies on knowing the input and output tensors and therefore may not get triggered properly.
[2022-05-18 14:33:12,828] [WARNING] [stage3.py:106:_apply_to_tensors_only] A module has unknown inputs or outputs type (<class 'torch.device'>) and the tensors embedded in it cannot be detected. The ZeRO-3 hooks designed to trigger before or after backward pass of the module relies on knowing the input and output tensors and therefore may not get triggered properly.

    self.model.backward(loss)
  File "/home/kyungmin.lee/DeepSpeed/deepspeed/utils/nvtx.py", line 11, in wrapped_fn
    return func(*args, **kwargs)
  File "/home/kyungmin.lee/DeepSpeed/deepspeed/runtime/engine.py", line 1726, in backward
    self.optimizer.backward(loss)
  File "/home/kyungmin.lee/DeepSpeed/deepspeed/utils/nvtx.py", line 11, in wrapped_fn
    return func(*args, **kwargs)
  File "/home/kyungmin.lee/DeepSpeed/deepspeed/runtime/zero/stage3.py", line 2538, in backward
    self._get_param_coordinator(training=True).reset_step()
  File "/home/kyungmin.lee/DeepSpeed/deepspeed/runtime/zero/partitioned_param_coordinator.py", line 194, in reset_step
    assert_ints_same_as_other_ranks([m.id for m in self.__submodule_order])
  File "/home/kyungmin.lee/DeepSpeed/deepspeed/utils/nvtx.py", line 11, in wrapped_fn
    return func(*args, **kwargs)
  File "/home/kyungmin.lee/DeepSpeed/deepspeed/runtime/zero/utils.py", line 86, in assert_ints_same_as_other_ranks
    raise RuntimeError(f"disagreement between rank0 and rank{dist.get_rank()}: "
RuntimeError: disagreement between rank0 and rank1: rank0: [0, 1, 3, 103, 103, 103, 103, 5, 75, 75, 75, 75, 75, 75, 75, 75, 75, 75, 75, 75, 75, 75, 75, 75, 75, 75, 75, 75, 75, 75, 75, 75, 75, 75, 75, 75, 75, 75, 75, 75, 75, 75, 75, 75, 75, 75, 75, 75, 75, 75, 75, 75, 75, 75, 75, 75, 75, 75, 76, 79, 80, 81, 82, 83, 79, 80, 81, 82, 83, 79, 80, 81, 82, 83, 79, 80, 81, 82, 83, 79, 80, 81, 82, 83, 85, 86, 87, 88, 89, 79, 80, 81, 82, 83, 97, 98, 99, 100, 101, 79, 80, 81, 82, 83, 97, 98, 99, 100, 101, 79, 80, 81, 82, 83, 79, 80, 81, 82, 83, 85, 86, 87, 88, 89, 97, 98, 99, 100, 101, 79, 80, 81, 82, 83, 97, 98, 99, 100, 101, 97, 98, 99, 100, 101, 97, 98, 99, 100, 101, 85, 86, 87, 88, 89, 85, 86, 87, 88, 89, 79, 80, 81, 82, 83, 85, 86, 87, 88, 89, 79, 80, 81, 82, 83, 85, 86, 87, 88, 89, 97, 98, 99, 100, 101, 85, 86, 87, 88, 89, 79, 80, 81, 82, 83, 97, 98, 99, 100, 101, 97, 98, 99, 100, 101, 79, 80, 81, 82, 83, 79, 80, 81, 82, 83, 97, 98, 99, 100, 101, 79, 80, 81, 82, 83, 79, 80, 81, 82, 83, 97, 98, 99, 100, 101, 85, 86, 87, 88, 89, 85, 86, 87, 88, 89, 79, 80, 81, 82, 83, 85, 86, 87, 88, 89, 79, 80, 81, 82, 83, 85, 86, 87, 88, 89, 79, 80, 81, 82, 83, 97, 98, 99, 100, 101, 79, 80, 81, 82, 83, 85, 86, 87, 88, 89, 97, 98, 99, 100, 101, 97, 98, 99, 100, 101, 97, 98, 99, 100, 101, 97, 98, 99, 100, 101, 79, 80, 81, 82, 83, 68, 69, 70, 71, 73, 72, 72, 72, 72, 72, 72, 72, 72, 72, 72, 72, 72, 72, 72, 72, 72, 72, 72, 72, 72, 72, 72, 72, 72, 72, 72, 72, 72, 72, 72, 72, 72, 72, 72, 72, 72, 72, 72, 72, 72, 72, 72, 72, 72, 72, 72, 72, 72, 72, 72, 67, 7, 8, 11, 9, 10, 12, 13, 14, 15, 16, 17, 18, 21, 19, 20, 22, 23, 24, 25, 26, 27, 28, 31, 29, 30, 32, 33, 34, 35, 36, 37, 38, 41, 39, 40, 42, 43, 44, 45, 46, 47, 48, 51, 49, 50, 52, 53, 54, 55, 56, 57, 58, 61, 59, 60, 62, 63, 64, 65, 66, 102, 103, 103, 103, 103, 104, 202, 106, 107, 110, 108, 109, 111, 112, 113, 116, 114, 115, 117, 118, 119, 120, 121, 122, 123, 126, 124, 125, 127, 128, 129, 132, 130, 131, 133, 134, 135, 136, 137, 138, 139, 142, 140, 141, 143, 144, 145, 148, 146, 147, 149, 150, 151, 152, 153, 154, 155, 158, 156, 157, 159, 160, 161, 164, 162, 163, 165, 166, 167, 168, 169, 170, 171, 174, 172, 173, 175, 176, 177, 180, 178, 179, 181, 182, 183, 184, 185, 186, 187, 190, 188, 189, 191, 192, 193, 196, 194, 195, 197, 198, 199, 200, 201, 203, 0, 203, 1, 102, 186, 187, 190, 188, 189, 191, 192, 193, 196, 194, 195, 197, 198, 199, 200, 201, 186, 201, 200, 199, 198, 193, 197, 195, 194, 196, 192, 187, 191, 189, 188, 190, 170, 171, 174, 172, 173, 175, 176, 177, 180, 178, 179, 181, 182, 183, 184, 185, 170, 185, 184, 183, 182, 177, 181, 179, 178, 180, 176, 171, 175, 173, 172, 174, 154, 155, 158, 156, 157, 159, 160, 161, 164, 162, 163, 165, 166, 167, 168, 169, 154, 169, 168, 167, 166, 161, 165, 163, 162, 164, 160, 155, 159, 157, 156, 158, 138, 139, 142, 140, 141, 143, 144, 145, 148, 146, 147, 149, 150, 151, 152, 153, 138, 153, 152, 151, 150, 145, 149, 147, 146, 148, 144, 139, 143, 141, 140, 142, 122, 123, 126, 124, 125, 127, 128, 129, 132, 130, 131, 133, 134, 135, 136, 137, 122, 137, 136, 135, 134, 129, 133, 131, 130, 132, 128, 123, 127, 125, 124, 126, 106, 107, 110, 108, 109, 111, 112, 113, 116, 114, 115, 117, 118, 119, 120, 121, 106, 121, 120, 119, 118, 113, 117, 115, 114, 3, 57, 58, 61, 59, 60, 62, 63, 64, 65, 66, 57, 66, 65, 64, 63, 58, 62, 60, 59, 61, 47, 48, 51, 49, 50, 52, 53, 54, 55, 56, 47, 56, 55, 54, 53, 48, 52, 50, 49, 51, 37, 38, 41, 39, 40, 42, 43, 44, 45, 46, 37, 46, 45, 44, 43, 38, 42, 40, 39, 41, 27, 28, 31, 29, 30, 32, 33, 34, 35, 36, 27, 36, 35, 34, 33, 28, 32, 30, 29, 31, 17, 18, 21, 19, 20, 22, 23, 24, 25, 26, 17, 26, 25, 24, 23, 18, 22, 20, 19, 21, 7, 8, 11, 9, 10, 12, 13, 14, 15, 16, 7, 16, 15, 14, 13, 8, 12, 10, 9, 11, 67, 68, 72, 72, 72, 72, 72, 72, 72, 72, 72, 72, 72, 72, 72, 72, 72, 72, 72, 72, 72, 72, 72, 72, 72, 72, 72, 72, 72, 72, 72, 72, 72, 72, 72, 72, 72, 72, 72, 72, 72, 72, 72, 72, 72, 72, 72, 72, 72, 72, 72, 72, 73, 69, 71, 70, 76, 79, 83, 82, 81, 80, 97, 101, 100, 99, 98, 97, 101, 100, 99, 98, 97, 101, 100, 99, 98, 97, 101, 100, 99, 98, 85, 89, 88, 87, 86, 79, 83, 82, 81, 80, 97, 101, 100, 99, 98, 79, 83, 82, 81, 80, 85, 89, 88, 87, 86, 79, 83, 82, 81, 80, 85, 89, 88, 87, 86, 79, 83, 82, 81, 80, 85, 89, 88, 87, 86, 85, 89, 88, 87, 86, 97, 101, 100, 99, 98, 79, 83, 82, 81, 80, 79, 83, 82, 81, 80, 97, 101, 100, 99, 98, 79, 83, 82, 81, 80, 79, 83, 82, 81, 80, 97, 101, 100, 99, 98, 97, 101, 100, 99, 98, 79, 83, 82, 81, 80, 85, 89, 88, 87, 86, 97, 101, 100, 99, 98, 85, 89, 88, 87, 86, 79, 83, 82, 81, 80, 85, 89, 88, 87, 86, 79, 83, 82, 81, 80, 85, 89, 88, 87, 86, 85, 89, 88, 87, 86, 97, 101, 100, 99, 98, 97, 101, 100, 99, 98, 97, 101, 100, 99, 98, 79, 83, 82, 81, 80, 97, 101, 100, 99, 98, 85, 89, 88, 87, 86, 79, 83, 82, 81, 80, 79, 83, 82, 81, 80, 97, 101, 100, 99, 98, 79, 83, 82, 81, 80, 97, 101, 100, 99, 98, 79, 83, 82, 81, 80, 85, 89, 88, 87, 86, 79, 83, 82, 81, 80, 79, 83, 82, 81, 80, 79, 83, 82, 81, 80, 79, 83, 82, 81, 80, 79, 83, 82, 81, 80], rank1: [0, 1, 3, 103, 103, 103, 103, 5, 75, 75, 75, 75, 75, 75, 75, 75, 75, 75, 75, 75, 75, 75, 75, 75, 75, 75, 75, 75, 75, 75, 75, 75, 75, 75, 75, 75, 75, 75, 75, 75, 75, 75, 75, 75, 75, 75, 75, 75, 75, 75, 75, 75, 75, 75, 75, 75, 75, 75, 76, 97, 98, 99, 100, 101, 97, 98, 99, 100, 101, 85, 86, 87, 88, 89, 79, 80, 81, 82, 83, 97, 98, 99, 100, 101, 79, 80, 81, 82, 83, 85, 86, 87, 88, 89, 85, 86, 87, 88, 89, 85, 86, 87, 88, 89, 79, 80, 81, 82, 83, 97, 98, 99, 100, 101, 85, 86, 87, 88, 89, 79, 80, 81, 82, 83, 85, 86, 87, 88, 89, 79, 80, 81, 82, 83, 97, 98, 99, 100, 101, 79, 80, 81, 82, 83, 79, 80, 81, 82, 83, 79, 80, 81, 82, 83, 85, 86, 87, 88, 89, 85, 86, 87, 88, 89, 85, 86, 87, 88, 89, 85, 86, 87, 88, 89, 79, 80, 81, 82, 83, 79, 80, 81, 82, 83, 79, 80, 81, 82, 83, 79, 80, 81, 82, 83, 79, 80, 81, 82, 83, 97, 98, 99, 100, 101, 97, 98, 99, 100, 101, 97, 98, 99, 100, 101, 85, 86, 87, 88, 89, 79, 80, 81, 82, 83, 85, 86, 87, 88, 89, 85, 86, 87, 88, 89, 85, 86, 87, 88, 89, 79, 80, 81, 82, 83, 79, 80, 81, 82, 83, 85, 86, 87, 88, 89, 79, 80, 81, 82, 83, 79, 80, 81, 82, 83, 79, 80, 81, 82, 83, 97, 98, 99, 100, 101, 85, 86, 87, 88, 89, 85, 86, 87, 88, 89, 79, 80, 81, 82, 83, 79, 80, 81, 82, 83, 79, 80, 81, 82, 83, 79, 80, 81, 82, 83, 97, 98, 99, 100, 101, 68, 69, 70, 71, 73, 72, 72, 72, 72, 72, 72, 72, 72, 72, 72, 72, 72, 72, 72, 72, 72, 72, 72, 72, 72, 72, 72, 72, 72, 72, 72, 72, 72, 72, 72, 72, 72, 72, 72, 72, 72, 72, 72, 72, 72, 72, 72, 72, 72, 72, 72, 72, 72, 72, 72, 67, 7, 8, 11, 9, 10, 12, 13, 14, 15, 16, 17, 18, 21, 19, 20, 22, 23, 24, 25, 26, 27, 28, 31, 29, 30, 32, 33, 34, 35, 36, 37, 38, 41, 39, 40, 42, 43, 44, 45, 46, 47, 48, 51, 49, 50, 52, 53, 54, 55, 56, 57, 58, 61, 59, 60, 62, 63, 64, 65, 66, 102, 103, 103, 103, 103, 104, 202, 106, 107, 110, 108, 109, 111, 112, 113, 116, 114, 115, 117, 118, 119, 120, 121, 122, 123, 126, 124, 125, 127, 128, 129, 132, 130, 131, 133, 134, 135, 136, 137, 138, 139, 142, 140, 141, 143, 144, 145, 148, 146, 147, 149, 150, 151, 152, 153, 154, 155, 158, 156, 157, 159, 160, 161, 164, 162, 163, 165, 166, 167, 168, 169, 170, 171, 174, 172, 173, 175, 176, 177, 180, 178, 179, 181, 182, 183, 184, 185, 186, 187, 190, 188, 189, 191, 192, 193, 196, 194, 195, 197, 198, 199, 200, 201, 203, 0, 203, 1, 102, 186, 187, 190, 188, 189, 191, 192, 193, 196, 194, 195, 197, 198, 199, 200, 201, 186, 201, 200, 199, 198, 193, 197, 195, 194, 196, 192, 187, 191, 189, 188, 190, 170, 171, 174, 172, 173, 175, 176, 177, 180, 178, 179, 181, 182, 183, 184, 185, 170, 185, 184, 183, 182, 177, 181, 179, 178, 180, 176, 171, 175, 173, 172, 174, 154, 155, 158, 156, 157, 159, 160, 161, 164, 162, 163, 165, 166, 167, 168, 169, 154, 169, 168, 167, 166, 161, 165, 163, 162, 164, 160, 155, 159, 157, 156, 158, 138, 139, 142, 140, 141, 143, 144, 145, 148, 146, 147, 149, 150, 151, 152, 153, 138, 153, 152, 151, 150, 145, 149, 147, 146, 148, 144, 139, 143, 141, 140, 142, 122, 123, 126, 124, 125, 127, 128, 129, 132, 130, 131, 133, 134, 135, 136, 137, 122, 137, 136, 135, 134, 129, 133, 131, 130, 132, 128, 123, 127, 125, 124, 126, 106, 107, 110, 108, 109, 111, 112, 113, 116, 114, 115, 117, 118, 119, 120, 121, 106, 121, 120, 119, 118, 113, 117, 115, 114, 3, 57, 58, 61, 59, 60, 62, 63, 64, 65, 66, 57, 66, 65, 64, 63, 58, 62, 60, 59, 61, 47, 48, 51, 49, 50, 52, 53, 54, 55, 56, 47, 56, 55, 54, 53, 48, 52, 50, 49, 51, 37, 38, 41, 39, 40, 42, 43, 44, 45, 46, 37, 46, 45, 44, 43, 38, 42, 40, 39, 41, 27, 28, 31, 29, 30, 32, 33, 34, 35, 36, 27, 36, 35, 34, 33, 28, 32, 30, 29, 31, 17, 18, 21, 19, 20, 22, 23, 24, 25, 26, 17, 26, 25, 24, 23, 18, 22, 20, 19, 21, 7, 8, 11, 9, 10, 12, 13, 14, 15, 16, 7, 16, 15, 14, 13, 8, 12, 10, 9, 11, 67, 68, 72, 72, 72, 72, 72, 72, 72, 72, 72, 72, 72, 72, 72, 72, 72, 72, 72, 72, 72, 72, 72, 72, 72, 72, 72, 72, 72, 72, 72, 72, 72, 72, 72, 72, 72, 72, 72, 72, 72, 72, 72, 72, 72, 72, 72, 72, 72, 72, 72, 72, 73, 69, 71, 70, 76, 97, 101, 100, 99, 98, 79, 83, 82, 81, 80, 79, 83, 82, 81, 80, 79, 83, 82, 81, 80, 79, 83, 82, 81, 80, 85, 89, 88, 87, 86, 85, 89, 88, 87, 86, 97, 101, 100, 99, 98, 79, 83, 82, 81, 80, 79, 83, 82, 81, 80, 79, 83, 82, 81, 80, 85, 89, 88, 87, 86, 79, 83, 82, 81, 80, 79, 83, 82, 81, 80, 85, 89, 88, 87, 86, 85, 89, 88, 87, 86, 85, 89, 88, 87, 86, 79, 83, 82, 81, 80, 85, 89, 88, 87, 86, 97, 101, 100, 99, 98, 97, 101, 100, 99, 98, 97, 101, 100, 99, 98, 79, 83, 82, 81, 80, 79, 83, 82, 81, 80, 79, 83, 82, 81, 80, 79, 83, 82, 81, 80, 79, 83, 82, 81, 80, 85, 89, 88, 87, 86, 85, 89, 88, 87, 86, 85, 89, 88, 87, 86, 85, 89, 88, 87, 86, 79, 83, 82, 81, 80, 79, 83, 82, 81, 80, 79, 83, 82, 81, 80, 97, 101, 100, 99, 98, 79, 83, 82, 81, 80, 85, 89, 88, 87, 86, 79, 83, 82, 81, 80, 85, 89, 88, 87, 86, 97, 101, 100, 99, 98, 79, 83, 82, 81, 80, 85, 89, 88, 87, 86, 85, 89, 88, 87, 86, 85, 89, 88, 87, 86, 79, 83, 82, 81, 80, 97, 101, 100, 99, 98, 79, 83, 82, 81, 80, 85, 89, 88, 87, 86, 97, 101, 100, 99, 98, 97, 101, 100, 99, 98]

Stage2 config

{
  "zero_optimization": {
    "stage": 2,
    "allgather_partitions": true,
    "allgather_bucket_size": 5e8,
    "overlap_comm": true,
    "reduce_scatter": true,
    "reduce_bucket_size": 5e8,
    "contiguous_gradients": true
  },
"gradient_accumulation_steps": 1,
"gradient_clipping": 1.0,
"train_batch_size": 100,
"train_micro_batch_size_per_gpu": 50,
"steps_per_print":10000
}

Stage3 config

{
  "zero_optimization": {
    "stage": 3,
    "allgather_partitions": true,
    "allgather_bucket_size": 5e8,
    "overlap_comm": true,
    "reduce_scatter": true,
    "reduce_bucket_size": 5e8,
    "contiguous_gradients": true,
    "stage3_max_live_parameters" : 1e9,
    "stage3_max_reuse_distance" : 1e9,
    "stage3_prefetch_bucket_size" : 5e8,
    "stage3_param_persistence_threshold" : 1e6,
    "sub_group_size" : 1e12,
    "elastic_checkpoint" : [true],
    "stage3_gather_16bit_weights_on_model_save": [false],
    "ignore_unused_parameters": [true],
    "round_robin_gradients": [false]
  },
"gradient_accumulation_steps": 1,
"gradient_clipping": 1.0,
"train_batch_size": 100,
"train_micro_batch_size_per_gpu": 50,
"steps_per_print":10000
}

System info (please complete the following information):

• OS: Ubuntu 20.04
• used 2 A100 GPUs, [one machines with x16 A100s each]
• Python: 3.8.13
• DeepSpeed: 0.6.5+5053217e
• Transformers: 4.2.1