Find decoder layers through PEFT wrappers and other layouts

AGENTS.md

@@ -13,7 +13,7 @@ A self-steering agent library: gives any local HF model the ability to steer its
 ```python
 from modulo import SteeringAgent
 
-agent = SteeringAgent(model, tokenizer)  # HF causal LM with model.model.layers
+agent = SteeringAgent(model, tokenizer)  # HF causal LM, plain or PEFT-wrapped
 result = agent.step(messages)             # -> StepResult
 messages += [{"role": "assistant", "content": result.raw_output}]
 messages += result.tool_messages()        # tool results back to the model
@@ -412,4 +412,4 @@ Run the GPU tests after touching anything in the steering path (`_SteeringHooks`
 - **An index vector doesn't load**: Check the `[registry]` lines printed at startup. A sha256 mismatch on a cached file means the file in `vectors/<model_slug>/` isn't the one the index pins; delete it to re-download. Unverified entries also need `include_community=True`.
 - **Steering seems to have no effect**: Run `python verify_vector.py --vector <name> --strength 12`. If the ± outputs match the baseline, the vector (not the steering) is the problem.
 - **Model downloaded twice**: `HF_HOME` is set to `/workspace/hf_cache` when `/workspace` exists, which only works if `config` is imported before `transformers` (the entry points do this). Older runs may have left a copy in `~/.cache/huggingface`.
-- **`model.model.layers` not found**: The library assumes standard HF architecture. Custom model classes may need adapter code.
+- **"Can't find the decoder layers"**: `modulo.layers.decoder_layers` looks through PEFT wrappers and tries the usual paths (`model.layers`, `model.language_model.layers`, `transformer.h`, `gpt_neox.layers`, ...), then any `ModuleList` as long as the config's `num_hidden_layers`. A custom model class that matches none of these needs its path added to `LAYER_PATHS`.

README.md

@@ -42,7 +42,7 @@ Probably not well. When we tested this, models said a vector had worked about as
 ```python
 from modulo import SteeringAgent
 
-agent = SteeringAgent(model, tokenizer)  # HF causal LM with model.model.layers (Llama/Qwen/Mistral-style)
+agent = SteeringAgent(model, tokenizer)  # HF causal LM (Llama/Qwen/Mistral-style), plain or PEFT-wrapped
 result = agent.step(messages)  # -> StepResult(prose, raw_output, tool_calls, tool_results, ...)
 messages += [{"role": "assistant", "content": result.raw_output}]
 messages += result.tool_messages()  # tool results back to the model

modulo/__init__.py

@@ -9,7 +9,7 @@
 
 Quick start:
     from modulo import SteeringAgent
-    agent = SteeringAgent(model, tokenizer)  # HF causal LM with model.model.layers
+    agent = SteeringAgent(model, tokenizer)  # HF causal LM, plain or PEFT-wrapped
     result = agent.step(messages)            # -> StepResult
     messages += [{"role": "assistant", "content": result.raw_output}]
     messages += result.tool_messages()       # tool results back to the model

modulo/agent.py

@@ -26,6 +26,7 @@
 import torch
 
 from modulo.chat import chat_prompt, encode
+from modulo.layers import decoder_layers, num_layers
 from modulo.registry import VectorRegistry, VectorEntry, DEFAULT_VECTORS_DIR
 from modulo.reports import make_report, new_session_id
 
@@ -313,7 +314,7 @@ def hook(module, inputs, output):
         return hook
 
     def __enter__(self):
-        layers = self.model.model.layers
+        layers = decoder_layers(self.model)
         for layer, vec in self.layer_vectors.items():
             self.handles.append(layers[layer].register_forward_hook(self._make_hook(vec)))
         return self
@@ -413,7 +414,8 @@ def __exit__(self, *exc):
 class SteeringAgent:
     """Gives a model the ability to steer its own activations.
 
-    Works with HF causal LMs that expose decoder layers as model.model.layers
+    Works with HF causal LMs, PEFT-wrapped or not, whose decoder layers
+    modulo.layers.decoder_layers can find (model.model.layers for most)
     (Llama/Qwen/Mistral-style architectures).
     """
 
