Allow streaming (datasets.IterableDataset)

[BramVanroy]

The motivation for this PR is given in this issue: #1455 (comment)

Currently using IterableDatasets in the SFTTrainer is not plausible because of two issues:

1. an IterableDataset is a subclass of a torch dataset, and therefore will not be considered when deciding whether to use packing

trl/trl/trainer/sft_trainer.py

Lines 404 to 405 in 1705aeb

	if isinstance(dataset, (torch.utils.data.IterableDataset, torch.utils.data.Dataset, ConstantLengthDataset)):
	return dataset

2. when packing, a given dataset is exhausted and loaded fully in-memory with Dataset.from_generator - which defeats the purpose of having an IterableDataset in the first place

trl/trl/trainer/sft_trainer.py

Lines 520 to 522 in 1705aeb

	packed_dataset = Dataset.from_generator(
	data_generator, gen_kwargs={"constant_length_iterator": constant_length_iterator}
	)

This PR remedies both those issues by explicitly checking whether the dataset is a datasets.IterableDataset in point 1, and also in point 2 in which case the packed ConstantLengthDataset is returned as-is (which is also an IterableDataset).

These changes should make the SFTTrainer better compatible with streaming datasets. The motivation to improve the situation is because in the alignment handbook we also use the SFTTrainer for continued pretraining where massive datasets (streamed) should be supported.

Note: seems that the epoch calculation does not happen correctly when max_steps is given though. With batch size 2, accum. steps 8 and optim steps 7153 with 16 gpus, I get this strange calculation for num examples and num epochs:

***** Running training *****
Num examples = 1,831,168
Num Epochs = 9,223,372,036,854,775,807
Instantaneous batch size per device = 2
Total train batch size (w. parallel, distributed & accumulation) = 256
Gradient Accumulation steps = 8
Total optimization steps = 7,153

closes #1455

[younesbelkada]

Thanks a lot @BramVanroy for the detailed work on this ! I think there should be no harm supporting this in SFTTrainer, can you run the styling checks? make precommit - then we can merge imo

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

[BramVanroy]

Thanks a lot @BramVanroy for the detailed work on this ! I think there should be no harm supporting this in SFTTrainer, can you run the styling checks? make precommit - then we can merge imo

Do you have any thoughts on the incorrect reporting of number of epochs?

[younesbelkada]

@BramVanroy hmmm not sure what we can do here to be honest :/

[BramVanroy]

@BramVanroy hmmm not sure what we can do here to be honest :/

Hm yeah, maybe let's keep it like that for now then.

[younesbelkada]

Thanks again !

[BramVanroy]

@younesbelkada I probably should have included it in this PR, too, but there is another useful (imo) change that is just three lines of code that could make the SFTTrainer even more flexible. I added it to a separate PR: #1520

[snow-kartikbagalore]

@BramVanroy @younesbelkada Hello guys, amazing work on this repo.
I wanted to know if it is possible to pass num_train_epochs now? Or is max_steps the recommended way to go about this for now?
I am working with a very large dataset, and am streaming the dataset by passing it as a ConstantLengthDataset.