Repository: huggingface/trl

Issue:
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.


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
@@ -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 and will be removed in a future version of TRL. Use `TrainingArguments.max_steps` or `TrainingArguments.num_train_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,50 @@ 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 `dataset_text_field` with `SFTTrainer`.")
+
+            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


Test patch:
diff --git a/tests/test_sft_trainer.py b/tests/test_sft_trainer.py
--- 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,