diff --git a/AGENTS.md b/AGENTS.md index 280176d..b866e9d 100644 --- a/AGENTS.md +++ b/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//` 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 --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`. diff --git a/README.md b/README.md index 129973d..5dc17ae 100644 --- a/README.md +++ b/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 diff --git a/modulo/__init__.py b/modulo/__init__.py index 66f08a2..7eb2808 100644 --- a/modulo/__init__.py +++ b/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 diff --git a/modulo/agent.py b/modulo/agent.py index 2956c10..80df2e2 100644 --- a/modulo/agent.py +++ b/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) diff --git a/modulo/cli.py b/modulo/cli.py index e663b5c..c96fb86 100644 --- a/modulo/cli.py +++ b/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) diff --git a/modulo/extraction.py b/modulo/extraction.py index 329e434..06a6609 100644 --- a/modulo/extraction.py +++ b/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"]))] diff --git a/modulo/layers.py b/modulo/layers.py new file mode 100644 index 0000000..8940f3b --- /dev/null +++ b/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)) diff --git a/tests/test_layers.py b/tests/test_layers.py new file mode 100644 index 0000000..cf40dca --- /dev/null +++ b/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) diff --git a/tests/test_steering_gpu.py b/tests/test_steering_gpu.py index aa7e65e..2d09d39 100644 --- a/tests/test_steering_gpu.py +++ b/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():