Unexpected ZeroDivisionError when calling model.prune_heads
Author: ghost
Pull request: no
Labels: name: wontfix

# 🐛 Bug
Traceback (most recent call last):
  File "/Users/user/anaconda3/lib/python3.7/runpy.py", line 193, in _run_module_as_main
    "__main__", mod_spec)
  File "/Users/user/anaconda3/lib/python3.7/runpy.py", line 85, in _run_code
    exec(code, run_globals)
  File "/Users/user/Desktop/org/text_models_branches/section_extraction/text_models/bert_text_classification/train_bert_ml_mc.py", line 609, in <module>
    masking_amount=args.masking_amount,
  File "/Users/user/Desktop/org/text_models_branches/section_extraction/text_models/bert_text_classification/train_bert_ml_mc.py", line 288, in train_model
    local_rank=transformer_args["local_rank"],
  File "/Users/user/Desktop/org/text_models_branches/section_extraction/text_models/bert_text_classification/prune_attention_heads.py", line 488, in prune_model_and_return
    metric=metric,
  File "/Users/user/Desktop/org/text_models_branches/section_extraction/text_models/bert_text_classification/prune_attention_heads.py", line 396, in prune_heads
    model.prune_heads(heads_to_prune)
  File "/Users/user/Desktop/org/text_models_branches/section_extraction/text_models/env/lib/python3.7/site-packages/transformers/modeling_utils.py", line 234, in prune_heads
    self.base_model._prune_heads(heads_to_prune)
  File "/Users/user/Desktop/org/text_models_branches/section_extraction/text_models/env/lib/python3.7/site-packages/transformers/modeling_bert.py", line 635, in _prune_heads
    self.encoder.layer[layer].attention.prune_heads(heads)
  File "/Users/user/Desktop/org/text_models_branches/section_extraction/text_models/env/lib/python3.7/site-packages/transformers/modeling_bert.py", line 301, in prune_heads
    self.output.dense = prune_linear_layer(self.output.dense, index, dim=1)
  File "/Users/user/Desktop/org/text_models_branches/section_extraction/text_models/env/lib/python3.7/site-packages/transformers/modeling_utils.py", line 824, in prune_linear_layer
    new_layer = nn.Linear(new_size[1], new_size[0], bias=layer.bias is not None).to(layer.weight.device)
  File "/Users/user/Desktop/org/text_models_branches/section_extraction/text_models/env/lib/python3.7/site-packages/torch/nn/modules/linear.py", line 81, in __init__
    self.reset_parameters()
  File "/Users/user/Desktop/org/text_models_branches/section_extraction/text_models/env/lib/python3.7/site-packages/torch/nn/modules/linear.py", line 84, in reset_parameters
    init.kaiming_uniform_(self.weight, a=math.sqrt(5))
  File "/Users/user/Desktop/org/text_models_branches/section_extraction/text_models/env/lib/python3.7/site-packages/torch/nn/init.py", line 325, in kaiming_uniform_
    std = gain / math.sqrt(fan)
ZeroDivisionError: float division by zero

Basically, the error is being thrown on calling model.prune_heads(heads_to_prune). This error is not coming every time I run the script. Not sure what is causing the error. 

## Information

Model I am using (Bert, XLNet ...): BERT

Language I am using the model on (English, Chinese ...): English

The problem arises when using:
* [ ] the official example scripts: (give details below) - No
* [x] my own modified scripts: (give details below) - Yes 
My script calls mask_heads and then prune_heads similar to the original bertology script

The tasks I am working on is:
* [ ] an official GLUE/SQUaD task: (give the name)
* [x] my own task or dataset: (give details below)
A multi-class classification task on a proprietary dataset

## To reproduce

Steps to reproduce the behavior:

Still don't know as the error seems to be unexpected. I don't get it every time I run the script.

<!-- If you have code snippets, error messages, stack traces please provide them here as well.
     Important! Use code tags to correctly format your code. See https://help.github.com/en/github/writing-on-github/creating-and-highlighting-code-blocks#syntax-highlighting
     Do not use screenshots, as they are hard to read and (more importantly) don't allow others to copy-and-paste your code.-->

## Expected behavior

<!-- A clear and concise description of what you would expect to happen. --> 
model.prune_heads(heads_to_prune) where heads_to_prune -> Dict[int, List] where key is the layer number and value is the list of heads to prune (Calculated by calling the mask_heads function). Expected behaviour is pruning off heads present in heads_to_prune for each layer

## Environment info
<!-- You can run the command `python transformers-cli env` and copy-and-paste its output below.
     Don't forget to fill out the missing fields in that output! -->
     
- `transformers` version: 2.2.2
- Platform: macOS Catalina 10.15
- Python version: Python 3.7.3
- PyTorch version (GPU?): 1.3.1 (No GPU)
- Tensorflow version (GPU?): 1.14.0
- Using GPU in script?: No
- Using distributed or parallel set-up in script?: No


Comments:
This issue has been automatically marked as stale because it has not had recent activity. It will be closed if no further activity occurs. Thank you for your contributions.


Found the solution to this. When you try to prune all the attention heads in a layer, you will run into this error. This is why it sometimes shows up and sometimes does not, because your pruning function may or may not decide to prune all the attention heads in some layer depending on how you are computing the importance of each attention head. If you try the example on the Hugging Face Model page for the prune_heads function, {1: [0, 2], 2: [2, 3]}, it should work without any error (at least, you will not end up with the ZeroDivisionError. 

I was able to debug this by printing out what my original heads_to_prune dictionary looked like, and therefore noticed the edge case. With this hunch, testing it out on other cases confirmed the cause. In the future, printing out your inputs to the function that is returning the error is a good practice! Especially when the function is implemented by some entity like Hugging Face and the only thing that could probably go wrong is the input you give it.

Hope this helps!