interp-engine 1.11.2__tar.gz → 1.11.3__tar.gz

{interp_engine-1.11.2 → interp_engine-1.11.3}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.5
 Name: interp-engine
-Version: 1.11.2
+Version: 1.11.3
 Summary: A fast, standardized, and easy to use interpretability engine.
 Project-URL: Homepage, https://github.com/decoderesearch/interp-engine
 Project-URL: Repository, https://github.com/decoderesearch/interp-engine

{interp_engine-1.11.2 → interp_engine-1.11.3}/interp_engine/vllm_capture/requests.py

@@ -65,7 +65,7 @@ from interp_engine.vllm_capture.capture import worker_addresses
 from interp_engine.vllm_capture.graphs import graph_debug, refuse_writes_reason
 from interp_engine.vllm_capture.lens import _make_lens_modifier
 from interp_engine.vllm_capture.mhc import mhc_taps, require_available, require_steerable
-from interp_engine.vllm_capture.steering import _make_steer_modifier
+from interp_engine.vllm_capture.steering import _compose_modifiers, _make_steer_modifiers
 
 # --- the per-request combined hook body ---------------------------------------
 
@@ -679,9 +679,11 @@ def worker_register_steering(
     demux.registered.add(req_id)
     mods = demux.steer_mods.setdefault(req_id, {})
     skip_set = {int(i) for i in (skip_positions or [])}
+    by_site: dict[Address, list[dict]] = {}
     for s in specs:
-        site = _write_site(worker, s)
-        mods[site] = (_make_steer_modifier(s, demux.dev, demux.dt), skip_set, int(prompt_len))
+        by_site.setdefault(_write_site(worker, s), []).append(s)
+    for site, group in by_site.items():
+        mods[site] = (_make_steer_modifiers(group, demux.dev, demux.dt), skip_set, int(prompt_len))
         _ensure_hook(worker, demux, site)
 
 
@@ -708,10 +710,12 @@ def worker_register_lens(
     demux.registered.add(req_id)
     lm = demux.lens_mods.setdefault(req_id, {})
     skip_set = {int(i) for i in (skip_positions or [])}
+    by_site: dict[Address, list[dict]] = {}
     for s in specs:
-        site = _write_site(worker, s)
+        by_site.setdefault(_write_site(worker, s), []).append(s)
+    for site, group in by_site.items():
         lm[site] = (
-            _make_lens_modifier(s, demux.dev, demux.dt),
+            _compose_modifiers([_make_lens_modifier(s, demux.dev, demux.dt) for s in group]),
             bool(steer_generated),
             skip_set,
             int(prompt_len),

{interp_engine-1.11.2 → interp_engine-1.11.3}/interp_engine/vllm_capture/static.py

@@ -1812,9 +1812,9 @@ def worker_set_static_delta(
 ) -> None:
     """Install static writes from worker specs. Zeros every write site first.
 
-    ``op="additive"`` without a lens scope fills the static ``delta`` buffer. Orthogonal,
-    projection_cap, and lens ops attach a live ``modify`` that reads the residual each
-    forward (breakable ``add_eager``). ``lens_scope`` is the jlens prefill/decode skip.
+    A site of only plain adds, with no lens scope, fills the static ``delta`` buffer with their
+    sum. Any other op attaches one live ``modify`` for the site's ops, in order, that reads the
+    residual each forward (breakable ``add_eager``). ``lens_scope`` is the jlens skip.
 
     A ``stream`` disqualifies the static buffer too. The buffer is added whole to a ``[tokens,
     streams, width]`` activation, so it has no way to say "this stream and not the others"; taking
@@ -1827,59 +1827,78 @@ def worker_set_static_delta(
     if static is None:
         raise RuntimeError("set_static_delta: this worker has no static wraps")
     worker_clear_static_delta(worker)
-    from interp_engine.vllm_capture.lens.intervene import _make_lens_modifier
-    from interp_engine.vllm_capture.steering import _make_steer_modifier
+    for site, group in _group_by_site(static, specs).items():
+        assert site.delta is not None
+        constant = _constant_delta(group, site.delta.device, site.delta.dtype) if lens_scope is None else None
+        if constant is not None:
+            site.delta.copy_(constant.expand_as(site.delta))
+            site.delta_set = True
+            continue
+        site.modify = _group_modifier(group, site.delta.device, site.delta.dtype)
+        site.lens_scope = dict(lens_scope) if lens_scope else None
+
 
+def _group_by_site(static: StaticState, specs: list[dict]) -> dict[_Site, list[dict]]:
+    """``specs`` per write site, in order, each op checked against what a static tap can serve."""
+    out: dict[_Site, list[dict]] = {}
     for spec in specs:
         op = str(spec.get("op", SteerMethod.ADDITIVE))
         if op not in STATIC_WRITE_OPS:
             raise ValueError(f"static taps cannot serve op={op!r}; supported ops are {sorted(STATIC_WRITE_OPS)}")
-        site = _write_site(static, Address(str(spec["point"]), int(spec["layer"])))
-        assert site.delta is not None
-        if op == SteerMethod.ADDITIVE and lens_scope is None and spec.get("stream") is None:
-            vec = torch.tensor(spec["vector"], dtype=torch.float32, device=site.delta.device)
-            vec = (vec * float(spec["coeff"])).to(dtype=site.delta.dtype)
-            site.delta.copy_(vec.reshape(1, -1).expand_as(site.delta))
-            site.delta_set = True
-            continue
-        device, dtype = site.delta.device, site.delta.dtype
-        if op in _LENS_OPS:
-            site.modify = _make_lens_modifier(spec, device, dtype)
-        else:
-            site.modify = _make_steer_modifier(spec, device, dtype)
-        site.lens_scope = dict(lens_scope) if lens_scope else None
+        site = _write_site(static, Address(str(spec.get("point") or "resid_post"), int(spec["layer"])))
+        out.setdefault(site, []).append(spec)
+    return out
+
+
+def _constant_delta(group: list[dict], device: torch.device, dtype: torch.dtype) -> torch.Tensor | None:
+    """The ``[1, width]`` sum of ``group``'s deltas when every op is a plain add; else ``None``.
+
+    A constant vector broadcasts over a stream axis and so cannot exclude one: a ``stream`` has to
+    go the modifier way.
+    """
+    if any(
+        str(s.get("op", SteerMethod.ADDITIVE)) != SteerMethod.ADDITIVE or s.get("stream") is not None for s in group
+    ):
+        return None
+    total = sum(torch.tensor(s["vector"], dtype=torch.float32, device=device) * float(s["coeff"]) for s in group)
+    return torch.as_tensor(total).to(dtype=dtype).reshape(1, -1)
+
+
+def _group_modifier(group: list[dict], device: torch.device, dtype: torch.dtype):
+    """One ``modify`` for every op in ``group``, steer and lens ops alike, applied in order."""
+    from interp_engine.vllm_capture.lens.intervene import _make_lens_modifier
+    from interp_engine.vllm_capture.steering import _compose_modifiers, _make_steer_modifier
+
+    return _compose_modifiers(
+        [
+            _make_lens_modifier(s, device, dtype)
+            if str(s.get("op", SteerMethod.ADDITIVE)) in _LENS_OPS
+            else _make_steer_modifier(s, device, dtype)
+            for s in group
+        ]
+    )
 
 
 def _compile_write_req(
-    spec: dict,
+    group: list[dict],
     site: _Site,
     *,
     skip_positions: tuple[int, ...],
     prompt_len: int,
     steer_generated: bool,
 ) -> _WriteReq:
-    op = str(spec.get("op", SteerMethod.ADDITIVE))
-    if op not in STATIC_WRITE_OPS:
-        raise ValueError(f"static taps cannot serve op={op!r}; supported ops are {sorted(STATIC_WRITE_OPS)}")
     assert site.delta is not None
     device, dtype = site.delta.device, site.delta.dtype
-    # A constant `[1, width]` vector broadcasts over a stream axis and so cannot exclude one; see
-    # `worker_set_static_delta` for why a `stream` therefore has to go the modifier way.
-    if op == SteerMethod.ADDITIVE and spec.get("stream") is None:
-        vec = torch.tensor(spec["vector"], dtype=torch.float32, device=device)
-        vec = (vec * float(spec["coeff"])).to(dtype=dtype).reshape(1, -1)
+    constant = _constant_delta(group, device, dtype)
+    if constant is not None:
         return _WriteReq(
-            vector=vec,
+            vector=constant,
             skip_positions=skip_positions,
             prompt_len=prompt_len,
             steer_generated=steer_generated,
         )
-    from interp_engine.vllm_capture.lens.intervene import _make_lens_modifier
-    from interp_engine.vllm_capture.steering import _make_steer_modifier
-
-    modify = _make_lens_modifier(spec, device, dtype) if op in _LENS_OPS else _make_steer_modifier(spec, device, dtype)
     return _WriteReq(
-        modify=modify,
+        modify=_group_modifier(group, device, dtype),
         skip_positions=skip_positions,
         prompt_len=prompt_len,
         steer_generated=steer_generated,
@@ -1906,10 +1925,9 @@ def worker_register_static_write(
         length = int(lens_scope.get("prompt_len") or length)
         generated = bool(lens_scope.get("steer_generated", False))
     by_site: dict[str, _WriteReq] = {}
-    for spec in specs:
-        site = _write_site(static, Address(str(spec.get("point") or "resid_post"), int(spec["layer"])))
+    for site, group in _group_by_site(static, specs).items():
         by_site[format_address(site.address)] = _compile_write_req(
-            spec, site, skip_positions=skip, prompt_len=length, steer_generated=generated
+            group, site, skip_positions=skip, prompt_len=length, steer_generated=generated
         )
     static.write_reqs[req_id] = by_site
     static.registered.add(req_id)

{interp_engine-1.11.2 → interp_engine-1.11.3}/interp_engine/vllm_capture/steering.py

@@ -51,6 +51,32 @@ def _one_stream(modify, stream: int):
     return _modify
 
 
+def _compose_modifiers(mods: list):
+    """One ``modify`` for every op at one site, applied in order as eager applies a layer's ops.
+
+    Each op reads the residual the ops before it wrote, and the delta returned is their sum, so a
+    site that holds one write slot still carries all of them.
+    """
+    if len(mods) == 1:
+        return mods[0]
+
+    def _modify_all(full: torch.Tensor) -> torch.Tensor:
+        total = mods[0](full)
+        out = full + total
+        for modify in mods[1:]:
+            delta = modify(out)
+            out = out + delta
+            total = total + delta
+        return total
+
+    return _modify_all
+
+
+def _make_steer_modifiers(specs: list[dict], dev, dt):
+    """:func:`_compose_modifiers` over ``specs``, which all write one site."""
+    return _compose_modifiers([_make_steer_modifier(s, dev, dt) for s in specs])
+
+
 def _make_steer_modifier(spec: dict, dev, dt):
     """Return ``modify(full_resid) -> delta`` (the tensor to ADD to the residual).
 

{interp_engine-1.11.2 → interp_engine-1.11.3}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "interp-engine"
-version = "1.11.2"
+version = "1.11.3"
 description = "A fast, standardized, and easy to use interpretability engine."
 readme = "README.md"
 license = "Apache-2.0"

{interp_engine-1.11.2 → interp_engine-1.11.3}/tests/test_steer_math_parity.py

@@ -320,3 +320,111 @@ def test_the_worker_refuses_the_old_spelling() -> None:
         _make_steer_modifier(
             {"op": "add", "vector": _vector().tolist(), "coeff": 1.0}, torch.device("cpu"), torch.float32
         )
+
+
+# --- several ops at one site --------------------------------------------------------------------
+#
+# Eager applies a layer's ops in order, each reading the residual the one before it wrote. The
+# worker holds one write per site, so it has to fold the ops into that one write. It used to keep
+# only the last, so two features on one layer steered with the second alone.
+
+
+def _other_vector(seed: int) -> torch.Tensor:
+    return torch.randn(D_MODEL, generator=torch.Generator().manual_seed(seed))
+
+
+def _three_ops_at_one_layer() -> SteeringSpec:
+    return SteeringSpec(
+        layers={
+            0: LayerSteeringSpec(
+                operations=[
+                    AddSpec(vector=_other_vector(1), scale=2.0),
+                    OrthogonalDecompSpec(vector=_other_vector(2), coeff=0.5),
+                    ProjectionCapSpec(vector=_other_vector(3), min=None, max=0.1),
+                ]
+            )
+        }
+    )
+
+
+def _eager_in_order(spec: SteeringSpec, residual: torch.Tensor) -> torch.Tensor:
+    out = residual
+    for one in steering_spec_to_eager_specs(spec):
+        out = out + steer_delta(one, out, one.vector)
+    return out
+
+
+def test_ops_at_one_site_compose_on_the_worker_as_eager_applies_them() -> None:
+    from interp_engine.vllm_capture.steering import _make_steer_modifiers
+
+    spec, residual = _three_ops_at_one_layer(), _residual()
+    modify = _make_steer_modifiers(steering_spec_to_worker_specs(spec), residual.device, residual.dtype)
+    torch.testing.assert_close(residual + modify(residual), _eager_in_order(spec, residual))
+
+
+def test_one_op_at_a_site_is_its_own_modifier() -> None:
+    from interp_engine.vllm_capture.steering import _make_steer_modifiers
+
+    spec = {"op": "orthogonal", "vector": _vector().tolist(), "coeff": 0.5}
+    residual = _residual()
+    torch.testing.assert_close(
+        _make_steer_modifiers([spec], residual.device, residual.dtype)(residual), _worker_delta(spec, residual)
+    )
+
+
+def test_plain_adds_at_a_static_site_are_summed_into_one_constant() -> None:
+    from interp_engine.vllm_capture.static import _constant_delta
+
+    a, b = _other_vector(1), _other_vector(2)
+    group = [
+        {"op": "additive", "vector": a.tolist(), "coeff": 2.0},
+        {"op": "additive", "vector": b.tolist(), "coeff": -1.0},
+    ]
+    got = _constant_delta(group, torch.device("cpu"), torch.float32)
+    assert got is not None
+    torch.testing.assert_close(got[0], a * 2.0 - b)
+
+
+@pytest.mark.parametrize(
+    "extra",
+    [
+        {"op": "orthogonal", "vector": [1.0] * D_MODEL, "coeff": 0.5},
+        {"op": "additive", "vector": [1.0] * D_MODEL, "coeff": 1.0, "stream": 0},
+    ],
+)
+def test_a_static_site_with_any_other_op_or_a_stream_needs_the_modifier(extra: dict) -> None:
+    from interp_engine.vllm_capture.static import _constant_delta
+
+    plain = {"op": "additive", "vector": [1.0] * D_MODEL, "coeff": 1.0}
+    assert _constant_delta([plain, extra], torch.device("cpu"), torch.float32) is None
+
+
+def _static_worker(rows: int):
+    """A worker with one static write site at ``resid_post.0``, and nothing else."""
+    from types import SimpleNamespace
+
+    from interp_engine.address import Address
+    from interp_engine.vllm_capture.static import StaticState, _Site
+
+    site = _Site(Address("resid_post", 0), delta=torch.zeros(rows, D_MODEL))
+    return SimpleNamespace(_ie_static=StaticState(writes={"resid_post.0": site}), model_runner=None), site
+
+
+@pytest.mark.parametrize("static_path", ["per_request", "global"])
+def test_ops_at_one_static_site_compose_as_eager_applies_them(static_path: str) -> None:
+    from interp_engine.vllm_capture.static import (
+        _apply_write,
+        worker_register_static_write,
+        worker_set_static_delta,
+    )
+
+    spec, residual = _three_ops_at_one_layer(), _residual()
+    specs = [{**s, "point": "resid_post"} for s in steering_spec_to_worker_specs(spec)]
+    worker, site = _static_worker(residual.shape[0])
+    if static_path == "per_request":
+        worker_register_static_write(worker, "r", specs)
+    else:
+        worker_set_static_delta(worker, specs)
+    hidden = residual.clone()
+    _apply_write(hidden, None, site, residual.shape[0], fused=False, worker=worker)
+    torch.testing.assert_close(hidden, _eager_in_order(spec, residual))

interp_engine-1.11.3/tests/test_steer_ops_per_site_gpu.py

@@ -0,0 +1,132 @@
+"""Several steer ops at one site: vLLM against eager, hooked and static.
+
+Eager applies a layer's ops in order. The vLLM worker holds one write per site, and used to keep
+only the last op there, so two features on one layer steered with the second alone. Hooked and
+static vLLM agreed with each other, so a hooked-vs-static parity check could not see it: the
+reference here is eager.
+
+Two vLLM engines cannot share this process (static sets ``VLLM_USE_BREAKABLE_CUDAGRAPH``
+process-wide), so hooked runs first, then shuts down, then static.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import gc
+
+import pytest
+import torch
+from harness import require_vllm
+
+from interp_engine.address import Address
+
+require_vllm()
+
+pytestmark = [
+    pytest.mark.gpu,
+    pytest.mark.skipif(not torch.cuda.is_available(), reason="the vLLM backend initializes on CUDA"),
+]
+
+MODEL = "openai-community/gpt2"
+PROMPT = "The capital of France is Paris, and the capital of Germany is"
+MID = Address("resid_post", 6)
+DOWN = Address("resid_post", 9)
+COSINE_MIN = 0.999
+LOAD_KW = {"dtype": "float32", "max_model_len": 512, "gpu_memory_utilization": 0.2}
+
+
+def _vector(width: int, seed: int, norm: float) -> torch.Tensor:
+    v = torch.randn(width, generator=torch.Generator().manual_seed(seed))
+    return v / v.norm() * norm
+
+
+def _specs(width: int) -> dict:
+    from interp_engine.steer_specs import (
+        AddSpec,
+        LayerSteeringSpec,
+        OrthogonalDecompSpec,
+        ProjectionCapSpec,
+        SteeringSpec,
+    )
+
+    def at_mid(*ops) -> SteeringSpec:
+        return SteeringSpec(layers={int(MID.layer): LayerSteeringSpec(operations=list(ops))}, point="resid_post")
+
+    a, b, c = _vector(width, 1, 60.0), _vector(width, 2, 60.0), _vector(width, 3, 1.0)
+    return {
+        # All plain adds: the static constant path.
+        "adds": at_mid(AddSpec(vector=a, scale=1.0), AddSpec(vector=b, scale=1.0)),
+        # An op that reads the residual: the static modifier path.
+        "mixed": at_mid(
+            AddSpec(vector=a, scale=1.0),
+            OrthogonalDecompSpec(vector=c, coeff=0.0),
+            ProjectionCapSpec(vector=b, min=None, max=-5.0),
+        ),
+        "last_add_only": at_mid(AddSpec(vector=b, scale=1.0)),
+    }
+
+
+def _cpu(captures: dict[Address, torch.Tensor]) -> torch.Tensor:
+    return captures[DOWN].detach().float().cpu()
+
+
+def _min_cosine(a: torch.Tensor, b: torch.Tensor) -> float:
+    assert a.shape == b.shape, f"{tuple(a.shape)} vs {tuple(b.shape)}"
+    return float(torch.nn.functional.cosine_similarity(a, b, dim=-1).min())
+
+
+@pytest.fixture(scope="module")
+def loop():
+    """One event loop for the module -- AsyncLLM dies if asyncio.run closes the loop."""
+    made = asyncio.new_event_loop()
+    asyncio.set_event_loop(made)
+    yield made
+    asyncio.set_event_loop(None)
+    made.close()
+
+
+@pytest.fixture(scope="module")
+def runs(loop):
+    from transformers import AutoTokenizer
+
+    from interp_engine import load_model
+
+    tokens = list(AutoTokenizer.from_pretrained(MODEL)(PROMPT)["input_ids"])
+
+    eager = load_model(MODEL, backend="eager", dtype="float32", device="cuda")
+    specs = _specs(int(eager.d_model))
+    want = {k: _cpu(loop.run_until_complete(eager.capture(tokens, [DOWN], steering_spec=s))) for k, s in specs.items()}
+    eager = None
+    gc.collect()
+    torch.cuda.empty_cache()
+
+    async def _vllm(model) -> dict:
+        return {k: _cpu(await model.capture(tokens, [DOWN], steering_spec=specs[k])) for k in ("adds", "mixed")}
+
+    hooked = load_model(MODEL, backend="vllm", **LOAD_KW)
+    loop.run_until_complete(hooked.warmup())
+    hooked_out = loop.run_until_complete(_vllm(hooked))
+    loop.run_until_complete(hooked.shutdown())
+    hooked = None
+    gc.collect()
+    torch.cuda.empty_cache()
+
+    static = load_model(MODEL, backend="vllm-static", static_points=[DOWN], static_writes=[MID], **LOAD_KW)
+    loop.run_until_complete(static.warmup())
+    try:
+        static_out = loop.run_until_complete(_vllm(static))
+    finally:
+        loop.run_until_complete(static.shutdown())
+    return {"eager": want, "hooked": hooked_out, "static": static_out}
+
+
+def test_the_reference_would_see_a_dropped_op(runs) -> None:
+    """Keeping only the last op moves the downstream residual well past the tolerance."""
+    assert _min_cosine(runs["eager"]["adds"], runs["eager"]["last_add_only"]) < 0.99
+
+
+@pytest.mark.parametrize("backend", ["hooked", "static"])
+@pytest.mark.parametrize("spec", ["adds", "mixed"])
+def test_every_op_at_a_site_applies_as_on_eager(runs, backend: str, spec: str) -> None:
+    worst = _min_cosine(runs[backend][spec], runs["eager"][spec])
+    assert worst >= COSINE_MIN, f"{backend} {spec}: min cosine {worst:.6f} against eager"

{interp_engine-1.11.2 → interp_engine-1.11.3}/tests/test_vllm_hyper_connections.py

@@ -1385,10 +1385,10 @@ def test_a_static_additive_steer_confined_to_one_stream_does_not_take_the_consta
 
     site = _Site(Address("resid_streams", 0), delta=torch.zeros(8, HC_MULT, D_MODEL))
     spec = {"op": "additive", "vector": [1.0] * D_MODEL, "coeff": 1.0}
-    plain = _compile_write_req(spec, site, skip_positions=(), prompt_len=0, steer_generated=True)
+    plain = _compile_write_req([spec], site, skip_positions=(), prompt_len=0, steer_generated=True)
     assert plain.vector is not None and plain.modify is None, "no stream is still the constant path"
 
-    scoped = _compile_write_req({**spec, "stream": 2}, site, skip_positions=(), prompt_len=0, steer_generated=True)
+    scoped = _compile_write_req([{**spec, "stream": 2}], site, skip_positions=(), prompt_len=0, steer_generated=True)
     assert scoped.vector is None and scoped.modify is not None
     delta = scoped.modify(torch.zeros(TOKENS, HC_MULT, D_MODEL))
     torch.testing.assert_close(delta[:, 2, :], torch.ones(TOKENS, D_MODEL))