Big Model Inference: OOM with simple forward pass
Author: ggbetz
Pull request: no
Labels: name: bug

### System Info

```shell
- `transformers` version: 4.20.0
- Platform: Linux-5.11.0-40-generic-x86_64-with-glibc2.29
- Python version: 3.8.10
- Huggingface_hub version: 0.4.0
- PyTorch version (GPU?): 1.11.0+cu102 (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



+-----------------------------------------------------------------------------+
| NVIDIA-SMI 460.91.03    Driver Version: 460.91.03    CUDA Version: 11.2     |
|-------------------------------+----------------------+----------------------+
| GPU  Name        Persistence-M| Bus-Id        Disp.A | Volatile Uncorr. ECC |
| Fan  Temp  Perf  Pwr:Usage/Cap|         Memory-Usage | GPU-Util  Compute M. |
|                               |                      |               MIG M. |
|===============================+======================+======================|
|   0  GeForce RTX 208...  Off  | 00000000:01:00.0 Off |                  N/A |
| 30%   50C    P8    13W / 250W |      5MiB / 11019MiB |      0%      Default |
|                               |                      |                  N/A |
+-------------------------------+----------------------+----------------------+
|   1  GeForce RTX 208...  Off  | 00000000:02:00.0 Off |                  N/A |
| 29%   42C    P8    27W / 250W |      5MiB / 11019MiB |      0%      Default |
|                               |                      |                  N/A |
+-------------------------------+----------------------+----------------------+
                                                                               
+-----------------------------------------------------------------------------+
| Processes:                                                                  |
|  GPU   GI   CI        PID   Type   Process name                  GPU Memory |
|        ID   ID                                                   Usage      |
|=============================================================================|
|    0   N/A  N/A      1237      G   /usr/lib/xorg/Xorg                  4MiB |
|    1   N/A  N/A      1237      G   /usr/lib/xorg/Xorg                  4MiB |
+-----------------------------------------------------------------------------+
```


### Who can help?

@sgugger 

### Information

- [ ] The official example scripts
- [X] My own modified scripts

### Tasks

- [ ] An officially supported task in the `examples` folder (such as GLUE/SQuAD, ...)
- [X] My own task or dataset (give details below)

### Reproduction

1. Loading `T03B` for big model inference:

```python
$ python3
Python 3.8.10 (default, Nov 26 2021, 20:14:08) 
[GCC 9.3.0] on linux
Type "help", "copyright", "credits" or "license" for more information.
>>> from transformers import AutoModelForSeq2SeqLM, AutoTokenizer
>>> model = AutoModelForSeq2SeqLM.from_pretrained("bigscience/T0_3B", device_map="auto")
>>> tokenizer = AutoTokenizer.from_pretrained("bigscience/T0_3B")

```
2. The model is distributed over two GPUs:

```python
>>> model.hf_device_map
{'shared': 0, 'decoder': 0, 'encoder.embed_tokens': 0, 'encoder.block.0': 0, 'encoder.block.1': 0, 'encoder.block.2': 0, 'encoder.block.3': 0, 'encoder.block.4': 0, 'encoder.block.5': 0, 'encoder.block.6': 0, 'encoder.block.7': 0, 'encoder.block.8': 0, 'encoder.block.9': 0, 'encoder.block.10': 0, 'encoder.block.11': 0, 'encoder.block.12': 0, 'encoder.block.13': 0, 'encoder.block.14': 0, 'encoder.block.15': 0, 'encoder.block.16': 0, 'encoder.block.17': 0, 'encoder.block.18': 0, 'encoder.block.19': 1, 'encoder.block.20': 1, 'encoder.block.21': 1, 'encoder.block.22': 1, 'encoder.block.23': 1, 'encoder.final_layer_norm': 1, 'encoder.dropout': 1, 'lm_head': 1}
```
3. Generation works fine:
```python
>>> inputs = tokenizer("Task: copy but say the opposite. PSG won its match against Barca.", return_tensors="pt")
>>> inputs = inputs.to(0)
>>> tokenizer.decode(model.generate(inputs['input_ids'])[0].tolist())
'<pad> Paris St-Germain beat Barcelona 1-0 in their Champions League Group B match on Tuesday.
'
```

4. But a simple forward pass throws an OOM error:
```python
>>> inputs = tokenizer("Task: copy but say the opposite. PSG won its match against Barca.", return_tensors="pt")
>>> inputs = inputs.to(0)
>>> labels = tokenizer("PSG lost its match against Barca.", return_tensors="pt")
>>> labels = labels.to(0)
>>> output = model(**inputs, labels=labels['input_ids'])
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/home/###-py3.8/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1110, in _call_impl
    return forward_call(*input, **kwargs)
  File "/home/###-py3.8/lib/python3.8/site-packages/accelerate/hooks.py", line 148, in new_forward
    output = old_forward(*args, **kwargs)
  File "/home/###-py3.8/lib/python3.8/site-packages/transformers/models/t5/modeling_t5.py", line 1601, in forward
    encoder_outputs = self.encoder(
  File "/home/###-py3.8/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1110, in _call_impl
    return forward_call(*input, **kwargs)
  File "/home/###-py3.8/lib/python3.8/site-packages/transformers/models/t5/modeling_t5.py", line 1033, in forward
    layer_outputs = layer_module(
  File "/home/###-py3.8/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1110, in _call_impl
    return forward_call(*input, **kwargs)
  File "/home/###-py3.8/lib/python3.8/site-packages/accelerate/hooks.py", line 148, in new_forward
    output = old_forward(*args, **kwargs)
  File "/home/###-py3.8/lib/python3.8/site-packages/transformers/models/t5/modeling_t5.py", line 716, in forward
    hidden_states = self.layer[-1](hidden_states)
  File "/home/###-py3.8/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1110, in _call_impl
    return forward_call(*input, **kwargs)
  File "/home/###-py3.8/lib/python3.8/site-packages/accelerate/hooks.py", line 148, in new_forward
    output = old_forward(*args, **kwargs)
  File "/home/###-py3.8/lib/python3.8/site-packages/transformers/models/t5/modeling_t5.py", line 326, in forward
    forwarded_states = self.DenseReluDense(forwarded_states)
  File "/home/###-py3.8/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1110, in _call_impl
    return forward_call(*input, **kwargs)
  File "/home/###-py3.8/lib/python3.8/site-packages/accelerate/hooks.py", line 148, in new_forward
    output = old_forward(*args, **kwargs)
  File "/home/###-py3.8/lib/python3.8/site-packages/transformers/models/t5/modeling_t5.py", line 305, in forward
    hidden_gelu = self.act(self.wi_0(hidden_states))
  File "/home/###-py3.8/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1110, in _call_impl
    return forward_call(*input, **kwargs)
  File "/home/###-py3.8/lib/python3.8/site-packages/transformers/activations.py", line 34, in forward
    return 0.5 * input * (1.0 + torch.tanh(math.sqrt(2.0 / math.pi) * (input + 0.044715 * torch.pow(input, 3.0))))
RuntimeError: CUDA out of memory. Tried to allocate 2.00 MiB (GPU 0; 10.76 GiB total capacity; 9.72 GiB already allocated; 3.69 MiB free; 9.73 GiB reserved in total by PyTorch) If reserved memory is >> allocated memory try setting max_split_size_mb to avoid fragmentation.  See documentation for Memory Management and PYTORCH_CUDA_ALLOC_CONF
```



### Expected behavior

```shell
There should be no OOM error, especially so as generation works fine.
```


Comments:
You need to put your forward pass inside a `torch.no_grad()` context manager, otherwise you will get memory used for the activations saved for the backward pass (which is not supported anyway for big model inference), which is why you get OOM.

Thanks for the quick reply, @sgugger! With the context manager I get another error (which seems to relate to the fact the model resides on two GPUs?):

```python
# inputs and labels are on cuda:0
>>> inputs
{'input_ids': tensor([[16107,    10,  2405,    68,   497,     8,  6401,     5,   276,  9945,
           751,   165,  1588,   581,  1386,   658,     5,     1]],
       device='cuda:0'), 'attention_mask': tensor([[1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]],
       device='cuda:0')}
>>> labels
{'input_ids': tensor([[ 276, 9945, 1513,  165, 1588,  581, 1386,  658,    5,    1]],
       device='cuda:0'), 'attention_mask': tensor([[1, 1, 1, 1, 1, 1, 1, 1, 1, 1]], device='cuda:0')}

# forward:
>>> with torch.no_grad():
...   output = model(**inputs, labels=labels['input_ids'])
... 
Traceback (most recent call last):
  File "<stdin>", line 2, in <module>
  File "/home/###-py3.8/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1110, in _call_impl
    return forward_call(*input, **kwargs)
  File "/home/###-py3.8/lib/python3.8/site-packages/accelerate/hooks.py", line 148, in new_forward
    output = old_forward(*args, **kwargs)
  File "/home/###-py3.8/lib/python3.8/site-packages/transformers/models/t5/modeling_t5.py", line 1671, in forward
    loss = loss_fct(lm_logits.view(-1, lm_logits.size(-1)), labels.view(-1))
  File "/home/###-py3.8/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1110, in _call_impl
    return forward_call(*input, **kwargs)
  File "/home/###-py3.8/lib/python3.8/site-packages/torch/nn/modules/loss.py", line 1163, in forward
    return F.cross_entropy(input, target, weight=self.weight,
  File "/home/###-py3.8/lib/python3.8/site-packages/torch/nn/functional.py", line 2996, in cross_entropy
    return torch._C._nn.cross_entropy_loss(input, target, weight, _Reduction.get_enum(reduction), ignore_index, label_smoothing)
RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:1 and cuda:0! (when checking argument for argument target in method wrapper_nll_loss_forward)
```

Indeed, your labels will need to be on the same device as the last layer of the model if you want to compute the loss inside the model (here device 1).

Thanks so much, that makes sense -- and works:

```python
>>> labels = labels.to(1)
>>> with torch.no_grad():
...   output = model(**inputs, labels=labels['input_ids'])
>>> output[0].tolist()
0.9372970461845398
```

(As far as I'm concerned, this is solved.)