source: rebindable controller object, the module's own forward as its body, and .source on callable instances

[JadenFiotto-Kaufman]

Two findings from the stress sweep, both in src/nnsight/intervention/source.py.

1. The controller closure written into module.__dict__["forward"] (item 4)

install_controller did state = State(type(module).forward) and
module.__dict__["forward"] = make_controller(module, state). Three measured symptoms:

• copy.deepcopy of a wrapped module computed with the original's weights. copy treats
  a function as atomic, so the copy's forward was the original's closure, dereferencing the
  original module. Setting the copy's weights to 9.0 still returned the original's answer.
• A wrapped module could never be pickled again: cannot pickle 'weakref.ReferenceType' object, from the weakly-held interleavers in State.routes (not from the closure).
• An instance-level forward was destroyed — self.forward = self._fast picked in
  __init__, a monkeypatched layer, torch.compile's OptimizedModule. A module returning
  300 before wrapping returned 3 after, or raised NotImplementedError, permanently, outside
  traces too.

What changed, as decided:

• make_controller → a Controller object holding the module weakref and the State.
• Controller.__deepcopy__ looks the copied module up in deepcopy's memo (the copy is
  registered there before its __dict__ is copied) and binds to it; __reduce__ pickles it
  as (module, state).
• State.__getstate__ drops the routes — unpicklable, and a copy belongs to whatever wraps it
  next — and the instrumented body with its Compiled, which are built at run time (a
  function pickle can't name, a code object it can't write at all). A new original slot
  keeps the pre-instrumentation body so a copy re-instruments from it on demand.
• module_body takes the body from the instance forward when there is one, wrapping an
  already-bound one so it takes and ignores the module the controller passes (a method of the
  module keeps its __func__, which avoids pinning the module). install_source now
  instruments that body rather than the class's.
• Replacing forward after wrapping warns on the next install_controller.

Two wrappers are deliberately not treated as bodies: accelerate's device-alignment
wrapper (run_body re-applies it from _hf_hook) and transformers' tensor-parallel wrapper
(_keep_tp_forward rebuilds it around the controller). Both sit in the same instance slot,
and taking either as the body would apply its transforms twice — the TP one would have broken
tests/tp outright. _framework_forward is the guard, and it also suppresses the new warning
for both, since nnsight is what put them back. That means source.py knows one string about
transformers' TP wrapper (.install_forward.<locals>., the same test fragments.py makes);
the alternative was a change in fragments.py, which isn't mine.

.source on a module whose forward is an instance-level plain function now raises
SourceNotAvailable rather than instrumenting the class's forward and installing it as the
body — i.e. it says it can't rather than quietly running something else.

2. .source refused callable instances (item 20)

instrument() handled functions and bound methods; a callable instance has no __code__,
so compiled() raised SourceNotAvailable("callable has no Python source (builtin or C function)") — untrue, since type(fn).__call__ is ordinary Python. instrument() now uses
fn.__call__ in that case (a bound method, which the existing receiver path already handled —
the same thing an agent did by hand with set_attn_processor(AttnProcessor().__call__)), and
the message names the real cases. get_attention_scores, and so attention probability maps,
are reachable on diffusers UNets again.

Testing

Run with /home/localjadenfk/wd/nnsight-stress/env (transformers 5.16.1, torch 2.14, Python
3.12) and PYTHONPATH at this worktree.

• tests/test_source.py: 69 passed. Ten new tests in TestInstall /TestRecursive —
  deepcopy independence and tracing a copy, a pickle round trip (including after .source,
  and re-sourcing the restored module), a patched instance forward and a self.forward = self._fast one, a torch.compiled module, the replaced-forward warning, an accelerate
  ModelHook firing exactly once through the controller, and drilling into a callable
  instance.
• Whole offline suite (tests/ --ignore=tests/vllm --ignore=tests/tp, CUDA_VISIBLE_DEVICES=""):
  997 passed, 13 errors — the same 13, identically, on a pristine git archive of HEAD. They
  are HF_HUB_OFFLINE=1 artifacts ("You cannot infer task automatically within pipeline when
  using offline mode"), not this change.
• tests/tp + tests/test_tensor_parallel_rules.py with the GPU visible: 111 passed, 69
  skipped, same as baseline. This is what covers the TP-wrapper exclusion.
• The three original repros (module-zoo-custom-models/skills/repro_2.py, repro_3.py,
  noskills/repro_5_instance_forward_clobbered.py) and
  diffusion-pipelines/skills/repro_2.py (tiny SD on CPU, prints the processor's __call__
  and returns attn_get_attention_scores_0.output, [16, 1024, 77]). The repros' stale
  .save() idiom was adapted; nothing else.
• tests/performance/interleave_bench.py against baseline: every row within run-to-run noise
  (two baseline runs differ by 6–12% on this shared box), so the object call costs nothing
  measurable on the pass-through path.

Not tested here: vLLM — not installed in this environment. The change is runtime-agnostic
(the vLLM runtime uses the same controller), but a vLLM module carrying an instance-level
forward would now have it as the body rather than losing it, and nothing here exercises that.

Deliberately left alone

• deepcopy(model) on the NNsight wrapper (rather than the raw module) still gives a model
  whose traces raise OutOfOrderError. That lives in the envoy tree, not this file; dropping
  stale routes in State.__getstate__ may help it, but I did not verify or claim it.
• nn.DataParallel._replicate_for_data_parallel copies __dict__ shallowly, so a replica
  still shares the original's controller. Same for copy.copy. Only deepcopy was decided.
• docs/usage/source.md's "When .source isn't available" list, per the decision note: it
  never claimed callable instances were excluded.

🤖 Generated with Claude Code