[SFT Trainer] precompute packed iterable into a dataset

[lvwerra]

With this PR we'll precompute the packed iterable dataloader and push it into a dataset. That will make sure the epoch estimate is right and we have a working progress bar.

Fixes #1008 #1004

[HuggingFaceDocBuilderDev]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint.

[lewtun]

Great fix @lvwerra - very clean and elegant 🔥 !

Have you run the SFT example script to verify the training steps are correct & learning rate scheduler etc all look OK?

[lewtun]

It would be nice to have a code comment that explains why you need to create the packed dataset via from_generator and data_generator

[chujiezheng]

Same question here

[younesbelkada]

Looking great, thanks @lvwerra !
I tested this PR with the script below:

from transformers import AutoModelForCausalLM, TrainingArguments
from datasets import load_dataset
from trl import SFTTrainer

# This should get ~700 examples
dataset = load_dataset("imdb", split="train[:5%]")

model = AutoModelForCausalLM.from_pretrained("facebook/opt-125m")

args = TrainingArguments(
    output_dir="./",
    per_device_train_batch_size=2,
    num_train_epochs=1,
    logging_steps=1,
    lr_scheduler_type="cosine",
)

trainer = SFTTrainer(
    model,
    args=args,
    train_dataset=dataset,
    dataset_text_field="text",
    max_seq_length=512,
    packing=True,
)

trainer.train()

And it gives me correct LR curves: