Repository: huggingface/trl
Issue category: Bug Report

Whether chatglm3 6b is supported by trl ?
### System Info

trl  0.11.4
transformers    4.46.1


File "/home/fjy/folders/ERNIE2.0/finetune_chatglm6b.py", line 42, in <module>
    model = AutoModelForCausalLMWithValueHead.from_pretrained(model.model, trust_remote_code=True)
            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/fjy/anaconda3/envs/env/lib/python3.12/site-packages/trl/models/modeling_base.py", line 233, in from_pretrained
    model = cls(pretrained_model, **trl_model_args)
            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/fjy/anaconda3/envs/env/lib/python3.12/site-packages/trl/models/modeling_value_head.py", line 107, in __init__
    raise ValueError("The model does not have a language model head, please use a model that has one.")
ValueError: The model does not have a language model head, please use a model that has one.

### Information

- [ ] The official example scripts
- [X] My own modified scripts

### Tasks

- [ ] An officially supported task in the `examples` folder
- [X] My own task or dataset (give details below)

### Reproduction

### loading tokenizer
tokenizer = AutoTokenizer.from_pretrained("THUDM/chatglm3-6b-128k", trust_remote_code=True)

### loading base model
model = AutoModelForCausalLM.from_pretrained("THUDM/chatglm3-6b-128k", trust_remote_code=True)

### loading adapter
model = PeftModel.from_pretrained(model, adapter_dir)

### wraping model with a value head. 
model = AutoModelForCausalLMWithValueHead.from_pretrained(model.model, trust_remote_code=True)

The system throws an error at AutoModelForCausalLMWithValueHead.from_pretrained, The model does not have a language model head, please use a model that has one. Is it because trl does not support chatglm3? Because chatglm3 should have a language header, right?

### Expected behavior

How can trl work with chatglm3. 


Functions to edit:
trl/models/modeling_value_head.py:AutoModelForCausalLMWithValueHead
Functions added: trl/models/modeling_value_head.py:AutoModelForCausalLMWithValueHead

Patch:
diff --git a/trl/models/modeling_value_head.py b/trl/models/modeling_value_head.py
index 75a8fcd157..0797794013 100644
--- a/trl/models/modeling_value_head.py
+++ b/trl/models/modeling_value_head.py
@@ -70,8 +70,8 @@ class AutoModelForCausalLMWithValueHead(PreTrainedModelWrapper):
         - **transformers_parent_class** (`transformers.PreTrainedModel`) -- The parent class of the wrapped model. This
             should be set to `transformers.AutoModelForCausalLM` for this class.
         - **lm_head_namings** (`tuple`) -- A tuple of strings that are used to identify the language model head of the
-            wrapped model. This is set to `("lm_head", "embed_out")` for this class but can be changed for other models
-            in the future
+            wrapped model. This is set to `("lm_head", "embed_out", "output_layer")` for this class but can be changed
+            for other models in the future
         - **supported_args** (`tuple`) -- A tuple of strings that are used to identify the arguments that are supported
             by the `ValueHead` class. Currently, the supported args are:
             - **summary_dropout_prob** (`float`, `optional`, defaults to `None`) -- The dropout probability for the
@@ -86,7 +86,7 @@ class AutoModelForCausalLMWithValueHead(PreTrainedModelWrapper):
     """
 
     transformers_parent_class = AutoModelForCausalLM
-    lm_head_namings = ["lm_head", "embed_out"]
+    lm_head_namings = ["lm_head", "embed_out", "output_layer"]
     supported_args = (
         "summary_dropout_prob",
         "v_head_initializer_range",