🧺 [1/N] Refactor _generate in GRPO/RLOO: list of ints instead of tensors

[qgallouedec]

This PR belongs to a sequence of PR that aims to refactor the generation part of GRPO/RLOO to allow for easier customization and ultimately tool calling

Previous:

• 🟩 Drop image_split_sizes in favour of image_grid_thw #4111
• 📽 Multi image support for GRPO/RLOO #4113
• 😷 Refactor GRPO/RLOO to isolate _generate #4114

Next:

• 🧺 [2/N] Refactor _generate in GRPO/RLOO: Use prompt_ids from generation #4152
• 🧺 [3/N] Refactor _generate in GRPO/RLOO: Rely on generator for prompt truncation #4153
• 🧺 [4/N] Refactor _generate in GRPO/RLOO: Move forward_kwargs outside generation method #4154
• 🧺 [5/N] Refactor _generate in GRPO/RLOO: Insert images in the prompt #4155

The idea with this PR is to make _generate return list of ints instead of tensors. This will help a lots when implementing tool calling.

Several modifications:

1. truncate_with_protected_tokens: instead of operating on 2D tensors (ids and mask), it will operate on sequence ids directly:

   before

   >>> ids = torch.tensor([[1, 2, 3], [4, 5, 6]])
   >>> mask = torch.tensor([[1, 1, 1], [0, 0, 1]])
   >>> truncate_with_protected_tokens(ids, mask, target_length=2, protected_tokens=[])
   tensor([[2, 3], [5, 6]]), tensor([[1, 1], [0, 1]])

   after

   >>> ids = [[1, 2, 3], [4, 5, 6]]
   >>> [truncate_with_protected_tokens(seq, target_length=2, protected_tokens=[]) for seq in ids]
   [[2, 3], [5, 6]]

2. _generate now returns list of ids, instead of tensor + mask
   conversion to tensor is handle in _generate_and_score_completions

3. The generation part is moved to a function _generate_single_turn, which is called by _generate

[qgallouedec]

the function must now return a list of ints, so we must remove padding

[qgallouedec]

prompt and completion masks are later inferred from the sequence lengths

[albertvillanova]

Thanks.

[albertvillanova]

The sequences name in the docstring is not aligned with the ids name in the signature.

[albertvillanova]

Additionally, before it accepted batch_size sequences (within the tensor) and now it accepts a single sequence (list[int]). Isn't this breaking something? Some tests should be failing because of the new behavior?

[qgallouedec]

Some tests should be failing because of the new behavior?

yes, tests have been updated as well, see TruncateWithProtectedTokensTester

[qgallouedec]

The sequences name in the docstring is not aligned with the ids name in the signature.

thanks! fixed in c570fb0