How to use model.prunes when you are using transformers.T5ForConditionalGeneration
Author: CaffreyR
Pull request: no

### System Info

- `transformers` version: 4.20.1
- Platform: macOS-12.4-arm64-arm-64bit
- Python version: 3.9.10
- Huggingface_hub version: 0.8.1
- PyTorch version (GPU?): 1.13.0.dev20220709 (False)
- Tensorflow version (GPU?): 2.8.0 (True)
- Flax version (CPU?/GPU?/TPU?): not installed (NA)
- Jax version: not installed
- JaxLib version: not installed
- Using GPU in script?: <fill in>
- Using distributed or parallel set-up in script?: <fill in>

### 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

I use this code to prune the model from `T5ForConditionalGeneration`, but it went wrong. Many thanks for your time!:)
```
from transformers import T5ForConditionalGeneration
model = T5ForConditionalGeneration.from_pretrained('t5-base')
prune_heads = {}
prune_heads[0] = [0,1]
model.prune_heads(prune_heads)
```

### Expected behavior

```
Traceback (most recent call last):
  File "/Users/caffrey/Documents/research/FiD/prunetest.py", line 8, in <module>
    model.prune_heads(prune_heads)
  File "/Users/caffrey/miniforge3/envs/tongji/lib/python3.9/site-packages/transformers/modeling_utils.py", line 1507, in prune_heads
    self.base_model._prune_heads(heads_to_prune)
  File "/Users/caffrey/miniforge3/envs/tongji/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1261, in __getattr__
    raise AttributeError("'{}' object has no attribute '{}'".format(
AttributeError: 'T5ForConditionalGeneration' object has no attribute '_prune_heads'
```

Comments:
Gently pinging @ArthurZucker 

Thanks @ArthurZucker

Gently pinging @LysandreJik

Hey! I'll have a look sorry for the long wait 😃 

Anyone can help? Thanks! @gante @ArthurZucker 

Hey! It seems that the `T5ForConditionalGeneration` is simply missing the `_prune_heads`.
We should add the following few lines : 
```

def _prune_heads(self, heads_to_prune):
        """
        Prunes heads of the model. heads_to_prune: dict of {layer_num: list of heads to prune in this layer} See base
        class PreTrainedModel
        """
        for layer, heads in heads_to_prune.items():
            self.modules()[layer].layer[0].SelfAttention.prune_heads(heads)

```
Or something along those lines! 
Would you like to open a PR for that? Otherwise I will take care of it 🤗

Hi @ArthurZucker, I will give it a try :)



Hi @ArthurZucker , there seems a little more problem. Since in this line 
https://github.com/huggingface/transformers/blob/bd469c40659ce76c81f69c7726759d249b4aef49/src/transformers/models/t5/modeling_t5.py#L533
The model report error when I successfully pruned some heads, it says
```
  File "/home/user/anaconda3/envs/uw/lib/python3.7/site-packages/transformers/models/t5/modeling_t5.py", line 533, in forward
    mask[list(self.pruned_heads)] = 0
IndexError: index 8 is out of bounds for dimension 0 with size 8
```

Are we gonna modify the `forward`?

Since I try to print the `mask.shape` and `self.pruned_heads)`

It says
```
torch.Size([12])
{8, 2, 10, 6}
torch.Size([8])
{8, 2, 10, 6}
Traceback (most recent call last):

```


Hi @ArthurZucker , I open a PR here. https://github.com/huggingface/transformers/pull/19975

We can see the test on a colab https://colab.research.google.com/drive/1b9mHjtn2UxuHU_Sb_RXts12rDzbebBX0#scrollTo=hUSe4a1oOp6D

I use `opendelta` to visualize the pruning process.

But we seems to be a forward problem

We can conclude that the problem is between L531~L548

The difference is whether use head_mask

> Hi @ArthurZucker , there seems a little more problem. Since in this line
> 
> https://github.com/huggingface/transformers/blob/bd469c40659ce76c81f69c7726759d249b4aef49/src/transformers/models/t5/modeling_t5.py#L533
> 
> 
> The model report error when I successfully pruned some heads, it says
> ```
>   File "/home/user/anaconda3/envs/uw/lib/python3.7/site-packages/transformers/models/t5/modeling_t5.py", line 533, in forward
>     mask[list(self.pruned_heads)] = 0
> IndexError: index 8 is out of bounds for dimension 0 with size 8
> ```
> 
> Are we gonna modify the `forward`?
> 
> Since I try to print the `mask.shape` and `self.pruned_heads)`
> 
> It says
> 
> ```
> torch.Size([12])
> {8, 2, 10, 6}
> torch.Size([8])
> {8, 2, 10, 6}
> Traceback (most recent call last):
> ```


This one occur in this code
```
outputs = model.forward(
                input_ids=context_ids,
                attention_mask=context_mask,
                labels=labels,
                return_dict=True,
                # head_mask=head_mask,
                # decoder_head_mask=decoder_head_mask
            )
```




> Hi @ArthurZucker , I open a PR here. #19975
> 
> We can see the test on a colab https://colab.research.google.com/drive/1b9mHjtn2UxuHU_Sb_RXts12rDzbebBX0#scrollTo=hUSe4a1oOp6D
> 
> I use `opendelta` to visualize the pruning process.
> 
> But we seems to be a forward problem

This one use the code 
```
outputs = model.forward(
                input_ids=context_ids,
                attention_mask=context_mask,
                labels=labels,
                return_dict=True,
                head_mask=head_mask,
                decoder_head_mask=decoder_head_mask
            )
```

Hi @ArthurZucker , what is the function of `position_bias` ？ It seems in this line 

https://github.com/huggingface/transformers/blob/bd469c40659ce76c81f69c7726759d249b4aef49/src/transformers/models/t5/modeling_t5.py#L513

We calculate it, and it seems only modify `position_bias` only when the first block. I want to delete `if` or put the `scores += position_bias_masked` into `if`, which means we calculate `position_bias` only in the first block or we calculate `position_bias` all the blocks and it should be the same to the score.

You can see the code here
```
scores = torch.matmul(
            query_states, key_states.transpose(3, 2)
        )  # equivalent of torch.einsum("bnqd,bnkd->bnqk", query_states, key_states), compatible with onnx op>9
        print("Score",scores.shape)
        if position_bias is None:
            print("A", position_bias)
            if not self.has_relative_attention_bias:
                position_bias = torch.zeros(
                    (1, self.n_heads, real_seq_length, key_length), device=scores.device, dtype=scores.dtype
                )
                print("B",position_bias.shape)
                if self.gradient_checkpointing and self.training:
                    position_bias.requires_grad = True
            else:
                position_bias = self.compute_bias(real_seq_length, key_length, device=scores.device)
                print("C", position_bias.shape)
            # if key and values are already calculated
            # we want only the last query position bias
            if past_key_value is not None:
                position_bias = position_bias[:, :, -hidden_states.size(1) :, :]

            if mask is not None:
                position_bias = position_bias + mask  # (batch_size, n_heads, seq_length, key_length)

        if self.pruned_heads:
            position_bias_masked = position_bias
            # print(self.pruned_heads)
            # mask = torch.ones(position_bias.shape[1])
            # mask[list(self.pruned_heads)] = 0
            # print("Position bias",position_bias.shape)
            # position_bias_masked = position_bias[:, mask.bool()]
            # print("Position bias masked",position_bias_masked.shape)
        else:
            position_bias_masked = position_bias


        scores += position_bias_masked
```

And output is here
```
Score torch.Size([1, 8, 2, 2])
A None
C torch.Size([1, 8, 2, 2])
query_states torch.Size([1, 8, 2, 64])
key_states.transpose(3, 2) torch.Size([1, 8, 64, 200])
Score torch.Size([1, 8, 2, 200])
A None
B torch.Size([1, 8, 2, 200])
query_states torch.Size([1, 9, 2, 64])
key_states.transpose(3, 2) torch.Size([1, 9, 64, 2])
Score torch.Size([1, 9, 2, 2])
Traceback (most recent call last):
  File "/Users/caffrey/Documents/research/FiD/prunetest2.py", line 72, in <module>
    outputs = model.forward(
  File "/Users/caffrey/miniforge3/envs/huggingface/lib/python3.9/site-packages/transformers/models/t5/modeling_t5.py", line 1658, in forward
    decoder_outputs = self.decoder(
  File "/Users/caffrey/miniforge3/envs/huggingface/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1190, in _call_impl
    return forward_call(*input, **kwargs)
  File "/Users/caffrey/miniforge3/envs/huggingface/lib/python3.9/site-packages/transformers/models/t5/modeling_t5.py", line 1050, in forward
    layer_outputs = layer_module(
  File "/Users/caffrey/miniforge3/envs/huggingface/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1190, in _call_impl
    return forward_call(*input, **kwargs)
  File "/Users/caffrey/miniforge3/envs/huggingface/lib/python3.9/site-packages/transformers/models/t5/modeling_t5.py", line 683, in forward
    self_attention_outputs = self.layer[0](
  File "/Users/caffrey/miniforge3/envs/huggingface/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1190, in _call_impl
    return forward_call(*input, **kwargs)
  File "/Users/caffrey/miniforge3/envs/huggingface/lib/python3.9/site-packages/transformers/models/t5/modeling_t5.py", line 589, in forward
    attention_output = self.SelfAttention(
  File "/Users/caffrey/miniforge3/envs/huggingface/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1190, in _call_impl
    return forward_call(*input, **kwargs)
  File "/Users/caffrey/miniforge3/envs/huggingface/lib/python3.9/site-packages/transformers/models/t5/modeling_t5.py", line 548, in forward
    scores += position_bias_masked
RuntimeError: The size of tensor a (9) must match the size of tensor b (8) at non-singleton dimension 1

```
So the `position_bias` can not have the same shape as `score`

I also add a code in _prune_head to re-define the `self.relative_attention_bias`


In the PR https://github.com/huggingface/transformers/pull/19975/ , I delete `if` so that it could match with the shape of score. (It could run the colab notebook), but I do not know whether its meaning is right! Basically, the only problem is to make `position_bias ` have the same shape as `score`

gently pin @patrickvonplaten @ArthurZucker Many thanks

I will have a look 🤗

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.

Hi @CaffreyR

I am also facing the same issue. Did you find a fix yet?

Really appreciate any help! Raised an issue recently, so am willing to close if there's some progress :)