👮 Deprecate policy in favor of model in PPOTrainer (huggingface#2386)

Author: qgallouedec

docs/source/detoxifying_a_lm.mdx

@@ -105,8 +105,8 @@ and the optimizer will take care of computing the gradients in `bfloat16` precis
 </div>
 
 ```python
-ref_policy = create_reference_model(model, num_shared_layers=6)
-trainer = PPOTrainer(..., ref_policy=ref_policy)
+ref_model = create_reference_model(model, num_shared_layers=6)
+trainer = PPOTrainer(..., ref_model=ref_model)
 ```
 
 In the example above this means that the model have the 4 first layers frozen (i.e. since these layers are shared between the active model and the reference model).

examples/scripts/ppo/ppo.py

@@ -154,8 +154,8 @@ def tokenize(element):
     trainer = PPOTrainer(
         args=training_args,
         processing_class=tokenizer,
-        policy=policy,
-        ref_policy=ref_policy,
+        model=policy,
+        ref_model=ref_policy,
         reward_model=reward_model,
         value_model=value_model,
         train_dataset=train_dataset,

examples/scripts/ppo/ppo_tldr.py

@@ -165,8 +165,8 @@ def tokenize(element):
     trainer = PPOTrainer(
         args=training_args,
         processing_class=tokenizer,
-        policy=policy,
-        ref_policy=ref_policy,
+        model=policy,
+        ref_model=ref_policy,
         reward_model=reward_model,
         value_model=value_model,
         train_dataset=train_dataset,

trl/trainer/ppo_trainer.py

@@ -88,9 +88,7 @@ def __init__(self, policy, value_model) -> None:
         self.critic_backbone = getattr(value_model, value_model.base_model_prefix)
 
     def forward(self, **kwargs):
-        output = self.critic_backbone(
-            **kwargs,
-        )
+        output = self.critic_backbone(**kwargs)
         logits = self.value_model.score(output.hidden_states[-1])
         return self.policy(**kwargs), logits
 
@@ -100,14 +98,18 @@ class PPOTrainer(Trainer):
 
     @deprecate_kwarg("config", new_name="args", version="0.15.0", raise_if_both_names=True)
     @deprecate_kwarg("tokenizer", new_name="processing_class", version="0.15.0", raise_if_both_names=True)
+    @deprecate_kwarg("policy", "0.15.0", "model", warn_if_greater_or_equal_version=True, raise_if_both_names=True)
+    @deprecate_kwarg(
+        "ref_policy", "0.15.0", "ref_model", warn_if_greater_or_equal_version=True, raise_if_both_names=True
+    )
     def __init__(
         self,
         args: PPOConfig,
         processing_class: Optional[
             Union[PreTrainedTokenizerBase, BaseImageProcessor, FeatureExtractionMixin, ProcessorMixin]
         ],
-        policy: nn.Module,
-        ref_policy: Optional[nn.Module],
+        model: nn.Module,
+        ref_model: Optional[nn.Module],
         reward_model: nn.Module,
         train_dataset: Dataset,
         value_model: Optional[nn.Module] = None,
@@ -118,24 +120,24 @@ def __init__(
         callbacks: Optional[List[TrainerCallback]] = None,
         peft_config: Optional["PeftConfig"] = None,
     ) -> None:
-        if ref_policy is policy:
+        if ref_model is model:
             raise ValueError(
-                "`policy` and `ref_policy` cannot be the same object. If you want `ref_policy` to be the "
-                "same as `policy`, you must make a copy of it, or `None` if you use peft."
+                "`model` and `ref_model` cannot be the same object. If you want `ref_model` to be the "
+                "same as `model`, you must make a copy of it, or `None` if you use peft."
             )
 
         self.args = args
         self.processing_class = processing_class
-        self.policy = policy
+        self.model = model
 
         # Define the collator if not provided
         if data_collator is None:
             data_collator = DataCollatorWithPadding(self.processing_class)
 
-        self.policy.generation_config.eos_token_id = (
+        self.model.generation_config.eos_token_id = (
             None  # disable `pad_token_id` and `eos_token_id` because we just want to
         )
-        self.policy.generation_config.pad_token_id = None  # generate tokens without truncation / padding
+        self.model.generation_config.pad_token_id = None  # generate tokens without truncation / padding
 
         # peft support
         if not is_peft_available() and peft_config is not None:
@@ -144,24 +146,24 @@ def __init__(
             )
         elif is_peft_available() and peft_config is not None:
             # if model is a peft model and we have a peft_confg, we merge and unload it first
-            if isinstance(self.policy, PeftModel):
-                self.policy = self.policy.merge_and_unload()
+            if isinstance(self.model, PeftModel):
+                self.model = self.model.merge_and_unload()
 
             # get peft model with the given config
-            self.policy = get_peft_model(self.policy, peft_config)
-            if args.bf16 and getattr(self.policy, "is_loaded_in_4bit", False):
-                peft_module_casting_to_bf16(self.policy)
+            self.model = get_peft_model(self.model, peft_config)
+            if args.bf16 and getattr(self.model, "is_loaded_in_4bit", False):
+                peft_module_casting_to_bf16(self.model)
 
-        self.is_peft_model = is_peft_available() and isinstance(self.policy, PeftModel)
+        self.is_peft_model = is_peft_available() and isinstance(self.model, PeftModel)
         self.model_adapter_name = args.model_adapter_name
         self.ref_adapter_name = args.ref_adapter_name
 
-        if ref_policy:
-            self.ref_policy = ref_policy
+        if ref_model:
+            self.ref_model = ref_model
         elif self.is_peft_model:
-            self.ref_policy = None
+            self.ref_model = None
         else:
-            self.ref_policy = create_reference_model(self.policy)
+            self.ref_model = create_reference_model(self.model)
 
         self.reward_model = reward_model
         self.train_dataset = train_dataset
@@ -211,13 +213,13 @@ def __init__(
         #########
         # setup model, optimizer, and others
         #########
-        for module in [self.policy, self.ref_policy, self.value_model, self.reward_model]:
+        for module in [self.model, self.ref_model, self.value_model, self.reward_model]:
             if module is not None:
                 disable_dropout_in_model(module)
         if args.stop_token and args.stop_token == "eos":
             args.stop_token_id = processing_class.eos_token_id
-        self.model = PolicyAndValueWrapper(self.policy, self.value_model)
-        self.model.config = self.policy.config  # needed for pushing to hub
+        self.policy_and_value = PolicyAndValueWrapper(self.model, self.value_model)
+        self.policy_and_value.config = self.model.config  # needed for pushing to hub
         self.create_optimizer_and_scheduler(
             num_training_steps=args.num_total_batches
         )  # note that we are calling `self.lr_scheduler.step()` manually only at the batch level
@@ -228,7 +230,7 @@ def __init__(
         default_callbacks = DEFAULT_CALLBACKS + get_reporting_integration_callbacks(self.args.report_to)
         self.callbacks = default_callbacks if callbacks is None else default_callbacks + callbacks
         self.callback_handler = CallbackHandler(
-            self.callbacks, self.model, self.processing_class, self.optimizer, self.lr_scheduler
+            self.callbacks, self.policy_and_value, self.processing_class, self.optimizer, self.lr_scheduler
         )
         self.add_callback(PrinterCallback if self.args.disable_tqdm else DEFAULT_PROGRESS_CALLBACK)
         self.control = TrainerControl()
@@ -251,8 +253,8 @@ def __init__(
             os.makedirs(self.args.output_dir, exist_ok=True)
 
         # Add tags for models that have been loaded with the correct transformers version
-        if hasattr(self.model, "add_model_tags"):
-            self.model.add_model_tags(self._tag_names)
+        if hasattr(self.policy_and_value, "add_model_tags"):
+            self.policy_and_value.add_model_tags(self._tag_names)
 
         #########
         ### setup dataloader
@@ -267,7 +269,9 @@ def __init__(
         # sync random states for DataLoader(shuffle=True) before `accelerator.prepare`
         # see https://gist.github.com/vwxyzjn/2581bff1e48e185e0b85b6dfe1def79c
         torch.manual_seed(args.seed)
-        self.model, self.optimizer, self.dataloader = accelerator.prepare(self.model, self.optimizer, self.dataloader)
+        self.policy_and_value, self.optimizer, self.dataloader = accelerator.prepare(
+            self.policy_and_value, self.optimizer, self.dataloader
+        )
         torch.manual_seed(self.local_seed)  # reset the local seed again
 
         self.eval_dataloader = DataLoader(
@@ -283,19 +287,19 @@ def __init__(
                 self.reward_model, args.per_device_train_batch_size, args.fp16, args.bf16
             )
 
-            if self.ref_policy is None:
+            if self.ref_model is None:
                 if not self.is_peft_model:
                     raise ValueError("No reference model and model is not a Peft model.")
             else:
-                self.ref_policy = prepare_deepspeed(
-                    self.ref_policy, args.per_device_train_batch_size, args.fp16, args.bf16
+                self.ref_model = prepare_deepspeed(
+                    self.ref_model, args.per_device_train_batch_size, args.fp16, args.bf16
                 )
         else:
-            if self.ref_policy is None:
+            if self.ref_model is None:
                 if not self.is_peft_model:
                     raise ValueError("No reference model and model is not a Peft model.")
             else:
-                self.ref_policy = self.ref_policy.to(self.accelerator.device)
+                self.ref_model = self.ref_model.to(self.accelerator.device)
             self.reward_model = self.reward_model.to(self.accelerator.device)
 
     def get_train_dataloader(self) -> DataLoader:
@@ -308,25 +312,25 @@ def get_eval_dataloader(self) -> DataLoader:
     def null_ref_context(self):
         """Context manager for handling null reference model (that is, peft adapter manipulation)."""
         with self.accelerator.unwrap_model(
-            self.model.policy
+            self.policy_and_value.policy
         ).disable_adapter() if self.is_peft_model and not self.ref_adapter_name else nullcontext():
             if self.ref_adapter_name:
-                self.model.policy.set_adapter(self.ref_adapter_name)
+                self.policy_and_value.policy.set_adapter(self.ref_adapter_name)
             yield
             if self.ref_adapter_name:
-                self.model.policy.set_adapter(self.model_adapter_name or "default")
+                self.policy_and_value.policy.set_adapter(self.model_adapter_name or "default")
 
     def save_model(self, output_dir: Optional[str] = None, _internal_call: bool = False):
-        backup_model = self.model
-        self.model = self.model.policy  # save only the policy
+        backup_model = self.policy_and_value
+        self.policy_and_value = self.policy_and_value.policy  # save only the policy
 
         if self.is_deepspeed_enabled:
             backup_deepspeed = self.deepspeed
-            self.deepspeed = self.model
+            self.deepspeed = self.policy_and_value
 
         super().save_model(output_dir, _internal_call)
 
-        self.model = backup_model
+        self.policy_and_value = backup_model
 
         if self.is_deepspeed_enabled:
             self.deepspeed = backup_deepspeed
@@ -335,8 +339,8 @@ def train(self):
         args = self.args
         accelerator = self.accelerator
         optimizer = self.optimizer
-        model = self.model
-        ref_policy = self.ref_policy
+        model = self.policy_and_value
+        ref_policy = self.ref_model
         reward_model = self.reward_model
         processing_class = self.processing_class
         dataloader = self.dataloader
@@ -392,8 +396,8 @@ def repeat_generator():
 
         # backward compatibility
         if self.is_deepspeed_enabled:
-            self.deepspeed = self.model
-            self.model_wrapped = self.model
+            self.deepspeed = self.policy_and_value
+            self.model_wrapped = self.policy_and_value
 
         for update in range(1, args.num_total_batches + 1):
             self.state.episode += 1 * args.batch_size
@@ -680,7 +684,7 @@ def generate_completions(self, sampling: bool = False):
         )
 
         table = defaultdict(list)
-        with unwrap_model_for_generation(self.model, self.accelerator) as unwrapped_model:
+        with unwrap_model_for_generation(self.policy_and_value, self.accelerator) as unwrapped_model:
             for batch in self.eval_dataloader:
                 query = batch["input_ids"]
                 with torch.no_grad():
@@ -743,16 +747,18 @@ def create_model_card(
         if not self.is_world_process_zero():
             return
 
-        if hasattr(self.model.config, "_name_or_path") and not os.path.isdir(self.model.config._name_or_path):
-            base_model = self.model.config._name_or_path
+        if hasattr(self.policy_and_value.config, "_name_or_path") and not os.path.isdir(
+            self.policy_and_value.config._name_or_path
+        ):
+            base_model = self.policy_and_value.config._name_or_path
         else:
             base_model = None
 
         tags = tags or []
         if isinstance(tags, str):
             tags = [tags]
 
-        if hasattr(self.model.config, "unsloth_version"):
+        if hasattr(self.policy_and_value.config, "unsloth_version"):
             tags.append("unsloth")
 
         citation = textwrap.dedent("""\