diff --git a/docs/usage/cache.md b/docs/usage/cache.md index 818f5484..7c9bafdf 100644 --- a/docs/usage/cache.md +++ b/docs/usage/cache.md @@ -63,6 +63,17 @@ with model.trace("The Eiffel Tower is in") as tracer: ]) ``` +A path string is the **whole** path, spelled the way the cache's keys are +(`"model.lm_head"`, not `"lm_head"`), and is resolved against the envoy tree when +the cache is created: one that names no module raises there rather than +returning an empty cache. Only envoys and paths are accepted — a glob +(`"model.transformer.h.*"`), a regex, or a predicate matches nothing and raises +too. To select modules by pattern, filter them yourself and pass the envoys: + +```python +mlps = [envoy for envoy in model.modules() if envoy.path.endswith(".mlp")] +``` + ### Include inputs ```python diff --git a/src/nnsight/intervention/cache.py b/src/nnsight/intervention/cache.py index 59aa810a..cec4ef9c 100644 --- a/src/nnsight/intervention/cache.py +++ b/src/nnsight/intervention/cache.py @@ -20,6 +20,8 @@ from __future__ import annotations +import inspect +from collections.abc import Iterable from dataclasses import dataclass from typing import TYPE_CHECKING, Any @@ -66,7 +68,9 @@ class Cache: Attributes: model: The root envoy, used to resolve paths / aliases for [`CacheView`][nnsight.intervention.cache.CacheView]. - targets: The module paths to keep, or ``None`` to keep every module. + targets: The module paths to keep, resolved against the envoy tree (see + [`_resolve_targets`][nnsight.intervention.cache.Cache._resolve_targets]), + or ``None`` to keep every module. entries: Recorded values, ``{module_path: [Entry, ...]}``. """ @@ -90,15 +94,88 @@ def __init__( self.include_inputs = include_inputs # None => every module; else the exact set of paths to keep. self.targets: set[str] | None = ( - None - if modules is None - else {m if isinstance(m, str) else m.path for m in modules} + None if modules is None else self._resolve_targets(modules) ) self.entries: dict[str, list[Entry]] = {} # Which slots of each path's newest Entry have been written this visit. # Recording-time bookkeeping only; see `_record`. self._open: dict[str, set[str]] = {} + def _resolve_targets(self, modules: Any) -> set[str]: + """The module paths ``modules=`` names, raising for anything that names none. + + A target is matched against the paths the run reports, so a string that + resolves to no module -- a typo, a path missing the model's own name, a + glob, a regex -- subscribed to a location nothing ever provides: an empty + cache and no error. Resolving each one against the envoy tree here gives + the same error reading that module by attribute does. + """ + if isinstance(modules, str) or not isinstance(modules, Iterable): + # A lone target rather than a list of them: a string would otherwise + # be taken a character at a time, and a callable or a regex is not + # iterable at all -- `_resolve` says what `modules=` accepts. + modules = [modules] + + return {self._resolve(module) for module in modules} + + def _resolve(self, module: Any) -> str: + """One ``modules=`` entry as a module path.""" + from .envoy import Envoy + + if isinstance(module, Envoy): + return module.path + if not isinstance(module, str): + raise TypeError( + f"cache modules= takes envoys (model.transformer.h[0]) or their " + f"paths ('model.transformer.h.0'), not {type(module).__name__}. " + f"Globs, regexes and predicates are not matched against the tree: " + f"select the modules yourself and pass them." + ) + + root = self.model.path + if module == root: + return root + if not module.startswith(f"{root}."): + # A target is spelled the way the cache's own keys are, from the + # model's own name down; the common miss is dropping that name, so + # say so when adding it back does resolve. + try: + suggestion = self._resolve(f"{root}.{module}") + except (AttributeError, TypeError): + suggestion = None + raise AttributeError( + f"{module!r} is not a module path" + + ( + f" -- did you mean {suggestion!r}?" + if suggestion is not None + else f": no module of {root!r} is named that." + ) + ) + + envoy: Any = self.model + for name in module[len(root) + 1 :].split("."): + # `.output` / `.input` are values the run serves, not modules; reading + # one here would request it from the run instead of resolving a path. + if inspect.isdatadescriptor(getattr(type(envoy), name, None)): + raise AttributeError( + f"{module!r} names {name!r}, which is a value nnsight serves " + f"on every module, not a module. Cache targets are modules." + ) + try: + envoy = getattr(envoy, name) + except AttributeError as error: + # The envoy's own miss, named against the target it came from — + # otherwise a list of paths doesn't say which one is wrong. + raise AttributeError( + f"{module!r} is not a module path: {error}" + ) from error + if not isinstance(envoy, Envoy): + raise AttributeError( + f"{module!r} is not a module: {name!r} is a " + f"{type(envoy).__name__}." + ) + return envoy.path + def __getstate__(self) -> dict: # `model` is a live Envoy used only for CacheView navigation, and it drags # in the interleaver's unpicklable instrumented forwards. Drop it so the diff --git a/src/nnsight/intervention/tracer.py b/src/nnsight/intervention/tracer.py index f6bc610e..8eac2af3 100644 --- a/src/nnsight/intervention/tracer.py +++ b/src/nnsight/intervention/tracer.py @@ -200,7 +200,11 @@ def cache( per step (``len(cache[path])`` is the step count). Args: - modules: Envoys or path strings to capture; ``None`` captures every module. + modules: Envoys or their paths (``"model.transformer.h.0"``, the + whole path, as the cache's keys are) to capture; ``None`` + captures every module. A path is resolved against the envoy + tree, so one that names no module raises rather than caching + nothing; globs, regexes and predicates are not matched. device: Device to move captured tensors to (default CPU); ``None`` leaves them. dtype: Optional dtype to cast captured tensors to. detach: Detach captured tensors from the autograd graph. diff --git a/tests/test_interleaving.py b/tests/test_interleaving.py index e05af66c..664ff2f1 100644 --- a/tests/test_interleaving.py +++ b/tests/test_interleaving.py @@ -1,4 +1,6 @@ +import re + import pytest import torch import nnsight @@ -753,6 +755,52 @@ def test_subset_by_reference(self, envoy, x): assert cache.keys() == ["model.l2"] assert "model.l1" not in cache + # -- modules= is resolved against the tree ----------------------------- + + def test_subset_by_path_string(self, envoy, x): + with envoy.trace(x) as tracer: + cache = tracer.cache(modules=["model.l2"]) + assert cache.keys() == ["model.l2"] + + def test_a_lone_path_string_is_one_target(self, envoy, x): + # Not a list: iterating the string would take it a character at a time. + with envoy.trace(x) as tracer: + cache = tracer.cache(modules="model.l2") + assert cache.keys() == ["model.l2"] + + def test_alias_target_resolves_to_the_real_path(self, model, x): + renamed = Envoy(model, rename={"l1": "first"}) + with renamed.trace(x) as tracer: + cache = tracer.cache(modules=["model.first"]) + assert cache.keys() == ["model.l1"] + + @pytest.mark.parametrize( + "path", ["model.ln", "model.l1.lin", "model.l*", r".*\.l1$", "model.l1.weight"] + ) + def test_a_path_naming_no_module_raises(self, envoy, x, path): + # Targets used to be taken as given, so a typo (or a glob, or a regex) + # subscribed to a location nothing provides: an empty cache, no error. + with pytest.raises(AttributeError): + with envoy.trace(x) as tracer: + tracer.cache(modules=[path]) + + def test_a_path_missing_the_model_name_suggests_the_whole_one(self, envoy, x): + with pytest.raises(AttributeError, match="did you mean 'model.l1'"): + with envoy.trace(x) as tracer: + tracer.cache(modules=["l1"]) + + def test_a_served_value_is_not_a_module(self, envoy, x): + with pytest.raises(AttributeError, match="not a module"): + with envoy.trace(x) as tracer: + tracer.cache(modules=["model.l1.output"]) + + @pytest.mark.parametrize("target", [re.compile("l1"), lambda envoy: True]) + def test_a_target_that_is_neither_envoy_nor_path_says_so(self, envoy, x, target): + # `modules=` raised "'function' object is not iterable". + with pytest.raises(TypeError, match="cache modules="): + with envoy.trace(x) as tracer: + tracer.cache(modules=target) + def test_navigation_matches_path(self, envoy, x): with envoy.trace(x) as tracer: cache = tracer.cache()