Info

Last Execution: 2026-07-27

Package	Version
nnsight	0.8
Python	3.12.13
torch	2.13.0+cu126
transformers	5.15.0

Loading a Model¶

Before you can trace anything you need a model. nnsight never runs a bare torch.nn.Module directly — it wraps one in an Envoy tree so every submodule becomes observable. This page is about getting to that wrapped model: constructing one from a HuggingFace repo id, understanding the meta build and when weights actually load, wrapping a model you already have in memory, and choosing where the weights live.

The workhorse wrapper is TransformersModel, which loads any HuggingFace transformers checkpoint through a pipeline so tokenization and generation come for free. For an arbitrary PyTorch module there is the thin NNsight wrapper.

In [1]:

from nnsight.modeling.transformers import TransformersModel
from nnsight import NNsight
import torch

from nnsight.modeling.transformers import TransformersModel from nnsight import NNsight import torch

Constructing from a Repo ID¶

The first argument is a HuggingFace repo id. Any keyword you would pass to transformers loading — device_map, dtype (the torch_dtype alias also works), revision, attn_implementation, and so on — is forwarded through to the underlying load, alongside a few nnsight-specific options like dispatch and rename. dtype also takes a quantization name such as "nf4".

In [2]:

model = TransformersModel("openai-community/gpt2")

# The tokenizer the model will actually use is exposed as an attribute.
print(model.tokenizer)

model = TransformersModel("openai-community/gpt2") # The tokenizer the model will actually use is exposed as an attribute. print(model.tokenizer)

GPT2Tokenizer(name_or_path='openai-community/gpt2', vocab_size=50257, model_max_length=1024, padding_side='left', truncation_side='right', special_tokens={'bos_token': '<|endoftext|>', 'eos_token': '<|endoftext|>', 'unk_token': '<|endoftext|>', 'pad_token': '<|endoftext|>'}, added_tokens_decoder={
	50256: AddedToken("<|endoftext|>", rstrip=False, lstrip=False, single_word=False, normalized=True, special=True),
})

That single line is enough to trace against. Because the model was built lazily (see the next section), the weights load automatically the first time you run it.

In [3]:

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

token = logits[0, -1].argmax(dim=-1)
print(f"Prediction: {model.tokenizer.decode(token)!r}")

with model.trace("The Eiffel Tower is in the city of"): logits = model.output.logits.save() token = logits[0, -1].argmax(dim=-1) print(f"Prediction: {model.tokenizer.decode(token)!r}")

Prediction: ' Paris'

The Meta Model and Dispatching¶

By default TransformersModel(repo_id) does not download or allocate any weights. It reads only the model's config and builds the full architecture on the meta device — every module and parameter exists with the right shape and dtype, but backed by no storage. This is cheap and instant, and it is all nnsight needs to let you write interventions against module paths: the tree is fully navigable before a single byte of weights is loaded.

The dispatched flag tells you whether real weights are in memory yet.

In [4]:

meta_model = TransformersModel("openai-community/gpt2")

print("dispatched:", meta_model.dispatched)
print("weight device:", next(meta_model._module.parameters()).device)

meta_model = TransformersModel("openai-community/gpt2") print("dispatched:", meta_model.dispatched) print("weight device:", next(meta_model._module.parameters()).device)

dispatched: False
weight device: meta

Because the architecture is fully present, you can already inspect it and even scan for activation shapes — that runs the forward under fake tensors, so it needs no weights and never dispatches.

In [5]:

with meta_model.scan("Hello world"):
    shape = meta_model.transformer.h[0].output[0].shape.save()

print("block 0 output shape:", tuple(shape))
print("still on meta, dispatched:", meta_model.dispatched)

with meta_model.scan("Hello world"): shape = meta_model.transformer.h[0].output[0].shape.save() print("block 0 output shape:", tuple(shape)) print("still on meta, dispatched:", meta_model.dispatched)

block 0 output shape: (2, 768)
still on meta, dispatched: False

Dispatching is the step that loads the real weights and swaps them into the existing tree (in place, so every module path you already referenced stays valid). It happens automatically on the first real trace or generation:

In [6]:

with meta_model.trace("Hello world"):
    pass

print("after first trace, dispatched:", meta_model.dispatched)
print("weight device:", next(meta_model._module.parameters()).device)

with meta_model.trace("Hello world"): pass print("after first trace, dispatched:", meta_model.dispatched) print("weight device:", next(meta_model._module.parameters()).device)

after first trace, dispatched: True
weight device: cuda:0

