Repository: huggingface/trl

Issue:
Error with official examples
Hi, I am running the official example script on t5 using sentiment as reward function, and after 97 iterations the following error is raised:

```
╭───────────────────── Traceback (most recent call last) ──────────────────────╮
│ /home/alejandro.vaca/trl/examples/sentiment/scripts/t5-sentiment.py:126 in   │
│ <module>                                                                     │
│                                                                              │
│   123 │   rewards = [torch.tensor(output[1]["score"]).to(device) for output  │
│   124 │                                                                      │
│   125 │   #### Run PPO step                                                  │
│ ❱ 126 │   stats = ppo_trainer.step(query_tensors, response_tensors, rewards) │
│   127 │   ppo_trainer.log_stats(stats, batch, rewards)                       │
│   128                                                                        │
│                                                                              │
│ /home/alejandro.vaca/trl/trl/trainer/ppo_trainer.py:432 in step              │
│                                                                              │
│   429 │   │   """                                                            │
│   430 │   │   bs = self.config.batch_size                                    │
│   431 │   │                                                                  │
│ ❱ 432 │   │   queries, responses, scores = self._step_safety_checker(bs, que │
│   433 │   │                                                                  │
│   434 │   │   timing = dict()                                                │
│   435 │   │   t0 = time.time()                                               │
│                                                                              │
│ /home/alejandro.vaca/trl/trl/trainer/ppo_trainer.py:392 in                   │
│ _step_safety_checker                                                         │
│                                                                              │
│   389 │   │   │   if not isinstance(tensor_list[0], torch.Tensor):           │
│   390 │   │   │   │   raise ValueError(f"Elements in {name} must tensors - g │
│   391 │   │   │   if batch_size is not None and len(tensor_list) != batch_si │
│ ❱ 392 │   │   │   │   raise ValueError(                                      │
│   393 │   │   │   │   │   f"Batch size ({batch_size}) does not match number  │
│   394 │   │   │   │   )                                                      │
│   395                                                                        │
╰──────────────────────────────────────────────────────────────────────────────╯
ValueError: Batch size (256) does not match number of examples - but got 63 for:
queries

```


Discussion:
hmm maybe we run out of data at this point. Maybe you could verify by just iterating over the dataset and not doing generation/training inside the loop. @younesbelkada maybe we can deal with this better inside TRL.
We could `drop_last=True` when we create the dataloader?
Agreed with @edbeeching ! We can also have a stronger safety checker, happy to have a look!

Resolving patch:
diff --git a/trl/trainer/ppo_trainer.py b/trl/trainer/ppo_trainer.py
--- a/trl/trainer/ppo_trainer.py
+++ b/trl/trainer/ppo_trainer.py
@@ -309,6 +309,7 @@ def prepare_dataloader(self, dataset: Union[torch.utils.data.Dataset, Dataset],
             batch_size=self.config.batch_size,
             collate_fn=data_collator,
             shuffle=True,
+            drop_last=True,
         )
         return dataloader
 


Test patch:
diff --git a/tests/trainer/test_ppo_trainer.py b/tests/trainer/test_ppo_trainer.py
--- a/tests/trainer/test_ppo_trainer.py
+++ b/tests/trainer/test_ppo_trainer.py
@@ -147,6 +147,22 @@ def _init_dummy_dataset(self):
 
         return dummy_dataset
 
+    def test_drop_last_dataloader(self):
+        self.ppo_config = PPOConfig(batch_size=3, forward_batch_size=1, log_with=None)
+
+        dummy_dataset = self._init_dummy_dataset()
+
+        ppo_trainer = PPOTrainer(
+            config=self.ppo_config,
+            model=self.gpt2_model,
+            ref_model=self.gpt2_model_ref,
+            tokenizer=self.gpt2_tokenizer,
+            dataset=dummy_dataset,
+        )
+        dummy_dataloader = ppo_trainer.dataloader
+
+        self.assertEqual(len(dummy_dataloader), 0)
+
     def test_ppo_step(self):
         # initialize dataset
         dummy_dataset = self._init_dummy_dataset()