OOM on summarization example

Hi there, I'm having OOM errors when running the summarization example on a 80GB A100 (CUDA 11.8).

I'm also getting some Tensorflow/TensorRT warnings, I'm wondering if it's related to that
```
2022-11-08 22:44:46.878785: I tensorflow/core/platform/cpu_feature_guard.cc:193] This TensorFlow binary is optimized with oneAPI Deep Neural Network Library (oneDNN) to use the following CPU instructions in performance-critical operations: AVX2 FMA
To enable them in other operations, rebuild TensorFlow with the appropriate compiler flags.
2022-11-08 22:44:47.016183: E tensorflow/stream_executor/cuda/cuda_blas.cc:2981] Unable to register cuBLAS factory: Attempting to register factory for plugin cuBLAS when one has already been registered
2022-11-08 22:44:47.979748: W tensorflow/stream_executor/platform/default/dso_loader.cc:64] Could not load dynamic library 'libnvinfer.so.7'; dlerror: libnvinfer.so.7: cannot open shared object file: No such file or directory; LD_LIBRARY_PATH: /usr/local/nvidia/lib:/usr/local/nvidia/lib64
2022-11-08 22:44:47.979824: W tensorflow/stream_executor/platform/default/dso_loader.cc:64] Could not load dynamic library 'libnvinfer_plugin.so.7'; dlerror: libnvinfer_plugin.so.7: cannot open shared object file: No such file or directory; LD_LIBRARY_PATH: /usr/local/nvidia/lib:/usr/local/nvidia/lib64
2022-11-08 22:44:47.979834: W tensorflow/compiler/tf2tensorrt/utils/py_utils.cc:38] TF-TRT Warning: Cannot dlopen some TensorRT libraries. If you would like to use Nvidia GPU with TensorRT, please make sure the missing libraries mentioned above are installed properly.
```

OOM error:
```
╭───────────────────── Traceback (most recent call last) ──────────────────────╮
│ │
│ /mnt/home/code/RL4LMs/scripts/training/train_text_generation.py:66 in │
│ │
│ │
│ 63 │ │ │ │ │ │ help="Whether to use wandb logging") │
│ 64 │ args = parser.parse_args() │
│ 65 │ │
│ ❱ 66 │ main(args.config_path, │
│ 67 │ │ args.project_name, │
│ 68 │ │ args.experiment_name, │
│ 69 │ │ args.base_path_to_store_results, │
│ /mnt/home/code/RL4LMs/scripts/training/train_text_generation.py:42 in main │
│ │
│ 39 │ │ │ │ │ │ │ │ on_policy_alg_config=config["alg"], │
│ 40 │ │ │ │ │ │ │ │ train_eval_config=config["train_evalu │
│ 41 │ │ │ │ │ │ │ │ tracker=tracker) │
│ ❱ 42 │ trainer.train_and_eval() │
│ 43 │
│ 44 │
│ 45 if __name__ == "__main__": │
│ │
│ /mnt/home/code/RL4LMs/rl4lms/envs/text_generation/training_utils.py:205 in │
│ train_and_eval │
│ │
│ 202 │ │ │ self._trainer_state["current_iter"] = epoch │
│ 203 │ │ │ │
│ 204 │ │ │ # inner rollout and learn loop for on-policy algorithm │
│ ❱ 205 │ │ │ self._alg.learn(self._n_steps_per_iter) │
│ 206 │ │ │ │
│ 207 │ │ │ # save the policy checkpoint │
│ 208 │ │ │ if (epoch + 1) % self._train_eval_config.get("save_every", │
│ │
│ /mnt/home/code/RL4LMs/rl4lms/algorithms/ppo/ppo.py:347 in learn │
│ │
│ 344 │ │ reset_num_timesteps: bool = True, │
│ 345 │ ) -> "PPO": │
│ 346 │ │ │
│ ❱ 347 │ │ return super().learn( │
│ 348 │ │ │ total_timesteps=total_timesteps, │
│ 349 │ │ │ callback=callback, │
│ 350 │ │ │ log_interval=log_interval, │
│ │
│ /mnt/home/miniconda3/lib/python3.9/site-packages/stable_baselines3/common/on │
│ _policy_algorithm.py:267 in learn │
│ │
│ 264 │ │ │ │ self.logger.record("time/total_timesteps", self.num_ti │
│ 265 │ │ │ │ self.logger.dump(step=self.num_timesteps) │
│ 266 │ │ │ │
│ ❱ 267 │ │ │ self.train() │
│ 268 │ │ │
│ 269 │ │ callback.on_training_end() │
│ 270 │
│ │
│ /mnt/home/code/RL4LMs/rl4lms/algorithms/ppo/ppo.py:224 in train │
│ │
│ 221 │ │ │ │ if self.use_sde: │
│ 222 │ │ │ │ │ self.policy.reset_noise(self.batch_size) │
│ 223 │ │ │ │ │
│ ❱ 224 │ │ │ │ values, log_prob, entropy = self.policy.evaluate_actio │
│ 225 │ │ │ │ │ rollout_data.observations, actions) │
│ 226 │ │ │ │ values = values.flatten() │
│ 227 │ │ │ │ # Normalize advantage │
│ │
│ /mnt/home/code/RL4LMs/rl4lms/envs/text_generation/policy.py:211 in │
│ evaluate_actions │
│ │
│ 208 │ │ │
│ 209 │ │ _, log_prob, entropy, _, _ = self.forward_policy(obs=obs, │
│ 210 │ │ │ │ │ │ │ │ │ │ │ │ │ │ actions=acti │
│ ❱ 211 │ │ values, _ = self.forward_value(obs) │
│ 212 │ │ │
│ 213 │ │ return values, log_prob, entropy │
│ 214 │
│ │
│ /mnt/home/code/RL4LMs/rl4lms/envs/text_generation/policy.py:447 in │
│ forward_value │
│ │
│ 444 │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │
│ 445 │ │ │
│ 446 │ │ # and forrward pass to get hidden states │
│ ❱ 447 │ │ outputs = self._value_model( │
│ 448 │ │ │ **model_inputs, │
│ 449 │ │ │ output_hidden_states=True, │
│ 450 │ │ │ decoder_attention_mask=decoder_attn_mask, │
│ │
│ /mnt/home/miniconda3/lib/python3.9/site-packages/torch/nn/modules/module.py: │
│ 1110 in _call_impl │
│ │
│ 1107 │ │ # this function, and just call forward. │
│ 1108 │ │ if not (self._backward_hooks or self._forward_hooks or self._ │
│ 1109 │ │ │ │ or _global_forward_hooks or _global_forward_pre_hooks │
│ ❱ 1110 │ │ │ return forward_call(*input, **kwargs) │
│ 1111 │ │ # Do not call functions when jit is used │
│ 1112 │ │ full_backward_hooks, non_full_backward_hooks = [], [] │
│ 1113 │ │ if self._backward_hooks or _global_backward_hooks: │
│ │
│ /mnt/home/miniconda3/lib/python3.9/site-packages/transformers/models/t5/mode │
│ ling_t5.py:1648 in forward │
│ │
│ 1645 │ │ │ │ decoder_attention_mask = decoder_attention_mask.to(se │
│ 1646 │ │ │
│ 1647 │ │ # Decode │
│ ❱ 1648 │ │ decoder_outputs = self.decoder( │
│ 1649 │ │ │ input_ids=decoder_input_ids, │
│ 1650 │ │ │ attention_mask=decoder_attention_mask, │
│ 1651 │ │ │ inputs_embeds=decoder_inputs_embeds, │
│ │
│ /mnt/home/miniconda3/lib/python3.9/site-packages/torch/nn/modules/module.py: │
│ 1110 in _call_impl │
│ │
│ 1107 │ │ # this function, and just call forward. │
│ 1108 │ │ if not (self._backward_hooks or self._forward_hooks or self._ │
│ 1109 │ │ │ │ or _global_forward_hooks or _global_forward_pre_hooks │
│ ❱ 1110 │ │ │ return forward_call(*input, **kwargs) │
│ 1111 │ │ # Do not call functions when jit is used │
│ 1112 │ │ full_backward_hooks, non_full_backward_hooks = [], [] │
│ 1113 │ │ if self._backward_hooks or _global_backward_hooks: │
│ │
│ /mnt/home/miniconda3/lib/python3.9/site-packages/transformers/models/t5/mode │
│ ling_t5.py:1040 in forward │
│ │
│ 1037 │ │ │ │ │ None, # past_key_value is always None with gradi │
│ 1038 │ │ │ │ ) │
│ 1039 │ │ │ else: │
│ ❱ 1040 │ │ │ │ layer_outputs = layer_module( │
│ 1041 │ │ │ │ │ hidden_states, │
│ 1042 │ │ │ │ │ attention_mask=extended_attention_mask, │
│ 1043 │ │ │ │ │ position_bias=position_bias, │
│ │
│ /mnt/home/miniconda3/lib/python3.9/site-packages/torch/nn/modules/module.py: │
│ 1110 in _call_impl │
│ │
│ 1107 │ │ # this function, and just call forward. │
│ 1108 │ │ if not (self._backward_hooks or self._forward_hooks or self._ │
│ 1109 │ │ │ │ or _global_forward_hooks or _global_forward_pre_hooks │
│ ❱ 1110 │ │ │ return forward_call(*input, **kwargs) │
│ 1111 │ │ # Do not call functions when jit is used │
│ 1112 │ │ full_backward_hooks, non_full_backward_hooks = [], [] │
│ 1113 │ │ if self._backward_hooks or _global_backward_hooks: │
│ │
│ /mnt/home/miniconda3/lib/python3.9/site-packages/transformers/models/t5/mode │
│ ling_t5.py:699 in forward │
│ │
│ 696 │ │ │ else: │
│ 697 │ │ │ │ query_length = None │
│ 698 │ │ │ │
│ ❱ 699 │ │ │ cross_attention_outputs = self.layer[1]( │
│ 700 │ │ │ │ hidden_states, │
│ 701 │ │ │ │ key_value_states=encoder_hidden_states, │
│ 702 │ │ │ │ attention_mask=encoder_attention_mask, │
│ │
│ /mnt/home/miniconda3/lib/python3.9/site-packages/torch/nn/modules/module.py: │
│ 1110 in _call_impl │
│ │
│ 1107 │ │ # this function, and just call forward. │
│ 1108 │ │ if not (self._backward_hooks or self._forward_hooks or self._ │
│ 1109 │ │ │ │ or _global_forward_hooks or _global_forward_pre_hooks │
│ ❱ 1110 │ │ │ return forward_call(*input, **kwargs) │
│ 1111 │ │ # Do not call functions when jit is used │
│ 1112 │ │ full_backward_hooks, non_full_backward_hooks = [], [] │
│ 1113 │ │ if self._backward_hooks or _global_backward_hooks: │
│ │
│ /mnt/home/miniconda3/lib/python3.9/site-packages/transformers/models/t5/mode │
│ ling_t5.py:613 in forward │
│ │
│ 610 │ │ output_attentions=False, │
│ 611 │ ): │
│ 612 │ │ normed_hidden_states = self.layer_norm(hidden_states) │
│ ❱ 613 │ │ attention_output = self.EncDecAttention( │
│ 614 │ │ │ normed_hidden_states, │
│ 615 │ │ │ mask=attention_mask, │
│ 616 │ │ │ key_value_states=key_value_states, │
│ │
│ /mnt/home/miniconda3/lib/python3.9/site-packages/torch/nn/modules/module.py: │
│ 1110 in _call_impl │
│ │
│ 1107 │ │ # this function, and just call forward. │
│ 1108 │ │ if not (self._backward_hooks or self._forward_hooks or self._ │
│ 1109 │ │ │ │ or _global_forward_hooks or _global_forward_pre_hooks │
│ ❱ 1110 │ │ │ return forward_call(*input, **kwargs) │
│ 1111 │ │ # Do not call functions when jit is used │
│ 1112 │ │ full_backward_hooks, non_full_backward_hooks = [], [] │
│ 1113 │ │ if self._backward_hooks or _global_backward_hooks: │
│ │
│ /mnt/home/miniconda3/lib/python3.9/site-packages/transformers/models/t5/mode │
│ ling_t5.py:509 in forward │
│ │
│ 506 │ │ ) │
│ 507 │ │ │
│ 508 │ │ # compute scores │
│ ❱ 509 │ │ scores = torch.matmul( │
│ 510 │ │ │ query_states, key_states.transpose(3, 2) │
│ 511 │ │ ) # equivalent of torch.einsum("bnqd,bnkd->bnqk", query_stat │
│ 512 │
╰──────────────────────────────────────────────────────────────────────────────╯
RuntimeError: CUDA out of memory. Tried to allocate 150.00 MiB (GPU 0; 79.35 GiB
total capacity; 76.08 GiB already allocated; 108.19 MiB free; 77.20 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
```

Any clues what's the issue? 80GB seems like a lot for just a T5-base model