If you would rather pay that cost up front — for instance to fail fast on an OOM, or to warm the model before timing anything — pass dispatch=True to load eagerly at construction. The meta phase is skipped entirely.

In [7]:

eager_model = TransformersModel("openai-community/gpt2", dispatch=True)

print("dispatched immediately:", eager_model.dispatched)

eager_model = TransformersModel("openai-community/gpt2", dispatch=True) print("dispatched immediately:", eager_model.dispatched)

dispatched immediately: True

Why a meta build at all?

Building on meta makes construction instant and memory-free, which matters most for remote execution: your client builds the weightless skeleton so it knows the module paths, while the real weights only ever live on the server. Locally it means you can hold a TransformersModel for a checkpoint you have not decided to load yet, and only pay for the weights when you actually trace.

Wrapping an Already-Loaded Model¶

You do not have to hand nnsight a repo id. If you already have a model object in memory, pass the object itself and nnsight wraps it as-is — no meta phase, already dispatched.

For an arbitrary torch.nn.Module, use NNsight. It builds a root Envoy mirroring the module tree, so children are reachable by index or attribute exactly as in PyTorch.

In [8]:

net = torch.nn.Sequential(
    torch.nn.Linear(5, 10),
    torch.nn.Linear(10, 2),
)
wrapped = NNsight(net)

with wrapped.trace(torch.rand(1, 5)):
    hidden = wrapped[0].output.save()

print("hidden shape:", tuple(hidden.shape))

net = torch.nn.Sequential( torch.nn.Linear(5, 10), torch.nn.Linear(10, 2), ) wrapped = NNsight(net) with wrapped.trace(torch.rand(1, 5)): hidden = wrapped[0].output.save() print("hidden shape:", tuple(hidden.shape))

hidden shape: (1, 10)

For a HuggingFace model you loaded yourself, pass the loaded module to TransformersModel instead of a repo id. It infers the pipeline task from the model's architecture and reuses your instance directly — handy when you have already customized the model (quantization, an adapter, edited weights) and want to trace that object.

In [9]:

from transformers import AutoModelForCausalLM

hf_model = AutoModelForCausalLM.from_pretrained("openai-community/gpt2")
model_from_object = TransformersModel(hf_model)

print("inferred task:", model_from_object.task)
print("dispatched:", model_from_object.dispatched)

with model_from_object.trace("Hello"):
    logits = model_from_object.output.logits.save()

print("logits shape:", tuple(logits.shape))

from transformers import AutoModelForCausalLM hf_model = AutoModelForCausalLM.from_pretrained("openai-community/gpt2") model_from_object = TransformersModel(hf_model) print("inferred task:", model_from_object.task) print("dispatched:", model_from_object.dispatched) with model_from_object.trace("Hello"): logits = model_from_object.output.logits.save() print("logits shape:", tuple(logits.shape))

inferred task: text-generation
dispatched: True
logits shape: (1, 1, 50257)

Task inference

The pipeline factory can infer a task from a repo-id string, but not from a bare module instance. nnsight infers it from the architecture — a generative model (can_generate()) becomes text-generation, otherwise the class-name suffix decides (*ForMaskedLM → fill-mask, and so on). If it can't be inferred, pass task=... explicitly, e.g. TransformersModel(hf_model, task="text-generation").

Inference from a repo id asks the Hub for the checkpoint's metadata, and a fully cached checkpoint does not change that: under HF_HUB_OFFLINE=1 it raises RuntimeError: You cannot infer task automatically within 'pipeline' when using offline mode. Pass task= on an air-gapped machine or a cluster node with no outbound network.

Device Placement¶

Where the weights land is decided at dispatch, by the same knobs transformers uses. With a GPU available, transformers places the model on it by default, so a plain TransformersModel(repo_id) runs on the GPU here without asking. To choose explicitly:

• device="cpu" / device="cuda" / device=0 — put the whole model on one device. This is how you pin to the CPU when a GPU would otherwise be picked.
• device_map="cuda" — the single GPU, via accelerate.
• device_map="auto" — let accelerate place and, if needed, shard the model across every available device (GPUs first, spilling to CPU/disk for a model too big to fit). This is what you want for large checkpoints.

device and device_map are HuggingFace loading options, so they only take effect once the model dispatches.

In [10]:

# Force the model onto the CPU, overriding the GPU-by-default behavior.
cpu_model = TransformersModel("openai-community/gpt2", device="cpu", dispatch=True)
print("device=cpu       ->", next(cpu_model._module.parameters()).device)

