👨‍👩‍👧 GRPO + PEFT + vLLM (#2818)

* peft + grpo + vllm

* test change

* support model alread peft

* Update tests/test_grpo_trainer.py

---------

Co-authored-by: Quentin Gallouédec <[email protected]>
Co-authored-by: Quentin Gallouédec <[email protected]>

Author: winglian

tests/test_grpo_trainer.py

@@ -125,7 +125,7 @@ def test_training_peft(self):
 
             self.assertIsNotNone(trainer.state.log_history[-1]["train_loss"])
 
-            # Check the peft params have changed and the base model params have not changed
+            # Check that the peft params have changed and the base model params have not changed
             for n, param in previous_trainable_params.items():
                 new_param = trainer.model.get_parameter(n)
                 if n in base_param_names:  # We expect the base model params to be the same
@@ -168,7 +168,7 @@ def test_training_different_reward_model(self):
 
             self.assertIsNotNone(trainer.state.log_history[-1]["train_loss"])
 
-            # Check the params have changed
+            # Check that the params have changed
             for n, param in previous_trainable_params.items():
                 new_param = trainer.model.get_parameter(n)
                 self.assertFalse(torch.equal(param, new_param), f"Parameter {n} has not changed.")
@@ -203,7 +203,7 @@ def reward_func(completions, **kwargs):
 
             self.assertIsNotNone(trainer.state.log_history[-1]["train_loss"])
 
-            # Check the params have changed
+            # Check that the params have changed
             for n, param in previous_trainable_params.items():
                 new_param = trainer.model.get_parameter(n)
                 self.assertFalse(torch.equal(param, new_param), f"Parameter {n} has not changed.")
@@ -239,7 +239,7 @@ def reward_func(completions, **kwargs):
 
             self.assertIsNotNone(trainer.state.log_history[-1]["train_loss"])
 
-            # Check the params have changed
+            # Check that the params have changed
             for n, param in previous_trainable_params.items():
                 new_param = trainer.model.get_parameter(n)
                 self.assertFalse(torch.equal(param, new_param), f"Parameter {n} has not changed.")
@@ -278,7 +278,7 @@ def reward_func2(completions, **kwargs):
 
             self.assertIsNotNone(trainer.state.log_history[-1]["train_loss"])
 
-            # Check the params have changed
+            # Check that the params have changed
             for n, param in previous_trainable_params.items():
                 new_param = trainer.model.get_parameter(n)
                 self.assertFalse(torch.equal(param, new_param), f"Parameter {n} has not changed.")
@@ -356,7 +356,7 @@ def reward_func(completions, **kwargs):
 
             self.assertIsNotNone(trainer.state.log_history[-1]["train_loss"])
 
-            # Check the params have changed
+            # Check that the params have changed
             for n, param in previous_trainable_params.items():
                 new_param = trainer.model.get_parameter(n)
                 self.assertFalse(torch.equal(param, new_param), f"Parameter {n} has not changed.")
@@ -395,7 +395,7 @@ def reward_func(completions, some_values, **kwargs):
 
             self.assertIsNotNone(trainer.state.log_history[-1]["train_loss"])
 
-            # Check the params have changed
+            # Check that the params have changed
             for n, param in previous_trainable_params.items():
                 new_param = trainer.model.get_parameter(n)
                 self.assertFalse(torch.equal(param, new_param), f"Parameter {n} has not changed.")
@@ -416,9 +416,10 @@ def test_training_vllm(self):
                 report_to="none",
                 use_vllm=True,
                 vllm_device="cuda:0",  # will raise a warning, but allows this test to work with only one GPU
+                vllm_gpu_memory_utilization=0.5,  # reduce since because we use the same device for training and vllm
             )
             trainer = GRPOTrainer(
-                model="trl-internal-testing/small-Qwen2ForCausalLM-2.5",
+                model="Qwen/Qwen2.5-0.5B-Instruct",  # tiny is too small for vLLM
                 reward_funcs="trl-internal-testing/tiny-Qwen2ForSequenceClassification-2.5",
                 args=training_args,
                 train_dataset=dataset,
@@ -504,6 +505,8 @@ def test_training_with_sync_ref_model(self):
     @require_peft
     def test_training_vllm_and_peft(self):
         """Test that training works with vLLM for generation."""
+        model = AutoModelForCausalLM.from_pretrained("Qwen/Qwen2.5-0.5B-Instruct")  # tiny model is too small for vLLM
+        base_param_names = [f"base_model.model.{n}" for n, _ in model.named_parameters()]
         dataset = load_dataset("trl-internal-testing/zen", "standard_prompt_only", split="train")
 
         with tempfile.TemporaryDirectory() as tmp_dir:
@@ -513,14 +516,22 @@ def test_training_vllm_and_peft(self):
                 per_device_train_batch_size=3,  # reduce the batch size to reduce memory usage
                 num_generations=3,  # reduce the number of generations to reduce memory usage
                 max_completion_length=32,  # reduce the completion length to reduce memory usage
-                use_vllm=True,
                 report_to="none",
+                use_vllm=True,
+                vllm_device="cuda:0",  # will raise a warning, but allows this test to work with only one GPU
+                vllm_gpu_memory_utilization=0.5,  # reduce since because we use the same device for training and vllm
+            )
+            lora_config = LoraConfig(
+                target_modules="all-linear",
+                # test with non-default modules as it add extra keys in state_dict tht we need to handle
+                modules_to_save=["embed_tokens", "lm_head"],
             )
             trainer = GRPOTrainer(
-                model="trl-internal-testing/small-Qwen2ForCausalLM-2.5",
+                model=model,
                 reward_funcs="trl-internal-testing/tiny-Qwen2ForSequenceClassification-2.5",
                 args=training_args,
                 train_dataset=dataset,
+                peft_config=lora_config,
             )
 
             previous_trainable_params = {n: param.clone() for n, param in trainer.model.named_parameters()}
