[Whisper] Gradient checkpointing fails
Author: sanchit-gandhi
Pull request: no

### System Info

- `transformers` version: 4.24.0.dev0
- Platform: Linux-4.19.0-22-cloud-amd64-x86_64-with-glibc2.10
- Python version: 3.8.0
- Huggingface_hub version: 0.10.1
- PyTorch version (GPU?): 1.12.1 (True)
- Tensorflow version (GPU?): not installed (NA)
- Flax version (CPU?/GPU?/TPU?): not installed (NA)
- Jax version: not installed
- JaxLib version: not installed
- Using GPU in script?: yes
- Using distributed or parallel set-up in script?: no

### Who can help?

cc @ArthurZucker for info

### Information

- [X] The official example scripts
- [ ] My own modified scripts

### Tasks

- [X] An officially supported task in the `examples` folder (such as GLUE/SQuAD, ...)
- [ ] My own task or dataset (give details below)

### Reproduction

```python
from transformers import WhisperFeatureExtractor, WhisperConfig, WhisperForConditionalGeneration
import numpy as np

feature_extractor = WhisperFeatureExtractor()
config = WhisperConfig()
model_encoder = WhisperForConditionalGeneration(config).model.encoder

# enable checkpointing
model_encoder.gradient_checkpointing_enable()

# create dummy audio input
sample = {"array": np.ones(1000), "sampling_rate": 16000}

# pre-process audio input
inputs = feature_extractor(sample["array"], sampling_rate=sample["sampling_rate"], return_tensors="pt").input_features

# forward pass
outputs = model_encoder(inputs)
```
```python
---------------------------------------------------------------------------
AttributeError                            Traceback (most recent call last)
Input In [1], in <cell line: 18>()
     15 inputs = feature_extractor(sample["array"], sampling_rate=sample["sampling_rate"], return_tensors="pt").input_features
     17 # forward pass
---> 18 outputs = model_encoder(inputs)

File ~/venv/lib/python3.8/site-packages/torch/nn/modules/module.py:1131, in Module._call_impl(self, *input, **kwargs)
   1127 # If we don't have any hooks, we want to skip the rest of the logic in
   1128 # this function, and just call forward.
   1129 if not (self._backward_hooks or self._forward_hooks or self._forward_pre_hooks or _global_backward_hooks
   1130         or _global_forward_hooks or _global_forward_pre_hooks):
-> 1131     return forward_call(*input, **kwargs)
   1132 # Do not call functions when jit is used
   1133 full_backward_hooks, non_full_backward_hooks = [], []

File ~/transformers/src/transformers/models/whisper/modeling_whisper.py:682, in WhisperEncoder.forward(self, input_features, head_mask, output_attentions, output_hidden_states, return_dict)
    678             return module(*inputs, output_attentions)
    680         return custom_forward
--> 682     layer_outputs = torch.utils.checkpoint.checkpoint(
    683         create_custom_forward(encoder_layer),
    684         hidden_states,
    685         None,
    686         (head_mask[idx] if head_mask is not None else None),
    687     )
    688 else:
    689     layer_outputs = encoder_layer(
    690         hidden_states,
    691         None,
    692         layer_head_mask=(head_mask[idx] if head_mask is not None else None),
    693         output_attentions=output_attentions,
    694     )

AttributeError: module 'torch.utils' has no attribute 'checkpoint'
```



### Expected behavior

Forward pass without any hitch-ups!

Comments:
Jumped the gun 😅 Doesn't quite yet work with the decoder!
```python
from transformers import WhisperFeatureExtractor, WhisperConfig, WhisperForConditionalGeneration
import numpy as np
import torch

feature_extractor = WhisperFeatureExtractor()
config = WhisperConfig()
model = WhisperForConditionalGeneration(config)

# enable checkpointing
model.gradient_checkpointing_enable()

# create dummy audio input
sample = {"array": np.ones(1000), "sampling_rate": 16000}

# pre-process audio input
inputs = feature_extractor(sample["array"], sampling_rate=sample["sampling_rate"], return_tensors="pt").input_features

# create dummy decoder input ids
decoder_input_ids = torch.arange(10).reshape(1, 10)  # bsz, seq_len = (1, 10)

# forward pass
outputs = model(inputs, decoder_input_ids=decoder_input_ids)
```

<details>
<summary> Traceback </summary>

```python
---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
Input In [2], in <cell line: 22>()
     19 decoder_input_ids = torch.arange(10).reshape(1, 10)  # bsz, seq_len = (1, 10)
     21 # forward pass
---> 22 outputs = model(inputs, decoder_input_ids=decoder_input_ids)

File ~/venv/lib/python3.8/site-packages/torch/nn/modules/module.py:1131, in Module._call_impl(self, *input, **kwargs)
   1127 # If we don't have any hooks, we want to skip the rest of the logic in
   1128 # this function, and just call forward.
   1129 if not (self._backward_hooks or self._forward_hooks or self._forward_pre_hooks or _global_backward_hooks
   1130         or _global_forward_hooks or _global_forward_pre_hooks):
-> 1131     return forward_call(*input, **kwargs)
   1132 # Do not call functions when jit is used
   1133 full_backward_hooks, non_full_backward_hooks = [], []

File ~/transformers/src/transformers/models/whisper/modeling_whisper.py:1168, in WhisperForConditionalGeneration.forward(self, input_features, decoder_input_ids, decoder_attention_mask, head_mask, decoder_head_mask, cross_attn_head_mask, encoder_outputs, past_key_values, decoder_inputs_embeds, labels, use_cache, output_attentions, output_hidden_states, return_dict)
   1163     if decoder_input_ids is None:
   1164         decoder_input_ids = shift_tokens_right(
   1165             labels, self.config.pad_token_id, self.config.decoder_start_token_id
   1166         )
-> 1168 outputs = self.model(
   1169     input_features,
   1170     decoder_input_ids=decoder_input_ids,
   1171     encoder_outputs=encoder_outputs,
   1172     decoder_attention_mask=decoder_attention_mask,
   1173     head_mask=head_mask,
   1174     decoder_head_mask=decoder_head_mask,
   1175     cross_attn_head_mask=cross_attn_head_mask,
   1176     past_key_values=past_key_values,
   1177     decoder_inputs_embeds=decoder_inputs_embeds,
   1178     use_cache=use_cache,
   1179     output_attentions=output_attentions,
   1180     output_hidden_states=output_hidden_states,
   1181     return_dict=return_dict,
   1182 )
   1183 lm_logits = self.proj_out(outputs[0])
   1185 loss = None

File ~/venv/lib/python3.8/site-packages/torch/nn/modules/module.py:1131, in Module._call_impl(self, *input, **kwargs)
   1127 # If we don't have any hooks, we want to skip the rest of the logic in
   1128 # this function, and just call forward.
   1129 if not (self._backward_hooks or self._forward_hooks or self._forward_pre_hooks or _global_backward_hooks
   1130         or _global_forward_hooks or _global_forward_pre_hooks):
-> 1131     return forward_call(*input, **kwargs)
   1132 # Do not call functions when jit is used
   1133 full_backward_hooks, non_full_backward_hooks = [], []

File ~/transformers/src/transformers/models/whisper/modeling_whisper.py:1044, in WhisperModel.forward(self, input_features, decoder_input_ids, decoder_attention_mask, head_mask, decoder_head_mask, cross_attn_head_mask, encoder_outputs, past_key_values, decoder_inputs_embeds, use_cache, output_attentions, output_hidden_states, return_dict)
   1037     encoder_outputs = BaseModelOutput(
   1038         last_hidden_state=encoder_outputs[0],
   1039         hidden_states=encoder_outputs[1] if len(encoder_outputs) > 1 else None,
   1040         attentions=encoder_outputs[2] if len(encoder_outputs) > 2 else None,
   1041     )
   1043 # decoder outputs consists of (dec_features, past_key_value, dec_hidden, dec_attn)
-> 1044 decoder_outputs = self.decoder(
   1045     input_ids=decoder_input_ids,
   1046     attention_mask=decoder_attention_mask,
   1047     encoder_hidden_states=encoder_outputs[0],
   1048     head_mask=decoder_head_mask,
   1049     cross_attn_head_mask=cross_attn_head_mask,
   1050     past_key_values=past_key_values,
   1051     inputs_embeds=decoder_inputs_embeds,
   1052     use_cache=use_cache,
   1053     output_attentions=output_attentions,
   1054     output_hidden_states=output_hidden_states,
   1055     return_dict=return_dict,
   1056 )
   1058 if not return_dict:
   1059     return decoder_outputs + encoder_outputs

File ~/venv/lib/python3.8/site-packages/torch/nn/modules/module.py:1131, in Module._call_impl(self, *input, **kwargs)
   1127 # If we don't have any hooks, we want to skip the rest of the logic in
   1128 # this function, and just call forward.
   1129 if not (self._backward_hooks or self._forward_hooks or self._forward_pre_hooks or _global_backward_hooks
   1130         or _global_forward_hooks or _global_forward_pre_hooks):
-> 1131     return forward_call(*input, **kwargs)
   1132 # Do not call functions when jit is used
   1133 full_backward_hooks, non_full_backward_hooks = [], []

File ~/transformers/src/transformers/models/whisper/modeling_whisper.py:912, in WhisperDecoder.forward(self, input_ids, attention_mask, encoder_hidden_states, head_mask, cross_attn_head_mask, past_key_values, inputs_embeds, use_cache, output_attentions, output_hidden_states, return_dict)
    908             return module(*inputs, output_attentions, use_cache)
    910         return custom_forward
--> 912     layer_outputs = torch.utils.checkpoint.checkpoint(
    913         create_custom_forward(decoder_layer),
    914         hidden_states,
    915         attention_mask,
    916         encoder_hidden_states,
    917         head_mask[idx] if head_mask is not None else None,
    918         cross_attn_head_mask[idx] if cross_attn_head_mask is not None else None,
    919         None,
    920     )
    921 else:
    923     layer_outputs = decoder_layer(
    924         hidden_states,
    925         attention_mask=attention_mask,
   (...)
    933         use_cache=use_cache,
    934     )

File ~/venv/lib/python3.8/site-packages/torch/utils/checkpoint.py:235, in checkpoint(function, use_reentrant, *args, **kwargs)
    232     raise ValueError("Unexpected keyword arguments: " + ",".join(arg for arg in kwargs))
    234 if use_reentrant:
--> 235     return CheckpointFunction.apply(function, preserve, *args)
    236 else:
    237     return _checkpoint_without_reentrant(
    238         function,
    239         preserve,
    240         *args
    241     )

File ~/venv/lib/python3.8/site-packages/torch/utils/checkpoint.py:96, in CheckpointFunction.forward(ctx, run_function, preserve_rng_state, *args)
     93 ctx.save_for_backward(*tensor_inputs)
     95 with torch.no_grad():
---> 96     outputs = run_function(*args)
     97 return outputs

File ~/transformers/src/transformers/models/whisper/modeling_whisper.py:908, in WhisperDecoder.forward.<locals>.create_custom_forward.<locals>.custom_forward(*inputs)
    906 def custom_forward(*inputs):
    907     # None for past_key_value
--> 908     return module(*inputs, output_attentions, use_cache)

File ~/venv/lib/python3.8/site-packages/torch/nn/modules/module.py:1131, in Module._call_impl(self, *input, **kwargs)
   1127 # If we don't have any hooks, we want to skip the rest of the logic in
   1128 # this function, and just call forward.
   1129 if not (self._backward_hooks or self._forward_hooks or self._forward_pre_hooks or _global_backward_hooks
   1130         or _global_forward_hooks or _global_forward_pre_hooks):
-> 1131     return forward_call(*input, **kwargs)
   1132 # Do not call functions when jit is used
   1133 full_backward_hooks, non_full_backward_hooks = [], []

File ~/transformers/src/transformers/models/whisper/modeling_whisper.py:397, in WhisperDecoderLayer.forward(self, hidden_states, attention_mask, encoder_hidden_states, encoder_attention_mask, layer_head_mask, cross_attn_layer_head_mask, past_key_value, output_attentions, use_cache)
    393 hidden_states = self.self_attn_layer_norm(hidden_states)
    395 # Self Attention
    396 # decoder uni-directional self-attention cached key/values tuple is at positions 1,2
--> 397 self_attn_past_key_value = past_key_value[:2] if past_key_value is not None else None
    398 # add present self-attn cache to positions 1,2 of present_key_value tuple
    399 hidden_states, self_attn_weights, present_key_value = self.self_attn(
    400     hidden_states=hidden_states,
    401     past_key_value=self_attn_past_key_value,
   (...)
    404     output_attentions=output_attentions,
    405 )

TypeError: 'bool' object is not subscriptable
```

</details>

Need to pass `encoder_attention_mask` and `past_key_value` correctly to the decoder layer...