handle native lora merge before shipping to vllm

Author: winglian

trl/trainer/grpo_trainer.py

@@ -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, get_lora_merged_state_dict, pad, selective_log_softmax
+from .utils import generate_model_card, get_comet_experiment_url, pad, selective_log_softmax
 
 
 if is_peft_available():
@@ -451,7 +451,18 @@ def _move_model_to_vllm(self):
             if is_compiled_module(unwrapped_model):
                 unwrapped_model = unwrapped_model._orig_mod
             if is_peft_model(unwrapped_model):
-                state_dict = get_lora_merged_state_dict(unwrapped_model)
+                unwrapped_model.merge_adapter()
+                state_dict = unwrapped_model.state_dict()
+                state_dict = {
+                    k.removeprefix("base_model.model.")
+                    .removeprefix("base_model.model.")
+                    .replace(".default", "")
+                    .replace(".base_layer", "")
+                    .replace(".modules_to_save", ""): v
+                    for k, v in state_dict.items()
+                    if unwrapped_model.prefix not in k and "original_module" not in k
+                }
+                unwrapped_model.unmerge_adapter()
             else:
                 state_dict = unwrapped_model.state_dict()
         if self.accelerator.is_main_process:

trl/trainer/utils.py

@@ -922,66 +922,6 @@ 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):
     """