@@ -529,7 +540,11 @@ def test_training_vllm_and_peft(self):
 
             self.assertIsNotNone(trainer.state.log_history[-1]["train_loss"])
 
-            # Check that the params have changed
+            # Check that the peft params have changed and the base model params have not changed
             for n, param in previous_trainable_params.items():
                 new_param = trainer.model.get_parameter(n)
-                self.assertFalse(torch.equal(param, new_param), f"Parameter {n} has not changed.")
+                if n in base_param_names:  # We expect the base model params to be the same
+                    self.assertTrue(torch.allclose(param, new_param), f"Parameter {n} has changed.")
+                elif "base_layer" not in n and "original_module" not in n:
+                    # We expect the peft params to be different (except for the base layer)
+                    self.assertFalse(torch.allclose(param, new_param), f"Parameter {n} has not changed.")

trl/trainer/grpo_trainer.py

@@ -22,7 +22,7 @@
 import torch
 import torch.utils.data
 import transformers
-from accelerate.utils import broadcast_object_list, gather, gather_object, set_seed
+from accelerate.utils import broadcast_object_list, gather, gather_object, is_peft_model, set_seed
 from accelerate.utils.other import is_compiled_module
 from datasets import Dataset, IterableDataset
 from packaging import version
@@ -51,7 +51,7 @@
 
 
 if is_peft_available():
-    from peft import PeftConfig, PeftModel, get_peft_model
+    from peft import PeftConfig, get_peft_model
 
 if is_vllm_available():
     from vllm import LLM, SamplingParams
@@ -249,7 +249,7 @@ def __init__(
         # Reference model
         if is_deepspeed_zero3_enabled():
             self.ref_model = AutoModelForCausalLM.from_pretrained(model_id, **model_init_kwargs)
-        elif peft_config is None:
+        elif not is_peft_model(model):
             # If PEFT configuration is not provided, create a reference model based on the initial model.
             self.ref_model = create_reference_model(model)
         else:
@@ -491,16 +491,18 @@ def _move_model_to_vllm(self):
             self.model, self.accelerator, gather_deepspeed3_params=self.args.ds3_gather_for_generation
         ) as unwrapped_model:
             if is_compiled_module(unwrapped_model):
-                state_dict = unwrapped_model._orig_mod.state_dict()
-            elif is_peft_available() and isinstance(unwrapped_model, PeftModel):
+                unwrapped_model = unwrapped_model._orig_mod
+            if is_peft_model(unwrapped_model):
                 unwrapped_model.merge_adapter()
                 state_dict = unwrapped_model.state_dict()
                 unwrapped_model.unmerge_adapter()
+                # Remove base_model and base_layer prefixes
                 state_dict = {
-                    k.removeprefix("base_model.model.").replace(".base_layer", ""): v
-                    for k, v in state_dict.items()
-                    if self.model.prefix not in k
+                    k.removeprefix("base_model.model.").replace(".base_layer", ""): v for k, v in state_dict.items()
                 }
+                # Remove values with adapter prefix (example: "_lora")
+                state_dict = {k: v for k, v in state_dict.items() if unwrapped_model.prefix not in k}
+                # When module to save, remove its prefix and discard the original module
                 state_dict = {
                     k.replace("modules_to_save.default.", ""): v
                     for k, v in state_dict.items()