Info

Last Execution: 2026-02-17

Package	Version
nnsight	0.5.15
Python	3.12.3
torch	2.10.0+cu128
transformers	5.2.0

Module Access¶

Inside a tracing context, you can call any module as a function on arbitrary inputs. This runs the module's forward method directly, without triggering nnsight's hooks or requiring execution order.

Setup¶

In [1]:

from nnsight import LanguageModel
import torch

model = LanguageModel("openai-community/gpt2", device_map="auto", dispatch=True)

from nnsight import LanguageModel import torch model = LanguageModel("openai-community/gpt2", device_map="auto", dispatch=True)

Calling a Module as a Function¶

Use any module as a callable to run its computation on values you already have. This calls .forward() directly and does not need to follow execution order.

In [2]:

with model.trace("The Eiffel Tower is in the city of"):

    hs = model.transformer.h[-1].output[0]

    # Apply layer norm then lm_head to decode hidden states
    logits = model.lm_head(model.transformer.ln_f(hs))
    token = logits[0, -1].argmax(dim=-1).save()

print(f"Decoded: {model.tokenizer.decode(token)}")

with model.trace("The Eiffel Tower is in the city of"): hs = model.transformer.h[-1].output[0] # Apply layer norm then lm_head to decode hidden states logits = model.lm_head(model.transformer.ln_f(hs)) token = logits[0, -1].argmax(dim=-1).save() print(f"Decoded: {model.tokenizer.decode(token)}")

Decoded:  Paris

Function calls vs attribute access

When you call a module as a function (e.g., model.lm_head(x)), it runs .forward() directly — no hooks fire and .output is not affected. When you access .output or .input, you get the values from the model's actual forward pass with full hook interleaving.

Logit Lens¶

A classic application: decode each layer's hidden states through the final layer norm and language model head to see what the model "thinks" at each layer.

In [3]:

with model.trace("The Eiffel Tower is in the city of"):

    predictions = list().save()

    for layer in model.transformer.h:
        hs = layer.output[0]
        logits = model.lm_head(model.transformer.ln_f(hs))
        predictions.append(logits[0, -1].argmax(dim=-1))

for i, tok in enumerate(predictions):
    print(f"Layer {i:2d}: {model.tokenizer.decode(tok)}")

with model.trace("The Eiffel Tower is in the city of"): predictions = list().save() for layer in model.transformer.h: hs = layer.output[0] logits = model.lm_head(model.transformer.ln_f(hs)) predictions.append(logits[0, -1].argmax(dim=-1)) for i, tok in enumerate(predictions): print(f"Layer {i:2d}: {model.tokenizer.decode(tok)}")

Layer  0:  the
Layer  1:  the
Layer  2:  the
Layer  3:  the
Layer  4:  the
Layer  5:  the
Layer  6:  East
Layer  7:  Ing
Layer  8:  Rome
Layer  9:  London
Layer 10:  Paris
Layer 11:  Paris

Notice how the prediction converges toward "Paris" in later layers — this is the logit lens in action. Each layer's intermediate representation is projected into vocabulary space using the same final-layer modules.

Combining Module Calls with Torch Operations¶

Since everything inside a trace is real PyTorch, you can freely mix module calls with standard torch operations.

In [4]:

with model.trace("The Eiffel Tower is in the city of"):

    hs = model.transformer.h[-1].output[0]
    normed = model.transformer.ln_f(hs)

    # Cosine similarity between the last token and all others
    similarity = torch.cosine_similarity(
        normed[:, -1:, :], normed[:, :-1, :], dim=-1
    ).save()

print(f"Cosine similarity with last token: {similarity[0]}")

with model.trace("The Eiffel Tower is in the city of"): hs = model.transformer.h[-1].output[0] normed = model.transformer.ln_f(hs) # Cosine similarity between the last token and all others similarity = torch.cosine_similarity( normed[:, -1:, :], normed[:, :-1, :], dim=-1 ).save() print(f"Cosine similarity with last token: {similarity[0]}")

Cosine similarity with last token: tensor([0.9684, 0.9578, 0.9684, 0.9637, 0.9910, 0.9959, 0.9829, 0.9934, 0.9821],
       device='cuda:0', grad_fn=<SelectBackward0>)