Support iterable datasets in GRPO

[wilrop]

What does this PR do?

This PR solves the issue described in #3213. Additionally, it avoids the need for the PR in #3216.

I implemented support for an IterableDataset by overriding the get_train_dataloader and get_eval_dataloader methods from the Trainer class. Now, when GRPO is given an iterable dataset, the batch size is divided by self.num_generations and the data collator takes care of duplicating the samples afterwards.

Fixes #3213

Before submitting

• This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
• Did you read the contributor guideline,
  Pull Request section?
• Was this discussed/approved via a GitHub issue? Please add a link
  to it if that's the case.
• Did you make sure to update the documentation with your changes? Here are the
  documentation guidelines.
• Did you write any new necessary tests?

Who can review?

Anyone in the community is free to review the PR once the tests have passed. Feel free to tag
members/contributors who may be interested in your PR.

[wilrop]

One thing I still want to mention, but didn't know how to handle, is that we need to set dispatch_batches=False in the GRPOConfig when using an iterable dataset. If this is not set, we get an error when torch tries to concatenate strings (i.e. the prompts). I think this is probably not a GRPO specific problem though, see e.g. here huggingface/transformers#26548 (comment)

[qgallouedec]

Thanks for the PR. I don't really understand why we can't use the same approach as for the regular Dataset. Adapting the sampler should be enough, no?

[wilrop]

Thanks for the PR. I don't really understand why we can't use the same approach as for the regular Dataset. Adapting the sampler should be enough, no?

The problem is that torch's IterableDataset does not support custom samplers, see: https://github.com/pytorch/pytorch/blob/1017927c83dd95a4be6074c48e0fb38f0a1bd8f3/torch/utils/data/dataloader.py#L301 with the following comments

# Arg-check dataset related before checking samplers because we want to
# tell users that iterable-style datasets are incompatible with custom
# samplers first, so that they don't learn that this combo doesn't work
# after spending time fixing the custom sampler errors.

I thought that updating the data collator to do the duplication on the fly would be the easiest solution for an iterable dataset instead.

I apologise if your comment was referring to something else. Let me know!