Fix PyTorch RAG tests GPU OOM
Author: ydshieh
Pull request: yes

# What does this PR do?

Fix PyTorch RAG tests GPU OOM.

The GPU OOM 
```
E     tensorflow.python.framework.errors_impl.ResourceExhaustedError: OOM when allocating tensor with shape[32,5,16,300,64] and type float on /job:localhost/replica:0/task:0/device:GPU:0 by allocator GPU_0_bfc [Op:GatherV2]
```
could be found in https://github.com/huggingface/transformers/runs/6100697349?check_suite_focus=true

## Results

- Without this PR, after the PyTorch RAG test, torch occupies about `9.5 GB` GPU memory. There are 10 TF RAG tests failed.
- With this PR, all PT/TF RAG tests pass without GPU OOM

Comments:
_The documentation is not available anymore as the PR was closed or merged._

Also cc @patil-suraj and @stas00 to see if they have suggestions

Good for me!

> What cache is this emptying since the model is not deleted? I think this is a symptom there is a memory leak in the model, which would need to be fixed.

From the documentations [torch.cuda.empty_cache](https://pytorch.org/docs/stable/generated/torch.cuda.empty_cache.html) and [Memory management](https://pytorch.org/docs/stable/notes/cuda.html#memory-management), in particular

```
PyTorch uses a caching memory allocator to speed up memory allocations.
This allows fast memory deallocation without device synchronizations.
However, the unused memory managed by the allocator will still show as if used in nvidia-smi. 
```
and 
```
Calling empty_cache() releases all unused cached memory from PyTorch so that those can be used by other GPU applications.
```

my understanding is: `PyTorch` will keep the allocated GPU memory for later use in order to avoid to reduce the number of times of memory allocation - the goal is to speed up some (memory) operations.

This doesn't mean those GPU memory are leaked - PyTorch still controls them. But for other applications (say `TensorFlow` or `nvidia-smi`), it means those GPU memory are not available. Use `empty_cache()` will release them.

Of course, the memory occupied by current available tensors won't be release.


if you're trying to use multiple programs concurrently accessing the same GPU (regardless if they all are pytorch, or mixed framework), `torch.cuda.empty_cache` is a definite help and is OK to use as long as it's inside the test only and not in `transformers`

But why when the pytorch test finishes the torch still has allocated tensors? Why not free them?

Often `import gc; gc.collect()` is needed to force garbage collection immediately. I'm not sure if this is the case.





@stas00 

My words in the previous comment might be a bit confusing: we don't have the issue of `the pytorch test finishes the torch still has allocated tensors`. I am just mentioning a general fact (which is quite trivial) that is mentioned in the PyTorch docs I linked.

`empty_cache()` helps, but not completely. There are still some GPU memory occupied even we leave the testing methods, but still in the same Python process (for example, entering the TF testing module which is launched together with the PT testing).

There are some discussions, like this one

https://discuss.pytorch.org/t/pytorch-do-not-clear-gpu-memory-when-return-to-another-function/125944/4

when you do `torch.ones(1)` it allocates 1-2GB of cuda kernels on the gpu and they remain allocated unless the program is shutdown.

In such a case the solution is not to run the program inside `pytest` but to use an external process. Once an external process finishes 100% of gpu memory is returned. (Except the tests are then much slower because it has to launch an external program)

I created a special framework for running external programs

https://github.com/huggingface/transformers/blob/6d90d76f5db344e333ec184cc6f414abe2aa6559/src/transformers/testing_utils.py#L1536

You can see it extensively used in the deepspeed and extended tests.

Yeah, I know this approach, but wasn't very sure how not use it in a good way with testing.
Maybe we can discuss this!

By the way: `torch.ones(1) it allocates 1-2GB of cuda kernels` --> I tried it and it seems a correct statement.

I am really surprised (and not very happy) that there is no way to free these kinds of memory allocation. 

Chances are is that there was no need for that until now and nobody asked for it. If I may propose you could create a feature request at pytorch asking for a feature that releases the cuda env completely. It's very possible that there is a C++ API to do that and it just wasn't made available in python.

The use case can be for example this exact situation, where the same program needs to alternate between different frameworks in the same run and needs to be able to access all of gpu's memory.

Does tf give the memory fully back when it's done and the process is still running?

If the goal is to recover as much memory as possible, shouldn't we delete the model before calling the `empty_cache` function?

There have been some requests in torch GH page without response

https://github.com/pytorch/pytorch/issues/28829 
(this one is on 2019/10)

Same situation for TF: not fully giving back GPU memory, and the requests are always without response

> If the goal is to recover as much memory as possible, shouldn't we delete the model before calling the empty_cache function?

@sgugger You are right! I tried it and just like you said.

Maybe I can just implement `tearDownModule()` which calls `empty_cache()`, so we don't need to `del models` + `empty_cache()` in all testing methods ..?

I am going to try this and see if how it goes.
(tried with toy examples, and works as expected)

Note that the `tearDown` is only called at the end of all tests of the module, so won't do the same thing you implemented (clean up at the end of each test).

Implement `tearDown` for `unittest.TestCase` subclass (and make sure to call its `super`) - this one will be called at the end of each test.

and before `empty_cache` it often helps to call `gc.collect()` to make it deterministic.

OK, I can do that.

But I feel that while we are in PyTorch test itself, we don't need to call `empty_cache()` --> because the occupied cache will be managed by `torch` and will be assigned to subsequential torch operations if they require GPU.

This `empty_cache()` is mainly for other applications to use `GPU`, like TF for example, in the same process.

And since the TF tests are in other modules, `tearDownModule()` in PT test module should be enough.

But again, I can go for `tearDown()`

Of course, we are discussing this particular test. I wasn't suggesting to do it to all tests.

The reason I suggested `gc.collect` before `empty_cache` is because when you free the model it's not guaranteed it'll be immediately freed due to how python's GC works. So if you want a reliable deterministic memory release inside a long running `pytest` process, `gc.collect` followed by `empty_cache` is how you make things deterministic.

@sgugger @stas00 

With the suggestions, all TF RAG tests pass now on GPU! 🔥 Thank you!

Unrelated to this PR, but since you work a lot with tests (thank you!), in case you're not aware of it, awhile ago I have developed:

https://github.com/huggingface/transformers/blob/6d90d76f5db344e333ec184cc6f414abe2aa6559/src/transformers/testing_utils.py#L998

which extends `unittest.TestCase` with various handy features - like automatic removal of temp dirs, accessors to file paths and many others. It's extensively documented in the module itself and also in https://huggingface.co/docs/transformers/testing

You don't need to do anything about it, other than perhaps I hope it'll save you time in the future.

Thank you, @stas00 . Maybe I can play with it, and at some point have a discussion with other team members to see if to use it by default!

And of course please feel free to extend it if there are other features that can be re-used.

Would like to have @sgugger and/or @LysandreJik opinion before merge :-) 

Merge now - we should have 13 test failure fewer (if this PR also works well on multiple GPUs too)