whoops. forgot the lora merge part

Author: winglian

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
+from accelerate.utils import broadcast_object_list, gather, gather_object, is_peft_model
 from accelerate.utils.other import is_compiled_module
 from datasets import Dataset, IterableDataset
 from packaging import version
@@ -47,7 +47,7 @@
 from ..models import create_reference_model, prepare_deepspeed, unwrap_model_for_generation
 from .callbacks import SyncRefModelCallback
 from .grpo_config import GRPOConfig
-from .utils import generate_model_card, get_comet_experiment_url, pad, selective_log_softmax
+from .utils import generate_model_card, get_comet_experiment_url, get_lora_merged_state_dict, pad, selective_log_softmax
 
 
 if is_peft_available():
@@ -449,7 +449,9 @@ 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()
+                unwrapped_model = unwrapped_model._orig_mod
+            if is_peft_model(unwrapped_model):
+                state_dict = get_lora_merged_state_dict(unwrapped_model)
             else:
                 state_dict = unwrapped_model.state_dict()
         if self.accelerator.is_main_process:

trl/trainer/utils.py

@@ -60,6 +60,8 @@
 
 if is_peft_available():
     from peft import LoraConfig, PeftConfig
+    from peft.tuners.tuners_utils import onload_layer
+    from peft.utils import ModulesToSaveWrapper, _get_submodules
 
 
 class DataCollatorForCompletionOnlyLM(DataCollatorForLanguageModeling):
@@ -920,6 +922,66 @@ def get_peft_config(model_args: ModelConfig) -> "Optional[PeftConfig]":
 
     return peft_config
 
+def get_lora_merged_state_dict(
+        model: torch.nn.Module,
+) -> dict:
+    r"""
+    Create and return a state_dict that has the LoRA deltas
+    merged into the base model’s weights, without modifying `model` in place.
+
+    Arguments:
+        model (torch.nn.Module): A model that has LoRA/PEFT adapters attached.
+
+    Returns:
+        dict: A state_dict of the merged parameters.
+    """
+
+    if not is_peft_available():
+        raise ValueError(
+            "You need to have PEFT library installed in your environment, make sure to install `peft`. "
+            "Make sure to run `pip install -U peft`."
+        )
+
+    base_model_prefix = "base_model.model."
+    state_dict = {}
+    key_list = [key for key, _ in model.named_modules() if model.prefix not in key]
+    for key in key_list:
+        try:
+            _, target, _ = _get_submodules(model, key)
+        except AttributeError:
+            continue
+        with onload_layer(target):
+            weight_key = key.replace(base_model_prefix, "") + ".weight"
+            bias_key = key.replace(base_model_prefix, "") + ".bias"
+            if hasattr(target, "base_layer"):
+                target.merge(safe_merge=True, adapter_names=None)
+                # get the state_dict of target.base_layer
+                layer_state_dict = target.base_layer.state_dict()
+                state_dict[weight_key] = layer_state_dict["weight"]
+            elif isinstance(target, ModulesToSaveWrapper):
+                # save any additional trainable modules part of `modules_to_save`
+                new_module = target.modules_to_save[target.active_adapter]
+                if hasattr(new_module, "base_layer"):
+                    # check if the module is itself a tuner layer
+                    new_module.merge(safe_merge=True, adapter_names=None)
+                layer_state_dict = new_module.state_dict()
+                state_dict[weight_key] = layer_state_dict["weight"]
+            elif hasattr(target, "weight"):
+                if any(
+                        skip in key
+                        for skip in [
+                            ".original_module",
+                            ".modules_to_save",
+                            ".base_layer",
+                        ]
+                ):
+                    continue
+                layer_state_dict = target.state_dict()
+                state_dict[weight_key] = layer_state_dict["weight"]
+                if hasattr(target, "bias") and "bias" in layer_state_dict.keys():
+                    state_dict[bias_key] = layer_state_dict["bias"]
+    return state_dict
+
 
 def get_exp_cap(value, decimal=4):
     """