add minibatching

[lvwerra]

Until now the PPO mini batch size has been hardcoded to 1. This PR aims to change it by refactoring the forward/backward passing logic.

In summary this PR does the following things:

1. The batched_forward_pass returns new a mask which can be used to mask parts of the sequence to be ignored
2. enable mini-batching of PPO by creating a small dataloader with the mini_batch_size to sample from the current PPO batch
3. In the loss method we replace all operations affected by masked parts of the sequence with masked ones (masked_mean, masked_whiten)
4. remove compute_logits_vpred and use batched_forward_pass for everything
5. extend testing and refactor it (i don't think we need subfolders for the 3 test files we have)

W&B logs:

• GPT-2 (mini_batch_size=1)
• GPT-2 (mini_batch_size=4)
• T5 (mini_batch_size=1)
• T5 (mini_batch_size=4)

[HuggingFaceDocBuilderDev]

The documentation is not available anymore as the PR was closed or merged.

[younesbelkada]

Thanks a lot for this great addition! I left few comments and questions as a first pass!

[younesbelkada]

Suggested change

	collate_fn=collator,
	collate_fn=collator,
	drop_last=True,

Maybe we can add this to avoid some corner-cases such as the one described on a previous issue

[lvwerra]

Sounds good, let's also set a warning if that's the case so the user knows that a batch will be dropped.

[younesbelkada]

So this is the case where the last element has less instances than the mini_batch_size or the case a users put a batch_size that is smaller than mini_batch_size on the config? If it's the second case we can maybe add a warning on the config, if the first case since we have drop_last=True set here I don't think we'll face this case but I am not sure

[lvwerra]

It's for the case where mini_batch_size is smaller than forward_batch_size during the forward passes inside the minibatch loop. I am also not quite happy with how we do it actually.

[younesbelkada]

Also, what about completely removing forward_batch_size from the config? I don't think this is a breaking change as the configs cannot be pushed on the Hub, just need to update the examples accodingly. I believe this can be done on a follow up PR too

[lvwerra]

The breaking change actually also happens for users who currently use the library with forward_batch_size. What do you think about setting it default to None and overwrite mini_batch_size if it's set to another value with a warning that it affects now also the mini_batch_size if set to a value?

[younesbelkada]

This solution makes a lot of sense yes!

[review-notebook-app]

Check out this pull request on 

See visual diffs & provide feedback on Jupyter Notebooks.

---

Powered by ReviewNB

[lvwerra]

Deprecated forward_batch_size: feel free to have a look!

[younesbelkada]

Thanks a lot for your great work on this! 💯