Repository: huggingface/trl

Issue:
Support streaming + packing in SFTTrainer
Now that the alignment handbook has support for continued pretraining (using SFTTrainer under the hood), it'd be great to account for those cases where massive datasets are used. In such scenarios we want to use packing + streaming. However, in my experience that does not work well.


From my experiments and digging I have uncovered the following. In a typical scenario we first load a dataset with a given text column e.g. `text` and with streaming enabled through load_dataset. The result is a `datasets.IterableDataset`. This is passed to the SFTTrainer. However, the first thing goes wrong here when deciding how to prepare the dataset (packing/non-packing):

https://github.com/huggingface/trl/blob/1705aebebafea53ac7ad3509ce6e5002e0f36240/trl/trainer/sft_trainer.py#L404

This if-statement will evaluate to true because a `datasets.IterableDataset` is [also a `torch.IterableDataset`](https://github.com/huggingface/datasets/blob/ca8409a8bec4508255b9c3e808d0751eb1005260/src/datasets/iterable_dataset.py#L1187-L1188). As such tokenization/packing never happens for "stream" (IterableDataset) datasets and the dataset remains as-is: a dataset with just a text field.

The first suggestion therefore is to check explicitly that this is NOT a `datasets.IterableDataset`, so replace the if-statement by:

```python
    if isinstance(dataset, (torch.utils.data.IterableDataset, torch.utils.data.Dataset, ConstantLengthDataset)) and not isinstance(dataset, datasets.IterableDataset):
        return dataset
```

Unfortunately even after that change we run into undesired behavior. Because packing is enabled, we are preparing a constantlengthdataset. That's fine - it's iter method is a nice generator that will not by itself process the whole massive streaming dataset. However, it seems that that behavior IS triggered here, where the constantlengthdataset is exhausted into a new, regular Dataset (which is not an `IterableDataset` anymore!)

https://github.com/huggingface/trl/blob/1705aebebafea53ac7ad3509ce6e5002e0f36240/trl/trainer/sft_trainer.py#L520-L522

For small datasets that is not an issue, and I imagine that it is useful then to have all data preprocessed in-memory. However, for streaming datasets that is exactly what we do not want! Both in terms of memory and in terms of how long this preprocessing takes. This problem can be resolved by simply returning the constantlengthdataset itself, without transforming it into a regular Dataset.

```python
if isinstance(dataset, IterableDataset):
    return constant_length_iterator
```

If you agree I can make a PR for these two changes. However, it is also possible that I missed something that motivates the reason why the current implementation does not deal well with streaming.

---

To repro the original issue, you can install the alignment handbook ([from this branch](https://github.com/BramVanroy/alignment-handbook/tree/allow_streaming)), use the config file below

```yaml
# Model arguments
model_name_or_path: gpt2
model_revision: main
torch_dtype: bfloat16
use_flash_attention_2: false

# Data training arguments
dataset_mixer:
  uonlp/CulturaX: 1.0
dataset_configs:
  - nl
dataset_splits:
  - train
use_streaming: true
preprocessing_num_workers: 8

# SFT trainer config
bf16: true
do_eval: False
evaluation_strategy: "no"
gradient_accumulation_steps: 8
gradient_checkpointing: true
gradient_checkpointing_kwargs:
  use_reentrant: False
hub_model_id: gpt2-cpt-dutch
learning_rate: 1.0e-05
weight_decay: 0.01
log_level: info
logging_steps: 5
logging_strategy: steps
lr_scheduler_type: cosine
max_seq_length: 480
max_steps: 917130
output_dir: data/gpt2-cpt-dutch
overwrite_output_dir: true
per_device_train_batch_size: 4
remove_unused_columns: true
save_strategy: "steps"
save_steps: 100000
save_total_limit: 3
seed: 42
warmup_ratio: 0.01
```

and launch with

```bash
python scripts/run_cpt.py config.yaml
```


Discussion:

Resolving patch:
diff --git a/trl/trainer/sft_trainer.py b/trl/trainer/sft_trainer.py
--- a/trl/trainer/sft_trainer.py
+++ b/trl/trainer/sft_trainer.py
@@ -17,6 +17,7 @@
 from functools import wraps
 from typing import Callable, Dict, List, Optional, Tuple, Union
 
+import datasets
 import torch
 import torch.nn as nn
 from accelerate.state import PartialState
@@ -401,7 +402,9 @@ def _prepare_dataset(
             raise ValueError("The dataset should not be None")
 
         # check if torch dataset / dataloader and do nothing
-        if isinstance(dataset, (torch.utils.data.IterableDataset, torch.utils.data.Dataset, ConstantLengthDataset)):
+        if isinstance(
+            dataset, (torch.utils.data.IterableDataset, torch.utils.data.Dataset, ConstantLengthDataset)
+        ) and not isinstance(dataset, datasets.IterableDataset):
             return dataset
 
         if not packing:
@@ -513,6 +516,9 @@ def _prepare_packed_dataloader(
                 add_special_tokens=add_special_tokens,
             )
 
+            if isinstance(dataset, datasets.IterableDataset):
+                return constant_length_iterator
+
             def data_generator(constant_length_iterator):
                 yield from constant_length_iterator
 


Test patch: