Skip to content

cache: resolve modules= against the tree instead of caching nothing - #716

Open
JadenFiotto-Kaufman wants to merge 1 commit into
0.8from
fix/cache
Open

JadenFiotto-Kaufman wants to merge 1 commit into
0.8from
fix/cache

Conversation

@JadenFiotto-Kaufman

Copy link
Copy Markdown
Member

What was broken

Cache.__init__ built self.targets from the strings it was given without
resolving them, then subscribed to locations nothing provides. Any string that
named no module produced an empty cache and no error:

modules=['model.transformer.h.0.atn']  -> keys=[]
modules=['transformer.h.0']            -> keys=[]
modules=['model.transformer.h.*']      -> keys=[]
modules=['.*\.mlp$']                   -> keys=[]

while the correct spellings worked. A callable or an re.Pattern passed as
modules= failed with TypeError: 'function' object is not iterable.

Repro: tasks/cache-probes-function-vectors/skills/repro_3_unknown_path_silent.py.

What changed

src/nnsight/intervention/cache.py — each modules= entry is resolved in
Cache.__init__:

  • an envoy contributes its .path, as before;
  • a string is walked down the envoy tree, so the miss is the same
    AttributeError reading that module by attribute gives, named against the
    target it came from. A path is the whole path (the way the cache's keys are);
    one missing the model's own name gets a "did you mean 'model.…'?" when adding
    it back resolves. Aliases resolve too, so a renamed spelling now normalises to
    the real path rather than caching nothing;
  • a name that lands on a Parameter, or on .output/.input (a value nnsight
    serves, not a module), says so rather than being subscribed to;
  • anything that is neither an envoy nor a string raises a TypeError saying
    what modules= accepts and that globs, regexes and predicates are not matched
    against the tree. A lone string is treated as one target — iterating it would
    otherwise take it a character at a time, which was another silent empty cache.

Docstrings for tracer.cache(modules=...) and Cache.targets, and one
paragraph in docs/usage/cache.md under "Cache a subset of modules", say what a
path has to be and how to select modules by pattern (filter model.modules()
yourself).

Tests

12 tests in tests/test_interleaving.py::TestCache (the tiny Envoy(Net())
fixture that class already uses): a path-string subset and a lone path string as
the positive controls, an alias target resolving to the real path, the four
misses from the repro plus model.l1.weight, the missing-model-name suggestion,
.output as a target, and re.compile(...) / a lambda as modules=.

Ran with PYTHONPATH=src:

  • tests/test_interleaving.py — 89 passed, 10 failed; the same 10 fail on
    origin/0.8 untouched (they are the .save()-on-a-Tensor family, i.e. the
    trace-body capture problem another item covers), 77 passed there.
  • tests/test_language.py -k cache — 7 passed, 1 pre-existing failure
    (test_cache_per_invoke_holds_its_own_rows), identical against a pristine
    src checkout.
  • tests/test_language.py tests/test_batching.py — 43 passed / 94 failed, byte
    for byte the same counts as the untouched branch.
  • tests/test_fragments.py tests/test_envoy.py — 92 passed / 8 failed, same as
    untouched.

No GPU or remote paths are touched by this change.

Deliberately left alone

  • A lone envoy passed as modules= (not in a list) still iterates its
    children, because Envoy.__iter__ exists: modules=model.transformer.h means
    the 12 blocks today, and a leaf envoy passed that way still caches nothing.
    Making a lone envoy one target would change that first, documented-by-use
    behaviour, so it stayed out of a validation fix.
  • docs/usage/cache.md's separate include_inputs=True / DynamicCache
    problem, which another change is addressing in the same file.

🤖 Generated with Claude Code

A target string was taken as given, so a typo, a path missing the model's own
name, a glob or a regex subscribed to a location nothing ever provides: an empty
cache and no error. Resolve each one in Cache.__init__ and let the envoy tree
raise, and say what modules= accepts for a callable or a regex, which failed
with "'function' object is not iterable".

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant