Make ConstantLengthDataset (or packing=True) shuffle examples before they are packed

[muupan]

What does this PR do?

To address #2030, this PR changes the behavior of ConstantLengthDataset when shuffle=True so that it shuffles the examples before they are packed. Shuffling after packing is kept, otherwise split examples occur in consecutive tensors.

Example code (same as #2030):

from datasets import Dataset
from trl.trainer import ConstantLengthDataset

# Dataset with "000", "111", ..., "777"
def gen():
    for i in range(8):
        yield {"text": f"{i}" * 3}
dataset = Dataset.from_generator(gen)

class FakeTokenizer:
    # Tokenizer that just converts "000" to [0, 0, 0], etc. EOS token is 8.
    def __init__(self):
        self.eos_token_id = 8
    
    def __call__(self, texts, **kwargs):
        return {"input_ids": [[int(x) for x in text] for text in texts]}

packed_dataset = ConstantLengthDataset(
    tokenizer=FakeTokenizer(),
    dataset=dataset,
    dataset_text_field="text",
    seq_length=7,
    infinite=False,
    chars_per_token=1,
    num_of_sequences=100,
    shuffle=True,
    append_concat_token=True,
    add_special_tokens=True,
)
print("First epoch")
for x in packed_dataset:
    print(x)
print("Second epoch")
for x in packed_dataset:
    print(x)

Output before this PR (trl==0.10.1):

First epoch
{'input_ids': tensor([0, 0, 0, 8, 1, 1, 1]), 'labels': tensor([0, 0, 0, 8, 1, 1, 1])}
{'input_ids': tensor([8, 2, 2, 2, 8, 3, 3]), 'labels': tensor([8, 2, 2, 2, 8, 3, 3])}
{'input_ids': tensor([5, 5, 8, 6, 6, 6, 8]), 'labels': tensor([5, 5, 8, 6, 6, 6, 8])}
{'input_ids': tensor([3, 8, 4, 4, 4, 8, 5]), 'labels': tensor([3, 8, 4, 4, 4, 8, 5])}
Second epoch
{'input_ids': tensor([3, 8, 4, 4, 4, 8, 5]), 'labels': tensor([3, 8, 4, 4, 4, 8, 5])}
{'input_ids': tensor([8, 2, 2, 2, 8, 3, 3]), 'labels': tensor([8, 2, 2, 2, 8, 3, 3])}
{'input_ids': tensor([0, 0, 0, 8, 1, 1, 1]), 'labels': tensor([0, 0, 0, 8, 1, 1, 1])}
{'input_ids': tensor([5, 5, 8, 6, 6, 6, 8]), 'labels': tensor([5, 5, 8, 6, 6, 6, 8])}

Output after this PR:

First epoch
{'input_ids': tensor([2, 2, 8, 5, 5, 5, 8]), 'labels': tensor([2, 2, 8, 5, 5, 5, 8])}
{'input_ids': tensor([1, 1, 1, 8, 7, 7, 7]), 'labels': tensor([1, 1, 1, 8, 7, 7, 7])}
{'input_ids': tensor([4, 8, 3, 3, 3, 8, 2]), 'labels': tensor([4, 8, 3, 3, 3, 8, 2])}
{'input_ids': tensor([8, 6, 6, 6, 8, 4, 4]), 'labels': tensor([8, 6, 6, 6, 8, 4, 4])}
Second epoch
{'input_ids': tensor([4, 4, 4, 8, 1, 1, 1]), 'labels': tensor([4, 4, 4, 8, 1, 1, 1])}
{'input_ids': tensor([2, 2, 8, 7, 7, 7, 8]), 'labels': tensor([2, 2, 8, 7, 7, 7, 8])}
{'input_ids': tensor([3, 8, 6, 6, 6, 8, 2]), 'labels': tensor([3, 8, 6, 6, 6, 8, 2])}
{'input_ids': tensor([8, 0, 0, 0, 8, 3, 3]), 'labels': tensor([8, 0, 0, 0, 8, 3, 3])}

Fixes #2030

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.
@qgallouedec

[lewtun]

Thanks for the clean solution @muupan - LGTM!

[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.

[muupan]

@lewtun Thanks for your approval! I have merged main via the "Update branch" button just now, and it still says "2 workflows awaiting approval". Is there anything I need to do to merge this PR?

[lewtun]

@lewtun Thanks for your approval! I have merged main via the "Update branch" button just now, and it still says "2 workflows awaiting approval". Is there anything I need to do to merge this PR?

I've just run the CI, so will merge if it's green :)