Repository: stanfordnlp/pyvene (language: python)

Issue:
[External]: Attention head intervention doesn't work for models other than GPT-2
### Contact Details

wangxz098@gmail.com

### What happened?

When intervening attention heads (i.e. set the intervention component as `head_attention_value_output`, `head_query_output`, `head_key_output`, or `head_value_output`), pyvene will crash. Only GPT-2 works, all the other LMs have the bug.

The stack log will be like:

```
Traceback (most recent call last):
  File "/data1/wxz/test_pyvene.py", line 21, in <module>
    _, counterfactual_outputs = intervenable(
  File "/home/wxz/miniconda3/envs/llm/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1518, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/home/wxz/miniconda3/envs/llm/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1527, in _call_impl
    return forward_call(*args, **kwargs)
  File "/home/wxz/workspace/pyvene/pyvene/models/intervenable_base.py", line 1465, in forward
    raise e
  File "/home/wxz/workspace/pyvene/pyvene/models/intervenable_base.py", line 1425, in forward
    self._wait_for_forward_with_parallel_intervention(
  File "/home/wxz/workspace/pyvene/pyvene/models/intervenable_base.py", line 1072, in _wait_for_forward_with_parallel_intervention
    _ = self.model(**sources[group_id])
  File "/home/wxz/miniconda3/envs/llm/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1518, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/home/wxz/miniconda3/envs/llm/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1527, in _call_impl
    return forward_call(*args, **kwargs)
  File "/home/wxz/miniconda3/envs/llm/lib/python3.10/site-packages/transformers/models/llama/modeling_llama.py", line 1192, in forward
    outputs = self.model(
  File "/home/wxz/miniconda3/envs/llm/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1518, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/home/wxz/miniconda3/envs/llm/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1527, in _call_impl
    return forward_call(*args, **kwargs)
  File "/home/wxz/miniconda3/envs/llm/lib/python3.10/site-packages/transformers/models/llama/modeling_llama.py", line 1019, in forward
    layer_outputs = decoder_layer(
  File "/home/wxz/miniconda3/envs/llm/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1518, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/home/wxz/miniconda3/envs/llm/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1527, in _call_impl
    return forward_call(*args, **kwargs)
  File "/home/wxz/miniconda3/envs/llm/lib/python3.10/site-packages/transformers/models/llama/modeling_llama.py", line 740, in forward
    hidden_states, self_attn_weights, present_key_value = self.self_attn(
  File "/home/wxz/miniconda3/envs/llm/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1518, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/home/wxz/miniconda3/envs/llm/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1527, in _call_impl
    return forward_call(*args, **kwargs)
  File "/home/wxz/miniconda3/envs/llm/lib/python3.10/site-packages/transformers/models/llama/modeling_llama.py", line 682, in forward
    attn_output = self.o_proj(attn_output)
  File "/home/wxz/miniconda3/envs/llm/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1518, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/home/wxz/miniconda3/envs/llm/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1547, in _call_impl
    args_kwargs_result = hook(self, args, kwargs)  # type: ignore[misc]
  File "/home/wxz/workspace/pyvene/pyvene/models/intervenable_base.py", line 747, in hook_callback
    selected_output = self._gather_intervention_output(
  File "/home/wxz/workspace/pyvene/pyvene/models/intervenable_base.py", line 658, in _gather_intervention_output
    selected_output = gather_neurons(
  File "/home/wxz/workspace/pyvene/pyvene/models/modeling_utils.py", line 276, in gather_neurons
    pos_tensor_input = bhsd_to_bs_hd(head_tensor_output)
  File "/home/wxz/workspace/pyvene/pyvene/models/modeling_utils.py", line 204, in bhsd_to_bs_hd
    b, h, s, d = tensor.shape
ValueError: not enough values to unpack (expected 4, got 3)
```
I've basically figured out how to fix this bug and will submit a PR later.

### Code to produce this issue.

```shell
#simplified from the nested intervention tutorial
import pyvene as pv
_, tokenizer, gpt = pv.create_llama("/data3/MODELS/llama2-hf/llama-2-7b") # every LM except for GPT-2 has the bug
pv_config = pv.IntervenableConfig(
        model_type=type(gpt),
        representations=[
            pv.RepresentationConfig(
                1,
                "head_attention_value_output",  # can also be changed to head_query_output, head_key_output, head_value_output, all don't work
                "h.pos",
                1, 
            ),
        ],
        intervention_types=pv.VanillaIntervention,
)
base = tokenizer("The capital of Spain is", return_tensors="pt")
sources = [tokenizer("The capital of Italy is", return_tensors="pt")]
intervenable = pv.IntervenableModel(pv_config, gpt)
_, counterfactual_outputs = intervenable(
    base,
    sources,
    {
        "sources->base": (
            [[[[0]], [[4]]]],
            [[[[0]], [[4]]]],
        )
    },
)
```



Discussion / hints:

Fix patch:
diff --git a/pyvene/models/gemma/modelings_intervenable_gemma.py b/pyvene/models/gemma/modelings_intervenable_gemma.py
index 09135bb8..31f771a7 100644
--- a/pyvene/models/gemma/modelings_intervenable_gemma.py
+++ b/pyvene/models/gemma/modelings_intervenable_gemma.py
@@ -20,41 +20,43 @@
     "mlp_output": ("layers[%s].mlp", CONST_OUTPUT_HOOK),
     "mlp_input": ("layers[%s].mlp", CONST_INPUT_HOOK),
     "attention_value_output": ("layers[%s].self_attn.o_proj", CONST_INPUT_HOOK),
-    "head_attention_value_output": ("layers[%s].self_attn.o_proj", CONST_INPUT_HOOK),
+    "head_attention_value_output": ("layers[%s].self_attn.o_proj", CONST_INPUT_HOOK, (split_head_and_permute, "n_head")),
     "attention_output": ("layers[%s].self_attn", CONST_OUTPUT_HOOK),
     "attention_input": ("layers[%s].self_attn", CONST_INPUT_HOOK),
     "query_output": ("layers[%s].self_attn.q_proj", CONST_OUTPUT_HOOK),
     "key_output": ("layers[%s].self_attn.k_proj", CONST_OUTPUT_HOOK),
     "value_output": ("layers[%s].self_attn.v_proj", CONST_OUTPUT_HOOK),
-    "head_query_output": ("layers[%s].self_attn.q_proj", CONST_OUTPUT_HOOK),
-    "head_key_output": ("layers[%s].self_attn.k_proj", CONST_OUTPUT_HOOK),
-    "head_value_output": ("layers[%s].self_attn.v_proj", CONST_OUTPUT_HOOK),
+    "head_query_output": ("layers[%s].self_attn.q_proj", CONST_OUTPUT_HOOK, (split_head_and_permute, "n_head")),
+    "head_key_output": ("layers[%s].self_attn.k_proj", CONST_OUTPUT_HOOK, (split_head_and_permute, "n_kv_head")),
+    "head_value_output": ("layers[%s].self_attn.v_proj", CONST_OUTPUT_HOOK, (split_head_and_permute, "n_kv_head")),
 }
 
 
 gemma_type_to_dimension_mapping = {
+    "n_head": ("num_attention_heads",),
+    "n_kv_head": ("num_key_value_heads",),
     "block_input": ("hidden_size",),
     "block_output": ("hidden_size",),
     "mlp_activation": ("intermediate_size",),
     "mlp_output": ("hidden_size",),
     "mlp_input": ("hidden_size",),
     "attention_value_output": ("hidden_size",),
-    "head_attention_value_output": ("hidden_size/num_attention_heads",),
+    "head_attention_value_output": ("head_dim",),
     "attention_output": ("hidden_size",),
     "attention_input": ("hidden_size",),
     "query_output": ("hidden_size",),
     "key_output": ("hidden_size",),
     "value_output": ("hidden_size",),
-    "head_query_output": ("hidden_size/num_attention_heads",),
-    "head_key_output": ("hidden_size/num_attention_heads",),
-    "head_value_output": ("hidden_size/num_attention_heads",),
+    "head_query_output": ("head_dim",),
+    "head_key_output": ("head_dim",),
+    "head_value_output": ("hhead_dim",),
 }
 
 
 """gemma model with LM head"""
 gemma_lm_type_to_module_mapping = {}
 for k, v in gemma_type_to_module_mapping.items():
-    gemma_lm_type_to_module_mapping[k] = (f"model.{v[0]}", v[1])
+    gemma_lm_type_to_module_mapping[k] = (f"model.{v[0]}", ) + v[1:]
 
 
 gemma_lm_type_to_dimension_mapping = gemma_type_to_dimension_mapping
@@ -63,7 +65,7 @@
 """gemma model with classifier head"""
 gemma_classifier_type_to_module_mapping = {}
 for k, v in gemma_type_to_module_mapping.items():
-    gemma_classifier_type_to_module_mapping[k] = (f"model.{v[0]}", v[1])
+    gemma_classifier_type_to_module_mapping[k] = (f"model.{v[0]}", ) + v[1:]
 
 
 gemma_classifier_type_to_dimension_mapping = gemma_type_to_dimension_mapping
diff --git a/pyvene/models/gpt_neo/modelings_intervenable_gpt_neo.py b/pyvene/models/gpt_neo/modelings_intervenable_gpt_neo.py
index 09f94cad..a0f221b6 100644
--- a/pyvene/models/gpt_neo/modelings_intervenable_gpt_neo.py
+++ b/pyvene/models/gpt_neo/modelings_intervenable_gpt_neo.py
@@ -58,7 +58,7 @@
 """gpt_neo model with LM head"""
 gpt_neo_lm_type_to_module_mapping = {}
 for k, v in gpt_neo_type_to_module_mapping.items():
-    gpt_neo_lm_type_to_module_mapping[k] = (f"transformer.{v[0]}", v[1])
+    gpt_neo_lm_type_to_module_mapping[k] = (f"transformer.{v[0]}", ) + v[1:]
 
 
 gpt_neo_lm_type_to_dimension_mapping = gpt_neo_type_to_dimension_mapping
diff --git a/pyvene/models/gpt_neox/modelings_intervenable_gpt_neox.py b/pyvene/models/gpt_neox/modelings_intervenable_gpt_neox.py
index c88eeaa3..666b1fd9 100644
--- a/pyvene/models/gpt_neox/modelings_intervenable_gpt_neox.py
+++ b/pyvene/models/gpt_neox/modelings_intervenable_gpt_neox.py
@@ -33,7 +33,7 @@
 
 
 gpt_neox_type_to_dimension_mapping = {
-    "n_head": "num_attention_heads",
+    "n_head": ("num_attention_heads",),
     "block_input": ("hidden_size",),
     "block_output": ("hidden_size",),
     "mlp_activation": (
@@ -58,7 +58,7 @@
 """gpt_neox model with LM head"""
 gpt_neox_lm_type_to_module_mapping = {}
 for k, v in gpt_neox_type_to_module_mapping.items():
-    gpt_neox_lm_type_to_module_mapping[k] = (f"gpt_neox.{v[0]}", v[1])
+    gpt_neox_lm_type_to_module_mapping[k] = (f"gpt_neox.{v[0]}", ) + v[1:]
 
 
 gpt_neox_lm_type_to_dimension_mapping = gpt_neox_type_to_dimension_mapping
diff --git a/pyvene/models/llama/modelings_intervenable_llama.py b/pyvene/models/llama/modelings_intervenable_llama.py
index 82817b9b..b99b56f8 100644
--- a/pyvene/models/llama/modelings_intervenable_llama.py
+++ b/pyvene/models/llama/modelings_intervenable_llama.py
@@ -20,19 +20,21 @@
     "mlp_output": ("layers[%s].mlp", CONST_OUTPUT_HOOK),
     "mlp_input": ("layers[%s].mlp", CONST_INPUT_HOOK),
     "attention_value_output": ("layers[%s].self_attn.o_proj", CONST_INPUT_HOOK),
-    "head_attention_value_output": ("layers[%s].self_attn.o_proj", CONST_INPUT_HOOK),
+    "head_attention_value_output": ("layers[%s].self_attn.o_proj", CONST_INPUT_HOOK, (split_head_and_permute, "n_head")),
     "attention_output": ("layers[%s].self_attn", CONST_OUTPUT_HOOK),
     "attention_input": ("layers[%s].self_attn", CONST_INPUT_HOOK),
     "query_output": ("layers[%s].self_attn.q_proj", CONST_OUTPUT_HOOK),
     "key_output": ("layers[%s].self_attn.k_proj", CONST_OUTPUT_HOOK),
     "value_output": ("layers[%s].self_attn.v_proj", CONST_OUTPUT_HOOK),
-    "head_query_output": ("layers[%s].self_attn.q_proj", CONST_OUTPUT_HOOK),
-    "head_key_output": ("layers[%s].self_attn.k_proj", CONST_OUTPUT_HOOK),
-    "head_value_output": ("layers[%s].self_attn.v_proj", CONST_OUTPUT_HOOK),
+    "head_query_output": ("layers[%s].self_attn.q_proj", CONST_OUTPUT_HOOK, (split_head_and_permute, "n_head")),
+    "head_key_output": ("layers[%s].self_attn.k_proj", CONST_OUTPUT_HOOK, (split_head_and_permute, "n_kv_head")),
+    "head_value_output": ("layers[%s].self_attn.v_proj", CONST_OUTPUT_HOOK, (split_head_and_permute, "n_kv_head")),
 }
 
 
 llama_type_to_dimension_mapping = {
+    "n_head": ("num_attention_heads",),
+    "n_kv_head": ("num_key_value_heads",),
     "block_input": ("hidden_size",),
     "block_output": ("hidden_size",),
     "mlp_activation": ("intermediate_size",),
@@ -54,7 +56,7 @@
 """llama model with LM head"""
 llama_lm_type_to_module_mapping = {}
 for k, v in llama_type_to_module_mapping.items():
-    llama_lm_type_to_module_mapping[k] = (f"model.{v[0]}", v[1])
+    llama_lm_type_to_module_mapping[k] = (f"model.{v[0]}", ) + v[1:]
 
 
 llama_lm_type_to_dimension_mapping = llama_type_to_dimension_mapping
@@ -63,25 +65,28 @@
 """llama model with classifier head"""
 llama_classifier_type_to_module_mapping = {}
 for k, v in llama_type_to_module_mapping.items():
-    llama_classifier_type_to_module_mapping[k] = (f"model.{v[0]}", v[1])
+    llama_classifier_type_to_module_mapping[k] = (f"model.{v[0]}", ) + v[1:]
 
 
 llama_classifier_type_to_dimension_mapping = llama_type_to_dimension_mapping
 
 
 def create_llama(
-    name="sharpbai/alpaca-7b-merged", cache_dir=None, dtype=torch.bfloat16
+    name="sharpbai/alpaca-7b-merged", cache_dir=None, dtype=torch.bfloat16, config=None
 ):
     """Creates a LLaMA Causal LM model, config, and tokenizer from the given name and revision"""
     from transformers import LlamaForCausalLM, LlamaTokenizer, LlamaConfig
-
-    config = LlamaConfig.from_pretrained(name, cache_dir=cache_dir)
-    tokenizer = LlamaTokenizer.from_pretrained(name, cache_dir=cache_dir)
-    llama = LlamaForCausalLM.from_pretrained(
-        name,
-        config=config,
-        cache_dir=cache_dir,
-        torch_dtype=dtype,  # save memory
-    )
+    if config is None:
+        config = LlamaConfig.from_pretrained(name, cache_dir=cache_dir)
+        llama = LlamaForCausalLM.from_pretrained(
+            name,
+            config=config,
+            cache_dir=cache_dir,
+            torch_dtype=dtype,  # save memory
+        )
+        tokenizer = LlamaTokenizer.from_pretrained(name, cache_dir=cache_dir)
+    else:
+        llama = LlamaForCausalLM(config)
+        tokenizer = LlamaTokenizer.from_pretrained(name, cache_dir=cache_dir)
     print("loaded model")
-    return config, tokenizer, llama
+    return config, tokenizer, llama
\ No newline at end of file
diff --git a/pyvene/models/llava/modelings_intervenable_llava.py b/pyvene/models/llava/modelings_intervenable_llava.py
index c37d019b..41a7ed57 100644
--- a/pyvene/models/llava/modelings_intervenable_llava.py
+++ b/pyvene/models/llava/modelings_intervenable_llava.py
@@ -19,19 +19,21 @@
     "mlp_output": ("language_model.model.layers[%s].mlp", CONST_OUTPUT_HOOK),
     "mlp_input": ("language_model.model.layers[%s].mlp", CONST_INPUT_HOOK),
     "attention_value_output": ("language_model.model.layers[%s].self_attn.o_proj", CONST_INPUT_HOOK),
-    "head_attention_value_output": ("language_model.model.layers[%s].self_attn.o_proj", CONST_INPUT_HOOK),
+    "head_attention_value_output": ("language_model.model.layers[%s].self_attn.o_proj", CONST_INPUT_HOOK, (split_head_and_permute, "n_head")),
     "attention_output": ("language_model.model.layers[%s].self_attn", CONST_OUTPUT_HOOK),
     "attention_input": ("language_model.model.layers[%s].self_attn", CONST_INPUT_HOOK),
     "query_output": ("language_model.model.layers[%s].self_attn.q_proj", CONST_OUTPUT_HOOK),
     "key_output": ("language_model.model.layers[%s].self_attn.k_proj", CONST_OUTPUT_HOOK),
     "value_output": ("language_model.model.layers[%s].self_attn.v_proj", CONST_OUTPUT_HOOK),
-    "head_query_output": ("language_model.model.layers[%s].self_attn.q_proj", CONST_OUTPUT_HOOK),
-    "head_key_output": ("language_model.model.layers[%s].self_attn.k_proj", CONST_OUTPUT_HOOK),
-    "head_value_output": ("language_model.model.layers[%s].self_attn.v_proj", CONST_OUTPUT_HOOK),
+    "head_query_output": ("language_model.model.layers[%s].self_attn.q_proj", CONST_OUTPUT_HOOK, (split_head_and_permute, "n_head")),
+    "head_key_output": ("language_model.model.layers[%s].self_attn.k_proj", CONST_OUTPUT_HOOK, (split_head_and_permute, "n_kv_head")),
+    "head_value_output": ("language_model.model.layers[%s].self_attn.v_proj", CONST_OUTPUT_HOOK, (split_head_and_permute, "n_kv_head")),
 }
 
 
 llava_type_to_dimension_mapping = {
+    "n_head": ("text_config.num_attention_heads",),
+    "n_kv_head": ("text_config.num_key_value_heads",),
     "block_input": ("text_config.hidden_size",),
     "block_output": ("text_config.hidden_size",),
     "mlp_activation": ("text_config.intermediate_size",),
@@ -53,7 +55,7 @@
 """llava model with LM head"""
 llava_lm_type_to_module_mapping = {}
 for k, v in llava_type_to_module_mapping.items():
-    llava_lm_type_to_module_mapping[k] = (f"model.{v[0]}", v[1])
+    llava_lm_type_to_module_mapping[k] = (f"model.{v[0]}", ) + v[1:]
 
 
 llava_lm_type_to_dimension_mapping = llava_type_to_dimension_mapping
@@ -62,7 +64,7 @@
 """llava model with classifier head"""
 llava_classifier_type_to_module_mapping = {}
 for k, v in llava_type_to_module_mapping.items():
-    llava_classifier_type_to_module_mapping[k] = (f"model.{v[0]}", v[1])
+    llava_classifier_type_to_module_mapping[k] = (f"model.{v[0]}", ) + v[1:]
 
 
 llava_classifier_type_to_dimension_mapping = llava_type_to_dimension_mapping
diff --git a/pyvene/models/mistral/modellings_intervenable_mistral.py b/pyvene/models/mistral/modellings_intervenable_mistral.py
index edbd6176..7975edd3 100644
--- a/pyvene/models/mistral/modellings_intervenable_mistral.py
+++ b/pyvene/models/mistral/modellings_intervenable_mistral.py
@@ -20,19 +20,21 @@
     "mlp_output": ("layers[%s].mlp", CONST_OUTPUT_HOOK),
     "mlp_input": ("layers[%s].mlp", CONST_INPUT_HOOK),
     "attention_value_output": ("layers[%s].self_attn.o_proj", CONST_INPUT_HOOK),
-    "head_attention_value_output": ("layers[%s].self_attn.o_proj", CONST_INPUT_HOOK),
+    "head_attention_value_output": ("layers[%s].self_attn.o_proj", CONST_INPUT_HOOK, (split_head_and_permute, "n_head")),
     "attention_output": ("layers[%s].self_attn", CONST_OUTPUT_HOOK),
     "attention_input": ("layers[%s].self_attn", CONST_INPUT_HOOK),
     "query_output": ("layers[%s].self_attn.q_proj", CONST_OUTPUT_HOOK),
     "key_output": ("layers[%s].self_attn.k_proj", CONST_OUTPUT_HOOK),
     "value_output": ("layers[%s].self_attn.v_proj", CONST_OUTPUT_HOOK),
-    "head_query_output": ("layers[%s].self_attn.q_proj", CONST_OUTPUT_HOOK),
-    "head_key_output": ("layers[%s].self_attn.k_proj", CONST_OUTPUT_HOOK),
-    "head_value_output": ("layers[%s].self_attn.v_proj", CONST_OUTPUT_HOOK),
+    "head_query_output": ("layers[%s].self_attn.q_proj", CONST_OUTPUT_HOOK, (split_head_and_permute, "n_head")),
+    "head_key_output": ("layers[%s].self_attn.k_proj", CONST_OUTPUT_HOOK, (split_head_and_permute, "n_kv_head")),
+    "head_value_output": ("layers[%s].self_attn.v_proj", CONST_OUTPUT_HOOK, (split_head_and_permute, "n_kv_head")),
 }
 
 
 mistral_type_to_dimension_mapping = {
+    "n_head": ("num_attention_heads",),
+    "n_kv_head": ("num_key_value_heads",),
     "block_input": ("hidden_size",),
     "block_output": ("hidden_size",),
     "mlp_activation": ("intermediate_size",),
@@ -54,7 +56,7 @@
 """mistral model with LM head"""
 mistral_lm_type_to_module_mapping = {}
 for k, v in mistral_type_to_module_mapping.items():
-    mistral_lm_type_to_module_mapping[k] = (f"model.{v[0]}", v[1])
+    mistral_lm_type_to_module_mapping[k] = (f"model.{v[0]}", ) + v[1:]
 
 
 mistral_lm_type_to_dimension_mapping = mistral_type_to_dimension_mapping


Test patch:
diff --git a/tests/integration_tests/InterventionWithLlamaTestCase.py b/tests/integration_tests/InterventionWithLlamaTestCase.py
new file mode 100644
index 00000000..aa586d24
--- /dev/null
+++ b/tests/integration_tests/InterventionWithLlamaTestCase.py
@@ -0,0 +1,234 @@
+import unittest
+from ..utils import *
+
+
+class InterventionWithLlamaTestCase(unittest.TestCase):
+    @classmethod
+    def setUpClass(self):
+        print("=== Test Suite: InterventionWithLlamaTestCase ===")
+        self.config, self.tokenizer, self.llama = create_llama(
+            config=LlamaConfig(
+                bos_token_id=1,
+                eos_token_id=2,
+                intermediate_size=11008,
+                max_position_embeddings=1024,
+                num_attention_heads=32,
+                num_hidden_layers=4,
+                num_key_value_heads=32,
+                hidden_size=4096,
+                rms_norm_eps=1e-5,
+            )
+        )
+        self.vanilla_block_output_config = IntervenableConfig(
+            model_type=type(self.llama),
+            representations=[
+                RepresentationConfig(
+                    0,
+                    "block_output",
+                    "pos",
+                    1,
+                ),
+            ],
+            intervention_types=VanillaIntervention,
+        )
+        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        self.llama = self.llama.to(self.device)
+
+        self.nonhead_streams = [
+            "block_output",
+            "block_input",
+            "mlp_activation",
+            "mlp_output",
+            "mlp_input",
+            "attention_value_output",
+            "attention_output",
+            "attention_input",
+            "query_output",
+            "key_output",
+            "value_output",
+        ]
+
+        self.head_streams = [
+            "head_attention_value_output",
+            "head_query_output",
+            "head_key_output",
+            "head_value_output",
+        ]
+
+    def _test_with_head_position_intervention(
+        self,
+        intervention_layer,
+        intervention_stream,
+        intervention_type,
+        heads=[0],
+        positions=[0],
+    ):
+        max_position = np.max(np.array(positions))
+        if isinstance(positions[0], list):
+            b_s = len(positions)
+        else:
+            b_s = 10
+        base = {
+            "input_ids": torch.randint(0, 10, (b_s, max_position + 1)).to(self.device)
+        }
+        source = {
+            "input_ids": torch.randint(0, 10, (b_s, max_position + 2)).to(self.device)
+        }
+
+        config = IntervenableConfig(
+            model_type=type(self.llama),
+            representations=[
+                RepresentationConfig(
+                    intervention_layer,
+                    intervention_stream,
+                    "h.pos",
+                    len(positions),
+                )
+            ],
+            intervention_types=intervention_type,
+        )
+        intervenable = IntervenableModel(config, self.llama)
+        intervention = list(intervenable.interventions.values())[0][0]
+
+        base_activations = {}
+        source_activations = {}
+        _ = Llama_RUN(self.llama, base["input_ids"], base_activations, {})
+        _ = Llama_RUN(self.llama, source["input_ids"], source_activations, {})
+        _key = f"{intervention_layer}.{intervention_stream}"
+        if isinstance(heads[0], list):
+            for i, head in enumerate(heads):
+                for h in head:
+                    for position in positions:
+                        base_activations[_key][:, head, position] = intervention(
+                            base_activations[_key][:, head, position],
+                            source_activations[_key][:, head, position],
+                        )
+        else:
+            for head in heads:
+                for position in positions:
+                    base_activations[_key][:, head, position] = intervention(
+                        base_activations[_key][:, head, position],
+                        source_activations[_key][:, head, position],
+                    )
+        golden_out = Llama_RUN(
+            self.llama, base["input_ids"], {}, {_key: base_activations[_key]}
+        )
+
+        if isinstance(positions[0], list):
+            _, out_output = intervenable(
+                base,
+                [source],
+                {"sources->base": ([[heads, positions]], [[heads, positions]])},
+            )
+        else:
+            _, out_output = intervenable(
+                base,
+                [source],
+                {
+                    "sources->base": (
+                        [[[heads] * b_s, [positions] * b_s]],
+                        [[[heads] * b_s, [positions] * b_s]],
+                    )
+                },
+            )
+
+        self.assertTrue(torch.allclose(out_output[0], golden_out))
+
+
+    def test_with_multiple_heads_positions_vanilla_intervention_positive(self):
+        """
+        Multiple head and position with vanilla intervention.
+        """
+        for stream in self.head_streams:
+            print(f"testing stream: {stream} with multiple heads positions")
+            self._test_with_head_position_intervention(
+                intervention_layer=random.randint(0, 3),
+                intervention_stream=stream,
+                intervention_type=VanillaIntervention,
+                heads=[3, 7],
+                positions=[2, 6],
+            )
+
+            self._test_with_head_position_intervention(
+                intervention_layer=random.randint(0, 3),
+                intervention_stream=stream,
+                intervention_type=VanillaIntervention,
+                heads=[4, 1],
+                positions=[7, 2],
+            )
+
+            
+def suite():
+    suite = unittest.TestSuite()
+    suite.addTest(InterventionWithLlamaTestCase("test_clean_run_positive"))
+    suite.addTest(
+        InterventionWithLlamaTestCase(
+            "test_invalid_unit_negative"
+        )
+    )
+    suite.addTest(
+        InterventionWithLlamaTestCase(
+            "test_with_single_position_vanilla_intervention_positive"
+        )
+    )
+    suite.addTest(
+        InterventionWithLlamaTestCase(
+            "test_with_multiple_position_vanilla_intervention_positive"
+        )
+    )
+    suite.addTest(
+        InterventionWithLlamaTestCase(
+            "test_with_complex_position_vanilla_intervention_positive"
+        )
+    )
+    suite.addTest(
+        InterventionWithLlamaTestCase(
+            "test_with_single_head_position_vanilla_intervention_positive"
+        )
+    )
+    suite.addTest(
+        InterventionWithLlamaTestCase(
+            "test_with_multiple_heads_positions_vanilla_intervention_positive"
+        )
+    )
+    suite.addTest(
+        InterventionWithLlamaTestCase(
+            "test_with_use_fast_vanilla_intervention_positive"
+        )
+    )
+    suite.addTest(
+        InterventionWithLlamaTestCase(
+            "test_with_location_broadcast_vanilla_intervention_positive"
+        )
+    ) 
+    suite.addTest(
+        InterventionWithLlamaTestCase(
+            "test_with_position_intervention_constant_source_vanilla_intervention_positive"
+        )
+    ) 
+    suite.addTest(
+        InterventionWithLlamaTestCase(
+            "test_with_position_intervention_constant_source_addition_intervention_positive"
+        )
+    ) 
+    suite.addTest(
+        InterventionWithLlamaTestCase(
+            "test_with_position_intervention_constant_source_subtraction_intervention_positive"
+        )
+    )
+    suite.addTest(
+        InterventionWithLlamaTestCase(
+            "test_with_position_intervention_constant_source_zero_intervention_positive"
+        )
+    ) 
+    suite.addTest(
+        InterventionWithLlamaTestCase(
+            "_test_with_long_sequence_position_intervention_constant_source_positive"
+        )
+    ) 
+    return suite
+
+
+if __name__ == "__main__":
+    runner = unittest.TextTestRunner()
+    runner.run(suite())
\ No newline at end of file
diff --git a/tests/utils.py b/tests/utils.py
index d005a01d..4d6a9076 100644
--- a/tests/utils.py
+++ b/tests/utils.py
@@ -6,7 +6,9 @@
 import os, shutil, torch, random, uuid
 import pandas as pd
 import numpy as np
-from transformers import GPT2Config
+from transformers import GPT2Config, LlamaConfig
+import math
+from torch import nn
 
 
 import subprocess
@@ -19,7 +21,7 @@ def is_package_installed(package_name):
         # Check if package_name is in the result
         return package_name in result.stdout
     except Exception as e:
-        print(f"An error occurred: {e}")
+        print("An error occurred:")
         return False
 
 # Replace 'pyvene' with the name of the package you want to check
@@ -41,6 +43,7 @@ def is_package_installed(package_name):
 from pyvene.models.mlp.modelings_mlp import MLPConfig
 from pyvene.models.mlp.modelings_intervenable_mlp import create_mlp_classifier
 from pyvene.models.gpt2.modelings_intervenable_gpt2 import create_gpt2_lm
+from pyvene.models.llama.modelings_intervenable_llama import create_llama
 
 
 ##################
@@ -77,7 +80,7 @@ def is_package_installed(package_name):
 """
 
 
-def DO_GPT2_INTERVENTION(name, orig_hidden_states, INTERVENTION_ACTIVATIONS):
+def DO_INTERVENTION(name, orig_hidden_states, INTERVENTION_ACTIVATIONS):
     if name in INTERVENTION_ACTIVATIONS:
         return INTERVENTION_ACTIVATIONS[name]
     return orig_hidden_states
@@ -90,26 +93,26 @@ def GPT2_SELF_ATTENTION_RUN(
         self_attn.split_size, dim=2
     )
 
-    query = DO_GPT2_INTERVENTION(f"{i}.query_output", query, INTERVENTION_ACTIVATIONS)
+    query = DO_INTERVENTION(f"{i}.query_output", query, INTERVENTION_ACTIVATIONS)
     CACHE_ACTIVATIONS[f"{i}.query_output"] = query
-    key = DO_GPT2_INTERVENTION(f"{i}.key_output", key, INTERVENTION_ACTIVATIONS)
+    key = DO_INTERVENTION(f"{i}.key_output", key, INTERVENTION_ACTIVATIONS)
     CACHE_ACTIVATIONS[f"{i}.key_output"] = key
-    value = DO_GPT2_INTERVENTION(f"{i}.value_output", value, INTERVENTION_ACTIVATIONS)
+    value = DO_INTERVENTION(f"{i}.value_output", value, INTERVENTION_ACTIVATIONS)
     CACHE_ACTIVATIONS[f"{i}.value_output"] = value
 
     head_query = self_attn._split_heads(query, self_attn.num_heads, self_attn.head_dim)
     head_key = self_attn._split_heads(key, self_attn.num_heads, self_attn.head_dim)
     head_value = self_attn._split_heads(value, self_attn.num_heads, self_attn.head_dim)
 
-    head_query = DO_GPT2_INTERVENTION(
+    head_query = DO_INTERVENTION(
         f"{i}.head_query_output", head_query, INTERVENTION_ACTIVATIONS
     )
     CACHE_ACTIVATIONS[f"{i}.head_query_output"] = head_query
-    head_key = DO_GPT2_INTERVENTION(
+    head_key = DO_INTERVENTION(
         f"{i}.head_key_output", head_key, INTERVENTION_ACTIVATIONS
     )
     CACHE_ACTIVATIONS[f"{i}.head_key_output"] = head_key
-    head_value = DO_GPT2_INTERVENTION(
+    head_value = DO_INTERVENTION(
         f"{i}.head_value_output", head_value, INTERVENTION_ACTIVATIONS
     )
     CACHE_ACTIVATIONS[f"{i}.head_value_output"] = head_value
@@ -118,7 +121,7 @@ def GPT2_SELF_ATTENTION_RUN(
         head_query, head_key, head_value
     )
 
-    head_attention_value_output = DO_GPT2_INTERVENTION(
+    head_attention_value_output = DO_INTERVENTION(
         f"{i}.head_attention_value_output",
         head_attention_value_output,
         INTERVENTION_ACTIVATIONS,
@@ -128,7 +131,7 @@ def GPT2_SELF_ATTENTION_RUN(
     attn_value_output = self_attn._merge_heads(
         head_attention_value_output, self_attn.num_heads, self_attn.head_dim
     )
-    attn_value_output = DO_GPT2_INTERVENTION(
+    attn_value_output = DO_INTERVENTION(
         f"{i}.attention_value_output", attn_value_output, INTERVENTION_ACTIVATIONS
     )
     CACHE_ACTIVATIONS[f"{i}.attention_value_output"] = attn_value_output
@@ -142,7 +145,7 @@ def GPT2_MLP_RUN(mlp, hidden_states, i, CACHE_ACTIVATIONS, INTERVENTION_ACTIVATI
     hidden_states_c_fc = mlp.c_fc(hidden_states)
     hidden_states_act = mlp.act(hidden_states_c_fc)
 
-    hidden_states_act = DO_GPT2_INTERVENTION(
+    hidden_states_act = DO_INTERVENTION(
         f"{i}.mlp_activation", hidden_states_act, INTERVENTION_ACTIVATIONS
     )
     CACHE_ACTIVATIONS[f"{i}.mlp_activation"] = hidden_states_act
@@ -159,7 +162,7 @@ def GPT2_BLOCK_RUN(
     residual = hidden_states
     hidden_states = block.ln_1(hidden_states)
 
-    hidden_states = DO_GPT2_INTERVENTION(
+    hidden_states = DO_INTERVENTION(
         f"{i}.attention_input", hidden_states, INTERVENTION_ACTIVATIONS
     )
     CACHE_ACTIVATIONS[f"{i}.attention_input"] = hidden_states
@@ -168,7 +171,7 @@ def GPT2_BLOCK_RUN(
         block.attn, hidden_states, i, CACHE_ACTIVATIONS, INTERVENTION_ACTIVATIONS
     )
 
-    attn_outputs = DO_GPT2_INTERVENTION(
+    attn_outputs = DO_INTERVENTION(
         f"{i}.attention_output", attn_outputs, INTERVENTION_ACTIVATIONS
     )
     CACHE_ACTIVATIONS[f"{i}.attention_output"] = attn_outputs
@@ -181,7 +184,7 @@ def GPT2_BLOCK_RUN(
     residual = hidden_states
     hidden_states = block.ln_2(hidden_states)
 
-    hidden_states = DO_GPT2_INTERVENTION(
+    hidden_states = DO_INTERVENTION(
         f"{i}.mlp_input", hidden_states, INTERVENTION_ACTIVATIONS
     )
     CACHE_ACTIVATIONS[f"{i}.mlp_input"] = hidden_states
@@ -190,7 +193,7 @@ def GPT2_BLOCK_RUN(
         block.mlp, hidden_states, i, CACHE_ACTIVATIONS, INTERVENTION_ACTIVATIONS
     )
 
-    feed_forward_hidden_states = DO_GPT2_INTERVENTION(
+    feed_forward_hidden_states = DO_INTERVENTION(
         f"{i}.mlp_output", feed_forward_hidden_states, INTERVENTION_ACTIVATIONS
     )
     CACHE_ACTIVATIONS[f"{i}.mlp_output"] = feed_forward_hidden_states
@@ -216,7 +219,7 @@ def GPT2_RUN(gpt2, input_ids, CACHE_ACTIVATIONS, INTERVENTION_ACTIVATIONS):
     hidden_states = inputs_embeds + position_embeds
 
     for i, block in enumerate(gpt2.transformer.h):
-        hidden_states = DO_GPT2_INTERVENTION(
+        hidden_states = DO_INTERVENTION(
             f"{i}.block_input", hidden_states, INTERVENTION_ACTIVATIONS
         )
         CACHE_ACTIVATIONS[f"{i}.block_input"] = hidden_states
@@ -225,7 +228,7 @@ def GPT2_RUN(gpt2, input_ids, CACHE_ACTIVATIONS, INTERVENTION_ACTIVATIONS):
             block, hidden_states, i, CACHE_ACTIVATIONS, INTERVENTION_ACTIVATIONS
         )
 
-        hidden_states = DO_GPT2_INTERVENTION(
+        hidden_states = DO_INTERVENTION(
             f"{i}.block_output", hidden_states, INTERVENTION_ACTIVATIONS
         )
         CACHE_ACTIVATIONS[f"{i}.block_output"] = hidden_states
@@ -233,3 +236,203 @@ def GPT2_RUN(gpt2, input_ids, CACHE_ACTIVATIONS, INTERVENTION_ACTIVATIONS):
     hidden_states = gpt2.transformer.ln_f(hidden_states)
     lm_logits = gpt2.lm_head(hidden_states)
     return lm_logits
+
+"""
+forward calls to fetch activations or run with cached activations for Llama
+"""
+
+def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor:
+    """
+    For Llama-like models' grouped query attention
+    This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch,
+    num_key_value_heads, seqlen, head_dim) to (batch, num_attention_heads, seqlen, head_dim)
+    """
+    batch, num_key_value_heads, slen, head_dim = hidden_states.shape
+    if n_rep == 1:
+        return hidden_states
+    hidden_states = hidden_states[:, :, None, :, :].expand(batch, num_key_value_heads, n_rep, slen, head_dim)
+    return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim)
+
+def rotate_half(x):
+    """Rotates half the hidden dims of the input."""
+    x1 = x[..., : x.shape[-1] // 2]
+    x2 = x[..., x.shape[-1] // 2 :]
+    return torch.cat((-x2, x1), dim=-1)
+
+
+def apply_rotary_pos_emb(q, k, cos, sin, position_ids=None, unsqueeze_dim=1):
+    """Applies Rotary Position Embedding to the query and key tensors.
+
+    Args:
+        q (`torch.Tensor`): The query tensor.
+        k (`torch.Tensor`): The key tensor.
+        cos (`torch.Tensor`): The cosine part of the rotary embedding.
+        sin (`torch.Tensor`): The sine part of the rotary embedding.
+        position_ids (`torch.Tensor`, *optional*):
+            Deprecated and unused.
+        unsqueeze_dim (`int`, *optional*, defaults to 1):
+            The 'unsqueeze_dim' argument specifies the dimension along which to unsqueeze cos[position_ids] and
+            sin[position_ids] so that they can be properly broadcasted to the dimensions of q and k. For example, note
+            that cos[position_ids] and sin[position_ids] have the shape [batch_size, seq_len, head_dim]. Then, if q and
+            k have the shape [batch_size, heads, seq_len, head_dim], then setting unsqueeze_dim=1 makes
+            cos[position_ids] and sin[position_ids] broadcastable to the shapes of q and k. Similarly, if q and k have
+            the shape [batch_size, seq_len, heads, head_dim], then set unsqueeze_dim=2.
+    Returns:
+        `tuple(torch.Tensor)` comprising of the query and key tensors rotated using the Rotary Position Embedding.
+    """
+    cos = cos.unsqueeze(unsqueeze_dim)
+    sin = sin.unsqueeze(unsqueeze_dim)
+    q_embed = (q * cos) + (rotate_half(q) * sin)
+    k_embed = (k * cos) + (rotate_half(k) * sin)
+    return q_embed, k_embed
+
+def Llama_SELF_ATTENTION_RUN(
+    self_attn, hidden_states, i, CACHE_ACTIVATIONS, INTERVENTION_ACTIVATIONS
+):
+    bsz, q_len, _ = hidden_states.size()
+
+    query = self_attn.q_proj(hidden_states)
+    key = self_attn.k_proj(hidden_states)
+    value = self_attn.v_proj(hidden_states)
+
+    query = DO_INTERVENTION(f"{i}.query_output", query, INTERVENTION_ACTIVATIONS)
+    CACHE_ACTIVATIONS[f"{i}.query_output"] = query
+    key = DO_INTERVENTION(f"{i}.key_output", key, INTERVENTION_ACTIVATIONS)
+    CACHE_ACTIVATIONS[f"{i}.key_output"] = key
+    value = DO_INTERVENTION(f"{i}.value_output", value, INTERVENTION_ACTIVATIONS)
+    CACHE_ACTIVATIONS[f"{i}.value_output"] = value
+
+    head_query = query.view(bsz, q_len, self_attn.num_heads, self_attn.head_dim).transpose(1, 2)
+    head_key = key.view(bsz, q_len, self_attn.num_key_value_heads, self_attn.head_dim).transpose(1, 2)
+    head_value = value.view(bsz, q_len, self_attn.num_key_value_heads, self_attn.head_dim).transpose(1, 2)
+
+    head_query = DO_INTERVENTION(
+        f"{i}.head_query_output", head_query, INTERVENTION_ACTIVATIONS
+    )
+    CACHE_ACTIVATIONS[f"{i}.head_query_output"] = head_query
+    head_key = DO_INTERVENTION(
+        f"{i}.head_key_output", head_key, INTERVENTION_ACTIVATIONS
+    )
+    CACHE_ACTIVATIONS[f"{i}.head_key_output"] = head_key
+    head_value = DO_INTERVENTION(
+        f"{i}.head_value_output", head_value, INTERVENTION_ACTIVATIONS
+    )
+    CACHE_ACTIVATIONS[f"{i}.head_value_output"] = head_value
+
+    position_ids = torch.arange(q_len, device=hidden_states.device).repeat(bsz, 1)
+    cos, sin = self_attn.rotary_emb(head_value, position_ids)
+    head_query, head_key = apply_rotary_pos_emb(head_query, head_key, cos, sin)
+
+    head_key = repeat_kv(head_key, self_attn.num_key_value_groups)
+    head_value = repeat_kv(head_value, self_attn.num_key_value_groups)
+    is_causal = True if q_len > 1 else False
+    head_attention_value_output = torch.nn.functional.scaled_dot_product_attention(
+            head_query,
+            head_key,
+            head_value,
+            attn_mask=None,
+            dropout_p=self_attn.attention_dropout if self_attn.training else 0.0,
+            is_causal=is_causal,
+        )
+    head_attention_value_output = DO_INTERVENTION(
+        f"{i}.head_attention_value_output",
+        head_attention_value_output,
+        INTERVENTION_ACTIVATIONS,
+    )
+    CACHE_ACTIVATIONS[f"{i}.head_attention_value_output"] = head_attention_value_output
+    attn_value_output = head_attention_value_output.transpose(1, 2).contiguous().reshape(bsz, q_len, self_attn.hidden_size)
+    attn_value_output = DO_INTERVENTION(
+        f"{i}.attention_value_output", attn_value_output, INTERVENTION_ACTIVATIONS
+    )
+    CACHE_ACTIVATIONS[f"{i}.attention_value_output"] = attn_value_output
+
+    attn_output = self_attn.o_proj(attn_value_output)
+
+    return attn_output
+
+def Llama_MLP_RUN(mlp, hidden_states, i, CACHE_ACTIVATIONS, INTERVENTION_ACTIVATIONS):
+    hidden_states_gate_proj = mlp.gate_proj(hidden_states)
+    hidden_states_act = mlp.act_fn(hidden_states_gate_proj)
+
+    hidden_states_act = DO_INTERVENTION(
+        f"{i}.mlp_activation", hidden_states_act, INTERVENTION_ACTIVATIONS
+    )
+    CACHE_ACTIVATIONS[f"{i}.mlp_activation"] = hidden_states_act
+
+    hidden_states_up_proj = mlp.up_proj(hidden_states)
+    hidden_states_down_proj = mlp.down_proj(hidden_states_act * hidden_states_up_proj)
+    return hidden_states_down_proj
+
+def Llama_BLOCK_RUN(
+    block, hidden_states, i, CACHE_ACTIVATIONS, INTERVENTION_ACTIVATIONS
+):
+    # self attention + residual
+    residual = hidden_states
+    hidden_states = block.input_layernorm(hidden_states)
+
+    hidden_states = DO_INTERVENTION(
+        f"{i}.attention_input", hidden_states, INTERVENTION_ACTIVATIONS
+    )
+    CACHE_ACTIVATIONS[f"{i}.attention_input"] = hidden_states
+
+    attn_outputs = Llama_SELF_ATTENTION_RUN(
+        block.self_attn, hidden_states, i, CACHE_ACTIVATIONS, INTERVENTION_ACTIVATIONS
+    )
+
+    attn_outputs = DO_INTERVENTION(
+        f"{i}.attention_output", attn_outputs, INTERVENTION_ACTIVATIONS
+    )
+    CACHE_ACTIVATIONS[f"{i}.attention_output"] = attn_outputs
+
+    attn_output = attn_outputs
+    # residual connection
+    hidden_states = attn_output + residual
+
+    # mlp + residual
+    residual = hidden_states
+    hidden_states = block.post_attention_layernorm(hidden_states)
+
+    hidden_states = DO_INTERVENTION(
+        f"{i}.mlp_input", hidden_states, INTERVENTION_ACTIVATIONS
+    )
+    CACHE_ACTIVATIONS[f"{i}.mlp_input"] = hidden_states
+
+    feed_forward_hidden_states = Llama_MLP_RUN(
+        block.mlp, hidden_states, i, CACHE_ACTIVATIONS, INTERVENTION_ACTIVATIONS
+    )
+
+    feed_forward_hidden_states = DO_INTERVENTION(
+        f"{i}.mlp_output", feed_forward_hidden_states, INTERVENTION_ACTIVATIONS
+    )
+    CACHE_ACTIVATIONS[f"{i}.mlp_output"] = feed_forward_hidden_states
+
+    # residual connection
+    hidden_states = residual + feed_forward_hidden_states
+
+    return hidden_states
+
+def Llama_RUN(llama, input_ids, CACHE_ACTIVATIONS, INTERVENTION_ACTIVATIONS):
+    """
+    We basically explicitly do the LLama forward here.
+    """
+    # embed
+    inputs_embeds = llama.model.embed_tokens(input_ids)
+    hidden_states = inputs_embeds
+    for i, block in enumerate(llama.model.layers):
+        hidden_states = DO_INTERVENTION(
+            f"{i}.block_input", hidden_states, INTERVENTION_ACTIVATIONS
+        )
+        CACHE_ACTIVATIONS[f"{i}.block_input"] = hidden_states
+
+        hidden_states = Llama_BLOCK_RUN(
+            block, hidden_states, i, CACHE_ACTIVATIONS, INTERVENTION_ACTIVATIONS
+        )
+
+        hidden_states = DO_INTERVENTION(
+            f"{i}.block_output", hidden_states, INTERVENTION_ACTIVATIONS
+        )
+        CACHE_ACTIVATIONS[f"{i}.block_output"] = hidden_states
+
+    hidden_states = llama.model.norm(hidden_states)
+    lm_logits = llama.lm_head(hidden_states).float()
+    return lm_logits
\ No newline at end of file