Repository: huggingface/trl
Pull request: 📉 Add PEFT support for `PPOTrainer`
# What does this PR do?
This PR provides an initial implementation of PEFT support for PPOTrainer, targeted towards the policy model. The existing functionality remains fully intact and unaffected by these modifications.

The following tasks have been completed:
- Add `peft_config` for `PPOTrainer`
- Add unit test for `tests/test_ppo_trainer.py` (next: implement unit tests to cover the new `peft_config` integration)
- Update `examples/scripts/ppo/ppo_tldr.py`
- Update `examples/scripts/ppo/ppo.py`

As of now, this implementation does not support **quantization** or **gradient checkpointing**.

If there’s anything inappropriate or suboptimal in my implementation, please feel free to let me know—I’m open to all feedback and eager to improve to align with the project standards.

Fixes #1916


## Before submitting
- [ ] This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
- [X] Did you read the [contributor guideline](https://github.com/huggingface/trl/blob/main/CONTRIBUTING.md#create-a-pull-request),
      Pull Request section?
- [X] Was this discussed/approved via a GitHub issue? Please add a link
      to it if that's the case.
- [ ] Did you make sure to update the documentation with your changes? Here are the
      [documentation guidelines](https://github.com/huggingface/trl/tree/main/docs).
- [X] Did you write any new necessary tests?


## Who can review?

Anyone in the community is free to review the PR once the tests have passed. Feel free to tag
members/contributors who may be interested in your PR.

Problem statement:
How to Add PEFT to PPO Trainer or PPO Config
I am trying to realize RLHF through PPO.

May I ask how can I realize PEFT in RLHF/PPO. I can see this parameter in DPOTrainer. However, I cannot see that in PPOTrainer.

Hints from issue discussion:
PPO does not yet support PEFT. But it may be a good enhancement.
Hi @qgallouedec, I'd love to contribute to this enhancement! I noticed that this issue has been open for three months, and I'd like to help bring it forward.

Is there anything specific I should be mindful of as I start on this? Any guidance or notes on the current status would be greatly appreciated.

Thank you!

Commit message: Merge branch 'main' into feature-ppo-peft

Patch to review:
diff --git a/examples/scripts/ppo/ppo.py b/examples/scripts/ppo/ppo.py
index 19036ca7c1..e0dd07bb5a 100644
--- a/examples/scripts/ppo/ppo.py
+++ b/examples/scripts/ppo/ppo.py
@@ -14,6 +14,7 @@
 
 import shutil
 
+import torch
 from accelerate import PartialState
 from datasets import load_dataset
 from transformers import (
@@ -23,7 +24,15 @@
     HfArgumentParser,
 )
 
-from trl import ModelConfig, PPOConfig, PPOTrainer, ScriptArguments
+from trl import (
+    ModelConfig,
+    PPOConfig,
+    PPOTrainer,
+    ScriptArguments,
+    get_kbit_device_map,
+    get_peft_config,
+    get_quantization_config,
+)
 from trl.trainer.utils import SIMPLE_CHAT_TEMPLATE
 
 
@@ -67,6 +76,20 @@
     ################
     # Model & Tokenizer
     ################
+    torch_dtype = (
+        model_config.torch_dtype
+        if model_config.torch_dtype in ["auto", None]
+        else getattr(torch, model_config.torch_dtype)
+    )
+    quantization_config = get_quantization_config(model_config)
+    model_kwargs = dict(
+        revision=model_config.model_revision,
+        attn_implementation=model_config.attn_implementation,
+        torch_dtype=torch_dtype,
+        device_map=get_kbit_device_map() if quantization_config is not None else None,
+        quantization_config=quantization_config,
+    )
+
     tokenizer = AutoTokenizer.from_pretrained(
         model_config.model_name_or_path,
         padding_side="left",
@@ -81,12 +104,18 @@
     reward_model = AutoModelForSequenceClassification.from_pretrained(
         training_args.reward_model_path, trust_remote_code=model_config.trust_remote_code, num_labels=1
     )
-    ref_policy = AutoModelForCausalLM.from_pretrained(
-        training_args.sft_model_path, trust_remote_code=model_config.trust_remote_code
-    )
     policy = AutoModelForCausalLM.from_pretrained(
         training_args.sft_model_path, trust_remote_code=model_config.trust_remote_code
     )
+
+    peft_config = get_peft_config(model_config)
+    if peft_config is None:
+        ref_policy = AutoModelForCausalLM.from_pretrained(
+            training_args.sft_model_path, trust_remote_code=model_config.trust_remote_code
+        )
+    else:
+        ref_policy = None
+
     ################
     # Dataset
     ################
@@ -131,6 +160,7 @@ def tokenize(element):
         value_model=value_model,
         train_dataset=train_dataset,
         eval_dataset=eval_dataset,
+        peft_config=peft_config,
     )
     trainer.train()
 
diff --git a/examples/scripts/ppo/ppo_tldr.py b/examples/scripts/ppo/ppo_tldr.py
index fd20d6693d..73ea5cd852 100644
--- a/examples/scripts/ppo/ppo_tldr.py
+++ b/examples/scripts/ppo/ppo_tldr.py
@@ -14,6 +14,7 @@
 
 import shutil
 
+import torch
 from accelerate import PartialState
 from datasets import load_dataset
 from transformers import (
@@ -23,16 +24,24 @@
     HfArgumentParser,
 )
 
-from trl import ModelConfig, PPOConfig, PPOTrainer, ScriptArguments
+from trl import (
+    ModelConfig,
+    PPOConfig,
+    PPOTrainer,
+    ScriptArguments,
+    get_kbit_device_map,
+    get_peft_config,
+    get_quantization_config,
+)
 from trl.trainer.utils import SIMPLE_CHAT_TEMPLATE
 
 
 """
 python examples/scripts/ppo/ppo_tldr.py \
-    --dataset_name trl-internal-testing/tldr-preference-sft-trl-style
+    --dataset_name trl-internal-testing/tldr-preference-sft-trl-style \
     --dataset_test_split validation \
     --learning_rate 3e-6 \
-    --output_dir models/minimal/ppo \
+    --output_dir models/minimal/ppo_tldr \
     --per_device_train_batch_size 1 \
     --gradient_accumulation_steps 64 \
     --total_episodes 30000 \
@@ -41,11 +50,13 @@
     --reward_model_path cleanrl/EleutherAI_pythia-1b-deduped__reward__tldr \
     --missing_eos_penalty 1.0 \
     --stop_token eos \
-    --response_length 53
+    --response_length 53 \
+    --eval_strategy steps \
+    --eval_steps 100
 
 accelerate launch --config_file examples/accelerate_configs/deepspeed_zero2.yaml \
     examples/scripts/ppo/ppo_tldr.py \
-    --dataset_name trl-internal-testing/tldr-preference-sft-trl-style
+    --dataset_name trl-internal-testing/tldr-preference-sft-trl-style \
     --dataset_test_split validation \
     --output_dir models/minimal/ppo_tldr \
     --learning_rate 3e-6 \
@@ -57,7 +68,9 @@
     --reward_model_path cleanrl/EleutherAI_pythia-1b-deduped__reward__tldr \
     --local_rollout_forward_batch_size 16 \
     --missing_eos_penalty 1.0 \
-    --stop_token eos
+    --stop_token eos \
+    --eval_strategy steps \
+    --eval_steps 100
 """
 
 
@@ -70,6 +83,20 @@
     ################
     # Model & Tokenizer
     ################
+    torch_dtype = (
+        model_config.torch_dtype
+        if model_config.torch_dtype in ["auto", None]
+        else getattr(torch, model_config.torch_dtype)
+    )
+    quantization_config = get_quantization_config(model_config)
+    model_kwargs = dict(
+        revision=model_config.model_revision,
+        attn_implementation=model_config.attn_implementation,
+        torch_dtype=torch_dtype,
+        device_map=get_kbit_device_map() if quantization_config is not None else None,
+        quantization_config=quantization_config,
+    )
+
     tokenizer = AutoTokenizer.from_pretrained(
         model_config.model_name_or_path,
         padding_side="left",
@@ -84,12 +111,18 @@
     reward_model = AutoModelForSequenceClassification.from_pretrained(
         training_args.reward_model_path, trust_remote_code=model_config.trust_remote_code, num_labels=1
     )
-    ref_policy = AutoModelForCausalLM.from_pretrained(
-        training_args.sft_model_path, trust_remote_code=model_config.trust_remote_code
-    )
     policy = AutoModelForCausalLM.from_pretrained(
         training_args.sft_model_path, trust_remote_code=model_config.trust_remote_code
     )
+
+    peft_config = get_peft_config(model_config)
+    if peft_config is None:
+        ref_policy = AutoModelForCausalLM.from_pretrained(
+            training_args.sft_model_path, trust_remote_code=model_config.trust_remote_code
+        )
+    else:
+        ref_policy = None
+
     ################
     # Dataset
     ################
@@ -138,6 +171,7 @@ def tokenize(element):
         value_model=value_model,
         train_dataset=train_dataset,
         eval_dataset=eval_dataset,
+        peft_config=peft_config,
     )
     trainer.train()
 
diff --git a/tests/test_ppo_trainer.py b/tests/test_ppo_trainer.py
index c530f05b86..aacf0f2a1c 100644
--- a/tests/test_ppo_trainer.py
+++ b/tests/test_ppo_trainer.py
@@ -14,6 +14,8 @@
 import platform
 import subprocess
 
+from transformers.testing_utils import require_peft
+
 
 def test():
     command = """\
@@ -65,3 +67,34 @@ def test_num_train_epochs():
         shell=True,
         check=True,
     )
+
+
+@require_peft
+def test_peft_support():
+    command = """\
+python examples/scripts/ppo/ppo.py \
+    --dataset_name trl-internal-testing/descriptiveness-sentiment-trl-style \
+    --dataset_train_split descriptiveness \
+    --learning_rate 3e-6 \
+    --output_dir models/minimal/ppo \
+    --per_device_train_batch_size 4 \
+    --gradient_accumulation_steps 1 \
+    --total_episodes 10 \
+    --model_name_or_path EleutherAI/pythia-14m \
+    --missing_eos_penalty 1.0 \
+    --save_strategy no \
+    --stop_token eos \
+    --use_peft \
+    --lora_r 32 \
+    --lora_alpha 16 \
+    --lora_target_modules query_key_value dense
+"""
+    if platform.system() == "Windows":
+        # windows CI does not work with subprocesses for some reason
+        # e.g., https://github.com/huggingface/trl/actions/runs/9600036224/job/26475286210?pr=1743
+        return
+    subprocess.run(
+        command,
+        shell=True,
+        check=True,
+    )
diff --git a/trl/trainer/ppo_config.py b/trl/trainer/ppo_config.py
index 25c2916025..465048184b 100644
--- a/trl/trainer/ppo_config.py
+++ b/trl/trainer/ppo_config.py
@@ -14,6 +14,7 @@
 
 import os
 from dataclasses import dataclass
+from typing import Optional
 
 from ..trainer.utils import OnPolicyConfig
 
@@ -32,6 +33,10 @@ class PPOConfig(OnPolicyConfig):
             Name of this experiment.
         reward_model_path (`str`, *optional*, defaults to `"EleutherAI/pythia-160m"`):
             Path to the reward model.
+        model_adapter_name (`Optional[str]`, *optional*, defaults to `None`):
+            Name of the train target PEFT adapter, when using LoRA with multiple adapters.
+        ref_adapter_name (`Optional[str]`, *optional*, defaults to `None`):
+            Name of the reference PEFT adapter, when using LoRA with multiple adapters.
         num_ppo_epochs (`int`, *optional*, defaults to `4`):
             Number of epochs to train.
         whiten_rewards (`bool`, *optional*, defaults to `False`):
@@ -52,6 +57,8 @@ class PPOConfig(OnPolicyConfig):
 
     exp_name: str = os.path.basename(__file__)[: -len(".py")]
     reward_model_path: str = "EleutherAI/pythia-160m"
+    model_adapter_name: Optional[str] = None
+    ref_adapter_name: Optional[str] = None
     num_ppo_epochs: int = 4
     whiten_rewards: bool = False
     kl_coef: float = 0.05
diff --git a/trl/trainer/ppo_trainer.py b/trl/trainer/ppo_trainer.py
index 8e2c514a54..0a950a23cd 100644
--- a/trl/trainer/ppo_trainer.py
+++ b/trl/trainer/ppo_trainer.py
@@ -18,6 +18,7 @@
 import textwrap
 import time
 from collections import defaultdict
+from contextlib import contextmanager, nullcontext
 from typing import Dict, List, Optional, Tuple, Union
 
 import numpy as np
@@ -44,9 +45,11 @@
 from transformers.integrations import get_reporting_integration_callbacks
 from transformers.trainer import DEFAULT_CALLBACKS, DEFAULT_PROGRESS_CALLBACK
 from transformers.trainer_callback import CallbackHandler, ExportableState, PrinterCallback
+from transformers.utils import is_peft_available
 from transformers.utils.deprecation import deprecate_kwarg
 
 from ..core import masked_mean, masked_whiten
+from ..models import create_reference_model
 from ..models.utils import unwrap_model_for_generation
 from ..trainer.utils import (
     OnlineTrainerState,
@@ -61,9 +64,12 @@
     truncate_response,
 )
 from .ppo_config import PPOConfig
-from .utils import generate_model_card
+from .utils import generate_model_card, peft_module_casting_to_bf16
 
 
+if is_peft_available():
+    from peft import PeftConfig, PeftModel, get_peft_model
+
 if is_wandb_available():
     import wandb
 
@@ -99,7 +105,7 @@ def __init__(
             Union[PreTrainedTokenizerBase, BaseImageProcessor, FeatureExtractionMixin, ProcessorMixin]
         ],
         policy: nn.Module,
-        ref_policy: nn.Module,
+        ref_policy: Optional[nn.Module],
         reward_model: nn.Module,
         train_dataset: Dataset,
         value_model: Optional[nn.Module] = None,
@@ -108,11 +114,12 @@ def __init__(
         # less commonly used
         optimizers: Tuple[torch.optim.Optimizer, torch.optim.lr_scheduler.LambdaLR] = (None, None),
         callbacks: Optional[List[TrainerCallback]] = None,
+        peft_config: Optional["PeftConfig"] = None,
     ) -> None:
         if ref_policy is policy:
             raise ValueError(
                 "`policy` and `ref_policy` cannot be the same object. If you want `ref_policy` to be the "
-                "same as `policy`, you must mass a copy of it, or `None` if you use peft."
+                "same as `policy`, you must make a copy of it, or `None` if you use peft."
             )
 
         self.args = config
@@ -125,7 +132,32 @@ def __init__(
         )
         self.policy.generation_config.pad_token_id = None  # generate tokens without truncation / padding
 
-        self.ref_policy = ref_policy
+        # peft support
+        if not is_peft_available() and peft_config is not None:
+            raise ValueError(
+                "PEFT is not installed and you passed a `peft_config` in the trainer's kwargs, please install it to use the PEFT models"
+            )
+        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()
+
+            # 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.is_peft_model = is_peft_available() and isinstance(self.policy, 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
+        elif self.is_peft_model:
+            self.ref_policy = None
+        else:
+            self.ref_policy = create_reference_model(self.policy)
+
         self.reward_model = reward_model
         self.train_dataset = train_dataset
         self.train_dataset_len = len(train_dataset)
@@ -174,12 +206,13 @@ def __init__(
         #########
         # setup model, optimizer, and others
         #########
-        for module in [policy, ref_policy, value_model, reward_model]:
-            disable_dropout_in_model(module)
+        for module in [self.policy, self.ref_policy, 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(policy, value_model)
-        self.model.config = policy.config  # needed for pushing to hub
+        self.model = PolicyAndValueWrapper(self.policy, self.value_model)
+        self.model.config = self.policy.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
@@ -244,11 +277,20 @@ def __init__(
             self.reward_model = prepare_deepspeed(
                 self.reward_model, args.per_device_train_batch_size, args.fp16, args.bf16
             )
-            self.ref_policy = prepare_deepspeed(
-                self.ref_policy, args.per_device_train_batch_size, args.fp16, args.bf16
-            )
+
+            if self.ref_policy 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
+                )
         else:
-            self.ref_policy = self.ref_policy.to(self.accelerator.device)
+            if self.ref_policy 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.reward_model = self.reward_model.to(self.accelerator.device)
 
     def get_train_dataloader(self) -> DataLoader:
@@ -257,6 +299,18 @@ def get_train_dataloader(self) -> DataLoader:
     def get_eval_dataloader(self) -> DataLoader:
         return self.eval_dataloader
 
+    @contextmanager
+    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
+        ).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)
+            yield
+            if self.ref_adapter_name:
+                self.model.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
@@ -368,7 +422,11 @@ def repeat_generator():
                     del logits, all_logprob
                     torch.cuda.empty_cache()
 
-                    ref_output = forward(ref_policy, query_response, processing_class.pad_token_id)
+                    if ref_policy is None:
+                        with self.null_ref_context():
+                            ref_output = forward(model.policy, query_response, processing_class.pad_token_id)
+                    else:
+                        ref_output = forward(ref_policy, query_response, processing_class.pad_token_id)
                     ref_logits = ref_output.logits[:, context_length - 1 : -1]
                     ref_logits /= args.temperature + 1e-7
                     ref_all_logprob = F.log_softmax(ref_logits, dim=-1)


Code review comments:
diff_hunk: @@ -125,7 +132,32 @@ def __init__(
         )
         self.policy.generation_config.pad_token_id = None  # generate tokens without truncation / padding
 
-        self.ref_policy = ref_policy
+        # peft support
+        if not is_peft_available() and peft_config is not None:
+            raise ValueError(
path: trl/trainer/ppo_trainer.py
text: @user1:
```suggestion
            raise ImportError(
```