# Pin to the single GPU.
gpu_model = TransformersModel("openai-community/gpt2", device="cuda", dispatch=True)
print("device=cuda      ->", next(gpu_model._module.parameters()).device)

# Let accelerate place (and shard, for big models) across available devices.
auto_model = TransformersModel("openai-community/gpt2", device_map="auto", dispatch=True)
print("device_map=auto  ->", next(auto_model._module.parameters()).device)

# Force the model onto the CPU, overriding the GPU-by-default behavior. cpu_model = TransformersModel("openai-community/gpt2", device="cpu", dispatch=True) print("device=cpu ->", next(cpu_model._module.parameters()).device) # Pin to the single GPU. gpu_model = TransformersModel("openai-community/gpt2", device="cuda", dispatch=True) print("device=cuda ->", next(gpu_model._module.parameters()).device) # Let accelerate place (and shard, for big models) across available devices. auto_model = TransformersModel("openai-community/gpt2", device_map="auto", dispatch=True) print("device_map=auto ->", next(auto_model._module.parameters()).device)

device=cpu       -> cpu

device=cuda      -> cuda:0

device_map=auto  -> cuda:0

Sharding Across GPUs (Tensor Parallelism)¶

device_map spreads whole layers over several GPUs and runs them one after another, which is enough when a model merely does not fit. Tensor parallelism splits every attention and MLP projection within each layer, so the GPUs work on the same layer at once. nnsight supports tracing a model loaded this way, and the sharding is invisible in your intervention code.

Pass distributed_config and launch the script with torchrun — tensor parallelism needs one process per GPU, and every rank runs the whole script:

# tp_trace.py  —  torchrun --nproc_per_node=4 tp_trace.py
import torch
from transformers.distributed import DistributedConfig
from nnsight.modeling.transformers import TransformersModel

model = TransformersModel(
    "meta-llama/Llama-3.2-3B",
    task="text-generation",
    dispatch=True,
    dtype=torch.bfloat16,
    distributed_config=DistributedConfig(tp_size=4),
)

with model.trace("The Eiffel Tower is in the city of"):
    # Each rank computes only 2048 of gate_proj's 8192 features...
    gate = model.model.layers[5].mlp.gate_proj.output.save()
    logits = model.lm_head.output.save()

print(gate.shape)  # ...but you read all 8192: (1, 11, 8192)

Under the hood a sharded value is gathered before your intervention sees it and re-split afterwards, so reads and edits behave exactly as they would on one GPU — including edits that span rank boundaries:

with model.trace(prompt):
    model.model.layers[5].mlp.gate_proj.output[..., :3000] = 0
    logits = model.lm_head.output.save()

Only two kinds of value are ever really a slice: a column-parallel output (q_proj, k_proj, v_proj, gate_proj, up_proj) and a row-parallel input (o_proj.input, down_proj.input). Everything most people read — a decoder layer, mlp, self_attn, lm_head — is already whole, because those layers all-reduce their output. The gather only runs where an intervention is actually looking.

Two rules when every rank runs your code

Do not branch on rank. All ranks must take the same path, or they stop agreeing on when to communicate and the run hangs.

Seed before you sample. If sampling diverges, the ranks generate different tokens, and the model then all-reduces activations computed from different sequences — the output is wrong on every rank, not just inconsistent. Use do_sample=False, or torch.manual_seed(...) with the same value on every rank. Many checkpoints ship do_sample: true in their generation_config.json, so this applies even if you never asked for sampling.

Saved values are identical on every rank, so guard printing and file writes with if int(os.environ["RANK"]) == 0: or you will get one copy per GPU.

tp_plan="auto" / tp_size= are not from_pretrained arguments in transformers 5.x, despite what its docstring says — use distributed_config. Mixture-of-experts sharding is not supported yet and raises UnsupportedParallelStyle rather than returning a fragment of a tensor.

Quantization¶

The other way to fit a model that does not fit: rather than splitting it across GPUs, hold each weight in fewer bits. Name the format where you would name a dtype, and nothing else about the constructor changes.

model = TransformersModel(
    "meta-llama/Llama-3.2-3B",
    task="text-generation",
    dtype="nf4",          # where you would write "bfloat16"
    dispatch=True,
)

Quantization has the accepted names, what a trace sees, and what the memory, accuracy and speed actually cost.

Once loaded, the model is an ordinary wrapped module: from here every other feature page — module access, batching, generation, editing — applies unchanged, regardless of how you constructed it or where it lives.