envoy: mirror every module entry, not every distinct module

[JadenFiotto-Kaufman]

Items 12 (all five parts), 19, and 23.4 from the stress-sweep next-steps, plus the
in-place auxiliary-module idiom in docs/concepts/envoy.md.

What was broken

Envoy.__init__ built the child tree from module.named_children(), which
deduplicates by module identity, and _wrap_envoy returned early for a module the
tree already wrapped without recording the name at all. Four symptoms, three silent:

• A ModuleList holding one module twice got fewer children than entries, so
  layers[2] raised IndexError and iteration skipped it. A container rebuilt from
  blocks the tree already wraps — model.transformer.h = nn.ModuleList(list(h)[:4]),
  an ordinary truncation — got zero children: h[3] raised, iteration yielded
  nothing, the repr was empty, and the blocks vanished from modules() (so from
  cache() targets too).
• Replacing a child did _children.remove(existing) then .append(...), shifting
  every later index: after setattr(m, "1", nn.Identity()) on a 4-entry Sequential,
  m[1] was the third module and m[3] was the Identity.
• The replaced module kept its interleaver.envoys registration, so re-attaching it
  under a new name bound it as an alias of the stale envoy and served the
  replacement's values.
• .base_model — a transformers property returning getattr(self, prefix, self),
  i.e. the module itself on a bare GPT2Model/LlamaModel — fell through the
  shared is not self clause and built a duplicate envoy of a module inside itself:
  RecursionError at construction with rename, and without it a re-pointed registry
  and a phantom tracer.cache() key. Plain gpt2, no adapter needed.

What changed (src/nnsight/intervention/envoy.py)

• The tree is built from module._modules.items() (skipping None entries), so no
  entry is deduplicated away.
• New _child_map: dict[str, Envoy] — every entry of the wrapped module's _modules,
  by name, in module order. _children keeps its meaning (one envoy per module, plus
  standalone children appended from outside, e.g. TransformersModel.generator), so
  modeling/transformers.py and the _children.append pattern documented in
  docs/usage/extending.md are untouched. _named_children() joins the two: map
  entries first, then any _children extras.
• __getitem__ resolves an int/str key by name (getattr(self, str(key)),
  negative ints normalised against len(self)); slices index the ordered children.
• __iter__, _update, modules() and the repr walk _named_children(). modules()
  gained an identity seen set so an envoy two paths reach is still listed once, the
  way named_modules() lists a shared module once — that is also what stops a
  self-referential spelling recursing.
• Replacement now writes at the replaced child's index and drops the replaced module's
  interleaver.envoys entry.
• shared is not self dropped, so a property returning the module itself resolves to
  this envoy like any other shared module.

One behaviour change worth naming: envoy[i] out of range now raises AttributeError
(naming the key) rather than IndexError, since the key is resolved by name.

Docs (docs/concepts/envoy.md)

• New "when not to use" bullet: mutate the tree through the envoy, not the wrapped
  module — direct assignment on _module is not tracked and leaves the envoy
  addressing the module it replaced (OutOfOrderError on read).
• "Shared modules" now says the container entry still counts, with the truncation case.
• The stale sentence at :167 ("child envoys are always built as the base Envoy, so
  the descriptor goes on the model subclass ... not an arbitrary submodule") predates
  envoys=, documented 30 lines above it; rewritten to say the descriptor goes on the
  class the envoy is built as.
• The auxiliary-module attach at :157 used layer.output[:] = aux(acts, hook=True),
  which writes the attachment's output into its own input and makes backward through
  it impossible; now the replacement form, matching
  docs/patterns/sae-and-auxiliary-modules.md:101. The same idiom appeared twice in
  Envoy.__call__'s docstring and is fixed there too.

Tests

tests/test_envoy.py gains four classes adapted from the repro scripts:
TestSharedEntries (a ModuleList holding one module twice), TestRebuiltContainer
(truncation, including the real gpt2 transformer.h[:4] case with a trace),
TestReplacement (index kept, old module deregistered), TestSelfNamingAttribute
(a property returning the module itself, including the rename RecursionError).
13 of the 14 fail on origin/0.8 and all pass here.

Ran with PYTHONPATH at this worktree's src:

• tests/test_envoy.py — 101 passed
• the whole suite minus tests/vllm and tests/performance (so including tests/tp)
  — 1044 passed, 72 skipped, 1 xfailed

tests/vllm was not run: vLLM is not installed in this environment. The change is
generic to Envoy, so the vLLM tree is affected in principle; the MLA shared-module
case it exercises is covered by TestSharedEntries on a plain module.

🤖 Generated with Claude Code