Repository: huggingface/trl

Issue:
`RichProgressCallback` would break model evaluation and prediction
Hi! It's awesome to have a CLI for `trl`.
However, there seems to be a problem with the newly introduced `ProgressCallback`. This issue affects both the evaluation and prediction stages.

To reproduce the issue, simply run
```
trl sft --model_name_or_path facebook/opt-125m --dataset_name imdb --output_dir opt-sft-imdb --evaluation_strategy steps --eval_steps 1
```
this would leads to the following error message:
```
Traceback (most recent call last):
  File "**************/lib/python3.11/site-packages/trl/commands/scripts/sft.py", line 148, in <module>
    trainer.train()
  File "**************/lib/python3.11/site-packages/trl/trainer/sft_trainer.py", line 360, in train
    output = super().train(*args, **kwargs)
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "**************/lib/python3.11/site-packages/transformers/trainer.py", line 1780, in train
    return inner_training_loop(
           ^^^^^^^^^^^^^^^^^^^^
  File "**************/lib/python3.11/site-packages/transformers/trainer.py", line 2193, in _inner_training_loop
    self._maybe_log_save_evaluate(tr_loss, grad_norm, model, trial, epoch, ignore_keys_for_eval)
  File "**************/lib/python3.11/site-packages/transformers/trainer.py", line 2577, in _maybe_log_save_evaluate
    metrics = self.evaluate(ignore_keys=ignore_keys_for_eval)
              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "**************/lib/python3.11/site-packages/transformers/trainer.py", line 3365, in evaluate
    output = eval_loop(
             ^^^^^^^^^^
  File "**************/lib/python3.11/site-packages/transformers/trainer.py", line 3586, in evaluation_loop
    self.control = self.callback_handler.on_prediction_step(args, self.state, self.control)
                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "**************/lib/python3.11/site-packages/transformers/trainer_callback.py", line 410, in on_prediction_step
    return self.call_event("on_prediction_step", args, state, control)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "**************/lib/python3.11/site-packages/transformers/trainer_callback.py", line 414, in call_event
    result = getattr(callback, event)(
             ^^^^^^^^^^^^^^^^^^^^^^^^^
  File "**************/lib/python3.11/site-packages/trl/trainer/utils.py", line 783, in on_prediction_step
    self.prediction_bar.update(self.prediction_task_id, advance=1, update=True)
  File "**************/lib/python3.11/site-packages/rich/progress.py", line 1425, in update
    task = self._tasks[task_id]
           ~~~~~~~~~~~^^^^^^^^^
KeyError: None
```


Discussion:
I’m working on a PR to address this issue.

For those with urgent needs, a simple workaround is to comment out this line.
https://github.com/huggingface/trl/blob/0ee349dcd43b0f4b3169449f16751c38ac4a609f/trl/commands/cli.py#L45

Resolving patch:
diff --git a/trl/trainer/utils.py b/trl/trainer/utils.py
--- a/trl/trainer/utils.py
+++ b/trl/trainer/utils.py
@@ -746,8 +746,8 @@ class RichProgressCallback(TrainerCallback):
     """
 
     def __init__(self):
-        self.training_bar = Progress()
-        self.prediction_bar = Progress()
+        self.training_bar = None
+        self.prediction_bar = None
 
         self.training_task_id = None
         self.prediction_task_id = None
@@ -755,8 +755,14 @@ def __init__(self):
         self.rich_group = None
         self.rich_console = None
 
+        self.training_status = None
+        self.current_step = None
+
     def on_train_begin(self, args, state, control, **kwargs):
         if state.is_world_process_zero:
+            self.training_bar = Progress()
+            self.prediction_bar = Progress()
+
             self.rich_console = Console()
 
             self.training_status = self.rich_console.status("Nothing to log yet ...")
@@ -764,9 +770,8 @@ def on_train_begin(self, args, state, control, **kwargs):
             self.rich_group = Live(Panel(Group(self.training_bar, self.prediction_bar, self.training_status)))
             self.rich_group.start()
 
-            # self.training_bar.start()
             self.training_task_id = self.training_bar.add_task("[blue]Training the model", total=state.max_steps)
-        self.current_step = 0
+            self.current_step = 0
 
     def on_step_end(self, args, state, control, **kwargs):
         if state.is_world_process_zero:
@@ -775,25 +780,23 @@ def on_step_end(self, args, state, control, **kwargs):
 
     def on_prediction_step(self, args, state, control, eval_dataloader=None, **kwargs):
         if state.is_world_process_zero and has_length(eval_dataloader):
-            if self.prediction_bar is None:
-                # self.prediction_bar.start()
+            if self.prediction_task_id is None:
                 self.prediction_task_id = self.prediction_bar.add_task(
-                    "[blue]Predicting on the evaluation dataset", total=state.max_steps
+                    "[blue]Predicting on the evaluation dataset", total=len(eval_dataloader)
                 )
             self.prediction_bar.update(self.prediction_task_id, advance=1, update=True)
 
     def on_evaluate(self, args, state, control, **kwargs):
         if state.is_world_process_zero:
-            if self.prediction_bar is not None:
-                self.prediction_bar.close()
-            self.prediction_bar = None
+            if self.prediction_task_id is not None:
+                self.prediction_bar.remove_task(self.prediction_task_id)
+                self.prediction_task_id = None
 
     def on_predict(self, args, state, control, **kwargs):
         if state.is_world_process_zero:
-            if self.prediction_bar is not None:
-                self.prediction_bar.stop()
+            if self.prediction_task_id is not None:
                 self.prediction_bar.remove_task(self.prediction_task_id)
-            self.prediction_bar = None
+                self.prediction_task_id = None
 
     def on_log(self, args, state, control, logs=None, **kwargs):
         if state.is_world_process_zero and self.training_bar is not None:
@@ -802,6 +805,13 @@ def on_log(self, args, state, control, logs=None, **kwargs):
 
     def on_train_end(self, args, state, control, **kwargs):
         if state.is_world_process_zero:
-            self.training_bar.stop()
             self.rich_group.stop()
+
             self.training_bar = None
+            self.prediction_bar = None
+            self.training_task_id = None
+            self.prediction_task_id = None
+            self.rich_group = None
+            self.rich_console = None
+            self.training_status = None
+            self.current_step = None


Test patch:
diff --git a/tests/test_rich_progress_callback.py b/tests/test_rich_progress_callback.py
new file mode 100644
--- /dev/null
+++ b/tests/test_rich_progress_callback.py
@@ -0,0 +1,53 @@
+import tempfile
+import unittest
+
+import torch
+import torch.nn as nn
+from datasets import Dataset
+from transformers import Trainer, TrainingArguments
+
+from trl.trainer.utils import RichProgressCallback
+
+
+class DummyModel(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.a = nn.Parameter(torch.tensor(1.0))
+
+    def forward(self, x):
+        return self.a * x
+
+
+class TestRichProgressCallback(unittest.TestCase):
+    @classmethod
+    def setUpClass(cls):
+        cls.dummy_model = DummyModel()
+        cls.dummy_train_dataset = Dataset.from_list([{"x": 1.0, "y": 2.0}] * 5)
+        cls.dummy_val_dataset = Dataset.from_list([{"x": 1.0, "y": 2.0}] * 101)
+
+    def test_rich_progress_callback_logging(self):
+        with tempfile.TemporaryDirectory() as tmp_dir:
+            training_args = TrainingArguments(
+                output_dir=tmp_dir,
+                per_device_eval_batch_size=2,
+                per_device_train_batch_size=2,
+                num_train_epochs=4,
+                evaluation_strategy="steps",
+                eval_steps=1,
+                logging_strategy="steps",
+                logging_steps=1,
+                save_strategy="no",
+                report_to="none",
+                disable_tqdm=True,
+            )
+            callbacks = [RichProgressCallback()]
+            trainer = Trainer(
+                model=self.dummy_model,
+                train_dataset=self.dummy_train_dataset,
+                eval_dataset=self.dummy_val_dataset,
+                args=training_args,
+                callbacks=callbacks,
+            )
+
+            trainer.train()
+            trainer.train()