Activation Catcher

Capture module inputs/outputs during a forward pass and select them by product_id. See Architecture (“Capturing activations as products”) for how this plugs into a teacher supplier.

class softs.catcher.ModelIOCatcher

Capture inputs and/or outputs of named submodules during forward passes.

Hooks are installed once on every submodule (attach_hooks). Each capture session is then scoped with a with block that arms a chosen subset of modules and collects their activations into a buffer; outside the block the hooks stay attached but inert.

Construction. model is the module to instrument - dict(model.named_modules()) defines the valid target names (e.g. "model.layers.0" or "model.layers.0.mlp"). device is where captured tensors (inputs and outputs) are moved via transfer_to_device(); pass None to leave them in place.

Selecting what to capture. inputs and outputs are iterables of module names - the dotted keys from model.named_modules(). They are independent; a name may appear in either, both, or neither:

• a name in inputs captures the args/kwargs the module is called with, via a forward pre-hook (before the module runs);

• a name in outputs captures whatever the module returns, via a forward hook (after the module and its whole subtree have run).

The buffer. __enter__ returns a buffer shaped buf[name][kind] -> list:

buf["model.layers.0"]["inputs"]   # list of {"args": ..., "kwargs": ...}
buf["model.layers.0"]["outputs"]  # list of module return values

One entry is appended per forward call, so running N forwards inside the same with accumulates N entries (e.g. to gather a batch of activations).

early_exit. With early_exit=True a forward is cut short the instant every requested input/output has been captured: the last capturing hook raises an internal _StopForward to skip the rest of the network. Drive forwards through run() (not the model directly) so that signal is caught and the budget re-seeded per call - then for batch in loader: catcher.run(**batch) runs every batch, each stopping after the deepest requested module. Output hooks fire in post-order (a parent fires after its children), so requesting the output of model.layers.0 stops after layer 0 finishes, not inside its mlp. A cut-short forward returns None (read results from buf); if a requested module never runs, no exit fires and the forward completes. With early_exit=False no stop signal is raised, so calling the model directly or via run() is interchangeable.

Example

>>> catcher = ModelIOCatcher(model, torch.device("cuda"))
>>>
>>> # Full forwards: capture layer 0's input+output over the loader.
>>> with catcher(inputs=["model.layers.0"],
...              outputs=["model.layers.0"]) as buf:
...     for batch in loader:
...         catcher.run(**batch)                 # or model(**batch) here
>>> len(buf["model.layers.0"]["outputs"])        # one entry per batch
>>> buf["model.layers.0"]["inputs"][0]["args"][0].shape   # layer 0 input
>>>
>>> # Only need layer 5's output? Skip layers 6.. on *every* batch:
>>> with catcher.for_product("outputs[model.layers.5]", early_exit=True) as buf:
...     for batch in loader:
...         catcher.run(**batch)                 # returns None, cut short
>>> hidden = buf["model.layers.5"]["outputs"]    # one tensor per batch

Notes

The model is never modified - run() is a plain model call wrapped in a try/except, so it composes with DDP, torch.compile, etc. Capturing is scoped to the with block: __exit__ disarms the flag sets, and the returned buf stays valid for reading afterwards. Hooks from attach_hooks stay registered until detach_hooks, so every forward pays a cheap per-module membership check.

__init__(model, device)

Parameters:

• model (Module)

• device (device)

attach_hooks()

detach_hooks()

for_product(product_id, early_exit=True)

Arm a session from a product_id capture spec (see parse_io_spec()).

product_id looks like "inputs[model.layers.0]|outputs[model.layers.0]". Returns self for use as with catcher.for_product(pid) as buf:. Raises ValueError if the spec selects nothing and KeyError for unknown module names.

Parameters:

• product_id (str)

• early_exit (bool)

run(*args, **kwargs)

Run one forward under the active capture session.

Use this instead of calling the model directly: it re-seeds the early-exit budget for this forward and absorbs the internal _StopForward signal, so for batch in loader: catcher.run(**batch) runs every batch. Returns the model’s output, or None if early exit cut the forward short. (With early_exit=False it is just the model call and always returns the real output.)

The model is never modified - this is a plain call wrapped in a try/except, so it composes with DDP, torch.compile, etc.

softs.catcher.parse_io_spec(product_id)

Parse a capture spec string into (inputs, outputs) module-name lists.

The spec names which modules to capture, e.g.:

"inputs[model.layers.0]|outputs[model.layers.0]"
"inputs[]|outputs[model.layers.0|model.layers.5]"
"outputs[model.layers.0]"                      # inputs section optional

Items inside the brackets are |-separated; the separator between the inputs[...] and outputs[...] blocks is ignored (the regex matches each block independently), so the inner and outer | never clash. Whitespace and empty items are dropped, and a missing section yields [].

Returns:

Two lists of module names (either may be empty).

Return type:

(inputs, outputs)

Parameters:

product_id (str)