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 = []