Repository: huggingface/trl

Issue:
dpo training with Lora can not save fine-tuned weights
### issue
The following script manage to train and save. However, the saved weights are incorrect.

### Training script
```
def main():
    parser = HfArgumentParser((ModelArguments, DataTrainingArguments, TrainingArguments))
    model_args, data_args, training_args = parser.parse_args_into_dataclasses()

    # load config and tokenziers
    config = LlamaConfig.from_pretrained(model_args.model_name_or_path)
    config.use_cache = False
    tokenizer = LlamaTokenizer.from_pretrained(model_args.model_name_or_path, truncation_side='left')

    # initialize modules
    model = LlamaForCausalLM.from_pretrained(model_args.model_name_or_path, config=config)
    model.enable_input_require_grads()

    # add pad token in tokenizer if needed
    if tokenizer.pad_token is None:
        tokenizer.add_special_tokens({"pad_token":"<pad>"})
        tokenizer.pad_token_id = 0

    # Setup seed
    set_seed(training_args.seed)
    embedding_size = model.get_input_embeddings().weight.shape[0]
    if len(tokenizer) > embedding_size:
        model.resize_token_embeddings(len(tokenizer))

    # Setup Trainer
    training_args = training_args.to_dict()
    training_args |= {'remove_unused_columns': False}
    training_args = TrainingArguments(**training_args)
    peft_config = LoraConfig(
        r=model_args.lora_r,
        lora_alpha=model_args.lora_alpha,
        lora_dropout=model_args.lora_dropout,
        target_modules=[
            "q_proj",
            "v_proj",
            "k_proj",
            "out_proj",
            "fc_in",
            "fc_out",
            "wte",
        ],
        bias="none",
        task_type="CAUSAL_LM",
    )
    model_peft = get_peft_model(model, peft_config)
    trainer = DPOTrainer(
        model=model_peft,
        ref_model=None,
        beta=0.1, # DPO temprature
        train_dataset=prepared_dataset["train"],
        eval_dataset=prepared_dataset["eval"],
        tokenizer=tokenizer,
        args=training_args,
        peft_config=peft_config,
        max_length=data_args.model_max_length,
        max_prompt_length=int(data_args.model_max_length) * 3 // 4,
    )

    # Training
    train_result = trainer.train()
    trainer.save_state()
    trainer.save_model()
```

### Training output
```
{'loss': 0.6934, 'learning_rate': 0.0, 'rewards/chosen': 0.0, 'rewards/rejected': 0.0, 'rewards/accuracies': 0.0, 'rewards/margins': 0.0, 'logps/rejected': -63.125, 'logps/chosen': -84.625, 'logits/rejected': -1.3544921875, 'logits/chosen': -1.353515625, 'epoch': 0.0}
{'loss': 0.6933, 'learning_rate': 0.0005, 'rewards/chosen': 0.0263671875, 'rewards/rejected': 0.0164031982421875, 'rewards/accuracies': 0.5625, 'rewards/margins': 0.00995635986328125, 'logps/rejected': -55.625, 'logps/chosen': -38.59375, 'logits/rejected': -1.24609375, 'logits/chosen': -1.2490234375, 'epoch': 0.0}
{'loss': 0.6938, 'learning_rate': 0.0005, 'rewards/chosen': 0.0035152435302734375, 'rewards/rejected': 0.006542205810546875, 'rewards/accuracies': 0.5, 'rewards/margins': -0.0030307769775390625, 'logps/rejected': -66.75, 'logps/chosen': -65.0625, 'logits/rejected': -1.296875, 'logits/chosen': -1.2890625, 'epoch': 0.0}
{'loss': 0.6814, 'learning_rate': 0.0003333333333333333, 'rewards/chosen': 0.05535888671875, 'rewards/rejected': -0.007030487060546875, 'rewards/accuracies': 0.75, 'rewards/margins': 0.062408447265625, 'logps/rejected': -70.9375, 'logps/chosen': -40.875, 'logits/rejected': -1.255859375, 'logits/chosen': -1.259765625, 'epoch': 0.01}
{'loss': 0.6477, 'learning_rate': 0.00016666666666666666, 'rewards/chosen': -0.0174407958984375, 'rewards/rejected': -0.07177734375, 'rewards/accuracies': 0.5, 'rewards/margins': 0.054351806640625, 'logps/rejected': -65.75, 'logps/chosen': -54.28125, 'logits/rejected': -1.2783203125, 'logits/chosen': -1.28515625, 'epoch': 0.01}
```

### Weight check when doing merge
```
# initialize modules
model = LlamaForCausalLM.from_pretrained(script_args.base_model_name, config=config)

print("-" * 20 + "Weight before merge" + "-" * 20)
print(model.get_output_embeddings().weight)

# Load the Lora model
model = PeftModel.from_pretrained(model, script_args.adapter_model_name)
model.eval()

# check Lora weights
print("-" * 20 + "Check Lora Weights" + "-" * 20)
print(model.model.model.layers[0].self_attn.q_proj.lora_B.default.weight)

# merge lora weight and base model
model = model.merge_and_unload()

print("-" * 20 + "Weight after merge" + "-" * 20)
print(model.get_output_embeddings().weight)
```

### Weight check outputs
```
--------------------Weight before merge--------------------
Parameter containing:
tensor([[-0.0027,  0.0020, -0.0072,  ...,  0.0034, -0.0074,  0.0074],
        [-0.0315,  0.0452, -0.0030,  ..., -0.0226,  0.0144,  0.0317],
        [-0.0127,  0.0016,  0.0189,  ..., -0.0264,  0.0157, -0.0071],
        ...,
        [ 0.0199,  0.0242,  0.0271,  ...,  0.0052, -0.0103, -0.0067],
        [ 0.0074, -0.0048,  0.0076,  ..., -0.0273, -0.0171,  0.0308],
        [ 0.0192,  0.0271,  0.0170,  ..., -0.0015, -0.0046, -0.0046]],
       requires_grad=True)
--------------------Check Lora Weights--------------------
Parameter containing:
tensor([[0., 0., 0.,  ..., 0., 0., 0.],
        [0., 0., 0.,  ..., 0., 0., 0.],
        [0., 0., 0.,  ..., 0., 0., 0.],
        ...,
        [0., 0., 0.,  ..., 0., 0., 0.],
        [0., 0., 0.,  ..., 0., 0., 0.],
        [0., 0., 0.,  ..., 0., 0., 0.]])
--------------------Weight after merge--------------------
Parameter containing:
tensor([[-0.0027,  0.0020, -0.0072,  ...,  0.0034, -0.0074,  0.0074],
        [-0.0315,  0.0452, -0.0030,  ..., -0.0226,  0.0144,  0.0317],
        [-0.0127,  0.0016,  0.0189,  ..., -0.0264,  0.0157, -0.0071],
        ...,
        [ 0.0199,  0.0242,  0.0271,  ...,  0.0052, -0.0103, -0.0067],
        [ 0.0074, -0.0048,  0.0076,  ..., -0.0273, -0.0171,  0.0308],
        [ 0.0192,  0.0271,  0.0170,  ..., -0.0015, -0.0046, -0.0046]])
```

I think even if the Lora fine-tuning is not convergent, the saved adpater value should not be zero?


Discussion:
I have the same question~
Maybe @younesbelkada and @kashif could have a look.
@LuJunru did you try to save the model via:

```
model_peft.save_pretrained("peft_checkpoint")
```
@kashif In fact, I had the same problem when training DPO (without accelerate). I believe the code here is what you described, but it doesn't work
![image](https://github.com/huggingface/trl/assets/44131880/19853a83-272e-41a9-a3e0-bc35f35fc8c9)
self.model.save_pretrained(output_dir)
@kashif @Moyhub @lvwerra Hi Guys, thank you for the feedback. When i directly check the weights in adapter bin, i found the weights was saved. However, the weight key there was base_model.model.base_model.model.xxxx, this maybe related to the deepspeed wrapper. I tried to rename the key to base_model.model.xxxx, the merging succeeded. You may have a check as well.
@younesbelkada something we could also test with #724. 
It seems like when you feed a `PEFTModel` and a `peft_config` to the `DPOTrainer`, the model to be trained gets nested into 2 `PEFTModel` so this is the reason why we get `base_model.model.base_model.model.xxx` as keys in the saved model and not `base_model.model.xxx`. 
In that case, to load the saved DPO model we must create a PEFTModel on top of the PEFTModel that was used as a reference model. Doing this, I managed to load the weights correctly. 
Also, when merging the weights, you must call it twice as follows:
`merged_model = dpo_model.merge_and_unload().merge_and_unload()`



> It seems like when you feed a `PEFTModel` and a `peft_config` to the `DPOTrainer`, the model to be trained gets nested into 2 `PEFTModel` so this is the reason why we get `base_model.model.base_model.model.xxx` as keys in the saved model and not `base_model.model.xxx`. In that case, to load the saved DPO model we must create a PEFTModel on top of the PEFTModel that was used as a reference model. Doing this, I managed to load the weights correctly. Also, when merging the weights, you must call it twice as follows: `merged_model = dpo_model.merge_and_unload().merge_and_unload()`

Can you provide a code example for your solution?
Seeing same issue. Model performance of training same as original model due to this even after dpo training. Any conclusion for this bug? 
@lvwerra  can you point to any other similar thread which has a solution for this ?
Same issue here
> @kashif @Moyhub @lvwerra Hi Guys, thank you for the feedback. When i directly check the weights in adapter bin, i found the weights was saved. However, the weight key there was base_model.model.base_model.model.xxxx, this maybe related to the deepspeed wrapper. I tried to rename the key to base_model.model.xxxx, the merging succeeded. You may have a check as well.

Could you share the example code? Thanks.

Resolving patch:
diff --git a/trl/trainer/dpo_trainer.py b/trl/trainer/dpo_trainer.py
--- a/trl/trainer/dpo_trainer.py
+++ b/trl/trainer/dpo_trainer.py
@@ -137,6 +137,10 @@ def __init__(
                 "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_config, we merge and unload it first
+            if isinstance(model, PeftModel):
+                model = model.merge_and_unload()
+
             if getattr(model, "is_loaded_in_8bit", False) or getattr(model, "is_loaded_in_4bit", False):
                 _support_gc_kwargs = hasattr(
                     args, "gradient_checkpointing_kwargs"
@@ -160,7 +164,10 @@ def make_inputs_require_grad(module, input, output):
                         output.requires_grad_(True)
 
                     model.get_input_embeddings().register_forward_hook(make_inputs_require_grad)
+
+            # get peft model with the given config
             model = get_peft_model(model, peft_config)
+
         # For models that use gradient_checkpoiting, we need to attach a hook that enables input
         # to explicitly have `requires_grad=True`, otherwise training will either silently
         # fail or completely fail.


Test patch:
diff --git a/tests/test_dpo_trainer.py b/tests/test_dpo_trainer.py
--- a/tests/test_dpo_trainer.py
+++ b/tests/test_dpo_trainer.py
@@ -240,3 +240,57 @@ def test_dpo_trainer_generate_during_eval_no_wandb(self):
                     eval_dataset=dummy_dataset,
                     generate_during_eval=True,
                 )
+
+    @require_peft
+    @mark.peft_test
+    def test_dpo_lora_save(self):
+        from peft import LoraConfig, get_peft_model
+
+        lora_config = LoraConfig(
+            r=16,
+            lora_alpha=32,
+            lora_dropout=0.05,
+            bias="none",
+            task_type="CAUSAL_LM",
+        )
+
+        # lora model
+        model = AutoModelForCausalLM.from_pretrained(self.model_id)
+        model_peft = get_peft_model(model, lora_config)
+
+        with tempfile.TemporaryDirectory() as tmp_dir:
+            training_args = TrainingArguments(
+                output_dir=tmp_dir,
+                per_device_train_batch_size=2,
+                max_steps=3,
+                remove_unused_columns=False,
+                gradient_accumulation_steps=4,
+                learning_rate=9e-1,
+                evaluation_strategy="steps",
+            )
+
+            dummy_dataset = self._init_dummy_dataset()
+
+            # dpo train lora model with a lora config
+            trainer = DPOTrainer(
+                model=model_peft,
+                ref_model=None,
+                beta=0.1,
+                args=training_args,
+                tokenizer=self.tokenizer,
+                train_dataset=dummy_dataset,
+                eval_dataset=dummy_dataset,
+                peft_config=lora_config,
+            )
+
+            # train the model
+            trainer.train()
+
+            # save peft adapter
+            trainer.save_model()
+
+            # assert that the model is loaded without giving OSError
+            try:
+                AutoModelForCausalLM.from_pretrained(tmp_dir)
+            except OSError:
+                self.fail("Loading the saved peft adapter failed")