[s2s finetune] huge increase in memory demands with --fp16 native amp
Author: stas00
Pull request: no

While working on https://github.com/huggingface/transformers/issues/8353 I discovered that `--fp16` causes a 10x+ increase in gpu memory demands.

e.g. I can run bs=12 w/o  `--fp16` 

```
cd examples/seq2seq
export BS=12; rm -rf distilbart-cnn-12-6; python finetune.py --learning_rate=3e-5 --gpus 1 \
--do_train --do_predict --val_check_interval 0.25 --n_val 500 --num_train_epochs 2 --freeze_encoder \
--freeze_embeds --data_dir cnn_dm --max_target_length 142 --val_max_target_length=142 \
--train_batch_size=$BS --eval_batch_size=$BS --gradient_accumulation_steps 1 \
--model_name_or_path sshleifer/student_cnn_12_6 --tokenizer_name facebook/bart-large \
--warmup_steps 500 --output_dir distilbart-cnn-12-6

```
But if I add:
```
--fp16
```

(w/ or w/o `--fp16_opt_level O1`)

I get OOM even with bs=1 on a 8GB card and it barely manages on a 24GB card - I think the increase in memory demand is more than 10x.

The OOM either right away when it does the sanity check step, or after just 10-20 batches - so within a few secs

This is with pytorch-1.6. Same goes for pytorch-1.7 and 1.8-nightly.

I wasn't able to test `--fp16` with pytorch-1.5, since I can't build apex on ubuntu-20.04. Without `--fp16` pytorch-1.5 works the same as pytorch-1.6 gpu memory-wise.

I tested with pytorch-1.5 + apex and there is no problem there. Memory consumption is about half.

Here is the table of the batch sizes that fit into a 8gb rtx-1070 (bigger BS leads to an instant OOM):

bs | version
---|--------
12 | pt15
20 | pt15+fp16
12 | pt16
1 | pt16+fp16



If you'd like to reproduce the problem here are the full steps:

```
# prep library
git clone https://github.com/huggingface/transformers
cd transformers
pip install -e .[dev]
pip install -r examples/requirements.txt
cd examples/seq2seq

# prep data
wget https://cdn-datasets.huggingface.co/summarization/cnn_dm_v2.tgz
tar -xzvf cnn_dm_v2.tgz  # empty lines removed
mv cnn_cln cnn_dm

# run
export BS=12; 
rm -rf distilbart-cnn-12-6
python finetune.py --learning_rate=3e-5 --gpus 1 \
--do_train --do_predict --val_check_interval 0.25 --n_val 500 --num_train_epochs 2 --freeze_encoder \
--freeze_embeds --data_dir cnn_dm --max_target_length 142 --val_max_target_length=142 \
--train_batch_size=$BS --eval_batch_size=$BS --gradient_accumulation_steps 1 \
--model_name_or_path sshleifer/student_cnn_12_6 --tokenizer_name facebook/bart-large \
--warmup_steps 500 --output_dir distilbart-cnn-12-6 
```

This issue is to track the problem and hopefully finding a solution.

@sshleifer 

Comments:
I managed to install `nvidia-apex` binary via conda:
```
conda install nvidia-apex -c conda-forge
``` 
So now I was able to validate that with pytorch-1.5 + nvdia-apex `--fp16` consumes less memory, than w/o  `--fp16`. 

I was able to squeeze bs=20 (!) onto a 8gb card.

So the problem has to do with pytorch's 1.6 fp16


Found another report of memory increase with fp16:
https://discuss.pytorch.org/t/fp16-training-with-feedforward-network-slower-time-and-no-memory-reduction/95560/


