🧮 Fix max_steps calculation in RLOOTrainer (#2433)

huggingface · Dec 3, 2024 · 52201d3 · 52201d3
1 parent 9ff79a6
commit 52201d3
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/trl/trainer/rloo_trainer.py b/trl/trainer/rloo_trainer.py
@@ -279,7 +279,7 @@ def repeat_generator():
         # trainer state initialization
         self.state.global_step = 0
         self.state.episode = 0
-        self.state.max_steps = args.num_total_batches * args.num_mini_batches
+        self.state.max_steps = (args.num_total_batches * args.num_mini_batches) // 2
         self.state.num_train_epochs = args.total_episodes / self.train_dataset_len
         # Compute absolute values for logging, eval, and save if given as ratio
         if args.logging_steps is not None: