Repository: huggingface/trl

Issue:
Llama Reward Model is incorrectly merged
As mentioned in #287, `merge_peft_adapter` saves the Llama RM as a `LlamaForCausalLM` see [here](https://github.com/lvwerra/trl/blob/main/examples/stack_llama/scripts/merge_peft_adapter.py#L35)

But the reward model is trained and should be a `LlamaForSequenceClassification` and running `rl_training.py` gives the obvious warnings
```
Some weights of the model checkpoint at ./llama-7b-se-rm were not used when initializing LlamaForSequenceClassification: ['lm_head.weight']
- This IS expected if you are initializing LlamaForSequenceClassification from the checkpoint of a model trained on another task or with another architecture (e.g. initializing a BertForSequenceClassification model from a BertForPreTraining model).
- This IS NOT expected if you are initializing LlamaForSequenceClassification from the checkpoint of a model that you expect to be exactly identical (initializing a BertForSequenceClassification model from a BertForSequenceClassification model).
Some weights of LlamaForSequenceClassification were not initialized from the model checkpoint at /home/toolkit/huggingface/llama-7b-rm and are newly initialized: ['score.weight']
You should probably TRAIN this model on a down-stream task to be able to use it for predictions and inference.
```

We should instead check whether we are merging the rm and then save as a the correct model 

Also the `score.weight` is not being loaded as mentioned in #297 , see more info below


--- update --

It seems that `merge_peft_adapter` should be using `merge_and_unload()` which correctly overrides the score. But I haven't yet managed to get good results using the adapter weights on the hub


Discussion:
A very not-clean way of doing this is 

```
if "rm" in script_args.adapter_model_name:
    model = AutoModelForSequenceClassification.from_pretrained(
        script_args.base_model_name, num_labels=1, torch_dtype=torch.bfloat16
    )
else:
    model = AutoModelForCausalLM.from_pretrained(
        script_args.base_model_name, return_dict=True, torch_dtype=torch.bfloat16
    )
```

It also looks like the reward model's `score.weight` isn't being loaded from the adapter weights. I've checked and this can either be an issue with the adapter config from `trl-lib/llama-7b-se-rm-peft` as it doesn't mention the score as one of the `modules_to_save` even though it is. It could also be an issue with `peft` or how the adapter was saved. In any case, anyone who currently wants to use the RM should overwrite the score.weight before calling `model.save_pretrained` 

