Cannot fine tune my distilbart-cnn-12-6 model because of cuda memory
Author: Hildweig
Pull request: no
Labels: name: wontfix

I'm trying to fine tune my model like this:

```
import os
os.environ['PYTHONPATH'] += ":/content/transformers/examples"
%cd "/content/transformers/examples"

!python /content/transformers/examples/seq2seq/finetune.py \
    --learning_rate=3e-5 \
    --fp16 \
    --gpus 1 \
    --do_train \
    --do_predict \
    --n_val 1000 \
    --val_check_interval 0.1 \
    --sortish_sampler \
    --data_dir '/content/dataset' \
    --train_batch_size=1\
    --eval_batch_size=1\
    --output_dir=distilbart_multi_news \
    --num_train_epochs 1 \
    --model_name_or_path /content/model/best_tfmr
```

But even with a batch size of 1 I get this error:
 File "/content/transformers/examples/seq2seq/finetune.py", line 344, in <module>
    main(args)
  File "/content/transformers/examples/seq2seq/finetune.py", line 322, in main
    logger=logger,
  File "/content/transformers/examples/lightning_base.py", line 330, in generic_train
    trainer.fit(model)
  File "/usr/local/lib/python3.6/dist-packages/pytorch_lightning/trainer/trainer.py", line 918, in fit
    self.single_gpu_train(model)
  File "/usr/local/lib/python3.6/dist-packages/pytorch_lightning/trainer/distrib_parts.py", line 176, in single_gpu_train
    self.run_pretrain_routine(model)
  File "/usr/local/lib/python3.6/dist-packages/pytorch_lightning/trainer/trainer.py", line 1076, in run_pretrain_routine
    False)
  File "/usr/local/lib/python3.6/dist-packages/pytorch_lightning/trainer/evaluation_loop.py", line 279, in _evaluate
    output = self.evaluation_forward(model, batch, batch_idx, dataloader_idx, test_mode)
  File "/usr/local/lib/python3.6/dist-packages/pytorch_lightning/trainer/evaluation_loop.py", line 452, in evaluation_forward
    output = model.validation_step(*args)
  File "/content/transformers/examples/seq2seq/finetune.py", line 136, in validation_step
    return self._generative_step(batch)
  File "/content/transformers/examples/seq2seq/finetune.py", line 163, in _generative_step
    generated_ids = self.model.generate(input_ids=source_ids, attention_mask=source_mask, use_cache=True,)
  File "/usr/local/lib/python3.6/dist-packages/torch/autograd/grad_mode.py", line 15, in decorate_context
    return func(*args, **kwargs)
  File "/usr/local/lib/python3.6/dist-packages/transformers/generation_utils.py", line 459, in generate
    model_specific_kwargs=model_specific_kwargs,
  File "/usr/local/lib/python3.6/dist-packages/transformers/generation_utils.py", line 638, in _generate_beam_search
    outputs = self(**model_inputs)  # (batch_size * num_beams, cur_len, vocab_size)
  File "/usr/local/lib/python3.6/dist-packages/torch/nn/modules/module.py", line 722, in _call_impl
    result = self.forward(*input, **kwargs)
  File "/usr/local/lib/python3.6/dist-packages/transformers/modeling_bart.py", line 1005, in forward
    lm_logits = F.linear(outputs[0], self.model.shared.weight, bias=self.final_logits_bias)
  File "/usr/local/lib/python3.6/dist-packages/torch/nn/functional.py", line 1676, in linear
    output = input.matmul(weight.t())
RuntimeError: CUDA out of memory. Tried to allocate 100.00 MiB (GPU 0; 11.17 GiB total capacity; 10.59 GiB already allocated; 91.81 MiB free; 10.66 GiB reserved in total by PyTorch)

Any idea what to do? 


Comments:
Hi - it looks like the memory blowup occurs during beam search (presumably during eval).

In `_generative_step` in `fine_tune.py`,  feel free to play around with the `num_beams` parameter in the `self.model.generate` call.  I can confirm that explicitly setting `num_beams=1` works on a V100 16GB.  I believe the default is 5.  beam search is very memory & computation intensive.  Maybe for your final model evaluations, you can find a bigger GPU or just do on CPU.  Greedy decoding should be OK during model development.

Also, if you look at the available distilled BART models (https://huggingface.co/sshleifer/distilbart-xsum-12-1), you'll see some options with fewer params than distilbart-cnn-12-6 (i.e., distilbart-cnn-12-1)

please let me know if this works!

This issue has been automatically marked as stale because it has not had recent activity. It will be closed if no further activity occurs. Thank you for your contributions.