@@ -425,7 +427,7 @@ def __init__(self, model, tokenizer, registry: Optional[VectorRegistry] = None,
                  pass_through_unknown_tools: bool = False):
         """
         Args:
-            model: HF causal LM (must have model.model.layers)
+            model: HF causal LM, optionally PEFT-wrapped (see modulo.layers)
             tokenizer: matching tokenizer
             registry: pre-built VectorRegistry. If None, auto-creates one for
                       the model (community index vectors, else auto-extracts).
@@ -585,7 +587,7 @@ def create(self, name: str, more: str, less: str) -> tuple[bool, str]:
             return False, f"'{name}' is active. Drop it before re-creating it."
         if name not in self.created and len(self.created) >= MAX_CREATED_VECTORS:
             return False, f"You have already made {MAX_CREATED_VECTORS} vectors this session, the limit."
-        layer = self.model.config.num_hidden_layers // 2
+        layer = num_layers(self.model) // 2
         axis, half = extract_vector(self.model, self.tokenizer, more, less, layer,
                                     questions=EXTRACTION_QUESTIONS, batch_size=2 * len(EXTRACTION_QUESTIONS),
                                     return_halves=True)
@@ -701,7 +703,7 @@ def hook(module, inp, output):
                 h = output[0] if isinstance(output, tuple) else output
                 captured["h"] = h[0].float()
 
-            handle = self.model.model.layers[layer].register_forward_hook(hook)
+            handle = decoder_layers(self.model)[layer].register_forward_hook(hook)
             try:
                 with torch.inference_mode():
                     self.model(out)

modulo/cli.py

@@ -193,8 +193,9 @@ def say(line=""):
 
 def cmd_extract(args) -> int:
     from modulo.extraction import EXTRACTION_QUESTIONS, extract_vector, save_and_register
+    from modulo.layers import num_layers
     model, tokenizer = _load_model(args.model)
-    layer = args.layer if args.layer is not None else model.config.num_hidden_layers // 2
+    layer = args.layer if args.layer is not None else num_layers(model) // 2
     questions = EXTRACTION_QUESTIONS[: args.n_questions]
     print(f"Extracting '{args.name}' at layer {layer} from {len(questions)} questions...", file=sys.stderr)
     axis = extract_vector(model, tokenizer, args.prompt_a, args.prompt_b, layer, questions, read=args.read)

modulo/extraction.py

@@ -16,6 +16,7 @@
 import torch
 
 from modulo.chat import chat_prompt, encode
+from modulo.layers import decoder_layers, num_layers
 from modulo.registry import VectorEntry, VectorRegistry
 
 
@@ -227,7 +228,7 @@ def hook(module, input, output):
         resp = act[0, prompt_len:].float()
         activation = (resp.mean(0) if read == "mean" else resp[-1]).detach().clone()
 
-    handle = model.model.layers[layer].register_forward_hook(hook)
+    handle = decoder_layers(model)[layer].register_forward_hook(hook)
     try:
         with torch.inference_mode():
             model(**full_inputs)
@@ -281,7 +282,7 @@ def hook(module, input, output):
         nonlocal acts
         acts = (output[0] if isinstance(output, tuple) else output)[:, width:].float()
 
-    handle = model.model.layers[layer].register_forward_hook(hook)
+    handle = decoder_layers(model)[layer].register_forward_hook(hook)
     try:
         with torch.inference_mode():
             model(input_ids=out, attention_mask=mask, position_ids=position_ids)
@@ -299,7 +300,7 @@ def extract_vector(model, tokenizer, prompt_a, prompt_b, layer, questions=None,
     """Extract a steering vector: axis = mean(activations under A) - mean(activations under B).
 
     Args:
-        model: HF causal LM with model.model.layers (standard HF architecture)
+        model: HF causal LM, optionally PEFT-wrapped (see modulo.layers)
         tokenizer: matching tokenizer
         prompt_a: system prompt for the positive direction
         prompt_b: system prompt for the negative direction
@@ -363,7 +364,7 @@ def extract_auto(model, tokenizer, model_name: str, vectors_dir: Path, registry,
     ...) using contrastive system prompts. Saves them and adds them to the
     registry. With only_missing=True, skips names the registry already has.
     """
-    n_layers = model.config.num_hidden_layers
+    n_layers = num_layers(model)
     if layer is None:
         layer = n_layers // 2  # mid-layer default
     specs = [s for s in STARTER_VECTORS if not (only_missing and registry.get(s["name"]))]

modulo/layers.py

@@ -0,0 +1,74 @@
+"""Finding a model's decoder layers, wherever its architecture or a wrapper puts them.
+
+Modulo hooks the output of each decoder layer. Most HF causal LMs keep those at
+model.model.layers, but a PEFT-wrapped model (e.g. one carrying an Afterlife
+LoRA adapter) keeps them at model.base_model.model.model.layers, multimodal
+models nest them under a language model, and older architectures use other
+names (transformer.h, gpt_neox.layers).
+"""
+from typing import Optional
+
+import torch.nn as nn
+
+# Tried in order, relative to the unwrapped model.
+LAYER_PATHS = (
+    "model.layers",                  # Llama, Qwen, Mistral, Gemma, OLMo, Phi-3, ...
+    "model.language_model.layers",   # newer multimodal wrappers (Gemma 3, Qwen2.5-VL, Llava)
+    "language_model.model.layers",   # older multimodal wrappers
+    "model.decoder.layers",          # OPT
+    "transformer.h",                 # GPT-2, Falcon, BLOOM
+    "gpt_neox.layers",               # Pythia, GPT-NeoX
+    "layers",                        # a bare decoder model
+)
+
+
+def unwrap(model):
+    """The underlying HF model, through PEFT wrappers (PeftModel.get_base_model)."""
+    for _ in range(4):  # wrappers don't nest deeper than this in practice
+        get_base = getattr(model, "get_base_model", None)
+        if not callable(get_base):
+            break
+        inner = get_base()
+        if inner is model:
+            break
+        model = inner
+    return model
+
+
+def _resolve(obj, path: str):
+    for part in path.split("."):
+        obj = getattr(obj, part, None)
+        if obj is None:
+            return None
+    return obj
+
+
+def _config_layers(model) -> Optional[int]:
+    config = getattr(model, "config", None)
+    for cfg in (config, getattr(config, "text_config", None)):
+        n = getattr(cfg, "num_hidden_layers", None)
+        if isinstance(n, int):
+            return n
+    return None
+
+
+def decoder_layers(model) -> nn.ModuleList:
+    """The model's decoder layers, in order. Raises ValueError if none are found."""
+    base = unwrap(model)
+    for path in LAYER_PATHS:
+        layers = _resolve(base, path)
+        if isinstance(layers, nn.ModuleList) and len(layers) > 0:
+            return layers
+    # Fall back to the ModuleList whose length matches the config's layer count.
+    n = _config_layers(base)
+    if n:
+        for _, module in base.named_modules():
+            if isinstance(module, nn.ModuleList) and len(module) == n:
+                return module
+    raise ValueError(f"Can't find the decoder layers of {type(base).__name__}; tried {', '.join(LAYER_PATHS)}")
+
+
+def num_layers(model) -> int:
+    """How many decoder layers the model has (from its config, else by counting)."""
+    n = _config_layers(unwrap(model))
+    return n if n is not None else len(decoder_layers(model))

tests/test_layers.py