Michael Carilli suggested I add this to the top of the script to find the problem:
```
import torch
torch.cuda.amp.autocast = "hfdj"
```
rerunning gives:
```
Validation sanity check: 0it [00:00, ?it/s]Traceback (most recent call last):
  File "finetune.py", line 449, in <module>
    main(args)
  File "finetune.py", line 416, in main
    trainer: pl.Trainer = generic_train(
  File "/mnt/nvme1/code/huggingface/transformers-watchdog/examples/lightning_base.py", line 395, in generic_train
    trainer.fit(model)
  File "/mnt/nvme1/code/github/00pytorch/pytorch-lightning/pytorch_lightning/trainer/trainer.py", line 445, in fit
    results = self.accelerator_backend.train()
  File "/mnt/nvme1/code/github/00pytorch/pytorch-lightning/pytorch_lightning/accelerators/gpu_accelerator.py", line 54, in train
    results = self.train_or_test()
  File "/mnt/nvme1/code/github/00pytorch/pytorch-lightning/pytorch_lightning/accelerators/accelerator.py", line 74, in train_or_test
    results = self.trainer.train()
  File "/mnt/nvme1/code/github/00pytorch/pytorch-lightning/pytorch_lightning/trainer/trainer.py", line 467, in train
    self.run_sanity_check(self.get_model())
  File "/mnt/nvme1/code/github/00pytorch/pytorch-lightning/pytorch_lightning/trainer/trainer.py", line 671, in run_sanity_check
    _, eval_results = self.run_evaluation(test_mode=False, max_batches=self.num_sanity_val_batches)
  File "/mnt/nvme1/code/github/00pytorch/pytorch-lightning/pytorch_lightning/trainer/trainer.py", line 591, in run_evaluation
    output = self.evaluation_loop.evaluation_step(test_mode, batch, batch_idx, dataloader_idx)
  File "/mnt/nvme1/code/github/00pytorch/pytorch-lightning/pytorch_lightning/trainer/evaluation_loop.py", line 176, in evaluation_step
    output = self.trainer.accelerator_backend.validation_step(args)
  File "/mnt/nvme1/code/github/00pytorch/pytorch-lightning/pytorch_lightning/accelerators/gpu_accelerator.py", line 75, in validation_step
    with torch.cuda.amp.autocast():
TypeError: 'str' object is not callable
```

I am on the master of PL (pytorch-lightning), but tried bisecting on earlier versions (3 months back) with no change in behavior.

The same issue, when I run the train_distilbart_xsum.sh, in sanity check step, meet OOM in 12G card

Based on some initial debugging Lightning is calling into `trainer,model.validation_step`, which is calling `BartForConditionalGeneration` with  fixed tensor sizes of:
```python
torch.Size([48, 1, 1024]) torch.Size([50264, 1024]) torch.Size([1, 50264])
```
In each iteration more tensors are allocated and never freed, which yields the OOM in the `sanity_check`.

AMP run before OOM:
```python
|===========================================================================|
|                  PyTorch CUDA memory summary, device ID 0                 |
|---------------------------------------------------------------------------|
|            CUDA OOMs: 0            |        cudaMalloc retries: 0         |
|===========================================================================|
|        Metric         | Cur Usage  | Peak Usage | Tot Alloc  | Tot Freed  |
|---------------------------------------------------------------------------|
| Allocated memory      |   13553 MB |   13555 MB |  139840 MB |  126286 MB |
|       from large pool |   13551 MB |   13551 MB |  137780 MB |  124228 MB |
|       from small pool |       2 MB |      25 MB |    2059 MB |    2057 MB |
|---------------------------------------------------------------------------|
| Active memory         |   13553 MB |   13555 MB |  139840 MB |  126286 MB |
|       from large pool |   13551 MB |   13551 MB |  137780 MB |  124228 MB |
|       from small pool |       2 MB |      25 MB |    2059 MB |    2057 MB |
|---------------------------------------------------------------------------|
| GPU reserved memory   |   13888 MB |   13888 MB |   13888 MB |       0 B  |
|       from large pool |   13858 MB |   13858 MB |   13858 MB |       0 B  |
|       from small pool |      30 MB |      30 MB |      30 MB |       0 B  |
|---------------------------------------------------------------------------|
| Non-releasable memory |  315554 KB |     913 MB |  109380 MB |  109072 MB |
|       from large pool |  313557 KB |     912 MB |  107043 MB |  106737 MB |
|       from small pool |    1997 KB |       8 MB |    2336 MB |    2334 MB |
|---------------------------------------------------------------------------|
| Allocations           |    3403    |    3410    |   29885    |   26482    |
|       from large pool |    3161    |    3161    |    7265    |    4104    |
|       from small pool |     242    |     267    |   22620    |   22378    |
|---------------------------------------------------------------------------|
| Active allocs         |    3403    |    3410    |   29885    |   26482    |
|       from large pool |    3161    |    3161    |    7265    |    4104    |
|       from small pool |     242    |     267    |   22620    |   22378    |
|---------------------------------------------------------------------------|
| GPU reserved segments |     368    |     368    |     368    |       0    |
|       from large pool |     353    |     353    |     353    |       0    |
|       from small pool |      15    |      15    |      15    |       0    |
|---------------------------------------------------------------------------|
| Non-releasable allocs |     204    |     210    |   14948    |   14744    |
|       from large pool |     200    |     205    |    3772    |    3572    |
|       from small pool |       4    |      20    |   11176    |   11172    |
|===========================================================================|
```

