Repository: huggingface/trl
Pull request: [SFT Trainer] precompute packed iterable into a dataset
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

Problem statement:
SFTTrainer training stops early?
It seems it just happen today. Has anyone met the error where the training stops early?  For example, it should have run 1000 steps but stops after 300 steps, and giving a output like:

`TrainOutput(global_step=100, training_loss=1.2991660499572755, metrics={'train_runtime': 855.9696, 'train_samples_per_second': 1.869, 'train_steps_per_second': 0.117, 'total_flos': 6.50352940548096e+16, 'train_loss': 1.2991660499572755, 'epoch': 0.16})`

I did my training on Colab, and i don't know if it is an internet issue. But my previous training never meets this kind of issue.

Hints from issue discussion:
seems caused by `packing = True`
Yes, when you pack it can happen that you end up with fewer samples if the average length of samples is shorter than the sequence length.
Should be fixed in #979.

Commit message: fix test

Patch to review:
diff --git a/tests/test_sft_trainer.py b/tests/test_sft_trainer.py
index 100aeed970..0fe8d1627a 100644
--- a/tests/test_sft_trainer.py
+++ b/tests/test_sft_trainer.py
@@ -178,9 +178,21 @@ def test_sft_trainer_uncorrect_data(self):
                 args=training_args,
                 train_dataset=self.dummy_dataset,
                 formatting_func=formatting_prompts_func,
+                max_seq_length=32,  # make sure there is at least 1 packed sequence
                 packing=True,
             )
 
+            with self.assertRaises(ValueError):
+                # This should not work because not enough data for one sample
+                _ = SFTTrainer(
+                    model=self.model,
+                    args=training_args,
+                    train_dataset=self.dummy_dataset,
+                    formatting_func=formatting_prompts_func,
+                    max_seq_length=1024,  # make sure there is NOT at least 1 packed sequence
+                    packing=True,
+                )
+
             # This should not work as well
             with self.assertRaises(ValueError):
                 _ = SFTTrainer(
@@ -191,7 +203,7 @@ def test_sft_trainer_uncorrect_data(self):
                     packing=False,
                 )
 