```
if "rm" in script_args.adapter_model_name:
    peft_state_dict = torch.load(
        "/path/to/huggingface/hub/models--trl-lib--llama-7b-se-rm-peft/snapshots/{some_number}/adapter_model.bin"
    )
    score_weight = peft_state_dict["base_model.model.base_model.model.score.weight"]
    model.score = torch.nn.Linear(4096, 1, bias=False)
    with torch.no_grad():
        model.score.weight.copy_(score_weight)

```
Even after these fixes, I've been unable to get the reward model loaded. I've tried merging the `trl-lib/llama-7b-se-rm-peft` adapter weights directly on llama 7b, as well as merging the `trl-lib/llama-7b-se-peft` weights first and then merging the rm weights. I'm not getting above 56% accuracy on the rm validation set. 
I agree with you, [@mnoukhov](https://github.com/mnoukhov). The problem is that `merge_peft_adapter` only saves the `base_model` (the llm part, not the score part wich is the head), even after the "fix". I'm encountering the same issue as you, by the way.

I'm planning to modify the pipeline definition. Instead of giving a string to the pipeline, I will reconstruct the entire reward model within the script and pass the model as an argument to the pipeline. I'll keep you updated on my progress.

The issue with that approach is that `PeftModelForSequenceClassification` isn't supported by `pipeline` so you'd need to modify that as well.

The correct thing to do is actually to use `peft`'s `merge_and_unload` but the weights on the hub don't seem to be working for me. For that reason I'm re-running the whole pipeline and will upload my adapter weights once its done
I agree with your assessment. The Pipeline doesn't seem to work well with PeftModelForSequenceClassification. However, I believe I've come up with a fix. I didn't use merge_peft_adapter for the reward, instead I modified the [rl_training.py](https://github.com/lvwerra/trl/blob/main/examples/stack_llama/scripts/rl_training.py) script. Here's how I reconstructed the reward model and changed the reward computation:

reward_tokenizer_name is the name of the reward tokenizer, in my case it's "EleutherAI/gpt-neo-125m".
reward_model_name is the location of the saved adapter of the reward model.

```
# Now we reconstruct the rewardmodel
reward_peft_config = PeftConfig.from_pretrained(script_args.reward_model_name)
reward_model = AutoModelForSequenceClassification.from_pretrained(script_args.reward_tokenizer_name, num_labels=1, return_dict=True) #, torch_dtype=torch.bfloat16)
reward_model = PeftModel.from_pretrained(reward_model, script_args.reward_model_name)
reward_model = reward_model.to(current_device)

reward_tokenizer = AutoTokenizer.from_pretrained(script_args.reward_tokenizer_name)
reward_tokenizer.pad_token = reward_tokenizer.eos_token
# Define the pad token in the model configuration
reward_model.config.pad_token_id = reward_tokenizer.pad_token_id

for epoch, batch in enumerate(ppo_trainer.dataloader):
    question_tensors = batch["input_ids"]

    response_tensors = ppo_trainer.generate(
        question_tensors,
        return_prompt=False,
        length_sampler=output_length_sampler,
        **generation_kwargs,
    )
    # tokenizer.batch_decode works on CPU. perhaps it is faster if
    # yhou use two models with same tokenizer and just works with tokens on the GPU

    batch["response"] = tokenizer.batch_decode(response_tensors, skip_special_tokens=True)

    # Prepare texts
    texts = [q + r for q, r in zip(batch["query"], batch["response"])]

    # Tokenize the texts in batch
    texts_tokenized = reward_tokenizer(texts, truncation=True, padding=True, return_tensors="pt")

    # Get the tensors
    input_ids = texts_tokenized["input_ids"].to(current_device)
    attention_mask = texts_tokenized["attention_mask"].to(current_device)
)
    # Forward pass through the reward_model
    pipe_outputs = reward_model(input_ids=input_ids, attention_mask=attention_mask)

    rewards = list((torch.sigmoid(pipe_outputs.logits) - script_args.reward_baseline).unbind(dim=0))

    # Run PPO step
    stats = ppo_trainer.step(question_tensors, response_tensors, rewards)
    ppo_trainer.log_stats(stats, batch, rewards)

    if script_args.save_freq and epoch and epoch % script_args.save_freq == 0:
        ppo_trainer.save_pretrained(script_args.output_dir + f"step_{epoch}")
```
Any feedback or suggestions would be helpful.
This makes sense but it doesn't solve the problem. If you have adapter weights that you've trained yourself you can correctly load them in with `model = model.merge_and_unload()` , replacing lines 54-62 of `merge_peft_adapter.py` (and also changing CausalLM for SequenceClassification) 

The remaining issue is the weights on the hub don't seem to get the correct reward modelling accuracy
> 

That sounds like a good idea！May I ask if you run it successfully? How is the effect? Is there any loss in the indicator? Thanks.
Thx. I have been able to get the training process to run, although I haven't completed training the entire model yet. I'm unsure regarding the calculation of rewards. I don't know if I should directly use the logits or if I need to apply a sigmoid function to them. I have to check what works best.

Also if you have the same tokenizer for the reward function and the model, I think that the code I made can be improved... I have to check that...
Here is another version for the training loop, wich is improved when the reward model and the model share the same tokenizer.
This is faster, the other one is more general
```
reward_mean = []
reward_std = []
for epoch, batch in enumerate(ppo_trainer.dataloader):
    question_tensors = batch["input_ids"]

    response_tensors = ppo_trainer.generate(
        question_tensors,
        return_prompt=False,
        length_sampler=output_length_sampler,
        **generation_kwargs,
    )
    # Prepare tensors directly, no decoding
    input_ids = [torch.cat([q, r], dim=-1) for q, r in zip(question_tensors, response_tensors)]
    
    # Pad the input sequences
    input_ids_padded = torch.nn.utils.rnn.pad_sequence(input_ids, batch_first=True)
    
    # Create attention_mask
    attention_mask = torch.ones_like(input_ids_padded)

    # Forward pass through the reward_model
    pipe_outputs = reward_model(input_ids=input_ids_padded, attention_mask=attention_mask)

    rewards = list((torch.sigmoid(pipe_outputs.logits) - script_args.reward_baseline).unbind(dim=0))

    # Run PPO step
    stats = ppo_trainer.step(question_tensors, response_tensors, rewards)
    ppo_trainer.log_stats(stats, batch, rewards)
    reward_mean.append(stats['ppo/mean_scores'])
    reward_std.append(stats['ppo/std_scores'])
    if script_args.save_freq and epoch and epoch % script_args.save_freq == 0:
        ppo_trainer.save_pretrained(script_args.output_dir + f"step_{epoch}")
```
> Here is another version for the training loop, wich is improved when the reward model and the model share the same tokenizer. This is faster, the other one is more general
> 
> ```
> reward_mean = []
> reward_std = []
> for epoch, batch in enumerate(ppo_trainer.dataloader):
>     question_tensors = batch["input_ids"]
> 
>     response_tensors = ppo_trainer.generate(
>         question_tensors,
>         return_prompt=False,
>         length_sampler=output_length_sampler,
>         **generation_kwargs,
>     )
>     # Prepare tensors directly, no decoding
>     input_ids = [torch.cat([q, r], dim=-1) for q, r in zip(question_tensors, response_tensors)]
>     
>     # Pad the input sequences
>     input_ids_padded = torch.nn.utils.rnn.pad_sequence(input_ids, batch_first=True)
>     
>     # Create attention_mask
>     attention_mask = torch.ones_like(input_ids_padded)
> 
>     # Forward pass through the reward_model
>     pipe_outputs = reward_model(input_ids=input_ids_padded, attention_mask=attention_mask)
> 
>     rewards = list((torch.sigmoid(pipe_outputs.logits) - script_args.reward_baseline).unbind(dim=0))
> 
>     # Run PPO step
>     stats = ppo_trainer.step(question_tensors, response_tensors, rewards)
>     ppo_trainer.log_stats(stats, batch, rewards)
>     reward_mean.append(stats['ppo/mean_scores'])
>     reward_std.append(stats['ppo/std_scores'])
>     if script_args.save_freq and epoch and epoch % script_args.save_freq == 0:
>         ppo_trainer.save_pretrained(script_args.output_dir + f"step_{epoch}")
> ```

![截屏2023-05-17 18 34 15](https://github.com/lvwerra/trl/assets/26736518/315b2003-f0b2-44f5-9c78-184108bdd8bd)
![截屏2023-05-17 18 36 01](https://github.com/lvwerra/trl/assets/26736518/2bb2be6a-33d4-4c25-86ba-b5772ba16644)

Hi，did you encounter this problem when training this script?
No... I have never seen that error before. It seems to me that is an enviroment error related to cuda version... but no, no idea.

Resolving patch:
diff --git a/examples/stack_llama/scripts/merge_peft_adapter.py b/examples/stack_llama/scripts/merge_peft_adapter.py
--- a/examples/stack_llama/scripts/merge_peft_adapter.py
+++ b/examples/stack_llama/scripts/merge_peft_adapter.py
@@ -1,17 +1,9 @@
 from dataclasses import dataclass, field
 from typing import Optional
 
-import peft
 import torch
 from peft import PeftConfig, PeftModel
-from peft.utils import _get_submodules
-from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer, HfArgumentParser
-
-
-DEFAULT_PAD_TOKEN = "[PAD]"
-DEFAULT_EOS_TOKEN = "
"
-DEFAULT_BOS_TOKEN = "
"
-DEFAULT_UNK_TOKEN = "
"
+from transformers import AutoModelForCausalLM, AutoModelForSequenceClassification, AutoTokenizer, HfArgumentParser
 
 
 @dataclass
@@ -32,34 +24,23 @@ class ScriptArguments:
 assert script_args.base_model_name is not None, "please provide the output name of the merged model"
 
 peft_config = PeftConfig.from_pretrained(script_args.adapter_model_name)
-model = AutoModelForCausalLM.from_pretrained(script_args.base_model_name, return_dict=True, torch_dtype=torch.bfloat16)
-tokenizer = AutoTokenizer.from_pretrained(script_args.base_model_name)
-config = AutoConfig.from_pretrained(script_args.base_model_name)
-architecture = config.architectures[0]
-if "Llama" in architecture:
-    print("Setting EOS, BOS, and UNK tokens for LLama tokenizer")
-    tokenizer.add_special_tokens(
-        {
-            "eos_token": DEFAULT_EOS_TOKEN,
-            "bos_token": DEFAULT_BOS_TOKEN,
-            "unk_token": DEFAULT_UNK_TOKEN,
-            "pad_token": DEFAULT_PAD_TOKEN,
-        }
+if peft_config.task_type == "SEQ_CLS":
+    # peft is for reward model so load sequence classification
+    model = AutoModelForSequenceClassification.from_pretrained(
+        script_args.base_model_name, num_labels=1, torch_dtype=torch.bfloat16
+    )
+else:
+    model = AutoModelForCausalLM.from_pretrained(
+        script_args.base_model_name, return_dict=True, torch_dtype=torch.bfloat16
     )
 
+tokenizer = AutoTokenizer.from_pretrained(script_args.base_model_name)
+
 # Load the Lora model
 model = PeftModel.from_pretrained(model, script_args.adapter_model_name)
 model.eval()
 
-key_list = [key for key, _ in model.base_model.model.named_modules() if "lora" not in key]
-for key in key_list:
-    parent, target, target_name = _get_submodules(model.base_model.model, key)
-    if isinstance(target, peft.tuners.lora.Linear):
-        bias = target.bias is not None
-        new_module = torch.nn.Linear(target.in_features, target.out_features, bias=bias)
-        model.base_model._replace_module(parent, target_name, new_module, target)
-
-model = model.base_model.model
+model = model.merge_and_unload()
 
 model.save_pretrained(f"{script_args.output_name}")
 tokenizer.save_pretrained(f"{script_args.output_name}")


Test patch: