Wav2Vec2 CUDA memory usage doubled in v4.11.3 compared to v4.10.3 with the same batch size
Author: MarktHart
Pull request: no

## Environment info
- `transformers` version: 4.11.3
- Platform: Linux-5.11.0-40-generic-x86_64-with-glibc2.29
- Python version: 3.8.10
- PyTorch version (GPU?): 1.8.1+cu111 (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, 3090
- Using distributed or parallel set-up in script?: No

### Who can help
@patrickvonplaten, @anton-l

## Information

When using Wav2vec2 the memory usage roughly doubles when going from Huggingface v4.10.3 to v4.11.3
Whereas my 3090 (24GB memory) in v4.10.3 could handle a batchsize of ~32, in 4.11.3 this is reduced to ~10.

The problem arises when using:
* my own modified scripts

The tasks I am working on is:
*  ASR

## To reproduce

Steps to reproduce the behavior:

1. Run script with v4.10 and v4.11 and watch CUDA memory usage

Reproduce script (relatively minimal):
```
from transformers import Wav2Vec2ForCTC, Wav2Vec2Processor, TrainingArguments
from transformers.trainer import Trainer
from torch.utils.data.dataset import Dataset
import numpy as np

class ProcessedDataset(Dataset):
    def __init__(self, processor):
        self.processor = processor

    def __getitem__(self, i):
        x = np.ones(16000 * 10) # 10 seconds
        y = "this is a random sentence"
        with self.processor.as_target_processor():
            batch= {"labels": self.processor(y).input_ids}
        batch["input_values"] = self.processor(x, sampling_rate=16000).input_values
        return batch

    def __len__(self):
        return 10000

class DataCollator:
    def __init__(self, processor):
        self.processor = processor

    def __call__(self, features):
        input_features = [{"input_values": feature["input_values"][0]} for feature in features]
        label_features = [{"input_ids": feature["labels"]} for feature in features]
        batch = self.processor.pad(
            input_features,
            padding=True,
            max_length=None,
            pad_to_multiple_of=None,
            return_tensors="pt",
        )
        with self.processor.as_target_processor():
            labels_batch = self.processor.pad(
                label_features,
                padding=True,
                max_length=None,
                pad_to_multiple_of=None,
                return_tensors="pt",
            )
        labels = labels_batch["input_ids"].masked_fill(labels_batch.attention_mask.ne(1), -100)
        batch["labels"] = labels
        return batch


proc = Wav2Vec2Processor.from_pretrained("wietsedv/wav2vec2-large-xlsr-53-dutch")
model = Wav2Vec2ForCTC.from_pretrained(
    "facebook/wav2vec2-large-nl-voxpopuli",
    attention_dropout=0,
    hidden_dropout=0,
    feat_proj_dropout=0,
    mask_time_prob=0,
    layerdrop=0,
    activation_dropout=0,
    gradient_checkpointing=True,
    ctc_loss_reduction="mean",
    pad_token_id=proc.tokenizer.pad_token_id,
    vocab_size=len(proc.tokenizer),
    ctc_zero_infinity=True
)
ds = ProcessedDataset(proc)
data_collator = DataCollator(processor=proc)
args = TrainingArguments(
    output_dir="/tmp/tmp_model",
    per_device_train_batch_size=8,
    gradient_accumulation_steps=1,
    do_eval=False,
    num_train_epochs=1,
    fp16=True,
    group_by_length=False,
    save_steps=-1,
    eval_steps=1024,
    logging_steps=1024,
    warmup_steps=128,
    save_total_limit=1,
    dataloader_num_workers=1,
    seed=11
)

trainer = Trainer(model=model, args=args, train_dataset=ds, data_collator=data_collator)
trainer.train()

```

## Expected behavior

Upgrading Huggingface Transformers from 4.10 to a later version should keep the memory usage in the same ballpark


Comments:
Looking into it now

Benchmarking your script on current master gives me a peak GPU mem usage of `20068MiB` .

And with `4.10` it gives me `10738MiB` => so this seems like a pretty heavy bug! Thanks for the heads-up!

Will investigate now

No problem at all! If there is anything I can do to assist I would be happy to help. 

Ok I think I already found one problem. It seems like the `gradient_checkpointing` [PR](https://github.com/huggingface/transformers/pull/13657) refactor wasn't 100% backward compatible. 

@MarktHart - could you add 

```python
model.gradient_checkpointing_enable()
```

before this line:

```
trainer = Trainer(model=model, args=args, train_dataset=ds, data_collator=data_collator)
```

this should more or less solve the problem

That does solve the issue. Thanks a bunch! 

@patrickvonplaten do you decide whether to close the issue or that backward compatibility should be restored? 

@sgugger - this is a weird issue. For some reason `from_pretrained(...)` doesn't currently set `gradient_checkpointing` to `True` at the first init since the main models does not have the `nn.Modules` attached yet. 

Will open a hacky PR to fix it

> Ok I think I already found one problem. It seems like the `gradient_checkpointing` [PR](https://github.com/huggingface/transformers/pull/13657) refactor wasn't 100% backward compatible.
> 
> @MarktHart - could you add
> 
> ```python
> model.gradient_checkpointing_enable()
> ```
> 
> before this line:
> 
> ```
> trainer = Trainer(model=model, args=args, train_dataset=ds, data_collator=data_collator)
> ```
> 
> this should more or less solve the problem

I have this issue in 4.14.1 when i set group_by_length=True. Adding model.gradient_checkpointing_enable() can't solve this problem. 

@voidful - can you provide a reproducible script here? :-) Thanks a lot!

> @voidful - can you provide a reproducible script here? :-) Thanks a lot!

It turn out to be length issue on my custom dataset, simplify apply .filter can solve this problem～～～～
Sorry for misleading.