support full-parameter PPO

Former-commit-id: 4af967d69475e1c9fdf1a7983cd6b83bd431abff

src/llmtuner/train/ppo/trainer.py

@@ -37,24 +37,44 @@ class CustomPPOTrainer(PPOTrainer, Trainer):
         finetuning_args: "FinetuningArguments",
         generating_args: "GeneratingArguments",
         callbacks: List["TrainerCallback"],
+        reward_model: "AutoModelForCausalLMWithValueHead",
         **kwargs
     ):
         PPOTrainer.__init__(self, **kwargs)
+
         self.args = training_args
         self.model_args = model_args
         self.finetuning_args = finetuning_args
+
         self.generation_config = GenerationConfig(
             pad_token_id=self.tokenizer.pad_token_id,
             eos_token_id=[self.tokenizer.eos_token_id] + self.tokenizer.additional_special_tokens_ids,
             **generating_args.to_dict()
         )
+
         self.state = TrainerState()
         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")
 
+        if reward_model is not None:
+            is_deepspeed_enabled = self.accelerator.distributed_type == "DEEPSPEED" and hasattr(
+                self.accelerator.state, "deepspeed_plugin"
+            )
+            if is_deepspeed_enabled:
+                if not (
+                    getattr(reward_model.pretrained_model, "is_loaded_in_8bit", False)
+                    or getattr(reward_model.pretrained_model, "is_loaded_in_4bit", False)
+                ): # quantized models are already set on the correct device
+                    self.reward_model = self._prepare_deepspeed(self.reward_model)
+            else:
+                self.reward_model = self.accelerator.prepare_model(self.reward_model, evaluation_mode=True)
+        else:
+            self.reward_model = None
+
     def ppo_train(self) -> None:
         r"""
         Implements training loop for the PPO stage, like _inner_training_loop() in Huggingface's Trainer.
@@ -213,11 +233,14 @@ class CustomPPOTrainer(PPOTrainer, Trainer):
         r"""
         Computes scores using given reward model.
         """
-        replace_model(unwrapped_model, target="reward")
+        if self.reward_model is None:
+            replace_model(unwrapped_model, target="reward")
+
         batch = self.prepare_model_inputs(queries, responses)
 
         with torch.cuda.amp.autocast(dtype=self.model_args.compute_dtype): # support bf16
-            _, _, values = self.model(**batch, output_hidden_states=True, return_dict=True)
+            reward_model = self.reward_model if self.reward_model is not None else self.model
+            _, _, values = reward_model(**batch, output_hidden_states=True, return_dict=True)
 
         if values.size(0) != batch["input_ids"].size(0): # adapt to chatglm2
             values = torch.transpose(values, 0, 1)
@@ -228,7 +251,9 @@ class CustomPPOTrainer(PPOTrainer, Trainer):
             end_index = end_indexes[-1].item() if len(end_indexes) else 0
             rewards.append(values[i, end_index].float().detach().cpu()) # use fp32 type
 
-        replace_model(unwrapped_model, target="default")
+        if self.reward_model is None:
+            replace_model(unwrapped_model, target="default")
+
         return rewards
 
     @PPODecorators.empty_device_cache()

src/llmtuner/train/ppo/workflow.py

@@ -9,8 +9,9 @@ from transformers.optimization import get_scheduler
 
 from llmtuner.data import get_dataset, preprocess_dataset
 from llmtuner.extras.callbacks import SavePeftModelCallback
+from llmtuner.extras.logging import get_logger
 from llmtuner.extras.ploting import plot_loss
-from llmtuner.model import load_model_and_tokenizer
+from llmtuner.model import create_ref_model, create_reward_model, load_model_and_tokenizer
 from llmtuner.train.ppo.trainer import CustomPPOTrainer
 
 if TYPE_CHECKING:
@@ -18,6 +19,9 @@ if TYPE_CHECKING:
     from llmtuner.hparams import ModelArguments, DataArguments, FinetuningArguments, GeneratingArguments
 
 
+logger = get_logger(__name__)
+
+
 def run_ppo(
     model_args: "ModelArguments",
     data_args: "DataArguments",
@@ -33,6 +37,11 @@ def run_ppo(
     tokenizer.padding_side = "left" # use left-padding in generation while using right-padding in training
     data_collator = DataCollatorWithPadding(tokenizer=tokenizer)
 
+    # Create reference model and reward model
+    ref_model = create_ref_model(model_args, finetuning_args, stage="ppo")
+    reward_model = create_reward_model(model, model_args, finetuning_args)
+
+    # Create ppo config
     ppo_config = PPOConfig(
         model_name=model_args.model_name_or_path,
         learning_rate=training_args.learning_rate,
@@ -47,9 +56,11 @@ def run_ppo(
         log_with=finetuning_args.ppo_logger,
         use_score_scaling=finetuning_args.ppo_score_norm,
         use_score_norm=finetuning_args.ppo_score_norm,
+        whiten_rewards=finetuning_args.ppo_whiten_rewards,
         accelerator_kwargs={"step_scheduler_with_optimizer": False}
     )
 
+    # Create optimizer and scheduler
     optimizer = AdamW(filter(lambda p: p.requires_grad, model.parameters()), lr=training_args.learning_rate)
     if training_args.max_steps > 0:
         num_training_steps = training_args.max_steps
@@ -73,9 +84,10 @@ def run_ppo(
         finetuning_args=finetuning_args,
         generating_args=generating_args,
         callbacks=callbacks + [SavePeftModelCallback()],
+        reward_model=reward_model,
         config=ppo_config,
         model=model,
-        ref_model=None,
+        ref_model=ref_model,
         tokenizer=tokenizer,
         dataset=dataset,
         data_collator=data_collator,
@@ -88,5 +100,5 @@ def run_ppo(
         ppo_trainer.ppo_train()
         ppo_trainer.save_model()
         ppo_trainer.save_state() # must be called after save_model to have a folder
-        if ppo_trainer.is_world_process_zero() and model_args.plot_loss:
+        if ppo_trainer.is_world_process_zero() and finetuning_args.plot_loss:
             plot_loss(training_args.output_dir, keys=["loss", "reward"])