From d435a9ea189652d6a07ea8f2c74037b04dee3d11 Mon Sep 17 00:00:00 2001 From: Jaden Fiotto-Kaufman Date: Tue, 8 Sep 2026 14:30:06 -0400 Subject: [PATCH] envoy: mirror every module entry, not every distinct module MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `Envoy.__init__` built the tree from `named_children()`, which deduplicates by module identity, and `_wrap_envoy` returned a shared module's envoy before recording the name at all. So a `ModuleList` holding one module twice lost an entry, and a container rebuilt from blocks the tree already wraps — truncating `transformer.h` to its first four blocks — got no children whatsoever: `h[3]` raised IndexError, iteration yielded nothing, the repr was empty. The tree is now built from `module._modules`, and every entry is recorded in a name-keyed `_child_map` alongside `_children`, which keeps holding one envoy per module for the recursive walks. Indexing, iteration, `modules()` and the repr go through the map, so an entry sharing its module with another keeps its own index while still naming the one envoy at the module's first path. Three more bugs in the same bookkeeping: replacing a child appended the new envoy instead of putting it at the replaced one's index, shifting every later index; the replaced module kept its `interleaver.envoys` registration, so re-attaching it served the replacement's values; and `base_model`, a property returning the module itself, fell through to build a duplicate envoy of a module inside itself — a RecursionError with `rename`, a phantom `cache()` key without it. Docs: mutate the tree through the envoy, not the wrapped module; the eproperty descriptor may go on a submodule's class, which `envoys=` has allowed since it landed; and the auxiliary-module attach uses the replacement form, since writing the attachment's output into its own input makes backward through it impossible. Co-Authored-By: Claude Opus 5 (1M context) --- docs/concepts/envoy.md | 17 +++- src/nnsight/intervention/envoy.py | 148 ++++++++++++++++++++-------- tests/test_envoy.py | 154 ++++++++++++++++++++++++++++++ 3 files changed, 272 insertions(+), 47 deletions(-) diff --git a/docs/concepts/envoy.md b/docs/concepts/envoy.md index 3d03a47c..d9373dd4 100644 --- a/docs/concepts/envoy.md +++ b/docs/concepts/envoy.md @@ -27,6 +27,7 @@ sources: [src/nnsight/intervention/envoy.py, src/nnsight/modeling/base.py] - Wrap any PyTorch model with `nnsight.NNsight(my_module)` and trace it directly. - Subclass `NNsight`/`Envoy` when you need model-specific loading, input prep, or new served values (see [Extension surface](#extension-surface)). - Don't replace `Envoy._module`'s `forward` by hand — reassign the module through the envoy so `instrument` re-runs and children are rebuilt. +- **Mutate the tree through the envoy, not the wrapped module.** `model.layers = nn.ModuleList(...)` and `model.block.mlp = new_mlp` rebuild the envoys and re-instrument. Assigning on `model._module` instead is not tracked: `model.block.mlp` goes on addressing the module it replaced, which the forward pass no longer runs, so reading it raises `OutOfOrderError`. ## Canonical pattern @@ -154,7 +155,7 @@ Add any `nn.Module` as an attribute; it's auto-wrapped as a child envoy. Apply i model.transformer.h[0].adapter = MyAdapter() with model.edit() as (tracer, edited): acts = edited.transformer.h[0].output - edited.transformer.h[0].output[:] = edited.transformer.h[0].adapter(acts, hook=True) + edited.transformer.h[0].output = edited.transformer.h[0].adapter(acts, hook=True) with edited.trace(prompt): inner = edited.transformer.h[0].adapter.inner.output.save() # now observable @@ -163,10 +164,10 @@ with edited.trace(prompt): ### 3. Custom served values via `eproperty` A new served value is an `eproperty` on the model/runtime class, served from the -driver side with its `.provide`. Because child envoys are always built as the base -`Envoy`, the descriptor goes on the model subclass (or the tracer), not an arbitrary -submodule. This is exactly how the vLLM wrapper adds `.logits` and `.samples` -(`modeling/vllm/vllm.py`): +driver side with its `.provide`. The descriptor goes on the class the envoy is +built as: the model subclass (or the tracer) by default, or — for a chosen +submodule — the `Envoy` subclass `envoys=` names for it. This is exactly how the +vLLM wrapper adds `.logits` and `.samples` (`modeling/vllm/vllm.py`): ```python class VLLM(Remotable): @@ -208,6 +209,12 @@ envoy, at the first path, and the later name is bound as an alias to it, the way one location (`model...self_attn.q_proj.output`), every spelling reaches the same envoy, and `_aliases` on the aliasing parent records where it points. +The container entry still counts: `layers[2]` and iteration walk every entry the +wrapped `ModuleList` holds — including one whose module an earlier entry already +brought in, and including a container rebuilt from blocks the tree already wraps +(`model.transformer.h = nn.ModuleList(list(model.transformer._module.h)[:4])`). +Each such entry names the one envoy at the module's first path. + ## Module renaming (aliases) `rename={...}` on `NNsight`/`Envoy` binds aliases pointing at the same child envoy (`_bind_aliases`, `envoy.py`). A single-component path (`{"transformer": "gpt"}`) binds wherever it resolves; a multi-component path (`{"transformer.h": "layers"}`) binds on the envoy it resolves from. Aliases are ordinary attributes referencing the same object, so they survive a dispatch re-point with no rebuild. diff --git a/src/nnsight/intervention/envoy.py b/src/nnsight/intervention/envoy.py index c4c3a21a..d9048540 100644 --- a/src/nnsight/intervention/envoy.py +++ b/src/nnsight/intervention/envoy.py @@ -124,7 +124,11 @@ class Envoy: _module: The wrapped `torch.nn.Module`. _edits: Default interventions registered by [`edit`][nnsight.intervention.envoy.Envoy.edit], replayed on every trace (a list of [`Mediator`][nnsight.intervention.interleaver.Mediator]). - _children: The direct child envoys, in module order. + _children: The child envoys this envoy owns, in module order — one per + module, so a module the tree already wraps elsewhere is not in it. + _child_map: Every entry of the wrapped module's ``_modules``, by name and + in module order, so an entry sharing its module with another keeps its + own name. What indexing, iteration and the repr walk. """ def __init__( @@ -160,9 +164,14 @@ def __init__( self._envoys = envoys self._children: list[Envoy] = [] + self._child_map: dict[str, Envoy] = {} - for name, child in module.named_children(): - self._add_envoy(name, child) + # `_modules`, not `named_children()`: the latter deduplicates by module + # identity, so a ModuleList holding one module twice — or rebuilt from + # blocks this tree already wraps — would lose entries torch still indexes. + for name, child in list(module._modules.items()): + if child is not None: + self._add_envoy(name, child) # Children exist now, so multi-component alias paths (e.g. "h.0") resolve. self._bind_aliases() @@ -187,38 +196,71 @@ def _wrap_envoy(self, name: str, module: torch.nn.Module) -> Envoy: ): attribute = f"{self.OVERLOAD_PREFIX}{name}" existing = self.__dict__.get(attribute) + index = None if isinstance(existing, Envoy): - self._children.remove(existing) + # Replace at the replaced child's own index: appending instead would + # shift every child after it, so `_children[i]` would name a different + # module than `i` does on the wrapped module. + if existing in self._children: + index = self._children.index(existing) + # The replaced module keeps its registration otherwise, so a later + # path to it would serve the replacement's values. + if self.interleaver.envoys.get(id(existing._module)) is existing: + del self.interleaver.envoys[id(existing._module)] child_path = f"{self.path}.{name}" # A module the tree already holds under another path (a layer that a # wrapper keeps a second reference to, tied weights) gets no second envoy: # this name is an alias of the first, as torch's own named_modules lists # it once, so the module has one location and either spelling reads it. + # `_child_map` still records the name, so the entry keeps its index. shared = self.interleaver.envoys.get(id(module)) - if shared is not None and shared is not self: - self._aliases[attribute] = shared.path - object.__setattr__(self, attribute, shared) - return shared - envoy = self._resolve_envoy_class(module, child_path)( - module, - path=child_path, - interleaver=self.interleaver, - rename=self._rename, - envoys=self._envoys, - ) - self._children.append(envoy) - if attribute != name: - # Record it the way `rename` records its own, so the repr labels the - # child `E_output/output` and an alias cannot later claim the name. - self._aliases[attribute] = name - warnings.warn( - f"Module '{self.path}' has a submodule named '{name}', which " - f"nnsight already serves on every module. The submodule is " - f"'.{attribute}' here; '.{name}' stays the module's output." + if shared is not None: + envoy = shared + if shared is not self: + self._aliases[attribute] = shared.path + if index is not None: + del self._children[index] + else: + envoy = self._resolve_envoy_class(module, child_path)( + module, + path=child_path, + interleaver=self.interleaver, + rename=self._rename, + envoys=self._envoys, ) + if index is None: + self._children.append(envoy) + else: + self._children[index] = envoy + if attribute != name: + # Record it the way `rename` records its own, so the repr labels the + # child `E_output/output` and an alias cannot later claim the name. + self._aliases[attribute] = name + warnings.warn( + f"Module '{self.path}' has a submodule named '{name}', which " + f"nnsight already serves on every module. The submodule is " + f"'.{attribute}' here; '.{name}' stays the module's output." + ) + # A name the module itself holds is an entry of the tree; one reached any + # other way (a property like transformers' `base_model`) is just a spelling. + if self._module._modules.get(name) is module: + self._child_map[name] = envoy object.__setattr__(self, attribute, envoy) return envoy + def _named_children(self) -> list[tuple[str, Envoy]]: + # (name, child) for everything under this envoy: the wrapped module's own + # entries in module order, then children attached outside it + # (TransformersModel's `generator`), which have no entry to be ordered by. + named = list(self._child_map.items()) + bound = {id(child) for child in self._child_map.values()} + named.extend( + (child.path.rsplit(".", 1)[-1], child) + for child in self._children + if id(child) not in bound + ) + return named + def _bind_aliases(self) -> None: """Bind each ``rename`` alias as an attribute pointing at the same Envoy. @@ -358,13 +400,12 @@ def _update(self, module: torch.nn.Module) -> None: # module (its own forward; the previous module's controller doesn't carry # over) — see Interleaver.instrument. self.interleaver.instrument(self) - children = dict(module.named_children()) - for child in self._children: - name = child.path.rsplit(".", 1)[-1] + children = module._modules + for name, child in self._named_children(): # A child that isn't a submodule of the new module — e.g. a standalone # module added to the tree (TransformersModel's `generator`) — has nothing # to re-point at, so leave it as-is (it keeps its own module and controller). - if name in children: + if children.get(name) is not None: child._update(children[name]) def trace( @@ -807,7 +848,7 @@ def __call__(self, *args: Any, hook: bool = False, **kwargs: Any) -> Any: model.transformer.h[0].adapter = MyAdapter() with model.edit() as (tracer, edited): acts = edited.transformer.h[0].output - edited.transformer.h[0].output[:] = \ + edited.transformer.h[0].output = \ edited.transformer.h[0].adapter(acts, hook=True) with edited.trace(prompt): inner = edited.transformer.h[0].adapter.inner.output.save() @@ -819,7 +860,7 @@ def __call__(self, *args: Any, hook: bool = False, **kwargs: Any) -> Any: with model.edit(inplace=True) as tracer: for _ in tracer.iter[:]: acts = model.transformer.h[0].output - model.transformer.h[0].output[:] = \ + model.transformer.h[0].output = \ model.transformer.h[0].adapter(acts, hook=True) Args: @@ -882,19 +923,29 @@ def __iter__(self) -> Iterator[Envoy]: for layer in model.model.layers: print(layer.path) """ - return iter(self._children) + return iter([child for _, child in self._named_children()]) def __getitem__(self, key: Any) -> Envoy: """Index into direct child envoys, e.g. for a `ModuleList`. + An int or str key resolves **by name**, the way the wrapped module + indexes it — ``layers[2]`` is the module `layers` holds at ``"2"``, even + when an earlier entry is a module the tree already wraps elsewhere and so + has no envoy of its own here. + Args: - key: Any index the underlying child list accepts (an int, or a slice). + key: An index the wrapped module accepts (an int, or a str), or a + slice over this envoy's children in module order. Returns: Envoy: The child envoy at ``key`` (e.g. ``model.layers[0]`` for the first block of a ``ModuleList``). """ - return self._children[key] + if isinstance(key, slice): + return [child for _, child in self._named_children()][key] + if isinstance(key, int) and key < 0: + key += len(self) + return getattr(self, str(key)) def __len__(self) -> int: """The number of entries in the wrapped module (e.g. a ``ModuleList``'s length).""" @@ -941,12 +992,23 @@ def modules( A list of [`Envoy`][nnsight.intervention.envoy.Envoy] (or ``(path, Envoy)`` tuples when ``names``). """ # Flatten the envoy tree (children first, then self), optionally filtered - # by include_fn and paired with each envoy's path when names=True. + # by include_fn and paired with each envoy's path when names=True. An envoy + # two paths reach is listed once, at the first, as `named_modules()` lists + # a shared module once — and `seen` is what keeps a self-referential + # spelling (transformers' `base_model` on a base model) from recursing. result: list[Any] = [] - for child in self._children: - result.extend(child.modules(include_fn=include_fn, names=names)) - if include_fn is None or include_fn(self): - result.append((self.path, self) if names else self) + seen: set[int] = set() + + def walk(envoy: Envoy) -> None: + if id(envoy) in seen: + return + seen.add(id(envoy)) + for _, child in envoy._named_children(): + walk(child) + if include_fn is None or include_fn(envoy): + result.append((envoy.path, envoy) if names else envoy) + + walk(self) return result def named_modules( @@ -967,7 +1029,7 @@ def _name(self) -> str: return self._module._get_name() def _repr_modulelist(self) -> str: - reprs = [repr(child) for child in self._children] + reprs = [repr(child) for child in self] start_end = [[0, 0]] blocks = [reprs[0]] @@ -989,7 +1051,7 @@ def _repr_modulelist(self) -> str: return self._name() + "(\n " + "\n ".join(lines) + "\n)" def __repr__(self) -> str: - if self._children and isinstance(self._module, torch.nn.ModuleList): + if self._child_map and isinstance(self._module, torch.nn.ModuleList): return self._repr_modulelist() extra_lines = [] @@ -1008,11 +1070,13 @@ def __repr__(self) -> str: direct.setdefault(path, []).append(alias) child_lines = [] - for child in self._children: - name = child.path.rsplit(".", 1)[-1] + for name, child in self._named_children(): label = "/".join([*direct.get(name, []), name]) child_lines.append(f"({label}): " + _addindent(repr(child), 2)) for alias in mounts: + # A mount that is also an entry of the module has a line already. + if alias in self._child_map: + continue child_lines.append(f"({alias}): " + _addindent(repr(getattr(self, alias)), 2)) # eproperties given a description surface as their own lines, so special diff --git a/tests/test_envoy.py b/tests/test_envoy.py index 78e8a227..0a535c9b 100644 --- a/tests/test_envoy.py +++ b/tests/test_envoy.py @@ -211,6 +211,160 @@ def test_replace_no_duplicate_paths(self, envoy): assert len(paths) == len(set(paths)) +class Stack(nn.Module): + def __init__(self, n=4): + super().__init__() + self.layers = nn.ModuleList([nn.Linear(8, 8) for _ in range(n)]) + + def forward(self, x): + for layer in self.layers: + x = layer(x) + return x + + +class SharedStack(nn.Module): + """A ModuleList holding one module twice — torch still indexes three entries.""" + + def __init__(self): + super().__init__() + self.shared = nn.Linear(8, 8) + self.layers = nn.ModuleList([self.shared, nn.Linear(8, 8), self.shared]) + + def forward(self, x): + for layer in self.layers: + x = layer(x) + return x + + +class TestSharedEntries: + @pytest.fixture + def shared(self): + return Envoy(SharedStack()) + + def test_every_entry_is_indexable(self, shared): + # `named_children()` deduplicates by identity, so entry 2 used to be + # missing entirely and `layers[2]` raised IndexError. + assert len(list(shared.layers)) == 3 + assert shared.layers[1]._module is shared._module.layers[1] + assert shared.layers[2]._module is shared._module.shared + + def test_a_shared_entry_is_the_one_envoy(self, shared): + assert shared.layers[0] is shared.shared + assert shared.layers[2] is shared.shared + + def test_a_shared_module_is_listed_once(self, shared): + # The same rule torch's named_modules() follows. + assert len(shared.modules()) == len(list(shared._module.named_modules())) + + def test_the_repr_shows_every_entry(self, shared): + assert "(0-2): 3 x Linear" in repr(shared.layers) + + def test_a_shared_list_traces(self, shared): + x = torch.randn(1, 8) + with shared.trace(x): + middle = shared.layers[1].output.save() + module = shared._module + assert torch.allclose(middle, module.layers[1](module.shared(x))) + + +class TestRebuiltContainer: + """A container rebuilt from modules the tree already wraps — truncating layers.""" + + @pytest.fixture + def stack(self): + return Envoy(Stack()) + + def test_the_entries_are_kept(self, stack): + stack.layers = nn.ModuleList(list(stack._module.layers)[:2]) + assert len(list(stack.layers)) == 2 + assert [layer.path for layer in stack.layers] == [ + "model.layers.0", + "model.layers.1", + ] + + def test_the_tree_still_mirrors_the_module(self, stack): + stack.layers = nn.ModuleList(list(stack._module.layers)[:2]) + assert {node.path for node in stack.modules()} == module_paths(stack._module) + + def test_a_truncated_stack_traces(self, stack): + stack.layers = nn.ModuleList(list(stack._module.layers)[:2]) + x = torch.randn(1, 8) + with stack.trace(x): + last = stack.layers[1].output.save() + assert torch.allclose(last, stack._module(x)) + + def test_gpt2_truncated_blocks(self): + # Keeping the first four blocks of a real model is the ordinary way a + # user hits this; every block is already wrapped, so the new ModuleList + # used to end up with no children at all. + from nnsight.modeling.transformers import TransformersModel + + model = TransformersModel( + "openai-community/gpt2", task="text-generation", dispatch=True + ) + model.transformer.h = nn.ModuleList(list(model.transformer._module.h)[:4]) + assert len(list(model.transformer.h)) == 4 + with model.trace("Hello"): + hidden = model.transformer.h[3].output.save() + assert hidden.shape[-1] == model._module.config.n_embd + + +class TestReplacement: + def test_a_replacement_keeps_its_index(self): + # Remove-then-append put the new child last, shifting every index after + # it, so `envoy[2]` named the module the wrapped module holds at 3. + net = nn.Sequential(nn.Linear(8, 8), nn.ReLU(), nn.Linear(8, 8), nn.Tanh()) + sequential = Envoy(net) + setattr(sequential, "1", nn.Identity()) # `sequential[1] = ...` is not a thing + assert [type(child._module) for child in sequential] == [type(m) for m in net] + assert isinstance(sequential[1]._module, nn.Identity) + assert isinstance(sequential[3]._module, nn.Tanh) + + def test_a_replaced_module_is_deregistered(self, envoy, module): + # Left registered, the replaced module would come back as an alias of + # its replacement and serve the replacement's values. + old_module = module.head + old_envoy = envoy.head # held, so the registry entry can't just be collected + envoy.head = nn.Identity() + assert id(old_module) not in envoy.interleaver.envoys + envoy.spare = old_module + assert envoy.spare is not old_envoy + assert envoy.spare.path == "model.spare" + + +class SelfNaming(nn.Module): + """A property returning the module itself, as `base_model` does on a base model.""" + + def __init__(self): + super().__init__() + self.layer = nn.Linear(8, 8) + + @property + def base_model(self): + return self + + def forward(self, x): + return self.layer(x) + + +class TestSelfNamingAttribute: + def test_it_resolves_to_this_envoy(self): + envoy = Envoy(SelfNaming()) + assert envoy.base_model is envoy + assert envoy.interleaver.envoys[id(envoy._module)] is envoy + + def test_it_adds_no_path_to_the_tree(self): + envoy = Envoy(SelfNaming()) + before = {node.path for node in envoy.modules()} + envoy.base_model + assert {node.path for node in envoy.modules()} == before + + def test_a_rename_through_it_resolves(self): + # A duplicate envoy re-ran `_bind_aliases` on the same module: RecursionError. + envoy = Envoy(SelfNaming(), rename={"base_model.layer": "inner"}) + assert envoy.inner is envoy.layer + + class TestTraceable: def _model(self): ran = []