🎲 [GRPO] Shuffle mini batches

[edbeeching]

This PR adds minibatch shuffling to ensure that the prompts are not ordered before the effective batch is split into chunks.

[HuggingFaceDocBuilderDev]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

[shirinyamani]

Im sure this shuffling prompt. still guarantees that prompts-completions are aligned, right?

[qgallouedec]

Looks good to me. To keep a detailed changelog, please either merge now to main (it should be possible), or wait for the other PR to be merged (but don't merge into the other branch)

[lewtun]

LGTM but would love to see shuffle_tensor_dict() unit tested since it's somewhat critical for training

[edbeeching]

LGTM but would love to see shuffle_tensor_dict() unit tested since it's somewhat critical for training

Sure, I will also add some tests for split tensor dict as I did not find any

[HaleyCH]

This PR doesn't work correctly when my inputs contain non-Tensor inputs. Specifically, I'm passing a List[Tuple[int, int]] in the inputs for interpolation because torch.nn.functional.interpolate only supports size inputs of type int or Tuple[int] or Tuple[int, int] or Tuple[int, int, int].
Converting the size in inputs to a Tensor and then converting it back feels a bit less elegant for my use case. Do you have any plans to support inputs containing non-Tensor elements?🥰

[qgallouedec]

You work with a fork on your specific use-case, right? because _prepare_input isn't part of the public API

[HaleyCH]

You work with a fork on your specific use-case, right? because _prepare_input isn't part of the public API

Thanks for your reply. Yes, for the purpose of training a LISA-like model architecture using GRPO, I inherited trl.trainer.GRPOTrainer and overrode some of its functions.

[rank0]: Traceback (most recent call last):
[rank0]:   File "/raid0/workspace/a/reasoning-cod/rcod/train/train_lisa_grpo.py", line 133, in <module>
[rank0]:     trainer.train(resume_from_checkpoint=False)
[rank0]:   File "/raid0/workspace/a/reasoning-cod/submodules/transformers/src/transformers/trainer.py", line 2245, in train
[rank0]:     return inner_training_loop(
[rank0]:            ^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/raid0/workspace/a/reasoning-cod/submodules/transformers/src/transformers/trainer.py", line 2556, in _inner_training_loop
[rank0]:     tr_loss_step = self.training_step(model, inputs, num_items_in_batch)
[rank0]:                    ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/raid0/workspace/a/reasoning-cod/submodules/transformers/src/transformers/trainer.py", line 3712, in training_step
[rank0]:     inputs = self._prepare_inputs(inputs)
[rank0]:              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/raid0/workspace/a/reasoning-cod/submodules/trl/trl/extras/profiling.py", line 96, in wrapper
[rank0]:     return func(self, *args, **kwargs)
[rank0]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/raid0/workspace/a/reasoning-cod/submodules/trl/trl/trainer/grpo_trainer.py", line 973, in _prepare_inputs
[rank0]:     generation_batch = shuffle_tensor_dict(generation_batch)
[rank0]:                        ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/raid0/workspace/a/reasoning-cod/submodules/trl/trl/trainer/grpo_trainer.py", line 247, in shuffle_tensor_dict
[rank0]:     return {key: tensor[permutation] if tensor is not None else None for key, tensor in tensor_dict.items()}
[rank0]:                  ~~~~~~^^^^^^^^^^^^^
[rank0]: TypeError: only integer tensors of a single element can be converted to an index

According to traceback，in trainer's training_step，_prepare_inputs is called by inputs = self._prepare_inputs(inputs), in which inputs contains None, Tensor and non-tensor items. Is this error caused by my use of inheritance?