AttributeError: 'NoneType' object has no attribute 'type' with overlap_comm=True and zero 2

[szhengac]

When I train the model with zero 2 and overlap_comm=True on one single P4d instance, I received the following error:

10.0.89.152: File "/usr/local/lib64/python3.7/site-packages/deepspeed/runtime/engine.py", line 850, in backward
10.0.89.152: self.allreduce_gradients()
10.0.89.152: File "/usr/local/lib64/python3.7/site-packages/deepspeed/runtime/engine.py", line 770, in allreduce_gradients
10.0.89.152: self.optimizer.overlapping_partition_gradients_reduce_epilogue()
10.0.89.152: File "/usr/local/lib64/python3.7/site-packages/deepspeed/runtime/zero/stage2.py", line 581, in overlapping_partition_gradients_reduce_epilogue
10.0.89.152: self.independent_gradient_partition_epilogue()
10.0.89.152: File "/usr/local/lib64/python3.7/site-packages/deepspeed/runtime/zero/stage2.py", line 470, in independent_gradient_partition_epilogue
10.0.89.152: self.reduce_ipg_grads()
10.0.89.152: File "/usr/local/lib64/python3.7/site-packages/deepspeed/runtime/zero/stage2.py", line 954, in reduce_ipg_grads
10.0.89.152: elements_per_buffer=self.elements_in_ipg_bucket)
10.0.89.152: File "/usr/local/lib64/python3.7/site-packages/deepspeed/runtime/zero/stage2.py", line 1123, in buffered_reduce_fallback
10.0.89.152: split_buckets = split_half_float_double(grads)
10.0.89.152: File "/usr/local/lib64/python3.7/site-packages/deepspeed/runtime/zero/stage2.py", line 40, in split_half_float_double
10.0.89.152: bucket = [t for t in tensors if t.type() == dtype]
10.0.89.152: File "/usr/local/lib64/python3.7/site-packages/deepspeed/runtime/zero/stage2.py", line 40, in <listcomp>
10.0.89.152: bucket = [t for t in tensors if t.type() == dtype]
10.0.89.152: AttributeError: 'NoneType' object has no attribute 'type'

The model config is:

    "bert_model_config": {
        "vocab_size_or_config_json_file": 32003,
        "hidden_size": 1024,
        "num_hidden_layers": 38,
        "num_attention_heads": 16,
        "intermediate_size": 4096,
        "hidden_act": "gelu",
        "hidden_dropout_prob": 0.1,
        "attention_probs_dropout_prob": 0.1,
        "max_position_embeddings": 512,
        "initializer_range": 0.02
    }

If I turn off overlap_comm, it works.

[szhengac]

After some debugging, I found that it is due to that deepspeed does not handle tied parameters properly. In my example, I use tied weight for the embedding and output layers. There are two issues:

1. The grad hook is called twice for the tied parameters. In the second hook call for the tied parameters, a None grad is added: https://github.com/microsoft/DeepSpeed/blob/e59ba12d4df2e61e85abbd2c82a6097ac3722d7a/deepspeed/runtime/zero/stage2.py#L641
   when the tied parameters are also in self.previous_reduced_grads and their gradients are set to None as in
   https://github.com/microsoft/DeepSpeed/blob/e59ba12d4df2e61e85abbd2c82a6097ac3722d7a/deepspeed/runtime/zero/stage2.py#L1089.
2. The gradient aggregation is wrong for the tied parameters. Due to the pre-division by world_size https://github.com/microsoft/DeepSpeed/blob/e59ba12d4df2e61e85abbd2c82a6097ac3722d7a/deepspeed/runtime/zero/stage2.py#L1066, the final gradient is (1/n^2\sum_i g_1^i) + (1/n\sum_i g_2^i), where g_1^i and g_2^i are the gradients on rank i in w.r.t. output and embedding layers, respectively. This means the gradients from the output layer are weighted by 1/n times more! The main reason we have this issue is that we only keep this partition of the gradients on one rank and discard it on other ranks.

[szhengac]

If I am right about what I mentioned above. This also results in wrong aggregated gradients when we use zero 2 and gradient accumulation. Suppose that the size of gradient accumulation is 3, then the final aggregated gradient is 1/n^3\sum_ig_{t-2, i} + 1/n^2\sum_ig_{t-1, i} + 1/n\sum_ig_{t, i}.

[szhengac]

@ShadenSmith Hi Shaden, would anyone in the deepspeed team be able to confirm if I am right about it? Thanks.

[tjruwase]

@szhengac, apologies for not responding to this earlier. I certainly see the potential issues you have raised with ZeRO-2. Specifically, the handling of tied weights. Can you share repro steps on this as well?

[szhengac]

@tjruwase Apart from the tied weights, this also seems problematic to me when we use gradient accumulation. As I mentioned above the final aggregated gradient is 1/n^3\sum_ig_{t-2, i} + 1/n^2\sum_ig_{t-1, i} + 1/n\sum_ig_{t, i} for an accumulation size of 3 based on the current zero-2 logic. We don't have the problem without the gradient partition just like what apex is doing.

[tjruwase]

@szhengac I definitely also see the issue with gradient accumulation, which is a challenging for ZeRO-2 in general. But I think we should start first with tied weights issue.

[szhengac]

@tjruwase To reproduce the Nonetype error for the tied parameters, you can use bing_bert in DeepSpeedExamples on a single node. I use script ds_train_bert_nvidia_data_bsz64k_seq128.sh and add

  "zero_optimization": {
          "stage": 2,
          "overlap_comm": true
  },
  "zero_allow_untested_optimizer": true,

to deepspeed_bsz64k_lamb_config_seq128.json. I use the following model config in bert_large_lamb_nvidia_data.json

    "bert_model_config": {
        "vocab_size_or_config_json_file": 119547,
        "hidden_size": 1024,
        "num_hidden_layers": 38,
        "num_attention_heads": 16,
        "intermediate_size": 4096,
        "hidden_act": "gelu",
        "hidden_dropout_prob": 0.1,
        "attention_probs_dropout_prob": 0.1,
        "max_position_embeddings": 512,
        "type_vocab_size": 2,
        "initializer_range": 0.02
    },

Also, you may want to notice a recent PR for fixing Nonetype in a wrong way: #728. With such fix, we will lose the gradient for the embedding layers.

[tjruwase]

@szhengac Thanks for sharing these repro steps. I will look into it.

Also, thanks for sharing feedback on #728, I will check into it.

[szhengac]

@tjruwase I realized that ZeRO 3 has come out. I wonder if the issues mentioned here also exist in ZeRO 3?

[tjruwase]

@szhengac, thanks for keeping this alive. I suspect the same issues with weight tying applies to ZeRO3. With the release completed, I am now looking into deepspeed's support for tied weights.

[tjruwase]

@szhengac, unfortunately I am not able to easily repro this AttributeError with the bing_bert even after reverting #728 in my branch. More specifically, my run has completed 1 epoch without the error, which means the multiple gradient reduction steps has already occurred.

Can you give more details on your bing_bert repro? In my zero_optimization config, I set only stage=2 and overlap_comm=True

[szhengac]

@tjruwase Do you use the model config that has "num_hidden_layers": 38 I give to you?

[tjruwase]

@szhengac , oops, sorry I missed that part. I will do so now.

[tjruwase]

@szhengac, I have repro'd successfully. Thanks so much.