nnsight_utils¶

nnsight_utils ¶

Collecting activations and next-token distributions over many prompts.

Written against the standard values: a layer's activation is model.layers[i].layer_output unless a get_activations(model, layer) of your own says otherwise.

GetActivations `module-attribute` ¶

GetActivations = Callable[[StandardizedTransformer, int], torch.Tensor]

layer_output ¶

layer_output(model: StandardizedTransformer, layer: int) -> Tensor

The default get_activations: the residual stream leaving block layer.

get_token_activations ¶

get_token_activations(model: StandardizedTransformer, prompts: str | list[str] | None = None, layers: list[int] | None = None, get_activations: GetActivations | None = None, remote: bool = False, idx: int | None = None, tracer: Tracer | None = None) -> Tensor

The activation at one token position of each prompt, at each layer.

PARAMETER	DESCRIPTION
`model`	The standardized model. TYPE: `StandardizedTransformer`
`prompts`	What to run; `None` when called inside an open `tracer`. TYPE: `str \| list[str] \| None` DEFAULT: `None`
`layers`	Which blocks; default all. TYPE: `list[int] \| None` DEFAULT: `None`
`get_activations`	`(model, layer) -> tensor` giving a `[batch, seq, ...]` value; default `layer_output`. TYPE: `GetActivations \| None` DEFAULT: `None`
`remote`	Run on NDIF. TYPE: `bool` DEFAULT: `False`
`idx`	The token position; default `-1`, the last token. A negative index needs left padding, a positive one right padding. TYPE: `int \| None` DEFAULT: `None`
`tracer`	An open trace to read from instead of opening one. TYPE: `Tracer \| None` DEFAULT: `None`

RETURNS	DESCRIPTION
`Tensor`	`[num_layers, num_prompts, hidden_size]` on the CPU (still on the
`Tensor`	model's device when read inside a caller's `tracer`).

collect_last_token_activations_session ¶

collect_last_token_activations_session(model: StandardizedTransformer, prompts: list[str], batch_size: int, layers: list[int] | None = None, get_activations: GetActivations | None = None, remote: bool = False, idx: int | None = None) -> Tensor

get_token_activations over batches inside one model.session, so a remote run is one request.

collect_token_activations_batched ¶

collect_token_activations_batched(model: StandardizedTransformer, prompts: list[str], batch_size: int, layers: list[int] | None = None, get_activations: GetActivations | None = None, remote: bool = False, idx: int | None = None, tqdm=None, use_session: bool = True) -> Tensor

get_token_activations over batches; [num_layers, num_prompts, hidden_size] on the CPU.

A remote run goes through collect_last_token_activations_session (one request) unless use_session is off. tqdm is a progress-bar factory to wrap the batch loop with, or None.

compute_next_token_probs ¶

compute_next_token_probs(model: StandardizedTransformer, prompt: str | list[str], remote: bool = False) -> Tensor

The next-token distribution of each prompt, [num_prompts, vocab_size] on the CPU.