Trace log probs on generation
Author: aced125
Pull request: yes
Labels: name: wontfix

This PR makes a **few code-line changes** to accomplish the following:

- We want to trace the log probabilities of tokens generated during generation, so that we can do policy gradient methods (e.g we want to improve ROUGE scores for summarization by using RL).

- This requires keeping track of the computation graph as well as the log probs.

- We remove the @torch.no_grad() decorator on the `generate` method in `modeling_utils.py`. We replace this with `torch.set_grad_enabled(False)` by default. At the end of the function, we do `torch.set_grad_enabled(True), to restore the original state.

- We use `torch.distributions.Categorical` to sample from the softmax. We can call `dist.sample()` and Torch will keep the gradients.

- We modify `top_k_top_p_filtering` slightly by adding `with torch.no_grad()` for parts of the code which unnecessarily trace the gradient.

## Tests

I have run the tests not including the slow ones and they all passed.

## Example:

```
tokenizer = AutoTokenizer.from_pretrained('distilgpt2')
model = AutoModelWithLMHead.from_pretrained('distilgpt2')   

outputs = model.generate(max_length=40, 
do_sample=True, 
trace_log_probs=True, 
eos_token_id=99999, 
num_beams=1, 
num_return_sequences=3
) 

tokens, log_probs = outputs
print(log_probs)
print(log_probs.shape)
print(tokens.shape)
```

We add error handling to disallow for configurations not supported:

- beam search not supported
```
outputs = model.generate(max_length=40, 
do_sample=True, 
trace_log_probs=True, 
eos_token_id=99999, 
num_beams=5, 
num_return_sequences=3
)  # throws an error
```

- trying to trace while not doing do_sample
```
outputs = model.generate(max_length=40, 
do_sample=False, 
trace_log_probs=True, 
eos_token_id=99999, 
num_beams=1, 
num_return_sequences=3
)  # throws an error
```



Comments:
Thanks for the PR @aced125 - could you run the slow tests as well and see whether they pass? 
I think checking these three tests should be good enough:

`RUN_SLOW=1 pytest tests/test_modeling_gpt2.py`
`RUN_SLOW=1 pytest tests/test_modeling_t5.py`
`RUN_SLOW=1 pytest tests/test_modeling_bart.py`

Yep @patrickvonplaten , done the above, all passed in a Google Colab notebook: https://colab.research.google.com/drive/12-WUburVlYHsrgKhMMt5MRXOPKbfaS3l

Hi @patrickvonplaten, wondering if you managed to take a look at this?

I'm wondering if there is any experimental results demonstrating that `trace_log_probs` is a helpful thing to have in the repo? 

Hi @sshleifer I'll be honest I haven't done any RL in the NLP domain (using transformers in the drug discovery domain) but I know people have tried to optimize ROUGE score for summarization and stuff like that in the past. I can try and maybe put something together for this though? 

I do think it is quite a useful feature to have in general though, will need it at some point IMO.

Cool, thanks for the transparency. From my seat, it would be preferable for you to experiment a bit on a branch to see what works before we merge this into master, as you suggest.

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.


> Hi @sshleifer I'll be honest I haven't done any RL in the NLP domain (using transformers in the drug discovery domain) but I know people have tried to optimize ROUGE score for summarization and stuff like that in the past. I can try and maybe put something together for this though?
> 
> I do think it is quite a useful feature to have in general though, will need it at some point IMO.

Hi, I have a similar use case where I require the log_probs from the `generate()` method. Did you find any solution for it? Was your PR merged?

+1 - I'd also find this feature useful

This feature would be useful for incorporating a popular technique called "unlikelihood_training".

See (https://github.com/facebookresearch/unlikelihood_training/blob/944465589c0fab534fe6d14a5db2850ddeee43ce/custom/gpt2/run_gpt2.py#L85) 

You have to sample from the model to produce negative candidates.

Once this feature is added; adding the unlikelihood loss becomes extremely easy and efficient.

I'd also want to get the gradients from the `generate` method!