@@ -0,0 +1,76 @@
+"""Finding decoder layers on plain, PEFT-wrapped and other architectures (CPU, tiny random models)."""
+from types import SimpleNamespace
+
+import pytest
+import torch
+import torch.nn as nn
+
+from modulo.layers import decoder_layers, num_layers, unwrap
+
+transformers = pytest.importorskip("transformers")
+
+
+def tiny_qwen():
+    from transformers import Qwen2Config, Qwen2ForCausalLM
+    torch.manual_seed(0)
+    cfg = Qwen2Config(vocab_size=64, hidden_size=16, intermediate_size=32, num_hidden_layers=3,
+                      num_attention_heads=2, num_key_value_heads=1, max_position_embeddings=64)
+    return Qwen2ForCausalLM(cfg).eval()
+
+
+def test_plain_model():
+    model = tiny_qwen()
+    assert decoder_layers(model) is model.model.layers
+    assert num_layers(model) == 3
+
+
+def test_peft_wrapped_model_finds_the_same_layers():
+    peft = pytest.importorskip("peft")
+    model = tiny_qwen()
+    layers = model.model.layers
+    wrapped = peft.get_peft_model(model, peft.LoraConfig(r=2, target_modules=["q_proj", "v_proj"]))
+    assert unwrap(wrapped) is model
+    assert decoder_layers(wrapped) is layers
+    assert num_layers(wrapped) == 3
+
+
+def test_hook_on_peft_layers_changes_output():
+    """A steering hook on the PEFT-wrapped model's layers reaches the logits."""
+    peft = pytest.importorskip("peft")
+    wrapped = peft.get_peft_model(tiny_qwen(), peft.LoraConfig(r=2, target_modules=["q_proj"]))
+    ids = torch.tensor([[1, 2, 3, 4]])
+    with torch.no_grad():
+        base = wrapped(ids).logits
+
+        def hook(module, inputs, output):
+            h = output[0] if isinstance(output, tuple) else output
+            h += 5.0
+            return output
+
+        handle = decoder_layers(wrapped)[1].register_forward_hook(hook)
+        try:
+            steered = wrapped(ids).logits
+        finally:
+            handle.remove()
+    assert not torch.allclose(base, steered)
+
+
+def test_other_layouts():
+    gpt2 = SimpleNamespace(transformer=SimpleNamespace(h=nn.ModuleList([nn.Linear(2, 2)] * 2)))
+    assert len(decoder_layers(gpt2)) == 2
+    mm = SimpleNamespace(model=SimpleNamespace(language_model=SimpleNamespace(layers=nn.ModuleList([nn.Linear(2, 2)]))),
+                         config=SimpleNamespace(text_config=SimpleNamespace(num_hidden_layers=1)))
+    assert len(decoder_layers(mm)) == 1 and num_layers(mm) == 1
+
+
+def test_fallback_by_layer_count_and_failure():
+    class Odd(nn.Module):
+        def __init__(self, n):
+            super().__init__()
+            self.blocks = nn.ModuleList([nn.Linear(2, 2) for _ in range(n)])
+            self.config = SimpleNamespace(num_hidden_layers=n)
+    odd = Odd(4)
+    assert decoder_layers(odd) is odd.blocks
+    odd.config.num_hidden_layers = 5
+    with pytest.raises(ValueError, match="Can't find the decoder layers"):
+        decoder_layers(odd)

tests/test_steering_gpu.py

@@ -16,6 +16,7 @@
 from load_model import load_model_and_tokenizer
 from modulo import SteeringAgent
 from modulo.chat import chat_prompt, encode
+from modulo.layers import decoder_layers
 
 VECTOR = "assistant_axis"
 STRENGTH = 12.0
@@ -68,7 +69,7 @@ def test_steering_only_touches_last_prompt_position_onward(agent):
     # before the steering hook (registration order) and see unsteered output.
     layer = agent.registry.get(VECTOR).layer + 1
     captured = []
-    h = agent.model.model.layers[layer].register_forward_hook(
+    h = decoder_layers(agent.model)[layer].register_forward_hook(
         lambda m, i, o: captured.append((o[0] if isinstance(o, tuple) else o)[0].float().clone()))
     try:
         with torch.inference_mode():
@@ -135,7 +136,7 @@ def test_residual_norm_matches_generated_tokens(agent):
     n = inputs.input_ids.shape[1]
     out = agent._generate_ids(inputs, 40, temperature=0, top_p=1.0)
     captured = {}
-    handle = agent.model.model.layers[layer].register_forward_hook(
+    handle = decoder_layers(agent.model)[layer].register_forward_hook(
         lambda m, i, o: captured.__setitem__("h", (o[0] if isinstance(o, tuple) else o)[0].float()))
     try:
         with torch.inference_mode():