FP32 run for the same step:
```python
|===========================================================================|
|                  PyTorch CUDA memory summary, device ID 0                 |
|---------------------------------------------------------------------------|
|            CUDA OOMs: 0            |        cudaMalloc retries: 0         |
|===========================================================================|
|        Metric         | Cur Usage  | Peak Usage | Tot Alloc  | Tot Freed  |
|---------------------------------------------------------------------------|
| Allocated memory      |    3948 MB |    4000 MB |  203229 MB |  199281 MB |
|       from large pool |    3946 MB |    3998 MB |  201302 MB |  197356 MB |
|       from small pool |       2 MB |      24 MB |    1927 MB |    1925 MB |
|---------------------------------------------------------------------------|
| Active memory         |    3948 MB |    4000 MB |  203229 MB |  199281 MB |
|       from large pool |    3946 MB |    3998 MB |  201302 MB |  197356 MB |
|       from small pool |       2 MB |      24 MB |    1927 MB |    1925 MB |
|---------------------------------------------------------------------------|
| GPU reserved memory   |    4226 MB |    4226 MB |    4226 MB |       0 B  |
|       from large pool |    4198 MB |    4198 MB |    4198 MB |       0 B  |
|       from small pool |      28 MB |      28 MB |      28 MB |       0 B  |
|---------------------------------------------------------------------------|
| Non-releasable memory |  259725 KB |     778 MB |  118090 MB |  117836 MB |
|       from large pool |  257727 KB |     777 MB |  115922 MB |  115670 MB |
|       from small pool |    1997 KB |       8 MB |    2168 MB |    2166 MB |
|---------------------------------------------------------------------------|
| Allocations           |     415    |     422    |   23263    |   22848    |
|       from large pool |     173    |     175    |    3796    |    3623    |
|       from small pool |     242    |     267    |   19467    |   19225    |
|---------------------------------------------------------------------------|
| Active allocs         |     415    |     422    |   23263    |   22848    |
|       from large pool |     173    |     175    |    3796    |    3623    |
|       from small pool |     242    |     267    |   19467    |   19225    |
|---------------------------------------------------------------------------|
| GPU reserved segments |      85    |      85    |      85    |       0    |
|       from large pool |      71    |      71    |      71    |       0    |
|       from small pool |      14    |      14    |      14    |       0    |
|---------------------------------------------------------------------------|
| Non-releasable allocs |      17    |      26    |   13706    |   13689    |
|       from large pool |      13    |      14    |    2264    |    2251    |
|       from small pool |       4    |      21    |   11442    |   11438    |
|===========================================================================|
```
The allocations increase using `AMP` by ~50 in each iteration, while FP32 increases them sometimes by ~1.
So far I wasn't able to isolate the memory increase. 

@ptrblck, and if the same is done with pt15/nvidia-apex - why doesn't the same happen there? How are the two different (native vs apex)

I don't know, why an older PyTorch version with apex/amp works, as I wasn't able to isolate the issue yet.
The native amp implementation differs in various ways from `apex/amp`.
Btw. do you only see the OOM, if you are using Lightning or also using the standalone model + amp?

@ptrblck,  I have tried the hugging face trainer and the problem doesn't seem to happen there. I don't think `--fp16` makes any difference in that trainer but it doesn't increase memory requirements. So is it possible this is specifically a PL issue?

For testing I used:

setup:
```
cd examples/seq2seq
wget https://cdn-datasets.huggingface.co/translation/wmt_en_ro.tar.gz
tar -xzvf wmt_en_ro.tar.gz
```
and then run:
```

bs=11; rm -rf tmpdir; PYTHONPATH="../../src" python ./finetune_trainer.py \
--model_name_or_path sshleifer/distill-mbart-en-ro-12-4 --data_dir wmt_en_ro --output_dir tmpdir \
--overwrite_output_dir --max_source_length 128 --max_target_length 128 --val_max_target_length 128 \
--do_train --do_eval --do_predict --num_train_epochs 10 --per_device_train_batch_size $bs \
--per_device_eval_batch_size $bs --learning_rate 3e-4 --warmup_steps 2 --evaluate_during_training \
--predict_with_generate --logging_steps 0 --save_steps 2 --eval_steps 2 --sortish_sampler \
--label_smoothing 0.1 --adafactor --task translation --tgt_lang ro_RO --src_lang en_XX --n_train 100 \
--n_val 50 --fp16
```
(same memory consumption w/o `--fp16`)

This command also uses `BartForConditionalGeneration`.

`bs=11` is the biggest batch size I could fit onto the 8GB card, `bs=12` OOMs


@ptrblck How did you make such a nice table?

@stas00 I will check your #s on my card. cc @patil-suraj 



I replicated the OOM and fixed by passing `--amp_backend='apex'` in my torch 1.6 environment on a 24GB card. Would still be good to see if there is any easy way to get native amp working well.

Thank you for validating that, @sshleifer. So it really has something to do with the native amp in pytorch.

Here is the summary of my experiments:

- pt15 + confa-force apex w/ `--fp16` works at the start - reduces memory consumption
- pt16 + conda-forge apex w/ `--fp16  --amp_backend='apex'` works at the start too!

but both fail at the end with:
```
  File "python3.8/site-packages/apex/amp/_amp_state.py", line 32, in warn_or_err
    raise RuntimeError(msg)
RuntimeError: Found param model.model.shared.weight with type torch.cuda.HalfTensor, expected torch.cuda.FloatTensor.
When using amp.initialize, you do not need to call .half() on your model
```
If I use `--fp16_opt_level O1`, the failure changes to:
```
  File "python3.8/site-packages/torch/optim/lr_scheduler.py", line 56, in with_counter
    instance_ref = weakref.ref(method.__self__)
AttributeError: 'function' object has no attribute '__self__'
```
both envs use the binary apex from `conda-forge`

- apex doesn't support cuda11 at the moment
- rtx-30* cards don't support cuda<11 (and really cuda<11.1).
- ubuntu-20.4 doesn't support cuda<11, since it dropped gcc7, so can't build apex from source even for cuda-10

Bottom line, apex is not a great option at the moment, but may work for some short term - need to sort out native amp.

I will poke at it with debugger today.

I'm running in parallel `pt15+apex` and `pt16+native` with debugger:

Found the first issue. 

At this point in stack (we are in PL domain) - both have about 1.7GB allocated on GPU:
```
restore_weights, checkpoint_connector.py:64
setup_training, training_loop.py:174
train, gpu_accelerator.py:51
fit, trainer.py:444
generic_train, lightning_base.py:398
main, finetune.py:413
<module>, finetune.py:446
```

next step: `torch.cuda.empty_cache()` frees about 0.6GB on pt15, but 0 on pt16.

Could `GradScaler` be holding onto objects and not letting them go? I commented out its init and there is no change. 

Perhaps there are some circular references preventing the objects from being cleared out.


@sshleifer I used [`torch.cuda.memory_summary()`](https://pytorch.org/docs/stable/cuda.html#torch.cuda.memory_summary) and added it for each iteration in the `BartForConditionalGeneration` model.

After some more debugging it seems that the `autocast` cache is blowing up.
As a workaround you can add `torch.clear_autocast_cache()` in [BartForConditionalGeneration.forward](https://github.com/huggingface/transformers/blob/eb3bd73ce35bfef56eeb722d697f2d39a06a8f8d/src/transformers/modeling_bart.py#L1036), which might slow down your code but should at least work.

Based on @stas00 's debugging it seems that PL is interacting with native AMP in a way that the cache is increasing.
CC @mcarilli

The autocast cache is [cleared automatically](https://github.com/pytorch/pytorch/blob/88ec72e1c2a4b1e2a15cbe4703b9567bf9369a09/torch/cuda/amp/autocast_mode.py#L127) every time you exit an autocast context, which is one reason autocast should wrap the forward pass then exit.

@ptrblck where in `BartConditionalGeneration.forward` did you call `clear_autocast_cache()` to resolve the memory blowup?

It's also helpful to write in the following sequence:
```
torch.autocast_increment_nesting()
print(torch.autocast_decrement_nesting())
```
to see how deeply we're nested in autocast contexts at that point in forward.  It should print `1`.

@ptrblck, if I use your suggestion to add `torch.clear_autocast_cache()` torch blows up:

```
Epoch 0:  50%|████████████████Traceback (most recent call last):                                | 1/2 [00:02<00:02,  2.37s/it, loss=2.934, v_num=156]
  File "finetune.py", line 447, in <module>
    main(args)
  File "finetune.py", line 414, in main
    trainer: pl.Trainer = generic_train(
  File "/mnt/nvme1/code/huggingface/transformers-watchdog/examples/lightning_base.py", line 403, in generic_train
    trainer.fit(model)
  File "/home/stas/anaconda3/envs/main-38/lib/python3.8/site-packages/pytorch_lightning/trainer/trainer.py", line 444, in fit
    results = self.accelerator_backend.train()
  File "/home/stas/anaconda3/envs/main-38/lib/python3.8/site-packages/pytorch_lightning/accelerators/gpu_accelerator.py", line 63, in train
    results = self.train_or_test()
  File "/home/stas/anaconda3/envs/main-38/lib/python3.8/site-packages/pytorch_lightning/accelerators/accelerator.py", line 74, in train_or_test
    results = self.trainer.train()
  File "/home/stas/anaconda3/envs/main-38/lib/python3.8/site-packages/pytorch_lightning/trainer/trainer.py", line 493, in train
    self.train_loop.run_training_epoch()
  File "/home/stas/anaconda3/envs/main-38/lib/python3.8/site-packages/pytorch_lightning/trainer/training_loop.py", line 589, in run_training_epoch
    self.trainer.run_evaluation(test_mode=False)
  File "/home/stas/anaconda3/envs/main-38/lib/python3.8/site-packages/pytorch_lightning/trainer/trainer.py", line 578, in run_evaluation
    output = self.evaluation_loop.evaluation_step(test_mode, batch, batch_idx, dataloader_idx)
  File "/home/stas/anaconda3/envs/main-38/lib/python3.8/site-packages/pytorch_lightning/trainer/evaluation_loop.py", line 171, in evaluation_step
    output = self.trainer.accelerator_backend.validation_step(args)
  File "/home/stas/anaconda3/envs/main-38/lib/python3.8/site-packages/pytorch_lightning/accelerators/gpu_accelerator.py", line 85, in validation_step
    output = self.__validation_step(args)
  File "/home/stas/anaconda3/envs/main-38/lib/python3.8/site-packages/pytorch_lightning/accelerators/gpu_accelerator.py", line 95, in __validation_step
    output = self.trainer.model.validation_step(*args)
  File "finetune.py", line 183, in validation_step
    return self._generative_step(batch)
  File "finetune.py", line 216, in _generative_step
    generated_ids = self.model.generate(
  File "/home/stas/anaconda3/envs/main-38/lib/python3.8/site-packages/torch/autograd/grad_mode.py", line 26, in decorate_context
    return func(*args, **kwargs)
  File "/mnt/nvme1/code/huggingface/transformers-watchdog/src/transformers/generation_utils.py", line 553, in generate
    return self.beam_search(
  File "/mnt/nvme1/code/huggingface/transformers-watchdog/src/transformers/generation_utils.py", line 950, in beam_search
    beam_scores = torch.zeros((batch_size, num_beams), dtype=torch.float, device=input_ids.device)
UnboundLocalError: local variable 'torch' referenced before assignment
Exception ignored in: <function tqdm.__del__ at 0x7f1ec3812ee0>
Traceback (most recent call last):
  File "/home/stas/anaconda3/envs/main-38/lib/python3.8/site-packages/tqdm/std.py", line 1122, in __del__
  File "/home/stas/anaconda3/envs/main-38/lib/python3.8/site-packages/tqdm/std.py", line 1335, in close
  File "/home/stas/anaconda3/envs/main-38/lib/python3.8/site-packages/tqdm/std.py", line 1514, in display
  File "/home/stas/anaconda3/envs/main-38/lib/python3.8/site-packages/tqdm/std.py", line 1125, in __repr__
  File "/home/stas/anaconda3/envs/main-38/lib/python3.8/site-packages/tqdm/std.py", line 1475, in format_dict
TypeError: cannot unpack non-iterable NoneType object
```

this is on pt-nightly, the only pt version I can use rtx-3090 24G card.

**correction**: on pt-nightly it blows up with this error w/ or w/o cache clearing - it just doesn't work. I checked `torch` is imported at the point it says it's not defined.

**edit2**: some problem related to `--warmup 500` - I set it to `1` and the above failure is gone - something to deal with separately.

Let me see if I can figure out how to test this on pt16 - the problem is the 8GB card still can't fit even a single sample.

> It's also helpful to write in the following sequence:
> 
> ```
> torch.autocast_increment_nesting()
> print(torch.autocast_decrement_nesting())
> ```
> 
> to see how deeply we're nested in autocast contexts at that point in forward. It should print `1`.

I validated - it prints `1`.

for calling `torch.clear_autocast_cache()` how does a forward know it's in the autocast context?
```
if torch.autocast_enabled:
    torch.clear_autocast_cache()
```
I don't see a public API to check that [here](https://pytorch.org/docs/stable/_modules/torch/cuda/amp/autocast_mode.html#autocast)


`torch.is_autocast_enabled()`

Thank you, I see it now in the source code - should it be in the docs too? https://pytorch.org/docs/stable/amp.html

I'm not sure, I want people to use autocast through the context manager interface.  I guess it's useful for debugging.

Oh, for sure. I don't think any of the `transformers` models should be made autocast-aware - this is the job of the trainer.

Hmm, now I'm able to use the 24GB card so it's much easier to debug as I don't hit OOM all the time. Though I can't compare with apex, as I can't build it for cuda-11 - but I hope it should still be OK.

What it appears to be is the beam search (size=4) consumes some 20GB just to search for 1 sample. It calls forward about 100 times each time allocating about 200MB - never releasing memory. I suppose because with apex the model is much more lean it requires much less memory.

If I add @ptrblck's suggestion to clear the autocast cache in forward it now consumes only 2.5GB - that's 1/8th of the same with the cache.

W/o fp16 it consumes 5GB for the same beam_seach operation.

Let's summarize:

type         | memory
----------|-------
w/o fp16 | 5GB
w/ fp16 | 20GB
w/ fp16 + cache flush | 2.5GB

So definitely this is at least a huge part of the culprit. I think we are making progress. Much appreciation for your suggestion, @ptrblck~

So how do you recommend to proceed? Should PL have an option to clear the autocast cache?

Should `autocast` cache be made smarter and flush automatically if gpu ram is 90% full and be called to check when this happens ala `gc.collect()`-timing and not just in the context of its use? Since now it created an additional cache in to add to to `cuda.cache`.

If you follow along this is the command I use now: `PYTHONPATH="../../src" CUDA_VISIBLE_DEVICES=0 python finetune.py --learning_rate 3e-5 --gpus 1 --do_train --val_check_interval 1 --num_train_epochs 1 --freeze_encoder --freeze_embeds --data_dir cnn_dm --max_target_length 142 --val_max_target_length 142 --train_batch_size 1 --eval_batch_size 1 --gradient_accumulation_steps 1 --model_name_or_path sshleifer/student_cnn_12_6 --tokenizer_name facebook/bart-large --warmup_steps 1 --output_dir distilbart-cnn-12-6 --overwrite_output_dir --num_sanity_val_steps 0 --n_train 1 --n_val 1 --fp16`

Don't think this is a trainer issue, I've been able to replicate this OOM crash putting autocast directly into the forward call, wrapping the code found here with `torch.cuda.amp.autocast`:

https://github.com/huggingface/transformers/blob/master/src/transformers/modeling_bart.py#L419

and turning off `--fp16`. Haven't been able to investigate further into why memory isn't being freed in this block!

Thank you for looking into it, @SeanNaren!

> Don't think this is a trainer issue

If it's an issue of `autocast` cache then it is by definition a PL issue - since it's the one managing it - but let's see what @ptrblck and @mcarilli say about what's the correct approach of taking advantage of native amp w/ incurring unreasonable increased memory demands. 

Surely there must be a way to manage it efficiently, and if caching is bad for whatever reason (application specific?) - there must be a way to turn it off in first place, rather than wasting resources/time copying/deleting it repeatedly.

> Haven't been able to investigate further into why memory isn't being freed in this block!

@ptrblck and I already did - it's the `autocast` cache (at least a huge part of it). See https://github.com/huggingface/transformers/issues/8403#issuecomment-725562117

Thanks @stas00! Any idea why within that block of code I posteted above that the autocast cache is not being freed? Btw if you wrap that particular section with `with torch.cuda.amp.autocast(enabled=False)`, the code runs (which I assume just turns off autocast functionality for that region of code). 

@SeanNaren, if PL still invokes the `autocast` context then this sub-context it won't free the cache on its exit. It will only free it when the most outer context will exit. You will have to remove the `autocast` call in PL for this to work.

Updates so far:

It appears that `autocast` is not designed to handle a massive number of `forward` calls in a single context - it caches them all! If I understand it correctly, it has to be called as close as possible to the first `forward` call that actually needs the casting. Currently, in `finetune.py` we end up with PL calling `autocast` on `SummarizationModule` which has nothing to do with pytorch math (i.e. needs no casting), which then goes through a massive logic of generate/beam_search which again has nothing to do with math, and only when it hits `BartForConditionalGeneration`'s `forward` we need the `autocast`.

The problem is that `generate` (which already runs under `autocast` via PL - so caching is on) ends up calling  `BartForConditionalGeneration`'s `forward` 100s of times, and every such call in the debug example was adding 200MB - so in 100 calls for beam search of size 4 it accumulated 20GB - we have a big problem.

So a workaround suggested by @ptrblck is to call 
```
if torch.autocast_enabled:
    torch.clear_autocast_cache()
```
inside `BartForConditionalGeneration.forward`  - but ideally it should be called in a more generic way in `generation_utils`  `while` loop ([here](https://github.com/huggingface/transformers/blob/121c24efa4453e4e726b5f0b2cf7095b14b7e74e/src/transformers/generation_utils.py#L954)), so that it works for any `transformers` model:

```
        while cur_len < max_length:
            model_inputs = self.prepare_inputs_for_generation(input_ids, **model_kwargs)
            if torch.is_autocast_enabled():
                torch.clear_autocast_cache()
            outputs = self(**model_inputs, return_dict=True)
```
but this is clearly far from optimal as a lot of resources will be wasted on filling the cache and immediately emptying it.

(Perhaps there is a way to disable the cache completely, but I don't know about its performance implications).

So a more efficient solution would be to `autocast` here instead:
```
        while cur_len < max_length:
            model_inputs = self.prepare_inputs_for_generation(input_ids, **model_kwargs)
            if somehow_we_know_autocast_should_be_used():
                with autocast(): 
                    outputs = self(**model_inputs, return_dict=True)
            else:
                outputs = self(**model_inputs, return_dict=True)          
```

but we have two issues here:

1. we have a problem if someone invoked `autocast()` sooner - e.g. as PL or HF trainer do it now. As this will defeat the purpose and the cache will blow up again. There is no cache-per-context, but only a single cache, regardless of whether the contexts are stacked. The outer call defines the scope for the cache and it'll clear only on the exit of that call. So the `autocast` call above for all means and purposes is a no-op if `autocast` has already been called in earlier frames.

2. now we are mixing trainer logic with the middle-layer (`generate` is neither a trainer nor a model - it's in between, and `SummarizationModule` in `finetune.py` is definitely not a model, but more of a trainer). How would `generation_utils` know about `somehow_we_know_autocast_should_be_used()`?

At this point since the problem is better understood I invite @LysandreJik, @patrickvonplaten, @sgugger and others to chime in and suggest how to move forward.

I found at least part of the culprit or trigger of the leak - it's `@torch.no_grad()` used for `generate` https://github.com/huggingface/transformers/blob/eb3bd73ce35bfef56eeb722d697f2d39a06a8f8d/src/transformers/generation_utils.py#L281-L282

Here is a short script that reproduces the leakage. It removes all the generate/search logic and feeds the same random input_ids to `BartForConditionalGeneration` pre-trained model.

Please first run:
```
pip install ipyexperiments 
```
to get the memory tracing, but feel free to disable it if for some reason it's not working for you. (it should)

```
#!/usr/bin/env python

import os
import sys
import torch
os.environ["USE_TF"] = "0"
sys.path.insert(1, "src")

# !pip install ipyexperiments 
from ipyexperiments.utils.mem import gpu_mem_get_used_mbs, gpu_mem_get_used_no_cache_mbs

from transformers import BartForConditionalGeneration

device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')

model = BartForConditionalGeneration.from_pretrained('sshleifer/student_cnn_12_6').to(device)
model.eval()

vocab_size = 50264 # model.config.vocab_size
length = 10

AUTOCAST = False if "-f" in sys.argv else True
print(f"autocast: {AUTOCAST}")

class MemReport():
    def __init__(self, gc_collect=True):
        self.get_mem = gpu_mem_get_used_no_cache_mbs if gc_collect else gpu_mem_get_used_mbs
        self.cur = self.get_mem()
    def delta(self, id):
        peak = torch.cuda.memory_stats()["allocated_bytes.all.peak"]
        print(f"{id}: {gpu_mem_get_used_mbs()-self.cur}MB (peak {peak>>20}MB)")
        self.cur = self.get_mem()
        
mr = MemReport(gc_collect=False)

### reproducible code starts here ###

@torch.no_grad()
def logic():
    input_ids = torch.randint(vocab_size, (1,length)).to(device)
    mr.delta(0)
    for i in range(1,10):
        outputs = model(input_ids)
        mr.delta(i)

if AUTOCAST:
    with torch.cuda.amp.autocast():
        logic()
else:
    logic()
```

So if I run it with `-f` which disables `autocast`, I get:

```
./reproduce.py -f
autocast: False
0: 0MB (peak 1165MB)
1: 12MB (peak 1167MB)
2: 0MB (peak 1169MB)
3: 0MB (peak 1169MB)
4: 0MB (peak 1169MB)
5: 0MB (peak 1169MB)
6: 0MB (peak 1169MB)
7: 0MB (peak 1169MB)
8: 0MB (peak 1169MB)
9: 0MB (peak 1169MB)
```
no leak.

If however I remove `-f` and `autocast` gets enabled, we get:
```
./reproduce.py -f
autocast: True
0: 0MB (peak 1165MB)
1: 592MB (peak 1744MB)
2: 580MB (peak 2324MB)
3: 580MB (peak 2902MB)
4: 580MB (peak 3480MB)
5: 580MB (peak 4058MB)
6: 580MB (peak 4636MB)
7: 580MB (peak 5214MB)
8: 580MB (peak 5793MB)
9: 580MB (peak 6371MB)
```
the memory logger prints the delta for each `forward` call in the loop and the peak memory.

You can see that we are leaking 600Mb per forward call here.

If I comment out `@torch.no_grad()`, the total memory usage doubles but there is no leak:

```
autocast: True
0: 0MB (peak 1165MB)
1: 602MB (peak 1754MB)
2: 590MB (peak 2343MB)
3: 0MB (peak 2343MB)
4: 0MB (peak 2343MB)
5: 0MB (peak 2343MB)
6: 0MB (peak 2343MB)
7: 0MB (peak 2343MB)
8: 0MB (peak 2343MB)
9: 0MB (peak 2343MB)
```

I was using pycharm to debug this and to write a small script and boy it got me so delayed as it leaks gpu ram on its own, since it has to save all those variables on cuda, but I wasn't aware of it. Well, now I know not to do that. Luckily I had https://github.com/stas00/ipyexperiments handy to give me easy memory tracing.

Note I'm importing two gpu mem tracking functions - one of them clears cuda cache - but here it appears it's better not use that version. 

One other issue to look into is that what happens under `autocast` to weights that are deterministic (such as positional weights [SinusoidalPositionalEmbedding](https://github.com/huggingface/transformers/blob/24184e73c441397edd51e9068e0f49c0418d25ab/src/transformers/modeling_bart.py#L1340)) as these are set with `requires_grad = False`.

Seeing how the caching logic [works](https://github.com/pytorch/pytorch/blob/21f447ee2c6ebbd72b6c3608c4df17c74edd4784/aten/src/ATen/autocast_mode.cpp#L69-L71):
```
  bool can_try_cache = (to_type == at::kHalf && arg.scalar_type() == at::kFloat && arg.requires_grad() && arg.is_leaf());
```
the conversion of these to fp16 will not be cached if I read the code correctly. This probably belongs to a separate issue though.

I think for it to get the full attention of pytorch devs it's the best for us to continue this discussion at pytorch, so I opened the ticket there: https://github.com/pytorch/pytorch/issues/48049

Great, thanks a lot for your thorough investigation @stas00 !

a fix has been applied to pytorch-nightly https://github.com/pytorch/pytorch/pull/48696 which fixes 
https://github.com/pytorch/pytorch/issues/48049 and I verified with pytorch-nightly this issue to no longer leak memory under native amp.

Please note that this change is going to be available in pytorch-1.8 - so until then native amp and transformers aren't going to play well at times. Until then the solution is to use apex.

**edit**: good news it seems that pytorch-1.7.1 will have this fix too! https://github.com/pytorch/pytorch/issues/48049#issuecomment-742790722


> Thank you for validating that, @sshleifer. So it really has something to do with the native amp in pytorch.
> 
> Here is the summary of my experiments:
> 
> * pt15 + confa-force apex w/ `--fp16` works at the start - reduces memory consumption
> * pt16 + conda-forge apex w/ `--fp16  --amp_backend='apex'` works at the start too!
> 
> but both fail at the end with:
> 
> ```
>   File "python3.8/site-packages/apex/amp/_amp_state.py", line 32, in warn_or_err
>     raise RuntimeError(msg)
> RuntimeError: Found param model.model.shared.weight with type torch.cuda.HalfTensor, expected torch.cuda.FloatTensor.
> When using amp.initialize, you do not need to call .half() on your model
> ```
> 
> If I use `--fp16_opt_level O1`, the failure changes to:
> 
> ```
>   File "python3.8/site-packages/torch/optim/lr_scheduler.py", line 56, in with_counter
>     instance_ref = weakref.ref(method.__self__)
> AttributeError: 'function' object has no attribute '__self__'
> ```
> 
> both envs use the binary apex from `conda-forge`
> 
> * apex doesn't support cuda11 at the moment
> * rtx-30* cards don't support cuda<11 (and really cuda<11.1).
> * ubuntu-20.4 doesn't support cuda<11, since it dropped gcc7, so can't build apex from source even for cuda-10
> 
> Bottom line, apex is not a great option at the moment, but may work for some short term - need to sort out native amp.
> 
> I will poke at it with debugger today.

I encountered the same issue with "AttributeError: 'function' object has no attribute '__self__'  when I tried to use "--fp16_opt_level O1" with pt15 + nvidia apex. Any solution to get around this? Or how to get it to work with nvidia-apex without installing torch-nightly? 

Could you please post a new issue, @XiangLi1999 - this one already has too many comments - and you're asking a related but a totally different question. Thank you!

You can tag me on it, and we can discuss it there.