fix #1325

Former-commit-id: 59f2cbbd52d4646fbd1ba83032bf522ecc49a50f

src/llmtuner/tuner/ppo/trainer.py

@@ -1,4 +1,5 @@
 import os
+import sys
 import math
 import torch
 from tqdm import tqdm
@@ -39,9 +40,6 @@ class CustomPPOTrainer(PPOTrainer, Trainer):
         **kwargs
     ):
         PPOTrainer.__init__(self, **kwargs)
-        if getattr(self.accelerator.state, "deepspeed_plugin", None) is not None:
-            raise ValueError("PPOTrainer is incompatible with DeepSpeed.")
-
         self.args = training_args
         self.model_args = model_args
         self.finetuning_args = finetuning_args
@@ -54,6 +52,8 @@ class CustomPPOTrainer(PPOTrainer, Trainer):
         self.control = TrainerControl()
         self.log_callback, self.save_callback = callbacks[0], callbacks[1]
         assert isinstance(self.log_callback, LogCallback) and isinstance(self.save_callback, SavePeftModelCallback)
+        if self.args.max_steps > 0:
+            logger.info("max_steps is given, it will override any value given in num_train_epochs")
 
     def ppo_train(self) -> None:
         r"""
@@ -62,10 +62,17 @@ class CustomPPOTrainer(PPOTrainer, Trainer):
         total_train_batch_size = (
             self.args.per_device_train_batch_size * self.args.gradient_accumulation_steps * self.args.world_size
         )
-        len_dataloader = len(self.dataloader)
-        num_examples = len(self.dataset)
-        num_train_epochs = self.args.num_train_epochs
-        max_steps = math.ceil(num_train_epochs * len_dataloader)
+        if self.args.max_steps > 0:
+            num_examples = total_train_batch_size * self.args.max_steps
+            num_train_epochs = sys.maxsize
+            max_steps = self.args.max_steps
+            steps_in_epoch = self.args.max_steps * self.args.gradient_accumulation_steps
+        else:
+            len_dataloader = len(self.dataloader)
+            num_examples = len(self.dataset)
+            num_train_epochs = self.args.num_train_epochs
+            max_steps = math.ceil(num_train_epochs * len_dataloader)
+            steps_in_epoch = len_dataloader
 
         self.state.max_steps = max_steps
         self.state.num_train_epochs = num_train_epochs
@@ -84,14 +91,16 @@ class CustomPPOTrainer(PPOTrainer, Trainer):
 
         unwrapped_model: "AutoModelForCausalLMWithValueHead" = self.accelerator.unwrap_model(self.model)
         dataiter = iter(self.dataloader)
-        steps_trained = 0
         loss_meter = AverageMeter()
         reward_meter = AverageMeter()
         self.log_callback.on_train_begin(self.args, self.state, self.control)
 
         for step in tqdm(range(max_steps), disable=not self.is_local_process_zero()):
-            batch = next(dataiter)
-            steps_trained += 1
+            try:
+                batch = next(dataiter)
+            except StopIteration:
+                dataiter = iter(self.dataloader)
+                batch = next(dataiter)
 
             # Cast to inference mode
             unwrapped_model.gradient_checkpointing_disable()
@@ -130,7 +139,7 @@ class CustomPPOTrainer(PPOTrainer, Trainer):
                     loss=round(loss_meter.avg, 4),
                     reward=round(reward_meter.avg, 4),
                     learning_rate=stats["ppo/learning_rate"],
-                    epoch=round(step / len_dataloader, 2)
+                    epoch=round(step / steps_in_epoch, 2)
                 )
                 tqdm.write(str(logs))
                 logs["step"] = step
@@ -150,10 +159,6 @@ class CustomPPOTrainer(PPOTrainer, Trainer):
             if self.control.should_epoch_stop or self.control.should_training_stop:
                 break
 
-            if steps_trained == len_dataloader:
-                dataiter = iter(self.dataloader)
-                steps_trained = 0
-
         self.log_callback.on_train_end(self.args, self.state, self.control)
         self.save_callback.on_train_end(
             self.args, self.state, self.control, model=self.accelerator.unwrap_model(self.model)