-            # but this shpuld work
+            # but this should work
             _ = SFTTrainer(
                 model=self.model,
                 args=training_args,
diff --git a/trl/trainer/sft_trainer.py b/trl/trainer/sft_trainer.py
index 87ba1d166c..d344cc1fcf 100644
--- a/trl/trainer/sft_trainer.py
+++ b/trl/trainer/sft_trainer.py
@@ -20,6 +20,8 @@
 import torch
 import torch.nn as nn
 from datasets import Dataset
+from datasets.arrow_writer import SchemaInferenceError
+from datasets.builder import DatasetGenerationError
 from transformers import (
     AutoModelForCausalLM,
     AutoTokenizer,
@@ -128,7 +130,7 @@ def __init__(
         packing: Optional[bool] = False,
         formatting_func: Optional[Callable] = None,
         max_seq_length: Optional[int] = None,
-        infinite: Optional[bool] = False,
+        infinite: Optional[bool] = None,
         num_of_sequences: Optional[int] = 1024,
         chars_per_token: Optional[float] = 3.6,
         dataset_num_proc: Optional[int] = None,
@@ -141,6 +143,11 @@ def __init__(
         elif not isinstance(model, str):
             raise ValueError("You passed model_kwargs to the SFTTrainer. But your model is already instantiated.")
 
+        if infinite is not None:
+            warnings.warn(
+                "The `infinite` argument is deprecated. Use `TrainingArgument.max_steps` or `TrainingArgument.num_epochs` instead to control training length."
+            )
+
         if isinstance(model, str):
             warnings.warn(
                 "You passed a model_id to the SFTTrainer. This will automatically create an "
@@ -226,7 +233,6 @@ def __init__(
                 dataset_text_field,
                 max_seq_length,
                 formatting_func,
-                infinite,
                 num_of_sequences,
                 chars_per_token,
             )
@@ -238,7 +244,6 @@ def __init__(
                 dataset_text_field,
                 max_seq_length,
                 formatting_func,
-                infinite,
                 num_of_sequences,
                 chars_per_token,
             )
@@ -301,7 +306,6 @@ def _prepare_dataset(
         dataset_text_field,
         max_seq_length,
         formatting_func,
-        infinite,
         num_of_sequences,
         chars_per_token,
     ):
@@ -317,30 +321,19 @@ def _prepare_dataset(
                 tokenizer, dataset, dataset_text_field, max_seq_length, formatting_func
             )
 
-        if dataset_text_field is not None or formatting_func is not None:
-            if tokenizer is None:
-                raise ValueError(
-                    "You need to pass a tokenizer when using the SFT Trainer when passing a `dataset_text_field`."
-                )
-
-            return ConstantLengthDataset(
+        else:
+            return self._prepare_packed_dataloader(
                 tokenizer,
                 dataset,
-                dataset_text_field=dataset_text_field,
-                formatting_func=formatting_func,
-                seq_length=max_seq_length,
-                infinite=infinite,
-                num_of_sequences=num_of_sequences,
-                chars_per_token=chars_per_token,
-                eos_token_id=tokenizer.eos_token_id,
+                dataset_text_field,
+                max_seq_length,
+                num_of_sequences,
+                chars_per_token,
+                formatting_func,
             )
 
-        raise ValueError(
-            "You need to pass a `dataset_text_field` or `formatting_func` argument to the SFTTrainer if you want to use the `ConstantLengthDataset`."
-        )
-
     def _prepare_non_packed_dataloader(
-        self, tokenizer, dataset, dataset_text_field, max_seq_len, formatting_func=None
+        self, tokenizer, dataset, dataset_text_field, max_seq_length, formatting_func=None
     ):
         use_formatting_func = formatting_func is not None and dataset_text_field is None
         self._dataset_sanity_checked = False
@@ -351,7 +344,7 @@ def tokenize(element):
                 element[dataset_text_field] if not use_formatting_func else formatting_func(element),
                 truncation=True,
                 padding=False,
-                max_length=max_seq_len,
+                max_length=max_seq_length,
                 return_overflowing_tokens=False,
                 return_length=False,
             )
@@ -376,6 +369,52 @@ def tokenize(element):
 
         return tokenized_dataset
 
+    def _prepare_packed_dataloader(
+        self,
+        tokenizer,
+        dataset,
+        dataset_text_field,
+        max_seq_length,
+        num_of_sequences,
+        chars_per_token,
+        formatting_func=None,
+    ):
+        if dataset_text_field is not None or formatting_func is not None:
+            if tokenizer is None:
+                raise ValueError(
+                    "You need to pass a tokenizer when using the SFT Trainer when passing a `dataset_text_field`."
+                )
+
+            constant_length_iterator = ConstantLengthDataset(
+                tokenizer,
+                dataset,
+                dataset_text_field=dataset_text_field,
+                formatting_func=formatting_func,
+                seq_length=max_seq_length,
+                infinite=False,
+                num_of_sequences=num_of_sequences,
+                chars_per_token=chars_per_token,
+                eos_token_id=tokenizer.eos_token_id,
+            )
+
+            def data_generator(constant_length_iterator):
+                for i in constant_length_iterator:
+                    yield i
+
+            try:
+                packed_dataset = Dataset.from_generator(
+                    data_generator, gen_kwargs={"constant_length_iterator": constant_length_iterator}
+                )
+            except (DatasetGenerationError, SchemaInferenceError):
+                raise ValueError(
+                    "Error occurred while packing the dataset. Make sure that your dataset has enough samples to at least yield one packed sequence."
+                )
+            return packed_dataset
+        else:
+            raise ValueError(
+                "You need to pass a `dataset_text_field` or `formatting_func` argument to the SFTTrainer if you want to use the `ConstantLengthDataset`."
+            )
+
     def _trl_activate_neftune(self, model):
         r"""
         Activates the neftune as presented in this code: https://github.com/neelsjain/NEFTune and paper: https://arxiv.org/abs/2310.05914


Code review comments:
diff_hunk: @@ -376,6 +369,52 @@ def tokenize(element):
 
         return tokenized_dataset
 
+    def _prepare_packed_dataloader(
+        self,
+        tokenizer,
+        dataset,
+        dataset_text_field,
+        max_seq_length,
+        num_of_sequences,
+        chars_per_token,
+        formatting_func=None,
+    ):
+        if dataset_text_field is not None or formatting_func is not None:
+            if tokenizer is None:
+                raise ValueError(
+                    "You need to pass a tokenizer when using the SFT Trainer when passing a `dataset_text_field`."
path: trl/trainer/sft_trainer.py
text: @user1:
nit:
```suggestion
                    "You need to pass a tokenizer when using `dataset_text_field` with `SFTTrainer`."
```

diff_hunk: @@ -141,6 +143,11 @@ def __init__(
         elif not isinstance(model, str):
             raise ValueError("You passed model_kwargs to the SFTTrainer. But your model is already instantiated.")
 
+        if infinite is not None:
+            warnings.warn(
+                "The `infinite` argument is deprecated. Use `TrainingArgument.max_steps` or `TrainingArgument.num_epochs` instead to control training length."
path: trl/trainer/sft_trainer.py
text: @user1:
turbo nit:
```suggestion
                "The `infinite` argument is deprecated and will be removed in a future version of TRL. Use `TrainingArguments.max_steps` or `TrainingArguments.num_train_epochs` instead to control training length."
```