Skip to content
Murano

Model

MuranoModel: nnterp-based model wrapper.

MuranoModel

Section titled “MuranoModel”

Thin wrapper around nnterp StandardizedTransformer for mechanistic interpretability.

Provides cross-architecture access to model layers, tokenizer, and metadata. All analysis logic lives in pipeline steps, not here.

Args:

  • model_id: HuggingFace model identifier.
  • device_map: Device placement strategy.
  • dtype: Model weight dtype.
  • enable_attention_probs: If True, load with eager attention so nnterp can
  • expose the per-head softmax attention weights via
  • “: attr:attention_probabilities. Off by default because eager
  • attention is slower and heavier than the fused SDPA/flash kernels;
  • turn it on only for attention-pattern analysis or intervention.
  • loader_kwargs: Extra keyword arguments forwarded to the nnterp loader
  • (e.g. attn_implementation="eager", check_renaming=False).
  • GPT-J needs both: its fused SDPA path traces to symbolic tensors, and
  • nnterp’s renaming-validation scan trips a data-dependent op under
  • fake tensors, so it loads with those two set.

Example:

  • model = MuranoModel(“meta-llama/Llama-3.2-1B-Instruct”)
  • print(model.n_layers, model.d_model)

__init__

Section titled “__init__”
def __init__(self, model_id: str, device_map: str = 'auto', dtype: TorchDtype = bfloat16, enable_attention_probs: bool = False, loader_kwargs: Any = {}):

layer

Section titled “layer”
def layer(self, idx: int):

Return the nnterp module proxy for a decoder layer.

resolve_module

Section titled “resolve_module”
def resolve_module(self, layer_idx: int, module: str):

Resolve a submodule proxy by name at a given layer.

Args:

  • layer_idx: Decoder layer index.
  • module: Module name to resolve, e.g. "residual", "mlp", or
  • a dotted path like "mlp.gate_proj".

Returns:

  • nnsight proxy for the requested submodule.

Raises:

  • ValueError: If any part of the dotted path does not exist.

attn_out_proj

Section titled “attn_out_proj”
def attn_out_proj(self, layer_idx: int, module: str):

Resolve an attention module’s output projection for per-head capture.

The input to this projection is the concatenated per-head outputs, so callers reshape it to recover per-head activations.

Args:

  • layer_idx: Decoder layer index.
  • module: Module name expected to resolve to an attention module.

Returns:

  • nnsight proxy for the attention output projection.

Raises:

  • NotImplementedError: If the module exposes no known output
  • projection (e.g. it is not attention, or the architecture is
  • unsupported for per-head capture).

raw_layer

Section titled “raw_layer”
def raw_layer(self, idx: int):

Return the raw torch module for decoder layer idx.

The resolvers above return nnsight proxies for tracing; this returns the module itself, for native torch hooks or weight access.

raw_module

Section titled “raw_module”
def raw_module(self, layer_idx: int, module: str):

Return the raw torch module named module at layer_idx.

raw_attn_out_proj

Section titled “raw_attn_out_proj”
def raw_attn_out_proj(self, layer_idx: int, module: str):

Return the raw torch module of the attention output projection.

project_on_vocab

Section titled “project_on_vocab”
def project_on_vocab(self, hidden: Tensor) -> Tensor:

Project hidden states onto the vocabulary.

Applies the standardized final norm and unembedding, lm_head(ln_final(hidden)), matching the logit-lens computation. hidden is cast to the unembedding’s device and dtype first, so a direction stored elsewhere (e.g. an fp32 SAE feature on CPU) projects cleanly against a model on GPU.

Args:

  • hidden: Hidden states [..., d_model].

Returns:

  • Vocabulary logits [..., vocab_size].

require_attention_probs

Section titled “require_attention_probs”
def require_attention_probs(self) -> None:

Raise if attention weights were not enabled at load.

Raises:

  • RuntimeError: If the model was loaded without
  • enable_attention_probs=True, so no attention-pattern hook
  • exists to read or write.

forward_logits_attention

Section titled “forward_logits_attention”
def forward_logits_attention(self, tokens: Any, edits: dict[int, tuple[Tensor, Tensor]]) -> Tensor:

Run a forward pass overwriting per-head attention weights, return logits.

The attention-pattern analogue of :meth:forward_logits: it edits the per-head softmax weights (which :meth:forward_logits cannot reach) through nnterp’s settable :attr:attention_probabilities accessor, then reads the model’s output logits under that intervention.

Each layer’s weights become pattern * (1 - mask) + replacement * mask, so only the masked heads and query positions change and every other head, position, and layer keeps its clean pattern.

Args:

  • tokens: Tokenizer output (or anything :attr:trace accepts) to run.
  • edits: {layer: (replacement, mask)}; both tensors must already sit
  • on: attr:device and broadcast to the “[batch, n_heads, query,
  • key]“ pattern.

Returns:

  • Output logits [batch, seq, vocab_size] on CPU as float32.

forward_logits

Section titled “forward_logits”
def forward_logits(self, tokens: Any, fn: Callable[[Tensor, Node], Tensor] | None = None, layers: list[int] | str = 'all', modules: str | list[str] = 'residual', per_head: bool = False) -> Tensor:

Run a forward pass and return the model’s output logits [B, S, V].

Returns the model’s true unembedding output via nnterp’s standardized .logits accessor, not a re-projection of a captured hidden state (contrast :meth:project_on_vocab). The result is detached, cast to float32, and moved to CPU to match the other stores and to keep bf16/fp16 models safe for downstream cross-entropy and argmax.

With fn given, the pass is intervened: fn rewrites each target module’s activation before the logits are read, so an ablation or patch flows through to the output (the forward-pass analogue of :meth:generate_with_hooks). fn=None is the plain forward pass.

Args:

  • tokens: Tokenizer output (or anything :attr:trace accepts) to run.
  • fn: (activation, key) -> activation applied to each target
  • module's activation, keyed by a: class:Node at
  • (layer, module); None runs unmodified. The function is
  • expected to no-op on sites it does not target.
  • layers: Layer indices to hook, or "all". Only used when fn
  • is given.
  • modules: Module name(s) to hook at each layer. Only used when fn
  • is given.
  • per_head: If True, hook the attention output projection’s input and
  • pass fn the per-head activation [B, S, n_heads, head_dim]
  • (the concatenated head outputs), so fn can rewrite a single
  • head’s slice. The rewritten tensor is reshaped back before the
  • projection runs. Requires tokens to expose input_ids for
  • the batch and sequence dimensions.

Returns:

  • Output logits [batch, seq, vocab_size] on CPU as float32.

generate_with_hooks

Section titled “generate_with_hooks”
def generate_with_hooks(self, text: str, fn: Callable[[Tensor, Node], Tensor] | None = None, layers: list[int] | str = 'all', modules: str | list[str] = 'residual', gen_kwargs: dict[str, Any] | None = None) -> str:

Generate from text, optionally applying fn per layer/module.

Args:

  • text: Prompt to generate from.
  • fn: (activation, key) -> activation applied to each target
  • module’s output on every generated token; None runs
  • unmodified.
  • layers: Layer indices to hook, or "all".
  • modules: Module name(s) to hook at each layer.
  • gen_kwargs: Forwarded to the underlying generation call.

Returns:

  • The decoded continuation, excluding the prompt.

record

Section titled “record”
def record(self, text: str | Sequence[str], layers: list[int] | str = 'all', modules: str | list[str] = 'residual', position: str | int = 'last', batch_size: int = 8, per_head: bool = False) -> ActivationStore:

Record activations on one or more texts.

Args:

  • text: Single string or sequence of strings to record from.
  • layers: Layer indices to record at, or "all" for every layer.
  • modules: Module name(s) to record at each layer.
  • position: Token position to record. One of "last", "first",
  • "mean", an integer index, or "none" to keep every
  • position.
  • batch_size: Forward-pass batch size.
  • per_head: If True, split attention activations per head (attention
  • modules only).

Returns:

  • ActivationStore with per-layer activations under positive;
  • negative is empty since this is a single-class call.

logits

Section titled “logits”
def logits(self, text: str | Sequence[str]) -> Tensor:

Return the model’s output logits [B, S, V] for one or more texts.

Quick-API counterpart to :meth:record and :meth:generate: tokenizes text and runs a single forward pass through the Logits step.

Args:

  • text: A single string or a sequence of strings.

Returns:

  • Output logits [batch, seq, vocab_size] on CPU as float32, one
  • row per input text.

find_direction

Section titled “find_direction”
def find_direction(self, positive: Sequence[str], negative: Sequence[str], layers: list[int] | str = 'all', modules: str | list[str] = 'residual', position: str | int = 'last', batch_size: int = 8, normalize: bool = True) -> SteeringResult:

Find a contrastive steering direction between two text sets.

Args:

  • positive: Texts in the positive class.
  • negative: Texts in the negative class.
  • layers: Layer indices to record at, or "all" for every layer.
  • modules: Module name(s) to record at each layer.
  • position: Token position to record. One of "last", "first",
  • "mean", or an integer index.
  • batch_size: Forward-pass batch size.
  • normalize: If True, normalize each per-layer direction to unit norm.

Returns:

  • SteeringResult with one direction per layer plus the best-scoring
  • layer.

generate

Section titled “generate”
def generate(self, text: str | Sequence[str], ablate: Any | None = None, steer: tuple[Any, float] | None = None, layers: list[int] | str = 'all', modules: str | list[str] = 'residual', gen_kwargs: dict[str, Any] | None = None) -> str | list[str]:

Generate text, optionally with activation-space steering or ablation.

Args:

  • text: Single prompt or sequence of prompts.
  • ablate: SteeringResult or {layer: tensor} mapping; if given, the
  • direction is projected out of the residual stream at each
  • target layer during generation.
  • steer: (direction_like, alpha) tuple; if given, “alpha *
  • direction“ is added to the residual stream at each target
  • layer. Pass either ablate or steer, not both.
  • layers: Layer indices to apply the intervention at, or "all".
  • modules: Module name(s) to apply the intervention at each layer.
  • gen_kwargs: Forwarded to the underlying generation call. Defaults
  • to ``{"max_new_tokens": 256, “do_sample”: False}“.

Returns:

  • A single string when text is a single string, otherwise a list.

Raises:

  • ValueError: If both ablate and steer are passed.

chat_template

Section titled “chat_template”
def chat_template(self, messages: list[dict]) -> str:

Apply the tokenizer’s chat template to a list of messages.

Args:

  • messages: List of message dicts with role and content keys.

Returns:

  • The rendered prompt string with a generation prompt appended.