Pruning function in T5Attention doesnt affect _relative_position_bucket
Author: hadaev8
Pull request: no
Labels: name: bug

### Who can help?


@patrickvonplaten


### Information

- [X] The official example scripts
- [ ] My own modified scripts

### Tasks

- [ ] An officially supported task in the `examples` folder (such as GLUE/SQuAD, ...)
- [ ] My own task or dataset (give details below)

### Reproduction

Run pruning function in t5 model, then run inference.

### Expected behavior

Relative position head should be pruned too.

Here it is
https://github.com/huggingface/transformers/blob/main/src/transformers/models/t5/modeling_t5.py#L355


Comments:
@hadaev8 

It is not clear to me about this.

`_relative_position_bucket` is a `staticmethod` without using any model weight in it, and IMO there is no need to do anything when pruning a model.

cc @patrickvonplaten 

@ydshieh 

Relative position bias have shape (dim, heads).
For example I have 6 heads and pruned one, would be mismatch, (dim, 5) + (dim, 6)

Here this line
https://github.com/huggingface/transformers/blob/3ccff0d400ffd1b0c5074e15afb2b1f2af0e7b44/src/transformers/models/t5/modeling_t5.py#L529

I realized all layers use same positional bias, so it should be masked in forward, not pruned.

After looking the 2 blocks below, I think there is indeed a shape issue when we prune the heads.

Would you like to try to make a minimal code snippet that could confirm the issue, @hadaev8?

https://github.com/huggingface/transformers/blob/3ccff0d400ffd1b0c5074e15afb2b1f2af0e7b44/src/transformers/models/t5/modeling_t5.py#L432

https://github.com/huggingface/transformers/blob/3ccff0d400ffd1b0c5074e15afb2b1f2af0e7b44/src/transformers/models/t5/modeling_t5.py#L351

@ydshieh 
Here it is
https://colab.research.google.com/drive/1HYu-yzmmbumbskGZExXlOP0WFmDYdgAp?usp=sharing

I fixed rel pos bias, but where is some other error

Hey @hadaev8,

This is quite an edge case and I don't think it'll be to find an easy fix here because usually one only prunes some heads of some layers (not of all layers), where as the same `position_bias` is applied to **all** layers. So pruning some heads of only some layers will necessarily lead to problems here. The solution I see it to dynamically discard the superfluous dimensions of `relative_attention_bias`at every attention layer if the corresponding head has been discarded. @hadaev8 would you be interested in opening a PR for this? I won't have the time to dive deeper here for this sadly in the near future, but more than happy to review! 

@patrickvonplaten 
My fix looks like this and seems to work, but I'm not satisfied with it, idk if it worth adding to codebase.

```
        if self.pruned_heads:
            mask = torch.ones(position_bias.shape[1])
            mask[list(self.pruned_heads)] = 0
            position_bias_masked = position_bias[:,mask.bool()]
        else:
            position_bias_masked = position_bias

        scores += position_bias_masked
```

Hey @hadaev8,

That's actually quite a smart fix :-) Think I'd be ok with adding this! Do you want to open a PR for it ? :-)

@patrickvonplaten 
Okay, if you think its ok, i will do pr tomorrow.

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the [contributing guidelines](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md) are likely to be ignored.