Repository: huggingface/trl
Pull request: Add LlaMa in tests + `create_reference_model`
# What does this PR do?

Fixes https://github.com/lvwerra/trl/issues/258
Before this PR, Llama was not supported by `create_reference_model`, hence users needed to manually define the pattern of the layers to properly create a reference model, and this can be error prone.

This PR simply adds the Llama layer pattern for `create_reference_model`, I also added a tiny Llama model for testing purpose

cc @lvwerra 

Problem statement:
Debugging `create_reference_model` + `ValueError: optimizer got an empty parameter list`
I'm using the `create_reference_model` function to debug using LLaMA models with shared parameters, but when I use the default optimizer, I get the error that `ValueError: optimizer got an empty parameter list`.

I looked up the pattern used in the llama models, and it should be correct here:
```
ref_model = create_reference_model(model, num_shared_layers=20, pattern="model.layers{layer}")
```

( **UPDATE**: I saw that the pattern had a typo, missing a ".", but that is still some weird behavior, the correct one is `model.layers.{layer}`. Would be helpful to improve this still )

Though, I'm not sure why it results in an empty parameter list passed to Adam. Could be maybe check this in the `create_reference_model` function to create a more useful error?

Here's a shortened snippet:
```
    model = AutoModelForCausalLMWithValueHead.from_pretrained(model, revision=model_args.model_revision)

    # create reference / base model for KL constraint in RL
    if training_args.num_shared_layers:
        # TODO make configurable (hard coded for LLaMA)
        ref_model = create_reference_model(model, num_shared_layers=training_args.num_shared_layers, pattern="model.layers{layer}")
        # ref model requires specific optimizer
        optimizer = Adam(filter(lambda p: p.requires_grad, model.parameters()), lr=training_args.learning_rate)

...

    ppo_trainer = PPOTrainer(
        rl_config,
        model,
        ref_model,
        optimizer=optimizer,
        tokenizer=tokenizer,
        dataset=tokenized_datasets["train"],
        data_collator=collator,
    )
```

Here's the llama dict for reference on `pattern`
```
  (model): LlamaModel(
    (embed_tokens): Embedding(32000, 4096, padding_idx=0)
    (layers): ModuleList(
      (0): LlamaDecoderLayer(
        (self_attn): LlamaAttention(
          (q_proj): Linear(in_features=4096, out_features=4096, bias=False)
          (k_proj): Linear(in_features=4096, out_features=4096, bias=False)
          (v_proj): Linear(in_features=4096, out_features=4096, bias=False)
          (o_proj): Linear(in_features=4096, out_features=4096, bias=False)
          (rotary_emb): LlamaRotaryEmbedding()
        )
        (mlp): LlamaMLP(
          (gate_proj): Linear(in_features=4096, out_features=11008, bias=False)
          (down_proj): Linear(in_features=11008, out_features=4096, bias=False)
          (up_proj): Linear(in_features=4096, out_features=11008, bias=False)
          (act_fn): SiLUActivation()
        )
        (input_layernorm): LlamaRMSNorm()
        (post_attention_layernorm): LlamaRMSNorm()
      )
      (1): LlamaDecoderLayer(
        (self_attn): LlamaAttention(
          (q_proj): Linear(in_features=4096, out_features=4096, bias=False)
          (k_proj): Linear(in_features=4096, out_features=4096, bias=False)
          (v_proj): Linear(in_features=4096, out_features=4096, bias=False)
          (o_proj): Linear(in_features=4096, out_features=4096, bias=False)
          (rotary_emb): LlamaRotaryEmbedding()
        )
        (mlp): LlamaMLP(
          (gate_proj): Linear(in_features=4096, out_features=11008, bias=False)
          (down_proj): Linear(in_features=11008, out_features=4096, bias=False)
          (up_proj): Linear(in_features=4096, out_features=11008, bias=False)
          (act_fn): SiLUActivation()
        )
        (input_layernorm): LlamaRMSNorm()
        (post_attention_layernorm): LlamaRMSNorm()
      )
...
      (31): LlamaDecoderLayer(
        (self_attn): LlamaAttention(
          (q_proj): Linear(in_features=4096, out_features=4096, bias=False)
          (k_proj): Linear(in_features=4096, out_features=4096, bias=False)
          (v_proj): Linear(in_features=4096, out_features=4096, bias=False)
          (o_proj): Linear(in_features=4096, out_features=4096, bias=False)
          (rotary_emb): LlamaRotaryEmbedding()
        )
        (mlp): LlamaMLP(
          (gate_proj): Linear(in_features=4096, out_features=11008, bias=False)
          (down_proj): Linear(in_features=11008, out_features=4096, bias=False)
          (up_proj): Linear(in_features=4096, out_features=11008, bias=False)
          (act_fn): SiLUActivation()
        )
        (input_layernorm): LlamaRMSNorm()
        (post_attention_layernorm): LlamaRMSNorm()
      )
    )
    (norm): LlamaRMSNorm()
  )
  (lm_head): Linear(in_features=4096, out_features=32000, bias=False)
)
```

Hints from issue discussion:

Commit message: add LlaMa in tests

Patch to review:
diff --git a/tests/test_modeling_value_head.py b/tests/test_modeling_value_head.py
index 25c7096225..ba44288cc6 100644
--- a/tests/test_modeling_value_head.py
+++ b/tests/test_modeling_value_head.py
@@ -30,6 +30,7 @@
     "trl-internal-testing/tiny-random-BloomForCausalLM",
     "trl-internal-testing/tiny-random-GPT2LMHeadModel",
     "trl-internal-testing/tiny-random-CodeGenForCausalLM-sharded",
+    "trl-internal-testing/tiny-random-LlamaForCausalLM",
 ]
 
 ALL_SEQ2SEQ_MODELS = [
diff --git a/trl/models/modeling_base.py b/trl/models/modeling_base.py
index f86dd7c1e3..d8f8697568 100644
--- a/trl/models/modeling_base.py
+++ b/trl/models/modeling_base.py
@@ -36,7 +36,12 @@
     )
 
 
-LAYER_PATTERNS = ["transformer.h.{layer}", "model.decoder.layers.{layer}", "gpt_neox.layers.{layer}"]
+LAYER_PATTERNS = [
+    "transformer.h.{layer}",
+    "model.decoder.layers.{layer}",
+    "gpt_neox.layers.{layer}",
+    "model.layers.{layer}",
+]
 
 
 class PreTrainedModelWrapper(nn.Module):


Code review comments:
diff_hunk: @@ -30,6 +30,7 @@
     "trl-internal-testing/tiny-random-BloomForCausalLM",
     "trl-internal-testing/tiny-random-GPT2LMHeadModel",
     "trl-internal-testing/tiny-random-CodeGenForCausalLM-sharded",
+    "trl-internal-testing/tiny-random-LlamaForCausalLM",
path: tests/test_modeling_value_head.py
text: @author:
```suggestion
    # "trl-internal-testing/tiny-random-LlamaForCausalLM", uncomment on the next transformers release
```