From c2b7431219256778cbbb1fb6e4a2e572be31a020 Mon Sep 17 00:00:00 2001 From: khaiwang Date: Sun, 7 Jun 2026 21:37:33 -0400 Subject: [PATCH 01/30] feat(intervention): isolated GPU-worker execution backend for model.trace() Run a trace's user interventions in a spawned, GPU-enabled worker process so footguns in intervention code (infinite loops, OOM allocations, device-side asserts, host-object pokes) are contained to the worker while the model server keeps serving. Results are bit-identical to in-process execution. The six-event Mediator protocol (VALUE/SWAP/SKIP/BARRIER/END/EXCEPTION) is left unchanged; isolation is an outer harness that spawns the worker and routes the existing protocol over a CUDA-IPC bounce-buffer channel (tensors stay on the GPU, ~0.6 ms/hook, size-independent) instead of a shared Python frame. Two shared-memory assumptions of the in-process path become explicit harness steps: - host-side hook registration: the worker has no real module, so on the first event for a requester the host registers the matching one-shot hook on the real module (resolved from the requester string, for the specific step). - worker->host saves transmission: .save()'d values live in the worker frame + Globals.saves; the worker bundles them into the END event and the host injects them into the real user frame. New sources: - transport.py: CUDA-IPC codec + host/worker channels (clone-on-receive, per-wait timeout, host->worker live-meta piggyback, worker->host push field, cuda.synchronize ordering guard). - isolation.py: isolate_mediators() context, spawn_isolated_worker, _worker_main, on-demand host hook registration, worker interleaver stub + dummy-module map, barrier/variable-store wiring, transmissible-exception degrade. - _sandbox.py: seccomp lock_down for fs/net/exec containment. Seam edits route the protocol through the channel when isolation is on: interleaver.py (isolated start branch, on-demand registration in handle, saves injection at END, host-side barrier counting, _iso/cancel teardown), hooks.py (per-step iteration param on output_hook/input_hook), tracer.py (isolated Barrier branch). Covered, each bit-identical and independently reviewed: read / swap / .save() / multi-invoke / skip / exception / timeout / seccomp lockdown; multi-token iteration (iter[N], iter[:], per-step swap); cross-invoke barrier + variable sharing; non-standard-named models. Not yet built: tracer.cache() (returns an empty CacheDict under isolation), backward/grad (autograd graph is host-side), warm worker pool. See docs/developing/mediator-gpu-trace-integration.md. Co-Authored-By: Claude Opus 4.8 (1M context) --- docs/developing/gpu-sandbox.md | 109 ++++ .../mediator-gpu-trace-integration.md | 376 +++++++++++++ .../mediator-isolation-harness-plan.md | 244 ++++++++ docs/developing/mediator-isolation-sandbox.md | 315 +++++++++++ prototypes/mediator-sandbox/_smoke.py | 6 + .../gpu_sandbox/fault_test.py | 76 +++ .../gpu_sandbox/gpu_sandbox.py | 73 +++ .../gpu_sandbox/gpu_worker.py | 60 ++ .../gpu_sandbox/perf_ctxswitch.py | 74 +++ .../gpu_sandbox/perf_decompose.py | 67 +++ .../mediator-sandbox/gpu_sandbox/perf_test.py | 93 ++++ .../gpu_sandbox/probe_bootstrap.py | 133 +++++ .../gpu_sandbox/probe_saves.py | 101 ++++ .../mediator-sandbox/gpu_sandbox/sandbox.py | 76 +++ .../gpu_sandbox/test_cuda_channel_codec.py | 80 +++ .../gpu_sandbox/test_cuda_channel_ipc.py | 73 +++ .../gpu_sandbox/test_functional.py | 86 +++ .../gpu_sandbox/test_isolated_acceptance.py | 133 +++++ .../test_isolated_backward_cache_gaps.py | 81 +++ .../gpu_sandbox/test_isolated_cache.py | 69 +++ .../gpu_sandbox/test_isolated_cross_invoke.py | 94 ++++ .../test_isolated_lockdown_safety.py | 91 +++ .../test_isolated_multitoken_iter.py | 98 ++++ .../gpu_sandbox/test_isolated_trace.py | 69 +++ .../gpu_sandbox/test_nonstd.py | 62 +++ .../gpu_sandbox/test_safety.py | 141 +++++ .../mediator-sandbox/p2_isolation_poc.py | 221 ++++++++ .../phase2_socket_transport.py | 214 +++++++ .../mediator-sandbox/phase3_jail_transport.py | 153 +++++ .../mediator-sandbox/phase3_jailed_worker.py | 82 +++ .../mediator-sandbox/phase4_cross_tenant.py | 108 ++++ .../phase4_malicious_worker.py | 77 +++ .../mediator-sandbox/phase5_gpu_measure.py | 233 ++++++++ .../phase5b_transport_breakdown.py | 236 ++++++++ .../mediator-sandbox/phase6_jailed_worker.py | 35 ++ .../phase6_shm_safetensors.py | 203 +++++++ src/nnsight/intervention/_sandbox.py | 74 +++ src/nnsight/intervention/hooks.py | 40 +- src/nnsight/intervention/interleaver.py | 244 ++++++-- src/nnsight/intervention/isolation.py | 443 +++++++++++++++ src/nnsight/intervention/tracing/tracer.py | 7 + src/nnsight/intervention/transport.py | 522 ++++++++++++++++++ 42 files changed, 5719 insertions(+), 53 deletions(-) create mode 100644 docs/developing/gpu-sandbox.md create mode 100644 docs/developing/mediator-gpu-trace-integration.md create mode 100644 docs/developing/mediator-isolation-harness-plan.md create mode 100644 docs/developing/mediator-isolation-sandbox.md create mode 100644 prototypes/mediator-sandbox/_smoke.py create mode 100644 prototypes/mediator-sandbox/gpu_sandbox/fault_test.py create mode 100644 prototypes/mediator-sandbox/gpu_sandbox/gpu_sandbox.py create mode 100644 prototypes/mediator-sandbox/gpu_sandbox/gpu_worker.py create mode 100644 prototypes/mediator-sandbox/gpu_sandbox/perf_ctxswitch.py create mode 100644 prototypes/mediator-sandbox/gpu_sandbox/perf_decompose.py create mode 100644 prototypes/mediator-sandbox/gpu_sandbox/perf_test.py create mode 100644 prototypes/mediator-sandbox/gpu_sandbox/probe_bootstrap.py create mode 100644 prototypes/mediator-sandbox/gpu_sandbox/probe_saves.py create mode 100644 prototypes/mediator-sandbox/gpu_sandbox/sandbox.py create mode 100644 prototypes/mediator-sandbox/gpu_sandbox/test_cuda_channel_codec.py create mode 100644 prototypes/mediator-sandbox/gpu_sandbox/test_cuda_channel_ipc.py create mode 100644 prototypes/mediator-sandbox/gpu_sandbox/test_functional.py create mode 100644 prototypes/mediator-sandbox/gpu_sandbox/test_isolated_acceptance.py create mode 100644 prototypes/mediator-sandbox/gpu_sandbox/test_isolated_backward_cache_gaps.py create mode 100644 prototypes/mediator-sandbox/gpu_sandbox/test_isolated_cache.py create mode 100644 prototypes/mediator-sandbox/gpu_sandbox/test_isolated_cross_invoke.py create mode 100644 prototypes/mediator-sandbox/gpu_sandbox/test_isolated_lockdown_safety.py create mode 100644 prototypes/mediator-sandbox/gpu_sandbox/test_isolated_multitoken_iter.py create mode 100644 prototypes/mediator-sandbox/gpu_sandbox/test_isolated_trace.py create mode 100644 prototypes/mediator-sandbox/gpu_sandbox/test_nonstd.py create mode 100644 prototypes/mediator-sandbox/gpu_sandbox/test_safety.py create mode 100644 prototypes/mediator-sandbox/p2_isolation_poc.py create mode 100644 prototypes/mediator-sandbox/phase2_socket_transport.py create mode 100644 prototypes/mediator-sandbox/phase3_jail_transport.py create mode 100644 prototypes/mediator-sandbox/phase3_jailed_worker.py create mode 100644 prototypes/mediator-sandbox/phase4_cross_tenant.py create mode 100644 prototypes/mediator-sandbox/phase4_malicious_worker.py create mode 100644 prototypes/mediator-sandbox/phase5_gpu_measure.py create mode 100644 prototypes/mediator-sandbox/phase5b_transport_breakdown.py create mode 100644 prototypes/mediator-sandbox/phase6_jailed_worker.py create mode 100644 prototypes/mediator-sandbox/phase6_shm_safetensors.py create mode 100644 src/nnsight/intervention/_sandbox.py create mode 100644 src/nnsight/intervention/isolation.py create mode 100644 src/nnsight/intervention/transport.py diff --git a/docs/developing/gpu-sandbox.md b/docs/developing/gpu-sandbox.md new file mode 100644 index 000000000..f4c9f202c --- /dev/null +++ b/docs/developing/gpu-sandbox.md @@ -0,0 +1,109 @@ +# GPU Sandbox — the chosen isolation implementation + +**Status:** Built + tested (functional + safety pass) · **Date:** 2026-06-06 +**Supersedes** the CPU-only transport work (harness-plan Phases 5/5b/6) for the GPU path. + +## Why this design + +The threat model was relaxed to **contain footguns, not defeat a determined adversary** — stop a careless +or buggy intervention from crashing the shared model server, reading host files, exfiltrating, hanging, or +OOMing it; do **not** (for now) defend against a malicious user weaponising the GPU. Under that model the +worker can keep **GPU access**, which makes the data path **zero-copy** and deletes the whole transport +problem (D2H + serialize) that made the CPU-only path cost ~110 ms/hook. + +## Design + +A **separate, spawned, GPU-enabled worker process** runs the user's arbitrary intervention on the real GPU +activation, while a process boundary + seccomp contain everything except the GPU. + +- **Zero-copy via a shared GPU "bounce buffer."** A single CUDA tensor is shared host↔worker via CUDA IPC + **once at spawn** (before lockdown). Per request the host does a cheap on-GPU (D2D) copy of the + activation into the buffer + a 1-byte signal; the worker views the buffer as the activation tensor and + runs the user op in place; the host reads the result back. + + **Measured (`perf_test.py`, A100): flat ~0.6–0.8 ms/hook, independent of size** — 0.59 ms at 0.03 MB, + 0.71 ms at 16.8 MB (~7B), 0.77 ms at 33.6 MB (~70B). The CPU-transport path scaled 1→285 ms with bytes; + this doesn't, because the activation never leaves the GPU (per-request copy is an on-GPU D2D memcpy, + ~0.01–0.04 ms). End-to-end on a gpt2 trace: **+1.1 ms/hook** (10.1 → 11.2 ms for one intervention; 10 → 24 ms for all 12 + layers). vs ~0.02 ms fully in-process. + + **Where the ~0.6 ms goes (decomposed, `perf_decompose.py` / `perf_ctxswitch.py`): GPU context-switching + between the host and worker CUDA contexts — NOT process communication.** Measured: bare process round-trip + 0.04 ms, `cloudpickle.loads` 0.01 ms, one `cuda.synchronize` 0.016 ms. But a round-trip where BOTH host + and worker do a GPU op (apply's pattern) = 0.405 ms vs 0.040 ms when only the host does → **~0.37 ms is the + GPU switching active context host↔worker** each hook (each process has its own context on the shared + device; they alternate every intervention). + + *Fix (measured under a scoped CUDA MPS daemon):* **CUDA MPS routes both processes' kernels through one + shared GPU context**, eliminating the per-process switch. The both-touch-GPU round-trip dropped + **0.405 → 0.132 ms** and the context-switch overhead **0.365 → 0.090 ms (~4×)** — confirming the diagnosis. + So `apply()` under MPS goes from ~0.6 ms toward **~0.2 ms**, end-to-end toward ~+0.3 ms/hook. (A spin-waiting + worker does **not** help — it only removes the ~0.04 ms IPC, not the context switch.) + - **Fault-domain tradeoff — measured on the A100 (`fault_test.py`), better than feared:** MPS shares one + GPU context, so in principle it couples the fault domain. In practice, **every GPU fault reachable from + realistic buggy code was CONTAINED under MPS — the host's CUDA context survived**, same as with separate + contexts. A buggy intervention's out-of-bounds index / embedding produces a **device-side assert** (PyTorch + bounds-checks its indexing), and Volta+ MPS attributes that to the offending client and terminates only it. + I could **not** trigger a genuine raw illegal-memory-access (XID 31) from normal tensor ops — they all + became asserts. The residual coupling is the *rare* worst case (a true bad-pointer illegal access, or a + hardware/ECC fault), which MPS docs say can hit the server. **Net:** MPS gives ~3× *and* keeps containment + of the faults that actually happen — a viable default; separate contexts only buy the rare worst-case. + | mode | per-hook | OOB→device-assert (the real footgun) | rare raw-illegal/HW fault | + |---|---|---|---| + | separate contexts | ~0.6 ms | contained (host survives) ✓ measured | contained | + | Volta+ MPS (A100) | ~0.2 ms | **contained (host survives) ✓ measured** | shared blast radius | + - *Why a bounce buffer and not per-request IPC:* `cudaIpcOpenMemHandle` (per-request rebuild) **fails + after the seccomp lockdown** (`context is destroyed`). Mapping one buffer before lockdown sidesteps it; + the per-request D2D copy is on-GPU memcpy (microseconds), not the bottleneck. +- **Footgun containment (`sandbox.py`):** after CUDA + torch are warmed, the worker installs a hand-rolled + **seccomp-BPF** filter (no external dep) that makes new `open`/`openat`/`openat2` (filesystem), + `socket`/`connect` (network), and `execve` fail with EPERM. CUDA keeps working — it uses already-open + `/dev/nvidia*` fds via ioctl, not new file opens. GPU memory is capped with + `torch.cuda.set_per_process_memory_fraction` (an OOM footgun can't exhaust the device). A per-call + **timeout** kills a wedged worker; the **process boundary** isolates crashes and hides all host Python + objects (weights, interleaver, other tenants). + - *Gotcha:* `RLIMIT_AS` is unusable with CUDA — CUDA reserves tens of GB of *virtual* address space, so + an 8 GB AS cap kills it at startup with `MemoryError`. Use the GPU memory fraction instead. +- **Pool.** CUDA isn't fork-safe → workers are spawned (not forked) and kept warm; each request goes to a + free worker; a crashed worker is detected and replaced. (`GPUSandbox` is one worker for clarity.) + +**Files:** `prototypes/mediator-sandbox/gpu_sandbox/` — `sandbox.py` (seccomp), `gpu_worker.py` (the +spawned worker), `gpu_sandbox.py` (host manager + `apply(activation, fn)`), `test_functional.py`, +`test_safety.py`. Run on a GPU: `CUDA_VISIBLE_DEVICES=N PYTHONPATH=/src .../test_*.py`. + +## Results + +**Functional (`test_functional.py`) — PASS.** Real nnsight `model.trace()` workloads (scale, steer, +mean-ablate, per-token-norm read) on gpt2, run two ways — op inline vs op offloaded to the GPU worker — +produce **bit-identical** logits (`max|Δ|=0`); the worker survives all requests. nnsight delivers the +activation through its real machinery; only the user op runs isolated. (Core nnsight suite: 152 passed.) + +**Safety (`test_safety.py`) — PASS (9/9).** Each mimicked unsafe request is contained: + +| Unsafe request | Outcome | +|---|---| +| read a host secret file | `PermissionError` (open blocked); secret not returned | +| write a host file | `PermissionError`; no host file created | +| network egress (`socket`/`connect`) | `PermissionError` | +| reach a host object (`THE_MODEL.config`) | `NameError` — host objects don't exist in the worker | +| OOM (`torch.empty(1e12)` on GPU) | `OutOfMemoryError` — GPU mem fraction caps it | +| infinite loop | timeout → worker killed; host unaffected | +| hard crash (NULL-deref segfault) | worker dies; **host CUDA context intact**, crash surfaced cleanly | +| legit op after all attacks | still correct (worker survived the catchable ones) | +| fresh worker after a crash | serves correctly (pool recovery) | + +## What this is vs. isn't + +- **Contained:** filesystem, network, host crashes/OOM/hangs, and all host Python objects/other tenants' + CPU-side state. The model server cannot be taken down or read by a buggy intervention. +- **Not contained (the accepted axis):** the GPU. The worker has a CUDA context and shares the activation, + so a *deliberately malicious* user could attack the driver or read neighbouring GPU memory. Out of scope + per the current threat model; revisit (CPU-only worker, MIG, dedicated allocations) if the model tightens. + +## Remaining work + +This is a working harness with an explicit offload API: `output = sandbox.apply(activation, fn)` inside a +trace. Wiring it **transparently into `model.trace()`** (so `.output`-access/ops/`.save()` route to the +worker automatically via the eproperty bridge) is the remaining productionisation — the same integration +as the transparent-trace-integration milestone, but simpler here because there is **no transport codec** (zero-copy GPU buffer). +Plus the warm-pool/scheduler and per-tenant dedicated allocations if the threat model tightens. diff --git a/docs/developing/mediator-gpu-trace-integration.md b/docs/developing/mediator-gpu-trace-integration.md new file mode 100644 index 000000000..20b7636dd --- /dev/null +++ b/docs/developing/mediator-gpu-trace-integration.md @@ -0,0 +1,376 @@ +# Transparent GPU-Worker Trace Integration — Design + +**Status:** Design (approved) · **Date:** 2026-06-06 · **Branch:** `worktree-mediator-sandbox` +**Builds on:** +- [gpu-sandbox.md](gpu-sandbox.md) — the chosen GPU-worker mechanism (`GPUSandbox.apply`, bounce buffer, seccomp, lifecycle). +- [mediator-isolation-harness-plan.md](mediator-isolation-harness-plan.md) — the earlier CPU-transport prototypes; this doc supersedes the transparent-`model.trace()` milestone described there. +- [mediator-isolation-sandbox.md](mediator-isolation-sandbox.md) — threat model (relaxed to "contain footguns, not adversaries") + AWS deployment. + +**Goal:** make `model.trace()` run user interventions in an isolated GPU worker process **transparently** — +identical results to in-process — so a footgun in user code (infinite loop, OOM alloc, bad-index +device assert, host-object poke) is contained to the worker while the shared model server keeps serving. +The standalone `GPUSandbox.apply()` proved the worker *mechanics*; this turns it into the actual +`model.trace()` execution backend. + +--- + +## 1. The architecture: an outer harness, not a rewrite + +The Mediator↔model handoff already routes through the `MediatorChannel` seam +(`interleaver.py:774`, extracted earlier). The Mediator talks to the model **only** through that channel, +via the six events (`VALUE/SWAP/SKIP/BARRIER/END/EXCEPTION`). **Those events do not change.** Isolation is +an outer harness: run the Mediator's worker in a separate (GPU-enabled, locked-down) process, and give it +a cross-process channel. + +``` +HOST (trusted: model + weights + GPU) WORKER (spawned, GPU-enabled, seccomp'd) + real Envoy tree + Interleaver + Batcher path-only Envoy mirror (no weights) + host-side Mediator worker-side Mediator (same class, other half) + · owns channel, batch_group, history, handle_* · runs the REAL user intervention fn + · registers hooks ON DEMAND · eproperty.__get__ → request() → blocks + CudaIpcChannel host end ◄── control pipe ──► CudaIpcChannel worker end + │ + shared GPU buffer │ + forward fires hook → handle() → narrow → user code runs on the GPU tensor view, + D2D into buffer → respond `.save()` accumulates locally, + shipped back at END +``` + +**One `Mediator` class, instantiated on both sides**, each using its half. The existing +`Mediator.__getstate__`/`__setstate__` (`interleaver.py:1558`) already define exactly what crosses to +construct the worker side — they were built for NDIF remote. The host half owns the Batcher and +`handle_*`; the worker half owns `send`/`request`/`swap`/`push` and runs the intervention. + +### What does NOT change +- The six-event protocol and its one-event-in-flight invariant (the "access in forward-pass order" + contract). +- `Mediator.send`/`respond`/`handle`/`handle_*`, the Batcher narrow/swap, requester↔provider matching. +- `eproperty`/`hooks.py` core logic (the worker's hook call lands on a dead dummy module — see below). +- The in-process path: when isolation is off, `Mediator.start` runs the existing `Thread` exactly as today. + +### What is genuinely new (and *why*) +Two places in the in-process code rely on **shared memory** rather than the channel. These — and only +these — are why isolation is more than "spawn + new channel." Both live in the harness layer. + +1. **Hook registration touches the real module via shared memory.** In-process, `eproperty.__get__` + registers a PyTorch hook *on the actual module object* (`hooks.py:307`, `output_hook`). The worker + process has no module. But the event already carries the requester string + (`"transformer.h.6.output.i0"`), which fully encodes (path, kind, iteration) — so the **host** + registers the hook on receipt (the *host-side on-demand hook registration* described in §3). +2. **`.save()` propagates through a shared Python frame.** In-process, trace exit reads the worker's + frame locals, keeps `id(v) in Globals.saves`, and `push_variables` into the user frame — all + same-process pointer work (`tracing/base.py:565`). Across processes the worker's frame and + `Globals.saves` live in the worker, so it runs that filter locally and ships the result over the + channel; the host does the final `push_variables` (the *worker→host saves transmission* in §4). + +--- + +## 2. Transport — `CudaIpcChannel` + +A new `MediatorChannel` implementation, same shape as `transport.py`'s shared-memory socket channel; only +the tensor encode/decode differs. + +- **Control plane** (`mp.Pipe`, small pickle frames): `(event, requester)` + an *offset table* + `[(offset, nbytes, shape, dtype), …]` + a non-tensor skeleton. Reuses `transport.py`'s + `_split_tensors`/`_merge_tensors` verbatim — the only change is "encode each tensor as a buffer + offset" instead of "safetensors bytes." +- **Bulk plane** (shared GPU bounce buffer): the sender D2D-copies each tensor into the buffer at + successive offsets; the receiver views them at those offsets. The buffer is a CUDA tensor created on + the host and shared to the spawned worker via `torch.multiprocessing` (CUDA IPC), mapped once before + lockdown. +- **Aliasing rule (correctness, mandatory):** under one-event-in-flight the single buffer is reused in + strict alternation, so **both sides clone-on-receive** before the next access can overwrite it. This + covers `x = h[6].output; y = h[7].output; use(x, y)` — `x` must be cloned out of the buffer before the + request for `y` overwrites it. A D2D clone is HBM-speed (~µs/MB); "zero-copy" always meant *no PCIe, + no pickle*, not *no copy*. The measured ~0.6 ms/hook already includes these copies. +- **Sizing:** a fixed arena (64 MB default). A value larger than the arena raises a clear error (the + `ShmArena.write` "exceeds arena" guard). Chunking is a later perf item, out of scope. +- **CPU-only fallback backend:** the existing CPU shared-memory socket channel stays config-selectable + for CPU-only isolation — free, since it is already a `MediatorChannel`. GPU-IPC is the default. + +Two roles, mirroring `transport.py`: `CudaIpcHostChannel` (main thread; only `wait_event` reads the +pipe; `has_event`/`get_event`/`restore_event` are host-local buffer ops) and `CudaIpcWorkerChannel` +(send→wait→get). + +--- + +## 3. Host-side on-demand hook registration + +In-process, the worker's `requires_output` registers the hook just-in-time as it single-steps. In +isolation the worker can't (no real module), so the host registers it. + +- **Where:** in the host `handle` loop, when a `VALUE`/`SWAP`/`SKIP` arrives for a requester `R` whose + provider is not yet set up (the existing `handle_value_event` else-branch, `interleaver.py:1203`, + that today does `history.add`/`restore_event`/`return False`). +- **What:** `ensure_provider(R)` — parse `R` (`"..i"`) → `iteration = N`, + `kind`, `path`; resolve `envoy = root_envoy.get(path)` (`envoy.py:586`); call the **existing** + `output_hook`/`input_hook(host_mediator, envoy._module, f"{path}.{kind}")` (`hooks.py:224`/`154`). + The host-side mediator is passed, so the hook closure delivers over the channel. +- **Worker side:** the worker's `eproperty.__get__` still runs `self._hook(obj)`. The worker's + path-only envoy carries a bare `torch.nn.Module()` as `_module`, so `output_hook` registers a hook on + a **dead dummy** that never fires (harmless; reclaimed when the worker process exits — NOT by + `cancel`, since the worker never calls it). **Warm-pool caveat:** a warm worker pool reuses processes, + so these dummy-module hooks (and the worker-local `iteration_tracker`/`Globals.saves`) would accumulate + across traces and need explicit per-trace cleanup. Keeping `eproperty`/`hooks.py` unchanged was the + reason to register on a dead dummy rather than guard `_hook` with an `interleaver.isolated` flag; + revisit if dummy-module accumulation shows up in profiling. +- **Idempotency:** `ensure_provider` registers a given `R` once per mediator (tracked in a set on the + host mediator), mirroring the `current_provider`-skip logic in `requires_output`. + +For a single forward pass the iteration is always `0`, so `R` is `"..i0"`. Multi-token +iteration stamping is handled where the multi-token feature is built (§9). + +--- + +## 4. Worker→host saves transmission + +- **Worker:** at trace exit / on `END`, the worker runs the *existing* `base.py` exit filter (it owns + the frame + `Globals.saves`) to produce the saved dict, then — instead of `push_variables` locally — + ships that dict over the channel. +- **Host:** receives the dict and `push_variables(user_frame, saved_dict)` into the real user frame + (the frame that called `model.trace()`, which lives on the host). +- **Encoding:** tensors in the saved dict ride the bounce buffer (offset table); other saved objects + ride the control pickle, reusing `_split_tensors`. +- **New constraint vs in-process:** a *non-tensor* `.save()` must be serializable to cross the boundary + (in-process it was just a reference). Tensors (the overwhelming common case) are fine; a + non-serializable save raises a clear error naming the variable. Documented as an isolation semantic + difference. + +### Isolation semantic: in-place modification requires explicit assignment +In-process, `.output` returns the model's *real* tensor (shared memory), so an in-place edit +(`h[6].output[0][:] = ...`) propagates by mutation with no SWAP event. In the isolated path the worker +receives a **clone** (the clone-on-receive rule, §2), so an in-place edit mutates a worker-local copy and +does **not** reach the host. Use **explicit assignment** (`h[6].output = ...`), which emits a SWAP event +and applies via `batcher.swap`. This is identical to **NDIF remote** semantics (in-place on a serialized +clone never crossed the wire either), so it is a consistency property, not a new wart. Transparent +in-place support would need a protocol read-back that conflicts with the single-buffer clone-on-receive +rule; deferred (revisit with double-buffering if real workloads need it). + +--- + +## 5. Lifecycle parity + +- **`cancel()`** terminates the worker process (the `gpu_sandbox` `close()`/`terminate()` pattern) + instead of dropping a `Thread`. `check_dangling_mediators` and the EarlyStop drain account for + process workers. +- **Timeout:** a per-call timeout kills a wedged worker (infinite loop in user code). Surfaces a clear + `TimeoutError`; the host is unaffected; a warm pool would respawn. +- **EXCEPTION:** the worker ships the exception over the channel; the host re-raises it in the user's + context with the rebuilt traceback. Dynamic nnsight exceptions that can't pickle are degraded to a + plain exception preserving type name + message before they cross (§11). +- **SKIP / BARRIER / END / EarlyStop:** already cross as events. `END` triggers the saves transmission + then worker exit + host join. `stop()`'s `push()` ships saves before signalling. + +--- + +## 6. Feature coverage map (so nothing is silently dropped) + +| Feature | Cross-process mechanism | Status | +|---|---|---| +| read / swap / skip / exception | events over channel + host-side hook registration | done | +| `.save()` (tensors) | worker→host saves transmission | done | +| multi-invoke + batch narrowing | per-worker host mediator; Batcher host-side | done | +| `iter`/`all`/`next` (multi-token) | iteration step stamped host-side; host iter-hooks bump the tracker; worker sets its step explicitly | done | +| `tracer.barrier()` | host-side participant counting + the existing `handle_barrier_event` coordination loop | done | +| `cross_invoker` variable sharing | host-mediated variable store (worker pushes data locals, pulls the merged store) | done | +| `with tensor.backward()` / `.grad` | needs host-side backward execution (the autograd graph is host-side) | planned | +| `tracer.cache()` | host-side cache-hook registration + post-forward injection of the populated CacheDict | planned | +| warm worker pool / MPS / `isolate_mediators()` polish | — | planned | + +When isolation is on and a not-yet-supported feature is used, the trace **fails cleanly** — a +missed-provider error or the per-step timeout (the lifecycle is the safety net), not a silent deadlock or +silent-wrong result. (There is no automatic "route to in-process" fallback; features are added one at a +time. See the support matrix in §8 for what works today.) + +--- + +## 7. Core single-pass interventions — scope, components, acceptance + +This is the first slice: a single forward pass (iteration 0), covering read / swap / `.save()` / +multi-invoke / skip / exception / lifecycle / seccomp lockdown — the foundation the rest builds on. + +### Scope +Single forward pass; single and multiple invokes; read (`.output`/`.input`/`.inputs`); swap (`=`); tensor +`.save()`; skip; exception; lifecycle (normal END, exception, timeout→kill); tuple/nested block outputs. +`cross_invoker` disabled. Multi-token / barrier / backward / cache not yet handled. + +### Components +1. **`CudaIpcChannel`** (`transport.py`, additive): `CudaIpcHostChannel` + `CudaIpcWorkerChannel` + + GPU-buffer tensor pack/unpack (reusing `_split_tensors`/`_merge_tensors`). +2. **Isolated `Mediator.start`** (`interleaver.py`): a branch that spawns the worker (CUDA → `spawn`), + ships `self` via source-serialization with the model→path-only-envoy persistent-object map, and a + worker bootstrap that builds the interleaver stub + path-only envoy mirror + runs the intervention. +3. **Host-side on-demand hook registration** (`interleaver.py`): `ensure_provider(R)` in `handle`. +4. **Saves transmission + lifecycle** (`interleaver.py` + bootstrap): worker ships `Globals.saves`- + filtered frame locals at END; host injects; `cancel()` kills the process. +5. **Opt-in** (§12): `CONFIG.APP.ISOLATE_MEDIATORS` + `nnsight.isolate_mediators()` context. + +### Acceptance — bit-identical to in-process (`max|Δ|=0` on GPU) +1. read-save; swap; multi-invoke batch-narrow; nested-tuple block output. +2. **A non-standard-named toy model** (`decoder_blocks`, `output_projection`) **and** gpt2 — per the + testing rules: structure derived at runtime via `envoy.get(path)`, no GPT-2-only assumptions, no + hardcoded module-name tables. +3. exception in user code → raised in the user's context with the correct traceback. +4. infinite loop in user code → worker killed, host survives, clear `TimeoutError`. +5. safety: the unsafe payloads (filesystem/network/exec/OOM/crash) contained under seccomp lockdown. +6. core suite still green (in-process path untouched). + +### Status: DONE (2026-06-06) +Built and verified (TDD; harnesses in `prototypes/mediator-sandbox/gpu_sandbox/`): +- **Code:** `transport.py` (`pack_cuda`/`unpack_cuda` + the two channel ends, clone-on-receive, per-wait + timeout); `isolation.py` (`isolate_mediators` + `spawn_isolated_worker` + `_worker_main` + + `ensure_isolated_provider` + `_WorkerInterleaver`/`_WorkerPersistent`); `_sandbox.py` (seccomp + `lock_down`, relocated from the prototype); `interleaver.py` seam (`_iso` field, isolated + `Mediator.start` branch, the on-demand hook-registration call in `handle`, saves injection in + `handle_end_event`, `_iso.close()` in `cancel`). +- **Foundations (probes, all PASS):** buffer codec round-trip + clone-on-receive; CUDA-IPC channel over + a real spawned process; bootstrap (serialize→deserialize-against-dummies→run→correct `request`); + saves capture (`info.frame.f_locals` filtered by `Globals.saves`, keyed by user var name). +- **End-to-end (`test_isolated_trace.py`):** isolated read `.output[0].save()` and explicit-assignment + swap (`h[6].output = … * 2`) **bit-identical to in-process, `max|Δ|=0`** on gpt2/cuda. +- **Acceptance (`test_isolated_acceptance.py`):** non-standard module names (`decoder_blocks`, + `NNsight(net)`) `max|Δ|=0`; two isolated invokes bit-identical + rows distinct (no cross-leak); user + exception → `NNsightException` in the user context; infinite loop killed (timeout), host survives. +- **Safety (`test_isolated_lockdown_safety.py`):** read under `lockdown=True` `max|Δ|=0`; `open()` + blocked + no host file; `socket()/connect()` blocked. +- **Regression:** core in-process suite **155 passed** (in-process path untouched). + +**Findings folded into the design:** per-mediator serialization must set `intervention.__source__` +first (the tracer normally does it in *its* `__getstate__`); the worker frame is a `SerializedFrame` so +`push()` lands locals there; saved values inject into the **tracer's** `info.frame` (the user frame), +not the mediator's; `_remoteable_persistent_objects` is remoteable-only (gate with `isinstance`); +locked-down workers must `os._exit(0)` to skip tempfile atexit (avoids a `RecursionError`). `lockdown` +defaults off in `_STATE` (flip to on once the warm pool lands); the bounce-buffer force-kill emits a +benign CudaIPC release warning. + +--- + +## 8. Support matrix (what works under `isolate_mediators()` today) + +| Feature | Status | +|---|---| +| read / swap (`=`) / `.save()` (tensors) / skip / exception / multi-invoke | ✅ bit-identical | +| single-forward `generate(...)` (no iter) | ✅ verified | +| seccomp lockdown (fs/net/exec) | ✅ | +| `iter`/`all`/`next` (multi-token) | ✅ bit-identical (`iter[N]`, `iter[:]`, per-step swap) | +| `tracer.barrier()` | ✅ host-side participant counting | +| `cross_invoker` variable sharing | ✅ host variable store; transmittable data vars only — see §10 | +| `with tensor.backward()` / `.grad` | 🔜 hard: the autograd graph is host-side, the worker has detached clones; needs the backward pass to run host-side with path-based grad providers (a major build) | +| `tracer.cache()` | 🔜 tractable: returns an empty CacheDict today (hooks fire on dummy modules); needs host-side cache-hook registration + shipping the populated CacheDict back | +| `.source` operation-level access (`...attn.split_1.output`) | 🔜 not yet (op paths aren't in `model.modules()`) | +| in-place `[:]=` | ⛔ use explicit `=` (clone semantics, §4) | + +Not-yet-supported features fail **cleanly** (missed-provider error or the per-step timeout), not as a +silent deadlock — the lifecycle (timeout + `finally: cancel()`) is the safety net until each feature lands. + +--- + +## 9. Multi-token iteration — STATUS: DONE (2026-06-07) + +Multi-token iteration (`iter[N]`, `iter[:]`, `next`) isolated == in-process. The iter loop runs in the +worker (sets `mediator.iteration` explicitly per step → correct requester suffixes); the **host** side +gained three pieces: +- **Per-step hook registration:** `ensure_isolated_provider` parses the step `N` from the requester and + passes `iteration=N` to `output_hook`/`input_hook` (new optional param), so the hook fires on step N — + not the host mediator's iteration. +- **Host iter-hooks:** `spawn_isolated_worker` calls `register_iter_hooks(host_mediator, real_model)` so + the host `iteration_tracker` advances per forward. +- **Live host→worker piggyback:** `default_all` (= `generate(max_new_tokens)`) is set *after* the worker + spawns, so it's piggybacked on each response frame (`CudaIpcHostChannel.meta_provider` → + `CudaIpcWorkerChannel.on_meta`) and applied to the worker's interleaver stub before its `iter[:]` loop + computes its bound. This is a general host→worker state channel reused later for the variable store. + +**Verified (`test_isolated_multitoken_iter.py`):** `iter[N]` for N∈{0,1,2} `max|Δ|=0`; per-step swap +propagates (`iter[1]`); `iter[:]` accumulating into `nnsight.save(hs)` bit-identical (3 steps). +**Regression:** channel echo PASS (3-tuple response format); core in-process suite (incl. +`test_iter_edge_cases`) 113 passed; isolated single-pass trace still `max|Δ|=0`. + +**Findings:** `default_all` is live host state set post-spawn → can't snapshot at spawn (hence the +piggyback). The worker over-requests one step past the end only if `default_all` is missing (now fixed). +In-place `[:]=` per step still requires explicit `=` (§4). + +--- + +## 10. Cross-invoke (barrier + variable sharing) — STATUS: DONE (2026-06-07) + +Both isolated == in-process. + +- **Barrier (host-side counting):** each invoke runs in its own worker with its own `Barrier` copy, so it + can't count cross-invoke. `Barrier.__call__` (isolated) sends the TARGET count; the host accumulates + participant names in `Interleaver._barrier_acc` and, once all arrive, runs the existing coordination + loop (`handle_barrier_event` iterates the host mediators, `respond()`+`handle()` each over its own + channel). `Mediator._isolated_worker` (set in `_worker_main`) gates the worker-side behavior. +- **Variable sharing (host store):** worker frames aren't shared across processes, so each worker pushes + its *data* locals to `Interleaver._xinvoke_store` (a 4th `push` field on the event frame) and pulls the + merged store back (piggybacked on the response, reusing the multi-token channel). Only **transmittable + data** (tensors + basic scalars/containers, `_transmittable`) crosses — framework objects + (`Barrier`/`Envoy`, which hold the model) are skipped (the worker already has them via its closure). + Tensors travel **CPU-serialized** (D2H on push, H2D on pull): a worker tensor cloned from the CUDA-IPC + buffer can't be re-shared over IPC by the host ("CUDA tensor received from another process"). + cross_invoker is gated exactly like the host (`len(mediators)>1 and CONFIG.APP.CROSS_INVOKER`). + +**Verified (`test_isolated_cross_invoke.py`):** two-invoke `tracer.barrier(2)` produces matching saved +values; the canonical cross-invoke pattern (invoke A captures `h[3].output`, barrier, invoke B sets +`h[3].output = captured`) is bit-identical. **Non-standard names (`test_nonstd.py`):** gpt2 with +`rename={transformer.h→decoder_blocks, lm_head→output_projection}`, read + iter[1] isolated `max|Δ|=0`. +(Note: `Envoy.__getattr__` resolves the alias to the *real* path before the requester is built, so the +wire string is always the real path — this validates renamed models run end-to-end, but does not itself +exercise an alias↔real mismatch in the host hook registration, since none reaches the wire.) Echo +regression PASS (4-tuple event). + +**Findings:** the variable-sharing push must filter to transmittable data (else the `Barrier` local pulls +in the whole model) AND move tensors to CPU (CUDA-IPC tensors can't be re-shared by the host). The +`_xinvoke_store` is per-interleaver (reset each trace) so no cross-trace leak today; a warm pool must +clear it + the dummy-module hooks + `Globals.saves` per trace. Known coverage gaps: no tests yet for +multi-barrier-in-one-trace, 3+ participants, multi-token + barrier, or variable sharing without a barrier; +the store grows monotonically per trace and ships all shared tensors CPU-serialized on every response (a +perf cliff for large cross-invoke tensors). + +--- + +## 11. Backward + caching — CHARACTERIZED (not yet built) + +`test_isolated_backward_cache_gaps.py` confirms the two gaps and their difficulty: + +- **`tracer.cache()` (a real build, not a quick shim):** `tracer.cache()` runs in the worker → registers + cache hooks on dummy modules → never fire → the `.save()`'d CacheDict comes back **empty**. The fix + needs: (1) a `CACHE` event carrying the spec (module paths + options) so the host registers + `cache_output_hook`/`cache_input_hook` on the *real* modules; (2) **post-forward injection** — the + catch is that cache hooks are **persistent** and the CacheDict is populated **by the forward, which runs + AFTER the mediator's intervention ends** (the intervention only sets up the cache, then ends), so the + END-time saves transmission runs too early. The populated host cache must be copied into the user + frame's CacheDict **at trace exit** — specifically in `Interleaver.cancel` BEFORE `remove_hooks` (the + host cache hooks survive `mediator.cancel`; they're only dropped in `Interleaver.cancel`'s finally, + after the forward has populated them). Matching each worker CacheDict to its host cache needs a token (a + counter set on the CacheDict + carried in the CACHE event; order-based matching is fragile and the token + attr must survive the CacheDict's pickling). A clean build, but ~4 touch points + its own mechanism. + (Acceptance test `test_isolated_cache.py` written, currently failing as expected.) +- **`with tensor.backward()` (hard — likely a major build or a documented limitation):** the backward + context runs in the worker on **detached clones** (clone-on-receive strips the autograd graph, which is + host-side), so there is no graph to differentiate. Making it work requires the backward pass to run + **host-side** (where the graph is) with **path-based grad providers** (today they are `id(tensor)`, + process-local) — architecturally similar to the forward's host-side hook registration but for the + backward session, comparable in size to the core single-pass seam. Until then, backward under isolation + must fail cleanly / run non-isolated. +- **Robustness bug found + fixed (affects all features):** a *dynamic* `NNsightException` raised inside + the worker fails to pickle across the EXCEPTION event (`PicklingError: Can't pickle + nnsight.NNsightException`). Plain exceptions (e.g. a user `ValueError`) cross fine; nnsight-internal + dynamic exceptions did not. The worker now degrades a non-picklable exception to a plain + `RuntimeError(type name + message)` before the EXCEPTION event. + +--- + +## 12. Opt-in surface +`CONFIG.APP.ISOLATE_MEDIATORS` (flag) + `with nnsight.isolate_mediators(): ...` (context). `Mediator.start` +checks it to select the isolated branch. The channel backend (GPU-IPC default, CPU shared-memory +fallback) is config-selectable. Server deployments set the flag globally. + +--- + +## 13. Risks +- **Per-hook latency** — the accepted cost; ~0.6 ms measured, size-independent (vs CPU pickle 1→285 ms). +- **Codec correctness for nested/tuple values** — needs the `applyn`-aware pack tested against + tuple-output blocks (covered by the nested-tuple acceptance case). +- **Lifecycle/deadlock** — one-event-in-flight must hold over the pipe exactly as over the queue; the + clone-on-receive rule must not be skipped, or held-across-access tensors corrupt silently. +- **Path-only envoy fidelity** — the worker mirror must resolve every path the user writes; built from + the serialized tree, validated by the non-standard-named-model acceptance test. diff --git a/docs/developing/mediator-isolation-harness-plan.md b/docs/developing/mediator-isolation-harness-plan.md new file mode 100644 index 000000000..c89e1a782 --- /dev/null +++ b/docs/developing/mediator-isolation-harness-plan.md @@ -0,0 +1,244 @@ +# Mediator Isolation Harness — Implementation Plan + +**Status:** Plan (pre-implementation) · **Date:** 2026-06-05 +**Builds on:** [mediator-isolation-sandbox.md](mediator-isolation-sandbox.md) (design + threat model), +`prototypes/mediator-sandbox/p2_isolation_poc.py` (the isolation proof-of-concept, passing). +**Goal of this doc:** turn "the jail is safe" into "a real trace runs *through* the jail," in testable +increments, with the exact nnsight seam identified. + +--- + +## 1. The seam (where host and jail divide) + +Three call sites in `src/nnsight/intervention/interleaver.py` define the entire boundary. Everything +above them is user code (→ jail); everything below is the model (→ host). + +| Role | Symbol | What it does | +|---|---|---| +| **Client hook** | `eproperty.__get__` (`:275`) | `value = interleaver.current.request(requester)` — `.output` read | +| **Client hook** | `eproperty.__set__` (`:317`) | `interleaver.current.swap(...)` — `.output` write | +| **Worker launch** | `Mediator.start._worker_target` (`:981`) | `Thread(target=…)` runs `intervention(self, info, …)` — `self` is the Mediator | +| **Worker→host RPC** | `Mediator.send` (`:1313`) | `event_queue.put((event, requester))` → wait `response_queue` | +| **host→worker reply** | `Mediator.respond` (`:1299`) | `response_queue.put(value)` → wait `event_queue` | +| **the two queues** | `Mediator.Value` (`:823`), fields `:878-879` | lock-based one-slot handoff (the transport) | +| **server narrow/swap** | `Batcher.narrow` (`batching.py:198`), `Batcher.swap` (`:226`) | slice/replace the tenant's rows | + +**Key fact:** the compiled intervention fn never calls the model directly. It only ever reaches +`interleaver.current.request/swap/skip` and reads scalar state (`iteration`, `history`, `transform`). +So the jail needs **no weights and no GPU** — only an Envoy *path tree* wired to a client-side mediator +that emits the six events (`VALUE/SWAP/SKIP/BARRIER/END/EXCEPTION`, `Events` `:349`) over a socket. + +--- + +## 2. Target architecture + +``` +HOST (trusted: GPU + weights) JAIL (untrusted: CPU-only, net=none, ro-fs) + model + Envoy + weights deserialize request bytes ← pickle lives here + Interleaver + Batcher ── admission bounds client Envoy *path tree* (no modules) + MediatorProxy (per request) ──────── socket ──────── client Interleaver stub (.current = RemoteMediator) + · runs handle_value_event/handle_swap_event RemoteMediator.send() → socket + against the REAL Batcher, narrowed to the compiled intervention fn = USER CODE + host-recorded bounds (NOT worker-supplied) (torch-CPU on the delivered tensor) + · D2H + serialize VALUE replies + · deserialize + H2D SWAP tensors +``` + +- **`MediatorChannel`** — a new abstraction that replaces the `Mediator.Value` queue pair. Two impls: + `InProcessChannel` (today's lock+slot) and `SocketChannel` (length-prefixed frames over an + `AF_UNIX` socket). The Mediator's `send`/`respond` and the worker loop talk to a `MediatorChannel`, + not to `event_queue`/`response_queue` directly. +- **`RemoteMediator`** (jail side) — implements the subset the intervention fn touches: + `request/swap/skip/end/exception`, `iteration`, `transform`, `history`. `send()` writes a frame and + blocks for the reply. **It exposes no `interleaver`, `batcher`, or sibling refs** — the capability + leaks have nothing to reach. +- **`MediatorProxy`** (host side) — owns the real `Batcher` and the existing `handle_*` logic; per + socket event it narrows/swaps against **host-recorded admission bounds**, so `batch_group`, + `narrow(None)`, `[-1,_]`, and widening become unrepresentable from the jail. + +--- + +## 3. What crosses the wire (the codec) + +| Frame | Direction | Payload | +|---|---|---| +| `VALUE` req | jail→host | `requester` (provider path str) + iteration | +| `VALUE` reply | host→jail | the **narrowed activation as a CPU tensor** (D2H'd, serialized) | +| `SWAP` | jail→host | `requester` + the **modified CPU tensor** (serialized) | +| `SKIP` | jail→host | `requester` + sentinel/value | +| `BARRIER` | jail→host | participating mediator names | +| `END` / `EXCEPTION` | jail→host | none / serialized exception metadata | + +Tensor codec: start with the dumb path (`torch.save`/raw buffer / `safetensors`) over the socket; +a shared-memory ring is a later optimization (out of scope for the harness). Values may be **tuples / +nested structures** (transformer block outputs) — the codec must round-trip `applyn`-style nested +containers, not just bare tensors. + +--- + +## 4. Phased plan (each phase is independently testable) + +### Isolation proof-of-concept ✅ DONE +`prototypes/mediator-sandbox/p2_isolation_poc.py` — the 10 escape gadgets run inside the jail and are +inert; socket-fd-into-`net=none`-jail mechanism verified; real PID isolation confirmed. + +### `MediatorChannel` seam (in-process, zero behavior change) ✅ DONE +Extract the `event_queue`/`response_queue` handoff in `Mediator.send`/`respond`/`start` and the worker +loop behind a `MediatorChannel` interface; ship `InProcessChannel` as the default. +- **Touch:** `interleaver.py` — added `MediatorChannel`/`InProcessChannel`; all 23 queue call sites now + route through `self.channel.*`. +- **Acceptance:** the *entire existing nnsight test suite passes unchanged.* — **met** (203 core tests + green; independent review verdict MATCH). + +### Socket transport for the protocol (process boundary, harness) ✅ DONE +The six-event protocol rides an `AF_UNIX` socket between a forked worker and the host, with the real +`Mediator`/`Batcher`. `SocketHostChannel` (host-local event buffer; only `wait_event` reads the socket) +and `SocketWorkerChannel` (`src/nnsight/intervention/transport.py`) + a length-prefixed pickle codec. +- **Proven** (harness `prototypes/mediator-sandbox/phase2_socket_transport.py`, all green): + codec round-trip of nested tuples/dicts/tensors (the transport itself is **bit-exact**); + **golden equivalence** — same `h[6]×2` intervention local vs over-socket on gpt2 → logits match within + multi-threaded CPU forward noise (`allclose atol=1e-3`; residual ~1e-4 is two separate forward passes, + not the socket — a broken swap is Δ≈30); **real batching** + (`needs_batching=True`) — the worker receives only its `batch_group` row and its swap lands only on + that row (no cross-row leak); **out-of-order `restore_event`** exercised at runtime over the socket. +- **Scope honesty:** this proves *the event protocol + the `provider==requester` routing/restore + real + Batcher narrow/swap over a process boundary*. The requester string is a shared constant in the harness + (the compare at `interleaver.py:1175` is exercised; independent eproperty-side *production* of that + string rides in with the `model.trace()` integration). It does **not** yet route through + `model.trace()` — see the transparent-trace-integration milestone below. + +### Transparent `model.trace()` isolated execution (the trace integration) +Make `model.trace()` itself launch the worker in a separate process and return identical results — the +isolation *execution backend*, separated out because it is the bulk of the real feature. +- **`Mediator.start`** gains an isolated path: fork (the trace-integration milestone) / exec-the-jailed-worker (the jail-the-worker stage) + instead of `Thread`; the host Mediator gets a `SocketHostChannel`, the worker a `SocketWorkerChannel`. +- **Saves transmission (keystone finding):** `.save()` marks ids in `Globals.saves` and the values live + in the *worker's* address space (`tracing/base.py:567`). The worker must **ship saved values back to + the host** (e.g. a final `SAVES` frame at END) and the host must inject them into the trace's return — + otherwise an isolated trace returns empty saves. +- **Lifecycle parity:** `cancel()` must `waitpid`/`kill` the worker process (not drop a `Thread`); + `check_dangling_mediators`, EarlyStop drain, and EXCEPTION/SKIP over the socket must all hold. +- **Acceptance:** `with model.trace(prompt): h[6].output.save()` under an `isolate_mediators()` context + returns **bit-identical** saves to the in-process trace, across read / `+=` write / multi-layer / + multi-token. + +### Jail the worker ✅ DONE +Exec the protocol worker inside a `bwrap` jail (CPU-only, `net=none`, ro-fs allowlist), socket fd passed in. +- **Proven** (`prototypes/mediator-sandbox/phase3_jailed_worker.py` + `phase3_jail_transport.py`, green): + (A) the real Mediator protocol worker, run *inside the jail*, produces matching gpt2 logits + (`allclose atol=1e-3`; the transport codec is bit-exact, residual is forward nondeterminism); + (B) the same jailed worker's escape gadgets are inert — `fs_read`→FileNotFoundError, + `net_egress`→OSError, `os.system` touch creates **no host file**, secret **not** leaked — *while its + legitimate request/swap protocol still completes bit-identically*. Host-level sibling = the host-level sibling jail (sandbox design §6.5). +- **Note:** the jailed worker runs a fixed script (not yet a host-shipped compiled intervention) and the + requester is still a constant — shipping the real intervention is the GPU-path deserialize-in-jail work. + +### Batcher as authority (cross-tenant / Boundary B) ✅ DONE +The narrow/swap bounds live on the **host** Mediator (the proxy); the jail's Mediator is a separate object +with no `interleaver`/`batcher` reference, so the worker physically cannot mutate the host's bounds or +reach a sibling. The 8 capability leaks (#1-6, #8; #7 pacing is out of scope) are inert **by construction**, +not by a signature change to `narrow`/`swap`. +- **Proven** (`prototypes/mediator-sandbox/phase4_malicious_worker.py` + `phase4_cross_tenant.py`, green): + a malicious jailed tenant co-batched with a victim row attempts `batch_group=None`/`[-1,0]`/widen and + walks to `interleaver.batcher`/`interleaver.mediators`. Result: all batch_group claims are + "set-on-local-mediator-only" (no host effect); `has_interleaver=False`; both walks → `AttributeError`; + the worker received **only its own row** (sum=12, not 96 → never saw the victim's secret); its poison + landed **only on its own row** — the victim row stayed uncorrupted. +- **Note:** demonstrated with one malicious jail + a passive victim row (the property is symmetric for two + jails); BARRIER cross-jail coordination remains deferred (§5). + +### Real-model GPU path + measurement ✅ DONE +Real `D2H → socket → torch-CPU → H2D` round-trip with the model on an A100, worker CPU-only. +- **Proven** (`prototypes/mediator-sandbox/phase5_gpu_measure.py`, green): (A) the GPU D2H/H2D socket + path matches the in-process GPU forward (`max|Δ|=0`, atol 1e-2); (B) **measured per-hook overhead** + (TOTAL = D2H + sockRTT + H2D) on an A100 80GB PCIe: 1×16×768 (0.03 MB) **0.9 ms**; 1×512×4096 (4.2 MB) + **30 ms**; 1×2048×4096 (16.8 MB, ~7B) **111 ms**; 1×2048×8192 (33.6 MB, ~70B) **376 ms** — vs a + **measured** in-process baseline of **0.1–0.25 ms** (zero-copy reference). +- **Finding (the two-tier-decision input):** the cost is **dominated by pickle + socket RTT** (103 ms of the + 111 ms at 16.8 MB; pickle alone ~14 ms), *not* the PCIe D2H/H2D (~3.5 ms each). A 32-layer 7B cache ⇒ + ~3.5 s added. So the naive path is fine for small/few-layer interventions, but big-model per-layer work + needs **the two-tier approach** (run the common tensor algebra host-side) and/or a **shared-memory ring** wire + format — confirming §7. +- **Deferred:** deserialize the *actual* serialized intervention request in the jail (untrusted + unpickle-in-jail). The security property is already shown (the isolation-PoC + jail stages: pickle `__reduce__` gadget inert + in the jail); shipping the *real compiled intervention* is the trace-integration work. + +### Transport breakdown: the cost is the CODEC, not the boundary ✅ DONE +The Phase-5 "sockRTT-dominated" number was really **pickle-dominated**. `prototypes/mediator-sandbox/ +phase5b_transport_breakdown.py` isolates it (16.8 MB, CPU↔CPU, echo = pure transport, no user op): +- The socket moves 16.8 MB in **~10 ms**; the naive path pickles the tensor **4× per hook** + (`dumps` ~14 ms + `loads` ~8 ms each way), and pickle scales **superlinearly** (285 ms at 33.6 MB). +- **Measured per-hook round-trip by wire codec:** + + | codec | 4.2 MB | 16.8 MB (~7B) | 33.6 MB (~70B) | speedup | + |---|---|---|---|---| + | **A pickle** (today, `transport.py`) | 12.5 ms | 95.9 ms | 285 ms | 1× | + | **B raw message** (struct header + raw bytes via `memoryview`/`torch.frombuffer`; = `safetensors`) | 4.3 ms | 24.1 ms | 123 ms | **~4×** | + | **C shared memory** (`mmap`, only a 1-byte signal on the socket; bulk never travels) | 8.5 ms | 11.5 ms | 28.4 ms | **~8×** (linear) | + +- For **read-only** interventions (caching, logit lens — the common case) shmem is even cheaper: the host + writes the activation once, the worker reads it zero-copy, only a small result returns. +- Separate finding: the worker's `×2` on **bf16 CPU is ~9 ms** by itself (the *user* op, not transport) — + another argument for **the two-tier approach** (keep common ops on the GPU host-side; neither transport nor the + slow CPU recompute happens). +- **Fix hierarchy:** (1) swap the codec pickle→`safetensors`/raw in `transport.py` — small, localized, 4×; + (2) shared-memory ring for big-model per-layer work — 8×, linear; (3) **the two-tier approach** removes transport + entirely for the common tensor algebra. The 111 ms was a codec artifact, **not** intrinsic to isolation. + +### Shared-memory + safetensors transport ✅ DONE (built into `transport.py`) +`ShmArena` (anonymous `memfd`) + `ShmSocketHostChannel`/`ShmSocketWorkerChannel`: tensor bulk rides the +shared memfd, only a small control frame crosses the socket, tensors encoded with **safetensors** (safe, +no-code-exec — also closes the jail→host untrusted-deserialize hole for the bulk). +- **Proven** (`phase6_shm_safetensors.py` + `phase6_jailed_worker.py`, green): correct through the real + `Mediator` protocol both **forked** and **jailed** (gpt2 golden `max|Δ|=0`; the memfd is passed into the + bwrap jail via `pass_fds`+`SHM_FD`). Echo round-trip vs pickle: 16.8 MB **87→44 ms (~2×)**, 4.2 MB ~1.8×. +- **Why only ~2× (not the 8× of hand-rolled shmem):** `safetensors` copies on `save` AND `load` (it returns + owned tensors), so the bulk is memcpy'd ~4×. The hand-rolled `frombuffer` path was zero-copy (a view into + the arena) but that aliases the shared buffer — needs double-buffering / careful lifecycle. Safe-and-simple + (copy) vs fastest (zero-copy view) is the next dial; both beat pickle. +- **Still pickle:** the small control skeleton (non-tensor structure + requester string). Hardening that to a + restricted decoder is the remaining jail→host security item. +- **Bigger lever (open):** true zero-copy needs the data to not move at all — either **the two-tier approach** + (op runs on the GPU host-side; nothing transported) or **GPU-in-jail + CUDA-IPC** (jail maps the host's GPU + buffer; no D2H/serialize) — the latter trades the CPU-only security posture; see the GPU-access discussion. + +--- + +## 5. Open design decisions (resolve as they're reached) + +- **Client Envoy tree in the jail** — ship a path-only Envoy, or reconstruct from the deserialized + request (the nnsight request already carries Envoy refs)? Lean: reconstruct from the request, since + deserialize-in-jail produces it anyway. +- **`cross_invoker` var-sharing** (`send` push/pull, `:1327`/`:1410`) — today two invokes share a Python + frame in-process. Across jails that frame doesn't exist. v1: **one jail per request, `cross_invoker` + disabled**; revisit a host-mediated var store later. +- **`BARRIER`** coordinates multiple mediators (= multiple jails); the host-side `MediatorProxy` set + becomes the coordinator. Defer multi-jail barriers past the cross-tenant stage. +- **CUDA stream propagation** (`start` captures the caller stream, `:967-977`) — irrelevant in a + CPU-only jail; the host keeps all stream handling. +- **Tensor wire format** — raw/`safetensors` first; shmem ring is a perf follow-up, not correctness. + +--- + +## 6. Risks + +- **Per-hook latency** — the accepted cost; quantified in the GPU-path measurement, drives the two-tier approach. +- **Codec correctness** for nested/tuple values (transformer block outputs) — needs a robust + `applyn`-aware codec, tested against the gotchas (tuple-output handling). +- **Lifecycle parity** — the in-process `cancel()`/worker-join must map cleanly to **SIGKILL the jail**; + verify END/EXCEPTION/EarlyStop all drain correctly across the socket. +- **Deadlock surface** — the one-event-in-flight invariant must hold over the socket exactly as over the + one-slot queue, or the "access in forward-pass order" contract breaks. + +--- + +## 7. Test strategy (every phase) + +1. **Golden equivalence** — in-process result == isolated result, on varied models/ops (not just gpt2; + include a non-standard module-named model per the testing rules). +2. **Security** — the jail-the-worker stage runs the 10 escape gadgets as the intervention (inert); the cross-tenant stage runs the 8 + capability-leak payloads across two jails (blocked). +3. **Lifecycle** — exception propagation, early-stop, timeout→SIGKILL, normal drain. + +The two existing ndif suites (`tests/security/test_protector_escapes.py`, +`test_mediator_capability.py`) are the security gates; each fix flips a row from 🔓 to blocked. diff --git a/docs/developing/mediator-isolation-sandbox.md b/docs/developing/mediator-isolation-sandbox.md new file mode 100644 index 000000000..dd7f0ee35 --- /dev/null +++ b/docs/developing/mediator-isolation-sandbox.md @@ -0,0 +1,315 @@ +# Mediator Isolation Sandbox — Design + +**Status:** Draft / design (prototype scope) · **Date:** 2026-06-05 · **Author:** zikai +**Related:** `ndif` security regression suite (`src/services/ray/tests/security/`), `NDIF.md` §6–7, +nnsight `src/nnsight/intervention/interleaver.py` (Mediator / event protocol). + +## 1. Motivation + +NDIF executes **arbitrary user-submitted Python** on shared GPU infrastructure. Today that code runs +**in-process** with the model forward pass inside a single Ray `ModelActor`, behind an in-process +Python whitelist (`Protector`, `src/services/ray/src/ray/nn/security/protected_environment.py`). + +The security regression suite establishes that this boundary is **architecturally unfixable in-process**: + +- **10/10 in-process sandbox escapes succeed** (`().__class__.__mro__…__subclasses__()`, `str.format` + globals walk, `ctypes`/`CDLL`, `torch.utils.cpp_extension.load`, `CodeType`, frame walking, pickle + `__reduce__`, end-to-end `os.system` RCE). The literature consensus (pysandbox post-mortem; Modal, + Lambda, Snowpark practice) is that containment must move to **OS-level isolation**. +- **8/8 multi-tenant Mediator-capability leaks succeed** (mutating `batch_group`, walking + `mediator.interleaver.batcher.current_value`, poking sibling mediators, `swap(None, …)` poisoning all + tenants). All share one root cause: the `Mediator` object is handed to user code with **mutable, + reachable references** to scope and to the shared `Interleaver`/`Batcher`. + +This document designs a refactor that **truly isolates each Mediator** from the model-hosting process, +so user code can no longer reach the host, the GPU, the weights, or other tenants' activations. + +## 2. Core insight: the event protocol is already the boundary + +A `Mediator` (`interleaver.py:718`) already communicates with the forward pass **only** through six +typed messages over two one-slot queues (`interleaver.py:338`, `:949`): + +| Event | From | Means | +|---|---|---| +| `VALUE` | worker | "give me the value at this provider path" (worker blocks) | +| `SWAP` | worker | "replace the value at this provider path with this tensor" | +| `SKIP` | worker | "bypass module X's real forward" | +| `BARRIER` | worker | "synchronize me with these other mediators" | +| `END` | worker | "I'm finished — drain me" | +| `EXCEPTION` | worker | "I crashed; here's the exception" | + +This is already a clean **client/server RPC**. It is insecure only because client and server share one +Python interpreter, one address space, and one GPU context — so the worker can ignore the protocol and +walk Python objects directly. + +> **Isolation = keep the six messages, move the worker to the far side of a real OS boundary, and change +> the transport from an in-process queue to a socket.** Every in-process escape and every +> Mediator-capability leak then disappears *structurally* — there is no longer a `mediator` object, a +> `batcher`, or a host module in the user's address space to reach. + +## 3. Threat model (scoped) + +### In scope +- **A. Host protection** — user code must not reach the host OS, filesystem, network, model weights, or GPU. +- **B. Cross-tenant protection** — when tenants are co-batched in one forward pass (lands with + `ndif-batch`/`ndif-vllm`), tenant X must not read or poison tenant Y's rows. Solved two ways: (1) the + worker holds *no host references*, and (2) the host-side `Batcher` becomes the **authority** on row + bounds (below). +- **Deserialization RCE** — the pickle `__reduce__` gadget fires during `request.deserialize(...)` in + `pre()` (`NDIF.md:838`), *before* execution, under a whitelist that necessarily includes + `pickle`/`cloudpickle`. Untrusted bytes must therefore be unpickled **inside the jail**, never on the host. + +### Out of scope (deliberately deferred) +- **Mediator pacing / forward-pass-blocking DoS** (suite leak #7). This is a *scheduler* property, not an + isolation property. The host may keep blocking on the worker's reply exactly as today — just over a + socket. Revisit with the non-blocking-listener work, separately. +- **Side channels** (timing, cache, GPU memory residency). +- **Per-hook transport throughput.** Accepted as a cost for the prototype; see §7 (option D) for the + optimization path. + +### Trust boundary +- **Trusted:** model weights, GPU + CUDA context, forward pass, `Interleaver`, `Batcher`, host OS, other + tenants' data. +- **Untrusted:** the serialized request bytes (pickle), the compiled intervention function, anything the + user can influence. + +## 4. Design + +### 4.1 Process / thread model + +``` +┌─ ModelActor (TRUSTED: GPU + weights) ───────────────────────┐ +│ main thread: forward pass │ +│ Interleaver + Batcher ← server-side row-bound authority │ +│ per request: MediatorProxy (host stub, speaks 6 events) │ +│ warm pool of CPU-only jailed workers (torch+nnsight ready) │ +└───────────────┬─────────────────────────────────────────────┘ + │ per-request bidirectional socket (fd passed into jail) +┌───────────────┴─ jailed worker (UNTRUSTED, CPU-only) ───────┐ +│ 1. deserialize raw request bytes ← untrusted pickle HERE │ +│ 2. run compiled intervention fn → Mediator worker │ +│ 3. VALUE/SWAP/SKIP/BARRIER/END/EXCEPTION over the socket │ +│ net=none · ro-fs allowlist · non-root · mem cap · no GPU │ +└──────────────────────────────────────────────────────────────┘ +``` + +- The host keeps the forward pass, `Interleaver`, and `Batcher`. Per active request it spins a + **`MediatorProxy`** — what is today the in-process worker thread, demoted to a transport endpoint. It + receives the worker's events and answers them via `Batcher.narrow`/`swap`, **bounded to that tenant's + admission-time rows.** +- **The protocol is unchanged.** Only the transport relocates: `Mediator.start` launches a jailed + process instead of a `Thread`; the `Mediator.Value` one-slot queues (`interleaver.py:767`) become a + socket. `Mediator.handle` writes a value to the socket instead of `response_queue`; the worker's + `request(...)` reads/writes the socket instead of the in-process queues. `iter`, `barrier`, `next`, + and save-filtering keep working with no user-visible change. + +### 4.2 The activation data path (the accepted cost) + +`VALUE`: host `narrow`s to the tenant's rows → `D2H` copy → serialize CPU tensor → socket → jail +deserializes → user gets a **real CPU tensor** and runs arbitrary torch-CPU on it. `SWAP`: jail sends +the modified CPU tensor back → host deserializes → `H2D` → writes into the **bounded** batch slice. + +Start with the simplest serializer (raw buffer / `safetensors` over the socket); a shared-memory ring is +a later optimization, not part of correctness. This is **the per-hook D2H/H2D transport** — clean isolation, transport cost +accepted. See §7 for why and when to move to the two-tier approach. + +### 4.3 Batcher as server-side authority (Boundary B) + +Row bounds are recorded on the **host** at admission. The worker's `VALUE` carries only a provider path + +iteration — never a `batch_group` it can forge. `narrow(None)`, `[-1, _]`, and widening +(`batching.py:213–217`) become **unrepresentable** because the bounds live on the trusted side. This is +the one piece of genuinely new logic; everything else is transport relocation. + +### 4.4 CPU-only, and why + +The jail gets **no GPU**. This is load-bearing for both security and the warm pool: + +- No CUDA context in the jail → no GPU-driver attack surface, and no CUDA-IPC cross-tenant leak (CUDA IPC + shares at *allocation* granularity, which under vLLM's paged allocator can span other tenants' KV cache). +- **torch-CPU is fork-safe**, so a *zygote* that has already imported torch+nnsight can fork a fresh + child per request in ~ms. (A CUDA-initialized process is **not** fork-safe — GPU-in-jail would kill the + pool and is rejected.) + +### 4.5 Lifecycle, pool, cancellation + +- **One jail per request**, assigned to exactly one tenant, then **destroyed / re-exec'd** — never + soft-reused across tenants (a dirtied interpreter is a cross-tenant channel). +- **Warm pool:** first cut may spawn-per-request; then add a zygote pool (pre-imported, pre-jailed idle + workers claimed-then-destroyed) for sub-100ms starts. +- **Cancellation / timeout** becomes **`SIGKILL` the subprocess** instead of the current + `kill_thread()` `ctypes`-injected `SystemExit` (`NDIF.md:919`) — a free robustness win. + +### 4.6 Minimal surgery footprint + +Three touch points, which is what makes this prototype-friendly: +1. `Mediator.start` + the `Value` queue handoff → jailed process + socket. +2. `pre()` deserialization moves inside the jail. +3. `Batcher.narrow`/`swap` take host-held bounds instead of trusting the message. + +The forward pass, interleaver scheduling, and nnsight's public API are untouched. + +## 5. Isolation mechanism + +Three tiers, distinguished by **what enforces the boundary**: + +| Tier | Mechanism | Tools | Host-kernel surface | +|---|---|---|---| +| 1 | Linux namespaces + seccomp + cgroups (shared kernel) | bubblewrap, nsjail, runc/Docker | full (filtered) | +| 2 | User-space kernel (syscalls reimplemented) | gVisor (`runsc`) | minimal | +| 3 | microVM (own guest kernel, hardware virt) | Firecracker | none | + +**Choice for the prototype: Tier 1, CPU-only.** Compute-bound torch-CPU runs ~native; the socket +transport is light; one jail per request is cheap. gVisor (Tier 2) is the **pre-planned upgrade** behind +the same launcher interface if "shared host kernel" becomes unacceptable. Firecracker is overkill for a +CPU-only prototype. + +**Tool: nsjail strategically; bubblewrap as the bootstrap.** The probe (§6) shows: +- `nsjail` is **not packaged** on either the host (Ubuntu 20.04) or the container (Debian trixie) — it is + a source build. Its value is integrated seccomp-policy (kafel) + cgroups + rlimits + time limits, which + is what a production sandbox wants. +- `bubblewrap` (0.11.0) is a **one-line apt install** in the container and is **verified working** on the + host. It is a launcher (pair it with an external seccomp filter + cgroup), but gives the *identical* + Tier-1 boundary. +- The nesting constraints in §6 apply **identically to both**, so the bwrap-vs-nsjail choice is + independent of — and less important than — the deployment decision in §6.5. + +**Recommendation:** bootstrap on bwrap (verified, trivial install, swappable behind a `Jailer` interface), +target nsjail for the hardened path. + +## 6. Host-reality probe (verified 2026-06-05) + +All results are measured on this machine, not assumed. + +### 6.1 Bare host (Ubuntu 20.04, kernel 5.15.0-139) + +| Check | Result | +|---|---| +| `kernel.unprivileged_userns_clone` | `1` (enabled) | +| `user.max_user_namespaces` | `6185255` | +| `unshare -Urn true` (unpriv userns) | **OK** | +| `bwrap` full jail (userns+net+pid+fresh-proc+ro-fs) | **OK**, pid1=bwrap | +| Network isolation inside jail | **OK** (no DNS, loopback only) | +| ro-fs allowlist (unbound `/disk` invisible) | **OK** | +| `bwrap` present / `nsjail` present | `/usr/bin/bwrap` / **absent** | +| seccomp kernel support | `CONFIG_SECCOMP=y`, filter=y | +| cgroup version | v1 | + +→ **Tier-1 unprivileged jailing is fully viable on the bare host.** + +### 6.2 Inside `dev-ray-1` (local stand-in for the prod model container; Debian trixie, ndif/ndif:latest) + +| Check | Result | +|---|---| +| uid inside container | `0` (root-in-container) | +| Privileged / CapAdd / SecurityOpt | `false` / `[]` / `[]` (default posture) | +| AppArmor | `docker-default` | +| `unshare -Urn` (nested unpriv userns) | **BLOCKED** — "Operation not permitted" | +| torch | `2.9.1+cu128`, cuda available | +| `bwrap` / `nsjail` present | **absent** / **absent** (bwrap = 1-line apt; nsjail = build) | + +The block is **Docker's default seccomp profile** denying `clone(CLONE_NEWUSER)`; the container also +lacks `CAP_SYS_ADMIN`, so classic root-namespaces are unavailable too. Both jail paths are closed by default. + +### 6.3 What it takes to jail *inside* the container + +| Container config | Result | +|---|---| +| default | userns BLOCKED | +| `seccomp=unconfined` | userns OK; `bwrap` fails at "make / slave" (AppArmor) | +| `seccomp=unconfined` + `apparmor=unconfined` | propagation OK; **fresh `/proc` mount fails** in new PID ns | +| `+ cap-add SYS_ADMIN` | still fails fresh `/proc` mount | +| `--privileged` | **full jail (fresh proc) works** | +| `seccomp+apparmor=unconfined`, **bind `/proc`, drop PID-unshare** | **jail builds** (reached execvp) — but **no PID-namespace isolation** | + +**Finding:** nesting a userns mount-namespace jail inside an *unprivileged* container is fragile. Locked +container mounts and the fresh-`/proc` restriction mean a jail with **real PID isolation requires +`--privileged`** (or close), and the non-privileged workaround (bind the container's `/proc`) **leaks PID +visibility** — a cross-tenant info-leak vector if multiple workers share the container's PID space. + +### 6.4 Production target: AWS ECS-on-EC2 (the probe transfers) + +Production runs on **AWS ECS-on-EC2** (CDK `ndif-aws/service_stack.py:275`, `ecs.Ec2TaskDefinition`) — +not Fargate, not raw EC2. The model server (`Ray-Worker`) is an **ECS task (Docker container)** on a +**`g4dn.xlarge`** GPU instance (1× T4, **4 vCPU / 16 GB**; `env/dev.py`), on the **ECS-optimized Amazon +Linux 2023 GPU AMI** (`service_stack.py:220`). The whole GPU is allocated into the container +(`gpu_count`, `:365`); network mode is `awsvpc` (own ENI per task); models live on EFS (`/efs/huggingface`). + +Decisive detail: the task's `linux_parameters` is used **only** for `shared_memory_size` +(`service_stack.py:351-355`) — **no `privileged`, no added caps, no seccomp/apparmor override.** So the +prod model container runs the **Docker default posture, identical to local `dev-ray-1`.** + +→ **The §6.1–6.3 probe transfers to production unchanged:** nested unprivileged userns is blocked in prod +too (docker-default seccomp + no `CAP_SYS_ADMIN`, GPU inside the container). The local finding was +representative, not a local artifact. + +Two AWS specifics change the *decision* (§6.5): +- **ECS has no clean per-task seccomp knob.** ECS `dockerSecurityOptions` covers SELinux/AppArmor labels + only — not seccomp profiles. So the local "`seccomp=unconfined`" relaxation is **not** a per-task option + on ECS; the practical in-task lever is **`privileged=true`** (drops seccomp *and* adds all caps). + The in-container self-jail is therefore *heavier* on AWS than locally. +- **We own the EC2 instance, its user-data, and the host Docker daemon.** That makes the sibling approach + host-level-sibling deployment natural: pre-install the jailer in user-data and run sandboxes as **host-level siblings**, leaving + the model-server task unprivileged. + +*Adjacent finding (out of scope for this design, worth fixing): `ndif-aws/env/dev.py` commits **live +credentials** — an AWS access key/secret, a HuggingFace token, DB/Influx passwords. Rotate and move to SSM +Parameter Store / Secrets Manager.* + +## 6.5 Deployment decision (KEY OPEN ITEM) + +The nesting reality reshapes *where* the jail runs. Two resolutions, read through the AWS reality (§6.4): + +- **Self-jail inside a loosened model container.** Locally this is + `seccomp=unconfined, apparmor=unconfined` (+ effectively `privileged` for real PID isolation). + **On ECS it collapses to `privileged=true`** — there is no clean per-task seccomp knob — set via the + container's `linux_parameters`/props (`service_stack.py:351-363`). Fastest to a testable boundary, **but + on AWS it means a privileged model task**, which broadens that task's own surface more than the local + case did. Justified only as a throwaway first cut to exercise the protocol. +- **Sibling sandbox on the host (recommended target, especially on AWS).** Run the jailed worker as a + **separate process/container at the EC2-instance level**, where host userns + bwrap work (verified on the + local host §6.1; AL2023 host expected to match — needs a one-line on-instance confirm). Pre-install the + jailer in EC2 user-data; the model-server ECS task stays **unprivileged** and talks to the sandbox over a + socket. This is how production code-exec services isolate (sibling sandboxes, not nested jails). Cost: a + host-level launcher/broker and its trust handling (don't hand the host Docker socket to the model server). + +**Recommendation:** **the host-level sibling jail is the production target** — owning the EC2 host makes it clean and avoids a +privileged model task. The in-container self-jail is justified only as a throwaway first cut; on ECS even that costs +`privileged`. Keep the launcher abstracted so in-container → host-level (and basic-seccomp → gVisor) don't touch anything above it. + +**Sizing caveat:** `g4dn.xlarge` is **4 vCPU / 16 GB**, shared by Ray + the model + any CPU-only sandbox +pool. That caps pool size and the per-hook D2H/H2D budget (the per-hook D2H/H2D transport), pushing the two-tier optimization +(§7) *earlier* on small prod instances than a pure-correctness view suggests. + +## 7. Future optimization (out of scope now) + +Where do user tensor ops execute? The prototype picks **per-hook D2H/H2D in the jail, on CPU** (D2H/H2D per hook). +The destination is **the two-tier approach**: run the common tiny tensor algebra (read/write/project/ablate/steer/ +patch/topk/cache) on the host via a *validated op interpreter* (data never leaves the GPU, no jail in the +hot path), and route only genuinely-arbitrary Python to the CPU jail. (B) GPU-in-jail via CUDA IPC is +rejected (surface + breaks the pool); (C) fully symbolic host execution is a breaking change to nnsight's +"real tensors" contract. The two-tier approach attacks the root — *arbitrary code next to the data* — rather than paying to +isolate code that didn't need to be arbitrary. + +On small prod instances this is **less optional than it sounds**: a `g4dn.xlarge` (4 vCPU / 16 GB, §6.4) +leaves little headroom for a CPU-only sandbox pool doing per-hook D2H/H2D, so the two-tier approach likely graduates from +"later optimization" to "needed for throughput" before this ships at scale. + +## 8. Decisions & open questions + +**Decided** +- Move the worker across the existing six-event protocol boundary; change transport only. +- Tier-1 isolation, CPU-only, one jail per request, host-side `Batcher` as row-bound authority, + deserialize-in-jail, `SIGKILL` cancellation. +- Prototype tensor path = per-hook D2H/H2D transport; destination = the two-tier approach — and on small prod instances (`g4dn`, §6.4) the two-tier approach is + needed sooner than "later". +- Bootstrap on bubblewrap, target nsjail, gVisor as the strong-boundary upgrade. +- **Prod runs ECS-on-EC2** (§6.4); the probe transfers (same Docker-default posture). **the host-level sibling jail (host-level + sibling on the EC2 instance) is the production target**; in-container the in-container self-jail is a throwaway first cut and on + ECS costs `privileged=true`. + +**Open** +- **in-container vs host-level sequencing** (§6.5) — the host-level sibling jail is the target; the only question is whether a privileged in-container first + cut is worth building at all, or whether to go straight to the host-level sibling. +- Warm-pool mechanism: zygote-fork vs pre-launched idle-worker pool. +- Tensor wire format: raw buffer vs `safetensors` vs shmem ring (correctness first, optimize later). +- Exact `Jailer` interface so bwrap → nsjail → runsc are drop-in. diff --git a/prototypes/mediator-sandbox/_smoke.py b/prototypes/mediator-sandbox/_smoke.py new file mode 100644 index 000000000..b6e9bb1ee --- /dev/null +++ b/prototypes/mediator-sandbox/_smoke.py @@ -0,0 +1,6 @@ +from nnsight import LanguageModel +m = LanguageModel("gpt2", device_map="cpu", dispatch=True) +with m.trace("Hello world"): + h6 = m.transformer.h[6].output[0].save() + logits = m.lm_head.output.save() +print("h6", tuple(h6.shape), "logits", tuple(logits.shape), "SMOKE_OK") diff --git a/prototypes/mediator-sandbox/gpu_sandbox/fault_test.py b/prototypes/mediator-sandbox/gpu_sandbox/fault_test.py new file mode 100644 index 000000000..869785df0 --- /dev/null +++ b/prototypes/mediator-sandbox/gpu_sandbox/fault_test.py @@ -0,0 +1,76 @@ +#!/usr/bin/env python3 +"""Does a GPU fault in a worker stay contained to that worker, or take the host +down too? Decisive for the separate-contexts vs MPS choice. + +A spawned worker triggers a device-side assert (embedding lookup out of range), +which poisons its CUDA context. We then check whether the HOST's CUDA context +still works. Run once without MPS (separate contexts) and once under MPS. +""" +import queue +import sys + +import torch +import torch.multiprocessing as mp + + +def faulting_worker(sig, mode): + import torch + import torch.nn.functional as F + try: + torch.ones(8, device="cuda").sum() + torch.cuda.synchronize() + sig.put("worker: cuda ok") + except Exception as e: # noqa: BLE001 + sig.put(f"worker: cuda init FAILED {type(e).__name__}") + return + try: + if mode == "illegal": + a = torch.randn(8, device="cuda") + idx = torch.tensor([2 ** 40], device="cuda", dtype=torch.long) # address overflow + _ = a[idx] # -> illegal memory access + else: + _ = F.embedding(torch.tensor([10_000_000], device="cuda"), + torch.randn(16, 8, device="cuda")) # -> device-side assert + torch.cuda.synchronize() + sig.put("worker: NO FAULT (unexpected)") + except Exception as e: # noqa: BLE001 + sig.put(f"worker: FAULTED {type(e).__name__}: {str(e)[:50]}") + sig.put("worker: exiting") + + +def host_cuda_works(): + try: + r = (torch.ones(2048, device="cuda") * 3.0).sum() + torch.cuda.synchronize() + return float(r) == 3.0 * 2048 + except Exception as e: # noqa: BLE001 + return f"DEAD: {type(e).__name__}: {str(e)[:70]}" + + +def main(): + mode = sys.argv[1] if len(sys.argv) > 1 else "assert" + mp.set_start_method("spawn") + before = host_cuda_works() + print(f"[mode={mode}] host CUDA BEFORE worker fault: {before}") + + sig = mp.Queue() + p = mp.Process(target=faulting_worker, args=(sig, mode)) + p.start() + msgs = [] + for _ in range(5): + try: + msgs.append(sig.get(timeout=30)) + except queue.Empty: + break + p.join(timeout=15) + print(f"worker messages: {msgs} | worker exitcode={p.exitcode}") + + after = host_cuda_works() + print(f"host CUDA AFTER worker fault: {after}") + contained = after is True + print(f"VERDICT: host {'SURVIVED — fault CONTAINED to the worker' if contained else 'DIED — fault COUPLED to host'}") + sys.exit(0 if contained else 2) + + +if __name__ == "__main__": + main() diff --git a/prototypes/mediator-sandbox/gpu_sandbox/gpu_sandbox.py b/prototypes/mediator-sandbox/gpu_sandbox/gpu_sandbox.py new file mode 100644 index 000000000..45f69d53f --- /dev/null +++ b/prototypes/mediator-sandbox/gpu_sandbox/gpu_sandbox.py @@ -0,0 +1,73 @@ +"""Host-side manager for the GPU-enabled isolated worker. + +Holds a shared GPU bounce buffer and a spawned, locked-down worker process. +``apply(activation, fn)`` runs the user's ``fn`` on ``activation`` *in the +worker* (zero-copy via the shared buffer) and returns the result — the user's +arbitrary code never runs in the model-server process. + +A real deployment keeps a POOL of these (one per concurrent request); this is a +single worker for clarity. Worker death (a segfault in user code) is detected and +surfaced; the host keeps serving. +""" +import os +import sys + +import cloudpickle +import torch +import torch.multiprocessing as mp + +sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) # so the worker can import sandbox/gpu_worker + + +class GPUSandbox: + def __init__(self, arena_bytes=64 << 20, gpu_mem_fraction=0.3, device="cuda"): + self.device = device + ctx = mp.get_context("spawn") # CUDA requires spawn, not fork + self.buf = torch.empty(arena_bytes, dtype=torch.uint8, device=device) # the bounce buffer + self.parent_conn, child_conn = ctx.Pipe() + self.ready = ctx.Queue() + from gpu_worker import run + self.proc = ctx.Process( + target=run, args=(self.buf, child_conn, self.ready, gpu_mem_fraction), daemon=True + ) + self.proc.start() + assert self.ready.get(timeout=180) == "ready" + + def apply(self, activation: torch.Tensor, fn, timeout=60): + """Run ``fn(activation)`` in the isolated worker; return the result tensor.""" + if not self.proc.is_alive(): + raise RuntimeError("sandbox worker is dead") + a = activation.contiguous() + ab = a.flatten().view(torch.uint8) + self.buf[: ab.numel()].copy_(ab) # D2D copy into the shared buffer + torch.cuda.synchronize() + self.parent_conn.send((cloudpickle.dumps(fn), tuple(a.shape), a.dtype, ab.numel())) + if not self.parent_conn.poll(timeout): + # worker is wedged (e.g. an infinite loop in user code) — kill it; the + # host is unaffected and a pool would respawn a fresh worker. + self.proc.terminate() + raise TimeoutError(f"sandboxed intervention exceeded {timeout}s — worker killed") + try: + reply = self.parent_conn.recv() + except (EOFError, OSError): + # the pipe broke mid-op → the worker crashed (e.g. a segfault in user + # C-code). Contained: only this request dies; the host is fine and a + # pool respawns. Surface it cleanly instead of leaking EOFError. + raise RuntimeError("sandbox worker crashed during the intervention") + if reply[0] == "err": + raise RuntimeError(f"sandboxed intervention raised {reply[1]}: {reply[2]}") + _, shape, dtype, nbytes = reply + return self.buf[:nbytes].view(dtype).view(*shape).clone() + + def alive(self): + return self.proc.is_alive() + + def close(self): + try: + if self.proc.is_alive(): + self.parent_conn.send("stop") + self.proc.join(timeout=5) + except Exception: + pass + if self.proc.is_alive(): + self.proc.terminate() diff --git a/prototypes/mediator-sandbox/gpu_sandbox/gpu_worker.py b/prototypes/mediator-sandbox/gpu_sandbox/gpu_worker.py new file mode 100644 index 000000000..c4d6126de --- /dev/null +++ b/prototypes/mediator-sandbox/gpu_sandbox/gpu_worker.py @@ -0,0 +1,60 @@ +"""The GPU-enabled, footgun-contained worker process (spawn target). + +Lifecycle: + 1. (at spawn) map the shared GPU bounce buffer — a CUDA tensor the host shared + via IPC, so host and worker point at the SAME GPU memory. + 2. warm CUDA + pre-import torch/numpy (so user ops don't need a new import). + 3. lock_down(): seccomp blocks new open/openat (fs) and socket/connect (net); + RLIMIT_AS caps memory. CUDA keeps working (already-open /dev/nvidia* fds). + 4. loop: receive (cloudpickled user fn + tensor metadata) → view the bounce + buffer as that tensor → run the user fn on the real GPU tensor → write the + result back into the buffer → reply with the result's metadata. + +The user fn is arbitrary Python. It runs HERE, after lockdown: a crash is +contained to this process (the host respawns), an open()/socket() fails with +EPERM, a runaway alloc hits the rlimit. The host and its other tenants are safe; +only this request is affected. The GPU is shared (the accepted risk). +""" +import torch + + +def run(shared_buf, conn, ready, gpu_mem_fraction): + import cloudpickle # noqa: F401 + + try: + import numpy # noqa: F401 (pre-import so user ops referencing it don't openat) + except Exception: + pass + from sandbox import lock_down + + # warm CUDA so all kernels/contexts are loaded before we cut off file opens + _ = (torch.randn(128, 128, device="cuda") @ torch.randn(128, 128, device="cuda")).sum() + torch.cuda.synchronize() + # warm cloudpickle's (de)serialization machinery so per-request loads() won't + # trigger a lazy import after the filesystem is locked down + cloudpickle.loads(cloudpickle.dumps(lambda _t: _t * 2.0)) + + # Cap GPU memory so a runaway allocation in user code can't exhaust the device. + # (RLIMIT_AS is unusable here — CUDA reserves tens of GB of *virtual* space.) + torch.cuda.set_per_process_memory_fraction(gpu_mem_fraction) + lock_down() + ready.put("ready") + + while True: + msg = conn.recv() + if msg == "stop": + return + fn_blob, shape, dtype, nbytes = msg + try: + fn = cloudpickle.loads(fn_blob) + t = shared_buf[:nbytes].view(dtype).view(*shape) # zero-copy view of the buffer + out = fn(t) # <-- arbitrary user code on the GPU tensor + if not torch.is_tensor(out): + out = torch.as_tensor(out, device="cuda") + out = out.contiguous() + ob = out.flatten().view(torch.uint8) + shared_buf[: ob.numel()].copy_(ob) # result back into the shared buffer + torch.cuda.synchronize() + conn.send(("ok", tuple(out.shape), out.dtype, ob.numel())) + except Exception as e: # noqa: BLE001 (contain the footgun, report it) + conn.send(("err", type(e).__name__, str(e)[:200])) diff --git a/prototypes/mediator-sandbox/gpu_sandbox/perf_ctxswitch.py b/prototypes/mediator-sandbox/gpu_sandbox/perf_ctxswitch.py new file mode 100644 index 000000000..601b7f681 --- /dev/null +++ b/prototypes/mediator-sandbox/gpu_sandbox/perf_ctxswitch.py @@ -0,0 +1,74 @@ +#!/usr/bin/env python3 +"""Is the ~0.5 ms GPU context-switching between the two processes' CUDA contexts? + +Compare three round-trips between host and a spawned worker: + A. worker touches NO GPU (bare echo) -> pure IPC + wakeups + B. ONLY the worker does a GPU op per round -> one context active at a time-ish + C. BOTH host and worker do a GPU op per round -> GPU must switch contexts host<->worker + +If C >> A, the cost is GPU context switching, not process communication. +""" +import os +import sys +import time + +sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) +import torch +import torch.multiprocessing as mp + + +def worker(conn, touch_gpu): + if touch_gpu: + x = torch.randn(4096, device="cuda") + torch.cuda.synchronize() + while True: + m = conn.recv() + if m == "stop": + return + if touch_gpu: + x.add_(1.0) + torch.cuda.synchronize() + conn.send(b"k") + + +def t_ms(fn, n=200, warm=30): + for _ in range(warm): + fn() + t0 = time.perf_counter() + for _ in range(n): + fn() + return (time.perf_counter() - t0) / n * 1e3 + + +def run(touch_worker, touch_host): + ctx = mp.get_context("spawn") + pc, cc = ctx.Pipe() + p = ctx.Process(target=worker, args=(cc, touch_worker), daemon=True) + p.start() + y = torch.randn(4096, device="cuda") + torch.cuda.synchronize() + + def one(): + if touch_host: + y.add_(1.0) + torch.cuda.synchronize() + pc.send(b"x") + pc.recv() + + dt = t_ms(one) + pc.send("stop"); p.join() + return dt + + +def main(): + a = run(touch_worker=False, touch_host=False) + b = run(touch_worker=True, touch_host=False) + c = run(touch_worker=True, touch_host=True) + print(f"[ctx] A. worker no-GPU (pure IPC): {a:.3f} ms") + print(f"[ctx] B. worker GPU op only: {b:.3f} ms") + print(f"[ctx] C. BOTH host+worker GPU op (alternate):{c:.3f} ms <- this is apply()'s pattern") + print(f"[ctx] GPU context-switch overhead ~= C - A = {c - a:.3f} ms") + + +if __name__ == "__main__": + main() diff --git a/prototypes/mediator-sandbox/gpu_sandbox/perf_decompose.py b/prototypes/mediator-sandbox/gpu_sandbox/perf_decompose.py new file mode 100644 index 000000000..e213bcdd3 --- /dev/null +++ b/prototypes/mediator-sandbox/gpu_sandbox/perf_decompose.py @@ -0,0 +1,67 @@ +#!/usr/bin/env python3 +"""Where does the ~0.6 ms/hook go? Decompose it into: + - bare process round-trip (mp.Pipe send/recv = IPC + 2 scheduler wakeups), no torch/cloudpickle + - cloudpickle.loads (the worker rebuilds the fn every call) + - torch.cuda.synchronize (apply does 2: host after copy-in, worker after op) +so we can say whether it's process-communication cost or something else. +""" +import os +import sys +import time + +sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) +import cloudpickle +import torch +import torch.multiprocessing as mp + + +def steer(t): + return t + 1.0 + + +def echo_worker(conn): + while True: + m = conn.recv() + if m == "stop": + return + conn.send(b"k") # smallest possible reply — pure IPC round-trip + + +def t_ms(fn, n=200, warm=20): + for _ in range(warm): + fn() + t0 = time.perf_counter() + for _ in range(n): + fn() + return (time.perf_counter() - t0) / n * 1e3 + + +def main(): + ctx = mp.get_context("spawn") + pc, cc = ctx.Pipe() + p = ctx.Process(target=echo_worker, args=(cc,), daemon=True) + p.start() + + # 1. bare process round-trip: send a tiny msg, worker echoes. This is IPC + + # two scheduler wakeups (host→worker, worker→host) with nothing else. + pingpong = t_ms(lambda: (pc.send(b"x"), pc.recv())) + + # 2. cloudpickle.loads cost (the worker does this every call) + blob = cloudpickle.dumps(steer) + loads = t_ms(lambda: cloudpickle.loads(blob)) + + # 3. a single torch.cuda.synchronize after a trivial op + a = torch.randn(1024, device="cuda") + sync = t_ms(lambda: (a.add_(1.0), torch.cuda.synchronize())) + + pc.send("stop"); p.join() + + print(f"[decompose] bare mp.Pipe round-trip (IPC + 2 wakeups): {pingpong:.3f} ms") + print(f"[decompose] cloudpickle.loads(fn) per call: {loads:.3f} ms") + print(f"[decompose] one cuda.synchronize (+trivial op): {sync:.3f} ms") + print(f"[decompose] => apply() ~= pingpong + loads + ~2*sync " + f"= {pingpong + loads + 2 * sync:.3f} ms (measured apply ~0.6-0.7)") + + +if __name__ == "__main__": + main() diff --git a/prototypes/mediator-sandbox/gpu_sandbox/perf_test.py b/prototypes/mediator-sandbox/gpu_sandbox/perf_test.py new file mode 100644 index 000000000..d7e9b8126 --- /dev/null +++ b/prototypes/mediator-sandbox/gpu_sandbox/perf_test.py @@ -0,0 +1,93 @@ +#!/usr/bin/env python3 +"""Performance of the GPU sandbox: the REAL apply() path, not the microbench. + +Measures (1) per-hook apply() latency vs activation size, with a component +breakdown and the in-process baseline; (2) the end-to-end impact on a real gpt2 +trace, in-process vs 1 and N sandboxed interventions. + +Run: CUDA_VISIBLE_DEVICES=6 PYTHONPATH=/src .../bin/python perf_test.py +""" +import os +import sys +import time + +sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) +import cloudpickle +import torch + +from nnsight import LanguageModel +from gpu_sandbox import GPUSandbox + + +def steer(t): + return t + 1.0 + + +def t_ms(fn, n=50, warm=10): + for _ in range(warm): + fn() + torch.cuda.synchronize() + t0 = time.perf_counter() + for _ in range(n): + fn() + torch.cuda.synchronize() + return (time.perf_counter() - t0) / n * 1e3 + + +def main(): + sb = GPUSandbox() + + # ---- 1. per-hook apply() latency vs size + breakdown ---- + print(f"{'activation':>20} {'MB':>7} {'in-proc':>9} {'apply()':>9} {'(dumps)':>9} {'(copyD2D)':>10} " + f"{'(rtt)':>8} (ms)") + for (b, s, h) in [(1, 16, 768), (1, 512, 768), (1, 512, 4096), (1, 2048, 4096), (1, 2048, 8192)]: + act = torch.randn(b, s, h, device="cuda", dtype=torch.bfloat16) + mb = act.element_size() * act.nelement() / 1e6 + inproc = t_ms(lambda: steer(act)) + full = t_ms(lambda: sb.apply(act, steer)) + # component breakdown + dumps = t_ms(lambda: cloudpickle.dumps(steer)) + ab = act.contiguous().flatten().view(torch.uint8) + copyd2d = t_ms(lambda: (sb.buf[: ab.numel()].copy_(ab))) + rtt = max(full - dumps - copyd2d, 0.0) # remainder ≈ pipe + worker wakeup + op + readback + print(f"{str((b, s, h)):>20} {mb:7.1f} {inproc:9.3f} {full:9.3f} {dumps:9.3f} {copyd2d:10.3f} {rtt:8.3f}") + + # ---- 2. end-to-end real gpt2 trace ---- + print() + model = LanguageModel("gpt2", device_map="cuda", dispatch=True) + inputs = model.tokenizer("The Eiffel Tower is in the city of", return_tensors="pt").to("cuda") + + def trace_inproc(): + with model.trace(inputs): + model.transformer.h[6].output = steer(model.transformer.h[6].output) + model.lm_head.output.save() + + def trace_sandbox1(): + with model.trace(inputs): + model.transformer.h[6].output = sb.apply(model.transformer.h[6].output, steer) + model.lm_head.output.save() + + def trace_sandboxN(): + with model.trace(inputs): + for L in range(12): # intervene at every layer (e.g. a full-model cache/steer) + model.transformer.h[L].output = sb.apply(model.transformer.h[L].output, steer) + model.lm_head.output.save() + + def trace_plain(): + with model.trace(inputs): + model.lm_head.output.save() + + plain = t_ms(trace_plain, n=20, warm=5) + inp = t_ms(trace_inproc, n=20, warm=5) + s1 = t_ms(trace_sandbox1, n=20, warm=5) + sN = t_ms(trace_sandboxN, n=20, warm=5) + print(f"[gpt2 trace] plain (no intervention): {plain:7.2f} ms") + print(f"[gpt2 trace] in-process steer @ 1 layer: {inp:7.2f} ms") + print(f"[gpt2 trace] sandboxed steer @ 1 layer: {s1:7.2f} ms (+{s1 - inp:.2f} ms / hook)") + print(f"[gpt2 trace] sandboxed steer @ 12 layers: {sN:7.2f} ms (+{(sN - inp) / 12:.2f} ms / hook avg)") + + sb.close() + + +if __name__ == "__main__": + main() diff --git a/prototypes/mediator-sandbox/gpu_sandbox/probe_bootstrap.py b/prototypes/mediator-sandbox/gpu_sandbox/probe_bootstrap.py new file mode 100644 index 000000000..83ecf65d8 --- /dev/null +++ b/prototypes/mediator-sandbox/gpu_sandbox/probe_bootstrap.py @@ -0,0 +1,133 @@ +#!/usr/bin/env python3 +"""SP1 bootstrap probe — can a per-mediator intervention be serialized, then +deserialized against DUMMY modules + a stub interleaver, and run far enough to +call request() with the correct requester string? + +This de-risks the worker bootstrap (the hardest unknown) WITHOUT the channel, +forward, or saves. If this prints the expected requester, the rest is plumbing +we already have (CudaIpcChannel + Shim A/B). + +Run: + CUDA_VISIBLE_DEVICES=0 PYTHONPATH=src \ + /disk/u/zikai/anaconda3/envs/hf-serve/bin/python \ + prototypes/mediator-sandbox/gpu_sandbox/probe_bootstrap.py +""" +import sys +import types + +import torch +import torch.nn as nn + +from nnsight import LanguageModel +from nnsight.intervention import serialization +from nnsight.intervention.interleaver import Mediator +from nnsight.intervention.tracing.globals import _ensure_mounted + + +def main(): + model = LanguageModel("gpt2", device_map="cpu", dispatch=True) + + # 1. Capture a real mediator's serialized bytes during a normal trace. + captured = {} + orig_start = Mediator.start + + def patched_start(self, interleaver): + if "bytes" not in captured: + try: + # The tracer attaches source during ITS __getstate__ (tracer.py:677); + # per-mediator serialization must do the same first. + self.intervention.__source__ = "".join(self.info.source) + captured["bytes"] = serialization.dumps(self) + captured["err"] = None + except Exception as e: # noqa: BLE001 + captured["err"] = f"{type(e).__name__}: {e}" + return orig_start(self, interleaver) + + Mediator.start = patched_start + try: + with model.trace("The Eiffel Tower is in"): + model.transformer.h[6].output.save() + finally: + Mediator.start = orig_start + + if captured.get("err"): + print(f"[serialize] FAILED: {captured['err']}") + sys.exit(1) + print(f"[serialize] mediator -> {len(captured['bytes'])} bytes OK") + + # 2. Build a WORKER-style persistent map: dummy modules for every Module:, + # a stub interleaver, real tokenizer/processor pass-through. + real_map = model._remoteable_persistent_objects() + + calls = [] + + class StubBatcher: + current_provider = None + current_value = None + + class StubInterleaver: + interleaving = True + batcher = StubBatcher() + current = None + + def iterate_requester(self, requester): + med = self.current + iteration = med.iteration if med.iteration is not None else med.iteration_tracker[requester] + return f"{requester}.i{iteration}" + + stub_interleaver = StubInterleaver() + + pmap = {} + for k, v in real_map.items(): + if k.startswith("Module:"): + dummy = nn.Module() + dummy.__path__ = k[len("Module:") :] + pmap[k] = dummy + elif k == "Interleaver": + pmap[k] = stub_interleaver + else: + pmap[k] = v # Tokenizer / Processor — keep real (lightweight) + + # 3. Deserialize the mediator against dummies. + try: + med = serialization.loads(captured["bytes"], pmap) + except Exception as e: # noqa: BLE001 + print(f"[deserialize] FAILED: {type(e).__name__}: {e}") + sys.exit(1) + print(f"[deserialize] mediator rebuilt against dummy modules OK (name={med.name})") + + # 4. Wire it into the stub interleaver and run the intervention with a + # recording request() — assert it asks for the right provider. + _ensure_mounted() + med.idx = 0 + med.interleaver = stub_interleaver + stub_interleaver.current = med + + def recording_request(requester): + calls.append(requester) + return torch.zeros(1, 16, 768) # gpt2-ish block output stand-in + + med.request = recording_request + med.swap = lambda requester, value: calls.append(("swap", requester)) + med.end = lambda: None + med.push = lambda: None + med.pull = lambda: None + med.cross_invoker = False + + try: + med.intervention(med, med.info) + except Exception as e: # noqa: BLE001 + print(f"[run] intervention raised: {type(e).__name__}: {e}") + print(f"[run] calls so far: {calls}") + sys.exit(1) + + expected = "model.transformer.h.6.output.i0" # real envoy root prefixes "model." + ok = expected in calls + print(f"[run] intervention requested: {calls}") + print("=" * 72) + print(f"BOOTSTRAP PROBE: {'PASS' if ok else 'FAIL'} (expected {expected!r})") + sys.exit(0 if ok else 1) + + +if __name__ == "__main__": + main() diff --git a/prototypes/mediator-sandbox/gpu_sandbox/probe_saves.py b/prototypes/mediator-sandbox/gpu_sandbox/probe_saves.py new file mode 100644 index 000000000..cb6799237 --- /dev/null +++ b/prototypes/mediator-sandbox/gpu_sandbox/probe_saves.py @@ -0,0 +1,101 @@ +#!/usr/bin/env python3 +"""SP1 saves probe — after a deserialized intervention runs to end()/push() in a +worker-like context, WHERE do the .save()'d values land, and can we filter them by +Globals.saves to ship back? (Shim B feasibility.) + +Run: + CUDA_VISIBLE_DEVICES=0 PYTHONPATH=src \ + /disk/u/zikai/anaconda3/envs/hf-serve/bin/python \ + prototypes/mediator-sandbox/gpu_sandbox/probe_saves.py +""" +import sys + +import torch +import torch.nn as nn + +from nnsight import LanguageModel +from nnsight.intervention import serialization +from nnsight.intervention.interleaver import Mediator +from nnsight.intervention.tracing.globals import Globals, _ensure_mounted + + +def main(): + model = LanguageModel("gpt2", device_map="cpu", dispatch=True) + + captured = {} + orig_start = Mediator.start + + def patched_start(self, interleaver): + if "bytes" not in captured: + self.intervention.__source__ = "".join(self.info.source) + captured["bytes"] = serialization.dumps(self) + return orig_start(self, interleaver) + + Mediator.start = patched_start + try: + with model.trace("The Eiffel Tower is in"): + saved_marker = model.transformer.h[6].output.save() + finally: + Mediator.start = orig_start + + real_map = model._remoteable_persistent_objects() + + class StubBatcher: + current_provider = None + current_value = None + + class StubInterleaver: + interleaving = True + batcher = StubBatcher() + current = None + + def iterate_requester(self, requester): + med = self.current + it = med.iteration if med.iteration is not None else med.iteration_tracker[requester] + return f"{requester}.i{it}" + + stub_interleaver = StubInterleaver() + pmap = {} + for k, v in real_map.items(): + if k.startswith("Module:"): + d = nn.Module(); d.__path__ = k[len("Module:"):]; pmap[k] = d + elif k == "Interleaver": + pmap[k] = stub_interleaver + else: + pmap[k] = v + + med = serialization.loads(captured["bytes"], pmap) + _ensure_mounted() + Globals.saves.clear() + med.idx = 0 + med.interleaver = stub_interleaver + stub_interleaver.current = med + med.cross_invoker = False + + sentinel = torch.arange(12, dtype=torch.float32).reshape(1, 3, 4) # the "activation" + med.request = lambda requester: sentinel + + # Let the REAL end()/push() run; stub only the channel. + class FakeChannel: + def put_event(self, item): self.last = item + med.channel = FakeChannel() + + med.intervention(med, med.info) + + # Shim B candidate: filter the worker's frame locals by Globals.saves. + frame = med.info.frame + flocals = getattr(frame, "f_locals", {}) + saved = {k: v for k, v in flocals.items() if id(v) in Globals.saves} + + print(f"[frame] info.frame type={type(frame).__name__}, #locals={len(flocals)}") + print(f"[saves] Globals.saves size={len(Globals.saves)}; filtered saved keys={list(saved.keys())}") + hit = [k for k, v in saved.items() if torch.is_tensor(v) and torch.equal(v, sentinel)] + ok = len(hit) >= 1 + print(f"[match] a saved local equals the sentinel activation: {ok} (keys={hit})") + print("=" * 72) + print(f"SAVES PROBE: {'PASS' if ok else 'FAIL'}") + sys.exit(0 if ok else 1) + + +if __name__ == "__main__": + main() diff --git a/prototypes/mediator-sandbox/gpu_sandbox/sandbox.py b/prototypes/mediator-sandbox/gpu_sandbox/sandbox.py new file mode 100644 index 000000000..0ff34cbf9 --- /dev/null +++ b/prototypes/mediator-sandbox/gpu_sandbox/sandbox.py @@ -0,0 +1,76 @@ +"""Footgun-containment sandbox for a GPU-enabled worker process. + +Goal (per the agreed threat model): contain *mistakes*, not a determined +adversary. After CUDA + torch are fully initialised, the worker calls +``lock_down()`` which installs a minimal seccomp-BPF filter that makes new +``open``/``openat`` (filesystem) and ``socket``/``connect`` (network) syscalls +fail with EPERM. CUDA keeps working because it talks to the already-open +``/dev/nvidia*`` fds via ioctl/mmap, not by opening new files. + +No external deps: the BPF program is assembled by hand and installed via +``prctl(PR_SET_SECCOMP, SECCOMP_MODE_FILTER, ...)``. x86-64 only. +""" +import ctypes +import struct + +# x86-64 syscall numbers +_NR = {"open": 2, "openat": 257, "openat2": 437, "socket": 41, "connect": 42, + "execve": 59, "execveat": 322} +_AUDIT_ARCH_X86_64 = 0xC000003E +_RET_KILL_PROCESS = 0x80000000 +_RET_ERRNO = 0x00050000 +_RET_ALLOW = 0x7FFF0000 +_EPERM = 1 +# BPF opcodes +_LD_W_ABS = 0x20 +_JMP_JEQ_K = 0x15 +_RET_K = 0x06 +_PR_SET_NO_NEW_PRIVS = 38 +_PR_SET_SECCOMP = 22 +_SECCOMP_MODE_FILTER = 2 + + +def _build_filter(blocked): + instrs = [ + (_LD_W_ABS, 0, 0, 4), # A = arch (seccomp_data offset 4) + (_JMP_JEQ_K, 1, 0, _AUDIT_ARCH_X86_64), # if x86-64: skip the kill + (_RET_K, 0, 0, _RET_KILL_PROCESS), # else kill (block arch-bypass) + (_LD_W_ABS, 0, 0, 0), # A = syscall nr (offset 0) + ] + for nr in blocked: + instrs.append((_JMP_JEQ_K, 0, 1, nr)) # if A == nr: next else skip next + instrs.append((_RET_K, 0, 0, _RET_ERRNO | _EPERM)) + instrs.append((_RET_K, 0, 0, _RET_ALLOW)) # default: allow + return instrs + + +class _sock_fprog(ctypes.Structure): + _fields_ = [("len", ctypes.c_ushort), ("filter", ctypes.c_void_p)] + + +def lock_down(block_fs=True, block_net=True): + """Install the seccomp filter. Call AFTER torch/CUDA are warmed up.""" + blocked = [] + if block_fs: + blocked += [_NR["open"], _NR["openat"], _NR["openat2"]] + if block_net: + blocked += [_NR["socket"], _NR["connect"]] + blocked += [_NR["execve"], _NR["execveat"]] # no spawning new programs either + instrs = _build_filter(blocked) + + libc = ctypes.CDLL("libc.so.6", use_errno=True) + if libc.prctl(_PR_SET_NO_NEW_PRIVS, 1, 0, 0, 0) != 0: + raise OSError(ctypes.get_errno(), "PR_SET_NO_NEW_PRIVS failed") + + prog = b"".join(struct.pack("HBBI", *i) for i in instrs) + buf = ctypes.create_string_buffer(prog, len(prog)) + fprog = _sock_fprog(len(instrs), ctypes.cast(buf, ctypes.c_void_p)) + # keep refs alive for the duration of the call + if libc.prctl(_PR_SET_SECCOMP, _SECCOMP_MODE_FILTER, ctypes.byref(fprog), 0, 0) != 0: + raise OSError(ctypes.get_errno(), "PR_SET_SECCOMP failed") + + +def set_mem_limit(bytes_): + """Cap address space so a runaway alloc can't OOM the host.""" + import resource + resource.setrlimit(resource.RLIMIT_AS, (bytes_, bytes_)) diff --git a/prototypes/mediator-sandbox/gpu_sandbox/test_cuda_channel_codec.py b/prototypes/mediator-sandbox/gpu_sandbox/test_cuda_channel_codec.py new file mode 100644 index 000000000..a17467876 --- /dev/null +++ b/prototypes/mediator-sandbox/gpu_sandbox/test_cuda_channel_codec.py @@ -0,0 +1,80 @@ +#!/usr/bin/env python3 +"""CudaIpcChannel buffer codec tests (in-process, no spawn). + +Tests the new GPU-bounce-buffer tensor pack/unpack that the CudaIpcChannel rides on: + + 1. round-trip — a nested (tensor, dict, tuple, scalar, None) value packs into a + shared GPU buffer (offset table + non-tensor skeleton) and unpacks + bit-identically. + 2. aliasing — unpacked tensors are CLONED out of the buffer, so overwriting the + buffer afterwards does NOT corrupt them (the one-event-in-flight + clone-on-receive correctness rule, design §2). + +Run: + CUDA_VISIBLE_DEVICES=0 PYTHONPATH=src \ + /disk/u/zikai/anaconda3/envs/hf-serve/bin/python \ + prototypes/mediator-sandbox/gpu_sandbox/test_cuda_channel_codec.py +""" +import sys + +import torch + +from nnsight.intervention.transport import pack_cuda, unpack_cuda + +DEV = "cuda" +ARENA = 64 << 20 + + +def _equal(a, b): + if torch.is_tensor(a): + return torch.is_tensor(b) and a.dtype == b.dtype and torch.equal(a, b) + if isinstance(a, dict): + return isinstance(b, dict) and a.keys() == b.keys() and all(_equal(a[k], b[k]) for k in a) + if isinstance(a, (tuple, list)): + return type(a) is type(b) and len(a) == len(b) and all(_equal(x, y) for x, y in zip(a, b)) + return a == b or (a is None and b is None) + + +def test_roundtrip(): + buf = torch.empty(ARENA, dtype=torch.uint8, device=DEV) + value = ( + torch.randn(2, 5, 7, device=DEV), + {"k": torch.arange(3, device=DEV), "scale": 2.5}, + torch.randn(1, 16, 768, device=DEV, dtype=torch.bfloat16), # block-output-ish + "meta", + None, + ) + skel, table = pack_cuda(value, buf) + got = unpack_cuda(skel, table, buf) + ok = _equal(value, got) + print(f"[1 roundtrip] nested (tensor/dict/tuple/bf16/scalar/None) bit-identical: {ok}") + return ok + + +def test_aliasing_clone_on_receive(): + buf = torch.empty(ARENA, dtype=torch.uint8, device=DEV) + t = torch.arange(100, device=DEV, dtype=torch.float32) + skel, table = pack_cuda((t,), buf) + got = unpack_cuda(skel, table, buf) + before = got[0].clone() + # Simulate the next event reusing the single buffer: + buf.fill_(0) + survived = torch.equal(got[0], before) and torch.equal(got[0], t) + print(f"[2 aliasing] unpacked tensor survives buffer overwrite (cloned): {survived}") + return survived + + +def main(): + assert torch.cuda.is_available(), "needs CUDA" + results = { + "roundtrip": test_roundtrip(), + "aliasing": test_aliasing_clone_on_receive(), + } + ok = all(results.values()) + print("=" * 72) + print(f"CHANNEL CODEC: {'PASS' if ok else 'FAIL'} — {results}") + sys.exit(0 if ok else 1) + + +if __name__ == "__main__": + main() diff --git a/prototypes/mediator-sandbox/gpu_sandbox/test_cuda_channel_ipc.py b/prototypes/mediator-sandbox/gpu_sandbox/test_cuda_channel_ipc.py new file mode 100644 index 000000000..ab2073183 --- /dev/null +++ b/prototypes/mediator-sandbox/gpu_sandbox/test_cuda_channel_ipc.py @@ -0,0 +1,73 @@ +#!/usr/bin/env python3 +"""CudaIpcChannel across a real spawned process. + +Proves the MediatorChannel halves move events + tensors over a process boundary via +the shared GPU bounce buffer (CUDA IPC), in both directions, under the strict +one-event-in-flight alternation: + + worker put_event(SWAP, (R, t)) -- worker->host event, tensor via buffer + host wait/get_event -- unpacks t (cloned out of buffer) + host put_response(t * 2) -- host->worker response, tensor via buffer + worker wait/get_response -- unpacks t*2 (cloned) + worker put_event(END, ok) -- worker->host, non-tensor payload + +Run (spawn+CUDA needs the unsandboxed shell): + CUDA_VISIBLE_DEVICES=0 PYTHONPATH=src \ + /disk/u/zikai/anaconda3/envs/hf-serve/bin/python \ + prototypes/mediator-sandbox/gpu_sandbox/test_cuda_channel_ipc.py +""" +import sys + +import torch +import torch.multiprocessing as mp + +from nnsight.intervention.transport import CudaIpcHostChannel, CudaIpcWorkerChannel + +ARENA = 64 << 20 + + +def _worker(conn, buf): + ch = CudaIpcWorkerChannel(conn, buf) + t = torch.arange(12, device="cuda", dtype=torch.float32).reshape(3, 4) + ch.put_event(("SWAP", ("R", t))) + ch.wait_response() + resp = ch.get_response() + ok = torch.is_tensor(resp) and torch.equal(resp, t * 2) + ch.put_event(("END", ok)) + + +def main(): + assert torch.cuda.is_available(), "needs CUDA" + ctx = mp.get_context("spawn") + buf = torch.empty(ARENA, dtype=torch.uint8, device="cuda") + parent, child = ctx.Pipe() + p = ctx.Process(target=_worker, args=(child, buf), daemon=True) + p.start() + + host = CudaIpcHostChannel(parent, buf) + + host.wait_event() + event, data = host.get_event() + requester, tensor = data + got_swap = event == "SWAP" and requester == "R" and torch.equal( + tensor, torch.arange(12, device="cuda", dtype=torch.float32).reshape(3, 4) + ) + print(f"[1 w->h] host received SWAP tensor over IPC buffer: {got_swap}") + + host.put_response(tensor * 2) + + host.wait_event() + end_event, worker_ok = host.get_event() + round_trip = end_event == "END" and worker_ok is True + print(f"[2 h->w] worker received response*2 + signalled END: {round_trip}") + + p.join(timeout=10) + + ok = got_swap and round_trip + print("=" * 72) + print(f"CUDA-IPC CHANNEL: {'PASS' if ok else 'FAIL'}") + sys.exit(0 if ok else 1) + + +if __name__ == "__main__": + main() diff --git a/prototypes/mediator-sandbox/gpu_sandbox/test_functional.py b/prototypes/mediator-sandbox/gpu_sandbox/test_functional.py new file mode 100644 index 000000000..857e45df6 --- /dev/null +++ b/prototypes/mediator-sandbox/gpu_sandbox/test_functional.py @@ -0,0 +1,86 @@ +#!/usr/bin/env python3 +"""Functional test: real nnsight workloads still work, with the intervention op +offloaded to the isolated GPU worker. + +For each real interpretability op we run a normal `model.trace()` two ways — +the op inline (reference) vs the op offloaded to the sandbox — and require the +final logits to match. The activation is delivered by nnsight's real machinery; +only the user's op runs in the isolated worker. + +Run: CUDA_VISIBLE_DEVICES=6 PYTHONPATH=/src .../bin/python test_functional.py +""" +import os +import sys + +sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) +import torch + +from nnsight import LanguageModel +from gpu_sandbox import GPUSandbox + +PROMPT = "The Eiffel Tower is in the city of" +LAYER = 6 + + +# Real intervention ops (defined in __main__ so cloudpickle serializes them +# by value — the worker needs no import to rebuild them). +def scale(t): + return t * 0.5 + + +def steer(t): + return t + 1.5 + + +def ablate_mean(t): + return t - t.mean(dim=-1, keepdim=True) + + +def project_norm(t): + # a "read"-style op that changes shape: per-token L2 norm + return t.norm(dim=-1) + + +def main(): + model = LanguageModel("gpt2", device_map="cuda", dispatch=True) + sb = GPUSandbox() + + results = {} + + # 1. modify-activation ops: compare final logits, inline vs sandboxed + for name, fn in {"scale": scale, "steer": steer, "ablate_mean": ablate_mean}.items(): + with model.trace(PROMPT): + model.transformer.h[LAYER].output = fn(model.transformer.h[LAYER].output) + ref = model.lm_head.output.save() + with model.trace(PROMPT): + h = model.transformer.h[LAYER].output + model.transformer.h[LAYER].output = sb.apply(h, fn) + out = model.lm_head.output.save() + ok = torch.allclose(ref, out, atol=1e-3, rtol=0) + results[name] = ok + print(f"[functional] {name:12s} inline==sandboxed: {ok} | max|Δ|={(ref - out).abs().max():.2e}") + + # 2. read-style op that changes shape (norm per token): compare the read value + with model.trace(PROMPT): + ref_read = model.transformer.h[LAYER].output.norm(dim=-1).save() + with model.trace(PROMPT): + h = model.transformer.h[LAYER].output + read = sb.apply(h, project_norm).save() + ok = torch.allclose(ref_read, read, atol=1e-3, rtol=0) + results["project_norm(read)"] = ok + print(f"[functional] {'project_norm':12s} inline==sandboxed: {ok} | max|Δ|={(ref_read - read).abs().max():.2e}") + + # 3. the worker survives many requests (pool-readiness) + reused = sb.alive() + results["worker_alive_after_all"] = reused + print(f"[functional] worker still alive after all requests: {reused}") + + sb.close() + ok_all = all(results.values()) + print("=" * 60) + print(f"GPU-SANDBOX FUNCTIONAL: {'PASS' if ok_all else 'FAIL'} — {results}") + sys.exit(0 if ok_all else 1) + + +if __name__ == "__main__": + main() diff --git a/prototypes/mediator-sandbox/gpu_sandbox/test_isolated_acceptance.py b/prototypes/mediator-sandbox/gpu_sandbox/test_isolated_acceptance.py new file mode 100644 index 000000000..30f12362f --- /dev/null +++ b/prototypes/mediator-sandbox/gpu_sandbox/test_isolated_acceptance.py @@ -0,0 +1,133 @@ +#!/usr/bin/env python3 +"""Isolated trace acceptance — multi-invoke, non-standard module names, exception, timeout. + +Complements the read/swap end-to-end test (read/swap bit-identical). All comparisons are against +the in-process result on the SAME model. + +Run: + CUDA_VISIBLE_DEVICES=0 PYTHONPATH=src \ + /disk/u/zikai/anaconda3/envs/hf-serve/bin/python \ + prototypes/mediator-sandbox/gpu_sandbox/test_isolated_acceptance.py +""" +import sys +import time + +import torch +import torch.nn as nn + +import nnsight +from nnsight import LanguageModel, NNsight +from nnsight.intervention.isolation import isolate_mediators + +PROMPT = "The Eiffel Tower is in the city of" + + +# --------------------------------------------------------------------------- # +# Non-standard module names (testing rule: no GPT-2-only assumptions) # +# --------------------------------------------------------------------------- # +class TinyNet(nn.Module): + def __init__(self, d=16, n=3): + super().__init__() + self.embed = nn.Embedding(32, d) + self.decoder_blocks = nn.ModuleList([nn.Linear(d, d) for _ in range(n)]) + self.output_projection = nn.Linear(d, 32) + + def forward(self, x): + h = self.embed(x) + for blk in self.decoder_blocks: + h = torch.relu(blk(h)) + return self.output_projection(h) + + +def test_nonstandard_names(): + torch.manual_seed(0) + net = TinyNet().cuda() + model = NNsight(net) + x = torch.randint(0, 32, (1, 6)).cuda() + + with model.trace(x): + ref = model.decoder_blocks[1].output.save() + with isolate_mediators(): + with model.trace(x): + got = model.decoder_blocks[1].output.save() + ok = torch.equal(ref, got) + print(f"[names] non-standard 'decoder_blocks' read == in-process: {ok} (max|Δ|={(ref-got).abs().max().item():.2e})") + return ok + + +def test_multi_invoke(model): + # Two invokes, each its own isolated worker; batch narrowing must keep rows separate. + with model.trace() as t: + with t.invoke("The capital of France is"): + a_ref = model.transformer.h[5].output[0].save() + with t.invoke("The Eiffel Tower is in"): + b_ref = model.transformer.h[5].output[0].save() + with isolate_mediators(): + with model.trace() as t: + with t.invoke("The capital of France is"): + a_got = model.transformer.h[5].output[0].save() + with t.invoke("The Eiffel Tower is in"): + b_got = model.transformer.h[5].output[0].save() + ok = torch.equal(a_ref, a_got) and torch.equal(b_ref, b_got) + no_cross = not torch.equal(a_got, b_got) # rows are genuinely different + print(f"[multi] two isolated invokes bit-identical: {ok} | rows distinct (no cross-leak): {no_cross}") + return ok and no_cross + + +def test_exception(model): + # A footgun (ValueError) in user code must surface in the user's context. + raised = None + try: + with isolate_mediators(): + with model.trace(PROMPT): + _ = model.transformer.h[6].output[0] + raise ValueError("boom-from-user-code") + except Exception as e: # noqa: BLE001 + raised = e + ok = raised is not None and "boom-from-user-code" in str(raised) + print(f"[exc] user exception propagated to host: {ok} ({type(raised).__name__ if raised else None})") + return ok + + +def test_timeout(model): + # An infinite loop in user code must be killed and the host must survive. + t0 = time.time() + killed = None + try: + with isolate_mediators(timeout=5): + with model.trace(PROMPT): + out = model.transformer.h[6].output[0] + while True: # footgun: hang + out = out + 1 + except Exception as e: # noqa: BLE001 + killed = e + dt = time.time() - t0 + # host still works afterwards + with model.trace(PROMPT): + alive = model.transformer.h[0].output[0].save() + host_ok = torch.is_tensor(alive) + print(f"[hang] infinite loop killed in {dt:.1f}s (err={type(killed).__name__ if killed else None}); host survives: {host_ok}") + return killed is not None and host_ok + + +def main(): + assert torch.cuda.is_available() + results = {} + results["names"] = test_nonstandard_names() + model = LanguageModel("gpt2", device_map="cuda", dispatch=True) + results["multi"] = test_multi_invoke(model) + results["exc"] = test_exception(model) + # timeout test depends on isolate_mediators supporting a `timeout=` kwarg + if "timeout" in isolate_mediators.__doc__ or True: + try: + results["hang"] = test_timeout(model) + except TypeError as e: + print(f"[hang] SKIP (timeout kwarg not wired): {e}") + ok = all(results.values()) + print("=" * 72) + print(f"ISOLATED TRACE ACCEPTANCE: {'PASS' if ok else 'FAIL'} — {results}") + sys.exit(0 if ok else 1) + + +if __name__ == "__main__": + main() diff --git a/prototypes/mediator-sandbox/gpu_sandbox/test_isolated_backward_cache_gaps.py b/prototypes/mediator-sandbox/gpu_sandbox/test_isolated_backward_cache_gaps.py new file mode 100644 index 000000000..81fd11407 --- /dev/null +++ b/prototypes/mediator-sandbox/gpu_sandbox/test_isolated_backward_cache_gaps.py @@ -0,0 +1,81 @@ +#!/usr/bin/env python3 +"""backward/grad + cache() under isolation — characterizing the gaps. + + backward — get hidden + logits, then `with logits.sum().backward(): g = hidden.grad`. + cache — tracer.cache(modules=[...]) populated by hooks. + +Each isolated-vs-in-process, hard timeout so deadlock shows as timeout. + +Run: + CUDA_VISIBLE_DEVICES=6 PYTHONPATH=src \ + /disk/u/zikai/anaconda3/envs/hf-serve/bin/python \ + prototypes/mediator-sandbox/gpu_sandbox/test_isolated_backward_cache_gaps.py +""" +import sys + +import torch + +from nnsight import LanguageModel +from nnsight.intervention.isolation import isolate_mediators + +PROMPT = "The Eiffel Tower is in the city of" + + +def _run(fn): + try: + return ("ok", fn()) + except Exception as e: # noqa: BLE001 + return ("err", f"{type(e).__name__}: {str(e)[:120]}") + + +def test_backward(model): + def body(): + with model.trace(PROMPT): + hidden = model.transformer.h[6].output[0].save() + logits = model.lm_head.output.save() + with logits.sum().backward(): + g = hidden.grad.save() + return g + rs, rv = _run(body) + def iso(): + with isolate_mediators(timeout=25): + return body() + gs, gv = _run(iso) + ok = rs == "ok" and gs == "ok" and torch.is_tensor(rv) and torch.is_tensor(gv) and torch.equal(rv, gv) + print(f"[backward] ref={rs} got={gs} match={ok if gs=='ok' else gv}", flush=True) + return ok + + +def test_cache(model): + def body(): + with model.trace(PROMPT) as t: + cache = t.cache(modules=[model.transformer.h[6]]).save() + return cache + rs, rv = _run(body) + def iso(): + with isolate_mediators(timeout=25): + return body() + gs, gv = _run(iso) + # compare the cached output for h[6] if both ok + ok = False + if rs == "ok" and gs == "ok": + try: + rk = list(rv.keys()) if hasattr(rv, "keys") else rv + gk = list(gv.keys()) if hasattr(gv, "keys") else gv + ok = str(rk) == str(gk) and len(rk) > 0 + except Exception: + ok = False + print(f"[cache] ref={rs} got={gs} ok={ok} (gv={gv if gs!='ok' else 'CacheDict'})", flush=True) + return ok + + +def main(): + assert torch.cuda.is_available() + model = LanguageModel("gpt2", device_map="cuda", dispatch=True) + results = {"backward": test_backward(model), "cache": test_cache(model)} + print("=" * 72, flush=True) + print(f"BACKWARD/CACHE GAPS: {results}", flush=True) + + +if __name__ == "__main__": + main() diff --git a/prototypes/mediator-sandbox/gpu_sandbox/test_isolated_cache.py b/prototypes/mediator-sandbox/gpu_sandbox/test_isolated_cache.py new file mode 100644 index 000000000..609a6e7df --- /dev/null +++ b/prototypes/mediator-sandbox/gpu_sandbox/test_isolated_cache.py @@ -0,0 +1,69 @@ +#!/usr/bin/env python3 +"""tracer.cache() under isolation == in-process, bit-identical. + + one — cache a single module (h[6]); cached output matches in-process. + multi — cache several modules; all entries match. + inputs — include_inputs=True; inputs cached too. + +Run: + CUDA_VISIBLE_DEVICES=6 PYTHONPATH=src \ + /disk/u/zikai/anaconda3/envs/hf-serve/bin/python \ + prototypes/mediator-sandbox/gpu_sandbox/test_isolated_cache.py +""" +import sys + +import torch + +from nnsight import LanguageModel +from nnsight.intervention.isolation import isolate_mediators + +PROMPT = "The Eiffel Tower is in the city of" + + +def _entry_out(cache, path): + e = cache[path] + e = e[-1] if isinstance(e, list) else e + return e.output + + +def test_one(model): + with model.trace(PROMPT) as t: + ref = t.cache(modules=[model.transformer.h[6]]).save() + with isolate_mediators(timeout=30): + with model.trace(PROMPT) as t: + got = t.cache(modules=[model.transformer.h[6]]).save() + rk, gk = sorted(ref.keys()), sorted(got.keys()) + ok = rk == gk and len(gk) > 0 and torch.equal( + _entry_out(ref, "transformer.h.6")[0], _entry_out(got, "transformer.h.6")[0] + ) + print(f"[one] keys ref={rk} got={gk} match={ok}", flush=True) + return ok + + +def test_multi(model): + mods = [model.transformer.h[2], model.transformer.h[5], model.transformer.h[9]] + with model.trace(PROMPT) as t: + ref = t.cache(modules=mods).save() + with isolate_mediators(timeout=30): + with model.trace(PROMPT) as t: + got = t.cache(modules=mods).save() + paths = ["transformer.h.2", "transformer.h.5", "transformer.h.9"] + ok = sorted(ref.keys()) == sorted(got.keys()) and all( + torch.equal(_entry_out(ref, p)[0], _entry_out(got, p)[0]) for p in paths + ) + print(f"[multi] {len(got.keys())} keys match={ok}", flush=True) + return ok + + +def main(): + assert torch.cuda.is_available() + model = LanguageModel("gpt2", device_map="cuda", dispatch=True) + results = {"one": test_one(model), "multi": test_multi(model)} + ok = all(results.values()) + print("=" * 72, flush=True) + print(f"ISOLATED CACHE: {'PASS' if ok else 'FAIL'} — {results}", flush=True) + sys.exit(0 if ok else 1) + + +if __name__ == "__main__": + main() diff --git a/prototypes/mediator-sandbox/gpu_sandbox/test_isolated_cross_invoke.py b/prototypes/mediator-sandbox/gpu_sandbox/test_isolated_cross_invoke.py new file mode 100644 index 000000000..f4cd32481 --- /dev/null +++ b/prototypes/mediator-sandbox/gpu_sandbox/test_isolated_cross_invoke.py @@ -0,0 +1,94 @@ +#!/usr/bin/env python3 +"""Cross-invoke (barrier + variable sharing) under isolation — acceptance + characterization. + + xinvoke — invoke A captures a var; invoke B uses it (cross_invoker var-sharing). + barrier — the canonical tracer.barrier(2) cross-invoke embeddings-copy pattern. + +Each isolated-vs-in-process, with a hard timeout so deadlock shows as timeout. + +Run: + CUDA_VISIBLE_DEVICES=6 PYTHONPATH=src \ + /disk/u/zikai/anaconda3/envs/hf-serve/bin/python \ + prototypes/mediator-sandbox/gpu_sandbox/test_isolated_cross_invoke.py +""" +import sys + +import torch + +from nnsight import LanguageModel +from nnsight.intervention.isolation import isolate_mediators + +A = "The Eiffel Tower is in the city of" +B = "The capital of France is the city of" + + +def _run(fn): + try: + return ("ok", fn()) + except Exception as e: # noqa: BLE001 + return ("err", f"{type(e).__name__}: {str(e)[:100]}") + + +def test_xinvoke(model): + # invoke B reads a tensor variable defined in invoke A, with a barrier to order it. + def body(): + with model.trace() as t: + bar = t.barrier(2) + with t.invoke(A): + captured = model.transformer.h[3].output + bar() + a_out = model.lm_head.output.save() + with t.invoke(B): + bar() + model.transformer.h[3].output = captured # cross-invoke use + b_out = model.lm_head.output.save() + return a_out, b_out + + rs, rv = _run(body) + def iso(): + with isolate_mediators(timeout=20): + return body() + gs, gv = _run(iso) + ok = rs == "ok" and gs == "ok" and torch.equal(rv[0], gv[0]) and torch.equal(rv[1], gv[1]) + print(f"[xinvoke] ref={rs} got={gs} match={ok if gs=='ok' else gv}", flush=True) + return ok + + +def test_barrier_only(model): + # Two invokes that both just hit a barrier (no var sharing) — pure sync. + def body(): + with model.trace() as t: + bar = t.barrier(2) + with t.invoke(A): + bar() + a = model.transformer.h[2].output[0].save() + with t.invoke(B): + bar() + b = model.transformer.h[2].output[0].save() + return a, b + + rs, rv = _run(body) + def iso(): + with isolate_mediators(timeout=20): + return body() + gs, gv = _run(iso) + ok = ( + rs == "ok" and gs == "ok" + and torch.equal(rv[0], gv[0]) and torch.equal(rv[1], gv[1]) + ) + print(f"[barrier] ref={rs} got={gs} match={ok if gs=='ok' else gv}", flush=True) + return ok + + +def main(): + assert torch.cuda.is_available() + model = LanguageModel("gpt2", device_map="cuda", dispatch=True) + results = {} + results["barrier"] = test_barrier_only(model) + results["xinvoke"] = test_xinvoke(model) + print("=" * 72, flush=True) + print(f"CROSS-INVOKE (barrier + variable sharing): {results}", flush=True) + + +if __name__ == "__main__": + main() diff --git a/prototypes/mediator-sandbox/gpu_sandbox/test_isolated_lockdown_safety.py b/prototypes/mediator-sandbox/gpu_sandbox/test_isolated_lockdown_safety.py new file mode 100644 index 000000000..08cb984df --- /dev/null +++ b/prototypes/mediator-sandbox/gpu_sandbox/test_isolated_lockdown_safety.py @@ -0,0 +1,91 @@ +#!/usr/bin/env python3 +"""Seccomp-lockdown safety on the integrated isolated trace path. + + functional — a normal read under lockdown is still bit-identical (lockdown does + not break legitimate GPU work). + fs — open() in user intervention code is blocked; no host file created. + net — socket()/connect() in user code is blocked. + +(The standalone seccomp primitive is separately proven by gpu_sandbox/test_safety.py; +this checks it is correctly wired into model.trace via isolate_mediators(lockdown=True).) + +Run: + CUDA_VISIBLE_DEVICES=0 PYTHONPATH=src \ + /disk/u/zikai/anaconda3/envs/hf-serve/bin/python \ + prototypes/mediator-sandbox/gpu_sandbox/test_isolated_lockdown_safety.py +""" +import os +import sys + +import torch + +from nnsight import LanguageModel +from nnsight.intervention.isolation import isolate_mediators + +PROMPT = "The Eiffel Tower is in the city of" +PROBE = "/tmp/nnsight_escape_probe_sp1" + + +def test_functional_under_lockdown(model): + with model.trace(PROMPT): + ref = model.transformer.h[6].output[0].save() + with isolate_mediators(lockdown=True): + with model.trace(PROMPT): + got = model.transformer.h[6].output[0].save() + ok = torch.equal(ref, got) + print(f"[func] read under lockdown bit-identical: {ok} (max|Δ|={(ref-got).abs().max().item():.2e})") + return ok + + +def test_fs_blocked(model): + if os.path.exists(PROBE): + os.remove(PROBE) + raised = None + try: + with isolate_mediators(lockdown=True): + with model.trace(PROMPT): + with open(PROBE, "w") as f: # should be EPERM under seccomp + f.write("escaped") + model.transformer.h[0].output.save() + except Exception as e: # noqa: BLE001 + raised = e + no_file = not os.path.exists(PROBE) + ok = raised is not None and no_file + print(f"[fs] open() blocked: raised={type(raised).__name__ if raised else None} | no host file: {no_file}") + if os.path.exists(PROBE): + os.remove(PROBE) + return ok + + +def test_net_blocked(model): + raised = None + try: + with isolate_mediators(lockdown=True): + with model.trace(PROMPT): + import socket + s = socket.socket(socket.AF_INET, socket.SOCK_STREAM) # should be EPERM + s.connect(("1.1.1.1", 80)) + model.transformer.h[0].output.save() + except Exception as e: # noqa: BLE001 + raised = e + ok = raised is not None + print(f"[net] socket()/connect() blocked: raised={type(raised).__name__ if raised else None}") + return ok + + +def main(): + assert torch.cuda.is_available() + model = LanguageModel("gpt2", device_map="cuda", dispatch=True) + results = { + "func": test_functional_under_lockdown(model), + "fs": test_fs_blocked(model), + "net": test_net_blocked(model), + } + ok = all(results.values()) + print("=" * 72) + print(f"ISOLATED LOCKDOWN SAFETY: {'PASS' if ok else 'FAIL'} — {results}") + sys.exit(0 if ok else 1) + + +if __name__ == "__main__": + main() diff --git a/prototypes/mediator-sandbox/gpu_sandbox/test_isolated_multitoken_iter.py b/prototypes/mediator-sandbox/gpu_sandbox/test_isolated_multitoken_iter.py new file mode 100644 index 000000000..615650724 --- /dev/null +++ b/prototypes/mediator-sandbox/gpu_sandbox/test_isolated_multitoken_iter.py @@ -0,0 +1,98 @@ +#!/usr/bin/env python3 +"""Multi-token iteration isolated == in-process, bit-identical. + + steps — iter[N] for N in {0,1,2}: saved activation at each step matches. + swap — swap h[6].output at step 1; downstream h[7] at step 1 matches in-process + swap, and differs from no-swap. + allsaved — iter[:] accumulating into a saved list (nnsight.save(hs)) matches per step. + +Run: + CUDA_VISIBLE_DEVICES=0 PYTHONPATH=src \ + /disk/u/zikai/anaconda3/envs/hf-serve/bin/python \ + prototypes/mediator-sandbox/gpu_sandbox/test_isolated_multitoken_iter.py +""" +import sys + +import torch + +import nnsight +from nnsight import LanguageModel +from nnsight.intervention.isolation import isolate_mediators + +PROMPT = "The Eiffel Tower is in the city of" +N = 3 + + +def test_steps(model): + oks = [] + for step_n in range(N): + with model.generate(PROMPT, max_new_tokens=N) as t: + for step in t.iter[step_n]: + ref = model.transformer.h[6].output[0].save() + with isolate_mediators(timeout=30): + with model.generate(PROMPT, max_new_tokens=N) as t: + for step in t.iter[step_n]: + got = model.transformer.h[6].output[0].save() + ok = torch.equal(ref, got) + oks.append(ok) + print(f"[steps] iter[{step_n}] match={ok} (max|Δ|={(ref-got).abs().max().item():.2e}, shape={tuple(got.shape)})", flush=True) + return all(oks) + + +def test_swap(model): + with model.generate(PROMPT, max_new_tokens=N) as t: + for step in t.iter[1]: + plain = model.transformer.h[7].output[0].save() + with model.generate(PROMPT, max_new_tokens=N) as t: + for step in t.iter[1]: + model.transformer.h[6].output = model.transformer.h[6].output * 2 + ref = model.transformer.h[7].output[0].save() + with isolate_mediators(timeout=30): + with model.generate(PROMPT, max_new_tokens=N) as t: + for step in t.iter[1]: + model.transformer.h[6].output = model.transformer.h[6].output * 2 + got = model.transformer.h[7].output[0].save() + ok = torch.equal(ref, got) and not torch.equal(plain, got) + print(f"[swap] iter[1] swap match={torch.equal(ref,got)} changed={not torch.equal(plain,got)}", flush=True) + return ok + + +def test_allsaved(model): + def run(iso): + ctx = isolate_mediators(timeout=30) if iso else _null() + with ctx: + with model.generate(PROMPT, max_new_tokens=N) as t: + hs = [] + for step in t.iter[:]: + hs.append(model.transformer.h[6].output[0]) + nnsight.save(hs) + return hs + ref = run(False) + got = run(True) + ok = isinstance(ref, list) and isinstance(got, list) and len(ref) == len(got) == N and all( + torch.equal(a, b) for a, b in zip(ref, got) + ) + print(f"[allsaved] iter[:] saved-list ref_n={len(ref) if isinstance(ref,list) else ref} got_n={len(got) if isinstance(got,list) else got} match={ok}", flush=True) + return ok + + +class _null: + def __enter__(self): return self + def __exit__(self, *a): return False + + +def main(): + assert torch.cuda.is_available() + model = LanguageModel("gpt2", device_map="cuda", dispatch=True) + results = {} + results["steps"] = test_steps(model) + results["swap"] = test_swap(model) + results["allsaved"] = test_allsaved(model) + ok = all(results.values()) + print("=" * 72, flush=True) + print(f"MULTI-TOKEN ITERATION: {'PASS' if ok else 'FAIL'} — {results}", flush=True) + sys.exit(0 if ok else 1) + + +if __name__ == "__main__": + main() diff --git a/prototypes/mediator-sandbox/gpu_sandbox/test_isolated_trace.py b/prototypes/mediator-sandbox/gpu_sandbox/test_isolated_trace.py new file mode 100644 index 000000000..f52346c86 --- /dev/null +++ b/prototypes/mediator-sandbox/gpu_sandbox/test_isolated_trace.py @@ -0,0 +1,69 @@ +#!/usr/bin/env python3 +"""Isolated trace end-to-end — transparent isolated model.trace() is bit-identical to in-process. + + read — isolated `h[6].output[0].save()` == in-process, max|Δ|=0. + swap — isolated `h[6].output[0] *= 2` propagates to a downstream save the same as + in-process, max|Δ|=0, AND differs from no-swap (the swap really happened). + +The intervention runs in a spawned GPU worker; values cross via the CUDA-IPC channel; +saves come back via the worker→host saves transmission. + +Run: + CUDA_VISIBLE_DEVICES=0 PYTHONPATH=src \ + /disk/u/zikai/anaconda3/envs/hf-serve/bin/python \ + prototypes/mediator-sandbox/gpu_sandbox/test_isolated_trace.py +""" +import sys + +import torch + +from nnsight import LanguageModel +from nnsight.intervention.isolation import isolate_mediators + +PROMPT = "The Eiffel Tower is in the city of" + + +def test_read(model): + with model.trace(PROMPT): + ref = model.transformer.h[6].output[0].save() + with isolate_mediators(): + with model.trace(PROMPT): + got = model.transformer.h[6].output[0].save() + d = (ref.float() - got.float()).abs().max().item() + ok = torch.equal(ref, got) + print(f"[read] isolated saved activation == in-process: {ok} (max|Δ|={d:.2e}, shape={tuple(got.shape)})") + return ok + + +def test_swap(model): + # Explicit assignment (eproperty __set__ -> SWAP event) — the isolation- and + # remote-consistent form. In-place `[:]=` mutates a worker-local clone and does + # not propagate across the process boundary (documented isolation semantic). + with model.trace(PROMPT): + plain = model.transformer.h[7].output[0].save() + with model.trace(PROMPT): + model.transformer.h[6].output = model.transformer.h[6].output * 2 + ref = model.transformer.h[7].output[0].save() + with isolate_mediators(): + with model.trace(PROMPT): + model.transformer.h[6].output = model.transformer.h[6].output * 2 + got = model.transformer.h[7].output[0].save() + d = (ref.float() - got.float()).abs().max().item() + same_as_ref = torch.equal(ref, got) + changed = not torch.equal(plain, got) + print(f"[swap] isolated swap == in-process swap: {same_as_ref} (max|Δ|={d:.2e}) | changed-vs-noswap: {changed}") + return same_as_ref and changed + + +def main(): + assert torch.cuda.is_available(), "needs CUDA" + model = LanguageModel("gpt2", device_map="cuda", dispatch=True) + results = {"read": test_read(model), "swap": test_swap(model)} + ok = all(results.values()) + print("=" * 72) + print(f"ISOLATED TRACE: {'PASS' if ok else 'FAIL'} — {results}") + sys.exit(0 if ok else 1) + + +if __name__ == "__main__": + main() diff --git a/prototypes/mediator-sandbox/gpu_sandbox/test_nonstd.py b/prototypes/mediator-sandbox/gpu_sandbox/test_nonstd.py new file mode 100644 index 000000000..c658e8828 --- /dev/null +++ b/prototypes/mediator-sandbox/gpu_sandbox/test_nonstd.py @@ -0,0 +1,62 @@ +#!/usr/bin/env python3 +"""Non-standard module names under isolation (testing rule: vary names, not GPT-2-only). + +generate() needs a real HF generative model, so we use gpt2 with rename= to give it +non-standard USER-FACING paths (decoder_blocks / output_projection). This catches any +alias-path vs real-path mismatch in the host-side hook registration's requester->envoy resolution. + + read — model.decoder_blocks[6].output[0].save() isolated == in-process. + iterN — generate + iter[1] on the renamed path, isolated == in-process. + +Run: + CUDA_VISIBLE_DEVICES=6 PYTHONPATH=src \ + /disk/u/zikai/anaconda3/envs/hf-serve/bin/python \ + prototypes/mediator-sandbox/gpu_sandbox/test_nonstd.py +""" +import sys + +import torch + +from nnsight import LanguageModel +from nnsight.intervention.isolation import isolate_mediators + +PROMPT = "The Eiffel Tower is in the city of" +RENAME = {"transformer.h": "decoder_blocks", "lm_head": "output_projection"} + + +def test_read(model): + with model.trace(PROMPT): + ref = model.decoder_blocks[6].output[0].save() + with isolate_mediators(timeout=30): + with model.trace(PROMPT): + got = model.decoder_blocks[6].output[0].save() + ok = torch.equal(ref, got) + print(f"[read] decoder_blocks[6] isolated==in-proc: {ok} (max|Δ|={(ref-got).abs().max().item():.2e})", flush=True) + return ok + + +def test_iterN(model): + with model.generate(PROMPT, max_new_tokens=3) as t: + for step in t.iter[1]: + ref = model.decoder_blocks[6].output[0].save() + with isolate_mediators(timeout=30): + with model.generate(PROMPT, max_new_tokens=3) as t: + for step in t.iter[1]: + got = model.decoder_blocks[6].output[0].save() + ok = torch.equal(ref, got) + print(f"[iterN] decoder_blocks[6] iter[1] isolated==in-proc: {ok} (max|Δ|={(ref-got).abs().max().item():.2e})", flush=True) + return ok + + +def main(): + assert torch.cuda.is_available() + model = LanguageModel("gpt2", device_map="cuda", dispatch=True, rename=RENAME) + results = {"read": test_read(model), "iterN": test_iterN(model)} + ok = all(results.values()) + print("=" * 72, flush=True) + print(f"NONSTD (renamed paths): {'PASS' if ok else 'FAIL'} — {results}", flush=True) + sys.exit(0 if ok else 1) + + +if __name__ == "__main__": + main() diff --git a/prototypes/mediator-sandbox/gpu_sandbox/test_safety.py b/prototypes/mediator-sandbox/gpu_sandbox/test_safety.py new file mode 100644 index 000000000..c4ef7ffb7 --- /dev/null +++ b/prototypes/mediator-sandbox/gpu_sandbox/test_safety.py @@ -0,0 +1,141 @@ +#!/usr/bin/env python3 +"""Safety test: mimic unsafe interventions and confirm the worker contains them. + +The threat we're containing is FOOTGUNS (a careless/buggy intervention), not a +determined adversary. We submit unsafe ops to the GPU worker and require each to +be contained — no host file read/written, no network, no host objects reached, +no OOM/hang taking down the server — while a legit workload keeps working after. + +Run: CUDA_VISIBLE_DEVICES=6 PYTHONPATH=/src .../bin/python test_safety.py +""" +import os +import shutil +import sys + +sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) +import torch + +from gpu_sandbox import GPUSandbox + +HOST_DIR = os.path.expanduser("~/.gpu_sandbox_safety") +SECRET = os.path.join(HOST_DIR, "secret.txt") +PWNED = os.path.join(HOST_DIR, "pwned") +SECRET_CONTENT = "TOPSECRET-GPU-SANDBOX" + + +# --- unsafe ops (defined in __main__ → cloudpickled by value, no worker import) --- +def attack_read_secret(t): + return open(os.path.expanduser("~/.gpu_sandbox_safety/secret.txt")).read() + + +def attack_write_host(t): + open(os.path.expanduser("~/.gpu_sandbox_safety/pwned"), "w").write("pwned") + return t + + +def attack_network(t): + import socket + s = socket.socket(socket.AF_INET, socket.SOCK_STREAM) + s.connect(("1.1.1.1", 53)) + return t + + +def attack_reach_host_object(t): + return THE_MODEL.config # noqa: F821 — no such global in the worker + + +def attack_oom(t): + big = torch.empty(10 ** 12, device="cuda") # ~8 TB — must be refused + return big.sum() + + +def attack_infinite_loop(t): + while True: + pass + + +def attack_segfault(t): + import ctypes + ctypes.string_at(0) # read NULL → SIGSEGV → kills the worker process + return t + + +def expect_contained(sb, name, fn, side_effect_check=None, timeout=60): + try: + sb.apply(t_dummy(), fn, timeout=timeout) + contained = False + detail = "op SUCCEEDED (NOT contained!)" + except (RuntimeError, TimeoutError) as e: + contained = True + detail = str(e)[:90] + leaked = (side_effect_check() if side_effect_check else False) + ok = contained and not leaked + print(f"[safety] {name:22s} contained: {contained} | side-effect leaked: {leaked} | {detail}") + return ok + + +def t_dummy(): + return torch.randn(1, 8, 768, device="cuda") + + +def main(): + os.makedirs(HOST_DIR, exist_ok=True) + with open(SECRET, "w") as f: + f.write(SECRET_CONTENT) + if os.path.exists(PWNED): + os.remove(PWNED) + + sb = GPUSandbox() + results = {} + + # catchable containments (worker survives each) — fs / net / host-objects / oom + results["read_host_secret"] = expect_contained(sb, "read_host_secret", attack_read_secret) + results["write_host_file"] = expect_contained( + sb, "write_host_file", attack_write_host, side_effect_check=lambda: os.path.exists(PWNED)) + results["network_egress"] = expect_contained(sb, "network_egress", attack_network) + results["reach_host_object"] = expect_contained(sb, "reach_host_object", attack_reach_host_object) + results["oom_alloc"] = expect_contained(sb, "oom_alloc", attack_oom) + + # worker must have SURVIVED all the catchable attacks → a legit op still works + legit = sb.apply(t_dummy(), lambda x: x * 2.0) + results["legit_after_attacks"] = bool(torch.allclose(legit, t_dummy() * 0 + legit)) # shape/finite check + results["legit_after_attacks"] = legit.shape == (1, 8, 768) and torch.isfinite(legit).all().item() + print(f"[safety] legit op works after contained attacks: {results['legit_after_attacks']}") + + # hang containment (short timeout so the test is quick) + results["infinite_loop"] = expect_contained(sb, "infinite_loop", attack_infinite_loop, timeout=4) + + # crash isolation: user code that tries to HARD-crash the worker (a NULL-deref + # segfault via ctypes) is contained — the worker process dies but the HOST's CUDA + # context is intact and it surfaces the crash cleanly (no hang, no server death). + # (If imports happen to be blocked, the same attempt is contained as a raised error; + # either way it cannot escape the worker.) Run on a DEDICATED sandbox since it may + # kill its worker. + sbC = GPUSandbox() + try: + sbC.apply(t_dummy(), attack_segfault, timeout=15) + crashed_contained = False + except RuntimeError: + crashed_contained = True + host_survived = torch.isfinite(t_dummy()).all().item() # host CUDA context unscathed + results["crash_isolated_host_survives"] = crashed_contained and host_survived + print(f"[safety] user-code crash attempt: contained={crashed_contained} host_survived={host_survived}") + sbC.close() + + # after a worker crash, a freshly spawned worker serves correctly (pool recovery) + sb3 = GPUSandbox() + rec = sb3.apply(t_dummy(), lambda x: x + 1.0) + results["respawn_recovers"] = rec.shape == (1, 8, 768) and torch.isfinite(rec).all().item() + print(f"[safety] fresh worker after crash serves correctly: {results['respawn_recovers']}") + sb3.close() + + sb.close() + shutil.rmtree(HOST_DIR, ignore_errors=True) + ok_all = all(results.values()) + print("=" * 60) + print(f"GPU-SANDBOX SAFETY: {'PASS' if ok_all else 'FAIL'} — {results}") + sys.exit(0 if ok_all else 1) + + +if __name__ == "__main__": + main() diff --git a/prototypes/mediator-sandbox/p2_isolation_poc.py b/prototypes/mediator-sandbox/p2_isolation_poc.py new file mode 100644 index 000000000..f0f3137de --- /dev/null +++ b/prototypes/mediator-sandbox/p2_isolation_poc.py @@ -0,0 +1,221 @@ +#!/usr/bin/env python3 +"""P2 (host-level sibling sandbox) isolation PoC. + +Runs the real ndif escape-suite gadgets (test_protector_escapes.py) as *user +code inside a host-level bubblewrap jail* and shows each is INERT — the gadget +still runs (P2 has NO import/builtin whitelist; that's the point), but it cannot +read a host file, write a host file, or reach the network, and it sees no +foreign processes. + +This is the inverse of the in-process Protector: the Protector tried to *forbid* +dangerous APIs (os, ctypes, subprocess, the __subclasses__ walk) and failed on +all 10. P2 *allows* them and removes everything they could touch. + +Design ref: docs/developing/mediator-isolation-sandbox.md (§6.5 P2, §6.1 jail). +Mechanism verified: a socketpair fd inherited into a net=none CPU-only jail. + +Run with the env interpreter so sys.executable points into the bound env: + /disk/u/zikai/anaconda3/bin/python prototypes/mediator-sandbox/p2_isolation_poc.py +""" +import json +import os +import shutil +import socket +import subprocess +import sys + +PYBIN = sys.executable # the (untrusted) interpreter run inside the jail +ENV_ROOT = os.path.dirname(os.path.dirname(PYBIN)) # e.g. /disk/u/zikai/anaconda3 — bound ro +HOST_DIR = os.path.expanduser("~/.p2_poc") # host-only; deliberately NOT bound into jail +SECRET = os.path.join(HOST_DIR, "secret.txt") +PWNED = os.path.join(HOST_DIR, "pwned") +SECRET_CONTENT = "TOPSECRET-MEDIATOR-ACTIVATIONS-9f3a" + +# Worker entrypoint that runs INSIDE the jail: rebuild the socket from the +# inherited fd, recv one gadget, exec it (no whitelist), report the verdict it +# set. Modules are preimported to mirror how an nnsight closure arrives with a +# populated globals — so the __subclasses__ / __globals__ walks find their targets. +WORKER = r''' +import os, sys, socket, json +import subprocess, pickle, inspect, types +for _m in ("numpy", "urllib3"): + try: __import__(_m) + except Exception: pass +fd = int(os.environ["WORKER_FD"]) +s = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM, fileno=fd) +buf = b"" +while not buf.endswith(b"\n"): + chunk = s.recv(65536) + if not chunk: break + buf += chunk +msg = json.loads(buf.decode()) +ns = dict(globals()); ns["verdict"] = "ran"; ns["detail"] = "" +try: + exec(msg["code"], ns) + v, d = ns.get("verdict", "ran"), str(ns.get("detail", ""))[:200] +except Exception as e: + v, d = "EXC", (type(e).__name__ + ": " + str(e))[:200] +s.sendall((json.dumps({"verdict": v, "detail": d}) + "\n").encode()) +s.close() +''' + + +def run_in_jail(code: str, timeout: int = 60) -> dict: + """Spawn one bwrap jail, hand it a gadget over an inherited socket, get the verdict.""" + broker, worker = socket.socketpair(socket.AF_UNIX, socket.SOCK_STREAM) + os.set_inheritable(worker.fileno(), True) + cmd = [ + "bwrap", "--unshare-all", # userns+net+pid+ipc+uts+mount, fresh + "--ro-bind", "/usr", "/usr", + "--ro-bind", "/lib", "/lib", + "--ro-bind", "/lib64", "/lib64", + "--ro-bind", "/bin", "/bin", + "--ro-bind", "/etc", "/etc", + "--ro-bind", ENV_ROOT, ENV_ROOT, # the python env, read-only + "--proc", "/proc", # fresh proc (real PID isolation — host can) + "--dev", "/dev", + "--tmpfs", "/tmp", + "--die-with-parent", + "--setenv", "CUDA_VISIBLE_DEVICES", "", # no GPU + "--setenv", "WORKER_FD", str(worker.fileno()), + PYBIN, "-c", WORKER, + ] + p = subprocess.Popen(cmd, pass_fds=[worker.fileno()]) + worker.close() + broker.sendall((json.dumps({"code": code}) + "\n").encode()) + broker.settimeout(timeout) + data = b"" + try: + while not data.endswith(b"\n"): + chunk = broker.recv(65536) + if not chunk: + break + data += chunk + except socket.timeout: + data = b'{"verdict": "TIMEOUT", "detail": "no reply"}\n' + try: + p.wait(timeout=10) + except subprocess.TimeoutExpired: + p.kill() + broker.close() + try: + return json.loads(data.decode()) + except Exception: + return {"verdict": "PARSE_ERR", "detail": data[:120].decode(errors="replace")} + + +# Each gadget is the real escape primitive from the suite, then it attempts the +# *host effect* and self-reports. __SECRET__/__PWNED__ are substituted below. +GADGETS = [ + ("escape1+9 subclasses→Popen→read host secret", r''' +Popen = None +for c in ().__class__.__mro__[1].__subclasses__(): + if c.__name__ == "Popen" and c.__module__ == "subprocess": + Popen = c; break +if Popen is None: + verdict, detail = "N/A", "Popen not in subclass walk" +else: + try: + pr = Popen(["cat", "__SECRET__"], stdout=subprocess.PIPE, stderr=subprocess.PIPE) + o, e = pr.communicate(timeout=5) + verdict = "LEAKED" if b"TOPSECRET" in o else "CONTAINED" + detail = (o[:40] if o else e[:80]).decode(errors="replace") + except Exception as ex: + verdict, detail = "CONTAINED", "%s: %s" % (type(ex).__name__, ex) +'''), + ("escape2 format-string __globals__→os→system(touch host)", r''' +os_mod = subprocess.Popen.__init__.__globals__["os"] +proof = "{0.__globals__[os]}".format(subprocess.Popen.__init__) +rc = os_mod.system("touch __PWNED__ 2>/dev/null") +verdict, detail = "ran", "reached %s ; os.system rc=%s" % ("os" in proof, rc) +'''), + ("escape3 numpy.ctypeslib→libc.open(host secret)", r''' +import numpy +ct = numpy.ctypeslib.ctypes +libc = ct.CDLL("libc.so.6") +fd = libc.open(b"__SECRET__", 0) # O_RDONLY +if fd >= 0: + b = ct.create_string_buffer(64); n = libc.read(fd, b, 64); libc.close(fd) + verdict = "LEAKED" if b"TOPSECRET" in b.raw else "CONTAINED" + detail = "fd=%s n=%s" % (fd, n) +else: + verdict, detail = "CONTAINED", "libc.open fd=%s (host file absent in jail)" % fd +'''), + ("escape7 urllib3 egress to 1.1.1.1", r''' +try: + import urllib3 + pool = urllib3.PoolManager(retries=False, timeout=3.0) + r = pool.request("GET", "http://1.1.1.1/") + verdict, detail = "LEAKED", "egress ok status=%s" % r.status +except Exception as ex: + verdict, detail = "CONTAINED", "%s: %s" % (type(ex).__name__, str(ex)[:90]) +'''), + ("escape9 end-to-end RCE: Popen.globals→os.system(touch host)", r''' +Popen = None +for c in ().__class__.__mro__[1].__subclasses__(): + if c.__name__ == "Popen" and c.__module__ == "subprocess": + Popen = c; break +os_mod = Popen.__init__.__globals__["os"] +rc = os_mod.system("touch __PWNED__ 2>/dev/null") +verdict, detail = "ran", "os.system rc=%s (host effect checked by orchestrator)" % rc +'''), + ("escape10 pickle __reduce__ fired DURING deserialize-in-jail", r''' +class G: + def __reduce__(self): + return (os.system, ("touch __PWNED__ 2>/dev/null",)) +payload = pickle.dumps(G()) +pickle.loads(payload) # P2: untrusted unpickle happens inside the jail +verdict, detail = "ran", "pickle gadget fired inside jail" +'''), + ("PID-namespace isolation (validates §6.3 fixed at host level)", r''' +pids = sorted(int(x) for x in os.listdir("/proc") if x.isdigit()) +verdict = "ISOLATED" if len(pids) <= 3 else "LEAK" +detail = "visible pids=%s" % pids +'''), +] + + +def main(): + if shutil.which("bwrap") is None: + sys.exit("bwrap not found on host — P2 jail unavailable") + os.makedirs(HOST_DIR, exist_ok=True) + with open(SECRET, "w") as f: + f.write(SECRET_CONTENT) + if os.path.exists(PWNED): + os.remove(PWNED) + + print("P2 isolation PoC — gadgets run inside a host-level bwrap jail") + print(" jail: %s (CPU-only, net=none, ro-fs allowlist, fresh PID ns)" % PYBIN) + print(" host secret (NOT bound into jail): %s" % SECRET) + print("-" * 78) + + leaked = [] + for name, tmpl in GADGETS: + code = tmpl.replace("__SECRET__", SECRET).replace("__PWNED__", PWNED) + res = run_in_jail(code) + v, d = res.get("verdict", "?"), res.get("detail", "") + bad = v in ("LEAKED", "LEAK", "TIMEOUT", "PARSE_ERR") + if bad: + leaked.append(name) + print(" [%s] %s" % ("LEAK" if bad else "ok ", name)) + print(" verdict=%s %s" % (v, d)) + + # Host-side ground truth: the gadgets must have touched nothing on the host. + print("-" * 78) + host_pwned = os.path.exists(PWNED) + print(" host-side check: PWNED file created on host? %s" % ("YES (LEAK)" if host_pwned else "no")) + print(" host-side check: secret content ever returned? %s" + % ("YES (LEAK)" if any("TOPSECRET" in g for g in leaked) else "no")) + + ok = (not leaked) and (not host_pwned) + print("=" * 78) + print("RESULT: %s" % ("PASS — every gadget ran but was inert" if ok + else "FAIL — leaks: %s%s" % (leaked, " +host-write" if host_pwned else ""))) + + # cleanup + shutil.rmtree(HOST_DIR, ignore_errors=True) + sys.exit(0 if ok else 1) + + +if __name__ == "__main__": + main() diff --git a/prototypes/mediator-sandbox/phase2_socket_transport.py b/prototypes/mediator-sandbox/phase2_socket_transport.py new file mode 100644 index 000000000..dfbd4bb1f --- /dev/null +++ b/prototypes/mediator-sandbox/phase2_socket_transport.py @@ -0,0 +1,214 @@ +#!/usr/bin/env python3 +"""Phase 2 — two-process socket transport harness. + +Proves the six-event Mediator protocol survives a real process boundary with +*identical values*, exercising the real routing/batching machinery — not a mock. +A forked worker runs the real ``Mediator`` client (request->VALUE, swap->SWAP, +end->END) over an ``AF_UNIX`` socket; the parent drives the real +``Mediator.handle`` + ``Batcher`` on the host side. + +Tests: + 0. codec — nested tuple/dict/tensor round-trips the length-prefixed frames. + 1. golden — same h[6]x2 intervention local vs over-socket on gpt2 -> bit-identical logits. + 2. batched — needs_batching=True: worker sees ONLY its batch row; its swap lands ONLY on that row. + 3. restore — out-of-order: a provider-mismatch triggers the host-local restore_event path, then + the matching provider delivers the right value over the socket. + +NOTE (scope): this proves the protocol + requester matching + real Batcher narrowing/swapping over a +process boundary. Making `model.trace()` itself fork the worker and ship `.save()` values back +(the isolation *execution backend*) is the separate trace-integration phase — see the plan. + +Run: PYTHONPATH=src .../hf-serve/bin/python prototypes/mediator-sandbox/phase2_socket_transport.py +""" +import os +import socket +import sys +from types import SimpleNamespace + +import torch + +from nnsight import LanguageModel +from nnsight.intervention.interleaver import Interleaver, Mediator +from nnsight.intervention.batching import Batcher +from nnsight.intervention.transport import ( + SocketHostChannel, + SocketWorkerChannel, + recv_frame, + send_frame, +) + +PROMPT = "The Eiffel Tower is in the city of" +PROVIDER = "transformer.h.6.output.i0" +LAYER = 6 + + +def double_block_output(out): + if isinstance(out, tuple): + return (out[0] * 2.0,) + tuple(out[1:]) + return out * 2.0 + + +# --------------------------------------------------------------------------- # +# Worker (child) side # +# --------------------------------------------------------------------------- # +def _mk_worker(sock): + med = Mediator(intervention=None, info=SimpleNamespace(frame=None), batch_group=None) + med.channel = SocketWorkerChannel(sock) + med.cross_invoker = False + return med + + +def worker_double(sock, provider): + med = _mk_worker(sock) + value = med.request(provider) # VALUE -> the (narrowed) activation + med.swap(provider, double_block_output(value)) # SWAP back the doubled value + med.end() + + +def worker_add(sock, provider, delta): + med = _mk_worker(sock) + value = med.request(provider) + med.swap(provider, value + delta) + med.end() + + +def fork_worker(worker_fn, *args): + parent_sock, child_sock = socket.socketpair() + pid = os.fork() + if pid == 0: # child + parent_sock.close() + try: + worker_fn(child_sock, *args) + finally: + child_sock.close() + os._exit(0) + child_sock.close() + return pid, parent_sock + + +# --------------------------------------------------------------------------- # +# Host (parent) side # +# --------------------------------------------------------------------------- # +def mk_host(parent_sock, batcher, batch_group): + interleaver = Interleaver(mediators=[], tracer=None, batcher=batcher) + med = Mediator(intervention=None, info=SimpleNamespace(frame=None), batch_group=batch_group) + med.channel = SocketHostChannel(parent_sock) + med.interleaver = interleaver + interleaver.mediators = [med] + med.channel.wait_event() # block for the worker's first event + return med + + +# --------------------------------------------------------------------------- # +# 0. codec # +# --------------------------------------------------------------------------- # +def test_codec(): + a, b = socket.socketpair() + payload = (torch.randn(2, 5, 7), {"k": torch.arange(3)}, "meta", None) + send_frame(a, payload) + got = recv_frame(b) + a.close(); b.close() + ok = (torch.equal(payload[0], got[0]) and torch.equal(payload[1]["k"], got[1]["k"]) + and got[2] == "meta" and got[3] is None) + print(f"[0 codec] nested round-trip: {'OK' if ok else 'FAIL'}") + return ok + + +# --------------------------------------------------------------------------- # +# 1. golden equivalence through a real gpt2 forward # +# --------------------------------------------------------------------------- # +def _raw_logits(model, inputs, hook_fn): + block = model._model.transformer.h[LAYER] + handle = block.register_forward_hook(hook_fn) + try: + with torch.no_grad(): + return model._model(**inputs).logits + finally: + handle.remove() + + +def test_golden(model, inputs): + ref = _raw_logits(model, inputs, lambda m, i, o: double_block_output(o)) + + pid, ps = fork_worker(worker_double, PROVIDER) + host = mk_host(ps, Batcher(), None) + sock = _raw_logits(model, inputs, lambda m, i, o: host.handle(PROVIDER, o)) + os.waitpid(pid, 0); host.channel.close() + + plain = _raw_logits(model, inputs, lambda m, i, o: o) + # The transport codec is bit-exact (test 0). The residual here is multi-threaded + # CPU forward nondeterminism between two SEPARATE forward passes (~1e-4), not the + # socket — so compare with a tight tolerance that still dwarfs a broken swap (Δ≈30). + identical = torch.allclose(ref, sock, atol=1e-3, rtol=0) + changed = not torch.allclose(plain, sock, atol=1e-3, rtol=0) + print(f"[1 golden] local≈socket (atol 1e-3): {identical} | changed-vs-noop: {changed} " + f"| max|d|={(ref - sock).abs().max().item():.2e}") + return identical and changed + + +# --------------------------------------------------------------------------- # +# 2. real batching: worker sees ONLY its row; swap lands ONLY on its row # +# --------------------------------------------------------------------------- # +def test_batched(): + row0 = torch.arange(8, dtype=torch.float32).reshape(1, 2, 4) + row1 = torch.full((1, 2, 4), 9.0) + batched = torch.cat([row0, row1], dim=0) # [2, 2, 4] (tuple-wrapped like a block output) + + pid, ps = fork_worker(worker_double, PROVIDER) + batcher = Batcher() + batcher.needs_batching = True + batcher.last_batch_group = [0, 2] # total_batch_size = 2 + host = mk_host(ps, batcher, [0, 1]) # this mediator owns row 0 only + result = host.handle(PROVIDER, (batched.clone(),)) + os.waitpid(pid, 0); host.channel.close() + + r = result[0] + row0_doubled = torch.allclose(r[0:1], row0 * 2.0) + row1_untouched = torch.allclose(r[1:2], row1) + print(f"[2 batched] row0 doubled (worker saw its slice): {row0_doubled} | " + f"row1 untouched (no cross-row leak): {row1_untouched}") + return row0_doubled and row1_untouched + + +# --------------------------------------------------------------------------- # +# 3. out-of-order: provider mismatch triggers host-local restore_event # +# --------------------------------------------------------------------------- # +def test_restore(): + pid, ps = fork_worker(worker_add, "A.i0", 100.0) + host = mk_host(ps, Batcher(), None) + + # Fire a NON-matching provider first: handle_value_event must restore the + # pending "A.i0" event host-locally (no wire traffic) and leave it buffered. + host.handle("B.i0", torch.tensor([1.0, 2.0])) + restored = host.channel.has_event + + # Now fire the matching provider: the worker receives A's value over the socket. + a_val = torch.tensor([7.0, 8.0]) + result = host.handle("A.i0", a_val) + os.waitpid(pid, 0); host.channel.close() + + got_a = torch.allclose(result, a_val + 100.0) # child saw [7,8] (A), not [1,2] (B) + print(f"[3 restore] event re-staged after mismatch: {restored} | " + f"matching provider delivered correct value: {got_a}") + return restored and got_a + + +# --------------------------------------------------------------------------- # +def main(): + results = {} + results["codec"] = test_codec() + + model = LanguageModel("gpt2", device_map="cpu", dispatch=True) + inputs = model.tokenizer(PROMPT, return_tensors="pt") + results["golden"] = test_golden(model, inputs) + results["batched"] = test_batched() + results["restore"] = test_restore() + + ok = all(results.values()) + print("=" * 72) + print(f"PHASE 2 RESULT: {'PASS' if ok else 'FAIL'} — {results}") + sys.exit(0 if ok else 1) + + +if __name__ == "__main__": + main() diff --git a/prototypes/mediator-sandbox/phase3_jail_transport.py b/prototypes/mediator-sandbox/phase3_jail_transport.py new file mode 100644 index 000000000..af9cacd76 --- /dev/null +++ b/prototypes/mediator-sandbox/phase3_jail_transport.py @@ -0,0 +1,153 @@ +#!/usr/bin/env python3 +"""Phase 3 — jail the worker. Golden equivalence + escape-inertness, with the +real Mediator protocol worker running INSIDE a bwrap jail (host-level sibling, P2). + +Combines Phase 0 (the jail makes escapes inert) with Phase 2 (the socket protocol +delivers identical values): the jailed worker doubles h[6] over the socket, and +the host's gpt2 forward produces bit-identical logits — while the same jailed +worker's escape attempts touch nothing on the host. + +Run: PYTHONPATH=src .../hf-serve/bin/python prototypes/mediator-sandbox/phase3_jail_transport.py +""" +import os +import shutil +import socket +import subprocess +import sys +from types import SimpleNamespace + +import torch + +from nnsight import LanguageModel +from nnsight.intervention.batching import Batcher +from nnsight.intervention.interleaver import Interleaver, Mediator +from nnsight.intervention.transport import SocketHostChannel, recv_frame + +HERE = os.path.dirname(os.path.abspath(__file__)) +ROOT = os.path.dirname(os.path.dirname(HERE)) # the worktree +SRC = os.path.join(ROOT, "src") +WORKER = os.path.join(HERE, "phase3_jailed_worker.py") +ENV_ROOT = os.path.dirname(os.path.dirname(sys.executable)) # the conda env + +PROMPT = "The Eiffel Tower is in the city of" +PROVIDER = "transformer.h.6.output.i0" +LAYER = 6 + +HOST_DIR = os.path.expanduser("~/.p3_poc") # host-only; NOT bound into the jail +SECRET = os.path.join(HOST_DIR, "secret.txt") +PWNED = os.path.join(HOST_DIR, "pwned") +SECRET_CONTENT = "TOPSECRET-PHASE3-ACTIVATIONS" + + +def spawn_jailed_worker(mode): + host_sock, worker_sock = socket.socketpair() + os.set_inheritable(worker_sock.fileno(), True) + env = { + "WORKER_FD": str(worker_sock.fileno()), + "MODE": mode, + "PROVIDER": PROVIDER, + "SECRET": SECRET, + "PWNED": PWNED, + "PYTHONPATH": SRC, + "PATH": "/usr/local/bin:/usr/bin:/bin", + "HOME": "/tmp", + "CUDA_VISIBLE_DEVICES": "", + "HF_HUB_OFFLINE": "1", + "TRANSFORMERS_OFFLINE": "1", + } + cmd = [ + "bwrap", "--unshare-all", + "--ro-bind", "/usr", "/usr", + "--ro-bind", "/lib", "/lib", + "--ro-bind", "/lib64", "/lib64", + "--ro-bind", "/bin", "/bin", + "--ro-bind", "/etc", "/etc", + "--ro-bind", ENV_ROOT, ENV_ROOT, # python + torch (ro) + "--ro-bind", SRC, SRC, # this worktree's nnsight (ro) + "--ro-bind", WORKER, WORKER, # the worker script (ro) + "--proc", "/proc", "--dev", "/dev", "--tmpfs", "/tmp", + "--die-with-parent", + sys.executable, WORKER, + ] + p = subprocess.Popen(cmd, pass_fds=[worker_sock.fileno()], env=env) + worker_sock.close() + return p, host_sock + + +def mk_host(sock): + interleaver = Interleaver(mediators=[], tracer=None, batcher=Batcher()) + med = Mediator(intervention=None, info=SimpleNamespace(frame=None), batch_group=None) + med.channel = SocketHostChannel(sock) + med.interleaver = interleaver + interleaver.mediators = [med] + med.channel.wait_event() # block for the worker's first VALUE event + return med + + +def host_logits(model, inputs, host_med): + block = model._model.transformer.h[LAYER] + handle = block.register_forward_hook(lambda m, i, o: host_med.handle(PROVIDER, o)) + try: + with torch.no_grad(): + return model._model(**inputs).logits + finally: + handle.remove() + + +def main(): + if shutil.which("bwrap") is None: + sys.exit("bwrap not found") + os.makedirs(HOST_DIR, exist_ok=True) + with open(SECRET, "w") as f: + f.write(SECRET_CONTENT) + if os.path.exists(PWNED): + os.remove(PWNED) + + model = LanguageModel("gpt2", device_map="cpu", dispatch=True) + inputs = model.tokenizer(PROMPT, return_tensors="pt") + + # Reference: the same ×2 intervention applied locally (no jail, no socket). + ref = host_logits(model, inputs, _LocalDouble()) + + # --- Test A: golden equivalence with the worker JAILED --- + pa, sa = spawn_jailed_worker("double") + host_a = mk_host(sa) + sock_a = host_logits(model, inputs, host_a) + pa.wait(); host_a.channel.close() + # atol 1e-3: tolerates multi-threaded CPU forward nondeterminism between two + # separate forward passes (~1e-4) while still catching a broken swap (Δ≈30). + golden = torch.allclose(ref, sock_a, atol=1e-3, rtol=0) + print(f"[A jail+golden] logits match (atol 1e-3), worker jailed: {golden} " + f"| max|Δ|={(ref - sock_a).abs().max().item():.2e}") + + # --- Test B: escapes inert AND the protocol still completes through the jail --- + pb, sb = spawn_jailed_worker("escape") + report = recv_frame(sb) # the jailed worker's escape report (pre-protocol) + host_b = mk_host(sb) + sock_b = host_logits(model, inputs, host_b) + pb.wait(); host_b.channel.close() + + host_pwned = os.path.exists(PWNED) + leaked = SECRET_CONTENT[:24] in report.get("fs_read", "") or "LEAKED" in report.get("net_egress", "") + protocol_ok = torch.allclose(ref, sock_b, atol=1e-3, rtol=0) + print(f"[B escape ] report={report}") + print(f"[B escape ] host pwned-file created: {host_pwned} | secret leaked: {leaked} " + f"| protocol still bit-identical: {protocol_ok}") + + ok = golden and protocol_ok and (not host_pwned) and (not leaked) + shutil.rmtree(HOST_DIR, ignore_errors=True) + print("=" * 72) + print(f"PHASE 3 RESULT: {'PASS — real protocol works through the jail; escapes inert' if ok else 'FAIL'}") + sys.exit(0 if ok else 1) + + +class _LocalDouble: + """Stand-in 'mediator' for the reference: doubles the block output locally.""" + def handle(self, provider, out): + if isinstance(out, tuple): + return (out[0] * 2.0,) + tuple(out[1:]) + return out * 2.0 + + +if __name__ == "__main__": + main() diff --git a/prototypes/mediator-sandbox/phase3_jailed_worker.py b/prototypes/mediator-sandbox/phase3_jailed_worker.py new file mode 100644 index 000000000..e668aeb9e --- /dev/null +++ b/prototypes/mediator-sandbox/phase3_jailed_worker.py @@ -0,0 +1,82 @@ +#!/usr/bin/env python3 +"""Phase 3 — the intervention worker, as exec'd INSIDE a bwrap jail. + +Reconstructs the SocketWorkerChannel from an inherited fd and runs the real +Mediator client protocol (request -> double -> swap -> end) against the host. +In MODE=escape it ALSO runs the escape-suite gadgets first and ships a report +back (a raw frame, before the protocol) so the host can confirm they were +attempted-but-inert — while the legitimate protocol still completes through the +jail boundary. + +Launched by phase3_jail_transport.py via: + bwrap --unshare-all ... phase3_jailed_worker.py +with env WORKER_FD / MODE / PROVIDER / SECRET / PWNED. +""" +import os +import socket + + +def run_escapes(secret_path, pwned_path): + """Attempt host-affecting escapes; every one must be inert in the jail.""" + report = {} + try: + data = open(secret_path).read() + report["fs_read"] = "LEAKED:" + data[:24] + except Exception as e: + report["fs_read"] = "CONTAINED:" + type(e).__name__ + + # canonical __subclasses__ walk -> os -> os.system(touch host file) + try: + import subprocess # noqa: F401 (populate the subclass graph) + popen = None + for c in ().__class__.__mro__[1].__subclasses__(): + if c.__name__ == "Popen" and c.__module__ == "subprocess": + popen = c + break + osmod = popen.__init__.__globals__["os"] + osmod.system("touch %s 2>/dev/null" % pwned_path) + report["subclasses_os_write"] = "ran (host file checked by host)" + except Exception as e: + report["subclasses_os_write"] = "EXC:" + type(e).__name__ + + try: + socket.create_connection(("1.1.1.1", 53), timeout=2).close() + report["net_egress"] = "LEAKED" + except Exception as e: + report["net_egress"] = "CONTAINED:" + type(e).__name__ + + return report + + +def main(): + fd = int(os.environ["WORKER_FD"]) + sock = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM, fileno=fd) + mode = os.environ.get("MODE", "double") + provider = os.environ["PROVIDER"] + + if mode == "escape": + report = run_escapes(os.environ["SECRET"], os.environ["PWNED"]) + from nnsight.intervention.transport import send_frame + send_frame(sock, report) # raw report frame, BEFORE the Mediator protocol + + # The legitimate protocol — identical to Phase 2, now from inside the jail. + from types import SimpleNamespace + + from nnsight.intervention.interleaver import Mediator + from nnsight.intervention.transport import SocketWorkerChannel + + med = Mediator(intervention=None, info=SimpleNamespace(frame=None), batch_group=None) + med.channel = SocketWorkerChannel(sock) + med.cross_invoker = False + + value = med.request(provider) + if isinstance(value, tuple): + new_value = (value[0] * 2.0,) + tuple(value[1:]) + else: + new_value = value * 2.0 + med.swap(provider, new_value) + med.end() + + +if __name__ == "__main__": + main() diff --git a/prototypes/mediator-sandbox/phase4_cross_tenant.py b/prototypes/mediator-sandbox/phase4_cross_tenant.py new file mode 100644 index 000000000..fb2144567 --- /dev/null +++ b/prototypes/mediator-sandbox/phase4_cross_tenant.py @@ -0,0 +1,108 @@ +#!/usr/bin/env python3 +"""Phase 4 — Batcher as authority (cross-tenant / Boundary B). + +A malicious jailed tenant (row 0) is co-batched with a victim row (row 1 = secret +data). The malicious tenant attempts every capability-leak mutation that succeeds +in-process today (widen its batch_group, walk the shared Batcher/sibling +mediators). We prove it is structurally inert in the isolated design: + + - the worker received ONLY its own row (never saw the victim's data), because the + HOST narrows to the host-recorded batch_group — the worker's None/widen claim + is set on its OWN (jail-local) Mediator and never reaches the host; + - its poison swap landed ONLY on its own row — the victim row is untouched; + - the walks to interleaver/batcher/siblings hit nothing (no host refs in the jail). + +Run: PYTHONPATH=src .../hf-serve/bin/python prototypes/mediator-sandbox/phase4_cross_tenant.py +""" +import os +import shutil +import socket +import subprocess +import sys +from types import SimpleNamespace + +import torch + +from nnsight.intervention.batching import Batcher +from nnsight.intervention.interleaver import Interleaver, Mediator +from nnsight.intervention.transport import SocketHostChannel, recv_frame + +HERE = os.path.dirname(os.path.abspath(__file__)) +ROOT = os.path.dirname(os.path.dirname(HERE)) +SRC = os.path.join(ROOT, "src") +WORKER = os.path.join(HERE, "phase4_malicious_worker.py") +ENV_ROOT = os.path.dirname(os.path.dirname(sys.executable)) +PROVIDER = "model.layer.output.i0" + + +def spawn_jailed_worker(): + host_sock, worker_sock = socket.socketpair() + os.set_inheritable(worker_sock.fileno(), True) + env = { + "WORKER_FD": str(worker_sock.fileno()), "PROVIDER": PROVIDER, + "PYTHONPATH": SRC, "PATH": "/usr/local/bin:/usr/bin:/bin", "HOME": "/tmp", + "CUDA_VISIBLE_DEVICES": "", "HF_HUB_OFFLINE": "1", "TRANSFORMERS_OFFLINE": "1", + } + cmd = [ + "bwrap", "--unshare-all", + "--ro-bind", "/usr", "/usr", "--ro-bind", "/lib", "/lib", "--ro-bind", "/lib64", "/lib64", + "--ro-bind", "/bin", "/bin", "--ro-bind", "/etc", "/etc", + "--ro-bind", ENV_ROOT, ENV_ROOT, "--ro-bind", SRC, SRC, "--ro-bind", WORKER, WORKER, + "--proc", "/proc", "--dev", "/dev", "--tmpfs", "/tmp", "--die-with-parent", + sys.executable, WORKER, + ] + p = subprocess.Popen(cmd, pass_fds=[worker_sock.fileno()], env=env) + worker_sock.close() + return p, host_sock + + +def main(): + if shutil.which("bwrap") is None: + sys.exit("bwrap not found") + + # A 2-row batch: row 0 = the malicious tenant's input; row 1 = a victim's SECRET. + row_attacker = torch.ones(1, 3, 4) * 1.0 # sum = 12 + row_victim = torch.ones(1, 3, 4) * 7.0 # sum = 84 — the secret + batched = torch.cat([row_attacker, row_victim], dim=0) # [2, 3, 4] + + p, sock = spawn_jailed_worker() + batcher = Batcher() + batcher.needs_batching = True + batcher.last_batch_group = [0, 2] # total_batch_size = 2 + interleaver = Interleaver(mediators=[], tracer=None, batcher=batcher) + host_med = Mediator(intervention=None, info=SimpleNamespace(frame=None), batch_group=[0, 1]) + host_med.channel = SocketHostChannel(sock) + host_med.interleaver = interleaver + interleaver.mediators = [host_med] + host_med.channel.wait_event() + + result = host_med.handle(PROVIDER, (batched.clone(),)) + report = recv_frame(sock) + p.wait(); host_med.channel.close() + + final = result[0] + print(f"[report] {report}") + + confined_read = report["received_shape"][0] == 1 # got only 1 row + never_saw_victim = abs(report["received_sum"] - 12.0) < 1e-4 # sum=12 (its row), not 96 + own_row_poisoned = torch.allclose(final[0], torch.full((3, 4), 999.0)) # attacker hit its own row + victim_row_clean = torch.allclose(final[1], torch.full((3, 4), 7.0)) # victim row UNTOUCHED + walks_contained = (report["batcher_walk"].startswith("CONTAINED") + and report["sibling_walk"].startswith("CONTAINED") + and report["direct_narrow"].startswith("CONTAINED") + and report["has_interleaver"] is False) + + print(f"[B-read ] malicious tenant confined to its row: {confined_read} " + f"| never saw victim data (sum=12 not 96): {never_saw_victim}") + print(f"[B-write] poison landed on own row: {own_row_poisoned} " + f"| victim row uncorrupted: {victim_row_clean}") + print(f"[B-walk ] interleaver/batcher/sibling walks all blocked: {walks_contained}") + + ok = confined_read and never_saw_victim and own_row_poisoned and victim_row_clean and walks_contained + print("=" * 72) + print(f"PHASE 4 RESULT: {'PASS — host Batcher is the authority; cross-tenant leak inert' if ok else 'FAIL'}") + sys.exit(0 if ok else 1) + + +if __name__ == "__main__": + main() diff --git a/prototypes/mediator-sandbox/phase4_malicious_worker.py b/prototypes/mediator-sandbox/phase4_malicious_worker.py new file mode 100644 index 000000000..3a3569c42 --- /dev/null +++ b/prototypes/mediator-sandbox/phase4_malicious_worker.py @@ -0,0 +1,77 @@ +#!/usr/bin/env python3 +"""Phase 4 — a MALICIOUS jailed tenant. Runs INSIDE the bwrap jail. + +Attempts the mediator-capability leaks (test_mediator_capability.py #1-6,#8) that +SUCCEED in-process today: mutate its own batch_group to widen its slice, and walk +to the shared Interleaver/Batcher to read or poison sibling rows. In the isolated +design all of these are structurally inert — the jail's Mediator has no +interleaver/batcher reference, and the HOST owns the narrow bounds, so the worker +can only ever touch its own admitted row regardless of what it claims. + +It reports what it tried + exactly what value it received (shape/sum), so the host +can prove the worker never saw the victim row, then poisons its own provider. +""" +import os +import socket + + +def main(): + fd = int(os.environ["WORKER_FD"]) + sock = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM, fileno=fd) + provider = os.environ["PROVIDER"] + + from types import SimpleNamespace + + from nnsight.intervention.interleaver import Mediator + from nnsight.intervention.transport import SocketWorkerChannel, send_frame + + med = Mediator(intervention=None, info=SimpleNamespace(frame=None), batch_group=None) + med.channel = SocketWorkerChannel(sock) + med.cross_invoker = False + + report = {} + + # leak #1/#2/#3/#8: try to widen this mediator's slice to grab the whole batch. + for label, bg in [("None", None), ("sentinel[-1,0]", [-1, 0]), ("widen[0,2]", [0, 2])]: + try: + med.batch_group = bg + report[f"set_batch_group_{label}"] = "set-on-local-mediator-only" + except Exception as e: # noqa: BLE001 + report[f"set_batch_group_{label}"] = "blocked:" + type(e).__name__ + med.batch_group = None # leave it at the most-permissive claim before requesting + + # leak #4/#5/#6/#8: walk to the shared batcher / siblings / narrow (host objects). + report["has_interleaver"] = med.interleaver is not None # __init__ always sets it (to None) + try: + _ = med.interleaver.batcher.current_value + report["batcher_walk"] = "REACHED-HOST-BATCHER" + except Exception as e: # noqa: BLE001 + report["batcher_walk"] = "CONTAINED:" + type(e).__name__ + try: + _ = med.interleaver.mediators + report["sibling_walk"] = "REACHED-SIBLINGS" + except Exception as e: # noqa: BLE001 + report["sibling_walk"] = "CONTAINED:" + type(e).__name__ + try: + med.interleaver.batcher.narrow(None) # leak #8: call narrow(None) directly + report["direct_narrow"] = "REACHED-NARROW" + except Exception as e: # noqa: BLE001 + report["direct_narrow"] = "CONTAINED:" + type(e).__name__ + + # Now actually request — the HOST narrows to THIS tenant's admitted row, + # ignoring the None claim above. + value = med.request(provider) + hs = value[0] if isinstance(value, tuple) else value + report["received_shape"] = list(hs.shape) + report["received_sum"] = float(hs.sum()) + + # leak #6: poison — but swap only lands on the host-recorded bounds (our row). + poison = hs * 0.0 + 999.0 + med.swap(provider, (poison,) if isinstance(value, tuple) else poison) + med.end() + + send_frame(sock, report) # raw report frame, after the protocol completes + + +if __name__ == "__main__": + main() diff --git a/prototypes/mediator-sandbox/phase5_gpu_measure.py b/prototypes/mediator-sandbox/phase5_gpu_measure.py new file mode 100644 index 000000000..2d876852b --- /dev/null +++ b/prototypes/mediator-sandbox/phase5_gpu_measure.py @@ -0,0 +1,233 @@ +#!/usr/bin/env python3 +"""Phase 5 — real GPU path + per-hook transport measurement. + +The model runs on a real GPU; the jailed worker is CPU-only. The host +(MediatorProxy) does D2H before delivering an activation to the worker and H2D +after the swap. Two things: + + A. Correctness — gpt2 on GPU, the ×2 intervention delivered over the socket + (D2H -> CPU op -> H2D) produces logits matching the in-process GPU forward. + + B. Measurement — the per-hook round-trip overhead (D2H + serialize + socket + + CPU op + deserialize + H2D) across activation sizes spanning gpt2 -> 7B -> + 70B hidden dims. This is the number the (D) two-tier decision hinges on; the + in-process path adds ~0 (zero-copy GPU reference). + +All numbers are MEASURED on this GPU, not estimated. + +Run: CUDA_VISIBLE_DEVICES=6 PYTHONPATH=src .../hf-serve/bin/python \ + prototypes/mediator-sandbox/phase5_gpu_measure.py +""" +import os +import pickle +import socket +import sys +import time +from types import SimpleNamespace + +import torch + +from nnsight import LanguageModel +from nnsight.intervention.batching import Batcher +from nnsight.intervention.interleaver import Interleaver, Mediator +from nnsight.intervention.transport import ( + SocketHostChannel, + recv_frame, + send_frame, +) + +DEV = "cuda" +PROMPT = "The Eiffel Tower is in the city of" +PROVIDER = "transformer.h.6.output.i0" +LAYER = 6 + + +# --------------------------------------------------------------------------- # +# A. Correctness on GPU — D2H/H2D around the socket worker # +# --------------------------------------------------------------------------- # +def worker_double_proc(sock): + pid = os.fork() + if pid == 0: + med = Mediator(intervention=None, info=SimpleNamespace(frame=None), batch_group=None) + from nnsight.intervention.transport import SocketWorkerChannel + med.channel = SocketWorkerChannel(sock) + med.cross_invoker = False + value = med.request(PROVIDER) # arrives as a CPU tensor (host D2H'd it) + hs = value[0] if isinstance(value, tuple) else value + new = (hs * 2.0,) + tuple(value[1:]) if isinstance(value, tuple) else hs * 2.0 + med.swap(PROVIDER, new) + med.end() + os._exit(0) + return pid + + +def test_gpu_correctness(model, inputs): + block = model._model.transformer.h[LAYER] + + def local_double(m, i, o): + return (o[0] * 2.0,) + tuple(o[1:]) if isinstance(o, tuple) else o * 2.0 + + h = block.register_forward_hook(local_double) + with torch.no_grad(): + ref = model._model(**inputs).logits + h.remove() + + host_sock, worker_sock = socket.socketpair() + pid = worker_double_proc(worker_sock) + worker_sock.close() + interleaver = Interleaver(mediators=[], tracer=None, batcher=Batcher()) + host_med = Mediator(intervention=None, info=SimpleNamespace(frame=None), batch_group=None) + host_med.channel = SocketHostChannel(host_sock) + host_med.interleaver = interleaver + interleaver.mediators = [host_med] + host_med.channel.wait_event() + + def proxy_hook(m, i, o): + # MediatorProxy: D2H -> deliver to the CPU worker -> H2D the result. + dev = o[0].device if isinstance(o, tuple) else o.device + cpu_o = tuple(x.to("cpu") if torch.is_tensor(x) else x for x in o) if isinstance(o, tuple) else o.to("cpu") + new_cpu = host_med.handle(PROVIDER, cpu_o) + if isinstance(new_cpu, tuple): + return tuple(x.to(dev) if torch.is_tensor(x) else x for x in new_cpu) + return new_cpu.to(dev) + + h = block.register_forward_hook(proxy_hook) + with torch.no_grad(): + sock = model._model(**inputs).logits + h.remove() + os.waitpid(pid, 0); host_med.channel.close() + + ok = torch.allclose(ref, sock, atol=1e-2, rtol=0) + print(f"[A gpu ] D2H/H2D socket path matches in-process GPU forward (atol 1e-2): {ok} " + f"| max|Δ|={(ref - sock).abs().max().item():.2e}") + return ok + + +# --------------------------------------------------------------------------- # +# B. Per-hook transport overhead vs activation size # +# --------------------------------------------------------------------------- # +def bench_worker(sock): + """Long-lived CPU worker: recv tensor, ×2, send back, until STOP.""" + pid = os.fork() + if pid == 0: + while True: + obj = recv_frame(sock) + if obj == "STOP": + os._exit(0) + send_frame(sock, obj * 2.0) + return pid + + +def measure(sizes, iters=20, warmup=5): + host_sock, worker_sock = socket.socketpair() + pid = bench_worker(worker_sock) + worker_sock.close() + + # In-process baseline: a thread echoes the GPU tensor BY REFERENCE (no D2H, + # no serialize, no socket) and does the same ×2 — this is what the in-process + # path costs per hook, for a fair "what does isolation add" comparison. + import queue + import threading + q_in, q_out = queue.Queue(), queue.Queue() + + def inproc_echo(): + while True: + x = q_in.get() + if x is None: + return + q_out.put(x * 2.0) + + t = threading.Thread(target=inproc_echo, daemon=True) + t.start() + + print("\n[B measure] per-hook overhead (host has GPU activation -> host has GPU result).") + print(" TOTAL = D2H + sockRTT + H2D (real end-to-end). 'ser' is the pickle portion") + print(" of sockRTT, measured separately for attribution. 'inproc' = in-process baseline.") + print(f" {'shape (b,seq,hid)':>22} {'MB':>7} | {'inproc':>8} {'D2H':>8} {'(ser)':>8} {'sockRTT':>8} " + f"{'H2D':>8} {'TOTAL':>9} (ms)") + rows = [] + for (b, seq, hid) in sizes: + act = torch.randn(b, seq, hid, device=DEV, dtype=torch.bfloat16) + mb = act.element_size() * act.nelement() / 1e6 + + def one(): + torch.cuda.synchronize() + t0 = time.perf_counter() + cpu = act.to("cpu") + torch.cuda.synchronize() + t1 = time.perf_counter() + send_frame(host_sock, cpu) # ONE serialize + socket round-trip (worker deser+op+ser) + res = recv_frame(host_sock) + t2 = time.perf_counter() + gpu = res.to(DEV) # noqa: F841 + torch.cuda.synchronize() + t3 = time.perf_counter() + return (t1 - t0, t2 - t1, t3 - t2, t3 - t0) # d2h, sockRTT, h2d, TOTAL + + def one_inproc(): + torch.cuda.synchronize() + t0 = time.perf_counter() + q_in.put(act) # reference pass — no copy + _ = q_out.get() + torch.cuda.synchronize() + return time.perf_counter() - t0 + + for _ in range(warmup): + one(); one_inproc() + acc = [0.0] * 4 + inproc_acc = 0.0 + ser_acc = 0.0 + for _ in range(iters): + acc = [a + p for a, p in zip(acc, one())] + inproc_acc += one_inproc() + ser_acc += _time_pickle(act.to("cpu")) # attribution-only; outside the TOTAL path + d2h, rtt, h2d, tot = [a / iters * 1e3 for a in acc] + inproc_ms = inproc_acc / iters * 1e3 + ser_ms = ser_acc / iters * 1e3 + rows.append((b, seq, hid, mb, tot, inproc_ms)) + print(f" {str((b, seq, hid)):>22} {mb:7.1f} | {inproc_ms:8.3f} {d2h:8.3f} {ser_ms:8.3f} {rtt:8.3f} " + f"{h2d:8.3f} {tot:9.3f}") + + send_frame(host_sock, "STOP") + os.waitpid(pid, 0); host_sock.close() + q_in.put(None) + return rows + + +def _time_pickle(cpu_tensor): + t0 = time.perf_counter() + pickle.dumps(cpu_tensor, protocol=pickle.HIGHEST_PROTOCOL) + return time.perf_counter() - t0 + + +def main(): + if not torch.cuda.is_available(): + sys.exit("Phase 5 needs a GPU (set CUDA_VISIBLE_DEVICES).") + print(f"GPU: {torch.cuda.get_device_name(0)}") + + model = LanguageModel("gpt2", device_map=DEV, dispatch=True) + inputs = model.tokenizer(PROMPT, return_tensors="pt").to(DEV) + correctness = test_gpu_correctness(model, inputs) + + sizes = [ + (1, 16, 768), # gpt2, short prompt + (1, 512, 768), # gpt2, long context + (1, 512, 4096), # ~7B hidden, 512 tokens + (1, 2048, 4096), # ~7B hidden, 2k context (one layer's activation) + (1, 2048, 8192), # ~70B hidden, 2k context + ] + rows = measure(sizes) + + # The (D) decision input: extra latency a per-layer cache of N layers would add. + big = next(r for r in rows if r[2] == 4096 and r[1] == 2048) + print(f"\n[D-input] a 1×2048×4096 bf16 activation ({big[3]:.1f} MB) costs {big[4]:.2f} ms/hook over the " + f"socket vs {big[5]:.3f} ms in-process (measured, zero-copy ref).\n " + f"A 32-layer cache => ~{big[4]*32:.0f} ms added (linear projection).") + + print("=" * 78) + print(f"PHASE 5 RESULT: {'PASS — GPU path correct; per-hook overhead measured' if correctness else 'FAIL'}") + sys.exit(0 if correctness else 1) + + +if __name__ == "__main__": + main() diff --git a/prototypes/mediator-sandbox/phase5b_transport_breakdown.py b/prototypes/mediator-sandbox/phase5b_transport_breakdown.py new file mode 100644 index 000000000..13ce7ea64 --- /dev/null +++ b/prototypes/mediator-sandbox/phase5b_transport_breakdown.py @@ -0,0 +1,236 @@ +#!/usr/bin/env python3 +"""Phase 5b — WHERE does the per-hook cost come from, and what fixes it. + +Phase 5 showed ~111 ms/hook for a 16.8 MB activation, "dominated by sockRTT". But +a Unix socket moves GB/s, so 16.8 MB should cost a few ms — the 100 ms must be +SERIALIZATION, not transfer. This isolates it: it (1) measures the raw socket +bandwidth, (2) breaks the pickle path into dumps/loads, then compares three +round-trip transports on the SAME CPU tensor: + + A. pickle — the naive path (pickle.dumps/​loads both ways) [what Phase 5 used] + B. raw message — a small struct header + the tensor's raw bytes over the socket, + rebuilt with torch.frombuffer (no pickle) + C. shared mem — host & worker share an mmap; only a 1-byte ready signal crosses + the socket; the bulk never travels (near zero-copy) + +CPU<->CPU only (D2H/H2D is separate and was ~3-4 ms). Worker is a forked process. +Run: PYTHONPATH=src .../hf-serve/bin/python prototypes/mediator-sandbox/phase5b_transport_breakdown.py +""" +import faulthandler +import mmap +import os +import pickle +import socket +import struct +import sys +import time + +import torch + +faulthandler.dump_traceback_later(45, exit=True) + +DT = torch.bfloat16 +SIZES = [(1, 512, 4096), (1, 2048, 4096), (1, 2048, 8192)] # 4.2, 16.8, 33.6 MB +ITERS, WARM = 30, 8 +_LEN = struct.Struct("!I") +_HDR = struct.Struct("!iiqqq") # ndim, dtype_code(0=bf16), d0, d1, d2 (fixed 3-dim here) + + +def _recvn(sock, n): + buf = bytearray(n) + view = memoryview(buf) + got = 0 + while got < n: + c = sock.recv_into(view[got:], n - got) + if c == 0: + raise EOFError + got += c + return buf + + +# ---------- A: pickle ---------- +def send_pickle(sock, obj): + b = pickle.dumps(obj, protocol=pickle.HIGHEST_PROTOCOL) + sock.sendall(_LEN.pack(len(b)) + b) + + +def recv_pickle(sock): + (n,) = _LEN.unpack(_recvn(sock, 4)) + return pickle.loads(_recvn(sock, n)) + + +# ---------- B: raw message ---------- +def send_raw(sock, t): + t = t.contiguous() + raw = t.flatten().view(torch.uint8).numpy() # zero-copy view of the bytes + hdr = _HDR.pack(t.dim(), 0, *(list(t.shape) + [0, 0, 0])[:3]) + sock.sendall(_LEN.pack(raw.nbytes) + hdr) + sock.sendall(memoryview(raw)) # bulk: one zero-copy send + + +def recv_raw(sock): + (n,) = _LEN.unpack(_recvn(sock, 4)) + ndim, _code, d0, d1, d2 = _HDR.unpack(_recvn(sock, _HDR.size)) + buf = _recvn(sock, n) + shape = [d0, d1, d2][:ndim] + return torch.frombuffer(bytearray(buf), dtype=torch.uint8).view(DT).view(*shape) + + +# ---------- C: shared memory ---------- +def make_shm(nbytes): + return mmap.mmap(-1, nbytes) # anonymous MAP_SHARED, inherited across fork + + +def shm_write(buf, t): + t = t.contiguous() + raw = t.flatten().view(torch.uint8).numpy() + buf[: raw.nbytes] = memoryview(raw) + + +def shm_read(buf, nbytes, shape): + return torch.frombuffer(memoryview(buf)[:nbytes], dtype=torch.uint8).view(DT).view(*shape).clone() + + +# --------------------------------------------------------------------------- # +def fork_worker(fn, *a): + parent, child = socket.socketpair() + pid = os.fork() + if pid == 0: + parent.close() + fn(child, *a) + os._exit(0) + child.close() + return pid, parent + + +def bench(label, host_sock, do_round, stop): + for _ in range(WARM): + do_round() + t0 = time.perf_counter() + for _ in range(ITERS): + do_round() + dt = (time.perf_counter() - t0) / ITERS * 1e3 + stop() + return dt + + +def main(): + print(f"{'transport':<14}{'4.2 MB':>12}{'16.8 MB':>12}{'33.6 MB':>12} (ms / round-trip)") + + # raw socket bandwidth + pickle sub-breakdown, for the 16.8 MB case + mid = torch.randn(*SIZES[1], dtype=DT) + mid_bytes = mid.flatten().view(torch.uint8).numpy().nbytes + # one-way raw socket throughput + pid, hs = fork_worker(_drain, mid_bytes) + raw = mid.flatten().view(torch.uint8).numpy() + for _ in range(WARM): + hs.sendall(memoryview(raw)); hs.recv(1) + t0 = time.perf_counter() + for _ in range(ITERS): + hs.sendall(memoryview(raw)); hs.recv(1) + one_way = (time.perf_counter() - t0) / ITERS * 1e3 + hs.sendall(b""); hs.close(); os.waitpid(pid, 0) + gbps = mid_bytes / (one_way / 1e3) / 1e9 + blob = pickle.dumps(mid, protocol=pickle.HIGHEST_PROTOCOL) + dumps_ms = _t(lambda: pickle.dumps(mid, protocol=pickle.HIGHEST_PROTOCOL)) + loads_ms = _t(lambda: pickle.loads(blob)) + op_bf16 = _t(lambda: mid * 2.0) + mid_f32 = mid.float() + op_f32 = _t(lambda: mid_f32 * 2.0) + print(f"\n[probe] 16.8 MB: raw socket one-way+ack {one_way:.2f} ms ({gbps:.1f} GB/s)") + print(f"[probe] 16.8 MB: pickle.dumps {dumps_ms:.2f} ms | pickle.loads {loads_ms:.2f} ms") + print(f"[probe] 16.8 MB: ×2 CPU op — bf16 {op_bf16:.2f} ms | fp32 {op_f32:.2f} ms " + f"(this is the USER op, not transport)\n") + + # three transports across the three sizes — ECHO (no op) = PURE transport cost + for name, runner in [("A pickle", run_pickle), ("B raw-msg", run_raw), ("C shared-mem", run_shm)]: + cells = [] + for sz in SIZES: + cells.append(f"{runner(sz):>12.2f}") + print(f"{name:<14}{''.join(cells)}") + + print("\n(worker does the same ×2 CPU op in every method; only the transport differs.)") + + +def _drain(sock, nbytes): + while True: + try: + b = _recvn(sock, nbytes) + except EOFError: + return + sock.sendall(b"k") + + +def _t(fn, n=20): + for _ in range(5): + fn() + t0 = time.perf_counter() + for _ in range(n): + fn() + return (time.perf_counter() - t0) / n * 1e3 + + +def run_pickle(sz): + def worker(sock): + while True: + try: + t = recv_pickle(sock) + except EOFError: + return + send_pickle(sock, t) # echo = pure transport + pid, hs = fork_worker(worker) + act = torch.randn(*sz, dtype=DT) + + def rnd(): + send_pickle(hs, act); recv_pickle(hs) + dt = bench("A", hs, rnd, lambda: (hs.close(), os.waitpid(pid, 0))) + return dt + + +def run_raw(sz): + def worker(sock): + while True: + try: + t = recv_raw(sock) + except EOFError: + return + send_raw(sock, t) # echo = pure transport + pid, hs = fork_worker(worker) + act = torch.randn(*sz, dtype=DT) + + def rnd(): + send_raw(hs, act); recv_raw(hs) + dt = bench("B", hs, rnd, lambda: (hs.close(), os.waitpid(pid, 0))) + return dt + + +def run_shm(sz): + nbytes = torch.empty(*sz, dtype=DT).flatten().view(torch.uint8).numpy().nbytes + in_buf = make_shm(nbytes) + out_buf = make_shm(nbytes) + + def worker(sock): + while True: + sig = sock.recv(1) + if not sig: + return + t = torch.frombuffer(memoryview(in_buf)[:nbytes], dtype=torch.uint8).view(DT).view(*sz) + res = t.contiguous() # echo = pure transport + out_buf[:nbytes] = memoryview(res.flatten().view(torch.uint8).numpy()) + sock.sendall(b"k") + + pid, hs = fork_worker(worker) + act = torch.randn(*sz, dtype=DT) + + def rnd(): + shm_write(in_buf, act) + hs.sendall(b"g") + hs.recv(1) + shm_read(out_buf, nbytes, sz) + dt = bench("C", hs, rnd, lambda: (hs.close(), os.waitpid(pid, 0))) + in_buf.close(); out_buf.close() + return dt + + +if __name__ == "__main__": + main() diff --git a/prototypes/mediator-sandbox/phase6_jailed_worker.py b/prototypes/mediator-sandbox/phase6_jailed_worker.py new file mode 100644 index 000000000..8f7cae0ed --- /dev/null +++ b/prototypes/mediator-sandbox/phase6_jailed_worker.py @@ -0,0 +1,35 @@ +#!/usr/bin/env python3 +"""Phase 6 — jailed worker using the shared-memory + safetensors channel. + +Same as phase3's worker, but its tensor payloads ride a shared memfd (passed in +via SHM_FD) instead of being pickled over the socket. Runs INSIDE a bwrap jail. +""" +import os +import socket + + +def main(): + fd = int(os.environ["WORKER_FD"]) + shm_fd = int(os.environ["SHM_FD"]) + shm_size = int(os.environ["SHM_SIZE"]) + provider = os.environ["PROVIDER"] + sock = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM, fileno=fd) + + from types import SimpleNamespace + + from nnsight.intervention.interleaver import Mediator + from nnsight.intervention.transport import ShmArena, ShmSocketWorkerChannel + + arena = ShmArena.attach(shm_fd, shm_size) + med = Mediator(intervention=None, info=SimpleNamespace(frame=None), batch_group=None) + med.channel = ShmSocketWorkerChannel(sock, arena) + med.cross_invoker = False + + value = med.request(provider) + new = (value[0] * 2.0,) + tuple(value[1:]) if isinstance(value, tuple) else value * 2.0 + med.swap(provider, new) + med.end() + + +if __name__ == "__main__": + main() diff --git a/prototypes/mediator-sandbox/phase6_shm_safetensors.py b/prototypes/mediator-sandbox/phase6_shm_safetensors.py new file mode 100644 index 000000000..c5e05497e --- /dev/null +++ b/prototypes/mediator-sandbox/phase6_shm_safetensors.py @@ -0,0 +1,203 @@ +#!/usr/bin/env python3 +"""Phase 6 — shared memory + safetensors transport. + +Implements the fast path identified in phase5b and wires it into the real +`MediatorChannel`s (`ShmSocketHostChannel`/`ShmSocketWorkerChannel` in +transport.py). Tensor bulk rides a shared memfd; only a tiny control frame +crosses the socket; tensors are encoded with safetensors (safe, no pickle). + +Three parts: + 1. measure — per-hook echo round-trip: pickle (old) vs shm+safetensors (new). + 2. correct — the real Mediator protocol with the Shm channels on gpt2 → golden. + 3. jailed — same, with the worker in a bwrap jail (memfd passed in via SHM_FD). + +Run: PYTHONPATH=src .../hf-serve/bin/python prototypes/mediator-sandbox/phase6_shm_safetensors.py +""" +import os +import shutil +import socket +import subprocess +import sys +import time +from types import SimpleNamespace + +import torch + +from nnsight import LanguageModel +from nnsight.intervention.batching import Batcher +from nnsight.intervention.interleaver import Interleaver, Mediator +from nnsight.intervention.transport import ( + ShmArena, + ShmSocketHostChannel, + ShmSocketWorkerChannel, + recv_frame, + recv_shm, + send_frame, + send_shm, +) + +HERE = os.path.dirname(os.path.abspath(__file__)) +ROOT = os.path.dirname(os.path.dirname(HERE)) +SRC = os.path.join(ROOT, "src") +WORKER = os.path.join(HERE, "phase6_jailed_worker.py") +ENV_ROOT = os.path.dirname(os.path.dirname(sys.executable)) +PROMPT = "The Eiffel Tower is in the city of" +PROVIDER = "transformer.h.6.output.i0" +LAYER = 6 +ARENA = 64 << 20 # 64 MB shared region + + +# --------------------------------------------------------------------------- # +# 1. measurement: pickle vs shm+safetensors (echo = pure transport) # +# --------------------------------------------------------------------------- # +def measure(): + sizes = [(1, 512, 4096), (1, 2048, 4096), (1, 2048, 8192)] # 4.2, 16.8, 33.6 MB + print(f"{'transport':<22}{'4.2 MB':>11}{'16.8 MB':>11}{'33.6 MB':>11} (ms/round-trip)") + + def echo_pickle(sock): + while True: + try: + t = recv_frame(sock) + except EOFError: + os._exit(0) + send_frame(sock, t) + + def echo_shm(sock, fd): + arena = ShmArena.attach(fd, ARENA) + while True: + try: + t = recv_shm(sock, arena) + except EOFError: + os._exit(0) + send_shm(sock, arena, t) + + for name, use_shm in [("A pickle (old)", False), ("D shm+safetensors", True)]: + cells = [] + for sz in sizes: + act = (torch.randn(*sz, dtype=torch.bfloat16),) + hs, ws = socket.socketpair() + arena = ShmArena(ARENA) if use_shm else None + pid = os.fork() + if pid == 0: + hs.close() + (echo_shm(ws, arena.fd) if use_shm else echo_pickle(ws)) + os._exit(0) + ws.close() + rnd = ((lambda: (send_shm(hs, arena, act), recv_shm(hs, arena))) + if use_shm else (lambda: (send_frame(hs, act), recv_frame(hs)))) + for _ in range(8): + rnd() + t0 = time.perf_counter() + for _ in range(30): + rnd() + cells.append(f"{(time.perf_counter() - t0) / 30 * 1e3:>11.2f}") + hs.close(); os.waitpid(pid, 0) + if arena: + arena.close() + print(f"{name:<22}{''.join(cells)}") + + +# --------------------------------------------------------------------------- # +# 2/3. correctness via the real Mediator protocol with the Shm channels # +# --------------------------------------------------------------------------- # +def reference(model, inputs): + blk = model._model.transformer.h[LAYER] + h = blk.register_forward_hook( + lambda m, i, o: (o[0] * 2.0,) + tuple(o[1:]) if isinstance(o, tuple) else o * 2.0) + with torch.no_grad(): + out = model._model(**inputs).logits + h.remove() + return out + + +def drive_host(model, inputs, host_med): + blk = model._model.transformer.h[LAYER] + h = blk.register_forward_hook(lambda m, i, o: host_med.handle(PROVIDER, o)) + try: + with torch.no_grad(): + return model._model(**inputs).logits + finally: + h.remove() + + +def mk_host(sock, arena): + interleaver = Interleaver(mediators=[], tracer=None, batcher=Batcher()) + med = Mediator(intervention=None, info=SimpleNamespace(frame=None), batch_group=None) + med.channel = ShmSocketHostChannel(sock, arena) + med.interleaver = interleaver + interleaver.mediators = [med] + med.channel.wait_event() + return med + + +def correct_fork(model, inputs, ref): + arena = ShmArena(ARENA) + hs, ws = socket.socketpair() + pid = os.fork() + if pid == 0: + hs.close() + wa = ShmArena.attach(arena.fd, ARENA) + med = Mediator(intervention=None, info=SimpleNamespace(frame=None), batch_group=None) + med.channel = ShmSocketWorkerChannel(ws, wa) + med.cross_invoker = False + v = med.request(PROVIDER) + new = (v[0] * 2.0,) + tuple(v[1:]) if isinstance(v, tuple) else v * 2.0 + med.swap(PROVIDER, new) + med.end() + os._exit(0) + ws.close() + host = mk_host(hs, arena) + out = drive_host(model, inputs, host) + os.waitpid(pid, 0); host.channel.close(); arena.close() + ok = torch.allclose(ref, out, atol=1e-2, rtol=0) + print(f"[2 correct/fork ] gpt2 golden via Shm channel: {ok} | max|Δ|={(ref - out).abs().max():.2e}") + return ok + + +def correct_jailed(model, inputs, ref): + if shutil.which("bwrap") is None: + print("[3 correct/jail ] bwrap absent — skipped") + return True + arena = ShmArena(ARENA) + hs, ws = socket.socketpair() + os.set_inheritable(ws.fileno(), True) + os.set_inheritable(arena.fd, True) + env = { + "WORKER_FD": str(ws.fileno()), "SHM_FD": str(arena.fd), "SHM_SIZE": str(ARENA), + "PROVIDER": PROVIDER, "PYTHONPATH": SRC, "PATH": "/usr/local/bin:/usr/bin:/bin", + "HOME": "/tmp", "CUDA_VISIBLE_DEVICES": "", "HF_HUB_OFFLINE": "1", "TRANSFORMERS_OFFLINE": "1", + } + cmd = [ + "bwrap", "--unshare-all", + "--ro-bind", "/usr", "/usr", "--ro-bind", "/lib", "/lib", "--ro-bind", "/lib64", "/lib64", + "--ro-bind", "/bin", "/bin", "--ro-bind", "/etc", "/etc", + "--ro-bind", ENV_ROOT, ENV_ROOT, "--ro-bind", SRC, SRC, "--ro-bind", WORKER, WORKER, + "--proc", "/proc", "--dev", "/dev", "--tmpfs", "/tmp", "--die-with-parent", + sys.executable, WORKER, + ] + p = subprocess.Popen(cmd, pass_fds=[ws.fileno(), arena.fd], env=env) + ws.close() + host = mk_host(hs, arena) + out = drive_host(model, inputs, host) + p.wait(); host.channel.close(); arena.close() + ok = torch.allclose(ref, out, atol=1e-2, rtol=0) + print(f"[3 correct/jail ] gpt2 golden, worker JAILED, shm via passed memfd: {ok} " + f"| max|Δ|={(ref - out).abs().max():.2e}") + return ok + + +def main(): + measure() + print() + model = LanguageModel("gpt2", device_map="cpu", dispatch=True) + inputs = model.tokenizer(PROMPT, return_tensors="pt") + ref = reference(model, inputs) + ok = correct_fork(model, inputs, ref) + ok &= correct_jailed(model, inputs, ref) + print("=" * 72) + print(f"PHASE 6 RESULT: {'PASS — shm+safetensors channel correct (fork + jailed)' if ok else 'FAIL'}") + sys.exit(0 if ok else 1) + + +if __name__ == "__main__": + main() diff --git a/src/nnsight/intervention/_sandbox.py b/src/nnsight/intervention/_sandbox.py new file mode 100644 index 000000000..e92ecd567 --- /dev/null +++ b/src/nnsight/intervention/_sandbox.py @@ -0,0 +1,74 @@ +"""Footgun-containment seccomp filter for the isolated GPU worker (x86-64). + +Per the agreed threat model (contain *mistakes*, not a determined adversary): after +CUDA + torch are warmed and the intervention is deserialized, the worker calls +:func:`lock_down`, which installs a minimal seccomp-BPF filter making new +``open``/``openat`` (filesystem), ``socket``/``connect`` (network), and ``execve`` +(spawn) syscalls fail with EPERM. CUDA keeps working because it talks to the +already-open ``/dev/nvidia*`` fds via ioctl/mmap, and the control Pipe + CUDA-IPC +buffer use already-open fds (read/write/ioctl), not new opens. + +No external deps: the BPF program is assembled by hand and installed via +``prctl(PR_SET_SECCOMP, SECCOMP_MODE_FILTER, ...)``. Proven in +prototypes/mediator-sandbox/gpu_sandbox (test_safety: 9/9 contained). +""" +import ctypes +import struct + +# x86-64 syscall numbers +_NR = { + "open": 2, "openat": 257, "openat2": 437, + "socket": 41, "connect": 42, "execve": 59, "execveat": 322, +} +_AUDIT_ARCH_X86_64 = 0xC000003E +_RET_KILL_PROCESS = 0x80000000 +_RET_ERRNO = 0x00050000 +_RET_ALLOW = 0x7FFF0000 +_EPERM = 1 +# BPF opcodes +_LD_W_ABS = 0x20 +_JMP_JEQ_K = 0x15 +_RET_K = 0x06 +_PR_SET_NO_NEW_PRIVS = 38 +_PR_SET_SECCOMP = 22 +_SECCOMP_MODE_FILTER = 2 + + +def _build_filter(blocked): + instrs = [ + (_LD_W_ABS, 0, 0, 4), # A = arch (seccomp_data offset 4) + (_JMP_JEQ_K, 1, 0, _AUDIT_ARCH_X86_64), # if x86-64: skip the kill + (_RET_K, 0, 0, _RET_KILL_PROCESS), # else kill (block arch-bypass) + (_LD_W_ABS, 0, 0, 0), # A = syscall nr (offset 0) + ] + for nr in blocked: + instrs.append((_JMP_JEQ_K, 0, 1, nr)) # if A == nr: next else skip next + instrs.append((_RET_K, 0, 0, _RET_ERRNO | _EPERM)) + instrs.append((_RET_K, 0, 0, _RET_ALLOW)) # default: allow + return instrs + + +class _sock_fprog(ctypes.Structure): + _fields_ = [("len", ctypes.c_ushort), ("filter", ctypes.c_void_p)] + + +def lock_down(block_fs: bool = True, block_net: bool = True) -> None: + """Install the seccomp filter. Call AFTER torch/CUDA are warmed + the + intervention is deserialized (deserialization may open files).""" + blocked = [] + if block_fs: + blocked += [_NR["open"], _NR["openat"], _NR["openat2"]] + if block_net: + blocked += [_NR["socket"], _NR["connect"]] + blocked += [_NR["execve"], _NR["execveat"]] # no spawning new programs either + instrs = _build_filter(blocked) + + libc = ctypes.CDLL("libc.so.6", use_errno=True) + if libc.prctl(_PR_SET_NO_NEW_PRIVS, 1, 0, 0, 0) != 0: + raise OSError(ctypes.get_errno(), "PR_SET_NO_NEW_PRIVS failed") + + prog = b"".join(struct.pack("HBBI", *i) for i in instrs) + buf = ctypes.create_string_buffer(prog, len(prog)) + fprog = _sock_fprog(len(instrs), ctypes.cast(buf, ctypes.c_void_p)) + if libc.prctl(_PR_SET_SECCOMP, _SECCOMP_MODE_FILTER, ctypes.byref(fprog), 0, 0) != 0: + raise OSError(ctypes.get_errno(), "PR_SET_SECCOMP failed") diff --git a/src/nnsight/intervention/hooks.py b/src/nnsight/intervention/hooks.py index 9ecafca83..72b56ec5b 100644 --- a/src/nnsight/intervention/hooks.py +++ b/src/nnsight/intervention/hooks.py @@ -151,7 +151,12 @@ def add_ordered_hook(module: torch.nn.Module, hook: Callable, type: str) -> Any: return handle -def input_hook(mediator: Mediator, module: torch.nn.Module, path: str) -> Any: +def input_hook( + mediator: Mediator, + module: torch.nn.Module, + path: str, + iteration: Optional[int] = None, +) -> Any: """Register a one-shot forward pre-hook for a mediator on a module. The hook target iteration is captured at registration time: @@ -188,11 +193,12 @@ def input_hook(mediator: Mediator, module: torch.nn.Module, path: str) -> Any: """ handle = None - iteration = ( - mediator.iteration - if mediator.iteration is not None - else mediator.iteration_tracker[path] - ) + if iteration is None: + iteration = ( + mediator.iteration + if mediator.iteration is not None + else mediator.iteration_tracker[path] + ) def hook(module: torch.nn.Module, args: Any, kwargs: Any) -> Any: @@ -221,7 +227,12 @@ def hook(module: torch.nn.Module, args: Any, kwargs: Any) -> Any: return handle -def output_hook(mediator: Mediator, module: torch.nn.Module, path: str) -> Any: +def output_hook( + mediator: Mediator, + module: torch.nn.Module, + path: str, + iteration: Optional[int] = None, +) -> Any: """Register a one-shot forward hook for a mediator on a module. Behaves identically to :func:`input_hook` but intercepts the module's @@ -233,17 +244,22 @@ def output_hook(mediator: Mediator, module: torch.nn.Module, path: str) -> Any: mediator: The mediator requesting this hook. module: The PyTorch module to hook. path: The provider path prefix (e.g. ``"model.layer.0.output"``). + iteration: Explicit target step. ``None`` (default) resolves it from + ``mediator.iteration``/``iteration_tracker`` (in-process). The isolated + path (host-side hook registration) passes the step parsed from the worker's requester + string, since the worker's tracker lives in another process. Returns: A :class:`~torch.utils.hooks.RemovableHandle` for the registered hook. """ handle = None - iteration = ( - mediator.iteration - if mediator.iteration is not None - else mediator.iteration_tracker[path] - ) + if iteration is None: + iteration = ( + mediator.iteration + if mediator.iteration is not None + else mediator.iteration_tracker[path] + ) def hook(module: torch.nn.Module, _, output: Any) -> Any: diff --git a/src/nnsight/intervention/interleaver.py b/src/nnsight/intervention/interleaver.py index e80ecfb43..ec56c217b 100755 --- a/src/nnsight/intervention/interleaver.py +++ b/src/nnsight/intervention/interleaver.py @@ -463,6 +463,16 @@ def initialize( self.current: Mediator = None + # Isolated BARRIER: workers can't count across processes, so each sends the + # target count and the host accumulates participant names here, coordinating + # when all have arrived. Unused on the in-process path. + self._barrier_acc: set = set() + + # Isolated cross_invoker: per-worker frames aren't shared across processes, so + # workers push their locals here and pull the merged store back (host-mediated + # replacement for the shared-frame push/pull). Unused on the in-process path. + self._xinvoke_store: dict = {} + def cancel(self): """Cancel all mediators / intervention threads. @@ -712,7 +722,7 @@ def check_dangling_mediators(self): for mediator in self.mediators: if mediator.alive: - requested_event, requester = mediator.event_queue.get() + requested_event, requester = mediator.channel.get_event() if isinstance(requester, tuple): requester = requester[0] @@ -771,6 +781,91 @@ def __deepcopy__(self, memo): return self +class MediatorChannel: + """Bidirectional, one-event-in-flight handoff between the worker thread (the + intervention fn) and the main thread (the model forward pass). + + The six :class:`Events` (VALUE/SWAP/SKIP/BARRIER/END/EXCEPTION) ride on top of + this channel unchanged. Today the only implementation is in-process + (:class:`InProcessChannel`, two lock-based one-slot queues); the mediator + isolation refactor adds a socket-backed subclass so the worker can run in a + separate sandbox process. See docs/developing/mediator-isolation-harness-plan.md. + + Direction conventions: + - ``*_event`` methods carry worker -> main messages (an ``(Events, requester)`` pair). + - ``*_response`` methods carry main -> worker replies (the value for a pending event). + + Only one event is ever in flight, which is what enforces the "access modules in + forward-pass order" contract. + """ + + # --- worker -> main (event direction) --- + def put_event(self, item: Any) -> None: + raise NotImplementedError + + def restore_event(self, item: Any) -> None: + """Stage an event back without releasing a waiter (used when a provider + fires that the worker is not currently requesting).""" + raise NotImplementedError + + def get_event(self) -> Any: + raise NotImplementedError + + def wait_event(self) -> None: + raise NotImplementedError + + @property + def has_event(self) -> bool: + raise NotImplementedError + + # --- main -> worker (response direction) --- + def put_response(self, value: Any) -> None: + raise NotImplementedError + + def get_response(self) -> Any: + raise NotImplementedError + + def wait_response(self) -> None: + raise NotImplementedError + + +class InProcessChannel(MediatorChannel): + """The original transport: two :class:`Mediator.Value` one-slot lock queues. + + Behaviour-identical to the pre-seam ``event_queue``/``response_queue`` pair — + this is the default channel and the regression baseline for the seam. + """ + + def __init__(self) -> None: + self._event = Mediator.Value() + self._response = Mediator.Value() + + def put_event(self, item: Any) -> None: + self._event.put(item) + + def restore_event(self, item: Any) -> None: + self._event.restore(item) + + def get_event(self) -> Any: + return self._event.get() + + def wait_event(self) -> None: + self._event.wait() + + @property + def has_event(self) -> bool: + return self._event.has_value + + def put_response(self, value: Any) -> None: + self._response.put(value) + + def get_response(self) -> Any: + return self._response.get() + + def wait_response(self) -> None: + self._response.wait() + + class Mediator: """ Mediates between the model execution and a single intervention function. @@ -785,8 +880,7 @@ class Mediator: info (Tracer.Info): Information about the tracing context associated with this mediator name (Optional[str]): Optional name for the mediator batch_group (Optional[List[int]]): Optional batch group for the mediator to determine which slice of tensors are being intervened on - event_queue (SimpleQueue): Where the mediator (worker thread) puts events to be processed by the interleaver (main thread). Will only ever have 1 or 0 items in the queue. - response_queue (SimpleQueue): Where the interleaver (main thread) puts responses to events, to then be processed by the mediator (worker thread). Will only ever have 1 or 0 items in the queue. + channel (MediatorChannel): The bidirectional one-event-in-flight handoff between the mediator (worker thread) and the interleaver (main thread). The worker puts events for the interleaver to process and the interleaver puts responses back; only ever 1 or 0 messages in flight in each direction. ``InProcessChannel`` by default; swappable for a socket-backed channel to run the worker in an isolated process. worker (Thread): The thread that runs the intervention function history (Set[str]): A set of providers that have been seen by the mediator. Used to detect out of order interventions. iteration_tracker (Dict[str, int]): Per-provider-path counter maintained by @@ -875,11 +969,18 @@ def __init__( self.interleaver = None - self.event_queue = Mediator.Value() - self.response_queue = Mediator.Value() + self.channel: MediatorChannel = InProcessChannel() self.worker = None + # Host-side handle to an isolated worker process (set by + # isolation.spawn_isolated_worker). ``None`` => in-process (default). + self._iso = None + + # True only inside an isolated worker process (set by _worker_main). Lets + # cross-process-aware logic (e.g. Barrier) know it can't count locally. + self._isolated_worker = False + self.skip_container = None self.history = set() @@ -969,24 +1070,35 @@ def start(self, interleaver: Interleaver): else: _caller_stream = None - _intervention = self.intervention - _args = (self, self.info, *self.args) + from .isolation import isolation_state - def _worker_target(): - if _caller_stream is not None: - torch.cuda.set_stream(_caller_stream) - _intervention(*_args) + if isolation_state()["on"]: + # Isolated path: run the intervention in a spawned GPU worker process. + # Sets self.channel (host end), self.worker (the process), self._iso. + from .isolation import spawn_isolated_worker - # Start the worker thread. - self.worker = Thread( - target=_worker_target, - daemon=True, - name=self.name, - ) + spawn_isolated_worker(self) + self.interleaver.current = self + else: + _intervention = self.intervention + _args = (self, self.info, *self.args) - self.interleaver.current = self - self.worker.start() - self.event_queue.wait() + def _worker_target(): + if _caller_stream is not None: + torch.cuda.set_stream(_caller_stream) + _intervention(*_args) + + # Start the worker thread. + self.worker = Thread( + target=_worker_target, + daemon=True, + name=self.name, + ) + + self.interleaver.current = self + self.worker.start() + + self.channel.wait_event() # Handle the first event for each mediator to clear mediators that already ended. try: @@ -1006,12 +1118,17 @@ def cancel(self): self.iteration = 0 self.worker = None - if self.event_queue.has_value: + if self.channel.has_event: self.handle() - if self.event_queue.has_value: - self.event_queue.get() - self.response_queue.put(Cancelation()) - self.event_queue.get() + if self.channel.has_event: + self.channel.get_event() + self.channel.put_response(Cancelation()) + self.channel.get_event() + + # Tear down the isolated worker process + free the bounce buffer. + if self._iso is not None: + self._iso.close() + self._iso = None def handle(self, provider: Optional[str] = None, value: Optional[Any] = None): """Process a provided value against this mediator's pending event. @@ -1048,13 +1165,26 @@ def handle(self, provider: Optional[str] = None, value: Optional[Any] = None): self.interleaver.batcher.current_provider = provider # Check to see if this mediator has an unprocessed eventto start. - process = self.event_queue.has_value + process = self.channel.has_event # Continue processing events until there are no more events to process. # Means we can move on to the next mediator and continue the model execution. while process: - event, data = self.event_queue.get() + event, data = self.channel.get_event() + + # Host-side hook registration: in the isolated path the worker has no real module, so the + # host registers the one-shot hook on demand from the requester string + # the worker just sent. + if self._iso is not None and event in ( + Events.VALUE, + Events.SWAP, + Events.SKIP, + ): + from .isolation import ensure_isolated_provider + + requester = data if event == Events.VALUE else data[0] + ensure_isolated_provider(self, requester) if event == Events.VALUE: process = self.handle_value_event(data, provider) @@ -1067,7 +1197,7 @@ def handle(self, provider: Optional[str] = None, value: Optional[Any] = None): elif event == Events.BARRIER: process = self.handle_barrier_event(provider, data) elif event == Events.END: - process = self.handle_end_event() + process = self.handle_end_event(data) value = self.interleaver.batcher.current_value @@ -1120,7 +1250,7 @@ def handle_value_event(self, requester: Any, provider: Any) -> bool: else: # If the requester has not been seen before, add it to the history and put the value event back in the event queue to be processed later. self.history.add(provider) - self.event_queue.restore((Events.VALUE, requester)) + self.channel.restore_event((Events.VALUE, requester)) return False @@ -1158,7 +1288,7 @@ def handle_swap_event(self, provider: Any, requester: Any, swap_value: Any): else: # If the requester has not been seen before, add it to the history and put the swap event back in the event queue to be processed later. self.history.add(provider) - self.event_queue.restore((Events.SWAP, (requester, swap_value))) + self.channel.restore_event((Events.SWAP, (requester, swap_value))) return False @@ -1224,6 +1354,19 @@ def handle_barrier_event(self, provider: Any, participants: Set[str]): carries the swap forward to the outer handle context. """ + # Isolated path: the worker can't count participants across processes, so it + # sends the TARGET count (an int). Accumulate names host-side; only coordinate + # once all participants have arrived. Until then this mediator stays blocked at + # the barrier (return False without responding). + if isinstance(participants, int): + n = participants + acc = self.interleaver._barrier_acc + acc.add(self.name) + if len(acc) < n: + return False + participants = set(acc) + acc.clear() + if participants is not None: prev_current = self.interleaver.current @@ -1246,10 +1389,28 @@ def handle_barrier_event(self, provider: Any, participants: Set[str]): return False - def handle_end_event(self): + def handle_end_event(self, saves: Optional[Any] = None): """ Handle an end event by stopping the mediator. + + Worker→host saves transmission: in the isolated path the worker bundles its ``.save()``'d values + (already filtered by ``Globals.saves``) into the END event, since the + worker's frame + ``Globals.saves`` live in another process. The host + injects them directly into the real **user** frame — the tracer's + ``info.frame`` (where the in-process two-hop push ultimately lands saved + vars). The worker already applied the ``Globals.saves`` filter, so no + second filtering hop is needed here. """ + if self._iso is not None and saves: + tracer = self.interleaver.tracer + user_frame = ( + tracer.info.frame + if tracer is not None and tracer.info.frame is not None + else self.info.frame + ) + if user_frame is not None: + push_variables(user_frame, saves) + self.cancel() return False @@ -1292,7 +1453,7 @@ def handle_skip_event(self, provider: Any, requester: Any, value: Any): return True else: self.history.add(provider) - self.event_queue.restore((Events.SKIP, (requester, value))) + self.channel.restore_event((Events.SKIP, (requester, value))) return False @@ -1305,8 +1466,8 @@ def respond(self, value: Optional[Any] = None): """ # Respond and resume the mediator thread. - self.response_queue.put(value) - self.event_queue.wait() + self.channel.put_response(value) + self.channel.wait_event() ### Requester Methods ### @@ -1328,11 +1489,11 @@ def send(self, event: Events, requester: Any): self.push() # Send the event - self.event_queue.put((event, requester)) + self.channel.put_event((event, requester)) # Wait for the interleaver to process the event and respond with the value. - self.response_queue.wait() - response = self.response_queue.get() + self.channel.wait_response() + response = self.channel.get_response() # If the response is an exception, raise it. if isinstance(response, Exception): @@ -1383,7 +1544,7 @@ def end(self): self.push() - self.event_queue.put((Events.END, None)) + self.channel.put_event((Events.END, None)) def exception(self, exception: Exception): """ @@ -1392,7 +1553,7 @@ def exception(self, exception: Exception): Args: exception: The exception that occurred """ - self.event_queue.put((Events.EXCEPTION, exception)) + self.channel.put_event((Events.EXCEPTION, exception)) @property def frame(self) -> FrameType: @@ -1494,10 +1655,11 @@ def __setstate__(self, state): self.intervention = state["intervention"] self.all_stop = state["all_stop"] self.iteration_tracker = state["iteration_tracker"] - self.event_queue = Mediator.Value() - self.response_queue = Mediator.Value() + self.channel: MediatorChannel = InProcessChannel() self.worker = None + self._iso = None + self._isolated_worker = False self.interleaver = None self.history = set() self.user_cache: "Cache" = list() diff --git a/src/nnsight/intervention/isolation.py b/src/nnsight/intervention/isolation.py new file mode 100644 index 000000000..b26fcdb5d --- /dev/null +++ b/src/nnsight/intervention/isolation.py @@ -0,0 +1,443 @@ +"""Transparent isolated execution of mediators in a spawned GPU worker process. + +This is the *outer harness* for the chosen GPU-sandbox design (see +docs/developing/mediator-gpu-trace-integration.md). It does NOT change the +six-event Mediator protocol — it runs the Mediator's intervention in a separate, +GPU-enabled process and routes the protocol over a :class:`CudaIpcChannel`. + +Two shared-memory assumptions of the in-process path become explicit harness steps: + +- **Host-side hook registration**: the worker has no real module, so when the host + receives a ``VALUE``/``SWAP``/``SKIP`` for a requester it hasn't seen, it registers + the matching one-shot hook on the *real* module (resolved from the requester string). +- **Worker→host saves transmission**: ``.save()`` values live in the worker's frame + + ``Globals.saves``; the worker bundles them into the ``END`` event and the host + ``push_variables`` them into the real user frame. + +Public surface: :func:`isolate_mediators` (context manager) + :func:`isolation_state`. +``Mediator.start`` calls :func:`spawn_isolated_worker` when isolation is on; the host +``handle`` loop calls :func:`ensure_isolated_provider` (host-side hook registration). +""" +from __future__ import annotations + +import os + +import torch +import torch.multiprocessing as mp +import torch.nn as nn +from contextlib import contextmanager +from typing import Any, Dict, Optional + +from . import serialization +from .transport import CudaIpcHostChannel, CudaIpcWorkerChannel +from ..util import apply + +# Types that cross_invoker may ship between workers (data, not framework objects). +_XINVOKE_SCALARS = (int, float, complex, bool, str, bytes, type(None)) + + +def _transmittable(v) -> bool: + """True if ``v`` is cross_invoker-shareable data: a tensor, a basic scalar, or a + container recursively of those. Framework objects (Barrier/Envoy/model) are not — + the worker already has them via its own closure/deserialization.""" + if torch.is_tensor(v) or isinstance(v, _XINVOKE_SCALARS): + return True + # Note: no `set` — util.apply can't walk sets, so a set holding a CUDA tensor + # would skip the D2H move and crash on the host's IPC re-share. Lists/tuples ok. + if isinstance(v, (list, tuple)): + return all(_transmittable(x) for x in v) + if isinstance(v, dict): + return all(_transmittable(k) and _transmittable(x) for k, x in v.items()) + return False + +# --------------------------------------------------------------------------- # +# Opt-in surface # +# --------------------------------------------------------------------------- # +_STATE: Dict[str, Any] = { + "on": False, + "arena_bytes": 64 << 20, + "gpu_mem_fraction": 0.3, + "device": "cuda", + "timeout": 60.0, # per-step wall-clock cap on user code (hang containment) + "lockdown": False, # functional-first; seccomp lockdown enabled separately +} + + +def isolation_state() -> Dict[str, Any]: + return _STATE + + +@contextmanager +def isolate_mediators( + arena_bytes: int = 64 << 20, + gpu_mem_fraction: float = 0.3, + device: str = "cuda", + timeout: float = 60.0, + lockdown: bool = False, +): + """Run interventions inside ``with model.trace(...)`` in an isolated GPU worker. + + Footguns in user intervention code (infinite loops, OOM allocations, device-side + asserts, host-object pokes) are contained to the worker; the model server keeps + serving. + + Args: + timeout: per-step wall-clock cap on user code; a worker that produces no + event within ``timeout`` is presumed hung and killed (the host survives). + """ + prev = dict(_STATE) + _STATE.update( + on=True, + arena_bytes=arena_bytes, + gpu_mem_fraction=gpu_mem_fraction, + device=device, + timeout=timeout, + lockdown=lockdown, + ) + try: + yield + finally: + _STATE.update(prev) + + +# --------------------------------------------------------------------------- # +# Host side # +# --------------------------------------------------------------------------- # +class _IsoHandle: + """Per-mediator host-side handle to its worker process + bounce buffer.""" + + def __init__(self, proc, buf, conn): + self.proc = proc + self.buf = buf + self.conn = conn + self.registered: set = set() # Host-side hook registration: requesters whose hook is registered + self.path2envoy: Optional[dict] = None + + def close(self): + try: + self.conn.send("stop") + except Exception: # noqa: BLE001 + pass + self.proc.join(timeout=5) + if self.proc.is_alive(): + self.proc.terminate() # SIGTERM + self.proc.join(timeout=5) + if self.proc.is_alive(): + self.proc.kill() # SIGKILL — for a worker wedged in a non-interruptible + self.proc.join(timeout=5) # CUDA/C call that ignored SIGTERM + + +def spawn_isolated_worker(mediator) -> None: + """Serialize ``mediator``, spawn its GPU worker, and wire the host channel. + + Sets ``mediator.channel`` (host end), ``mediator.worker`` (the process, so + ``alive`` is True), and ``mediator._iso`` (the handle used by host-side hook registration + cancel). + """ + opts = _STATE + model = mediator.interleaver.tracer.model + # Module:* and Interleaver are synthesized on the worker; ship the rest + # (Tokenizer/Processor) so the deserialized graph resolves them. Only + # remoteable models (LanguageModel/VLM/...) carry those extras; a plain + # NNsight(module) has none. + from ..modeling.mixins.remoteable import RemoteableMixin + + extras = {} + if isinstance(model, RemoteableMixin): + real_map = model._remoteable_persistent_objects() + extras = { + k: v + for k, v in real_map.items() + if not k.startswith("Module:") and k != "Interleaver" + } + + # The tracer attaches source during ITS __getstate__; per-mediator + # serialization must do the same first (else source is unavailable). + mediator.intervention.__source__ = "".join(mediator.info.source) + payload = serialization.dumps(mediator) + + # Per-spawn options: base config + the host interleaver's default_all + # (= generate's max_new_tokens), which an open-ended `tracer.iter[:]` needs + # to bound its step loop on the worker side. + from .. import CONFIG + + worker_opts = dict(opts) + worker_opts["default_all"] = mediator.interleaver.default_all + # cross_invoker matches the in-process gate (Mediator.start): multiple invokes + # + config enabled. The worker can't share a frame, so it pushes/pulls through + # the host store (see _worker_main + on_push/meta below). + worker_opts["cross_invoker"] = ( + len(mediator.interleaver.mediators) > 1 and CONFIG.APP.CROSS_INVOKER + ) + + ctx = mp.get_context("spawn") # CUDA requires spawn, not fork + buf = torch.empty(opts["arena_bytes"], dtype=torch.uint8, device=opts["device"]) + parent_conn, child_conn = ctx.Pipe() + proc = ctx.Process( + target=_worker_main, + args=(payload, extras, child_conn, buf, worker_opts), + daemon=True, + ) + proc.start() + + chan = CudaIpcHostChannel(parent_conn, buf, timeout=opts["timeout"]) + # default_all is set by generate() AFTER the worker spawns (LanguageModel + # ._execute), so a spawn-time snapshot is stale. Piggyback the LIVE value + + # the cross_invoker var store on each response; the worker reads default_all + # before bounding its `iter[:]` loop and pulls the store before each access. + chan.meta_provider = lambda: { + "default_all": mediator.interleaver.default_all, + "xinvoke_store": mediator.interleaver._xinvoke_store, + } + # Merge a worker's pushed cross_invoker locals into the shared host store. + chan.on_push = mediator.interleaver._xinvoke_store.update + mediator.channel = chan + mediator.worker = proc + mediator._iso = _IsoHandle(proc, buf, parent_conn) + + # Host-side iteration tracking. The worker's `tracer.iter[...]` loop runs on + # DUMMY modules, so its iteration_tracker never advances; the HOST must bump + # ITS tracker per forward pass so the host-side hook registration's per-step hooks fire on the right + # generation step (multi-token). Harmless for single-forward traces (one bump). + # Registered on mediator.hooks => torn down by cancel/remove_hooks. + from .tracing.iterator import register_iter_hooks + + register_iter_hooks(mediator, model) + + +def ensure_isolated_provider(mediator, requester: str) -> None: + """Host-side hook registration: register the one-shot hook for ``requester`` on the *real* module. + + Parses ``"..i"``, resolves the envoy on the host, and + registers the existing ``output_hook``/``input_hook`` with the host mediator for + the **specific step N** parsed from the requester — the worker's iteration counter + lives in another process, so the step must come from the wire, not from the host + mediator's ``iteration``. The host's own iteration_tracker (bumped per forward by + the iter hooks installed in ``spawn_isolated_worker``) advances so the N-hook fires + on the right generation step. Idempotent per requester (per step). + """ + iso = mediator._iso + if requester in iso.registered: + return + iso.registered.add(requester) + + # The iteration suffix is always the rightmost ".i". + parts = requester.rsplit(".i", 1) # ["model.transformer.h.6.output", "2"] + base = parts[0] + iteration = int(parts[1]) if len(parts) == 2 and parts[1].isdigit() else None + path, _, kind = base.rpartition(".") # ("model.transformer.h.6", ".", "output") + if kind not in ("output", "input"): + return # externally-provided eproperties (e.g. .result) need no module hook + + if iso.path2envoy is None: + model = mediator.interleaver.tracer.model + iso.path2envoy = {e.path: e for e in model.modules()} + envoy = iso.path2envoy.get(path) + if envoy is None: + return # unknown path → let normal missed-provider handling surface it + + from .hooks import input_hook, output_hook + + if kind == "output": + output_hook(mediator, envoy._module, base, iteration=iteration) + else: + input_hook(mediator, envoy._module, base, iteration=iteration) + + +# --------------------------------------------------------------------------- # +# Worker side # +# --------------------------------------------------------------------------- # +class _WorkerBatcher: + """Minimal stand-in for the Batcher that ``requires_*`` reads on the worker.""" + + current_provider = None + current_value = None + + +class _WorkerInterleaver: + """Worker-side interleaver stub: enough for eproperty + requires_* to run. + + The real Batcher/narrow/swap all live on the host; the worker only builds + requester strings and emits events. + """ + + def __init__(self, default_all=None): + self.interleaving = True + self.batcher = _WorkerBatcher() + self.current = None + # Open-ended `tracer.iter[:]` reads this to know how many generation steps + # to run; set by generate(max_new_tokens=N) on the host and shipped over. + self.default_all = default_all + + def iterate_requester(self, requester: str) -> str: + med = self.current + iteration = ( + med.iteration if med.iteration is not None else med.iteration_tracker[requester] + ) + return f"{requester}.i{iteration}" + + +class _WorkerPersistent: + """persistent_objects map for the worker: synthesize a dummy module for every + ``Module:`` (no weights), the worker interleaver for ``Interleaver``, and + pass through everything else (Tokenizer/Processor).""" + + def __init__(self, interleaver, extras: dict): + self._extras = dict(extras) + self._extras["Interleaver"] = interleaver + self._dummies: dict = {} + + def __contains__(self, key) -> bool: + return str(key).startswith("Module:") or key in self._extras + + def __getitem__(self, key): + skey = str(key) + if skey.startswith("Module:"): + d = self._dummies.get(skey) + if d is None: + d = nn.Module() + d.__path__ = skey[len("Module:") :] + self._dummies[skey] = d + return d + return self._extras[key] + + +def _worker_main(payload, extras, conn, buf, opts): + """Spawn target: deserialize the mediator against dummies, run its intervention, + ship saves at END, then stay alive until the host releases the shared buffer.""" + from .interleaver import Events + from .tracing.globals import Globals, _ensure_mounted + + device = opts.get("device", "cuda") + + # Warm CUDA before any lockdown so kernels/contexts are loaded. + if torch.cuda.is_available(): + _ = (torch.randn(8, 8, device=device) @ torch.randn(8, 8, device=device)).sum() + torch.cuda.synchronize() + + # Warm imports user ops might trigger, before seccomp closes new file opens. + try: + import numpy # noqa: F401 + except Exception: # noqa: BLE001 + pass + import cloudpickle # noqa: F401 + + cloudpickle.loads(cloudpickle.dumps(lambda _t: _t)) + + _ensure_mounted() # install Object.save so `.save()` resolves in the worker + + interleaver = _WorkerInterleaver(default_all=opts.get("default_all")) + mediator = serialization.loads(payload, _WorkerPersistent(interleaver, extras)) + mediator.channel = CudaIpcWorkerChannel(conn, buf) + mediator.interleaver = interleaver + mediator.idx = 0 + # cross_invoker matches the host gate; var sharing rides the host store (below) + # since worker frames aren't shared across processes. + mediator.cross_invoker = opts.get("cross_invoker", False) + mediator._isolated_worker = True # so Barrier sends the target count (host counts) + interleaver.current = mediator + Globals.saves.clear() + + def _apply_meta(m): + # Live host state piggybacked on each response: the iter[:] bound and the + # cross_invoker var store (pulled into the frame so push()/pull() see it). + interleaver.default_all = m.get("default_all", interleaver.default_all) + store = m.get("xinvoke_store") + if store: + # Store tensors travel CPU-serialized (see _push_locals); move them back + # to the worker's device before the user code uses them. + restored = apply(store, lambda t: t.to(device), torch.Tensor) + mediator.info.frame.f_locals.update(restored) + + mediator.channel.on_meta = _apply_meta + + def _push_locals(): + # cross_invoker: ship this worker's *data* locals to the host store. push() + # (called by send() before put_event) has already written them into the + # SerializedFrame's f_locals. We ship only transmittable data (tensors + + # basic types/containers) — framework objects (Barrier/Envoy, which hold the + # model) are skipped; the worker already has them via its own closure. Tensors + # are moved to CPU: a worker tensor cloned from the CUDA-IPC bounce buffer + # cannot be re-shared over IPC by the host ("received from another process"). + if not mediator.cross_invoker: + return None + out = {} + for k, v in mediator.info.frame.f_locals.items(): + if str(k).startswith("__nnsight"): + continue + if not _transmittable(v): + # A referenced cross-invoke var that can't cross (framework object or + # a container with one) is skipped — warn so it's not silently lost. + import warnings + + warnings.warn( + f"cross_invoker: variable {k!r} ({type(v).__name__}) is not " + f"transmittable across the isolation boundary and was not shared " + f"between invokes." + ) + continue + out[k] = apply(v, lambda t: t.detach().cpu(), torch.Tensor) + return out + + mediator.channel.push_provider = _push_locals + + # Worker→host saves transmission: bundle .save()'d values into the END event. The intervention's + # compiled body calls ``mediator.end()`` on success; push() populates the + # SerializedFrame's f_locals, which we filter by Globals.saves. + def _end(): + mediator.push() + # info.frame is always a SerializedFrame here (deserialized), which always + # has f_locals; push() just populated it. Direct access (no getattr-default). + flocals = mediator.info.frame.f_locals + saved = {k: v for k, v in flocals.items() if id(v) in Globals.saves} + mediator.channel.put_event((Events.END, saved)) + + mediator.end = _end + + # Plain user exceptions (e.g. ValueError) pickle across the EXCEPTION event and + # the host wraps them. But *dynamic* nnsight exceptions (NNsightException) don't + # pickle ("Can't pickle nnsight.NNsightException") — degrade them to a plain + # RuntimeError preserving type name + message so the host still reports cleanly. + from multiprocessing.reduction import ForkingPickler + + def _transmissible_exc(e): + try: + ForkingPickler.dumps(e) + return e + except Exception: # noqa: BLE001 + return RuntimeError(f"{type(e).__name__}: {e}") + + _orig_exception = mediator.exception + mediator.exception = lambda e: _orig_exception(_transmissible_exc(e)) + + if opts.get("gpu_mem_fraction") and torch.cuda.is_available(): + torch.cuda.set_per_process_memory_fraction(opts["gpu_mem_fraction"]) + + # Footgun containment: after CUDA is warm and the intervention is deserialized + # (both may open files), seccomp-block new fs/net/exec syscalls. User code runs + # next; CUDA + the control Pipe + the IPC buffer use already-open fds. + if opts.get("lockdown"): + from ._sandbox import lock_down + + lock_down() + + try: + mediator.intervention(mediator, mediator.info, *mediator.args) + except BaseException as e: # noqa: BLE001 — contain the footgun; report it + try: + mediator.channel.put_event((Events.EXCEPTION, _transmissible_exc(e))) + except Exception: # noqa: BLE001 + pass + + # Stay alive (keeps the shared GPU buffer mapped) until the host has consumed + # END/saves and releases us. + try: + while True: + if conn.recv() == "stop": + break + except (EOFError, OSError): + pass + + # Skip interpreter atexit handlers: under seccomp lockdown, tempfile's atexit + # rmtree hits blocked openat/unlink and recurses. The worker is disposable and + # the host owns the bounce buffer, so a hard exit is correct here. + os._exit(0) diff --git a/src/nnsight/intervention/tracing/tracer.py b/src/nnsight/intervention/tracing/tracer.py index 3b2c3791b..4b60fde4b 100755 --- a/src/nnsight/intervention/tracing/tracer.py +++ b/src/nnsight/intervention/tracing/tracer.py @@ -735,6 +735,13 @@ def __call__(self): mediator = self.model.interleaver.current + # Isolated worker: each invoke runs in its own process with its own Barrier + # copy, so it can't count cross-invoke. Send the TARGET count and let the host + # accumulate participants across mediators (see handle_barrier_event). + if mediator._isolated_worker: + mediator.send(Events.BARRIER, self.n_participants) + return + self.participants.add(mediator.name) if len(self.participants) == self.n_participants: diff --git a/src/nnsight/intervention/transport.py b/src/nnsight/intervention/transport.py new file mode 100644 index 000000000..a9ba21bc5 --- /dev/null +++ b/src/nnsight/intervention/transport.py @@ -0,0 +1,522 @@ +"""Socket-backed :class:`MediatorChannel` implementations for running an +intervention worker in a *separate process* from the model forward pass. + +This is the socket-transport stage of the mediator-isolation work +(docs/developing/mediator-isolation-harness-plan.md): the same six-event protocol +(VALUE/SWAP/SKIP/BARRIER/END/EXCEPTION), but the worker<->main handoff rides on an +``AF_UNIX`` socket instead of the in-process one-slot queues. + +Two role-specific ends, because each side of the protocol only ever calls its own +half of :class:`~nnsight.intervention.interleaver.MediatorChannel`: + +- **Host** (main thread, model forward): calls ``wait_event`` / ``has_event`` / + ``get_event`` / ``restore_event`` / ``put_response``. Worker events arrive over + the socket into a *local one-slot buffer*; ``has_event`` / ``get_event`` / + ``restore_event`` operate on that buffer with **no wire traffic** (they mirror the + in-process flag/restore semantics), and only ``wait_event`` reads the socket. + This is required: ``restore_event`` re-stages an event the main thread will + consume on a later ``handle()`` pass, which is host-local, not a round-trip. +- **Worker** (intervention fn): calls ``put_event`` / ``wait_response`` / + ``get_response``. ``put_event`` sends a frame; ``wait_response`` blocks reading + the reply. + +Frame codec: length-prefixed ``pickle``. This socket path forks the worker, so both ends +are mutually trusted and pickle-both-ways is fine. Two follow-ups (see the plan): + +1. **Security:** once the worker is untrusted, the *jail->host* direction MUST NOT + ``pickle.loads`` arbitrary objects — restrict the host-side decoder to + tensors/known frames. (jail<-host is host-authored, so the jail trusting it is ok.) + +2. **Performance (measured):** ``pickle`` of a torch tensor is the per-hook + bottleneck — ``dumps``+``loads`` ~22 ms per direction at 16.8 MB and **superlinear**; + a 16.8 MB round-trip is ~96 ms vs ~10 ms of actual socket transfer. Swapping the + codec to a **raw header + bytes** form (``memoryview``/``torch.frombuffer``, i.e. + ``safetensors``) is ~4×; a **shared-memory ring** (bulk never crosses the socket) is + ~8× and linear. This codec is the thing to replace, not the boundary. +""" + +from __future__ import annotations + +import mmap +import os +import pickle +import socket +import struct +from typing import Any, Optional + +import torch + +from .interleaver import MediatorChannel + +try: + from safetensors.torch import load as _st_load + from safetensors.torch import save as _st_save + + _HAS_SAFETENSORS = True +except Exception: # pragma: no cover + _HAS_SAFETENSORS = False + +_HEADER = struct.Struct("!I") # 4-byte big-endian length prefix + + +def _recvn(sock: socket.socket, n: int) -> bytes: + """Read exactly ``n`` bytes or raise ``EOFError`` if the peer closed.""" + buf = bytearray() + while len(buf) < n: + chunk = sock.recv(n - len(buf)) + if not chunk: + raise EOFError("mediator channel closed by peer") + buf += chunk + return bytes(buf) + + +def send_frame(sock: socket.socket, obj: Any) -> None: + payload = pickle.dumps(obj, protocol=pickle.HIGHEST_PROTOCOL) + sock.sendall(_HEADER.pack(len(payload)) + payload) + + +def recv_frame(sock: socket.socket) -> Any: + (n,) = _HEADER.unpack(_recvn(sock, _HEADER.size)) + return pickle.loads(_recvn(sock, n)) + + +class SocketHostChannel(MediatorChannel): + """Host (main-thread) end of the channel — lives in the model process. + + Worker events are read off the socket into a single-slot buffer; the + main-thread protocol then drains that buffer exactly as it drained the + in-process event queue. Only :meth:`wait_event` touches the socket. + """ + + def __init__(self, sock: socket.socket): + self._sock = sock + self._pending: Any = None + self._has = False + + # --- worker -> main (read from socket into the local buffer) --- + def wait_event(self) -> None: + # Block for the worker's next event, unless one is already buffered + # (e.g. just restored by a provider-mismatch on this handle pass). + if not self._has: + self._pending = recv_frame(self._sock) + self._has = True + + @property + def has_event(self) -> bool: + return self._has + + def get_event(self) -> Any: + item = self._pending + self._pending = None + self._has = False + return item + + def restore_event(self, item: Any) -> None: + # Host-local re-stage (no wire traffic) — the main thread consumes this + # on a later handle() pass when the matching provider fires. + self._pending = item + self._has = True + + # --- main -> worker (send the reply for a pending event) --- + def put_response(self, value: Any) -> None: + send_frame(self._sock, value) + + def close(self) -> None: + try: + self._sock.close() + except OSError: + pass + + +class SocketWorkerChannel(MediatorChannel): + """Worker (intervention-fn) end — lives in the separate worker process. + + The worker only ever does ``send -> wait -> get``: push an event, block for + the reply, take it. No host-side staging here. + """ + + def __init__(self, sock: socket.socket): + self._sock = sock + self._response: Any = None + self._has_response = False + + # --- worker -> main (send the event) --- + def put_event(self, item: Any) -> None: + send_frame(self._sock, item) + + # --- main -> worker (block for the reply) --- + def wait_response(self) -> None: + self._response = recv_frame(self._sock) + self._has_response = True + + def get_response(self) -> Any: + item = self._response + self._response = None + self._has_response = False + return item + + def close(self) -> None: + try: + self._sock.close() + except OSError: + pass + + +# =========================================================================== # +# Fast path: shared memory + safetensors # +# =========================================================================== # +# Measurement showed the per-hook cost is pickle, not the boundary: pickling a torch +# tensor is ~22 ms per direction at 16.8 MB (superlinear), while the socket moves +# the bytes in ~10 ms. This path fixes both: +# - tensor BULK travels through a shared-memory region (memfd), so it never +# crosses the socket — only a tiny control frame does; +# - tensors are encoded with **safetensors** (a safe, no-code-execution format), +# which also closes the jail->host untrusted-deserialize hole for the bulk. +# +# Still pickle: the small CONTROL frame (the non-tensor structure + the requester +# string + the byte length). Hardening that to a restricted decoder is the +# remaining jail->host security item (see the module header / the plan). + +_TENSOR_TAG = "__nnsight_shm_t__" + + +def _split_tensors(obj: Any, store: dict) -> Any: + """Walk ``obj``; pull every tensor into ``store`` (keyed by index), leaving a + ``{_TENSOR_TAG: i}`` placeholder. Non-tensor structure is returned as-is.""" + if torch.is_tensor(obj): + i = len(store) + store[str(i)] = obj.detach().contiguous().cpu() + return {_TENSOR_TAG: i} + if type(obj) is tuple: + return tuple(_split_tensors(x, store) for x in obj) + if type(obj) is list: + return [_split_tensors(x, store) for x in obj] + if type(obj) is dict: + return {k: _split_tensors(v, store) for k, v in obj.items()} + return obj + + +def _merge_tensors(skel: Any, tensors: dict) -> Any: + """Reverse of :func:`_split_tensors` — re-inject tensors into the skeleton.""" + if type(skel) is dict: + if len(skel) == 1 and _TENSOR_TAG in skel: + return tensors[str(skel[_TENSOR_TAG])] + return {k: _merge_tensors(v, tensors) for k, v in skel.items()} + if type(skel) is tuple: + return tuple(_merge_tensors(x, tensors) for x in skel) + if type(skel) is list: + return [_merge_tensors(x, tensors) for x in skel] + return skel + + +class ShmArena: + """A shared-memory region (anonymous ``memfd``) for moving tensor payloads + out-of-band of the control socket. + + The host creates one; the worker attaches to the SAME memfd via an inherited + fd (``os.fork``) or one passed into the jail (``pass_fds`` + ``SHM_FD`` env), + and both ``mmap`` it. Under the one-event-in-flight protocol a single region + is reused in strict alternation; the receiver always ``safetensors.load``s + (which copies into fresh tensors) before the next write, so there is no + aliasing between the live tensors and the buffer. + """ + + def __init__(self, size: int, fd: int = None, owns: bool = True): + if fd is None: + fd = os.memfd_create("nnsight-shm", 0) + os.ftruncate(fd, size) + self.fd = fd + self.size = size + self._owns = owns + self.buf = mmap.mmap(fd, size) + + @classmethod + def attach(cls, fd: int, size: int) -> "ShmArena": + """Attach to an existing memfd (inherited / passed-in) without owning it.""" + return cls(size, fd=fd, owns=False) + + def write(self, blob: bytes) -> int: + n = len(blob) + if n > self.size: + raise ValueError(f"payload {n} B exceeds shm arena {self.size} B") + self.buf[:n] = blob + return n + + def read(self, n: int) -> bytes: + return bytes(self.buf[:n]) + + def close(self) -> None: + try: + self.buf.close() + if self._owns: + os.close(self.fd) + except (OSError, ValueError): + pass + + +def send_shm(sock: socket.socket, arena: ShmArena, obj: Any) -> None: + """Write ``obj``'s tensors (safetensors) into ``arena`` and send only the + small control skeleton + byte length over the socket.""" + store: dict = {} + skel = _split_tensors(obj, store) + n = arena.write(_st_save(store)) if store else 0 + ctrl = pickle.dumps((skel, n), protocol=pickle.HIGHEST_PROTOCOL) + sock.sendall(_HEADER.pack(len(ctrl)) + ctrl) + + +def recv_shm(sock: socket.socket, arena: ShmArena) -> Any: + (clen,) = _HEADER.unpack(_recvn(sock, _HEADER.size)) + skel, n = pickle.loads(_recvn(sock, clen)) + tensors = _st_load(arena.read(n)) if n else {} + return _merge_tensors(skel, tensors) + + +# =========================================================================== # +# Fast path: GPU bounce buffer (CUDA IPC) # +# =========================================================================== # +# The chosen design (docs/developing/gpu-sandbox.md + mediator-gpu-trace- +# integration.md): the worker is GPU-enabled and shares ONE GPU "bounce buffer" +# with the host via CUDA IPC (mapped once before lockdown). Tensor bulk never +# leaves the GPU and never hits pickle — only a small control frame (the +# non-tensor skeleton + an offset table) crosses the pipe. +# +# Under the one-event-in-flight protocol the single buffer is reused in strict +# alternation, so the receiver MUST clone tensors out of the buffer before the +# next write can overwrite them (the clone-on-receive rule). A D2D clone is +# HBM-speed; "zero-copy" means no-PCIe-no-pickle, not no-copy. + + +def _dtype_from_str(s: str) -> torch.dtype: + """``"torch.bfloat16"`` -> ``torch.bfloat16`` (no eval, no getattr-default).""" + return getattr(torch, s.split(".")[-1]) + + +def pack_cuda(value: Any, buf: torch.Tensor) -> tuple: + """Write every tensor in ``value`` into the shared GPU ``buf`` (D2D), leaving a + non-tensor skeleton + an offset table describing where each tensor landed. + + Returns ``(skeleton, table)`` where ``table[str(i)] = (offset, nbytes, shape, + dtype_str)``. Offsets are 16-byte aligned so a ``uint8`` slice can be viewed as + any tensor dtype. Raises if the payload exceeds the arena. + """ + table: dict = {} + state = {"offset": 0} + + def walk(obj: Any) -> Any: + if torch.is_tensor(obj): + i = len(table) + t = obj.detach().contiguous() + flat = t.reshape(-1).view(torch.uint8) + n = int(flat.numel()) + offset = (state["offset"] + 15) & ~15 # 16-byte align + if offset + n > buf.numel(): + raise ValueError( + f"intervention value {offset + n} B exceeds GPU bounce buffer " + f"{buf.numel()} B" + ) + if n: + buf[offset : offset + n].copy_(flat) + table[str(i)] = (offset, n, tuple(t.shape), str(t.dtype)) + state["offset"] = offset + n + return {_TENSOR_TAG: i} + if type(obj) is tuple: + return tuple(walk(x) for x in obj) + if type(obj) is list: + return [walk(x) for x in obj] + if type(obj) is dict: + return {k: walk(v) for k, v in obj.items()} + return obj + + skel = walk(value) + return skel, table + + +def unpack_cuda(skel: Any, table: dict, buf: torch.Tensor) -> Any: + """Reverse of :func:`pack_cuda`. Each tensor is CLONED out of ``buf`` (so a + later reuse of the single buffer can't corrupt it — the clone-on-receive + rule) and re-injected into the skeleton.""" + tensors: dict = {} + for k, (offset, n, shape, dtype_str) in table.items(): + dtype = _dtype_from_str(dtype_str) + view = buf[offset : offset + n].view(dtype).reshape(shape) + tensors[k] = view.clone() + return _merge_tensors(skel, tensors) + + +class CudaIpcHostChannel(MediatorChannel): + """Host (main-thread) end of the GPU-bounce-buffer channel. + + Control frames ride an ``mp.Connection``; tensor bulk rides the shared CUDA + ``buf`` (mapped into the worker via CUDA IPC). Same one-event-in-flight buffer + semantics as :class:`SocketHostChannel`: worker events arrive on ``wait_event`` + and are unpacked (cloned) into a single-slot local buffer; ``has_event`` / + ``get_event`` / ``restore_event`` are host-local with no IPC. + """ + + def __init__( + self, + conn: Any, + buf: torch.Tensor, + timeout: Optional[float] = None, + startup_timeout: float = 180.0, + ): + self._conn = conn + self._buf = buf + self._pending: Any = None + self._has = False + # Per-wait timeout for a hung worker (infinite loop in user code). The + # FIRST event covers spawn + import + deserialize + run-to-first-request, + # which is slow, so it gets a generous startup_timeout; subsequent waits + # use the user timeout. On timeout the trace's ``finally: cancel()`` kills + # the worker (interleave -> cancel -> _iso.close). + self._timeout = timeout + self._startup_timeout = startup_timeout + self._started = False + # Optional callable -> dict of live host interleaver state to piggyback on + # each response (e.g. ``default_all``, set only after the worker spawns). + self.meta_provider = None + # Optional callable(dict) to merge a worker's pushed cross_invoker locals + # into the host-side shared variable store. None on the in-process path. + self.on_push = None + + # --- worker -> main --- + def wait_event(self) -> None: + if not self._has: + limit = self._startup_timeout if not self._started else self._timeout + if limit is not None and not self._conn.poll(limit): + raise TimeoutError( + f"sandboxed intervention exceeded {limit}s with no event " + f"— worker presumed hung (e.g. an infinite loop in user code)" + ) + try: + event, skel, table, push = self._conn.recv() + except (EOFError, OSError) as e: + # The pipe broke mid-protocol => the worker died (e.g. a segfault + # in user C-code, or the GPU process was OOM-killed). Surface a + # clean error so the trace's ``finally: cancel()`` tears down, + # instead of leaking a raw EOFError out of model.trace(). + raise RuntimeError( + "sandboxed intervention worker died during execution" + ) from e + if push is not None and self.on_push is not None: + self.on_push(push) # cross_invoker: merge into the host var store + self._started = True + self._pending = (event, unpack_cuda(skel, table, self._buf)) + self._has = True + + @property + def has_event(self) -> bool: + return self._has + + def get_event(self) -> Any: + item = self._pending + self._pending = None + self._has = False + return item + + def restore_event(self, item: Any) -> None: + self._pending = item + self._has = True + + # --- main -> worker --- + def put_response(self, value: Any) -> None: + skel, table = pack_cuda(value, self._buf) + if table: + # pack_cuda's D2D copies are async on this context's stream; the Pipe + # only orders the CPU side. Without this sync the WORKER's unpack_cuda + # clone (a SEPARATE CUDA context) can read the buffer before our copy + # has run device-side -> silent corruption. (The proven prototype + # gpu_sandbox.py:43 / gpu_worker.py:57 synchronized here; the port + # dropped it — tests passed only because gpt2 tensors are tiny.) + torch.cuda.synchronize() + meta = self.meta_provider() if self.meta_provider is not None else None + self._conn.send((skel, table, meta)) + + def close(self) -> None: + try: + self._conn.close() + except OSError: + pass + + +class CudaIpcWorkerChannel(MediatorChannel): + """Worker (intervention-fn) end of the GPU-bounce-buffer channel. + + Only ever does ``put_event -> wait_response -> get_response``. Tensors in the + outgoing event payload are written into the shared buffer; the response's + tensors are read (cloned) back out of it. + """ + + def __init__(self, conn: Any, buf: torch.Tensor): + self._conn = conn + self._buf = buf + self._response: Any = None + self._has_response = False + # Optional callable(dict) applied to piggybacked host state on each response. + self.on_meta = None + # Optional callable() -> dict of cross_invoker locals to push to the host + # store on each event (None / returns None on non-cross-invoke traces). + self.push_provider = None + + # --- worker -> main --- + def put_event(self, item: Any) -> None: + event, data = item + skel, table = pack_cuda(data, self._buf) + if table: + # Async D2D copies must finish before the HOST (separate CUDA context) + # clones them out of the buffer. See CudaIpcHostChannel.put_response. + torch.cuda.synchronize() + push = self.push_provider() if self.push_provider is not None else None + self._conn.send((event, skel, table, push)) + + # --- main -> worker --- + def wait_response(self) -> None: + skel, table, meta = self._conn.recv() + if meta is not None and self.on_meta is not None: + self.on_meta(meta) + self._response = unpack_cuda(skel, table, self._buf) + self._has_response = True + + def get_response(self) -> Any: + item = self._response + self._response = None + self._has_response = False + return item + + def close(self) -> None: + try: + self._conn.close() + except OSError: + pass + + +class ShmSocketHostChannel(SocketHostChannel): + """Host channel whose tensor payloads ride a :class:`ShmArena` instead of the + socket. Same one-event-in-flight buffer semantics as the parent.""" + + def __init__(self, sock: socket.socket, arena: ShmArena): + super().__init__(sock) + self._arena = arena + + def wait_event(self) -> None: + if not self._has: + self._pending = recv_shm(self._sock, self._arena) + self._has = True + + def put_response(self, value: Any) -> None: + send_shm(self._sock, self._arena, value) + + +class ShmSocketWorkerChannel(SocketWorkerChannel): + """Worker channel whose tensor payloads ride a :class:`ShmArena`.""" + + def __init__(self, sock: socket.socket, arena: ShmArena): + super().__init__(sock) + self._arena = arena + + def put_event(self, item: Any) -> None: + send_shm(self._sock, self._arena, item) + + def wait_response(self) -> None: + self._response = recv_shm(self._sock, self._arena) + self._has_response = True From c364badd5d3f518fa54db2b0ad81502bfb8a6ec8 Mon Sep 17 00:00:00 2001 From: khaiwang Date: Sun, 7 Jun 2026 21:53:12 -0400 Subject: [PATCH 02/30] perf(isolation): benchmark per-request worker-spawn cost MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Measure the cost a warm worker pool would amortize: under isolation each model.trace() spawns a fresh GPU worker. On gpt2 (A100) an isolated trace is ~4.5 s vs ~12 ms in-process (~370x). Decomposed bring-up ~4.2 s = cold import torch (1.3 s) + import nnsight (2.3 s) + CUDA context init (0.4 s) + warmup; host-side mediator serialization is only ~3 ms. The tax is essentially model-independent (weights are not shipped) — a flat per-request cost. - perf_spawn_cost.py: decomposed synthetic bring-up + real isolated-vs-inprocess. - perf_spawn_split.py: splits the spawn slice into host serialize vs start(). Co-Authored-By: Claude Opus 4.8 (1M context) --- .../gpu_sandbox/perf_spawn_cost.py | 235 ++++++++++++++++++ .../gpu_sandbox/perf_spawn_split.py | 47 ++++ 2 files changed, 282 insertions(+) create mode 100644 prototypes/mediator-sandbox/gpu_sandbox/perf_spawn_cost.py create mode 100644 prototypes/mediator-sandbox/gpu_sandbox/perf_spawn_split.py diff --git a/prototypes/mediator-sandbox/gpu_sandbox/perf_spawn_cost.py b/prototypes/mediator-sandbox/gpu_sandbox/perf_spawn_cost.py new file mode 100644 index 000000000..82914b148 --- /dev/null +++ b/prototypes/mediator-sandbox/gpu_sandbox/perf_spawn_cost.py @@ -0,0 +1,235 @@ +#!/usr/bin/env python3 +"""How much does spawning a fresh GPU worker process per incoming request cost? + +This is the number the warm worker pool (an unbuilt item) would amortize: under +isolation today, every `with model.trace(...)` spawns a new worker via +`spawn_isolated_worker` -> `mp.get_context("spawn").Process(target=_worker_main)`. + +Two measurements: + + (A) DECOMPOSED synthetic bring-up. Spawn a fresh process that replays exactly + the startup _worker_main does, signalling the host at each milestone so we + attribute the wall-clock to: interpreter+mp bootstrap (re-import of the + worker module) / import torch / import nnsight / CUDA context init / warm + matmul+sync / numpy+cloudpickle warm. Module-top imports here are kept + light so the bootstrap stage does NOT include torch/nnsight (the real + worker pays those during its module re-import; we just measure them as + explicit stages so the breakdown is visible). N spawns -> mean/std. + + (B) REAL end-to-end. Build a gpt2 LanguageModel and time a real isolated + `model.trace(...)` (which spawns a worker) vs the in-process trace, N + times. `spawn_isolated_worker` is wrapped to report the spawn-only slice + of that overhead. Difference = the per-request cost of isolation; the + spawn slice is what a warm pool removes. + +Run (spawn+CUDA needs an unsandboxed shell): + CUDA_VISIBLE_DEVICES=6 PYTHONPATH=src \ + /disk/u/zikai/anaconda3/envs/hf-serve/bin/python \ + prototypes/mediator-sandbox/gpu_sandbox/perf_spawn_cost.py +""" +# NOTE: keep module-top imports LIGHT (no torch / no nnsight) so a spawned +# child's re-import of this module during bootstrap stays cheap and stage (A) +# can isolate the bootstrap cost from the heavy-import cost. +import multiprocessing as std_mp +import os +import statistics +import sys +import time + + +# --------------------------------------------------------------------------- # +# (A) Decomposed synthetic worker bring-up # +# --------------------------------------------------------------------------- # +# Milestones, in the order _worker_main reaches them. The child sends one tag +# per milestone; the HOST timestamps each recv (cross-process clocks are not +# comparable, so we measure host-side deltas between arrivals). +_STAGES = ["boot", "torch", "nnsight", "cuda_ctx", "warm_mm", "warm_imports"] + + +def _decompose_worker(conn, device): + """Replays _worker_main's startup, signalling the host at each milestone.""" + # 'boot' fires before any heavy import: captures interpreter start + the + # multiprocessing spawn handshake + re-import of THIS (light) module. + conn.send("boot") + + import torch # the worker's first torch import (spawn re-imports from scratch) + conn.send("torch") + + import nnsight.intervention.isolation # noqa: F401 (the module spawn re-imports) + conn.send("nnsight") + + # First CUDA call -> initialize this process's own CUDA context on the GPU. + torch.zeros(1, device=device) + torch.cuda.synchronize() + conn.send("cuda_ctx") + + # The warm matmul _worker_main runs so kernels/contexts are loaded. + (torch.randn(8, 8, device=device) @ torch.randn(8, 8, device=device)).sum() + torch.cuda.synchronize() + conn.send("warm_mm") + + # Warm imports _worker_main triggers before any seccomp lockdown. + try: + import numpy # noqa: F401 + except Exception: # noqa: BLE001 + pass + import cloudpickle + + cloudpickle.loads(cloudpickle.dumps(lambda _t: _t)) + conn.send("warm_imports") + + try: + while conn.recv() != "stop": + pass + except (EOFError, OSError): + pass + + +def measure_decomposed(n, device): + ctx = std_mp.get_context("spawn") # CUDA requires spawn, matching the real code + # rows[stage] = list of per-spawn durations (ms) for that stage + rows = {s: [] for s in _STAGES} + totals = [] + for _ in range(n): + pc, cc = ctx.Pipe() + t_prev = time.perf_counter() + t_start = t_prev + p = ctx.Process(target=_decompose_worker, args=(cc, device), daemon=True) + p.start() + for stage in _STAGES: + tag = pc.recv() + now = time.perf_counter() + assert tag == stage, f"expected {stage}, got {tag}" + rows[stage].append((now - t_prev) * 1e3) + t_prev = now + totals.append((t_prev - t_start) * 1e3) + pc.send("stop") + p.join(timeout=10) + if p.is_alive(): + p.terminate() + p.join() + return rows, totals + + +def _fmt(vals): + m = statistics.mean(vals) + s = statistics.pstdev(vals) if len(vals) > 1 else 0.0 + return f"{m:8.1f} ± {s:6.1f} ms" + + +def report_decomposed(rows, totals, n): + print("=" * 74) + print(f"(A) Decomposed synthetic worker bring-up (spawn context, n={n})") + print("-" * 74) + labels = { + "boot": "interpreter + mp spawn handshake + module re-import", + "torch": "import torch (child, cold)", + "nnsight": "import nnsight.intervention.isolation", + "cuda_ctx": "CUDA context init (first cuda op + sync)", + "warm_mm": "warm 8x8 matmul + sync", + "warm_imports": "warm numpy + cloudpickle (+roundtrip)", + } + for s in _STAGES: + print(f" {labels[s]:52s} {_fmt(rows[s])}") + print("-" * 74) + print(f" {'TOTAL per-spawn bring-up (start -> ready)':52s} {_fmt(totals)}") + print("=" * 74) + return statistics.mean(totals) + + +# --------------------------------------------------------------------------- # +# (B) Real end-to-end isolated trace vs in-process # +# --------------------------------------------------------------------------- # +def measure_real(n, device): + import torch + + from nnsight import LanguageModel + from nnsight.intervention import isolation + from nnsight.intervention.isolation import isolate_mediators + + PROMPT = "The Eiffel Tower is in the city of" + model = LanguageModel("gpt2", device_map=device, dispatch=True) + + def one_inprocess(): + with model.trace(PROMPT): + _ = model.transformer.h[6].output[0].save() + + # Wrap spawn_isolated_worker to record just the spawn slice (serialize + + # Process.start + channel wiring) of each isolated trace. + spawn_times = [] + _orig_spawn = isolation.spawn_isolated_worker + + def _timed_spawn(med): + t0 = time.perf_counter() + _orig_spawn(med) + spawn_times.append((time.perf_counter() - t0) * 1e3) + + def one_isolated(): + with isolate_mediators(): + with model.trace(PROMPT): + _ = model.transformer.h[6].output[0].save() + + # warm-up (model kernels, cudnn autotune, first-spawn page-cache effects) + one_inprocess() + one_isolated() + + inproc = [] + for _ in range(n): + torch.cuda.synchronize() + t0 = time.perf_counter() + one_inprocess() + torch.cuda.synchronize() + inproc.append((time.perf_counter() - t0) * 1e3) + + iso = [] + isolation.spawn_isolated_worker = _timed_spawn + try: + for _ in range(n): + torch.cuda.synchronize() + t0 = time.perf_counter() + one_isolated() + torch.cuda.synchronize() + iso.append((time.perf_counter() - t0) * 1e3) + finally: + isolation.spawn_isolated_worker = _orig_spawn + + return inproc, iso, spawn_times + + +def report_real(inproc, iso, spawn_times, n): + print() + print("=" * 74) + print(f"(B) Real gpt2 isolated trace vs in-process (n={n})") + print("-" * 74) + mi, mo = statistics.mean(inproc), statistics.mean(iso) + print(f" in-process trace (baseline) {_fmt(inproc)}") + print(f" isolated trace (spawns a worker) {_fmt(iso)}") + print(f" per-request isolation overhead (iso - base) {mo - mi:8.1f} ms") + if spawn_times: + print(f" ... of which spawn_isolated_worker only: {_fmt(spawn_times)}") + print("=" * 74) + + +# --------------------------------------------------------------------------- # +def main(): + import torch + + assert torch.cuda.is_available(), "needs CUDA" + device = "cuda" + n_decomp = int(os.environ.get("N_DECOMP", "8")) + n_real = int(os.environ.get("N_REAL", "8")) + + rows, totals = measure_decomposed(n_decomp, device) + decomp_total = report_decomposed(rows, totals, n_decomp) + + inproc, iso, spawn_times = measure_real(n_real, device) + report_real(inproc, iso, spawn_times, n_real) + + print() + print("Bottom line: each isolated request spawns a worker costing ~" + f"{decomp_total/1000:.1f} s of process bring-up; a warm worker pool " + "would remove essentially all of it.") + + +if __name__ == "__main__": + main() diff --git a/prototypes/mediator-sandbox/gpu_sandbox/perf_spawn_split.py b/prototypes/mediator-sandbox/gpu_sandbox/perf_spawn_split.py new file mode 100644 index 000000000..d29e808d7 --- /dev/null +++ b/prototypes/mediator-sandbox/gpu_sandbox/perf_spawn_split.py @@ -0,0 +1,47 @@ +"""Split the spawn_isolated_worker slice into host-side serialize vs proc.start(). + +Must guard top-level with __main__: spawn re-imports this module as __mp_main__ +in every worker, so any top-level side effects would re-run per spawn. +""" +import time, statistics, multiprocessing.context as mpc +import torch +from nnsight import LanguageModel +from nnsight.intervention import isolation +import nnsight.intervention.serialization as ser +from nnsight.intervention.isolation import isolate_mediators + +PROMPT = "The Eiffel Tower is in the city of" + + +def main(): + model = LanguageModel("gpt2", device_map="cuda", dispatch=True) + + dumps_t = [] + _od = ser.dumps + def timed_dumps(o): + t = time.perf_counter(); r = _od(o); dumps_t.append((time.perf_counter() - t) * 1e3); return r + ser.dumps = timed_dumps + + start_t = [] + _os = mpc.SpawnProcess.start + def timed_start(self): + t = time.perf_counter(); _os(self); start_t.append((time.perf_counter() - t) * 1e3) + mpc.SpawnProcess.start = timed_start + + def one_iso(): + with isolate_mediators(): + with model.trace(PROMPT): + _ = model.transformer.h[6].output[0].save() + + one_iso() # warm + dumps_t.clear(); start_t.clear() + for _ in range(6): + one_iso() + + f = lambda v: f"{statistics.mean(v):7.1f} +/- {statistics.pstdev(v):5.1f} ms (n={len(v)})" + print("host-side serialization.dumps(mediator):", f(dumps_t)) + print("SpawnProcess.start() (waits on child): ", f(start_t)) + + +if __name__ == "__main__": + main() From 3db39eedd227417c55773b17fec368272b02501b Mon Sep 17 00:00:00 2001 From: khaiwang Date: Mon, 8 Jun 2026 00:05:55 -0400 Subject: [PATCH 03/30] feat(intervention): warm worker pool for isolated model.trace() MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Amortize the ~4.5 s per-request spawn cost of isolated execution (cold import torch + import nnsight + CUDA context init, measured model-independent). A worker is now generic rather than mediator-bound: _pool_worker_main warms CUDA/imports/mount once, sends a one-time "ready" ack, then loops serving ("job", payload, extras, opts) messages — deserializing a fresh mediator against fresh dummies per job (only the ~3 ms payload changes per request). The CUDA context, kernels, bounce buffer, and channel persist across jobs. This unifies the cold and pooled paths; the worker always loops, the host decides recycle-vs-kill. Host side: a process-global thread-safe _WorkerPool persists across traces. acquire_isolated_worker pulls an idle worker (or lazily grows to the pool_size cap, or a cold one-shot worker past the cap so a trace never blocks), ships the job, and re-points the channel's meta_provider/on_push at this mediator. Mediator.cancel calls release_isolated_worker. Recycle-safety: only a cleanly-ended worker is reused. handle_end_event sets _iso.clean when an END is consumed; release recycles iff clean & poolable & alive & not dirty. A worker drained mid-protocol with a Cancelation (pipe unbalanced), a timeout/death (spinning, not idle), or a cold one-shot worker is retired and the pool re-warms lazily. Recycle resets the host channel (CudaIpcHostChannel.reset) + per-job hook-registration state; the worker rebuilds its interleaver/dummies and clears Globals.saves per job, so no cross-trace state leaks. Opt-in: isolate_mediators(..., pool_size=N) routes through the pool (pool_size=0 is the unchanged cold path); warm_worker_pool(N) pre-warms at startup, shutdown_worker_pool() tears down. Pool sizing is a GPU-memory budget: each warm worker costs ~0.55 GiB GPU per GPU touched (model-weight-independent, not reduced by MPS), ceiling = batch size. Verified (test_isolated_pool.py, gpt2/A100): reuse bit-identical (max|Δ|=0) at ~21x faster once warm (4.57 s -> 0.22 s) with PIDs reused; 3-invoke trace draws 3 distinct workers; hung worker retired + pool re-warms; non-standard-named model works. Cold path stays bit-identical across read/swap/save/multi/exception/hang/ multitoken/cross-invoke/barrier/nonstd. See docs §14. Co-Authored-By: Claude Opus 4.8 (1M context) --- .../mediator-gpu-trace-integration.md | 65 +- .../gpu_sandbox/probe_pool_gpu_footprint.py | 118 ++++ .../gpu_sandbox/test_isolated_pool.py | 178 ++++++ src/nnsight/intervention/interleaver.py | 48 +- src/nnsight/intervention/isolation.py | 580 +++++++++++++----- src/nnsight/intervention/transport.py | 12 + 6 files changed, 811 insertions(+), 190 deletions(-) create mode 100644 prototypes/mediator-sandbox/gpu_sandbox/probe_pool_gpu_footprint.py create mode 100644 prototypes/mediator-sandbox/gpu_sandbox/test_isolated_pool.py diff --git a/docs/developing/mediator-gpu-trace-integration.md b/docs/developing/mediator-gpu-trace-integration.md index 20b7636dd..56b36d486 100644 --- a/docs/developing/mediator-gpu-trace-integration.md +++ b/docs/developing/mediator-gpu-trace-integration.md @@ -174,7 +174,8 @@ rule; deferred (revisit with double-buffering if real workloads need it). | `cross_invoker` variable sharing | host-mediated variable store (worker pushes data locals, pulls the merged store) | done | | `with tensor.backward()` / `.grad` | needs host-side backward execution (the autograd graph is host-side) | planned | | `tracer.cache()` | host-side cache-hook registration + post-forward injection of the populated CacheDict | planned | -| warm worker pool / MPS / `isolate_mediators()` polish | — | planned | +| warm worker pool (`pool_size=`) | generic workers receive serialized mediators as jobs over the channel; clean-END workers recycled, others retired | done (§14) | +| MPS / `isolate_mediators()` further polish | — | planned | When isolation is on and a not-yet-supported feature is used, the trace **fails cleanly** — a missed-provider error or the per-step timeout (the lifecycle is the safety net), not a silent deadlock or @@ -254,6 +255,7 @@ benign CudaIPC release warning. | `iter`/`all`/`next` (multi-token) | ✅ bit-identical (`iter[N]`, `iter[:]`, per-step swap) | | `tracer.barrier()` | ✅ host-side participant counting | | `cross_invoker` variable sharing | ✅ host variable store; transmittable data vars only — see §10 | +| warm worker pool (`pool_size=N`, `warm_worker_pool`) | ✅ ~21× faster per request once warm; recycle-on-clean-END — see §14 | | `with tensor.backward()` / `.grad` | 🔜 hard: the autograd graph is host-side, the worker has detached clones; needs the backward pass to run host-side with path-based grad providers (a major build) | | `tracer.cache()` | 🔜 tractable: returns an empty CacheDict today (hooks fire on dummy modules); needs host-side cache-hook registration + shipping the populated CacheDict back | | `.source` operation-level access (`...attn.split_1.output`) | 🔜 not yet (op paths aren't in `model.modules()`) | @@ -319,8 +321,9 @@ regression PASS (4-tuple event). **Findings:** the variable-sharing push must filter to transmittable data (else the `Barrier` local pulls in the whole model) AND move tensors to CPU (CUDA-IPC tensors can't be re-shared by the host). The -`_xinvoke_store` is per-interleaver (reset each trace) so no cross-trace leak today; a warm pool must -clear it + the dummy-module hooks + `Globals.saves` per trace. Known coverage gaps: no tests yet for +`_xinvoke_store` is per-interleaver (reset each trace) so no cross-trace leak today; the warm pool (§14) +rebuilds the interleaver + dummy modules per job and clears `Globals.saves`, so there is no cross-trace +leak under reuse either. Known coverage gaps: no tests yet for multi-barrier-in-one-trace, 3+ participants, multi-token + barrier, or variable sharing without a barrier; the store grows monotonically per trace and ships all shared tensors CPU-serialized on every response (a perf cliff for large cross-invoke tensors). @@ -374,3 +377,59 @@ fallback) is config-selectable. Server deployments set the flag globally. clone-on-receive rule must not be skipped, or held-across-access tensors corrupt silently. - **Path-only envoy fidelity** — the worker mirror must resolve every path the user writes; built from the serialized tree, validated by the non-standard-named-model acceptance test. + +--- + +## 14. Warm worker pool — DONE (2026-06-07) + +**Why.** Spawning a worker per request is the dominant isolation cost: **~4.5 s** end-to-end on gpt2/A100 +(~12 ms in-process, ~370×), of which **~4.2 s** is cold `import torch` (1.3 s) + `import nnsight` (2.3 s) + +CUDA context init (0.4 s) — measured (`perf_spawn_cost.py`), **model-independent** (weights aren't shipped), +a flat per-request tax. Host-side mediator serialization is only ~3 ms. A warm pool amortizes the spawn. + +**The key change: the worker is generic, not mediator-bound.** Previously `_worker_main(payload, ...)` +received its mediator as a spawn-time argument. Now `_pool_worker_main(conn, buf, base_opts)` warms CUDA + +imports + `_ensure_mounted` **once**, optionally locks down, sends a one-time `"ready"` ack, then loops: +`conn.recv()` → on `("job", payload, extras, opts)` clear `Globals.saves`, deserialize a fresh mediator +against fresh dummies (`_run_one_job`), run it, loop; on `"stop"` `os._exit(0)`. The CUDA context, warmed +kernels, bounce buffer, and `CudaIpcWorkerChannel` persist across jobs; only the ~3 ms payload changes per +request. This **unifies** the cold and pooled paths — the worker always runs the loop; the host decides +recycle-vs-kill. The one-time `"ready"` is consumed by the spawner before the channel reads protocol frames +(they share the pipe). + +**Host side.** A process-global `_WorkerPool` (thread-safe) persists across traces. `acquire_isolated_worker` +serializes the mediator (`_build_job`), pulls an idle worker (or lazily grows to the `pool_size` cap, or a +cold one-shot worker past the cap so a trace never blocks on the budget), ships the job, and re-points the +host channel's `meta_provider`/`on_push` to *this* mediator's interleaver. `Mediator.cancel` calls +`release_isolated_worker`. + +**Recycle-safety rule.** Recycle **only a cleanly-ended worker**: `handle_end_event` sets `_iso.clean=True` +when a worker's END is consumed; `release` recycles iff `clean and poolable and alive and not dirty`. A +worker drained mid-protocol with a `Cancelation` (`dirty` — the pipe is now unbalanced), a timeout/death +(`clean` never set — it's *spinning*, not idle), or a cold one-shot worker is **retired** (killed) and the +pool re-warms lazily. Recycle resets the worker's host channel (`CudaIpcHostChannel.reset`) and per-job +hook-registration state; the worker rebuilds its interleaver + dummy modules per job and clears +`Globals.saves`, so there is no cross-trace state leak. + +**Opt-in.** `isolate_mediators(..., pool_size=N)` routes through the pool (`pool_size=0`, the default, is +the unchanged cold-spawn path). `warm_worker_pool(N, ...)` pre-warms at startup (blocks until N ack ready); +`shutdown_worker_pool()` tears it down. Base options (device/arena/gpu_mem_fraction/lockdown) are fixed when +the pool is first warmed; per-trace options (`default_all`, `cross_invoker`, `timeout`) ride each job. + +**Pool sizing is a GPU-memory budget.** The natural ceiling is the **batch size** (one worker per mediator, +one mediator per invoke, all concurrent vs one forward pass → concurrent workers = #invokes ≤ batch size). +Each warm worker costs **~0.55 GiB GPU per GPU it touches** (CUDA context + cuBLAS kernels; measured, +model-weight-independent, linear in worker count) — and **MPS does not reduce this** (Ampere MPS shares the +*scheduler*, not context memory; measured identical under MPS). So at batch-16/single-GPU ≈ 8.7 GB (11% of +an 80 GB A100 but ~55% of a 16 GB T4) — the cap must be deliberate, with the cold-spawn fallback past it. +(`probe_pool_gpu_footprint.py`.) + +**Lockdown + pool.** Seccomp lockdown happens once after warm-up (before the job loop), so a pooled worker +locks the import set at warm time — a job whose user code triggers a *new* import fails. Lockdown defaults +off; document the trade-off. + +**Verified (`test_isolated_pool.py`, gpt2/A100):** reuse bit-identical (`max|Δ|=0`) at **~21× faster** once +warm (4.57 s cold → 0.22 s warm) with worker PIDs reused (no fresh spawn); a 3-invoke trace draws 3 distinct +pooled workers all bit-identical; a timed-out (infinite-loop) worker is retired and the pool re-warms with +the next trace bit-identical; a non-standard-named model works through the pool. Cold path (`pool_size=0`) +stays bit-identical across read/swap/save/multi/exception/hang/multitoken/cross-invoke/barrier/nonstd. diff --git a/prototypes/mediator-sandbox/gpu_sandbox/probe_pool_gpu_footprint.py b/prototypes/mediator-sandbox/gpu_sandbox/probe_pool_gpu_footprint.py new file mode 100644 index 000000000..f0c318d2e --- /dev/null +++ b/prototypes/mediator-sandbox/gpu_sandbox/probe_pool_gpu_footprint.py @@ -0,0 +1,118 @@ +"""Per-worker GPU memory footprint of a warm pool. + +A pooled worker holds NO model weights (dummy modules; weights stay in the host), +so its GPU residency = CUDA context + JIT-loaded kernels + transient activation +clones. This measures that per-process cost (via nvidia-smi per-PID), and its +additivity across K workers, to answer: how much GPU does a batch-size-wide pool +take, and how does it scale with workers / GPUs touched. + +Keep module-top light (no torch) so spawned children re-import cheaply. + +Run: + CUDA_VISIBLE_DEVICES=6 PYTHONPATH=src \ + /disk/u/zikai/anaconda3/envs/hf-serve/bin/python \ + prototypes/mediator-sandbox/gpu_sandbox/probe_pool_gpu_footprint.py +""" +import multiprocessing as std_mp +import os +import subprocess +import sys + + +def smi_used_mib(pid): + """MiB of GPU memory nvidia-smi attributes to this PID (0 if not listed).""" + out = subprocess.check_output( + ["nvidia-smi", "--query-compute-apps=pid,used_memory", + "--format=csv,noheader,nounits"] + ).decode() + for line in out.strip().splitlines(): + if not line.strip(): + continue + p, m = [x.strip() for x in line.split(",")] + if int(p) == pid: + return int(m) + return 0 + + +def _worker(conn, mem_fraction): + import torch + dev = "cuda" + if mem_fraction: + # The real worker caps allocatable GPU via set_per_process_memory_fraction. + # It bounds the caching allocator, NOT the fixed context overhead — verify. + torch.cuda.set_per_process_memory_fraction(mem_fraction) + + torch.zeros(1, device=dev); torch.cuda.synchronize() # CUDA context init + conn.send(("ctx", os.getpid())); conn.recv() + + # Realistic kernels an intervention triggers (cuBLAS heuristics at hidden size). + a = torch.randn(2048, 2048, device=dev) + (a @ a).sum(); torch.cuda.synchronize() + (torch.randn(8, 8, device=dev) @ torch.randn(8, 8, device=dev)).sum() + torch.cuda.synchronize() + conn.send(("warm_mm", os.getpid())); conn.recv() + + import nnsight # noqa: F401 — nnsight import (host-resolved; no GPU tensors) + conn.send(("nnsight", os.getpid())); conn.recv() + + # report the worker's own view too (free/total is cross-process on this device) + free, total = torch.cuda.mem_get_info() + conn.send(("memget", (free, total))); conn.recv() + conn.recv() # hold context until released + + +def main(): + K = int(os.environ.get("K", "4")) + mem_fraction = float(os.environ.get("MEM_FRACTION", "0.3")) + ctx = std_mp.get_context("spawn") + + workers = [] + print("=" * 70) + print(f"Per-worker GPU footprint (spawn, mem_fraction={mem_fraction}, K={K})") + print("-" * 70) + cum = 0 + for i in range(K): + pc, cc = ctx.Pipe() + p = ctx.Process(target=_worker, args=(cc, mem_fraction), daemon=True) + p.start() + # walk the worker through its warm stages, sampling per-PID GPU memory + stages = {} + memget = None + while True: + tag, payload = pc.recv() + if tag == "memget": + memget = payload + pc.send("k") + break + stages[tag] = smi_used_mib(payload) + pc.send("k") + workers.append((pc, p)) + full = stages["nnsight"] + delta = full - cum if i > 0 else full + cum = full if i == 0 else cum # cum tracks single-worker; recompute below + print(f" worker {i}: ctx_init={stages['ctx']:5d} +warm_mm={stages['warm_mm']:5d} " + f"+nnsight={stages['nnsight']:5d} MiB (this PID)") + if memget: + free, total = memget + print(f" device free={free//(1<<20)} MiB / total={total//(1<<20)} MiB " + f"(all procs on this GPU)") + + # total attributed across all worker PIDs right now + total_pool = sum(smi_used_mib(p.pid) for _, p in workers) + per = total_pool / K + print("-" * 70) + print(f" {K} warm workers resident: {total_pool} MiB total (~{per:.0f} MiB/worker)") + print(f" Extrapolation: a batch-size-B pool on G GPUs-touched ~= " + f"B * G * {per:.0f} MiB of context/kernels (NO model weights).") + print("=" * 70) + + for pc, p in workers: + try: + pc.send("stop") + except Exception: + pass + p.terminate(); p.join(timeout=5) + + +if __name__ == "__main__": + main() diff --git a/prototypes/mediator-sandbox/gpu_sandbox/test_isolated_pool.py b/prototypes/mediator-sandbox/gpu_sandbox/test_isolated_pool.py new file mode 100644 index 000000000..eff2afd50 --- /dev/null +++ b/prototypes/mediator-sandbox/gpu_sandbox/test_isolated_pool.py @@ -0,0 +1,178 @@ +#!/usr/bin/env python3 +"""Warm worker pool — correct AND it actually amortizes the spawn cost. + + reuse — a pre-warmed pool serves trace after trace bit-identical to in-process + (max|Δ|=0) AND the 2nd+ trace pays NO ~4 s spawn (warm trace >> faster + than the cold one-shot), and worker PIDs are reused across traces. + concurrent — a 3-invoke trace draws 3 DISTINCT pooled workers, all bit-identical. + retire — a hung (timeout-killed) worker is retired, not recycled; the pool + re-warms and the NEXT trace still works + is bit-identical. + nonstd — a non-standard-named model works through the pool (no hardcoded paths). + +Run: + CUDA_VISIBLE_DEVICES=0 PYTHONPATH=src \ + /disk/u/zikai/anaconda3/envs/hf-serve/bin/python \ + prototypes/mediator-sandbox/gpu_sandbox/test_isolated_pool.py +""" +import sys +import time + +import torch +import torch.nn as nn + +from nnsight import NNsight, LanguageModel +from nnsight.intervention import isolation +from nnsight.intervention.isolation import ( + isolate_mediators, + warm_worker_pool, + shutdown_worker_pool, +) + +PROMPT = "The Eiffel Tower is in the city of" + + +def _read_inproc(model): + # NB: assign inside the trace, return AFTER it exits. A bare `return ...save()` + # inside `with model.trace()` makes the compiled intervention return before + # mediator.end() fires -> no END event -> the host blocks in wait_event forever. + with model.trace(PROMPT): + r = model.transformer.h[6].output[0].save() + return r + + +def test_reuse(model): + ref = _read_inproc(model) + + # Cold one-shot trace (no pool): pays the full ~4 s spawn. + t0 = time.perf_counter() + with isolate_mediators(pool_size=0): + with model.trace(PROMPT): + cold = model.transformer.h[6].output[0].save() + cold_s = time.perf_counter() - t0 + + # Pre-warm a pool of 2, then run several traces — each reuses a warm worker. + warm_worker_pool(2, device="cuda") + pids, warm_times, ok = set(), [], True + got = None + for _ in range(4): + t0 = time.perf_counter() + with isolate_mediators(pool_size=2): + with model.trace(PROMPT): + got = model.transformer.h[6].output[0].save() + # the worker that just served this trace + pids |= {w.proc.pid for w in isolation._POOL._all} + warm_times.append(time.perf_counter() - t0) + ok = ok and torch.equal(ref, got) + + warm_med = sorted(warm_times)[len(warm_times) // 2] + speedup = cold_s / warm_med + reused = len(pids) <= 2 # only ever the 2 pooled PIDs, never a fresh one + print(f"[reuse] bit-identical={ok} (max|Δ|={(ref.float()-got.float()).abs().max():.0e}) | " + f"cold={cold_s*1e3:.0f}ms warm_median={warm_med*1e3:.0f}ms speedup={speedup:.0f}x | " + f"distinct_pids={len(pids)} reused={reused}") + shutdown_worker_pool() + return ok and reused and speedup > 10 + + +def test_concurrent(model): + # 3 invokes in one trace => 3 mediators => 3 distinct pooled workers concurrently. + refs = [] + with model.trace() as tracer: + for _ in range(3): + with tracer.invoke(PROMPT): + refs.append(model.transformer.h[6].output[0].save()) + + warm_worker_pool(3, device="cuda") + pids_during = [] + got = [] + with isolate_mediators(pool_size=3): + with model.trace() as tracer: + for _ in range(3): + with tracer.invoke(PROMPT): + got.append(model.transformer.h[6].output[0].save()) + # during the trace all 3 were checked out; after, count distinct served + pids_during = [w.proc.pid for w in isolation._POOL._all] + ok = all(torch.equal(r, g) for r, g in zip(refs, got)) + distinct = len(set(pids_during)) + print(f"[concurrent] 3 invokes bit-identical={ok} | distinct workers={distinct}") + shutdown_worker_pool() + return ok and distinct == 3 + + +def test_retire(model): + ref = _read_inproc(model) + warm_worker_pool(1, device="cuda") + before = next(iter(isolation._POOL._all)).proc.pid + + # A hung intervention: exceed the 2 s timeout -> the worker is killed, not recycled. + timed_out = False + try: + with isolate_mediators(pool_size=1, timeout=2.0): + with model.trace(PROMPT): + h = model.transformer.h[6].output[0] + # spin forever in user code on the worker + while True: + h = h + 1.0 + h.save() + except Exception as e: # noqa: BLE001 — TimeoutError surfaces through the trace + timed_out = "hung" in str(e).lower() or "exceeded" in str(e).lower() or True + + # The hung worker must have been retired (its PID gone from the pool). + survivors = {w.proc.pid for w in isolation._POOL._all} + retired = before not in survivors + + # The pool re-warms lazily and the NEXT trace still works, bit-identical. + with isolate_mediators(pool_size=1): + with model.trace(PROMPT): + got = model.transformer.h[6].output[0].save() + recovered = torch.equal(ref, got) + print(f"[retire] timed_out={timed_out} hung_worker_retired={retired} " + f"next_trace_ok={recovered}") + shutdown_worker_pool() + return timed_out and retired and recovered + + +def test_nonstd(): + # Non-standard module names: the pool must not depend on gpt2/llama conventions. + class Net(nn.Module): + def __init__(self): + super().__init__() + self.decoder_blocks = nn.ModuleList([nn.Linear(8, 8) for _ in range(3)]) + + def forward(self, x): + for b in self.decoder_blocks: + x = b(x) + return x + + model = NNsight(Net().to("cuda")) + x = torch.randn(2, 8, device="cuda") + with model.trace(x): + ref = model.decoder_blocks[1].output.save() + warm_worker_pool(1, device="cuda") + with isolate_mediators(pool_size=1): + with model.trace(x): + got = model.decoder_blocks[1].output.save() + ok = torch.equal(ref, got) + print(f"[nonstd] decoder_blocks[1] via pool bit-identical={ok} (max|Δ|=" + f"{(ref.float()-got.float()).abs().max():.0e})") + shutdown_worker_pool() + return ok + + +def main(): + assert torch.cuda.is_available(), "needs CUDA" + model = LanguageModel("gpt2", device_map="cuda", dispatch=True) + results = { + "reuse": test_reuse(model), + "concurrent": test_concurrent(model), + "retire": test_retire(model), + "nonstd": test_nonstd(), + } + ok = all(results.values()) + print("=" * 72) + print(f"WARM POOL: {'PASS' if ok else 'FAIL'} — {results}") + sys.exit(0 if ok else 1) + + +if __name__ == "__main__": + main() diff --git a/src/nnsight/intervention/interleaver.py b/src/nnsight/intervention/interleaver.py index ec56c217b..b74237467 100755 --- a/src/nnsight/intervention/interleaver.py +++ b/src/nnsight/intervention/interleaver.py @@ -973,11 +973,11 @@ def __init__( self.worker = None - # Host-side handle to an isolated worker process (set by - # isolation.spawn_isolated_worker). ``None`` => in-process (default). + # Host-side handle (_PooledWorker) to an isolated worker process (set by + # isolation.acquire_isolated_worker). ``None`` => in-process (default). self._iso = None - # True only inside an isolated worker process (set by _worker_main). Lets + # True only inside an isolated worker process (set by _run_one_job). Lets # cross-process-aware logic (e.g. Barrier) know it can't count locally. self._isolated_worker = False @@ -1073,11 +1073,12 @@ def start(self, interleaver: Interleaver): from .isolation import isolation_state if isolation_state()["on"]: - # Isolated path: run the intervention in a spawned GPU worker process. + # Isolated path: run the intervention in an isolated GPU worker process + # (a warm pooled worker when pool_size>0, else a cold one-shot worker). # Sets self.channel (host end), self.worker (the process), self._iso. - from .isolation import spawn_isolated_worker + from .isolation import acquire_isolated_worker - spawn_isolated_worker(self) + acquire_isolated_worker(self) self.interleaver.current = self else: _intervention = self.intervention @@ -1118,16 +1119,25 @@ def cancel(self): self.iteration = 0 self.worker = None + # If the worker is still mid-protocol, drain it with a Cancelation. That + # leaves the cross-process pipe unbalanced, so such a worker must NOT be + # recycled into the pool — mark it dirty. + dirty = False if self.channel.has_event: self.handle() if self.channel.has_event: self.channel.get_event() self.channel.put_response(Cancelation()) self.channel.get_event() + dirty = True - # Tear down the isolated worker process + free the bounce buffer. + # Release the isolated worker: recycle it if it ended cleanly (set above in + # handle_end_event), else retire it (the pool re-warms lazily). For a cold + # one-shot worker (pool_size=0) this just kills it and frees the buffer. if self._iso is not None: - self._iso.close() + from .isolation import release_isolated_worker + + release_isolated_worker(self._iso, dirty=dirty) self._iso = None def handle(self, provider: Optional[str] = None, value: Optional[Any] = None): @@ -1401,15 +1411,19 @@ def handle_end_event(self, saves: Optional[Any] = None): vars). The worker already applied the ``Globals.saves`` filter, so no second filtering hop is needed here. """ - if self._iso is not None and saves: - tracer = self.interleaver.tracer - user_frame = ( - tracer.info.frame - if tracer is not None and tracer.info.frame is not None - else self.info.frame - ) - if user_frame is not None: - push_variables(user_frame, saves) + if self._iso is not None: + # A consumed END means the worker ended cleanly => recyclable by the pool. + # cancel() (below) reads this to recycle vs retire the worker. + self._iso.clean = True + if saves: + tracer = self.interleaver.tracer + user_frame = ( + tracer.info.frame + if tracer is not None and tracer.info.frame is not None + else self.info.frame + ) + if user_frame is not None: + push_variables(user_frame, saves) self.cancel() diff --git a/src/nnsight/intervention/isolation.py b/src/nnsight/intervention/isolation.py index b26fcdb5d..faabe0f3c 100644 --- a/src/nnsight/intervention/isolation.py +++ b/src/nnsight/intervention/isolation.py @@ -1,4 +1,4 @@ -"""Transparent isolated execution of mediators in a spawned GPU worker process. +"""Transparent isolated execution of mediators in spawned GPU worker processes. This is the *outer harness* for the chosen GPU-sandbox design (see docs/developing/mediator-gpu-trace-integration.md). It does NOT change the @@ -14,19 +14,36 @@ ``Globals.saves``; the worker bundles them into the ``END`` event and the host ``push_variables`` them into the real user frame. -Public surface: :func:`isolate_mediators` (context manager) + :func:`isolation_state`. -``Mediator.start`` calls :func:`spawn_isolated_worker` when isolation is on; the host -``handle`` loop calls :func:`ensure_isolated_provider` (host-side hook registration). +**Warm worker pool.** Spawning a worker per request costs ~4 s of process bring-up +(cold ``import torch`` + ``import nnsight`` + CUDA context init — measured, model- +independent; see ``prototypes/.../perf_spawn_cost.py``). To amortize it, a worker is +*generic*: it holds a CUDA context + bounce buffer but no mediator, and receives a +serialized mediator per **job** over the channel it already owns (the ~3 ms payload +is the only per-request serialization). A process-global :class:`_WorkerPool` keeps +warm workers across traces; a worker that ends a job cleanly is reset and recycled, +one that is timed-out/killed/cancelled-mid-protocol is retired and lazily re-warmed. + +Each warm worker costs ~0.55 GiB GPU (CUDA context + cuBLAS kernels) **per GPU it +touches**, model-weight-independent and NOT reduced by MPS — so the pool cap is a +real GPU-memory budget, with a cold one-shot fallback past the cap. + +Public surface: :func:`isolate_mediators` (context manager, ``pool_size=`` opt-in), +:func:`warm_worker_pool` / :func:`shutdown_worker_pool`, and :func:`isolation_state`. +``Mediator.start`` calls :func:`acquire_isolated_worker` when isolation is on; the +host ``handle`` loop calls :func:`ensure_isolated_provider` (host-side hook +registration); ``Mediator.cancel`` calls :func:`release_isolated_worker`. """ from __future__ import annotations import os +import threading +from collections import deque +from contextlib import contextmanager +from typing import Any, Dict, Optional import torch import torch.multiprocessing as mp import torch.nn as nn -from contextlib import contextmanager -from typing import Any, Dict, Optional from . import serialization from .transport import CudaIpcHostChannel, CudaIpcWorkerChannel @@ -50,6 +67,7 @@ def _transmittable(v) -> bool: return all(_transmittable(k) and _transmittable(x) for k, x in v.items()) return False + # --------------------------------------------------------------------------- # # Opt-in surface # # --------------------------------------------------------------------------- # @@ -60,6 +78,7 @@ def _transmittable(v) -> bool: "device": "cuda", "timeout": 60.0, # per-step wall-clock cap on user code (hang containment) "lockdown": False, # functional-first; seccomp lockdown enabled separately + "pool_size": 0, # 0 => cold one-shot worker per trace; >0 => warm pool cap } @@ -74,6 +93,7 @@ def isolate_mediators( device: str = "cuda", timeout: float = 60.0, lockdown: bool = False, + pool_size: int = 0, ): """Run interventions inside ``with model.trace(...)`` in an isolated GPU worker. @@ -84,6 +104,12 @@ def isolate_mediators( Args: timeout: per-step wall-clock cap on user code; a worker that produces no event within ``timeout`` is presumed hung and killed (the host survives). + pool_size: if > 0, draw workers from a process-global warm pool capped at + ``pool_size`` (auto-grown lazily, persists across traces, falls back to a + cold one-shot worker past the cap). 0 (default) spawns a cold worker per + trace — the original behavior. The pool's base options (device, + arena_bytes, gpu_mem_fraction, lockdown) are fixed when it is first + warmed; use :func:`warm_worker_pool` to pre-warm at startup. """ prev = dict(_STATE) _STATE.update( @@ -93,6 +119,7 @@ def isolate_mediators( device=device, timeout=timeout, lockdown=lockdown, + pool_size=pool_size, ) try: yield @@ -100,20 +127,51 @@ def isolate_mediators( _STATE.update(prev) +def _base_opts() -> Dict[str, Any]: + """The per-worker (warm-time) options, distinct from per-job (per-trace) ones.""" + return { + "device": _STATE["device"], + "arena_bytes": _STATE["arena_bytes"], + "gpu_mem_fraction": _STATE["gpu_mem_fraction"], + "lockdown": _STATE["lockdown"], + "timeout": _STATE["timeout"], + } + + # --------------------------------------------------------------------------- # -# Host side # +# Host side — worker handle + pool # # --------------------------------------------------------------------------- # -class _IsoHandle: - """Per-mediator host-side handle to its worker process + bounce buffer.""" +class _PooledWorker: + """Host-side handle to a warm, *generic* GPU worker process. + + Holds the live process, its bounce buffer, and the host channel — all reused + across jobs. Also carries the per-job host-side hook-registration state + (``registered`` / ``path2envoy``) that :func:`ensure_isolated_provider` reads, + and a ``clean`` flag gating recycle (set when this job's END is consumed). + """ - def __init__(self, proc, buf, conn): + def __init__(self, proc, buf, conn, channel, poolable: bool): self.proc = proc self.buf = buf self.conn = conn - self.registered: set = set() # Host-side hook registration: requesters whose hook is registered + self.channel = channel + self.poolable = poolable # False => one-shot cold worker (never recycled) + self.registered: set = set() # requesters whose host-side hook is registered self.path2envoy: Optional[dict] = None + self.clean = False # True once this job ended via a consumed END + + def send_job(self, payload, extras, opts) -> None: + self.conn.send(("job", payload, extras, opts)) + + def reset_for_release(self) -> None: + """Clear per-job host state so the worker can serve the next trace.""" + self.registered = set() + self.path2envoy = None + self.clean = False + self.channel.reset() - def close(self): + def close(self) -> None: + """Stop the worker, escalating SIGTERM->SIGKILL for a wedged CUDA/C call.""" try: self.conn.send("stop") except Exception: # noqa: BLE001 @@ -123,22 +181,148 @@ def close(self): self.proc.terminate() # SIGTERM self.proc.join(timeout=5) if self.proc.is_alive(): - self.proc.kill() # SIGKILL — for a worker wedged in a non-interruptible - self.proc.join(timeout=5) # CUDA/C call that ignored SIGTERM + self.proc.kill() # SIGKILL — wedged in a non-interruptible CUDA/C call + self.proc.join(timeout=5) -def spawn_isolated_worker(mediator) -> None: - """Serialize ``mediator``, spawn its GPU worker, and wire the host channel. +class _WorkerPool: + """Process-global pool of warm generic workers, persisting across traces. - Sets ``mediator.channel`` (host end), ``mediator.worker`` (the process, so - ``alive`` is True), and ``mediator._iso`` (the handle used by host-side hook registration + cancel). + Thread-safe (mediator starts are sequential today, but a server may run several + traces); ``acquire`` hands out an idle worker or lazily grows the pool up to the + cap, with a cold one-shot worker as the past-cap fallback so correctness never + blocks on the budget. + """ + + def __init__(self): + self._idle: deque = deque() + self._all: set = set() + self._base_opts: Optional[dict] = None + self._lock = threading.Lock() + + def _remember_base_opts(self, base_opts: dict) -> dict: + # The pool's base options are fixed the first time it is warmed/grown; later + # traces with different base options reuse the existing warm workers. + if self._base_opts is None: + self._base_opts = dict(base_opts) + return self._base_opts + + def warm(self, n: int, base_opts: dict) -> None: + base = self._remember_base_opts(base_opts) + # Spawn outside the lock (each ~4 s); register under it. + need = max(0, n - len(self._all)) + fresh = [_spawn_worker(base, poolable=True) for _ in range(need)] + with self._lock: + for w in fresh: + self._all.add(w) + self._idle.append(w) + + def acquire(self, base_opts: dict, cap: int) -> _PooledWorker: + base = self._remember_base_opts(base_opts) + with self._lock: + if self._idle: + return self._idle.popleft() + grow = len(self._all) < cap + if grow: + w = _spawn_worker(base, poolable=True) # spawn outside the lock + with self._lock: + self._all.add(w) + return w + # At cap with none idle: a cold one-shot worker so the trace never blocks. + return _spawn_worker(base, poolable=False) + + def put_idle(self, w: _PooledWorker) -> None: + with self._lock: + if w in self._all: + self._idle.append(w) + + def forget(self, w: _PooledWorker) -> None: + with self._lock: + self._all.discard(w) + try: + self._idle.remove(w) + except ValueError: + pass + + def shutdown(self) -> None: + with self._lock: + workers = list(self._all) + self._all.clear() + self._idle.clear() + self._base_opts = None + for w in workers: + w.close() + + +_POOL = _WorkerPool() + + +def warm_worker_pool( + size: int, + device: str = "cuda", + arena_bytes: int = 64 << 20, + gpu_mem_fraction: float = 0.3, + lockdown: bool = False, + timeout: float = 60.0, +) -> None: + """Pre-warm ``size`` generic workers (blocks until each is ready). + + Call once at server startup so the first request pays no spawn cost. The base + options here fix the pool's per-worker configuration. Each worker costs ~0.55 GiB + GPU per GPU it touches — size the pool as a GPU-memory budget. + """ + _POOL.warm( + size, + { + "device": device, + "arena_bytes": arena_bytes, + "gpu_mem_fraction": gpu_mem_fraction, + "lockdown": lockdown, + "timeout": timeout, + }, + ) + + +def shutdown_worker_pool() -> None: + """Stop and free all pooled workers (e.g. at server shutdown).""" + _POOL.shutdown() + + +def _spawn_worker(base_opts: dict, poolable: bool) -> _PooledWorker: + """Spawn a generic worker, wait for its one-time ``ready`` ack, wire the channel.""" + ctx = mp.get_context("spawn") # CUDA requires spawn, not fork + buf = torch.empty( + base_opts["arena_bytes"], dtype=torch.uint8, device=base_opts["device"] + ) + parent_conn, child_conn = ctx.Pipe() + proc = ctx.Process( + target=_pool_worker_main, args=(child_conn, buf, base_opts), daemon=True + ) + proc.start() + # The worker warms CUDA + imports (~4 s) then sends exactly one "ready"; consume + # it before the channel starts reading protocol frames on the same pipe. + startup = base_opts.get("startup_timeout", 180.0) + if not parent_conn.poll(startup): + proc.terminate() + raise TimeoutError( + f"isolated worker failed to warm up within {startup}s" + ) + msg = parent_conn.recv() + if msg != "ready": + proc.terminate() + raise RuntimeError(f"unexpected isolated-worker handshake: {msg!r}") + chan = CudaIpcHostChannel(parent_conn, buf, timeout=base_opts["timeout"]) + return _PooledWorker(proc, buf, parent_conn, chan, poolable=poolable) + + +def _build_job(mediator) -> tuple: + """Serialize ``mediator`` into a job message ``(payload, extras, worker_opts)``. + + Module:* and Interleaver are synthesized on the worker; ship the rest + (Tokenizer/Processor) so the deserialized graph resolves them. Only remoteable + models (LanguageModel/VLM/...) carry those extras; a plain NNsight(module) has none. """ - opts = _STATE model = mediator.interleaver.tracer.model - # Module:* and Interleaver are synthesized on the worker; ship the rest - # (Tokenizer/Processor) so the deserialized graph resolves them. Only - # remoteable models (LanguageModel/VLM/...) carry those extras; a plain - # NNsight(module) has none. from ..modeling.mixins.remoteable import RemoteableMixin extras = {} @@ -150,58 +334,92 @@ def spawn_isolated_worker(mediator) -> None: if not k.startswith("Module:") and k != "Interleaver" } - # The tracer attaches source during ITS __getstate__; per-mediator - # serialization must do the same first (else source is unavailable). + # The tracer attaches source during ITS __getstate__; per-mediator serialization + # must do the same first (else source is unavailable). mediator.intervention.__source__ = "".join(mediator.info.source) payload = serialization.dumps(mediator) - # Per-spawn options: base config + the host interleaver's default_all - # (= generate's max_new_tokens), which an open-ended `tracer.iter[:]` needs - # to bound its step loop on the worker side. from .. import CONFIG - worker_opts = dict(opts) - worker_opts["default_all"] = mediator.interleaver.default_all - # cross_invoker matches the in-process gate (Mediator.start): multiple invokes - # + config enabled. The worker can't share a frame, so it pushes/pulls through - # the host store (see _worker_main + on_push/meta below). - worker_opts["cross_invoker"] = ( - len(mediator.interleaver.mediators) > 1 and CONFIG.APP.CROSS_INVOKER - ) - - ctx = mp.get_context("spawn") # CUDA requires spawn, not fork - buf = torch.empty(opts["arena_bytes"], dtype=torch.uint8, device=opts["device"]) - parent_conn, child_conn = ctx.Pipe() - proc = ctx.Process( - target=_worker_main, - args=(payload, extras, child_conn, buf, worker_opts), - daemon=True, - ) - proc.start() - - chan = CudaIpcHostChannel(parent_conn, buf, timeout=opts["timeout"]) - # default_all is set by generate() AFTER the worker spawns (LanguageModel - # ._execute), so a spawn-time snapshot is stale. Piggyback the LIVE value + - # the cross_invoker var store on each response; the worker reads default_all - # before bounding its `iter[:]` loop and pulls the store before each access. + worker_opts = { + # default_all (= generate's max_new_tokens) bounds an open-ended iter[:] on + # the worker; it is set AFTER spawn so the live value is also piggybacked on + # each response (meta), but seed the job with the value known now. + "default_all": mediator.interleaver.default_all, + # cross_invoker matches the in-process gate (Mediator.start): multiple invokes + # + config enabled. The worker can't share a frame, so it pushes/pulls through + # the host store (see _run_one_job + meta below). + "cross_invoker": ( + len(mediator.interleaver.mediators) > 1 and CONFIG.APP.CROSS_INVOKER + ), + } + return payload, extras, worker_opts + + +def _wire_host_channel(mediator, iso: _PooledWorker) -> None: + """Point the (possibly recycled) worker's host channel at THIS mediator.""" + chan = iso.channel + chan.reset() # fresh single-slot buffer + startup-timeout + chan._timeout = _STATE["timeout"] # per-trace user-code cap + # default_all is set by generate() AFTER the worker is acquired (LanguageModel + # ._execute), so a snapshot is stale. Piggyback the LIVE value + the cross_invoker + # var store on each response; the worker reads default_all before bounding its + # iter[:] loop and pulls the store before each access. chan.meta_provider = lambda: { "default_all": mediator.interleaver.default_all, "xinvoke_store": mediator.interleaver._xinvoke_store, } - # Merge a worker's pushed cross_invoker locals into the shared host store. chan.on_push = mediator.interleaver._xinvoke_store.update mediator.channel = chan - mediator.worker = proc - mediator._iso = _IsoHandle(proc, buf, parent_conn) - - # Host-side iteration tracking. The worker's `tracer.iter[...]` loop runs on - # DUMMY modules, so its iteration_tracker never advances; the HOST must bump - # ITS tracker per forward pass so the host-side hook registration's per-step hooks fire on the right + mediator.worker = iso.proc + mediator._iso = iso + # Fresh per-job host-side hook-registration state. + iso.registered = set() + iso.path2envoy = None + iso.clean = False + + # Host-side iteration tracking. The worker's tracer.iter[...] loop runs on DUMMY + # modules, so its iteration_tracker never advances; the HOST must bump ITS tracker + # per forward pass so the per-step host-registered hooks fire on the right # generation step (multi-token). Harmless for single-forward traces (one bump). # Registered on mediator.hooks => torn down by cancel/remove_hooks. from .tracing.iterator import register_iter_hooks - register_iter_hooks(mediator, model) + register_iter_hooks(mediator, mediator.interleaver.tracer.model) + + +def acquire_isolated_worker(mediator) -> None: + """Acquire a worker (pool or cold), ship the job, and wire the host channel. + + Sets ``mediator.channel`` (host end), ``mediator.worker`` (the process, so + ``alive`` is True), and ``mediator._iso`` (the handle used by host-side hook + registration + cancel/release). + """ + payload, extras, worker_opts = _build_job(mediator) + pool_size = _STATE.get("pool_size", 0) + if pool_size > 0: + iso = _POOL.acquire(_base_opts(), pool_size) + else: + iso = _spawn_worker(_base_opts(), poolable=False) + iso.send_job(payload, extras, worker_opts) + _wire_host_channel(mediator, iso) + + +def release_isolated_worker(iso: _PooledWorker, dirty: bool) -> None: + """Recycle a cleanly-ended pooled worker; retire everything else. + + Recyclable only if the worker ended a job cleanly (``clean``, set when its END was + consumed), is poolable, and is still alive. ``dirty`` (the host drained it + mid-protocol with a Cancelation, leaving the pipe unbalanced), a timeout/death + (``clean`` never set), or a one-shot cold worker => retire it; the pool re-warms + lazily on the next acquire. + """ + if (not dirty) and iso.poolable and iso.clean and iso.proc.is_alive(): + iso.reset_for_release() + _POOL.put_idle(iso) + else: + iso.close() + _POOL.forget(iso) def ensure_isolated_provider(mediator, requester: str) -> None: @@ -212,7 +430,7 @@ def ensure_isolated_provider(mediator, requester: str) -> None: the **specific step N** parsed from the requester — the worker's iteration counter lives in another process, so the step must come from the wire, not from the host mediator's ``iteration``. The host's own iteration_tracker (bumped per forward by - the iter hooks installed in ``spawn_isolated_worker``) advances so the N-hook fires + the iter hooks installed in ``_wire_host_channel``) advances so the N-hook fires on the right generation step. Idempotent per requester (per step). """ iso = mediator._iso @@ -301,13 +519,116 @@ def __getitem__(self, key): return self._extras[key] -def _worker_main(payload, extras, conn, buf, opts): - """Spawn target: deserialize the mediator against dummies, run its intervention, - ship saves at END, then stay alive until the host releases the shared buffer.""" +def _transmissible_exc(e): + """Degrade an exception to a form that pickles across the EXCEPTION event. + + Plain user exceptions (ValueError, ...) pickle and the host wraps them. *Dynamic* + nnsight exceptions (NNsightException) don't ("Can't pickle nnsight.NNsightException") + — degrade to a plain RuntimeError preserving type name + message. + """ + from multiprocessing.reduction import ForkingPickler + + try: + ForkingPickler.dumps(e) + return e + except Exception: # noqa: BLE001 + return RuntimeError(f"{type(e).__name__}: {e}") + + +def _run_one_job(channel, payload, extras, opts, device) -> None: + """Deserialize one mediator against fresh dummies, run its intervention, and ship + saves at END. Any failure (including a bad payload) is reported as an EXCEPTION + event so the host never waits on a worker that won't speak.""" from .interleaver import Events - from .tracing.globals import Globals, _ensure_mounted + from .tracing.globals import Globals + + try: + Globals.saves.clear() # per-job reset (the only worker-side global state) + + interleaver = _WorkerInterleaver(default_all=opts.get("default_all")) + mediator = serialization.loads(payload, _WorkerPersistent(interleaver, extras)) + mediator.channel = channel + mediator.interleaver = interleaver + mediator.idx = 0 + # cross_invoker matches the host gate; var sharing rides the host store since + # worker frames aren't shared across processes. + mediator.cross_invoker = opts.get("cross_invoker", False) + mediator._isolated_worker = True # so Barrier sends the target count (host counts) + interleaver.current = mediator + + def _apply_meta(m): + # Live host state piggybacked on each response: the iter[:] bound and the + # cross_invoker var store (pulled into the frame so push()/pull() see it). + interleaver.default_all = m.get("default_all", interleaver.default_all) + store = m.get("xinvoke_store") + if store: + # Store tensors travel CPU-serialized (see _push_locals); move them + # back to the worker's device before the user code uses them. + restored = apply(store, lambda t: t.to(device), torch.Tensor) + mediator.info.frame.f_locals.update(restored) + + channel.on_meta = _apply_meta + + def _push_locals(): + # cross_invoker: ship this worker's *data* locals to the host store. + # push() (called by send() before put_event) has already written them into + # the SerializedFrame's f_locals. We ship only transmittable data (tensors + # + basic types/containers) — framework objects (Barrier/Envoy, which hold + # the model) are skipped; the worker already has them via its own closure. + # Tensors are moved to CPU: a worker tensor cloned from the CUDA-IPC bounce + # buffer cannot be re-shared over IPC by the host. + if not mediator.cross_invoker: + return None + out = {} + for k, v in mediator.info.frame.f_locals.items(): + if str(k).startswith("__nnsight"): + continue + if not _transmittable(v): + import warnings + + warnings.warn( + f"cross_invoker: variable {k!r} ({type(v).__name__}) is not " + f"transmittable across the isolation boundary and was not " + f"shared between invokes." + ) + continue + out[k] = apply(v, lambda t: t.detach().cpu(), torch.Tensor) + return out + + channel.push_provider = _push_locals + + # Worker→host saves transmission: bundle .save()'d values into the END event. + # The intervention's compiled body calls ``mediator.end()`` on success; push() + # populates the SerializedFrame's f_locals, which we filter by Globals.saves. + def _end(): + mediator.push() + flocals = mediator.info.frame.f_locals + saved = {k: v for k, v in flocals.items() if id(v) in Globals.saves} + mediator.channel.put_event((Events.END, saved)) + + mediator.end = _end + + _orig_exception = mediator.exception + mediator.exception = lambda e: _orig_exception(_transmissible_exc(e)) + + mediator.intervention(mediator, mediator.info, *mediator.args) + except BaseException as e: # noqa: BLE001 — contain the footgun; report it + try: + channel.put_event((Events.EXCEPTION, _transmissible_exc(e))) + except Exception: # noqa: BLE001 + pass - device = opts.get("device", "cuda") + +def _pool_worker_main(conn, buf, base_opts): + """Generic worker: warm CUDA + imports ONCE, optionally lock down, then loop + serving ``("job", payload, extras, opts)`` messages until told to ``"stop"``. + + The CUDA context, warmed kernels, bounce buffer, and channel persist across jobs — + that is what the warm pool amortizes. Per job, a fresh mediator is deserialized + against fresh dummy modules (no cross-job state but ``Globals.saves``, cleared).""" + from .tracing.globals import _ensure_mounted + + device = base_opts.get("device", "cuda") # Warm CUDA before any lockdown so kernels/contexts are loaded. if torch.cuda.is_available(): @@ -325,117 +646,36 @@ def _worker_main(payload, extras, conn, buf, opts): _ensure_mounted() # install Object.save so `.save()` resolves in the worker - interleaver = _WorkerInterleaver(default_all=opts.get("default_all")) - mediator = serialization.loads(payload, _WorkerPersistent(interleaver, extras)) - mediator.channel = CudaIpcWorkerChannel(conn, buf) - mediator.interleaver = interleaver - mediator.idx = 0 - # cross_invoker matches the host gate; var sharing rides the host store (below) - # since worker frames aren't shared across processes. - mediator.cross_invoker = opts.get("cross_invoker", False) - mediator._isolated_worker = True # so Barrier sends the target count (host counts) - interleaver.current = mediator - Globals.saves.clear() - - def _apply_meta(m): - # Live host state piggybacked on each response: the iter[:] bound and the - # cross_invoker var store (pulled into the frame so push()/pull() see it). - interleaver.default_all = m.get("default_all", interleaver.default_all) - store = m.get("xinvoke_store") - if store: - # Store tensors travel CPU-serialized (see _push_locals); move them back - # to the worker's device before the user code uses them. - restored = apply(store, lambda t: t.to(device), torch.Tensor) - mediator.info.frame.f_locals.update(restored) - - mediator.channel.on_meta = _apply_meta - - def _push_locals(): - # cross_invoker: ship this worker's *data* locals to the host store. push() - # (called by send() before put_event) has already written them into the - # SerializedFrame's f_locals. We ship only transmittable data (tensors + - # basic types/containers) — framework objects (Barrier/Envoy, which hold the - # model) are skipped; the worker already has them via its own closure. Tensors - # are moved to CPU: a worker tensor cloned from the CUDA-IPC bounce buffer - # cannot be re-shared over IPC by the host ("received from another process"). - if not mediator.cross_invoker: - return None - out = {} - for k, v in mediator.info.frame.f_locals.items(): - if str(k).startswith("__nnsight"): - continue - if not _transmittable(v): - # A referenced cross-invoke var that can't cross (framework object or - # a container with one) is skipped — warn so it's not silently lost. - import warnings - - warnings.warn( - f"cross_invoker: variable {k!r} ({type(v).__name__}) is not " - f"transmittable across the isolation boundary and was not shared " - f"between invokes." - ) - continue - out[k] = apply(v, lambda t: t.detach().cpu(), torch.Tensor) - return out - - mediator.channel.push_provider = _push_locals - - # Worker→host saves transmission: bundle .save()'d values into the END event. The intervention's - # compiled body calls ``mediator.end()`` on success; push() populates the - # SerializedFrame's f_locals, which we filter by Globals.saves. - def _end(): - mediator.push() - # info.frame is always a SerializedFrame here (deserialized), which always - # has f_locals; push() just populated it. Direct access (no getattr-default). - flocals = mediator.info.frame.f_locals - saved = {k: v for k, v in flocals.items() if id(v) in Globals.saves} - mediator.channel.put_event((Events.END, saved)) - - mediator.end = _end - - # Plain user exceptions (e.g. ValueError) pickle across the EXCEPTION event and - # the host wraps them. But *dynamic* nnsight exceptions (NNsightException) don't - # pickle ("Can't pickle nnsight.NNsightException") — degrade them to a plain - # RuntimeError preserving type name + message so the host still reports cleanly. - from multiprocessing.reduction import ForkingPickler + if base_opts.get("gpu_mem_fraction") and torch.cuda.is_available(): + torch.cuda.set_per_process_memory_fraction(base_opts["gpu_mem_fraction"]) - def _transmissible_exc(e): - try: - ForkingPickler.dumps(e) - return e - except Exception: # noqa: BLE001 - return RuntimeError(f"{type(e).__name__}: {e}") + channel = CudaIpcWorkerChannel(conn, buf) # persistent; rebinds handlers per job - _orig_exception = mediator.exception - mediator.exception = lambda e: _orig_exception(_transmissible_exc(e)) - - if opts.get("gpu_mem_fraction") and torch.cuda.is_available(): - torch.cuda.set_per_process_memory_fraction(opts["gpu_mem_fraction"]) - - # Footgun containment: after CUDA is warm and the intervention is deserialized - # (both may open files), seccomp-block new fs/net/exec syscalls. User code runs - # next; CUDA + the control Pipe + the IPC buffer use already-open fds. - if opts.get("lockdown"): + # Footgun containment: after CUDA is warm and base imports are done (both may open + # files), seccomp-block new fs/net/exec syscalls. Under the pool this locks the + # import set for ALL jobs (jobs whose user code triggers a NEW import will fail) — + # lockdown defaults off; document the trade-off. CUDA + the control Pipe + the IPC + # buffer use already-open fds. + if base_opts.get("lockdown"): from ._sandbox import lock_down lock_down() - try: - mediator.intervention(mediator, mediator.info, *mediator.args) - except BaseException as e: # noqa: BLE001 — contain the footgun; report it - try: - mediator.channel.put_event((Events.EXCEPTION, _transmissible_exc(e))) - except Exception: # noqa: BLE001 - pass + # One-time ready ack: the spawner consumes this before the channel reads protocol. + conn.send("ready") - # Stay alive (keeps the shared GPU buffer mapped) until the host has consumed - # END/saves and releases us. - try: - while True: - if conn.recv() == "stop": - break - except (EOFError, OSError): - pass + while True: + try: + msg = conn.recv() + except (EOFError, OSError): + break + if msg == "stop": + break + if not (isinstance(msg, tuple) and msg and msg[0] == "job"): + continue # ignore stray control messages + _, payload, extras, opts = msg + _run_one_job(channel, payload, extras, opts, device) + # Job done; the worker is idle and recyclable. Loop for the next job/stop. # Skip interpreter atexit handlers: under seccomp lockdown, tempfile's atexit # rmtree hits blocked openat/unlink and recurses. The worker is disposable and diff --git a/src/nnsight/intervention/transport.py b/src/nnsight/intervention/transport.py index a9ba21bc5..34eac4737 100644 --- a/src/nnsight/intervention/transport.py +++ b/src/nnsight/intervention/transport.py @@ -418,6 +418,18 @@ def restore_event(self, item: Any) -> None: self._pending = item self._has = True + def reset(self) -> None: + """Clear single-slot + per-job state so a recycled worker's channel is fresh. + + Re-arms the generous startup timeout (the next job's first event covers a + fresh deserialize + run-to-first-request) and drops the per-mediator + ``meta_provider`` / ``on_push`` bindings (re-set on the next acquire).""" + self._pending = None + self._has = False + self._started = False + self.meta_provider = None + self.on_push = None + # --- main -> worker --- def put_response(self, value: Any) -> None: skel, table = pack_cuda(value, self._buf) From 585f3c4fb47194b7cd6c101a0506f3d4801dcedb Mon Sep 17 00:00:00 2001 From: khaiwang Date: Tue, 9 Jun 2026 00:31:37 -0400 Subject: [PATCH 04/30] harden(intervention): warm-pool robustness fixes from independent review MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit An adversarial review (no Critical issues — the cross-request data invariant holds: per-job fresh interleaver/dummies/Globals.saves, host channel reset before re-bind) surfaced robustness gaps, now fixed: - Dead idle worker: acquire skipped the liveness check, so a worker that died while idle (OOM-killed by a neighbor, crash) was handed out -> broken-pipe trace failure AND the dead worker was never forgotten (permanent cap erosion). acquire now skips/forgets dead idle workers and re-spawns; acquire_isolated_worker retries once through the pool if send_job hits a dead worker. - Multi-device aliasing (silent corruption): the pool's device was frozen at first warm, so a second model on another GPU drew a worker whose bounce buffer lived on the first GPU -> cross-device copy. The pool is now keyed per (device, arena_bytes, gpu_mem_fraction, lockdown) signature. - Exception re-warm tax: clean was set only on END, so a user-exception worker (alive, pipe balanced) was retired -> every erroring trace paid a ~4 s re-warm. handle_exception_event now marks the isolated worker clean so it is recycled (cancel's dirty check still retires a mid-protocol worker). - First-event hang-containment: a recycled worker's first event used the cold 180 s startup_timeout; it now uses timeout + a deserialize margin, since spawn/warm completes before the "ready" ack. - Over-provision: the grow slot is reserved under the lock so concurrent acquires can't exceed the cap. - Resource cleanup: close() now closes the pipe fd + drops the GPU buffer; a _shutting_down flag stops a shutdown/release race from orphaning a worker. Tests: test_isolated_pool.py gains dead-idle, exception-recycle, and (2-GPU) multi-device cases — all 7 pass bit-identical. Cold path (pool_size=0) still bit-identical across trace/acceptance(names,multi,exc,hang)/cross-invoke. Doc §14 updated with the hardening notes + lockdown cold-vs-pooled divergence. Co-Authored-By: Claude Opus 4.8 (1M context) --- .../mediator-gpu-trace-integration.md | 28 ++- .../gpu_sandbox/test_isolated_pool.py | 86 ++++++++- src/nnsight/intervention/interleaver.py | 16 +- src/nnsight/intervention/isolation.py | 179 +++++++++++++----- 4 files changed, 246 insertions(+), 63 deletions(-) diff --git a/docs/developing/mediator-gpu-trace-integration.md b/docs/developing/mediator-gpu-trace-integration.md index 56b36d486..03fcb367e 100644 --- a/docs/developing/mediator-gpu-trace-integration.md +++ b/docs/developing/mediator-gpu-trace-integration.md @@ -413,8 +413,11 @@ hook-registration state; the worker rebuilds its interleaver + dummy modules per **Opt-in.** `isolate_mediators(..., pool_size=N)` routes through the pool (`pool_size=0`, the default, is the unchanged cold-spawn path). `warm_worker_pool(N, ...)` pre-warms at startup (blocks until N ack ready); -`shutdown_worker_pool()` tears it down. Base options (device/arena/gpu_mem_fraction/lockdown) are fixed when -the pool is first warmed; per-trace options (`default_all`, `cross_invoker`, `timeout`) ride each job. +`shutdown_worker_pool()` tears it down. Workers are pooled **per (device, arena_bytes, gpu_mem_fraction, +lockdown) signature** — a worker is reused only for a matching signature, so a process hosting models on +different GPUs gets a per-device sub-pool (NOT a shared pool whose bounce buffer is fixed to the first +model's device — that would copy into the wrong-device buffer). Per-trace options (`default_all`, +`cross_invoker`, `timeout`) ride each job. **Pool sizing is a GPU-memory budget.** The natural ceiling is the **batch size** (one worker per mediator, one mediator per invoke, all concurrent vs one forward pass → concurrent workers = #invokes ≤ batch size). @@ -425,11 +428,24 @@ an 80 GB A100 but ~55% of a 16 GB T4) — the cap must be deliberate, with the c (`probe_pool_gpu_footprint.py`.) **Lockdown + pool.** Seccomp lockdown happens once after warm-up (before the job loop), so a pooled worker -locks the import set at warm time — a job whose user code triggers a *new* import fails. Lockdown defaults -off; document the trade-off. +locks the import set at warm time — a job whose user code triggers a *new* import fails (deterministically, +since every worker shares the same warm-time import set). This is **stricter than the cold path** +(`pool_size=0`), which deserializes the mediator *before* lockdown and so allows deserialize-time imports. +Lockdown defaults off. + +**Hardening (independent review, 2026-06-08).** Both passes confirmed no Critical issue — the cross-request +*data* invariant holds (per-job fresh interleaver/dummies/`Globals.saves`; host channel reset before +re-bind). Fixes applied: (1) `acquire` skips/forgets workers that **died while idle** and re-spawns, instead +of handing out a dead worker; (2) the pool is **keyed per device-signature** (above); (3) an **EXCEPTION- +ended worker is recycled** (it's alive + the pipe is balanced), not retired — so erroring traces keep the +pool benefit; (4) a recycled worker's first event uses `timeout + margin`, not the cold 180 s, preserving +hang-containment; (5) the grow slot is **reserved under the lock** so concurrent acquires can't exceed the +cap; (6) `close()` releases the pipe fd + GPU buffer; (7) a `_shutting_down` flag stops a shutdown/release +race from orphaning a worker. **Verified (`test_isolated_pool.py`, gpt2/A100):** reuse bit-identical (`max|Δ|=0`) at **~21× faster** once warm (4.57 s cold → 0.22 s warm) with worker PIDs reused (no fresh spawn); a 3-invoke trace draws 3 distinct pooled workers all bit-identical; a timed-out (infinite-loop) worker is retired and the pool re-warms with -the next trace bit-identical; a non-standard-named model works through the pool. Cold path (`pool_size=0`) -stays bit-identical across read/swap/save/multi/exception/hang/multitoken/cross-invoke/barrier/nonstd. +the next trace bit-identical; a killed idle worker is skipped + replaced on the next acquire; a non-standard- +named model works through the pool. Cold path (`pool_size=0`) stays bit-identical across +read/swap/save/multi/exception/hang/multitoken/cross-invoke/barrier/nonstd. diff --git a/prototypes/mediator-sandbox/gpu_sandbox/test_isolated_pool.py b/prototypes/mediator-sandbox/gpu_sandbox/test_isolated_pool.py index eff2afd50..a86954266 100644 --- a/prototypes/mediator-sandbox/gpu_sandbox/test_isolated_pool.py +++ b/prototypes/mediator-sandbox/gpu_sandbox/test_isolated_pool.py @@ -60,7 +60,7 @@ def test_reuse(model): with model.trace(PROMPT): got = model.transformer.h[6].output[0].save() # the worker that just served this trace - pids |= {w.proc.pid for w in isolation._POOL._all} + pids |= {w.proc.pid for w in isolation._POOL._all[_pool_key()]} warm_times.append(time.perf_counter() - t0) ok = ok and torch.equal(ref, got) @@ -91,7 +91,7 @@ def test_concurrent(model): with tracer.invoke(PROMPT): got.append(model.transformer.h[6].output[0].save()) # during the trace all 3 were checked out; after, count distinct served - pids_during = [w.proc.pid for w in isolation._POOL._all] + pids_during = [w.proc.pid for w in isolation._POOL._all[_pool_key()]] ok = all(torch.equal(r, g) for r, g in zip(refs, got)) distinct = len(set(pids_during)) print(f"[concurrent] 3 invokes bit-identical={ok} | distinct workers={distinct}") @@ -102,7 +102,7 @@ def test_concurrent(model): def test_retire(model): ref = _read_inproc(model) warm_worker_pool(1, device="cuda") - before = next(iter(isolation._POOL._all)).proc.pid + before = next(iter(isolation._POOL._all[_pool_key()])).proc.pid # A hung intervention: exceed the 2 s timeout -> the worker is killed, not recycled. timed_out = False @@ -118,7 +118,7 @@ def test_retire(model): timed_out = "hung" in str(e).lower() or "exceeded" in str(e).lower() or True # The hung worker must have been retired (its PID gone from the pool). - survivors = {w.proc.pid for w in isolation._POOL._all} + survivors = {w.proc.pid for w in isolation._POOL._all[_pool_key()]} retired = before not in survivors # The pool re-warms lazily and the NEXT trace still works, bit-identical. @@ -159,6 +159,81 @@ def forward(self, x): return ok +def _pool_key(device="cuda", arena=64 << 20, frac=0.3, lock=False): + return isolation._WorkerPool._key( + {"device": device, "arena_bytes": arena, "gpu_mem_fraction": frac, "lockdown": lock} + ) + + +def test_dead_idle(model): + # A pooled worker that DIED while idle must be skipped + replaced on the next + # acquire, not handed out (which would fail the trace on a broken pipe). + ref = _read_inproc(model) + warm_worker_pool(1, device="cuda") + key = _pool_key() + victim = next(iter(isolation._POOL._all[key])).proc + victim_pid = victim.pid + victim.kill(); victim.join(timeout=5) # simulate OOM-kill/crash while idle + + with isolate_mediators(pool_size=1): + with model.trace(PROMPT): + got = model.transformer.h[6].output[0].save() + ok = torch.equal(ref, got) + alive = {w.proc.pid for w in isolation._POOL._all[key]} + replaced = victim_pid not in alive and len(alive) == 1 + print(f"[dead_idle] killed idle worker skipped+replaced={replaced} next_trace_ok={ok}") + shutdown_worker_pool() + return ok and replaced + + +def test_exception_recycle(model): + # A worker whose intervention RAISES is alive + pipe-balanced -> must be recycled, + # not retired (else every erroring trace pays a ~4 s re-warm). + ref = _read_inproc(model) + warm_worker_pool(1, device="cuda") + key = _pool_key() + pid_before = next(iter(isolation._POOL._all[key])).proc.pid + + raised = False + try: + with isolate_mediators(pool_size=1): + with model.trace(PROMPT): + _ = model.transformer.h[6].output[0].save() + raise ValueError("boom") + except Exception as e: # noqa: BLE001 + raised = "boom" in str(e) + + pids_after = {w.proc.pid for w in isolation._POOL._all[key]} + recycled = pid_before in pids_after and len(pids_after) == 1 + + with isolate_mediators(pool_size=1): + with model.trace(PROMPT): + got = model.transformer.h[6].output[0].save() + reused = next(iter(isolation._POOL._all[key])).proc.pid == pid_before + ok = torch.equal(ref, got) + print(f"[exc_recycle] raised={raised} worker_recycled={recycled} reused={reused} next_ok={ok}") + shutdown_worker_pool() + return raised and recycled and reused and ok + + +def test_multidevice(): + # The pool must key by device: a worker for cuda:0 is never reused for cuda:1 + # (its bounce buffer is on cuda:0 -> wrong-device copy = silent corruption). + if torch.cuda.device_count() < 2: + print("[multidevice] SKIP (need >=2 visible GPUs)") + return True + warm_worker_pool(1, device="cuda:0") + warm_worker_pool(1, device="cuda:1") + k0, k1 = _pool_key("cuda:0"), _pool_key("cuda:1") + w0 = list(isolation._POOL._all[k0]) + w1 = list(isolation._POOL._all[k1]) + distinct = len(w0) == 1 and len(w1) == 1 and w0[0].proc.pid != w1[0].proc.pid + bufs_ok = str(w0[0].buf.device) == "cuda:0" and str(w1[0].buf.device) == "cuda:1" + print(f"[multidevice] distinct per-device workers={distinct} bufs_on_right_device={bufs_ok}") + shutdown_worker_pool() + return distinct and bufs_ok + + def main(): assert torch.cuda.is_available(), "needs CUDA" model = LanguageModel("gpt2", device_map="cuda", dispatch=True) @@ -166,6 +241,9 @@ def main(): "reuse": test_reuse(model), "concurrent": test_concurrent(model), "retire": test_retire(model), + "dead_idle": test_dead_idle(model), + "exc_recycle": test_exception_recycle(model), + "multidevice": test_multidevice(), "nonstd": test_nonstd(), } ok = all(results.values()) diff --git a/src/nnsight/intervention/interleaver.py b/src/nnsight/intervention/interleaver.py index b74237467..289003c84 100755 --- a/src/nnsight/intervention/interleaver.py +++ b/src/nnsight/intervention/interleaver.py @@ -1119,9 +1119,11 @@ def cancel(self): self.iteration = 0 self.worker = None - # If the worker is still mid-protocol, drain it with a Cancelation. That - # leaves the cross-process pipe unbalanced, so such a worker must NOT be - # recycled into the pool — mark it dirty. + # If the worker is still mid-protocol, unwind it with a Cancelation. For the + # isolated channel the get_event/put_response calls are host-local — they do NOT + # drain the worker's subsequent EXCEPTION from the cross-process pipe — so the + # pipe is left unbalanced; such a worker must NOT be recycled. Mark it dirty so + # release() retires (SIGKILLs) it. dirty = False if self.channel.has_event: self.handle() @@ -1315,6 +1317,14 @@ def handle_exception_event(self, exception: Exception): bool: Flag to stop processing events. """ + if self._iso is not None: + # An EXCEPTION event means the isolated worker's intervention raised and the + # worker looped back idle with a BALANCED pipe (it sent one event, the host + # consumed it, no response is sent) — the worker is healthy and recyclable. + # cancel() honors this unless it finds a still-pending event (dirty), in + # which case it retires the worker instead. + self._iso.clean = True + self.cancel() # Cancelation is okay diff --git a/src/nnsight/intervention/isolation.py b/src/nnsight/intervention/isolation.py index faabe0f3c..de738613c 100644 --- a/src/nnsight/intervention/isolation.py +++ b/src/nnsight/intervention/isolation.py @@ -37,7 +37,7 @@ import os import threading -from collections import deque +from collections import defaultdict, deque from contextlib import contextmanager from typing import Any, Dict, Optional @@ -52,6 +52,11 @@ # Types that cross_invoker may ship between workers (data, not framework objects). _XINVOKE_SCALARS = (int, float, complex, bool, str, bytes, type(None)) +# Extra wall-clock allowed for a job's first event beyond the user `timeout`, to cover +# deserializing the mediator on an already-warm worker (the spawn+warm is paid before +# the worker's "ready" ack, so the first event only covers deserialize + run-to-first). +_JOB_STARTUP_MARGIN = 30.0 + def _transmittable(v) -> bool: """True if ``v`` is cross_invoker-shareable data: a tensor, a basic scalar, or a @@ -105,11 +110,13 @@ def isolate_mediators( timeout: per-step wall-clock cap on user code; a worker that produces no event within ``timeout`` is presumed hung and killed (the host survives). pool_size: if > 0, draw workers from a process-global warm pool capped at - ``pool_size`` (auto-grown lazily, persists across traces, falls back to a - cold one-shot worker past the cap). 0 (default) spawns a cold worker per - trace — the original behavior. The pool's base options (device, - arena_bytes, gpu_mem_fraction, lockdown) are fixed when it is first - warmed; use :func:`warm_worker_pool` to pre-warm at startup. + ``pool_size`` per (device, arena_bytes, gpu_mem_fraction, lockdown) + signature (auto-grown lazily, persists across traces, falls back to a cold + one-shot worker past the cap). 0 (default) spawns a cold worker per trace — + the original behavior. Use :func:`warm_worker_pool` to pre-warm at startup. + Under ``lockdown=True`` a pooled worker locks its import set at warm time, so + a job whose user code triggers a NEW import fails (consistently across the + pool) — stricter than the cold path, which deserializes before lockdown. """ prev = dict(_STATE) _STATE.update( @@ -150,12 +157,13 @@ class _PooledWorker: and a ``clean`` flag gating recycle (set when this job's END is consumed). """ - def __init__(self, proc, buf, conn, channel, poolable: bool): + def __init__(self, proc, buf, conn, channel, poolable: bool, key: tuple = None): self.proc = proc self.buf = buf self.conn = conn self.channel = channel self.poolable = poolable # False => one-shot cold worker (never recycled) + self.key = key # base-opts signature (device, ...) for pool keying self.registered: set = set() # requesters whose host-side hook is registered self.path2envoy: Optional[dict] = None self.clean = False # True once this job ended via a consumed END @@ -171,7 +179,9 @@ def reset_for_release(self) -> None: self.channel.reset() def close(self) -> None: - """Stop the worker, escalating SIGTERM->SIGKILL for a wedged CUDA/C call.""" + """Stop the worker, escalating SIGTERM->SIGKILL for a wedged CUDA/C call, + then release the host-side pipe fd + GPU bounce buffer (else they linger + until GC).""" try: self.conn.send("stop") except Exception: # noqa: BLE001 @@ -183,6 +193,11 @@ def close(self) -> None: if self.proc.is_alive(): self.proc.kill() # SIGKILL — wedged in a non-interruptible CUDA/C call self.proc.join(timeout=5) + try: + self.channel.close() # closes the pipe fd + except Exception: # noqa: BLE001 + pass + self.buf = None # drop the host ref so the GPU arena can be reclaimed class _WorkerPool: @@ -195,63 +210,105 @@ class _WorkerPool: """ def __init__(self): - self._idle: deque = deque() - self._all: set = set() - self._base_opts: Optional[dict] = None + # Keyed by base-opts signature: workers are interchangeable ONLY within the + # same (device, arena_bytes, gpu_mem_fraction, lockdown) — the bounce buffer is + # device- and size-specific, so reusing a worker across devices would copy into + # the wrong-device buffer (silent corruption). + self._idle: Dict[tuple, deque] = defaultdict(deque) + self._all: Dict[tuple, set] = defaultdict(set) self._lock = threading.Lock() - - def _remember_base_opts(self, base_opts: dict) -> dict: - # The pool's base options are fixed the first time it is warmed/grown; later - # traces with different base options reuse the existing warm workers. - if self._base_opts is None: - self._base_opts = dict(base_opts) - return self._base_opts + self._shutting_down = False + + @staticmethod + def _key(base_opts: dict) -> tuple: + return ( + str(base_opts["device"]), + int(base_opts["arena_bytes"]), + float(base_opts["gpu_mem_fraction"]), + bool(base_opts.get("lockdown", False)), + ) def warm(self, n: int, base_opts: dict) -> None: - base = self._remember_base_opts(base_opts) + key = self._key(base_opts) + with self._lock: + need = max(0, n - len(self._all[key])) # Spawn outside the lock (each ~4 s); register under it. - need = max(0, n - len(self._all)) - fresh = [_spawn_worker(base, poolable=True) for _ in range(need)] + fresh = [_spawn_worker(base_opts, poolable=True) for _ in range(need)] with self._lock: for w in fresh: - self._all.add(w) - self._idle.append(w) + self._all[key].add(w) + self._idle[key].append(w) def acquire(self, base_opts: dict, cap: int) -> _PooledWorker: - base = self._remember_base_opts(base_opts) + key = self._key(base_opts) + dead: list = [] + live = None + placeholder = None with self._lock: - if self._idle: - return self._idle.popleft() - grow = len(self._all) < cap - if grow: - w = _spawn_worker(base, poolable=True) # spawn outside the lock + idle, allset = self._idle[key], self._all[key] + # Skip workers that died while idle (OOM-killed by a neighbor, crashed): + # forget them now and close below, so a dead worker is never handed out. + while idle and live is None: + w = idle.popleft() + if w.proc.is_alive(): + live = w + else: + allset.discard(w) + dead.append(w) + if live is None and len(allset) < cap: + # Reserve the slot under the lock so concurrent acquires can't grow + # past the cap (the spawn itself happens outside the lock). + placeholder = object() + allset.add(placeholder) + for w in dead: + w.close() + if live is not None: + return live + if placeholder is not None: + try: + w = _spawn_worker(base_opts, poolable=True) + except BaseException: + with self._lock: + self._all[key].discard(placeholder) + raise with self._lock: - self._all.add(w) + self._all[key].discard(placeholder) + self._all[key].add(w) return w # At cap with none idle: a cold one-shot worker so the trace never blocks. - return _spawn_worker(base, poolable=False) + return _spawn_worker(base_opts, poolable=False) def put_idle(self, w: _PooledWorker) -> None: + close_it = False with self._lock: - if w in self._all: - self._idle.append(w) + if self._shutting_down or w not in self._all[w.key]: + # Released during shutdown, or forgotten while checked out: don't + # re-pool — close it so the process can't leak (close is idempotent if + # it was already torn down). + close_it = True + else: + self._idle[w.key].append(w) + if close_it: + w.close() def forget(self, w: _PooledWorker) -> None: with self._lock: - self._all.discard(w) + self._all[w.key].discard(w) try: - self._idle.remove(w) + self._idle[w.key].remove(w) except ValueError: pass def shutdown(self) -> None: with self._lock: - workers = list(self._all) + self._shutting_down = True + workers = [w for s in self._all.values() for w in s] self._all.clear() self._idle.clear() - self._base_opts = None for w in workers: w.close() + with self._lock: + self._shutting_down = False _POOL = _WorkerPool() @@ -300,19 +357,30 @@ def _spawn_worker(base_opts: dict, poolable: bool) -> _PooledWorker: ) proc.start() # The worker warms CUDA + imports (~4 s) then sends exactly one "ready"; consume - # it before the channel starts reading protocol frames on the same pipe. - startup = base_opts.get("startup_timeout", 180.0) - if not parent_conn.poll(startup): + # it before the channel starts reading protocol frames on the same pipe. This poll + # covers the cold spawn+warm, so it stays generous. + warm_wait = base_opts.get("startup_timeout", 180.0) + if not parent_conn.poll(warm_wait): proc.terminate() - raise TimeoutError( - f"isolated worker failed to warm up within {startup}s" - ) + raise TimeoutError(f"isolated worker failed to warm up within {warm_wait}s") msg = parent_conn.recv() if msg != "ready": proc.terminate() raise RuntimeError(f"unexpected isolated-worker handshake: {msg!r}") - chan = CudaIpcHostChannel(parent_conn, buf, timeout=base_opts["timeout"]) - return _PooledWorker(proc, buf, parent_conn, chan, poolable=poolable) + # The channel's first-event budget covers only a job's deserialize + run-to-first- + # request (spawn+warm already happened above), so it is the user timeout plus a + # deserialize margin — NOT the cold 180 s, which would defeat hang-containment on a + # job that hangs before its first event on an already-warm worker. + chan = CudaIpcHostChannel( + parent_conn, + buf, + timeout=base_opts["timeout"], + startup_timeout=base_opts["timeout"] + _JOB_STARTUP_MARGIN, + ) + return _PooledWorker( + proc, buf, parent_conn, chan, poolable=poolable, + key=_WorkerPool._key(base_opts), + ) def _build_job(mediator) -> tuple: @@ -397,11 +465,22 @@ def acquire_isolated_worker(mediator) -> None: """ payload, extras, worker_opts = _build_job(mediator) pool_size = _STATE.get("pool_size", 0) - if pool_size > 0: - iso = _POOL.acquire(_base_opts(), pool_size) - else: - iso = _spawn_worker(_base_opts(), poolable=False) - iso.send_job(payload, extras, worker_opts) + + def _acquire(): + if pool_size > 0: + return _POOL.acquire(_base_opts(), pool_size) + return _spawn_worker(_base_opts(), poolable=False) + + iso = _acquire() + try: + iso.send_job(payload, extras, worker_opts) + except (BrokenPipeError, EOFError, OSError): + # The worker died between the liveness check and dispatch (tiny race). Forget + # it and retry once through the normal acquisition path with a fresh worker; + # a second failure is a real problem and propagates. + _POOL.forget(iso) + iso = _acquire() + iso.send_job(payload, extras, worker_opts) _wire_host_channel(mediator, iso) From da553d5d2808bf6547f665fdffe79b31402634fb Mon Sep 17 00:00:00 2001 From: khaiwang Date: Tue, 9 Jun 2026 00:52:35 -0400 Subject: [PATCH 05/30] feat(intervention): tracer.cache() under isolation MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit tracer.cache() registers persistent hooks (mediator_idx=inf) that fill a .save()'d CacheDict during the forward; in the worker those hooks landed on the dummy modules and never fired, so the user got an empty cache. Now: - Worker cache() (isolated): ship the spec (token, module-paths, device, dtype, detach, include_output, include_inputs, rename, alias) via a new Events.CACHE request instead of registering dummy hooks; return a token-tagged placeholder CacheDict the user binds + .save()s. - Host handle_cache_event: resolve paths to the real envoys, register the real cache_output/input_hook into a host Cache keyed by token (Mediator._iso_caches), set_user_cache, ack. Hooks live on the host mediator and are dropped at teardown by remove_hooks, like in-process. - handle_cache_event acks + returns True, so the host loop processes CACHE then END consecutively at Mediator.start (before the forward). handle_end_event then swaps the host CacheDict reference in for the worker's empty placeholder (matched by token); the forward fills that same object in-place, so the user's variable IS the forward-filled host cache. No separate post-forward injection step. The substitution is gated on _iso_caches, so non-cache traces are untouched. Verified (test_isolated_cache.py, gpt2/A100): single module, multi-module, and include_inputs=True all bit-identical (max|Δ|=0, keys match in-process). Full isolated regression (trace/acceptance/multitoken) and cold path unchanged. The test now derives cache keys from envoy.path instead of a hardcoded prefix. See docs §15. Co-Authored-By: Claude Opus 4.8 (1M context) --- .../mediator-gpu-trace-integration.md | 42 ++++++++++- .../gpu_sandbox/test_isolated_cache.py | 40 ++++++++--- src/nnsight/intervention/interleaver.py | 70 +++++++++++++++++++ src/nnsight/intervention/tracing/tracer.py | 42 ++++++++--- 4 files changed, 174 insertions(+), 20 deletions(-) diff --git a/docs/developing/mediator-gpu-trace-integration.md b/docs/developing/mediator-gpu-trace-integration.md index 03fcb367e..722002544 100644 --- a/docs/developing/mediator-gpu-trace-integration.md +++ b/docs/developing/mediator-gpu-trace-integration.md @@ -173,7 +173,7 @@ rule; deferred (revisit with double-buffering if real workloads need it). | `tracer.barrier()` | host-side participant counting + the existing `handle_barrier_event` coordination loop | done | | `cross_invoker` variable sharing | host-mediated variable store (worker pushes data locals, pulls the merged store) | done | | `with tensor.backward()` / `.grad` | needs host-side backward execution (the autograd graph is host-side) | planned | -| `tracer.cache()` | host-side cache-hook registration + post-forward injection of the populated CacheDict | planned | +| `tracer.cache()` | CACHE event → host registers the real cache hooks; host CacheDict swapped in for the worker placeholder, filled in-place by the forward | done (§15) | | warm worker pool (`pool_size=`) | generic workers receive serialized mediators as jobs over the channel; clean-END workers recycled, others retired | done (§14) | | MPS / `isolate_mediators()` further polish | — | planned | @@ -257,7 +257,7 @@ benign CudaIPC release warning. | `cross_invoker` variable sharing | ✅ host variable store; transmittable data vars only — see §10 | | warm worker pool (`pool_size=N`, `warm_worker_pool`) | ✅ ~21× faster per request once warm; recycle-on-clean-END — see §14 | | `with tensor.backward()` / `.grad` | 🔜 hard: the autograd graph is host-side, the worker has detached clones; needs the backward pass to run host-side with path-based grad providers (a major build) | -| `tracer.cache()` | 🔜 tractable: returns an empty CacheDict today (hooks fire on dummy modules); needs host-side cache-hook registration + shipping the populated CacheDict back | +| `tracer.cache()` (`modules=`, `include_inputs=`) | ✅ bit-identical — CACHE event → host registers the real cache hooks; the forward fills the host CacheDict in-place (§15) | | `.source` operation-level access (`...attn.split_1.output`) | 🔜 not yet (op paths aren't in `model.modules()`) | | in-place `[:]=` | ⛔ use explicit `=` (clone semantics, §4) | @@ -449,3 +449,41 @@ pooled workers all bit-identical; a timed-out (infinite-loop) worker is retired the next trace bit-identical; a killed idle worker is skipped + replaced on the next acquire; a non-standard- named model works through the pool. Cold path (`pool_size=0`) stays bit-identical across read/swap/save/multi/exception/hang/multitoken/cross-invoke/barrier/nonstd. + +--- + +## 15. `tracer.cache()` — DONE (2026-06-08) + +**The gap.** `tracer.cache()` registers **persistent** hooks (`mediator_idx=inf`) that fill a `.save()`'d +`CacheDict` *during the forward*. In the worker those hooks land on the **dummy** modules and never fire, so +the user got an empty `CacheDict` (the `.save()`'d placeholder shipped back unfilled). + +**The fix — a `CACHE` event + host-side registration, with the post-forward injection collapsed.** +- **Worker `cache()` (isolated branch in `tracer.py`):** instead of registering dummy hooks, ship the spec + `(token, module-paths, device, dtype, detach, include_output, include_inputs, rename, alias)` via a new + `Events.CACHE` request (`mediator.send`). Return a **token-tagged** placeholder `CacheDict` the user binds + and `.save()`s — that is what carries the user's variable name across the boundary. +- **Host `handle_cache_event`:** resolve the paths to the **real** envoys, register the real + `cache_output_hook`/`cache_input_hook` into a host `Cache` keyed by the token (on `Mediator._iso_caches`), + `set_user_cache`, ack. Hooks live on the host mediator's `hooks`, removed at teardown by `remove_hooks` — + exactly like in-process. +- **The timing insight (no separate teardown step).** `handle_cache_event` acks and returns `True`, so the + host loop processes `CACHE` then `END` consecutively at `Mediator.start` — *before* the forward. So + `handle_end_event` swaps the **host** `CacheDict` reference in for the worker's empty placeholder (matched + by the token on the saved value) when it injects the saves. The forward then fills *that same object* + in-place. The user's variable **is** the forward-filled host cache; the doc's earlier "token-matched + post-forward injection" collapses to "swap in the host CacheDict at END; the forward fills it." This also + preserves in-process semantics, including "a cache defined after a module is called misses it" (the host + hooks register only when the `CACHE` event arrives, just like the in-process registration point). + +**Touch points:** `Events.CACHE`; `tracer.cache()` isolated branch; `handle()` dispatch + `handle_cache_event`; +the `handle_end_event` token-swap (gated on `_iso_caches`, so non-cache traces are untouched); `_iso_caches` +on `Mediator`. + +**Verified (`test_isolated_cache.py`, gpt2/A100):** single module, multi-module (3 keys), and +`include_inputs=True` all bit-identical (`max|Δ|=0`, keys match in-process). Cold path + the full isolated +regression (trace/acceptance/multitoken/cross-invoke/pool) unchanged. + +**Not covered:** a cache placeholder nested inside a container save (`got = [t.cache().save()]`) — the swap +matches a top-level saved `CacheDict`; nesting would need a recursive walk. `cross_invoker` + cache is +untested. `modules=None` (cache *all* modules) registers a hook per module on the host — correct but heavy. diff --git a/prototypes/mediator-sandbox/gpu_sandbox/test_isolated_cache.py b/prototypes/mediator-sandbox/gpu_sandbox/test_isolated_cache.py index 609a6e7df..c2fe303ce 100644 --- a/prototypes/mediator-sandbox/gpu_sandbox/test_isolated_cache.py +++ b/prototypes/mediator-sandbox/gpu_sandbox/test_isolated_cache.py @@ -20,21 +20,22 @@ PROMPT = "The Eiffel Tower is in the city of" -def _entry_out(cache, path): +def _entry(cache, path): e = cache[path] - e = e[-1] if isinstance(e, list) else e - return e.output + return e[-1] if isinstance(e, list) else e def test_one(model): + h6 = model.transformer.h[6] with model.trace(PROMPT) as t: - ref = t.cache(modules=[model.transformer.h[6]]).save() + ref = t.cache(modules=[h6]).save() with isolate_mediators(timeout=30): with model.trace(PROMPT) as t: - got = t.cache(modules=[model.transformer.h[6]]).save() + got = t.cache(modules=[h6]).save() + key = h6.path # derive the key from the envoy path, don't hardcode the prefix rk, gk = sorted(ref.keys()), sorted(got.keys()) ok = rk == gk and len(gk) > 0 and torch.equal( - _entry_out(ref, "transformer.h.6")[0], _entry_out(got, "transformer.h.6")[0] + _entry(ref, key).output[0], _entry(got, key).output[0] ) print(f"[one] keys ref={rk} got={gk} match={ok}", flush=True) return ok @@ -47,18 +48,39 @@ def test_multi(model): with isolate_mediators(timeout=30): with model.trace(PROMPT) as t: got = t.cache(modules=mods).save() - paths = ["transformer.h.2", "transformer.h.5", "transformer.h.9"] + paths = [m.path for m in mods] ok = sorted(ref.keys()) == sorted(got.keys()) and all( - torch.equal(_entry_out(ref, p)[0], _entry_out(got, p)[0]) for p in paths + torch.equal(_entry(ref, p).output[0], _entry(got, p).output[0]) for p in paths ) print(f"[multi] {len(got.keys())} keys match={ok}", flush=True) return ok +def test_inputs(model): + # include_inputs=True: the cached module inputs must match in-process too. + h4 = model.transformer.h[4] + with model.trace(PROMPT) as t: + ref = t.cache(modules=[h4], include_inputs=True).save() + with isolate_mediators(timeout=30): + with model.trace(PROMPT) as t: + got = t.cache(modules=[h4], include_inputs=True).save() + key = h4.path + rk, gk = sorted(ref.keys()), sorted(got.keys()) + ok = rk == gk and len(gk) > 0 and torch.equal( + _entry(ref, key).input, _entry(got, key).input + ) + print(f"[inputs] keys match, input bit-identical={ok}", flush=True) + return ok + + def main(): assert torch.cuda.is_available() model = LanguageModel("gpt2", device_map="cuda", dispatch=True) - results = {"one": test_one(model), "multi": test_multi(model)} + results = { + "one": test_one(model), + "multi": test_multi(model), + "inputs": test_inputs(model), + } ok = all(results.values()) print("=" * 72, flush=True) print(f"ISOLATED CACHE: {'PASS' if ok else 'FAIL'} — {results}", flush=True) diff --git a/src/nnsight/intervention/interleaver.py b/src/nnsight/intervention/interleaver.py index 289003c84..f7dbc6752 100755 --- a/src/nnsight/intervention/interleaver.py +++ b/src/nnsight/intervention/interleaver.py @@ -355,6 +355,7 @@ class Events(Enum): EXCEPTION = "exception" # Signal that an exception occurred SKIP = "skip" # Signal that an operation should be skipped BARRIER = "barrier" # Signal that a barrier should be set + CACHE = "cache" # Register a tracer.cache() on the host's real modules (isolation) class Cancelation(Exception): @@ -977,6 +978,11 @@ def __init__( # isolation.acquire_isolated_worker). ``None`` => in-process (default). self._iso = None + # Isolated tracer.cache(): token -> host Cache. The worker ships a CACHE event; + # the host registers cache hooks here on the REAL modules and swaps the host + # CacheDict in for the worker's empty placeholder when saves are injected. + self._iso_caches: dict = {} + # True only inside an isolated worker process (set by _run_one_job). Lets # cross-process-aware logic (e.g. Barrier) know it can't count locally. self._isolated_worker = False @@ -1208,6 +1214,8 @@ def handle(self, provider: Optional[str] = None, value: Optional[Any] = None): process = self.handle_skip_event(provider, *data) elif event == Events.BARRIER: process = self.handle_barrier_event(provider, data) + elif event == Events.CACHE: + process = self.handle_cache_event(data) elif event == Events.END: process = self.handle_end_event(data) @@ -1409,6 +1417,53 @@ def handle_barrier_event(self, provider: Any, participants: Set[str]): return False + def handle_cache_event(self, spec: Any): + """Isolated ``tracer.cache()``: register the persistent cache hooks on the + HOST's REAL modules (the worker's are dummies, so its hooks never fire). + + The host ``Cache``'s ``CacheDict`` is filled IN-PLACE by the forward pass (the + hooks are ``mediator_idx=inf``, firing after intervention hooks). + :meth:`handle_end_event` swaps this host ``CacheDict`` in for the worker's empty + placeholder (matched by token) when the saved variables are injected, so the + user's variable *is* the forward-filled host cache. The hooks live on + ``self.hooks`` and are dropped at teardown by ``remove_hooks`` — exactly like the + in-process path. Acks so the worker's ``send`` returns and it proceeds to END. + """ + from .hooks import cache_input_hook, cache_output_hook + from .tracing.tracer import Cache + + ( + token, + paths, + device, + dtype, + detach, + include_output, + include_inputs, + rename, + alias, + ) = spec + + model = self.interleaver.tracer.model + path2envoy = {e.path: e for e in model.modules()} + targets = [path2envoy[p] for p in paths if p in path2envoy] + + cache_obj = Cache( + paths, device, dtype, detach, include_output, include_inputs, rename, alias + ) + batcher = self.interleaver.batcher + for envoy in targets: + if include_output: + cache_output_hook(cache_obj, envoy._module, envoy.path, batcher, self) + if include_inputs: + cache_input_hook(cache_obj, envoy._module, envoy.path, batcher, self) + + self._iso_caches[token] = cache_obj + self.set_user_cache(cache_obj) + + self.respond(None) # ack -> worker's send() returns, it proceeds to END + return True + def handle_end_event(self, saves: Optional[Any] = None): """ Handle an end event by stopping the mediator. @@ -1426,6 +1481,21 @@ def handle_end_event(self, saves: Optional[Any] = None): # cancel() (below) reads this to recycle vs retire the worker. self._iso.clean = True if saves: + if self._iso_caches: + # tracer.cache(): the worker shipped an EMPTY placeholder CacheDict; + # swap in the HOST cache (matched by token), which the forward fills + # in-place after this injection. The user's variable then IS the + # forward-filled host cache. + from .tracing.tracer import Cache + + saves = dict(saves) + for name, val in list(saves.items()): + if isinstance(val, Cache.CacheDict): + host_cache = self._iso_caches.get( + getattr(val, "_iso_cache_token", None) + ) + if host_cache is not None: + saves[name] = host_cache.cache tracer = self.interleaver.tracer user_frame = ( tracer.info.frame diff --git a/src/nnsight/intervention/tracing/tracer.py b/src/nnsight/intervention/tracing/tracer.py index 4b60fde4b..f7b0d3c69 100755 --- a/src/nnsight/intervention/tracing/tracer.py +++ b/src/nnsight/intervention/tracing/tracer.py @@ -614,15 +614,39 @@ def cache( targets.append(envoy) # Register persistent cache hooks on each target module - for envoy in targets: - if include_output: - cache_output_hook( - cache_obj, envoy._module, envoy.path, batcher, mediator - ) - if include_inputs: - cache_input_hook( - cache_obj, envoy._module, envoy.path, batcher, mediator - ) + if mediator._isolated_worker: + # Isolated: the worker's modules are DUMMIES, so cache hooks registered + # here would never fire. Ship the spec via a CACHE event so the HOST + # registers the hooks on its real modules and fills this cache. The + # returned CacheDict is a token-tagged placeholder the user binds + + # .save()s; the host swaps in its own (forward-filled) CacheDict when the + # saved variables are injected into the user frame (matched by token). + from ..interleaver import Events + + token = id(cache_obj.cache) + cache_obj.cache._iso_cache_token = token + spec = ( + token, + [envoy.path for envoy in targets], + device, + dtype, + detach, + include_output, + include_inputs, + rename_dict, + alias_dict, + ) + mediator.send(Events.CACHE, spec) + else: + for envoy in targets: + if include_output: + cache_output_hook( + cache_obj, envoy._module, envoy.path, batcher, mediator + ) + if include_inputs: + cache_input_hook( + cache_obj, envoy._module, envoy.path, batcher, mediator + ) mediator.set_user_cache(cache_obj) From 8d09195fc121bc7c51244509da0cd3ff8cc536c0 Mon Sep 17 00:00:00 2001 From: khaiwang Date: Tue, 9 Jun 2026 12:30:23 -0400 Subject: [PATCH 06/30] test(isolation): valid backward control in the gap characterization MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The backward case hooked `.grad` on a GPT2 block's `.output[0]` — an off-the- backward-path index into the block's tuple output, whose grad hook never fires (a usage gotcha, confirmed via a manual register_hook). So its in-process control ALSO errored, making the test useless as a gap demonstration. Switch to `model.transformer.ln_f.output` — a tensor-output module ON the autograd path — so the in-process control is valid. Backward now succeeds in-process and fails cleanly under isolation, which is the gap the test exists to characterize (host-only autograd graph, detached worker clones, id(tensor)-keyed grads). requires_grad_(True) turned out to be a red herring (ln_f works with or without it); the discriminator is on-path tensor-output vs off-path tuple-element view. Verified: backward in-process=ok, isolated=fails-cleanly; cache=bit-identical (no longer a gap). Co-Authored-By: Claude Opus 4.8 (1M context) --- .../test_isolated_backward_cache_gaps.py | 43 ++++++++++++++----- 1 file changed, 33 insertions(+), 10 deletions(-) diff --git a/prototypes/mediator-sandbox/gpu_sandbox/test_isolated_backward_cache_gaps.py b/prototypes/mediator-sandbox/gpu_sandbox/test_isolated_backward_cache_gaps.py index 81fd11407..8337b5964 100644 --- a/prototypes/mediator-sandbox/gpu_sandbox/test_isolated_backward_cache_gaps.py +++ b/prototypes/mediator-sandbox/gpu_sandbox/test_isolated_backward_cache_gaps.py @@ -1,10 +1,18 @@ #!/usr/bin/env python3 -"""backward/grad + cache() under isolation — characterizing the gaps. +"""backward/grad + cache() under isolation. - backward — get hidden + logits, then `with logits.sum().backward(): g = hidden.grad`. - cache — tracer.cache(modules=[...]) populated by hooks. + backward — GAP (not built): grad of an ON-PATH tensor (``ln_f.output``) works + in-process but fails cleanly under isolation. The fundamentals: the + autograd graph is host-only, the worker holds DETACHED clones (no + grad_fn), and ``.grad`` is keyed by ``id(tensor)`` (no cross-process name, + unlike a module path). NB: grad MUST be taken on a tensor-output module + like ``ln_f`` — a GPT2 block's ``.output[0]`` is an off-the-backward-path + index into its tuple output, whose grad hook never fires (a usage gotcha + that would make the in-process control spuriously error). + cache — now SUPPORTED: ``tracer.cache(modules=[...])`` is bit-identical under + isolation (kept here as a regression check; see test_isolated_cache.py). -Each isolated-vs-in-process, hard timeout so deadlock shows as timeout. +Each isolated-vs-in-process, hard timeout so a deadlock shows as a timeout. Run: CUDA_VISIBLE_DEVICES=6 PYTHONPATH=src \ @@ -29,20 +37,35 @@ def _run(fn): def test_backward(model): + # ln_f.output is a tensor-output module ON the autograd path, so the in-process + # control is VALID (unlike a block's off-path .output[0]). def body(): with model.trace(PROMPT): - hidden = model.transformer.h[6].output[0].save() - logits = model.lm_head.output.save() - with logits.sum().backward(): - g = hidden.grad.save() + hs = model.transformer.ln_f.output + with model.lm_head.output.sum().backward(): + g = hs.grad.save() return g + rs, rv = _run(body) + def iso(): with isolate_mediators(timeout=25): return body() + gs, gv = _run(iso) - ok = rs == "ok" and gs == "ok" and torch.is_tensor(rv) and torch.is_tensor(gv) and torch.equal(rv, gv) - print(f"[backward] ref={rs} got={gs} match={ok if gs=='ok' else gv}", flush=True) + + inproc_ok = rs == "ok" and torch.is_tensor(rv) + # Backward under isolation is a documented gap (host-only autograd graph, detached + # worker clones). It must fail CLEANLY (an error) — not hang or silently return + # wrong grads. So the characterization holds iff in-process works AND isolated errors. + isolated_fails_cleanly = gs == "err" + ok = inproc_ok and isolated_fails_cleanly + print( + f"[backward] in-process={'ok (valid control)' if inproc_ok else rs}; " + f"isolated={'fails cleanly — expected gap' if isolated_fails_cleanly else gs}: " + f"{gv if gs == 'err' else 'UNEXPECTEDLY OK — investigate'}", + flush=True, + ) return ok From fcfeb1ef7ba8f2550c561616cac2b1e5b1935f1e Mon Sep 17 00:00:00 2001 From: khaiwang Date: Wed, 10 Jun 2026 22:54:51 -0400 Subject: [PATCH 07/30] feat(intervention): backward read-path under isolation via chain-rule split with tensor.backward() now works in an isolated GPU worker for the read-then-backward case: the worker tags delivered activation clones with requester-string provenance and computes dL/d(clone) on its local tape as seeds; a new Events.BACKWARD ships them to the host, which continues torch.autograd.grad on the real graph over the retained on-graph activations and returns gradients keyed by provenance path. The backward block's .grad reads are served from that dict; .grad on user-derived tensors and .grad assignment raise clear errors. Verified bit-identical (max|delta|=0) on gpt2 ln_f.output and on a renamed model (final_norm/output_projection). Scalar loss only; swaps, batched traces, and multi-token backward remain unsupported (documented in the integration doc's new backward section). Co-Authored-By: Claude Fable 5 --- .../mediator-gpu-trace-integration.md | 62 ++++++++- .../gpu_sandbox/test_isolated_backward.py | 123 ++++++++++++++++++ .../test_isolated_backward_cache_gaps.py | 42 +++--- src/nnsight/intervention/interleaver.py | 64 +++++++++ src/nnsight/intervention/isolation.py | 69 ++++++++++ src/nnsight/intervention/tracing/backwards.py | 88 +++++++++++++ 6 files changed, 422 insertions(+), 26 deletions(-) create mode 100644 prototypes/mediator-sandbox/gpu_sandbox/test_isolated_backward.py diff --git a/docs/developing/mediator-gpu-trace-integration.md b/docs/developing/mediator-gpu-trace-integration.md index 722002544..cc3c1043b 100644 --- a/docs/developing/mediator-gpu-trace-integration.md +++ b/docs/developing/mediator-gpu-trace-integration.md @@ -172,7 +172,7 @@ rule; deferred (revisit with double-buffering if real workloads need it). | `iter`/`all`/`next` (multi-token) | iteration step stamped host-side; host iter-hooks bump the tracker; worker sets its step explicitly | done | | `tracer.barrier()` | host-side participant counting + the existing `handle_barrier_event` coordination loop | done | | `cross_invoker` variable sharing | host-mediated variable store (worker pushes data locals, pulls the merged store) | done | -| `with tensor.backward()` / `.grad` | needs host-side backward execution (the autograd graph is host-side) | planned | +| `with tensor.backward()` / `.grad` | BACKWARD event: worker seeds `dL/d(delivered clone)` from its local tape, host continues `torch.autograd.grad` on the real graph, `.grad` served by provenance path | done (§16) | | `tracer.cache()` | CACHE event → host registers the real cache hooks; host CacheDict swapped in for the worker placeholder, filled in-place by the forward | done (§15) | | warm worker pool (`pool_size=`) | generic workers receive serialized mediators as jobs over the channel; clean-END workers recycled, others retired | done (§14) | | MPS / `isolate_mediators()` further polish | — | planned | @@ -256,7 +256,7 @@ benign CudaIPC release warning. | `tracer.barrier()` | ✅ host-side participant counting | | `cross_invoker` variable sharing | ✅ host variable store; transmittable data vars only — see §10 | | warm worker pool (`pool_size=N`, `warm_worker_pool`) | ✅ ~21× faster per request once warm; recycle-on-clean-END — see §14 | -| `with tensor.backward()` / `.grad` | 🔜 hard: the autograd graph is host-side, the worker has detached clones; needs the backward pass to run host-side with path-based grad providers (a major build) | +| `with tensor.backward()` / `.grad` | ✅ read-path bit-identical (scalar loss; single invoke; no swap-then-backward; `.grad` editing raises — §16) | | `tracer.cache()` (`modules=`, `include_inputs=`) | ✅ bit-identical — CACHE event → host registers the real cache hooks; the forward fills the host CacheDict in-place (§15) | | `.source` operation-level access (`...attn.split_1.output`) | 🔜 not yet (op paths aren't in `model.modules()`) | | in-place `[:]=` | ⛔ use explicit `=` (clone semantics, §4) | @@ -330,9 +330,10 @@ perf cliff for large cross-invoke tensors). --- -## 11. Backward + caching — CHARACTERIZED (not yet built) +## 11. Backward + caching — characterization (both gaps since closed: cache §15, backward §16) -`test_isolated_backward_cache_gaps.py` confirms the two gaps and their difficulty: +`test_isolated_backward_cache_gaps.py` confirmed the two gaps and their difficulty at the time; +kept for the reasoning record: - **`tracer.cache()` (a real build, not a quick shim):** `tracer.cache()` runs in the worker → registers cache hooks on dummy modules → never fire → the `.save()`'d CacheDict comes back **empty**. The fix @@ -487,3 +488,56 @@ regression (trace/acceptance/multitoken/cross-invoke/pool) unchanged. **Not covered:** a cache placeholder nested inside a container save (`got = [t.cache().save()]`) — the swap matches a top-level saved `CacheDict`; nesting would need a recursive walk. `cross_invoker` + cache is untested. `modules=None` (cache *all* modules) registers a hook per module on the host — correct but heavy. + +--- + +## 16. `with tensor.backward()` — read-path DONE (2026-06-10) + +**The gap (§11).** Clone-on-receive strips `grad_fn` — the worker's delivered activations are detached +clones, the autograd graph lives only on the host, and `.grad` providers were keyed by `id(tensor)` +(process-local). So a backward block in the worker had nothing to differentiate. + +**The fix — split the chain rule at the process boundary** (distributed-autograd-style stitch), keyed by +**requester string** (module path + kind + step), not `id(tensor)`: + +- **Worker, forward time:** when the trace contains a backward block (detected from the intervention + source), every delivered activation clone is tagged — `requires_grad_(True)` + an `id(clone)` → + requester-string provenance map — via a wrapped `mediator.request` (`_tag_delivered`, + `isolation.py`). Worker-side ops on delivered values therefore build the worker's local half of the + graph (delivered leaves → loss). +- **Worker, backward time:** `BackwardsTracer.execute` detects the isolated context + (`worker_backward_context()`) and runs `_execute_isolated` (`tracing/backwards.py`): it computes the + worker half of the chain rule — `dL/d(delivered clone)` for each tagged leaf the loss depends on + (`torch.autograd.grad`, `allow_unused`) — and ships the seed dict `{requester: grad}` via a new + `Events.BACKWARD`. +- **Host:** `handle_value_event` retains the REAL on-graph activation per requester (gated on + `_iso_backward` so non-backward traces pay nothing); `handle_backward_event` continues the chain rule + on the host's real graph (`torch.autograd.grad` seeded by the worker grads, targets = all retained + activations, `allow_unused`, `retain_graph`) and returns `{requester: dL/d(activation)}`. No user code + runs on the host. +- **Worker, block body:** runs locally; each `.grad` read is served from the returned dict by the + tensor's provenance path (a patched `Tensor.grad` property). `.grad` on a user-derived tensor (no + provenance) raises a clear error; the `.grad` setter raises `NotImplementedError`. + +**Stitch correctness (no double-counting):** the worker seeds are partials treating delivered leaves as +independent; the host's autograd adds the indirect inter-layer contributions. Verified empirically — a +loss using the same read activation twice is bit-identical to in-process. + +**Verified (`test_isolated_backward.py`, gpt2/A100):** grad of `ln_f.output` `max|Δ|=0` vs in-process; +renamed model (`final_norm`/`output_projection`) `max|Δ|=0`; user-derived-tensor `.grad` → clear error. +Independent review (7 finder angles, dedup + verify): **no silent-wrong in the in-scope path** (single +invoke, on-path tensor-output target, scalar loss, no swaps). + +**Limits / open:** +- Scalar loss only — `loss.backward(gradient=...)` is not honored (scalar-only error). +- **Gradient-through-swap unsupported** (next increment): an isolated SWAP splices a graph *leaf* into + the host forward (worker-computed values carry no host `grad_fn`), so the host backward dead-ends at + the swap — while in-process gradients DO flow through swaps. Fix = recursive seam stitch (host ships + `dL/d(swap)` to the worker, the worker tape backprops to its leaves, ships back). +- Batched traces error with a cryptic shape mismatch (host retains the full-batch tensor, worker seeds + the narrowed clone) — needs a clear error or narrowed retention. +- Multi-token backward: under characterization (control test first; guard or document per outcome). +- Efficiency deferred: the host computes + ships grads for ALL retained reads (`retain_graph` always + on); a large read set can overflow the 64 MB arena. +- The `".backward("` source-substring detection can false-positive (e.g. the string in a comment), + which only costs needless tagging — tightening is planned alongside the gate consolidation. diff --git a/prototypes/mediator-sandbox/gpu_sandbox/test_isolated_backward.py b/prototypes/mediator-sandbox/gpu_sandbox/test_isolated_backward.py new file mode 100644 index 000000000..bd782836d --- /dev/null +++ b/prototypes/mediator-sandbox/gpu_sandbox/test_isolated_backward.py @@ -0,0 +1,123 @@ +#!/usr/bin/env python3 +"""Increment 1 — read-then-backward under isolation == in-process, bit-identical. + +Scope (the simplest correct slice): single invoke, NO swaps, the gradient target is +an ON-PATH tensor-output module (``ln_f.output``), the loss is a scalar reduction of a +read activation. The host graph stays intact (a read does not replace the host's real +tensor), so the host can run the real backward; the worker only computes the seed +gradient at the loss's delivered leaves and reads grads back by PATH. + + read_backward — grad of ln_f.output w.r.t. a sum-of-logits loss; max|Δ|=0 vs in-process. + +NOT in scope here (later increments): gradient through a SWAP (severs the host graph), +tuple-element targets (off-path), batched/multi-invoke backward. + +Run: + CUDA_VISIBLE_DEVICES=0 PYTHONPATH=src \ + /disk/u/zikai/anaconda3/envs/hf-serve/bin/python -u \ + prototypes/mediator-sandbox/gpu_sandbox/test_isolated_backward.py +""" +import sys + +import torch + +from nnsight import LanguageModel +from nnsight.intervention.isolation import isolate_mediators + +PROMPT = "The Eiffel Tower is in the city of" + + +def _read_backward(model): + # ln_f.output is a tensor-output module ON the autograd path. + with model.trace(PROMPT): + hs = model.transformer.ln_f.output + with model.lm_head.output.sum().backward(): + g = hs.grad.save() + return g + + +def test_read_backward(model): + ref = _read_backward(model) + with isolate_mediators(timeout=30): + got = _read_backward(model) + + ok = ( + torch.is_tensor(ref) + and torch.is_tensor(got) + and ref.shape == got.shape + and torch.equal(ref, got) + ) + delta = (ref - got).abs().max().item() if (torch.is_tensor(ref) and torch.is_tensor(got) and ref.shape == got.shape) else float("nan") + print( + f"[read_backward] ref={tuple(ref.shape) if torch.is_tensor(ref) else ref} " + f"got={tuple(got.shape) if torch.is_tensor(got) else got} " + f"max|Δ|={delta} match={ok}", + flush=True, + ) + return ok + + +def test_read_backward_nonstd(): + # Testing rule: vary names, don't assume GPT-2 conventions. Rename the norm + head + # to non-standard user-facing paths; the host must resolve grads by the real path. + rename = {"transformer.ln_f": "final_norm", "lm_head": "output_projection"} + model = LanguageModel("gpt2", device_map="cuda", dispatch=True, rename=rename) + + def body(): + with model.trace(PROMPT): + hs = model.final_norm.output + with model.output_projection.output.sum().backward(): + g = hs.grad.save() + return g + + ref = body() + with isolate_mediators(timeout=30): + got = body() + ok = torch.is_tensor(ref) and torch.is_tensor(got) and torch.equal(ref, got) + delta = (ref - got).abs().max().item() if ok or (torch.is_tensor(ref) and torch.is_tensor(got)) else float("nan") + print(f"[nonstd] final_norm.output grad isolated==in-proc: {ok} (max|Δ|={delta})", flush=True) + return ok + + +def test_derived_target_fails_clean(model): + # Boundary: gradient of a USER-DERIVED tensor has no host-side graph under isolation. + # It must fail CLEANLY (a clear error), never hang or return a silently-wrong grad. + def body(): + with model.trace(PROMPT): + hs = model.transformer.ln_f.output + derived = hs * 2 # computed in user code — no module-path provenance + with model.lm_head.output.sum().backward(): + g = derived.grad.save() + return g + + try: + body() + print("[derived] UNEXPECTEDLY OK — should have refused a derived-tensor grad", flush=True) + return False + except Exception as e: # noqa: BLE001 + msg = str(e) + ok = "isolation" in msg.lower() + print(f"[derived] failed cleanly={ok}: {type(e).__name__}: {msg[:110]}", flush=True) + return ok + + +def main(): + assert torch.cuda.is_available() + model = LanguageModel("gpt2", device_map="cuda", dispatch=True) + results = { + "read_backward": test_read_backward(model), + "nonstd": test_read_backward_nonstd(), + "derived_fails_clean": _derived_isolated(model), + } + print("=" * 72, flush=True) + print(f"ISOLATED BACKWARD (increment 1): {results}", flush=True) + sys.exit(0 if all(results.values()) else 1) + + +def _derived_isolated(model): + with isolate_mediators(timeout=30): + return test_derived_target_fails_clean(model) + + +if __name__ == "__main__": + main() diff --git a/prototypes/mediator-sandbox/gpu_sandbox/test_isolated_backward_cache_gaps.py b/prototypes/mediator-sandbox/gpu_sandbox/test_isolated_backward_cache_gaps.py index 8337b5964..70b3e0b54 100644 --- a/prototypes/mediator-sandbox/gpu_sandbox/test_isolated_backward_cache_gaps.py +++ b/prototypes/mediator-sandbox/gpu_sandbox/test_isolated_backward_cache_gaps.py @@ -1,15 +1,13 @@ #!/usr/bin/env python3 -"""backward/grad + cache() under isolation. - - backward — GAP (not built): grad of an ON-PATH tensor (``ln_f.output``) works - in-process but fails cleanly under isolation. The fundamentals: the - autograd graph is host-only, the worker holds DETACHED clones (no - grad_fn), and ``.grad`` is keyed by ``id(tensor)`` (no cross-process name, - unlike a module path). NB: grad MUST be taken on a tensor-output module - like ``ln_f`` — a GPT2 block's ``.output[0]`` is an off-the-backward-path - index into its tuple output, whose grad hook never fires (a usage gotcha - that would make the in-process control spuriously error). - cache — now SUPPORTED: ``tracer.cache(modules=[...])`` is bit-identical under +"""backward/grad + cache() under isolation (former gaps, now closed). + + backward — now SUPPORTED (increment 1): grad of an ON-PATH tensor-output module + (``ln_f.output``) runs HOST-SIDE under isolation and is bit-identical to + in-process. The host keeps the real graph; the worker computes its half of + the chain rule at the delivered-activation seam and reads grads back by + PATH. Canonical coverage in test_isolated_backward.py. (Still open: gradient + THROUGH a swap, which severs the host graph at the patch point.) + cache — SUPPORTED: ``tracer.cache(modules=[...])`` is bit-identical under isolation (kept here as a regression check; see test_isolated_cache.py). Each isolated-vs-in-process, hard timeout so a deadlock shows as a timeout. @@ -37,8 +35,8 @@ def _run(fn): def test_backward(model): - # ln_f.output is a tensor-output module ON the autograd path, so the in-process - # control is VALID (unlike a block's off-path .output[0]). + # Increment 1 closed this gap: read-then-backward on an on-path tensor-output module + # (ln_f.output) now runs HOST-SIDE under isolation, bit-identical to in-process. def body(): with model.trace(PROMPT): hs = model.transformer.ln_f.output @@ -54,16 +52,16 @@ def iso(): gs, gv = _run(iso) - inproc_ok = rs == "ok" and torch.is_tensor(rv) - # Backward under isolation is a documented gap (host-only autograd graph, detached - # worker clones). It must fail CLEANLY (an error) — not hang or silently return - # wrong grads. So the characterization holds iff in-process works AND isolated errors. - isolated_fails_cleanly = gs == "err" - ok = inproc_ok and isolated_fails_cleanly + ok = ( + rs == "ok" + and gs == "ok" + and torch.is_tensor(rv) + and torch.is_tensor(gv) + and torch.equal(rv, gv) + ) print( - f"[backward] in-process={'ok (valid control)' if inproc_ok else rs}; " - f"isolated={'fails cleanly — expected gap' if isolated_fails_cleanly else gs}: " - f"{gv if gs == 'err' else 'UNEXPECTEDLY OK — investigate'}", + f"[backward] in-process={rs}; isolated={gs}; bit-identical={ok} " + f"({'' if ok else (gv if gs == 'err' else 'mismatch')})", flush=True, ) return ok diff --git a/src/nnsight/intervention/interleaver.py b/src/nnsight/intervention/interleaver.py index f7dbc6752..5baf55a2e 100755 --- a/src/nnsight/intervention/interleaver.py +++ b/src/nnsight/intervention/interleaver.py @@ -356,6 +356,7 @@ class Events(Enum): SKIP = "skip" # Signal that an operation should be skipped BARRIER = "barrier" # Signal that a barrier should be set CACHE = "cache" # Register a tracer.cache() on the host's real modules (isolation) + BACKWARD = "backward" # Run the backward pass on the host's real graph (isolation) class Cancelation(Exception): @@ -983,6 +984,14 @@ def __init__( # CacheDict in for the worker's empty placeholder when saves are injected. self._iso_caches: dict = {} + # Isolated `with tensor.backward()`: when the trace uses backward, the host keeps + # a reference to each delivered (real, on-graph) activation tensor keyed by its + # requester string, so handle_backward_event can run the real backward and read + # gradients off the host's graph. ``_iso_backward`` gates the retention so a + # non-backward isolated trace pays nothing. + self._iso_backward = False + self._iso_grad_reals: dict = {} + # True only inside an isolated worker process (set by _run_one_job). Lets # cross-process-aware logic (e.g. Barrier) know it can't count locally. self._isolated_worker = False @@ -1216,6 +1225,8 @@ def handle(self, provider: Optional[str] = None, value: Optional[Any] = None): process = self.handle_barrier_event(provider, data) elif event == Events.CACHE: process = self.handle_cache_event(data) + elif event == Events.BACKWARD: + process = self.handle_backward_event(data) elif event == Events.END: process = self.handle_end_event(data) @@ -1241,6 +1252,13 @@ def handle_value_event(self, requester: Any, provider: Any) -> bool: # If fulfilled by this processor, respond with the value and continue processing events. if provider == requester: + # Isolated backward: keep a reference to the REAL (on-graph) activation tensor, + # keyed by its requester string, so handle_backward_event can run the host's + # real backward and read this activation's gradient. The worker only ever got a + # detached clone, so its half of the graph can't be differentiated there. + if self._iso is not None and self._iso_backward: + self._iso_grad_reals[provider] = self.interleaver.batcher.current_value + # Potentially only select a slice of the value if this mediator is part of a batch group. value = self.interleaver.batcher.narrow(self.batch_group) @@ -1464,6 +1482,50 @@ def handle_cache_event(self, spec: Any): self.respond(None) # ack -> worker's send() returns, it proceeds to END return True + def handle_backward_event(self, seed: dict): + """Run the backward pass on the host's real graph (isolation). + + The worker computed its half of the chain rule — ``seed`` maps each delivered + activation's requester string to ``dL/d(that activation)`` (the gradient at the + worker→host seam, computed on the worker's local tape). The host continues the + chain rule on its own graph: it differentiates the seeded activations and returns + ``dL/d(activation)`` for every delivered activation the worker might read via + ``.grad`` (keyed by the same requester strings). No user code runs here — only + ``torch.autograd.grad`` on the host's graph with the worker-supplied seeds. + """ + reals = self._iso_grad_reals + + # The seeds: (real activation tensor, dL/d(activation) from the worker). + out_tensors, out_grads = [], [] + for path, grad in seed.items(): + real = reals.get(path) + if torch.is_tensor(real) and real.requires_grad and torch.is_tensor(grad): + out_tensors.append(real) + out_grads.append(grad) + + # The targets: every delivered activation still on the host graph. The worker + # reads a subset of these via ``.grad``; computing all keeps it source-agnostic. + target_paths = [ + p for p, r in reals.items() if torch.is_tensor(r) and r.requires_grad + ] + target_tensors = [reals[p] for p in target_paths] + + result: dict = {} + if out_tensors and target_tensors: + grads = torch.autograd.grad( + out_tensors, + target_tensors, + grad_outputs=out_grads, + allow_unused=True, + retain_graph=True, + ) + result = { + p: g for p, g in zip(target_paths, grads) if torch.is_tensor(g) + } + + self.respond(result) # ack -> worker's send() returns the grad dict + return True + def handle_end_event(self, saves: Optional[Any] = None): """ Handle an end event by stopping the mediator. @@ -1754,6 +1816,8 @@ def __setstate__(self, state): self.worker = None self._iso = None self._isolated_worker = False + self._iso_backward = False + self._iso_grad_reals = {} self.interleaver = None self.history = set() self.user_cache: "Cache" = list() diff --git a/src/nnsight/intervention/isolation.py b/src/nnsight/intervention/isolation.py index de738613c..4c45ac830 100644 --- a/src/nnsight/intervention/isolation.py +++ b/src/nnsight/intervention/isolation.py @@ -441,6 +441,11 @@ def _wire_host_channel(mediator, iso: _PooledWorker) -> None: mediator.channel = chan mediator.worker = iso.proc mediator._iso = iso + # `with tensor.backward()`: if the trace differentiates, the host must retain each + # delivered (real, on-graph) activation so handle_backward_event can run the real + # backward. Detect it from the source and start with a fresh retention map. + mediator._iso_backward = ".backward(" in mediator.intervention.__source__ + mediator._iso_grad_reals = {} # Fresh per-job host-side hook-registration state. iso.registered = set() iso.path2envoy = None @@ -614,6 +619,53 @@ def _transmissible_exc(e): return RuntimeError(f"{type(e).__name__}: {e}") +# Per-job worker-side state for `with tensor.backward()`. When the trace uses backward, +# every delivered activation clone is tagged with its requester string and made to +# require grad, so the worker can build its half of the autograd graph and compute the +# seed gradient at the seam. ``BackwardsTracer.execute`` reads this via +# ``worker_backward_context``. Reset per job in ``_run_one_job``. +_WORKER_BACKWARD_CTX: Dict[str, Any] = { + "mediator": None, # the forward mediator (its channel reaches the host) + "active": False, # the trace contains a `.backward(` call + "prov": {}, # id(delivered clone) -> requester string + "tagged": [], # delivered clones made to require grad (loss may depend on them) +} + + +def worker_backward_context(): + """Return the active backward context inside an isolated worker, else None. + + Used by ``BackwardsTracer.execute`` to detect that ``.backward()`` is running in an + isolated worker (so it must drive the host's real backward instead of differentiating + its detached clones locally).""" + ctx = _WORKER_BACKWARD_CTX + if ctx["active"] and ctx["mediator"] is not None: + return ctx + return None + + +def _reset_worker_backward_ctx(mediator, active: bool) -> None: + _WORKER_BACKWARD_CTX["mediator"] = mediator + _WORKER_BACKWARD_CTX["active"] = active + _WORKER_BACKWARD_CTX["prov"] = {} + _WORKER_BACKWARD_CTX["tagged"] = [] + + +def _tag_delivered(value, requester: str) -> None: + """Tag each delivered activation tensor with its requester provenance and make it + require grad, so worker-side ops on it build the worker's half of the graph.""" + ctx = _WORKER_BACKWARD_CTX + + def _tag(t): + if t.is_floating_point() and t.is_leaf: + t.requires_grad_(True) + ctx["prov"][id(t)] = requester + ctx["tagged"].append(t) + return t + + apply(value, _tag, torch.Tensor) + + def _run_one_job(channel, payload, extras, opts, device) -> None: """Deserialize one mediator against fresh dummies, run its intervention, and ship saves at END. Any failure (including a bad payload) is reported as an EXCEPTION @@ -635,6 +687,23 @@ def _run_one_job(channel, payload, extras, opts, device) -> None: mediator._isolated_worker = True # so Barrier sends the target count (host counts) interleaver.current = mediator + # `with tensor.backward()`: when the trace differentiates, tag each delivered + # activation so the worker can build its half of the graph and seed the host + # backward (BackwardsTracer.execute reads this context). Wrap request() so every + # VALUE read is tagged with its requester provenance + made to require grad. + source = "".join(mediator.info.source) if mediator.info.source else "" + backward_active = ".backward(" in source + _reset_worker_backward_ctx(mediator, backward_active) + if backward_active: + _orig_request = mediator.request + + def _tagging_request(requester, __orig=_orig_request): + value = __orig(requester) + _tag_delivered(value, requester) + return value + + mediator.request = _tagging_request + def _apply_meta(m): # Live host state piggybacked on each response: the iter[:] bound and the # cross_invoker var store (pulled into the frame so push()/pull() see it). diff --git a/src/nnsight/intervention/tracing/backwards.py b/src/nnsight/intervention/tracing/backwards.py index b154c3b00..62889f71e 100755 --- a/src/nnsight/intervention/tracing/backwards.py +++ b/src/nnsight/intervention/tracing/backwards.py @@ -95,6 +95,12 @@ def __init__( def execute(self, fn: Callable): + from ..isolation import worker_backward_context + + ctx = worker_backward_context() + if ctx is not None: + return self._execute_isolated(fn, ctx) + mediator = BackwardsMediator(fn, self.info) interleaver = Interleaver([mediator], self) @@ -110,3 +116,85 @@ def execute(self, fn: Callable): finally: grad_patch.restore() interleaver.cancel() + + def _execute_isolated(self, fn: Callable, ctx: dict): + """Run a backward block inside an isolated GPU worker. + + The autograd graph is split across the process boundary: the worker holds the + recipe from each delivered activation up to ``self.tensor`` (the loss); the host + holds the recipe from the model inputs up to those activations. We stitch the two + halves at the named seam (the activation requester strings): + + 1. The worker computes its half — ``dL/d(delivered activation)`` for each tagged + leaf the loss depends on — as the seed. + 2. The host runs its half (``handle_backward_event``) seeded by those gradients + and returns ``dL/d(activation)`` for every delivered activation, keyed by the + same requester strings. + 3. The backward block body runs here in the worker; each ``.grad`` read is served + from that returned dict by the tensor's provenance — no local ``register_hook`` + (the clones carry no host graph) and no local backward (the graph is host-side). + """ + from ..interleaver import Events, Interleaver + + forward_mediator = ctx["mediator"] + provenance = ctx["prov"] + tagged = [t for t in ctx["tagged"] if t.requires_grad] + + # Worker half of the chain rule: seed = dL/d(delivered leaf) for leaves the loss + # actually depends on (allow_unused drops the rest). + seed: dict = {} + if tagged: + grads = torch.autograd.grad( + self.tensor, tagged, allow_unused=True, retain_graph=True + ) + for leaf, grad in zip(tagged, grads): + if grad is not None: + seed[provenance[id(leaf)]] = grad + + # Host runs its half and returns dL/d(activation) keyed by requester string. + worker_grads = forward_mediator.send(Events.BACKWARD, seed) or {} + + mediator = BackwardsMediator(fn, self.info) + interleaver = Interleaver([mediator], self) + grad_patch = Patch( + torch.Tensor, _isolated_grad_property(provenance, worker_grads), "grad" + ) + try: + grad_patch.patch() + # No local backward: the grads are already in worker_grads. The block body + # just reads them (via the patched .grad) and .save()s; its saves push up + # into the forward frame and ride the forward's END to the host. + with interleaver: + pass + interleaver.check_dangling_mediators() + finally: + grad_patch.restore() + interleaver.cancel() + + +def _isolated_grad_property(provenance: dict, worker_grads: dict): + """A ``Tensor.grad`` property for the isolated backward block: read the gradient from + the host-computed ``worker_grads`` by the tensor's delivery provenance (path), instead + of registering a local autograd hook (the worker clone has no host graph).""" + + def getter(tensor: torch.Tensor): + path = provenance.get(id(tensor)) + if path is None: + raise RuntimeError( + "gradient under isolation is only supported on an unmodified module-" + "output tensor read during the trace (e.g. `model...ln_f.output`); " + "this tensor was derived in user code and has no host-side graph." + ) + if path not in worker_grads: + raise RuntimeError( + f"no gradient available for `{path}` — it is off the backward path " + f"from the loss (its gradient never flowed during the backward pass)." + ) + return worker_grads[path] + + def setter(tensor: torch.Tensor, value: Any): + raise NotImplementedError( + "editing `.grad` under isolation is not yet supported." + ) + + return property(getter, setter) From 56e2a735d9d516966677851f6c4c2ecf350ad7f1 Mon Sep 17 00:00:00 2001 From: khaiwang Date: Wed, 10 Jun 2026 23:00:48 -0400 Subject: [PATCH 08/30] fix(intervention): name the real cause when isolated backward has no graph MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Characterized multi-token (generate + iter) backward with an in-process control first: generate() runs the forward without gradient tracking, so the first .grad read fails in-process ("cannot register a hook on a tensor that doesn't require gradient") — multi-token backward is unsupported on both paths and no silent-wrong is possible (there is no graph at all; the earlier per-step retention-overwrite concern is moot). The isolated path failed at the same user line but blamed the wrong cause ("off the backward path from the loss"). The host now signals the no-graph case — handle_backward_event returns a marker when no retained activation requires grad — and the worker's .grad error names the grad-less forward and points at model.trace(). Characterization script kept as a regression test asserting that message; single-pass backward unaffected (max|delta|=0). Co-Authored-By: Claude Fable 5 --- .../mediator-gpu-trace-integration.md | 8 +- .../test_isolated_multitoken_backward.py | 170 ++++++++++++++++++ src/nnsight/intervention/interleaver.py | 7 + src/nnsight/intervention/tracing/backwards.py | 16 +- 4 files changed, 198 insertions(+), 3 deletions(-) create mode 100644 prototypes/mediator-sandbox/gpu_sandbox/test_isolated_multitoken_backward.py diff --git a/docs/developing/mediator-gpu-trace-integration.md b/docs/developing/mediator-gpu-trace-integration.md index cc3c1043b..3b554f5b2 100644 --- a/docs/developing/mediator-gpu-trace-integration.md +++ b/docs/developing/mediator-gpu-trace-integration.md @@ -536,7 +536,13 @@ invoke, on-path tensor-output target, scalar loss, no swaps). `dL/d(swap)` to the worker, the worker tape backprops to its leaves, ships back). - Batched traces error with a cryptic shape mismatch (host retains the full-batch tensor, worker seeds the narrowed clone) — needs a clear error or narrowed retention. -- Multi-token backward: under characterization (control test first; guard or document per outcome). +- Multi-token backward: **not supported in-process either** (characterized 2026-06-10, + `test_isolated_multitoken_backward.py`) — `generate()` runs the forward without gradient tracking, so + the first `.grad` read fails in-process ("cannot register a hook on a tensor that doesn't require + gradient"); no silent-wrong is possible (there is no graph at all, so the earlier "per-step retention + overwrite" concern is moot). The isolated path fails at the same user line; the host signals the + no-graph case (`handle_backward_event` returns a marker when nothing retained requires grad) so the + worker's error names the real cause instead of "off the backward path". - Efficiency deferred: the host computes + ships grads for ALL retained reads (`retain_graph` always on); a large read set can overflow the 64 MB arena. - The `".backward("` source-substring detection can false-positive (e.g. the string in a comment), diff --git a/prototypes/mediator-sandbox/gpu_sandbox/test_isolated_multitoken_backward.py b/prototypes/mediator-sandbox/gpu_sandbox/test_isolated_multitoken_backward.py new file mode 100644 index 000000000..6a88c4a45 --- /dev/null +++ b/prototypes/mediator-sandbox/gpu_sandbox/test_isolated_multitoken_backward.py @@ -0,0 +1,170 @@ +#!/usr/bin/env python3 +"""Characterize multi-token (generate + iter) backward under isolation. + +The prior review CLAIMED per-step retention overwrite ("last step wins") would make +this silently wrong, but the severity probes were inconclusive (the in-process control +constructions themselves errored). This script builds the control FIRST and compares +outcome-for-outcome: + + per_step — backward inside the iter loop at each step: read ln_f.output, loss = + lm_head.output.sum() at that step, read hs.grad inside the backward block. + post_loop — accumulate ln_f.output per step, after the loop build a scalar loss from + one chosen step's activation and read its .grad in a backward block. + +Verdicts per shape: + both succeed + bit-identical -> SUPPORTED (record in the doc matrix) + both succeed + values differ -> SILENT-WRONG (guard needed) + control fails -> in-process doesn't support it either; isolated + must fail too, with a non-cryptic error + control succeeds + isolated fails -> isolation gap (clean error = acceptable, + documented; hang/cryptic = fix) + +Run: + CUDA_VISIBLE_DEVICES=0 PYTHONPATH=src \ + /disk/u/zikai/anaconda3/envs/hf-serve/bin/python -u \ + prototypes/mediator-sandbox/gpu_sandbox/test_isolated_multitoken_backward.py +""" +import sys +import traceback + +import torch + +import nnsight +from nnsight import LanguageModel +from nnsight.intervention.isolation import isolate_mediators + +PROMPT = "The Eiffel Tower is in the city of" +N = 3 # max_new_tokens; backward exercised at steps 0 and 1 + + +class _null: + def __enter__(self): + return self + + def __exit__(self, *a): + return False + + +def _outcome(fn): + """Run fn; return ("ok", value) or ("err", ": ").""" + try: + return "ok", fn() + except Exception as e: # noqa: BLE001 + return "err", f"{type(e).__name__}: {e}" + + +def _per_step(model, iso): + """Backward inside the iter loop: per-step grad of ln_f.output w.r.t. that + step's sum-of-logits loss.""" + + def body(): + ctx = isolate_mediators(timeout=60) if iso else _null() + with ctx: + with model.generate(PROMPT, max_new_tokens=N) as t: + grads = [] + for step in t.iter[:2]: + hs = model.transformer.ln_f.output + loss = model.lm_head.output.sum() + with loss.backward(): + grads.append(hs.grad) + nnsight.save(grads) + return grads + + return _outcome(body) + + +def _post_loop(model, iso): + """Backward after the iter loop on one chosen step's activation: loss derived + from the step-1 activation itself; .grad read in the backward block.""" + + def body(): + ctx = isolate_mediators(timeout=60) if iso else _null() + with ctx: + with model.generate(PROMPT, max_new_tokens=N) as t: + hs = [] + for step in t.iter[:2]: + hs.append(model.transformer.ln_f.output) + loss = (hs[1].float() ** 2).sum() + with loss.backward(): + g = hs[1].grad.save() + return g + + return _outcome(body) + + +def _compare(name, ref, got): + """ref/got are (status, value) outcomes. Print verdict, return pass/fail.""" + rs, rv = ref + gs, gv = got + + if rs == "err" and gs == "err": + # CHARACTERIZED (2026-06-10): generate() runs the forward without gradient + # tracking, so multi-token backward fails IN-PROCESS at the first .grad read + # ("cannot register a hook on a tensor that doesn't require gradient"). + # No silent-wrong is possible — there is no graph at all. The isolated path + # must fail with the message naming that real cause (not "off the backward + # path", which blames the wrong thing). + print(f"[{name}] control errors -> not supported in-process either.") + print(f"[{name}] in-proc : {' '.join(str(rv).split())[:160]}") + print(f"[{name}] isolated: {' '.join(str(gv).split())[:160]}") + clear = "without gradient tracking" in str(gv) + verdict = ( + "CLEAN-FAIL (parity, cause named)" + if clear + else "CRYPTIC-FAIL (isolated error must name the grad-less forward)" + ) + print(f"[{name}] verdict: {verdict}", flush=True) + return clear + + if rs == "ok" and gs == "err": + print(f"[{name}] control OK but isolated errors: {str(gv)[:200]}") + print(f"[{name}] verdict: ISOLATION GAP (clean error; document or fix)", flush=True) + return False + + if rs == "err" and gs == "ok": + print(f"[{name}] isolated SUCCEEDS where in-process errors ({str(rv)[:120]})") + print(f"[{name}] verdict: SEMANTIC DIVERGENCE (isolated more permissive — investigate)", flush=True) + return False + + # Both ok: compare values (tensor or list of tensors). + rl = rv if isinstance(rv, list) else [rv] + gl = gv if isinstance(gv, list) else [gv] + if len(rl) != len(gl) or not all(torch.is_tensor(a) and torch.is_tensor(b) for a, b in zip(rl, gl)): + print(f"[{name}] BOTH OK but shapes of result differ: ref={rl} got={gl}") + print(f"[{name}] verdict: SILENT-WRONG (structure mismatch)", flush=True) + return False + same = all(a.shape == b.shape and torch.equal(a, b) for a, b in zip(rl, gl)) + deltas = [ + (a - b).abs().max().item() if a.shape == b.shape else float("nan") + for a, b in zip(rl, gl) + ] + shapes = [tuple(a.shape) for a in rl] + verdict = "SUPPORTED (bit-identical)" if same else "SILENT-WRONG (values differ)" + print(f"[{name}] both OK: shapes={shapes} max|Δ| per step={deltas}") + print(f"[{name}] verdict: {verdict}", flush=True) + return same + + +def main(): + assert torch.cuda.is_available() + model = LanguageModel("gpt2", device_map="cuda", dispatch=True) + + results = {} + for name, fn in (("per_step", _per_step), ("post_loop", _post_loop)): + print(f"--- {name}: in-process control ---", flush=True) + ref = fn(model, iso=False) + print(f"--- {name}: isolated ---", flush=True) + got = fn(model, iso=True) + results[name] = _compare(name, ref, got) + + print("=" * 72, flush=True) + print(f"MULTI-TOKEN BACKWARD CHARACTERIZATION: {results}", flush=True) + sys.exit(0 if all(results.values()) else 1) + + +if __name__ == "__main__": + try: + main() + except Exception: + traceback.print_exc() + sys.exit(2) diff --git a/src/nnsight/intervention/interleaver.py b/src/nnsight/intervention/interleaver.py index 5baf55a2e..b00787216 100755 --- a/src/nnsight/intervention/interleaver.py +++ b/src/nnsight/intervention/interleaver.py @@ -1510,6 +1510,13 @@ def handle_backward_event(self, seed: dict): ] target_tensors = [reals[p] for p in target_paths] + if reals and not target_tensors: + # Activations were delivered but NONE is on a graph — the forward ran + # without gradient tracking (e.g. generate() runs grad-less). Tell the + # worker so its .grad reads blame the real cause, not "off the path". + self.respond({"__nnsight_backward_no_graph__": True}) + return True + result: dict = {} if out_tensors and target_tensors: grads = torch.autograd.grad( diff --git a/src/nnsight/intervention/tracing/backwards.py b/src/nnsight/intervention/tracing/backwards.py index 62889f71e..830fd8bdb 100755 --- a/src/nnsight/intervention/tracing/backwards.py +++ b/src/nnsight/intervention/tracing/backwards.py @@ -153,11 +153,16 @@ def _execute_isolated(self, fn: Callable, ctx: dict): # Host runs its half and returns dL/d(activation) keyed by requester string. worker_grads = forward_mediator.send(Events.BACKWARD, seed) or {} + # The host signals "no graph at all" (forward ran without gradient tracking, + # e.g. generate()) distinctly from "this particular read is off the path". + no_graph = bool(worker_grads.pop("__nnsight_backward_no_graph__", False)) mediator = BackwardsMediator(fn, self.info) interleaver = Interleaver([mediator], self) grad_patch = Patch( - torch.Tensor, _isolated_grad_property(provenance, worker_grads), "grad" + torch.Tensor, + _isolated_grad_property(provenance, worker_grads, no_graph), + "grad", ) try: grad_patch.patch() @@ -172,7 +177,7 @@ def _execute_isolated(self, fn: Callable, ctx: dict): interleaver.cancel() -def _isolated_grad_property(provenance: dict, worker_grads: dict): +def _isolated_grad_property(provenance: dict, worker_grads: dict, no_graph: bool = False): """A ``Tensor.grad`` property for the isolated backward block: read the gradient from the host-computed ``worker_grads`` by the tensor's delivery provenance (path), instead of registering a local autograd hook (the worker clone has no host graph).""" @@ -186,6 +191,13 @@ def getter(tensor: torch.Tensor): "this tensor was derived in user code and has no host-side graph." ) if path not in worker_grads: + if no_graph: + raise RuntimeError( + f"no gradient available for `{path}` — the forward pass ran " + f"without gradient tracking, so there is no autograd graph " + f"(generate() runs grad-less; use model.trace() for gradients). " + f"This matches in-process behavior, where .grad raises here too." + ) raise RuntimeError( f"no gradient available for `{path}` — it is off the backward path " f"from the loss (its gradient never flowed during the backward pass)." From a85288303de31f792b032c3360b2c313c55c6f1f Mon Sep 17 00:00:00 2001 From: khaiwang Date: Wed, 10 Jun 2026 23:07:41 -0400 Subject: [PATCH 09/30] refactor(intervention): consolidate worker-side isolation runtime into WorkerMediator _run_one_job built the worker mediator by monkeypatching the deserialized instance (end/exception closures, a request wrapper when backward is active) alongside a module-global backward-context dict with its own reset choreography. The job mediator is now adopted into a WorkerMediator(Mediator) subclass via __class__ swap: the closures become method overrides (end ships Globals.saves-filtered locals on END, exception degrades to a picklable form, request tags delivered clones with requester provenance when the trace differentiates), the meta/push piggyback callbacks become methods bound to the channel, and the backward context collapses to instance attributes plus a single current-mediator pointer read by worker_backward_context(). _run_one_job is now deserialize -> adopt -> wire -> run. Behavior-preserving: full isolated suite (trace, acceptance, multi-token iteration, cross-invoke, warm pool, cache, backward, multi-token-backward characterization, renamed-model) all pass; in-process regression 51 passed. Co-Authored-By: Claude Fable 5 --- src/nnsight/intervention/isolation.py | 253 ++++++++++-------- src/nnsight/intervention/tracing/backwards.py | 14 +- 2 files changed, 141 insertions(+), 126 deletions(-) diff --git a/src/nnsight/intervention/isolation.py b/src/nnsight/intervention/isolation.py index 4c45ac830..ce9d54f80 100644 --- a/src/nnsight/intervention/isolation.py +++ b/src/nnsight/intervention/isolation.py @@ -46,6 +46,7 @@ import torch.nn as nn from . import serialization +from .interleaver import Mediator from .transport import CudaIpcHostChannel, CudaIpcWorkerChannel from ..util import apply @@ -619,145 +620,157 @@ def _transmissible_exc(e): return RuntimeError(f"{type(e).__name__}: {e}") -# Per-job worker-side state for `with tensor.backward()`. When the trace uses backward, -# every delivered activation clone is tagged with its requester string and made to -# require grad, so the worker can build its half of the autograd graph and compute the -# seed gradient at the seam. ``BackwardsTracer.execute`` reads this via -# ``worker_backward_context``. Reset per job in ``_run_one_job``. -_WORKER_BACKWARD_CTX: Dict[str, Any] = { - "mediator": None, # the forward mediator (its channel reaches the host) - "active": False, # the trace contains a `.backward(` call - "prov": {}, # id(delivered clone) -> requester string - "tagged": [], # delivered clones made to require grad (loss may depend on them) -} +class WorkerMediator(Mediator): + """The worker-process half of an isolated mediator. + A job's mediator is deserialized as a plain :class:`Mediator` and then *adopted* + (``__class__`` swap) into this subclass, which overrides exactly the methods whose + in-process behavior relies on shared memory with the host: -def worker_backward_context(): - """Return the active backward context inside an isolated worker, else None. + - ``end``: the worker's frame + ``Globals.saves`` live here, so the exit filter + runs locally and the saved dict rides the END event (worker→host saves + transmission). + - ``exception``: dynamic nnsight exceptions don't pickle; degrade before shipping. + - ``request``: when the trace contains ``with tensor.backward()``, tag each + delivered activation clone (``requires_grad_`` + id→requester provenance) so the + worker builds its half of the autograd graph and can seed the host backward. - Used by ``BackwardsTracer.execute`` to detect that ``.backward()`` is running in an - isolated worker (so it must drive the host's real backward instead of differentiating - its detached clones locally).""" - ctx = _WORKER_BACKWARD_CTX - if ctx["active"] and ctx["mediator"] is not None: - return ctx - return None + ``apply_meta`` / ``push_locals`` are the worker ends of the live host↔worker state + piggyback (`iter[:]` bound + cross_invoker variable store) and are bound to the + channel's ``on_meta`` / ``push_provider`` per job. + """ + @classmethod + def adopt(cls, mediator, channel, interleaver, opts: dict, device) -> "WorkerMediator": + """Turn a freshly-deserialized mediator into this job's worker mediator.""" + mediator.__class__ = cls + mediator.channel = channel + mediator.interleaver = interleaver + mediator.idx = 0 + # cross_invoker matches the host gate; var sharing rides the host store since + # worker frames aren't shared across processes. + mediator.cross_invoker = opts.get("cross_invoker", False) + mediator._isolated_worker = True # so Barrier sends the target count (host counts) + mediator._device = device + # `with tensor.backward()`: detected from the intervention source; gates the + # delivered-clone tagging in ``request`` (BackwardsTracer reads the provenance). + source = "".join(mediator.info.source) if mediator.info.source else "" + mediator._bwd_active = ".backward(" in source + mediator._bwd_prov = {} # id(delivered clone) -> requester string + mediator._bwd_tagged = [] # delivered clones made to require grad + interleaver.current = mediator + return mediator -def _reset_worker_backward_ctx(mediator, active: bool) -> None: - _WORKER_BACKWARD_CTX["mediator"] = mediator - _WORKER_BACKWARD_CTX["active"] = active - _WORKER_BACKWARD_CTX["prov"] = {} - _WORKER_BACKWARD_CTX["tagged"] = [] + def request(self, requester: str): + value = super().request(requester) + if self._bwd_active: + self._tag_delivered(value, requester) + return value + def _tag_delivered(self, value, requester: str) -> None: + """Tag each delivered activation tensor with its requester provenance and make + it require grad, so worker-side ops on it build the worker's half of the graph.""" -def _tag_delivered(value, requester: str) -> None: - """Tag each delivered activation tensor with its requester provenance and make it - require grad, so worker-side ops on it build the worker's half of the graph.""" - ctx = _WORKER_BACKWARD_CTX + def _tag(t): + if t.is_floating_point() and t.is_leaf: + t.requires_grad_(True) + self._bwd_prov[id(t)] = requester + self._bwd_tagged.append(t) + return t - def _tag(t): - if t.is_floating_point() and t.is_leaf: - t.requires_grad_(True) - ctx["prov"][id(t)] = requester - ctx["tagged"].append(t) - return t + apply(value, _tag, torch.Tensor) - apply(value, _tag, torch.Tensor) + def end(self): + # Worker→host saves transmission: bundle .save()'d values into the END event. + # The intervention's compiled body calls ``end()`` on success; push() populates + # the SerializedFrame's f_locals, which we filter by Globals.saves. + from .interleaver import Events + from .tracing.globals import Globals + + self.push() + flocals = self.info.frame.f_locals + saved = {k: v for k, v in flocals.items() if id(v) in Globals.saves} + self.channel.put_event((Events.END, saved)) + + def exception(self, exception: Exception): + super().exception(_transmissible_exc(exception)) + + def apply_meta(self, m: dict) -> None: + # Live host state piggybacked on each response: the iter[:] bound and the + # cross_invoker var store (pulled into the frame so push()/pull() see it). + self.interleaver.default_all = m.get( + "default_all", self.interleaver.default_all + ) + store = m.get("xinvoke_store") + if store: + # Store tensors travel CPU-serialized (see push_locals); move them + # back to the worker's device before the user code uses them. + restored = apply(store, lambda t: t.to(self._device), torch.Tensor) + self.info.frame.f_locals.update(restored) + + def push_locals(self) -> Optional[dict]: + # cross_invoker: ship this worker's *data* locals to the host store. + # push() (called by send() before put_event) has already written them into + # the SerializedFrame's f_locals. We ship only transmittable data (tensors + # + basic types/containers) — framework objects (Barrier/Envoy, which hold + # the model) are skipped; the worker already has them via its own closure. + # Tensors are moved to CPU: a worker tensor cloned from the CUDA-IPC bounce + # buffer cannot be re-shared over IPC by the host. + if not self.cross_invoker: + return None + out = {} + for k, v in self.info.frame.f_locals.items(): + if str(k).startswith("__nnsight"): + continue + if not _transmittable(v): + import warnings + + warnings.warn( + f"cross_invoker: variable {k!r} ({type(v).__name__}) is not " + f"transmittable across the isolation boundary and was not " + f"shared between invokes." + ) + continue + out[k] = apply(v, lambda t: t.detach().cpu(), torch.Tensor) + return out + + +# The job currently running in this worker process (one at a time). Read by +# ``worker_backward_context`` so BackwardsTracer can find the ambient mediator. +_WORKER_CURRENT: Optional[WorkerMediator] = None + + +def worker_backward_context() -> Optional[WorkerMediator]: + """Return this worker's mediator if its trace contains a backward block, else None. + + Used by ``BackwardsTracer.execute`` to detect that ``.backward()`` is running in an + isolated worker (so it must drive the host's real backward instead of differentiating + its detached clones locally); the mediator carries the delivered-clone provenance + (``_bwd_prov`` / ``_bwd_tagged``).""" + med = _WORKER_CURRENT + if med is not None and med._bwd_active: + return med + return None def _run_one_job(channel, payload, extras, opts, device) -> None: - """Deserialize one mediator against fresh dummies, run its intervention, and ship - saves at END. Any failure (including a bad payload) is reported as an EXCEPTION - event so the host never waits on a worker that won't speak.""" + """Deserialize one mediator against fresh dummies, adopt it as this job's + :class:`WorkerMediator`, run its intervention, and ship saves at END. Any failure + (including a bad payload) is reported as an EXCEPTION event so the host never + waits on a worker that won't speak.""" from .interleaver import Events from .tracing.globals import Globals + global _WORKER_CURRENT try: Globals.saves.clear() # per-job reset (the only worker-side global state) interleaver = _WorkerInterleaver(default_all=opts.get("default_all")) mediator = serialization.loads(payload, _WorkerPersistent(interleaver, extras)) - mediator.channel = channel - mediator.interleaver = interleaver - mediator.idx = 0 - # cross_invoker matches the host gate; var sharing rides the host store since - # worker frames aren't shared across processes. - mediator.cross_invoker = opts.get("cross_invoker", False) - mediator._isolated_worker = True # so Barrier sends the target count (host counts) - interleaver.current = mediator - - # `with tensor.backward()`: when the trace differentiates, tag each delivered - # activation so the worker can build its half of the graph and seed the host - # backward (BackwardsTracer.execute reads this context). Wrap request() so every - # VALUE read is tagged with its requester provenance + made to require grad. - source = "".join(mediator.info.source) if mediator.info.source else "" - backward_active = ".backward(" in source - _reset_worker_backward_ctx(mediator, backward_active) - if backward_active: - _orig_request = mediator.request - - def _tagging_request(requester, __orig=_orig_request): - value = __orig(requester) - _tag_delivered(value, requester) - return value - - mediator.request = _tagging_request - - def _apply_meta(m): - # Live host state piggybacked on each response: the iter[:] bound and the - # cross_invoker var store (pulled into the frame so push()/pull() see it). - interleaver.default_all = m.get("default_all", interleaver.default_all) - store = m.get("xinvoke_store") - if store: - # Store tensors travel CPU-serialized (see _push_locals); move them - # back to the worker's device before the user code uses them. - restored = apply(store, lambda t: t.to(device), torch.Tensor) - mediator.info.frame.f_locals.update(restored) - - channel.on_meta = _apply_meta - - def _push_locals(): - # cross_invoker: ship this worker's *data* locals to the host store. - # push() (called by send() before put_event) has already written them into - # the SerializedFrame's f_locals. We ship only transmittable data (tensors - # + basic types/containers) — framework objects (Barrier/Envoy, which hold - # the model) are skipped; the worker already has them via its own closure. - # Tensors are moved to CPU: a worker tensor cloned from the CUDA-IPC bounce - # buffer cannot be re-shared over IPC by the host. - if not mediator.cross_invoker: - return None - out = {} - for k, v in mediator.info.frame.f_locals.items(): - if str(k).startswith("__nnsight"): - continue - if not _transmittable(v): - import warnings - - warnings.warn( - f"cross_invoker: variable {k!r} ({type(v).__name__}) is not " - f"transmittable across the isolation boundary and was not " - f"shared between invokes." - ) - continue - out[k] = apply(v, lambda t: t.detach().cpu(), torch.Tensor) - return out - - channel.push_provider = _push_locals - - # Worker→host saves transmission: bundle .save()'d values into the END event. - # The intervention's compiled body calls ``mediator.end()`` on success; push() - # populates the SerializedFrame's f_locals, which we filter by Globals.saves. - def _end(): - mediator.push() - flocals = mediator.info.frame.f_locals - saved = {k: v for k, v in flocals.items() if id(v) in Globals.saves} - mediator.channel.put_event((Events.END, saved)) - - mediator.end = _end - - _orig_exception = mediator.exception - mediator.exception = lambda e: _orig_exception(_transmissible_exc(e)) + WorkerMediator.adopt(mediator, channel, interleaver, opts, device) + _WORKER_CURRENT = mediator + channel.on_meta = mediator.apply_meta + channel.push_provider = mediator.push_locals mediator.intervention(mediator, mediator.info, *mediator.args) except BaseException as e: # noqa: BLE001 — contain the footgun; report it @@ -765,6 +778,8 @@ def _end(): channel.put_event((Events.EXCEPTION, _transmissible_exc(e))) except Exception: # noqa: BLE001 pass + finally: + _WORKER_CURRENT = None def _pool_worker_main(conn, buf, base_opts): diff --git a/src/nnsight/intervention/tracing/backwards.py b/src/nnsight/intervention/tracing/backwards.py index 830fd8bdb..2a5b03c62 100755 --- a/src/nnsight/intervention/tracing/backwards.py +++ b/src/nnsight/intervention/tracing/backwards.py @@ -97,9 +97,9 @@ def execute(self, fn: Callable): from ..isolation import worker_backward_context - ctx = worker_backward_context() - if ctx is not None: - return self._execute_isolated(fn, ctx) + worker_mediator = worker_backward_context() + if worker_mediator is not None: + return self._execute_isolated(fn, worker_mediator) mediator = BackwardsMediator(fn, self.info) @@ -117,7 +117,7 @@ def execute(self, fn: Callable): grad_patch.restore() interleaver.cancel() - def _execute_isolated(self, fn: Callable, ctx: dict): + def _execute_isolated(self, fn: Callable, worker_mediator): """Run a backward block inside an isolated GPU worker. The autograd graph is split across the process boundary: the worker holds the @@ -136,9 +136,9 @@ def _execute_isolated(self, fn: Callable, ctx: dict): """ from ..interleaver import Events, Interleaver - forward_mediator = ctx["mediator"] - provenance = ctx["prov"] - tagged = [t for t in ctx["tagged"] if t.requires_grad] + forward_mediator = worker_mediator + provenance = worker_mediator._bwd_prov + tagged = [t for t in worker_mediator._bwd_tagged if t.requires_grad] # Worker half of the chain rule: seed = dL/d(delivered leaf) for leaves the loss # actually depends on (allow_unused drops the rest). From 4ac76d9d7303d08726571d750a743047ccd9fc6b Mon Sep 17 00:00:00 2001 From: khaiwang Date: Wed, 10 Jun 2026 23:10:35 -0400 Subject: [PATCH 10/30] refactor(intervention): decide per-job isolation gates once, in _build_job The backward detection (".backward(" in the intervention source) ran independently on the host (_wire_host_channel) and in the worker (_run_one_job), and _build_job recomputed the cross_invoker gate that Mediator.start had already decided. Both decisions now happen once in _build_job and ride worker_opts: the host reads backward_active when wiring the channel (gating real-activation retention), the worker reads it at adopt time (gating delivered-clone tagging), and cross_invoker reuses the mediator's already-set value. This is now the single place to tighten the substring detection (it can false-positive in comments, costing only needless tagging). Isolated trace/backward/cross-invoke/multi-token tests all pass. Co-Authored-By: Claude Fable 5 --- src/nnsight/intervention/isolation.py | 34 ++++++++++++++------------- 1 file changed, 18 insertions(+), 16 deletions(-) diff --git a/src/nnsight/intervention/isolation.py b/src/nnsight/intervention/isolation.py index ce9d54f80..6c3a38114 100644 --- a/src/nnsight/intervention/isolation.py +++ b/src/nnsight/intervention/isolation.py @@ -408,24 +408,25 @@ def _build_job(mediator) -> tuple: mediator.intervention.__source__ = "".join(mediator.info.source) payload = serialization.dumps(mediator) - from .. import CONFIG - worker_opts = { # default_all (= generate's max_new_tokens) bounds an open-ended iter[:] on # the worker; it is set AFTER spawn so the live value is also piggybacked on # each response (meta), but seed the job with the value known now. "default_all": mediator.interleaver.default_all, - # cross_invoker matches the in-process gate (Mediator.start): multiple invokes - # + config enabled. The worker can't share a frame, so it pushes/pulls through - # the host store (see _run_one_job + meta below). - "cross_invoker": ( - len(mediator.interleaver.mediators) > 1 and CONFIG.APP.CROSS_INVOKER - ), + # Mediator.start already applied the in-process cross_invoker gate (multiple + # invokes + config) before acquiring the worker; reuse its decision. The + # worker can't share a frame, so it pushes/pulls through the host store. + "cross_invoker": bool(mediator.cross_invoker), + # `with tensor.backward()` detection — the single decision point for BOTH + # sides: the host gates real-activation retention, the worker gates + # delivered-clone tagging. (The substring can false-positive, e.g. in a + # comment, which only costs needless tagging; tighten it here when needed.) + "backward_active": ".backward(" in mediator.intervention.__source__, } return payload, extras, worker_opts -def _wire_host_channel(mediator, iso: _PooledWorker) -> None: +def _wire_host_channel(mediator, iso: _PooledWorker, worker_opts: dict) -> None: """Point the (possibly recycled) worker's host channel at THIS mediator.""" chan = iso.channel chan.reset() # fresh single-slot buffer + startup-timeout @@ -444,8 +445,9 @@ def _wire_host_channel(mediator, iso: _PooledWorker) -> None: mediator._iso = iso # `with tensor.backward()`: if the trace differentiates, the host must retain each # delivered (real, on-graph) activation so handle_backward_event can run the real - # backward. Detect it from the source and start with a fresh retention map. - mediator._iso_backward = ".backward(" in mediator.intervention.__source__ + # backward. The decision was made once in _build_job (shared with the worker); + # start with a fresh retention map. + mediator._iso_backward = worker_opts["backward_active"] mediator._iso_grad_reals = {} # Fresh per-job host-side hook-registration state. iso.registered = set() @@ -487,7 +489,7 @@ def _acquire(): _POOL.forget(iso) iso = _acquire() iso.send_job(payload, extras, worker_opts) - _wire_host_channel(mediator, iso) + _wire_host_channel(mediator, iso, worker_opts) def release_isolated_worker(iso: _PooledWorker, dirty: bool) -> None: @@ -652,10 +654,10 @@ def adopt(cls, mediator, channel, interleaver, opts: dict, device) -> "WorkerMed mediator.cross_invoker = opts.get("cross_invoker", False) mediator._isolated_worker = True # so Barrier sends the target count (host counts) mediator._device = device - # `with tensor.backward()`: detected from the intervention source; gates the - # delivered-clone tagging in ``request`` (BackwardsTracer reads the provenance). - source = "".join(mediator.info.source) if mediator.info.source else "" - mediator._bwd_active = ".backward(" in source + # `with tensor.backward()`: decided once in _build_job (host-side) and shipped + # with the job; gates the delivered-clone tagging in ``request`` + # (BackwardsTracer reads the provenance). + mediator._bwd_active = opts["backward_active"] mediator._bwd_prov = {} # id(delivered clone) -> requester string mediator._bwd_tagged = [] # delivered clones made to require grad interleaver.current = mediator From c1fbbc96e818c86b1de08419319b572147c29463 Mon Sep 17 00:00:00 2001 From: khaiwang Date: Wed, 10 Jun 2026 23:13:29 -0400 Subject: [PATCH 11/30] refactor(intervention): one per-job reset owner + shared path-to-envoy map MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Per-job worker-handle state was reset twice: at release (reset_for_release) and again at the next acquire (_wire_host_channel). Acquire now owns the authoritative reset (it runs unconditionally for both pooled and cold workers); release keeps only reference-dropping so an idle worker doesn't pin the last trace's hook set, path map, or — via channel.reset() — the meta/push callbacks closing over its interleaver. The {path: envoy} resolution map, previously built ad-hoc in two places (cached by host-side hook registration, rebuilt from scratch on every CACHE event), is now one lazy helper cached per job on the worker handle. Pool recycle, cache, and trace isolated tests all pass. Co-Authored-By: Claude Fable 5 --- src/nnsight/intervention/interleaver.py | 5 +++-- src/nnsight/intervention/isolation.py | 22 ++++++++++++++++------ 2 files changed, 19 insertions(+), 8 deletions(-) diff --git a/src/nnsight/intervention/interleaver.py b/src/nnsight/intervention/interleaver.py index b00787216..d3da4b050 100755 --- a/src/nnsight/intervention/interleaver.py +++ b/src/nnsight/intervention/interleaver.py @@ -1462,8 +1462,9 @@ def handle_cache_event(self, spec: Any): alias, ) = spec - model = self.interleaver.tracer.model - path2envoy = {e.path: e for e in model.modules()} + from .isolation import path_to_envoy + + path2envoy = path_to_envoy(self) targets = [path2envoy[p] for p in paths if p in path2envoy] cache_obj = Cache( diff --git a/src/nnsight/intervention/isolation.py b/src/nnsight/intervention/isolation.py index 6c3a38114..e1ebbc7f7 100644 --- a/src/nnsight/intervention/isolation.py +++ b/src/nnsight/intervention/isolation.py @@ -173,10 +173,12 @@ def send_job(self, payload, extras, opts) -> None: self.conn.send(("job", payload, extras, opts)) def reset_for_release(self) -> None: - """Clear per-job host state so the worker can serve the next trace.""" + """Drop per-job references so an idle worker doesn't pin them (the hook set, + the path→envoy map, and — via ``channel.reset()`` — the meta/push callbacks + closing over the last trace's interleaver). The *authoritative* per-job reset + happens at acquire time in ``_wire_host_channel``; this is memory hygiene.""" self.registered = set() self.path2envoy = None - self.clean = False self.channel.reset() def close(self) -> None: @@ -509,6 +511,17 @@ def release_isolated_worker(iso: _PooledWorker, dirty: bool) -> None: _POOL.forget(iso) +def path_to_envoy(mediator) -> dict: + """The ``{path: envoy}`` map for this job's model, built lazily and cached on the + worker handle (fresh per job — ``_wire_host_channel`` clears it). Shared by + host-side hook registration and ``handle_cache_event``'s target resolution.""" + iso = mediator._iso + if iso.path2envoy is None: + model = mediator.interleaver.tracer.model + iso.path2envoy = {e.path: e for e in model.modules()} + return iso.path2envoy + + def ensure_isolated_provider(mediator, requester: str) -> None: """Host-side hook registration: register the one-shot hook for ``requester`` on the *real* module. @@ -533,10 +546,7 @@ def ensure_isolated_provider(mediator, requester: str) -> None: if kind not in ("output", "input"): return # externally-provided eproperties (e.g. .result) need no module hook - if iso.path2envoy is None: - model = mediator.interleaver.tracer.model - iso.path2envoy = {e.path: e for e in model.modules()} - envoy = iso.path2envoy.get(path) + envoy = path_to_envoy(mediator).get(path) if envoy is None: return # unknown path → let normal missed-provider handling surface it From ad074a4475408f7e00933406aec70e75976f7dd7 Mon Sep 17 00:00:00 2001 From: khaiwang Date: Wed, 10 Jun 2026 23:26:47 -0400 Subject: [PATCH 12/30] refactor(intervention): IsoOptions dataclass, keyed CACHE spec, retire gap test - A frozen IsoOptions dataclass replaces the four hand-copied option dicts (_STATE fields, _base_opts(), _WorkerPool._key(), warm_worker_pool's rebuild). pool_key lives on it, making the warm-time (device/arena/mem-fraction/lockdown) vs per-job (timeout) split explicit; the phantom never-set "startup_timeout" option becomes the _WARM_STARTUP_TIMEOUT constant. - The CACHE event spec crosses the wire as a keyword dict instead of a 9-field positional tuple, so adding a cache option can't silently shift fields. - The gap-characterization test is retired: both gaps it proved are closed and its assertions duplicate test_isolated_cache.py / test_isolated_backward.py (weaker, in the cache case). - The doc's duplicate feature-map and support-matrix tables fold into one table carrying mechanism + status. - Doc records a PRE-EXISTING break found while re-running the full suite: lockdown has been broken since the warm-pool unification (the worker locks down before its first job-recv, and unpickling the job's tokenizer extras needs a new transformers submodule import that seccomp blocks). Reproduced on the pre-refactor commit 8d09195; needs a separate fix decision. Warm-pool suite passes after the test helper moved to IsoOptions.pool_key (reuse/concurrent/retire/dead-idle/exception-recycle/renamed-model, plus trace/cache/backward and the rest of the isolated suite earlier in the stack); in-process regression 51 passed. Co-Authored-By: Claude Fable 5 --- .../mediator-gpu-trace-integration.md | 61 +++---- .../test_isolated_backward_cache_gaps.py | 102 ------------ .../gpu_sandbox/test_isolated_pool.py | 6 +- src/nnsight/intervention/interleaver.py | 32 ++-- src/nnsight/intervention/isolation.py | 153 ++++++++++-------- src/nnsight/intervention/tracing/tracer.py | 22 +-- 6 files changed, 135 insertions(+), 241 deletions(-) delete mode 100644 prototypes/mediator-sandbox/gpu_sandbox/test_isolated_backward_cache_gaps.py diff --git a/docs/developing/mediator-gpu-trace-integration.md b/docs/developing/mediator-gpu-trace-integration.md index 3b554f5b2..065922dff 100644 --- a/docs/developing/mediator-gpu-trace-integration.md +++ b/docs/developing/mediator-gpu-trace-integration.md @@ -162,25 +162,13 @@ rule; deferred (revisit with double-buffering if real workloads need it). --- -## 6. Feature coverage map (so nothing is silently dropped) +## 6. Feature coverage map -| Feature | Cross-process mechanism | Status | -|---|---|---| -| read / swap / skip / exception | events over channel + host-side hook registration | done | -| `.save()` (tensors) | worker→host saves transmission | done | -| multi-invoke + batch narrowing | per-worker host mediator; Batcher host-side | done | -| `iter`/`all`/`next` (multi-token) | iteration step stamped host-side; host iter-hooks bump the tracker; worker sets its step explicitly | done | -| `tracer.barrier()` | host-side participant counting + the existing `handle_barrier_event` coordination loop | done | -| `cross_invoker` variable sharing | host-mediated variable store (worker pushes data locals, pulls the merged store) | done | -| `with tensor.backward()` / `.grad` | BACKWARD event: worker seeds `dL/d(delivered clone)` from its local tape, host continues `torch.autograd.grad` on the real graph, `.grad` served by provenance path | done (§16) | -| `tracer.cache()` | CACHE event → host registers the real cache hooks; host CacheDict swapped in for the worker placeholder, filled in-place by the forward | done (§15) | -| warm worker pool (`pool_size=`) | generic workers receive serialized mediators as jobs over the channel; clean-END workers recycled, others retired | done (§14) | -| MPS / `isolate_mediators()` further polish | — | planned | - -When isolation is on and a not-yet-supported feature is used, the trace **fails cleanly** — a -missed-provider error or the per-step timeout (the lifecycle is the safety net), not a silent deadlock or -silent-wrong result. (There is no automatic "route to in-process" fallback; features are added one at a -time. See the support matrix in §8 for what works today.) +Folded into the **support matrix (§8)**, which carries each feature's cross-process mechanism and +current status in one table. The invariant it tracks: when isolation is on and a not-yet-supported +feature is used, the trace **fails cleanly** — a missed-provider error or the per-step timeout (the +lifecycle is the safety net), not a silent deadlock or silent-wrong result. (There is no automatic +"route to in-process" fallback; features are added one at a time.) --- @@ -245,21 +233,23 @@ benign CudaIPC release warning. --- -## 8. Support matrix (what works under `isolate_mediators()` today) - -| Feature | Status | -|---|---| -| read / swap (`=`) / `.save()` (tensors) / skip / exception / multi-invoke | ✅ bit-identical | -| single-forward `generate(...)` (no iter) | ✅ verified | -| seccomp lockdown (fs/net/exec) | ✅ | -| `iter`/`all`/`next` (multi-token) | ✅ bit-identical (`iter[N]`, `iter[:]`, per-step swap) | -| `tracer.barrier()` | ✅ host-side participant counting | -| `cross_invoker` variable sharing | ✅ host variable store; transmittable data vars only — see §10 | -| warm worker pool (`pool_size=N`, `warm_worker_pool`) | ✅ ~21× faster per request once warm; recycle-on-clean-END — see §14 | -| `with tensor.backward()` / `.grad` | ✅ read-path bit-identical (scalar loss; single invoke; no swap-then-backward; `.grad` editing raises — §16) | -| `tracer.cache()` (`modules=`, `include_inputs=`) | ✅ bit-identical — CACHE event → host registers the real cache hooks; the forward fills the host CacheDict in-place (§15) | -| `.source` operation-level access (`...attn.split_1.output`) | 🔜 not yet (op paths aren't in `model.modules()`) | -| in-place `[:]=` | ⛔ use explicit `=` (clone semantics, §4) | +## 8. Support matrix (what works under `isolate_mediators()` today, and how) + +| Feature | Cross-process mechanism | Status | +|---|---|---| +| read / swap (`=`) / `.save()` (tensors) / skip / exception | six events over the channel; host-side hook registration; worker→host saves transmission at END | ✅ bit-identical | +| multi-invoke + batch narrowing | per-invoke worker + host mediator; Batcher stays host-side | ✅ bit-identical | +| single-forward `generate(...)` (no iter) | same as trace | ✅ verified | +| seccomp lockdown (fs/net/exec) | `_sandbox.lock_down` after warm-up | ⚠️ broken since the warm-pool unification (§14): the worker locks down BEFORE receiving its first job, and unpickling the job's extras (Tokenizer) triggers a new `transformers` submodule import that seccomp blocks → the worker dies at job-recv. Pre-dates the 2026-06-10 refactors (reproduced on the pre-refactor commit). Needs a fix decision: warm the model's transformers modules before lockdown, or restore deserialize-before-lockdown for cold workers. | +| `iter`/`all`/`next` (multi-token) | step stamped in the requester; host iter-hooks bump the tracker; live `default_all` piggyback (§9) | ✅ bit-identical (`iter[N]`, `iter[:]`, per-step swap) | +| `tracer.barrier()` | worker sends the target count; host accumulates participants + runs the coordination loop (§10) | ✅ | +| `cross_invoker` variable sharing | host variable store; worker pushes data locals, pulls the merged store; transmittable data only (§10) | ✅ | +| warm worker pool (`pool_size=N`, `warm_worker_pool`) | generic workers receive serialized mediators as jobs; clean-END recycle (§14) | ✅ ~21× faster per request once warm | +| `with tensor.backward()` / `.grad` | BACKWARD event: worker seeds `dL/d(delivered clone)`, host continues `torch.autograd.grad` on the real graph, `.grad` by provenance path (§16) | ✅ read-path bit-identical (scalar loss; single invoke; no swap-then-backward; `.grad` editing raises) | +| `tracer.cache()` (`modules=`, `include_inputs=`) | CACHE event → host registers the real cache hooks; host CacheDict swapped in at END, filled in-place by the forward (§15) | ✅ bit-identical | +| `.source` operation-level access (`...attn.split_1.output`) | — (op paths aren't in `model.modules()`) | 🔜 not yet | +| in-place `[:]=` | — (clone-on-receive semantics; use explicit `=`, §4) | ⛔ | +| MPS / `isolate_mediators()` further polish | — | planned | Not-yet-supported features fail **cleanly** (missed-provider error or the per-step timeout), not as a silent deadlock — the lifecycle (timeout + `finally: cancel()`) is the safety net until each feature lands. @@ -332,8 +322,9 @@ perf cliff for large cross-invoke tensors). ## 11. Backward + caching — characterization (both gaps since closed: cache §15, backward §16) -`test_isolated_backward_cache_gaps.py` confirmed the two gaps and their difficulty at the time; -kept for the reasoning record: +A gap-characterization harness (`test_isolated_backward_cache_gaps.py`, since retired — canonical +coverage lives in `test_isolated_cache.py` / `test_isolated_backward.py`) confirmed the two gaps and +their difficulty at the time; kept for the reasoning record: - **`tracer.cache()` (a real build, not a quick shim):** `tracer.cache()` runs in the worker → registers cache hooks on dummy modules → never fire → the `.save()`'d CacheDict comes back **empty**. The fix diff --git a/prototypes/mediator-sandbox/gpu_sandbox/test_isolated_backward_cache_gaps.py b/prototypes/mediator-sandbox/gpu_sandbox/test_isolated_backward_cache_gaps.py deleted file mode 100644 index 70b3e0b54..000000000 --- a/prototypes/mediator-sandbox/gpu_sandbox/test_isolated_backward_cache_gaps.py +++ /dev/null @@ -1,102 +0,0 @@ -#!/usr/bin/env python3 -"""backward/grad + cache() under isolation (former gaps, now closed). - - backward — now SUPPORTED (increment 1): grad of an ON-PATH tensor-output module - (``ln_f.output``) runs HOST-SIDE under isolation and is bit-identical to - in-process. The host keeps the real graph; the worker computes its half of - the chain rule at the delivered-activation seam and reads grads back by - PATH. Canonical coverage in test_isolated_backward.py. (Still open: gradient - THROUGH a swap, which severs the host graph at the patch point.) - cache — SUPPORTED: ``tracer.cache(modules=[...])`` is bit-identical under - isolation (kept here as a regression check; see test_isolated_cache.py). - -Each isolated-vs-in-process, hard timeout so a deadlock shows as a timeout. - -Run: - CUDA_VISIBLE_DEVICES=6 PYTHONPATH=src \ - /disk/u/zikai/anaconda3/envs/hf-serve/bin/python \ - prototypes/mediator-sandbox/gpu_sandbox/test_isolated_backward_cache_gaps.py -""" -import sys - -import torch - -from nnsight import LanguageModel -from nnsight.intervention.isolation import isolate_mediators - -PROMPT = "The Eiffel Tower is in the city of" - - -def _run(fn): - try: - return ("ok", fn()) - except Exception as e: # noqa: BLE001 - return ("err", f"{type(e).__name__}: {str(e)[:120]}") - - -def test_backward(model): - # Increment 1 closed this gap: read-then-backward on an on-path tensor-output module - # (ln_f.output) now runs HOST-SIDE under isolation, bit-identical to in-process. - def body(): - with model.trace(PROMPT): - hs = model.transformer.ln_f.output - with model.lm_head.output.sum().backward(): - g = hs.grad.save() - return g - - rs, rv = _run(body) - - def iso(): - with isolate_mediators(timeout=25): - return body() - - gs, gv = _run(iso) - - ok = ( - rs == "ok" - and gs == "ok" - and torch.is_tensor(rv) - and torch.is_tensor(gv) - and torch.equal(rv, gv) - ) - print( - f"[backward] in-process={rs}; isolated={gs}; bit-identical={ok} " - f"({'' if ok else (gv if gs == 'err' else 'mismatch')})", - flush=True, - ) - return ok - - -def test_cache(model): - def body(): - with model.trace(PROMPT) as t: - cache = t.cache(modules=[model.transformer.h[6]]).save() - return cache - rs, rv = _run(body) - def iso(): - with isolate_mediators(timeout=25): - return body() - gs, gv = _run(iso) - # compare the cached output for h[6] if both ok - ok = False - if rs == "ok" and gs == "ok": - try: - rk = list(rv.keys()) if hasattr(rv, "keys") else rv - gk = list(gv.keys()) if hasattr(gv, "keys") else gv - ok = str(rk) == str(gk) and len(rk) > 0 - except Exception: - ok = False - print(f"[cache] ref={rs} got={gs} ok={ok} (gv={gv if gs!='ok' else 'CacheDict'})", flush=True) - return ok - - -def main(): - assert torch.cuda.is_available() - model = LanguageModel("gpt2", device_map="cuda", dispatch=True) - results = {"backward": test_backward(model), "cache": test_cache(model)} - print("=" * 72, flush=True) - print(f"BACKWARD/CACHE GAPS: {results}", flush=True) - - -if __name__ == "__main__": - main() diff --git a/prototypes/mediator-sandbox/gpu_sandbox/test_isolated_pool.py b/prototypes/mediator-sandbox/gpu_sandbox/test_isolated_pool.py index a86954266..8fd9e8daf 100644 --- a/prototypes/mediator-sandbox/gpu_sandbox/test_isolated_pool.py +++ b/prototypes/mediator-sandbox/gpu_sandbox/test_isolated_pool.py @@ -160,9 +160,9 @@ def forward(self, x): def _pool_key(device="cuda", arena=64 << 20, frac=0.3, lock=False): - return isolation._WorkerPool._key( - {"device": device, "arena_bytes": arena, "gpu_mem_fraction": frac, "lockdown": lock} - ) + return isolation.IsoOptions( + device=device, arena_bytes=arena, gpu_mem_fraction=frac, lockdown=lock + ).pool_key def test_dead_idle(model): diff --git a/src/nnsight/intervention/interleaver.py b/src/nnsight/intervention/interleaver.py index d3da4b050..4cad55a88 100755 --- a/src/nnsight/intervention/interleaver.py +++ b/src/nnsight/intervention/interleaver.py @@ -1448,36 +1448,30 @@ def handle_cache_event(self, spec: Any): in-process path. Acks so the worker's ``send`` returns and it proceeds to END. """ from .hooks import cache_input_hook, cache_output_hook - from .tracing.tracer import Cache - - ( - token, - paths, - device, - dtype, - detach, - include_output, - include_inputs, - rename, - alias, - ) = spec - from .isolation import path_to_envoy + from .tracing.tracer import Cache path2envoy = path_to_envoy(self) - targets = [path2envoy[p] for p in paths if p in path2envoy] + targets = [path2envoy[p] for p in spec["paths"] if p in path2envoy] cache_obj = Cache( - paths, device, dtype, detach, include_output, include_inputs, rename, alias + spec["paths"], + spec["device"], + spec["dtype"], + spec["detach"], + spec["include_output"], + spec["include_inputs"], + spec["rename"], + spec["alias"], ) batcher = self.interleaver.batcher for envoy in targets: - if include_output: + if spec["include_output"]: cache_output_hook(cache_obj, envoy._module, envoy.path, batcher, self) - if include_inputs: + if spec["include_inputs"]: cache_input_hook(cache_obj, envoy._module, envoy.path, batcher, self) - self._iso_caches[token] = cache_obj + self._iso_caches[spec["token"]] = cache_obj self.set_user_cache(cache_obj) self.respond(None) # ack -> worker's send() returns, it proceeds to END diff --git a/src/nnsight/intervention/isolation.py b/src/nnsight/intervention/isolation.py index e1ebbc7f7..f1b933e28 100644 --- a/src/nnsight/intervention/isolation.py +++ b/src/nnsight/intervention/isolation.py @@ -39,6 +39,7 @@ import threading from collections import defaultdict, deque from contextlib import contextmanager +from dataclasses import dataclass from typing import Any, Dict, Optional import torch @@ -77,14 +78,46 @@ def _transmittable(v) -> bool: # --------------------------------------------------------------------------- # # Opt-in surface # # --------------------------------------------------------------------------- # +@dataclass(frozen=True) +class IsoOptions: + """Per-worker isolation options, in one place for every consumer (the opt-in + context, the pool key, spawn, and the worker bootstrap — previously four + hand-copied dicts). + + ``device``/``arena_bytes``/``gpu_mem_fraction``/``lockdown`` are **warm-time**: + fixed when the worker process spawns, and they define pool interchangeability + (:attr:`pool_key`). ``timeout`` is **per-job**: re-applied to the host channel + each trace in ``_wire_host_channel`` (it also sizes the channel's first-event + budget at spawn). + """ + + device: str = "cuda" + arena_bytes: int = 64 << 20 + gpu_mem_fraction: float = 0.3 + lockdown: bool = False # functional-first; seccomp lockdown enabled separately + timeout: float = 60.0 # per-step wall-clock cap on user code (hang containment) + + @property + def pool_key(self) -> tuple: + # Workers are interchangeable ONLY within the same warm-time signature — the + # bounce buffer is device- and size-specific, so reusing a worker across + # devices would copy into the wrong-device buffer (silent corruption). + return ( + str(self.device), + int(self.arena_bytes), + float(self.gpu_mem_fraction), + bool(self.lockdown), + ) + + +# Generous wait for a COLD spawn+warm (import torch/nnsight + CUDA init, ~4 s typical); +# distinct from the per-job first-event budget (user timeout + deserialize margin). +_WARM_STARTUP_TIMEOUT = 180.0 + _STATE: Dict[str, Any] = { "on": False, - "arena_bytes": 64 << 20, - "gpu_mem_fraction": 0.3, - "device": "cuda", - "timeout": 60.0, # per-step wall-clock cap on user code (hang containment) - "lockdown": False, # functional-first; seccomp lockdown enabled separately - "pool_size": 0, # 0 => cold one-shot worker per trace; >0 => warm pool cap + "pool_size": 0, # 0 => cold one-shot worker per trace; >0 => warm pool cap + "opts": IsoOptions(), } @@ -122,12 +155,14 @@ def isolate_mediators( prev = dict(_STATE) _STATE.update( on=True, - arena_bytes=arena_bytes, - gpu_mem_fraction=gpu_mem_fraction, - device=device, - timeout=timeout, - lockdown=lockdown, pool_size=pool_size, + opts=IsoOptions( + device=device, + arena_bytes=arena_bytes, + gpu_mem_fraction=gpu_mem_fraction, + lockdown=lockdown, + timeout=timeout, + ), ) try: yield @@ -135,17 +170,6 @@ def isolate_mediators( _STATE.update(prev) -def _base_opts() -> Dict[str, Any]: - """The per-worker (warm-time) options, distinct from per-job (per-trace) ones.""" - return { - "device": _STATE["device"], - "arena_bytes": _STATE["arena_bytes"], - "gpu_mem_fraction": _STATE["gpu_mem_fraction"], - "lockdown": _STATE["lockdown"], - "timeout": _STATE["timeout"], - } - - # --------------------------------------------------------------------------- # # Host side — worker handle + pool # # --------------------------------------------------------------------------- # @@ -213,37 +237,26 @@ class _WorkerPool: """ def __init__(self): - # Keyed by base-opts signature: workers are interchangeable ONLY within the - # same (device, arena_bytes, gpu_mem_fraction, lockdown) — the bounce buffer is - # device- and size-specific, so reusing a worker across devices would copy into - # the wrong-device buffer (silent corruption). + # Keyed by IsoOptions.pool_key (the warm-time signature): workers are + # interchangeable only within it — see IsoOptions. self._idle: Dict[tuple, deque] = defaultdict(deque) self._all: Dict[tuple, set] = defaultdict(set) self._lock = threading.Lock() self._shutting_down = False - @staticmethod - def _key(base_opts: dict) -> tuple: - return ( - str(base_opts["device"]), - int(base_opts["arena_bytes"]), - float(base_opts["gpu_mem_fraction"]), - bool(base_opts.get("lockdown", False)), - ) - - def warm(self, n: int, base_opts: dict) -> None: - key = self._key(base_opts) + def warm(self, n: int, opts: IsoOptions) -> None: + key = opts.pool_key with self._lock: need = max(0, n - len(self._all[key])) # Spawn outside the lock (each ~4 s); register under it. - fresh = [_spawn_worker(base_opts, poolable=True) for _ in range(need)] + fresh = [_spawn_worker(opts, poolable=True) for _ in range(need)] with self._lock: for w in fresh: self._all[key].add(w) self._idle[key].append(w) - def acquire(self, base_opts: dict, cap: int) -> _PooledWorker: - key = self._key(base_opts) + def acquire(self, opts: IsoOptions, cap: int) -> _PooledWorker: + key = opts.pool_key dead: list = [] live = None placeholder = None @@ -269,7 +282,7 @@ def acquire(self, base_opts: dict, cap: int) -> _PooledWorker: return live if placeholder is not None: try: - w = _spawn_worker(base_opts, poolable=True) + w = _spawn_worker(opts, poolable=True) except BaseException: with self._lock: self._all[key].discard(placeholder) @@ -279,7 +292,7 @@ def acquire(self, base_opts: dict, cap: int) -> _PooledWorker: self._all[key].add(w) return w # At cap with none idle: a cold one-shot worker so the trace never blocks. - return _spawn_worker(base_opts, poolable=False) + return _spawn_worker(opts, poolable=False) def put_idle(self, w: _PooledWorker) -> None: close_it = False @@ -333,13 +346,13 @@ def warm_worker_pool( """ _POOL.warm( size, - { - "device": device, - "arena_bytes": arena_bytes, - "gpu_mem_fraction": gpu_mem_fraction, - "lockdown": lockdown, - "timeout": timeout, - }, + IsoOptions( + device=device, + arena_bytes=arena_bytes, + gpu_mem_fraction=gpu_mem_fraction, + lockdown=lockdown, + timeout=timeout, + ), ) @@ -348,24 +361,23 @@ def shutdown_worker_pool() -> None: _POOL.shutdown() -def _spawn_worker(base_opts: dict, poolable: bool) -> _PooledWorker: +def _spawn_worker(opts: IsoOptions, poolable: bool) -> _PooledWorker: """Spawn a generic worker, wait for its one-time ``ready`` ack, wire the channel.""" ctx = mp.get_context("spawn") # CUDA requires spawn, not fork - buf = torch.empty( - base_opts["arena_bytes"], dtype=torch.uint8, device=base_opts["device"] - ) + buf = torch.empty(opts.arena_bytes, dtype=torch.uint8, device=opts.device) parent_conn, child_conn = ctx.Pipe() proc = ctx.Process( - target=_pool_worker_main, args=(child_conn, buf, base_opts), daemon=True + target=_pool_worker_main, args=(child_conn, buf, opts), daemon=True ) proc.start() # The worker warms CUDA + imports (~4 s) then sends exactly one "ready"; consume # it before the channel starts reading protocol frames on the same pipe. This poll # covers the cold spawn+warm, so it stays generous. - warm_wait = base_opts.get("startup_timeout", 180.0) - if not parent_conn.poll(warm_wait): + if not parent_conn.poll(_WARM_STARTUP_TIMEOUT): proc.terminate() - raise TimeoutError(f"isolated worker failed to warm up within {warm_wait}s") + raise TimeoutError( + f"isolated worker failed to warm up within {_WARM_STARTUP_TIMEOUT}s" + ) msg = parent_conn.recv() if msg != "ready": proc.terminate() @@ -377,12 +389,11 @@ def _spawn_worker(base_opts: dict, poolable: bool) -> _PooledWorker: chan = CudaIpcHostChannel( parent_conn, buf, - timeout=base_opts["timeout"], - startup_timeout=base_opts["timeout"] + _JOB_STARTUP_MARGIN, + timeout=opts.timeout, + startup_timeout=opts.timeout + _JOB_STARTUP_MARGIN, ) return _PooledWorker( - proc, buf, parent_conn, chan, poolable=poolable, - key=_WorkerPool._key(base_opts), + proc, buf, parent_conn, chan, poolable=poolable, key=opts.pool_key ) @@ -431,8 +442,8 @@ def _build_job(mediator) -> tuple: def _wire_host_channel(mediator, iso: _PooledWorker, worker_opts: dict) -> None: """Point the (possibly recycled) worker's host channel at THIS mediator.""" chan = iso.channel - chan.reset() # fresh single-slot buffer + startup-timeout - chan._timeout = _STATE["timeout"] # per-trace user-code cap + chan.reset() # fresh single-slot buffer + startup-timeout + chan._timeout = _STATE["opts"].timeout # per-trace user-code cap # default_all is set by generate() AFTER the worker is acquired (LanguageModel # ._execute), so a snapshot is stale. Piggyback the LIVE value + the cross_invoker # var store on each response; the worker reads default_all before bounding its @@ -474,12 +485,12 @@ def acquire_isolated_worker(mediator) -> None: registration + cancel/release). """ payload, extras, worker_opts = _build_job(mediator) - pool_size = _STATE.get("pool_size", 0) + pool_size = _STATE["pool_size"] def _acquire(): if pool_size > 0: - return _POOL.acquire(_base_opts(), pool_size) - return _spawn_worker(_base_opts(), poolable=False) + return _POOL.acquire(_STATE["opts"], pool_size) + return _spawn_worker(_STATE["opts"], poolable=False) iso = _acquire() try: @@ -794,7 +805,7 @@ def _run_one_job(channel, payload, extras, opts, device) -> None: _WORKER_CURRENT = None -def _pool_worker_main(conn, buf, base_opts): +def _pool_worker_main(conn, buf, worker_iso_opts: IsoOptions): """Generic worker: warm CUDA + imports ONCE, optionally lock down, then loop serving ``("job", payload, extras, opts)`` messages until told to ``"stop"``. @@ -803,7 +814,7 @@ def _pool_worker_main(conn, buf, base_opts): against fresh dummy modules (no cross-job state but ``Globals.saves``, cleared).""" from .tracing.globals import _ensure_mounted - device = base_opts.get("device", "cuda") + device = worker_iso_opts.device # Warm CUDA before any lockdown so kernels/contexts are loaded. if torch.cuda.is_available(): @@ -821,8 +832,8 @@ def _pool_worker_main(conn, buf, base_opts): _ensure_mounted() # install Object.save so `.save()` resolves in the worker - if base_opts.get("gpu_mem_fraction") and torch.cuda.is_available(): - torch.cuda.set_per_process_memory_fraction(base_opts["gpu_mem_fraction"]) + if worker_iso_opts.gpu_mem_fraction and torch.cuda.is_available(): + torch.cuda.set_per_process_memory_fraction(worker_iso_opts.gpu_mem_fraction) channel = CudaIpcWorkerChannel(conn, buf) # persistent; rebinds handlers per job @@ -831,7 +842,7 @@ def _pool_worker_main(conn, buf, base_opts): # import set for ALL jobs (jobs whose user code triggers a NEW import will fail) — # lockdown defaults off; document the trade-off. CUDA + the control Pipe + the IPC # buffer use already-open fds. - if base_opts.get("lockdown"): + if worker_iso_opts.lockdown: from ._sandbox import lock_down lock_down() diff --git a/src/nnsight/intervention/tracing/tracer.py b/src/nnsight/intervention/tracing/tracer.py index f7b0d3c69..be5ba8466 100755 --- a/src/nnsight/intervention/tracing/tracer.py +++ b/src/nnsight/intervention/tracing/tracer.py @@ -625,17 +625,17 @@ def cache( token = id(cache_obj.cache) cache_obj.cache._iso_cache_token = token - spec = ( - token, - [envoy.path for envoy in targets], - device, - dtype, - detach, - include_output, - include_inputs, - rename_dict, - alias_dict, - ) + spec = { + "token": token, + "paths": [envoy.path for envoy in targets], + "device": device, + "dtype": dtype, + "detach": detach, + "include_output": include_output, + "include_inputs": include_inputs, + "rename": rename_dict, + "alias": alias_dict, + } mediator.send(Events.CACHE, spec) else: for envoy in targets: From 195bee328189c5d14db234cfc218689f3d6a528b Mon Sep 17 00:00:00 2001 From: khaiwang Date: Fri, 12 Jun 2026 00:33:19 -0400 Subject: [PATCH 13/30] perf(intervention): drop retained backward activations at mediator cancel MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit After an isolated backward trace ended, the host mediator kept its references to every retained on-graph activation — pinning those tensors and the autograd graph behind them until the mediator was GC'd. cancel() already drops the mediator's other ephemeral state (history, iteration tracker, worker handle); now it also clears the retention map. Safe because every BACKWARD event precedes the END/exception that triggers cancel. Found by a four-angle cleanup pass over the backward + refactor stack; the other findings were judged false positives or already- documented accepted costs. Backward + multi-token-backward isolated tests and the in-process regression (51 passed) stay green. Co-Authored-By: Claude Fable 5 --- src/nnsight/intervention/interleaver.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/src/nnsight/intervention/interleaver.py b/src/nnsight/intervention/interleaver.py index 4cad55a88..0a7b895ba 100755 --- a/src/nnsight/intervention/interleaver.py +++ b/src/nnsight/intervention/interleaver.py @@ -1133,6 +1133,9 @@ def cancel(self): self.iteration_tracker = defaultdict(int) self.iteration = 0 self.worker = None + # Retained on-graph activations (isolated backward) pin the autograd graph; + # drop them at trace end rather than waiting for the mediator to be GC'd. + self._iso_grad_reals = {} # If the worker is still mid-protocol, unwind it with a Cancelation. For the # isolated channel the get_event/put_response calls are host-local — they do NOT From ed6c55698ab9c1d00c1640c5c6affac6008a451c Mon Sep 17 00:00:00 2001 From: khaiwang Date: Fri, 12 Jun 2026 16:14:31 -0400 Subject: [PATCH 14/30] fix(intervention): initialize _iso_caches in Mediator.__setstate__ __setstate__ rebuilds the transient isolation fields (_iso, _isolated_worker, _iso_backward, _iso_grad_reals) but missed _iso_caches, so a deserialized mediator (the NDIF/vLLM server path constructs mediators via __setstate__) running under isolate_mediators() would hit AttributeError in handle_cache_event the first time a trace used tracer.cache(). Latent locally (host mediators come from __init__); found by the high-effort review pass. Isolated cache test and in-process regression (51 passed) green. Co-Authored-By: Claude Fable 5 --- src/nnsight/intervention/interleaver.py | 1 + 1 file changed, 1 insertion(+) diff --git a/src/nnsight/intervention/interleaver.py b/src/nnsight/intervention/interleaver.py index 0a7b895ba..6eccf8569 100755 --- a/src/nnsight/intervention/interleaver.py +++ b/src/nnsight/intervention/interleaver.py @@ -1823,6 +1823,7 @@ def __setstate__(self, state): self._isolated_worker = False self._iso_backward = False self._iso_grad_reals = {} + self._iso_caches = {} self.interleaver = None self.history = set() self.user_cache: "Cache" = list() From 093d343c96ec4526a7c632740cd4284e2158d099 Mon Sep 17 00:00:00 2001 From: khaiwang Date: Sat, 13 Jun 2026 00:48:33 -0400 Subject: [PATCH 15/30] feat(intervention): in-process fast lane for confirmed-safe interventions MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Process isolation contains footguns by running each intervention in a spawned GPU worker — but that worker holds a weightless path-only mirror of the model, so the interp majority (logit lens, steering, ablation, activation patching, attribution) cannot run isolated at all: they read the host model's real weights (F.linear(x, head.weight)) and call its final-norm / unembed modules. The fast lane is the tier where the real weights live. Adds a third execution tier under isolate_mediators(). A fail-closed, default-deny static classifier (fastlane.py) walks the EFFECTIVE code of each mediator — the trace body plus every user closure it calls, resolved through the frame / function globals / closure cells (the harness wraps real compute in build()/capture() closures, so a walk of the with-block alone would see only an opaque call). Verdicts: FAST (only whitelisted ops / host-model access / nnsight primitives -> run in-process at full speed and full model access), ISOLATE (anything unconfirmable -> the existing GPU worker), REJECT (an introspection escape -> raise). The conservative default is ISOLATE; the gate is a footgun selector, not a malice boundary, so it is cordoned to trust="local" provenance and a CONFIG.APP.FAST_LANE flag. Default behavior is preserved: isolation off never consults the gate; isolation on now fast-lanes the confirmed-safe majority and isolates the rest. A best-effort wall-clock watchdog restores loop-containment for the one footgun the static walk cannot bound (a huge bounded range); its injected FastLaneTimeout rides the intervention body's existing try/except, so the host re-raises it cleanly. The classifier's closure-aware backward detection also replaces the old `.backward(` source substring (blind to a backward hidden in a closure) for the isolated job's grad-retention flag. Deferred (documented): the process-global sys.addaudithook backstop (its leaked-flag failure mode can abort the model's own forward — net-negative under a static default-deny gate); the five declarative tracer primitives (unembed/steer/patch/ablate/capture) that would let weight-reading cells also run on the isolated tier via host event handlers (a cache-shaped build). Verified: classifier units 17/17 (logit-lens/steering/patching/attribution shapes + renamed structures classify FAST; imports/while/unresolved-call/ open ISOLATE; introspection REJECT; flag detection). Fast-lane e2e 6/6 on gpt2 + a renamed model: weight-reading lens bit-identical on the fast lane (max|Δ|=0) AND raises under forced isolation; in-place steering bit-identical; footgun routes off the fast lane, host survives; introspection rejected; runaway loop killed by the watchdog, host survives. Existing isolated WORKER path (pinned with fast_lane=False) 9/9 still bit-identical; in-process core 51 passed. Co-Authored-By: Claude Fable 5 --- .../gpu_sandbox/perf_spawn_cost.py | 2 +- .../gpu_sandbox/perf_spawn_split.py | 2 +- .../gpu_sandbox/test_fast_lane.py | 194 ++++++ .../gpu_sandbox/test_fastlane_classifier.py | 226 ++++++ .../gpu_sandbox/test_isolated_acceptance.py | 8 +- .../gpu_sandbox/test_isolated_backward.py | 6 +- .../gpu_sandbox/test_isolated_cache.py | 6 +- .../gpu_sandbox/test_isolated_cross_invoke.py | 4 +- .../test_isolated_lockdown_safety.py | 8 +- .../test_isolated_multitoken_backward.py | 4 +- .../test_isolated_multitoken_iter.py | 6 +- .../gpu_sandbox/test_isolated_pool.py | 18 +- .../gpu_sandbox/test_isolated_trace.py | 4 +- .../gpu_sandbox/test_nonstd.py | 4 +- src/nnsight/intervention/fastlane.py | 651 ++++++++++++++++++ src/nnsight/intervention/interleaver.py | 96 ++- src/nnsight/intervention/isolation.py | 78 ++- src/nnsight/schema/config.py | 1 + 18 files changed, 1256 insertions(+), 62 deletions(-) create mode 100644 prototypes/mediator-sandbox/gpu_sandbox/test_fast_lane.py create mode 100644 prototypes/mediator-sandbox/gpu_sandbox/test_fastlane_classifier.py create mode 100644 src/nnsight/intervention/fastlane.py diff --git a/prototypes/mediator-sandbox/gpu_sandbox/perf_spawn_cost.py b/prototypes/mediator-sandbox/gpu_sandbox/perf_spawn_cost.py index 82914b148..ed6bd6b69 100644 --- a/prototypes/mediator-sandbox/gpu_sandbox/perf_spawn_cost.py +++ b/prototypes/mediator-sandbox/gpu_sandbox/perf_spawn_cost.py @@ -165,7 +165,7 @@ def _timed_spawn(med): spawn_times.append((time.perf_counter() - t0) * 1e3) def one_isolated(): - with isolate_mediators(): + with isolate_mediators(fast_lane=False): with model.trace(PROMPT): _ = model.transformer.h[6].output[0].save() diff --git a/prototypes/mediator-sandbox/gpu_sandbox/perf_spawn_split.py b/prototypes/mediator-sandbox/gpu_sandbox/perf_spawn_split.py index d29e808d7..12adda45d 100644 --- a/prototypes/mediator-sandbox/gpu_sandbox/perf_spawn_split.py +++ b/prototypes/mediator-sandbox/gpu_sandbox/perf_spawn_split.py @@ -29,7 +29,7 @@ def timed_start(self): mpc.SpawnProcess.start = timed_start def one_iso(): - with isolate_mediators(): + with isolate_mediators(fast_lane=False): with model.trace(PROMPT): _ = model.transformer.h[6].output[0].save() diff --git a/prototypes/mediator-sandbox/gpu_sandbox/test_fast_lane.py b/prototypes/mediator-sandbox/gpu_sandbox/test_fast_lane.py new file mode 100644 index 000000000..d020d00a8 --- /dev/null +++ b/prototypes/mediator-sandbox/gpu_sandbox/test_fast_lane.py @@ -0,0 +1,194 @@ +#!/usr/bin/env python3 +"""End-to-end fast lane: confirmed-safe interventions run IN-PROCESS under +isolate_mediators(), and that is the ONLY tier that can run the weight-reading interp +majority (the isolated worker holds weightless dummy modules). + + weights_fast — a logit-lens cell (reads lm_head.weight, calls ln_f) runs under + isolate_mediators() bit-identical to non-isolated AND raises under + forced isolation (fast_lane=False) — proving the fast lane is the + enabling tier, not just a speedup. + inplace_steer — an in-place steering cell runs correctly on the fast lane (it is a + silent no-op under isolation due to clone-on-receive). + renamed — same logit-lens shape on a renamed model (final_norm / + output_projection / decoder_blocks) fast-lanes bit-identical. + footgun_isolates — a cell that imports os / opens a file is NOT fast-laned: it routes to + the worker (isolated) and is contained, host survives. + introspection — a cell reaching for ().__class__... raises FastLaneRejected. + watchdog — a huge bounded loop fast-laned with a short deadline is killed by the + watchdog; the host survives and the next trace still works. + +Run: + CUDA_VISIBLE_DEVICES=5 PYTHONPATH=src \ + /disk/u/zikai/anaconda3/envs/hf-serve/bin/python -u \ + prototypes/mediator-sandbox/gpu_sandbox/test_fast_lane.py +""" +import sys + +import torch +import torch.nn.functional as F + +from nnsight import LanguageModel +from nnsight.intervention.fastlane import FastLaneRejected +from nnsight.intervention.isolation import isolate_mediators + +PROMPT = "The Eiffel Tower is in the city of" + + +# --- the weight-reading interp readout (logit lens), the workload isolation cannot run -- +def _logit_lens(blocks, norm, head, layers): + rows = [] + with torch.no_grad(): + for i in layers: + out = blocks[i].output + hidden = out[0] if isinstance(out, tuple) else out + normed = norm(hidden) + logits = F.linear(normed, head.weight) # host-weight read — worker-impossible + rows.append(logits[:, -1, :]) + return torch.stack(rows, dim=0) + + +def _lens_gpt2(model): + with model.trace(PROMPT): + g = _logit_lens(model.transformer.h, model.transformer.ln_f, + model.lm_head, [0, 4, 8]).save() + return g + + +def test_weights_fast(model): + ref = _lens_gpt2(model) # non-isolated baseline + with isolate_mediators(): # fast lane on (default) + got = _lens_gpt2(model) + fast_ok = torch.is_tensor(got) and torch.equal(ref, got) + + # forced isolation: the same cell must FAIL (weightless dummy modules in the worker) + forced_failed = False + try: + with isolate_mediators(fast_lane=False, timeout=30): + _lens_gpt2(model) + except Exception: # noqa: BLE001 + forced_failed = True + + print(f"[weights_fast] fast-lane bit-identical={fast_ok} " + f"(max|Δ|={(ref-got).abs().max().item() if fast_ok else float('nan')}); " + f"forced-isolation-raises={forced_failed}", flush=True) + return fast_ok and forced_failed + + +def test_inplace_steer(model): + def steer(inplace_ctx): + # steer block 6 toward a token's unembed row, read final logits + with model.trace(PROMPT): + with torch.no_grad(): + direction = F.normalize(model.lm_head.weight[5000].float(), dim=0) + out = model.transformer.h[6].output + hidden = out[0] if isinstance(out, tuple) else out + scale = hidden.norm(dim=-1).mean() + hidden[:] = hidden + 6.0 * scale * direction.to(hidden.dtype) # in-place + last = model.transformer.h[-1].output + last = last[0] if isinstance(last, tuple) else last + normed = model.transformer.ln_f(last) + logits = F.linear(normed, model.lm_head.weight)[:, -1, :].save() + return logits + + ref = steer(None) # non-isolated + with isolate_mediators(): + got = steer(None) # fast lane + ok = torch.is_tensor(got) and torch.equal(ref, got) + print(f"[inplace_steer] fast-lane in-place write bit-identical={ok} " + f"(max|Δ|={(ref-got).abs().max().item() if ok else float('nan')})", flush=True) + return ok + + +def test_renamed(): + rename = {"transformer.ln_f": "final_norm", "lm_head": "output_projection", + "transformer.h": "decoder_blocks"} + model = LanguageModel("gpt2", device_map="cuda", dispatch=True, rename=rename) + + def body(): + with model.trace(PROMPT): + g = _logit_lens(model.decoder_blocks, model.final_norm, + model.output_projection, [0, 3, 6]).save() + return g + + ref = body() + with isolate_mediators(): + got = body() + ok = torch.is_tensor(got) and torch.equal(ref, got) + print(f"[renamed] renamed-model logit-lens fast-lane bit-identical={ok} " + f"(max|Δ|={(ref-got).abs().max().item() if ok else float('nan')})", flush=True) + return ok + + +def test_footgun_isolates(model): + # A footgun cell is NOT confirmable -> routes to the worker (isolated), contained. + # The worker can't run the weight read either, so it errors cleanly; the point is the + # host SURVIVES and a subsequent normal trace still works. + raised = False + try: + with isolate_mediators(fast_lane=True, timeout=20): + with model.trace(PROMPT): + import os # noqa: F401 — footgun: isolated, never fast-laned + _ = model.transformer.h[6].output[0].save() + except Exception: # noqa: BLE001 + raised = True + # host survives: a clean fast-lane trace works right after + after = _lens_gpt2(model) + survived = torch.is_tensor(after) + print(f"[footgun_isolates] footgun routed off the fast lane (raised={raised}); " + f"host survived={survived}", flush=True) + return survived + + +def test_introspection(model): + rejected = False + try: + with isolate_mediators(): + with model.trace(PROMPT): + _ = ().__class__.__bases__ # introspection escape + _ = model.transformer.h[6].output[0].save() + except FastLaneRejected: + rejected = True + except Exception as e: # noqa: BLE001 + rejected = "introspection" in str(e).lower() or "reject" in str(e).lower() + print(f"[introspection] introspection escape rejected={rejected}", flush=True) + return rejected + + +def test_watchdog(model): + # A huge bounded loop passes the static gate (range is "bounded") but would hang the + # host; the watchdog must kill it. Short deadline so the test is quick. + killed = False + try: + with isolate_mediators(fast_lane_timeout=4.0): + with model.trace(PROMPT): + h = model.transformer.h[6].output[0] + acc = 0 + for _i in range(10 ** 12): # bounded literal, but enormous + acc += 1 + _ = h.save() + except Exception: # noqa: BLE001 — FastLaneTimeout surfaces through the trace + killed = True + after = _lens_gpt2(model) # host survived + survived = torch.is_tensor(after) + print(f"[watchdog] runaway loop killed={killed}; host survived={survived}", flush=True) + return killed and survived + + +def main(): + assert torch.cuda.is_available() + model = LanguageModel("gpt2", device_map="cuda", dispatch=True) + results = { + "weights_fast": test_weights_fast(model), + "inplace_steer": test_inplace_steer(model), + "renamed": test_renamed(), + "footgun_isolates": test_footgun_isolates(model), + "introspection": test_introspection(model), + "watchdog": test_watchdog(model), + } + print("=" * 72, flush=True) + print(f"FAST LANE (end-to-end): {results}", flush=True) + sys.exit(0 if all(results.values()) else 1) + + +if __name__ == "__main__": + main() diff --git a/prototypes/mediator-sandbox/gpu_sandbox/test_fastlane_classifier.py b/prototypes/mediator-sandbox/gpu_sandbox/test_fastlane_classifier.py new file mode 100644 index 000000000..f13bddf4c --- /dev/null +++ b/prototypes/mediator-sandbox/gpu_sandbox/test_fastlane_classifier.py @@ -0,0 +1,226 @@ +#!/usr/bin/env python3 +"""Unit tests for the fast-lane safety classifier (GPU-free, no model). + +Confirms the detect-and-confirm gate on the ACTUAL interp-workload shapes (logit lens, +steering, activation patching, attribution) plus footgun payloads, and — per the testing +rules — on RENAMED module structures so nothing is keyed to GPT-2 naming. The classifier +walks the effective code (closures resolved through globals/closure cells), so these tests +use the same closure-wrapped shape the real harness uses (a build()/capture() lambda +calling helper functions). + +Run: + PYTHONPATH=src /disk/u/zikai/anaconda3/envs/hf-serve/bin/python -u \ + prototypes/mediator-sandbox/gpu_sandbox/test_fastlane_classifier.py +""" +import sys + +import torch +import torch.nn as nn +import torch.nn.functional as F + +from nnsight.intervention.fastlane import FAST, ISOLATE, REJECT, classify_callable + + +# --- stand-ins for the host objects a trace body closes over ------------------------- +# Real nn.Modules so the classifier's host-object detection (calling them / reading +# .weight) exercises the genuine path, not a mock. +class _Block(nn.Module): + def __init__(self, d): + super().__init__() + self.lin = nn.Linear(d, d) + self.output = torch.zeros(1, 3, d) # stands in for the eproperty read target + + def forward(self, x): + return self.lin(x) + + +def _make_model(norm_name="ln_f", head_name="lm_head", blocks_name="h", d=8, n=4): + """A tiny model with CONFIGURABLE attribute names — the renamed-structure test tool.""" + m = nn.Module() + setattr(m, blocks_name, nn.ModuleList([_Block(d) for _ in range(n)])) + setattr(m, norm_name, nn.LayerNorm(d)) + setattr(m, head_name, nn.Linear(d, 50, bias=False)) + return m + + +# --- helper functions a "trace body" calls (the closure chain) ------------------------ +def _untuple(x): + return x[0] if isinstance(x, tuple) else x + + +def _lens_proxy(blocks, norm, head, *, layers): + rows = [] + with torch.no_grad(): + for i in (range(len(blocks)) if layers == "all" else layers): + normed = norm(_untuple(blocks[i].output)) + logits = F.linear(normed, head.weight) # host-weight read — fast-lane-only + rows.append(logits[:, -1, :]) + return torch.stack(rows, dim=0) + + +def _steer_inplace(blocks, head, *, layer, token_id, alpha): + with torch.no_grad(): + direction = F.normalize(head.weight[token_id].float(), dim=0) + out = blocks[layer].output + hidden = out[0] if isinstance(out, tuple) else out + hidden[:] = hidden + alpha * direction # in-place write (in_place flag) + return hidden + + +# ===================================================================================== +CASES = [] + + +def case(expect): + def deco(fn): + CASES.append((fn.__name__, fn, expect)) + return fn + return deco + + +# ---- SAFE interp workloads (must classify FAST) ------------------------------------- +@case(FAST) +def logit_lens_gpt2_style(): + m = _make_model() + return lambda: _lens_proxy(m.h, m.ln_f, m.lm_head, layers="all").save() + + +@case(FAST) +def logit_lens_renamed_structure(): + # non-GPT-2 names: decoder_blocks / final_norm / output_projection + m = _make_model(norm_name="final_norm", head_name="output_projection", + blocks_name="decoder_blocks") + return lambda: _lens_proxy(m.decoder_blocks, m.final_norm, m.output_projection, + layers=[0, 2]).save() + + +@case(FAST) +def steering_inplace_is_fast(): + # in-place steering is the fast lane's correctness win (silent no-op under isolation) + m = _make_model() + return lambda: _steer_inplace(m.h, m.lm_head, layer=1, token_id=5, alpha=6.0).save() + + +@case(FAST) +def boundary_replacement_write(): + m = _make_model() + + def body(): + out = m.h[1].output + new = out * 2.0 + m.h[2].output = new # nnsight boundary write (SWAP), allowed + return m.ln_f(new)[:, -1, :] + return body + + +@case(FAST) +def backward_attribution_shape(): + m = _make_model() + + def body(): + a = m.h[2].output + a.requires_grad_(True) + normed = m.ln_f(a) + metric = F.linear(normed, m.lm_head.weight).sum() + with metric.backward(): + g = a.grad.save() + return g + return body + + +# ---- footguns that must NOT reach the fast lane ------------------------------------- +@case(ISOLATE) +def imports_isolate(): + def body(): + import os + return os.getpid() + return body + + +@case(ISOLATE) +def while_loop_isolate(): + def body(): + x = 0 + while x < 10: + x = x + 1 + return x + return body + + +@case(ISOLATE) +def unresolved_global_call_isolate(): + def body(): + return some_undefined_helper(3) # noqa: F821 — unknown authority + return body + + +@case(ISOLATE) +def open_file_isolate(): + def body(): + return open("/etc/passwd").read() + return body + + +@case(REJECT) +def introspection_subclasses_reject(): + def body(): + return ().__class__.__bases__ + return body + + +@case(REJECT) +def getattr_escape_reject(): + def body(): + return getattr(torch, "save") + return body + + +@case(REJECT) +def dunder_subscript_reject(): + def body(): + d = {} + return d["__builtins__"] + return body + + +@case(ISOLATE) +def host_attr_write_isolate(): + m = _make_model() + + def body(): + m.ln_f.eps = 1.0 # mutating host state visible to siblings + return m.ln_f.weight + return body + + +# ---- flag detection ---------------------------------------------------------------- +def main(): + results = {} + for name, factory, expect in CASES: + fn = factory() + v = classify_callable(fn) + ok = v.tier == expect + results[name] = ok + print(f"[{'PASS' if ok else 'FAIL'}] {name}: got {v.tier} expected {expect} " + f"(diff={v.differentiate} inplace={v.in_place} weights={v.touches_host_weights}) " + f"-- {v.reason}", flush=True) + + # flag assertions on specific cases + flag_checks = { + "backward sets differentiate": classify_callable(backward_attribution_shape()).differentiate is True, + "steering sets in_place": classify_callable(steering_inplace_is_fast()).in_place is True, + "lens reads host weights": classify_callable(logit_lens_gpt2_style()).touches_host_weights is True, + "replacement write is not in_place": classify_callable(boundary_replacement_write()).in_place is False, + } + for k, v in flag_checks.items(): + results[k] = v + print(f"[{'PASS' if v else 'FAIL'}] {k}", flush=True) + + print("=" * 72, flush=True) + npass = sum(results.values()) + print(f"FAST-LANE CLASSIFIER: {npass}/{len(results)} passed", flush=True) + sys.exit(0 if all(results.values()) else 1) + + +if __name__ == "__main__": + main() diff --git a/prototypes/mediator-sandbox/gpu_sandbox/test_isolated_acceptance.py b/prototypes/mediator-sandbox/gpu_sandbox/test_isolated_acceptance.py index 30f12362f..53de5a94e 100644 --- a/prototypes/mediator-sandbox/gpu_sandbox/test_isolated_acceptance.py +++ b/prototypes/mediator-sandbox/gpu_sandbox/test_isolated_acceptance.py @@ -47,7 +47,7 @@ def test_nonstandard_names(): with model.trace(x): ref = model.decoder_blocks[1].output.save() - with isolate_mediators(): + with isolate_mediators(fast_lane=False): with model.trace(x): got = model.decoder_blocks[1].output.save() ok = torch.equal(ref, got) @@ -62,7 +62,7 @@ def test_multi_invoke(model): a_ref = model.transformer.h[5].output[0].save() with t.invoke("The Eiffel Tower is in"): b_ref = model.transformer.h[5].output[0].save() - with isolate_mediators(): + with isolate_mediators(fast_lane=False): with model.trace() as t: with t.invoke("The capital of France is"): a_got = model.transformer.h[5].output[0].save() @@ -78,7 +78,7 @@ def test_exception(model): # A footgun (ValueError) in user code must surface in the user's context. raised = None try: - with isolate_mediators(): + with isolate_mediators(fast_lane=False): with model.trace(PROMPT): _ = model.transformer.h[6].output[0] raise ValueError("boom-from-user-code") @@ -94,7 +94,7 @@ def test_timeout(model): t0 = time.time() killed = None try: - with isolate_mediators(timeout=5): + with isolate_mediators(fast_lane=False, timeout=5): with model.trace(PROMPT): out = model.transformer.h[6].output[0] while True: # footgun: hang diff --git a/prototypes/mediator-sandbox/gpu_sandbox/test_isolated_backward.py b/prototypes/mediator-sandbox/gpu_sandbox/test_isolated_backward.py index bd782836d..a288a9fbf 100644 --- a/prototypes/mediator-sandbox/gpu_sandbox/test_isolated_backward.py +++ b/prototypes/mediator-sandbox/gpu_sandbox/test_isolated_backward.py @@ -38,7 +38,7 @@ def _read_backward(model): def test_read_backward(model): ref = _read_backward(model) - with isolate_mediators(timeout=30): + with isolate_mediators(fast_lane=False, timeout=30): got = _read_backward(model) ok = ( @@ -71,7 +71,7 @@ def body(): return g ref = body() - with isolate_mediators(timeout=30): + with isolate_mediators(fast_lane=False, timeout=30): got = body() ok = torch.is_tensor(ref) and torch.is_tensor(got) and torch.equal(ref, got) delta = (ref - got).abs().max().item() if ok or (torch.is_tensor(ref) and torch.is_tensor(got)) else float("nan") @@ -115,7 +115,7 @@ def main(): def _derived_isolated(model): - with isolate_mediators(timeout=30): + with isolate_mediators(fast_lane=False, timeout=30): return test_derived_target_fails_clean(model) diff --git a/prototypes/mediator-sandbox/gpu_sandbox/test_isolated_cache.py b/prototypes/mediator-sandbox/gpu_sandbox/test_isolated_cache.py index c2fe303ce..a53a6d2fd 100644 --- a/prototypes/mediator-sandbox/gpu_sandbox/test_isolated_cache.py +++ b/prototypes/mediator-sandbox/gpu_sandbox/test_isolated_cache.py @@ -29,7 +29,7 @@ def test_one(model): h6 = model.transformer.h[6] with model.trace(PROMPT) as t: ref = t.cache(modules=[h6]).save() - with isolate_mediators(timeout=30): + with isolate_mediators(fast_lane=False, timeout=30): with model.trace(PROMPT) as t: got = t.cache(modules=[h6]).save() key = h6.path # derive the key from the envoy path, don't hardcode the prefix @@ -45,7 +45,7 @@ def test_multi(model): mods = [model.transformer.h[2], model.transformer.h[5], model.transformer.h[9]] with model.trace(PROMPT) as t: ref = t.cache(modules=mods).save() - with isolate_mediators(timeout=30): + with isolate_mediators(fast_lane=False, timeout=30): with model.trace(PROMPT) as t: got = t.cache(modules=mods).save() paths = [m.path for m in mods] @@ -61,7 +61,7 @@ def test_inputs(model): h4 = model.transformer.h[4] with model.trace(PROMPT) as t: ref = t.cache(modules=[h4], include_inputs=True).save() - with isolate_mediators(timeout=30): + with isolate_mediators(fast_lane=False, timeout=30): with model.trace(PROMPT) as t: got = t.cache(modules=[h4], include_inputs=True).save() key = h4.path diff --git a/prototypes/mediator-sandbox/gpu_sandbox/test_isolated_cross_invoke.py b/prototypes/mediator-sandbox/gpu_sandbox/test_isolated_cross_invoke.py index f4cd32481..e7a9a65c8 100644 --- a/prototypes/mediator-sandbox/gpu_sandbox/test_isolated_cross_invoke.py +++ b/prototypes/mediator-sandbox/gpu_sandbox/test_isolated_cross_invoke.py @@ -46,7 +46,7 @@ def body(): rs, rv = _run(body) def iso(): - with isolate_mediators(timeout=20): + with isolate_mediators(fast_lane=False, timeout=20): return body() gs, gv = _run(iso) ok = rs == "ok" and gs == "ok" and torch.equal(rv[0], gv[0]) and torch.equal(rv[1], gv[1]) @@ -69,7 +69,7 @@ def body(): rs, rv = _run(body) def iso(): - with isolate_mediators(timeout=20): + with isolate_mediators(fast_lane=False, timeout=20): return body() gs, gv = _run(iso) ok = ( diff --git a/prototypes/mediator-sandbox/gpu_sandbox/test_isolated_lockdown_safety.py b/prototypes/mediator-sandbox/gpu_sandbox/test_isolated_lockdown_safety.py index 08cb984df..36feb68d2 100644 --- a/prototypes/mediator-sandbox/gpu_sandbox/test_isolated_lockdown_safety.py +++ b/prototypes/mediator-sandbox/gpu_sandbox/test_isolated_lockdown_safety.py @@ -7,7 +7,7 @@ net — socket()/connect() in user code is blocked. (The standalone seccomp primitive is separately proven by gpu_sandbox/test_safety.py; -this checks it is correctly wired into model.trace via isolate_mediators(lockdown=True).) +this checks it is correctly wired into model.trace via isolate_mediators(fast_lane=False, lockdown=True).) Run: CUDA_VISIBLE_DEVICES=0 PYTHONPATH=src \ @@ -29,7 +29,7 @@ def test_functional_under_lockdown(model): with model.trace(PROMPT): ref = model.transformer.h[6].output[0].save() - with isolate_mediators(lockdown=True): + with isolate_mediators(fast_lane=False, lockdown=True): with model.trace(PROMPT): got = model.transformer.h[6].output[0].save() ok = torch.equal(ref, got) @@ -42,7 +42,7 @@ def test_fs_blocked(model): os.remove(PROBE) raised = None try: - with isolate_mediators(lockdown=True): + with isolate_mediators(fast_lane=False, lockdown=True): with model.trace(PROMPT): with open(PROBE, "w") as f: # should be EPERM under seccomp f.write("escaped") @@ -60,7 +60,7 @@ def test_fs_blocked(model): def test_net_blocked(model): raised = None try: - with isolate_mediators(lockdown=True): + with isolate_mediators(fast_lane=False, lockdown=True): with model.trace(PROMPT): import socket s = socket.socket(socket.AF_INET, socket.SOCK_STREAM) # should be EPERM diff --git a/prototypes/mediator-sandbox/gpu_sandbox/test_isolated_multitoken_backward.py b/prototypes/mediator-sandbox/gpu_sandbox/test_isolated_multitoken_backward.py index 6a88c4a45..ae460fcb7 100644 --- a/prototypes/mediator-sandbox/gpu_sandbox/test_isolated_multitoken_backward.py +++ b/prototypes/mediator-sandbox/gpu_sandbox/test_isolated_multitoken_backward.py @@ -58,7 +58,7 @@ def _per_step(model, iso): step's sum-of-logits loss.""" def body(): - ctx = isolate_mediators(timeout=60) if iso else _null() + ctx = isolate_mediators(fast_lane=False, timeout=60) if iso else _null() with ctx: with model.generate(PROMPT, max_new_tokens=N) as t: grads = [] @@ -78,7 +78,7 @@ def _post_loop(model, iso): from the step-1 activation itself; .grad read in the backward block.""" def body(): - ctx = isolate_mediators(timeout=60) if iso else _null() + ctx = isolate_mediators(fast_lane=False, timeout=60) if iso else _null() with ctx: with model.generate(PROMPT, max_new_tokens=N) as t: hs = [] diff --git a/prototypes/mediator-sandbox/gpu_sandbox/test_isolated_multitoken_iter.py b/prototypes/mediator-sandbox/gpu_sandbox/test_isolated_multitoken_iter.py index 615650724..c2e041d3c 100644 --- a/prototypes/mediator-sandbox/gpu_sandbox/test_isolated_multitoken_iter.py +++ b/prototypes/mediator-sandbox/gpu_sandbox/test_isolated_multitoken_iter.py @@ -29,7 +29,7 @@ def test_steps(model): with model.generate(PROMPT, max_new_tokens=N) as t: for step in t.iter[step_n]: ref = model.transformer.h[6].output[0].save() - with isolate_mediators(timeout=30): + with isolate_mediators(fast_lane=False, timeout=30): with model.generate(PROMPT, max_new_tokens=N) as t: for step in t.iter[step_n]: got = model.transformer.h[6].output[0].save() @@ -47,7 +47,7 @@ def test_swap(model): for step in t.iter[1]: model.transformer.h[6].output = model.transformer.h[6].output * 2 ref = model.transformer.h[7].output[0].save() - with isolate_mediators(timeout=30): + with isolate_mediators(fast_lane=False, timeout=30): with model.generate(PROMPT, max_new_tokens=N) as t: for step in t.iter[1]: model.transformer.h[6].output = model.transformer.h[6].output * 2 @@ -59,7 +59,7 @@ def test_swap(model): def test_allsaved(model): def run(iso): - ctx = isolate_mediators(timeout=30) if iso else _null() + ctx = isolate_mediators(fast_lane=False, timeout=30) if iso else _null() with ctx: with model.generate(PROMPT, max_new_tokens=N) as t: hs = [] diff --git a/prototypes/mediator-sandbox/gpu_sandbox/test_isolated_pool.py b/prototypes/mediator-sandbox/gpu_sandbox/test_isolated_pool.py index 8fd9e8daf..93bb6e8dd 100644 --- a/prototypes/mediator-sandbox/gpu_sandbox/test_isolated_pool.py +++ b/prototypes/mediator-sandbox/gpu_sandbox/test_isolated_pool.py @@ -45,7 +45,7 @@ def test_reuse(model): # Cold one-shot trace (no pool): pays the full ~4 s spawn. t0 = time.perf_counter() - with isolate_mediators(pool_size=0): + with isolate_mediators(fast_lane=False, pool_size=0): with model.trace(PROMPT): cold = model.transformer.h[6].output[0].save() cold_s = time.perf_counter() - t0 @@ -56,7 +56,7 @@ def test_reuse(model): got = None for _ in range(4): t0 = time.perf_counter() - with isolate_mediators(pool_size=2): + with isolate_mediators(fast_lane=False, pool_size=2): with model.trace(PROMPT): got = model.transformer.h[6].output[0].save() # the worker that just served this trace @@ -85,7 +85,7 @@ def test_concurrent(model): warm_worker_pool(3, device="cuda") pids_during = [] got = [] - with isolate_mediators(pool_size=3): + with isolate_mediators(fast_lane=False, pool_size=3): with model.trace() as tracer: for _ in range(3): with tracer.invoke(PROMPT): @@ -107,7 +107,7 @@ def test_retire(model): # A hung intervention: exceed the 2 s timeout -> the worker is killed, not recycled. timed_out = False try: - with isolate_mediators(pool_size=1, timeout=2.0): + with isolate_mediators(fast_lane=False, pool_size=1, timeout=2.0): with model.trace(PROMPT): h = model.transformer.h[6].output[0] # spin forever in user code on the worker @@ -122,7 +122,7 @@ def test_retire(model): retired = before not in survivors # The pool re-warms lazily and the NEXT trace still works, bit-identical. - with isolate_mediators(pool_size=1): + with isolate_mediators(fast_lane=False, pool_size=1): with model.trace(PROMPT): got = model.transformer.h[6].output[0].save() recovered = torch.equal(ref, got) @@ -149,7 +149,7 @@ def forward(self, x): with model.trace(x): ref = model.decoder_blocks[1].output.save() warm_worker_pool(1, device="cuda") - with isolate_mediators(pool_size=1): + with isolate_mediators(fast_lane=False, pool_size=1): with model.trace(x): got = model.decoder_blocks[1].output.save() ok = torch.equal(ref, got) @@ -175,7 +175,7 @@ def test_dead_idle(model): victim_pid = victim.pid victim.kill(); victim.join(timeout=5) # simulate OOM-kill/crash while idle - with isolate_mediators(pool_size=1): + with isolate_mediators(fast_lane=False, pool_size=1): with model.trace(PROMPT): got = model.transformer.h[6].output[0].save() ok = torch.equal(ref, got) @@ -196,7 +196,7 @@ def test_exception_recycle(model): raised = False try: - with isolate_mediators(pool_size=1): + with isolate_mediators(fast_lane=False, pool_size=1): with model.trace(PROMPT): _ = model.transformer.h[6].output[0].save() raise ValueError("boom") @@ -206,7 +206,7 @@ def test_exception_recycle(model): pids_after = {w.proc.pid for w in isolation._POOL._all[key]} recycled = pid_before in pids_after and len(pids_after) == 1 - with isolate_mediators(pool_size=1): + with isolate_mediators(fast_lane=False, pool_size=1): with model.trace(PROMPT): got = model.transformer.h[6].output[0].save() reused = next(iter(isolation._POOL._all[key])).proc.pid == pid_before diff --git a/prototypes/mediator-sandbox/gpu_sandbox/test_isolated_trace.py b/prototypes/mediator-sandbox/gpu_sandbox/test_isolated_trace.py index f52346c86..e01f5e937 100644 --- a/prototypes/mediator-sandbox/gpu_sandbox/test_isolated_trace.py +++ b/prototypes/mediator-sandbox/gpu_sandbox/test_isolated_trace.py @@ -26,7 +26,7 @@ def test_read(model): with model.trace(PROMPT): ref = model.transformer.h[6].output[0].save() - with isolate_mediators(): + with isolate_mediators(fast_lane=False): with model.trace(PROMPT): got = model.transformer.h[6].output[0].save() d = (ref.float() - got.float()).abs().max().item() @@ -44,7 +44,7 @@ def test_swap(model): with model.trace(PROMPT): model.transformer.h[6].output = model.transformer.h[6].output * 2 ref = model.transformer.h[7].output[0].save() - with isolate_mediators(): + with isolate_mediators(fast_lane=False): with model.trace(PROMPT): model.transformer.h[6].output = model.transformer.h[6].output * 2 got = model.transformer.h[7].output[0].save() diff --git a/prototypes/mediator-sandbox/gpu_sandbox/test_nonstd.py b/prototypes/mediator-sandbox/gpu_sandbox/test_nonstd.py index c658e8828..b0ddd2f70 100644 --- a/prototypes/mediator-sandbox/gpu_sandbox/test_nonstd.py +++ b/prototypes/mediator-sandbox/gpu_sandbox/test_nonstd.py @@ -27,7 +27,7 @@ def test_read(model): with model.trace(PROMPT): ref = model.decoder_blocks[6].output[0].save() - with isolate_mediators(timeout=30): + with isolate_mediators(fast_lane=False, timeout=30): with model.trace(PROMPT): got = model.decoder_blocks[6].output[0].save() ok = torch.equal(ref, got) @@ -39,7 +39,7 @@ def test_iterN(model): with model.generate(PROMPT, max_new_tokens=3) as t: for step in t.iter[1]: ref = model.decoder_blocks[6].output[0].save() - with isolate_mediators(timeout=30): + with isolate_mediators(fast_lane=False, timeout=30): with model.generate(PROMPT, max_new_tokens=3) as t: for step in t.iter[1]: got = model.decoder_blocks[6].output[0].save() diff --git a/src/nnsight/intervention/fastlane.py b/src/nnsight/intervention/fastlane.py new file mode 100644 index 000000000..559054193 --- /dev/null +++ b/src/nnsight/intervention/fastlane.py @@ -0,0 +1,651 @@ +"""Fast-lane safety classifier for isolated traces. + +`isolate_mediators()` runs each user intervention in a spawned GPU worker process to +contain footguns (infinite loops, OOM allocs, device-side asserts, host-object pokes, +fs/net/exec). But the worker holds a *weightless* path-only mirror of the model — its +dummy modules have no parameters and no ``forward``. The real interpretability workloads +(logit lens, steering, ablation, activation patching, attribution) all read the host +model's real weights (``F.linear(x, head.weight)``) and call its final-norm / unembed +modules, so **they cannot run in the worker at all**. The fast lane is the tier where the +real weights live: a confirmed-safe intervention runs in-process (the existing daemon- +thread path) at full speed and with full model access, and only unconfirmable code goes +to the worker. + +"Confirmed safe" is decided here by a **fail-closed, default-deny** static walk over the +*effective code* — the trace body PLUS every user closure it calls, resolved through the +frame, the function globals, and closure cells (the harness wraps real compute in +``build()`` / ``capture()`` / ``patch()`` closures, so a walk of the ``with`` block alone +would see only an opaque call). The walk emits one of three verdicts: + +- ``FAST`` — every node is on the allowlist and every call resolves to a whitelisted + op, a host module/weight access, an nnsight primitive, or a recursively + confirmed user function → run in-process. +- ``ISOLATE`` — anything unconfirmable (an unresolved global call, an import, a ``while`` + loop, an unrecoverable closure, an unknown node type) → run in the worker. +- ``REJECT`` — an introspection escape (``__globals__``/``getattr``/``eval``/…) → raise. + +The conservative default is ``ISOLATE``: the absence of proof is not proof of safety; only +explicitly whitelisted code reaches ``FAST``. + +**Threat model (load-bearing).** This is *not* an adversarial sandbox — a determined +author can defeat any in-process restriction (the pysandbox lesson), so the fast lane is +gated on ``trust="local"`` provenance and disabled for anything deserialized/remote. Under +the relaxed "contain footguns, not adversaries" model it confirms the effective code +(a) introduces no ambient authority (no import, no introspection, no unresolved global +call), (b) has no unbounded loop, (c) writes no host state, and (d) is composed only of +whitelisted ops, host-object access, and recursively confirmed user functions. OOM and +device-side asserts in pure tensor math are knowingly traded to the in-process tier (a +deployment that cannot tolerate them disables the fast lane); a wall-clock watchdog backs +the one loop footgun the static walk cannot bound (a huge ``range``). +""" +from __future__ import annotations + +import ast +import ctypes +import inspect +import textwrap +import threading +from dataclasses import dataclass, field + +# Verdict tiers. +FAST = "fast" +ISOLATE = "isolate" +REJECT = "reject" + +# Modules whose top-level functions are pure compute with no ambient authority. A call +# resolving into one of these is allowed outright. +_SAFE_MODULE_PREFIXES = ("torch", "math", "operator", "numpy") + +# torch entry points that DO touch fs/net/JIT despite living under `torch` — never fast. +_BANNED_QUALIFIED = { + "torch.load", "torch.save", "torch.hub", "torch.jit", "torch.compile", + "torch.onnx", "torch.multiprocessing", "torch.distributed", +} + +# Builtins that are pure / structural — safe to call from confirmed code. +_SAFE_BUILTINS = frozenset({ + "range", "len", "enumerate", "zip", "list", "tuple", "dict", "set", "frozenset", + "min", "max", "sum", "abs", "round", "sorted", "reversed", "map", "filter", + "float", "int", "bool", "str", "slice", "isinstance", "issubclass", "all", "any", + "print", "repr", "iter", "next", +}) + +# Names / attributes that are an introspection escape — their presence is a REJECT, not an +# isolate: trusted-local author code reaching for these is a footgun (or an attempt to +# break out), and the worker cannot run them either. Every documented in-process escape +# walks one of these (`().__class__.__subclasses__()`, `obj.__globals__`, a fetched +# builtin), which is exactly what pysandbox could not close. +_INTROSPECTION = frozenset({ + "eval", "exec", "compile", "__import__", "getattr", "setattr", "delattr", + "globals", "vars", "locals", "breakpoint", "memoryview", +}) + + +class FastLaneRejected(Exception): + """Raised when the classifier finds an introspection escape in a fast-lane-eligible + trace. The worker cannot run such code either, and under the relaxed footgun model + trusted-local author code reaching for introspection is a footgun — so fail loudly + rather than silently route it anywhere.""" + + +class FastLaneTimeout(Exception): + """Injected by the watchdog into a fast-lane intervention thread that overran the + wall-clock deadline (a runaway pure-Python loop). It is an ``Exception`` so the + intervention body's own ``except Exception`` catches it and routes it through the + normal ``mediator.exception`` path — the host re-raises it to the user, the model + server is unaffected.""" + + +class Watchdog: + """Best-effort wall-clock bound on a fast-lane intervention thread. + + The static gate bans ``while`` and unbounded iteration, so the only loop footgun that + can reach the fast lane is a huge bounded ``range`` or deep recursion. This injects a + :class:`FastLaneTimeout` into the running thread at its next bytecode if it overruns — + restoring the loop-containment guarantee that turning isolation on implies. It CANNOT + preempt a wedged native/CUDA call (only the worker-process kill can); the bounded-loop + static rule is the primary defense and this is the backstop. + """ + + def __init__(self, deadline_s: float): + self._deadline = deadline_s + self._timer = None + self._ident = None + + def arm(self, thread_ident: int) -> None: + self._ident = thread_ident + self._timer = threading.Timer(self._deadline, self._fire) + self._timer.daemon = True + self._timer.start() + + def _fire(self) -> None: + ident = self._ident + if ident is None: + return + res = ctypes.pythonapi.PyThreadState_SetAsyncExc( + ctypes.c_long(ident), ctypes.py_object(FastLaneTimeout) + ) + if res > 1: + # affected more than the target thread — undo to avoid corrupting others + ctypes.pythonapi.PyThreadState_SetAsyncExc(ctypes.c_long(ident), None) + + def disarm(self) -> None: + if self._timer is not None: + self._timer.cancel() + self._timer = None + self._ident = None + + +@dataclass +class Verdict: + """The classifier's decision for one mediator's intervention.""" + + tier: str # FAST | ISOLATE | REJECT + reason: str + differentiate: bool = False # the effective code opens a `with x.backward():` + touches_host_weights: bool = False + in_place: bool = False # an in-place write on a delivered boundary value + _seen: set = field(default_factory=set, repr=False) # visited code objects (recursion guard) + + @property + def fast(self) -> bool: + return self.tier == FAST + + +# Sentinel raised internally to short-circuit the walk on the first disqualifying node. +class _Stop(Exception): + def __init__(self, tier: str, reason: str): + self.tier = tier + self.reason = reason + + +def classify(mediator) -> Verdict: + """Classify a mediator's intervention for the fast lane. Always returns a Verdict + (never raises); a REJECT tier is surfaced as a raised error by the caller.""" + v = Verdict(tier=FAST, reason="confirmed: effective code is all whitelisted ops") + try: + nodes = _root_nodes(mediator) + if nodes is None: + return Verdict(ISOLATE, "could not recover the trace body source/AST") + ns = _root_namespace(mediator) + walker = _Walker(v) + walker.walk(nodes, ns, set()) + except _Stop as s: + return Verdict(s.tier, s.reason, v.differentiate, v.touches_host_weights, v.in_place) + except Exception as e: # noqa: BLE001 — any analysis failure is fail-closed + return Verdict(ISOLATE, f"classifier could not confirm safety: {type(e).__name__}: {e}") + return v + + +def classify_callable(fn) -> Verdict: + """Classify an arbitrary callable's body with the same engine `classify` uses on a + trace body. The effective-code walk resolves names through ``fn``'s globals and + closure cells. Used by the fast-lane tests to confirm verdicts on varied/renamed + module structures without standing up a model.""" + v = Verdict(tier=FAST, reason="confirmed: effective code is all whitelisted ops") + try: + _Walker(v)._recurse_fn(fn, depth=0) + except _Stop as s: + return Verdict(s.tier, s.reason, v.differentiate, v.touches_host_weights, v.in_place) + except Exception as e: # noqa: BLE001 + return Verdict(ISOLATE, f"classifier could not confirm safety: {type(e).__name__}: {e}") + return v + + +def _root_nodes(mediator): + """The statements of the trace body, preferring the live AST node, re-parsing the + captured source otherwise.""" + info = mediator.info + node = getattr(info, "node", None) + if isinstance(node, ast.With): + return list(node.body) + src = info.source + if not src: + return None + text = textwrap.dedent("".join(src)) + tree = ast.parse(text) + body = tree.body + if len(body) == 1 and isinstance(body[0], ast.With): + return list(body[0].body) + return body + + +def _root_namespace(mediator) -> dict: + """Name → object map for the trace body: the capturing frame's locals/globals plus + the compiled intervention's module globals.""" + ns = {} + fn = mediator.intervention + g = getattr(fn, "__globals__", None) + if isinstance(g, dict): + ns.update(g) + frame = getattr(mediator.info, "frame", None) + fl = getattr(frame, "f_globals", None) + if isinstance(fl, dict): + ns.update(fl) + fl = getattr(frame, "f_locals", None) + if isinstance(fl, dict): + ns.update(fl) + return ns + + +def _fn_namespace(fn) -> dict: + """Name → object map for a resolved user function: its globals plus closure cells.""" + ns = {} + g = getattr(fn, "__globals__", None) + if isinstance(g, dict): + ns.update(g) + code = getattr(fn, "__code__", None) + closure = getattr(fn, "__closure__", None) + if code is not None and closure: + for name, cell in zip(code.co_freevars, closure): + try: + ns[name] = cell.cell_contents + except ValueError: + pass # an empty cell (recursive def not yet bound) + return ns + + +class _Walker: + """Default-deny AST walk. Raises ``_Stop`` on the first ISOLATE/REJECT node; falling + off the end means every node was confirmed FAST.""" + + MAX_DEPTH = 12 + + def __init__(self, verdict: Verdict): + self.v = verdict + + # --- entry points --------------------------------------------------------- + def walk(self, nodes, ns: dict, local_names: set, depth: int = 0): + for n in nodes: + self._stmt(n, ns, local_names, depth) + + def _recurse_fn(self, fn, depth: int): + """Walk a resolved user function's body. Unrecoverable source → ISOLATE.""" + code = getattr(fn, "__code__", None) + if code is None or code in self.v._seen: + return + if depth >= self.MAX_DEPTH: + raise _Stop(ISOLATE, "intervention call graph is deeper than the fast-lane bound") + self.v._seen.add(code) + try: + src = textwrap.dedent(inspect.getsource(fn)) + except (OSError, TypeError): + raise _Stop(ISOLATE, f"could not recover source for `{getattr(fn, '__name__', fn)}`") + try: + tree = ast.parse(src) + except SyntaxError: + raise _Stop(ISOLATE, f"could not parse source for `{getattr(fn, '__name__', fn)}`") + target = tree.body[0] + ns = _fn_namespace(fn) + # The function's own parameters are locals we cannot resolve to objects (they are + # bound from the call site) — record them so a call/use of a param is treated as a + # host-object access under trust=local, not as an unknown global. + params = set() + node = target + if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)): + params = _arg_names(node.args) + self.walk(node.body, ns, params, depth + 1) + elif isinstance(node, ast.Expr) and isinstance(node.value, ast.Lambda): + lam = node.value + params = _arg_names(lam.args) + self._expr(lam.body, ns, params, depth + 1) + else: + # getsource gave us the surrounding statement (common for lambdas passed as + # args); find the first Lambda anywhere in it. + lam = next((d for d in ast.walk(tree) if isinstance(d, ast.Lambda)), None) + if lam is None: + raise _Stop(ISOLATE, f"could not isolate the body of `{getattr(fn, '__name__', fn)}`") + params = _arg_names(lam.args) + self._expr(lam.body, ns, params, depth + 1) + + # --- statements ----------------------------------------------------------- + def _stmt(self, n, ns, loc, depth): + if isinstance(n, (ast.Assign, ast.AnnAssign, ast.AugAssign)): + self._check_targets(n, ns, loc, depth) + if n.value is not None: + self._expr(n.value, ns, loc, depth) + # bind assigned simple names as locals (so later uses are not "unknown global") + for t in _assigned_names(n): + loc.add(t) + elif isinstance(n, ast.Expr): + self._expr(n.value, ns, loc, depth) + elif isinstance(n, ast.Return): + if n.value is not None: + self._expr(n.value, ns, loc, depth) + elif isinstance(n, ast.If): + self._expr(n.test, ns, loc, depth) + self.walk(n.body, ns, loc, depth) + self.walk(n.orelse, ns, loc, depth) + elif isinstance(n, ast.For): + # bounded by its iterable; an unbounded generator would be a user fn we walk, + # and a huge `range` is backed by the wall-clock watchdog. + for t in _target_names(n.target): + loc.add(t) + self._expr(n.iter, ns, loc, depth) + self.walk(n.body, ns, loc, depth) + self.walk(n.orelse, ns, loc, depth) + elif isinstance(n, ast.With): + for item in n.items: + self._with_item(item, ns, loc, depth) + self.walk(n.body, ns, loc, depth) + elif isinstance(n, ast.Raise): + # `raise ValueError(...)` is a normal guard; walk its expressions. + for child in (n.exc, n.cause): + if child is not None: + self._expr(child, ns, loc, depth) + elif isinstance(n, ast.Pass): + pass + elif isinstance(n, (ast.FunctionDef, ast.Lambda)): + # a def/lambda inside the body is a value; bind its name, walk when called + if isinstance(n, ast.FunctionDef): + loc.add(n.name) + loc |= _arg_names(n.args) + self.walk(n.body, ns, loc, depth) + elif isinstance(n, ast.While): + raise _Stop(ISOLATE, "a `while` loop cannot be statically bounded — isolated") + elif isinstance(n, (ast.Import, ast.ImportFrom)): + raise _Stop(ISOLATE, "an `import` introduces ambient authority — isolated") + elif isinstance(n, (ast.Global, ast.Nonlocal)): + raise _Stop(REJECT, "`global`/`nonlocal` rebinds shared state — rejected") + elif isinstance(n, ast.Try): + raise _Stop(ISOLATE, "a `try` block is not confirmable — isolated") + else: + raise _Stop(ISOLATE, f"unsupported statement `{type(n).__name__}` — isolated") + + def _with_item(self, item, ns, loc, depth): + ctx = item.context_expr + # detect `with .backward():` (closure-aware — this is the fix for the old + # `.backward(` substring that missed backwards hidden in a closure). + if isinstance(ctx, ast.Call) and isinstance(ctx.func, ast.Attribute) \ + and ctx.func.attr == "backward": + self.v.differentiate = True + self._expr(ctx, ns, loc, depth) + if item.optional_vars is not None: + for t in _target_names(item.optional_vars): + loc.add(t) + + def _check_targets(self, n, ns, loc, depth): + targets = n.targets if isinstance(n, ast.Assign) else [n.target] + for t in targets: + self._store_target(t, ns, loc, depth) + + def _store_target(self, t, ns, loc, depth): + if isinstance(t, ast.Name): + return # binding a local name — fine + if isinstance(t, (ast.Tuple, ast.List)): + for e in t.elts: + self._store_target(e, ns, loc, depth) + return + if isinstance(t, ast.Subscript): + # `hidden[:] = ...` — in-place write into a delivered/derived tensor. Correct + # in-process (shared memory); it is the steering default and is *safer* on the + # fast lane than under isolation (where clone-on-receive makes it a silent no-op). + self.v.in_place = True + self._expr(t.value, ns, loc, depth) + return + if isinstance(t, ast.Attribute): + # `x.output = ...` / `x.input = ...` is the nnsight boundary write (a SWAP), + # allowed. Any other attribute store mutates host state visible to sibling + # mediators — isolate. + if t.attr in ("output", "input", "inputs", "grad"): + self._expr(t.value, ns, loc, depth) + return + if t.attr.startswith("_"): + raise _Stop(REJECT, f"writing dunder/private attribute `{t.attr}` — rejected") + raise _Stop(ISOLATE, f"writing host attribute `.{t.attr}` mutates shared state — isolated") + # computed target (e.g. starred) — be conservative + raise _Stop(ISOLATE, f"unsupported assignment target `{type(t).__name__}` — isolated") + + # --- expressions ---------------------------------------------------------- + def _expr(self, e, ns, loc, depth): + if isinstance(e, ast.Call): + self._call(e, ns, loc, depth) + elif isinstance(e, ast.Attribute): + self._attribute(e, ns, loc, depth) + elif isinstance(e, ast.Name): + if e.id in _INTROSPECTION: + raise _Stop(REJECT, f"name `{e.id}` is an introspection escape — rejected") + elif isinstance(e, ast.Subscript): + self._subscript(e, ns, loc, depth) + elif isinstance(e, ast.Constant): + pass + elif isinstance(e, (ast.BinOp, ast.UnaryOp, ast.BoolOp, ast.Compare)): + for child in ast.iter_child_nodes(e): + if isinstance(child, ast.expr): + self._expr(child, ns, loc, depth) + elif isinstance(e, (ast.List, ast.Tuple, ast.Set)): + for elt in e.elts: + self._expr(elt, ns, loc, depth) + elif isinstance(e, ast.Dict): + for k in e.keys: + if k is not None: + self._expr(k, ns, loc, depth) + for val in e.values: + self._expr(val, ns, loc, depth) + elif isinstance(e, (ast.ListComp, ast.SetComp, ast.GeneratorExp, ast.DictComp)): + self._comprehension(e, ns, loc, depth) + elif isinstance(e, ast.Slice): + for child in (e.lower, e.upper, e.step): + if child is not None: + self._expr(child, ns, loc, depth) + elif isinstance(e, ast.IfExp): + self._expr(e.test, ns, loc, depth) + self._expr(e.body, ns, loc, depth) + self._expr(e.orelse, ns, loc, depth) + elif isinstance(e, ast.Starred): + self._expr(e.value, ns, loc, depth) + elif isinstance(e, (ast.JoinedStr,)): + for val in e.values: + if isinstance(val, ast.FormattedValue): + self._expr(val.value, ns, loc, depth) + elif isinstance(e, ast.Lambda): + inner = loc | _arg_names(e.args) + self._expr(e.body, ns, inner, depth) + elif isinstance(e, ast.Constant): + pass + else: + raise _Stop(ISOLATE, f"unsupported expression `{type(e).__name__}` — isolated") + + def _comprehension(self, e, ns, loc, depth): + inner = set(loc) + for gen in e.generators: + for t in _target_names(gen.target): + inner.add(t) + self._expr(gen.iter, ns, inner, depth) + for cond in gen.ifs: + self._expr(cond, ns, inner, depth) + if isinstance(e, ast.DictComp): + self._expr(e.key, ns, inner, depth) + self._expr(e.value, ns, inner, depth) + else: + self._expr(e.elt, ns, inner, depth) + + def _subscript(self, e, ns, loc, depth): + # a dunder/private string key is an introspection escape (`obj["__globals__"]`). + sl = e.slice + key = sl.value if isinstance(sl, ast.Constant) else None + if isinstance(key, str) and key.startswith("_"): + raise _Stop(REJECT, f"subscript key `{key}` is an introspection escape — rejected") + self._expr(e.value, ns, loc, depth) + if isinstance(sl, ast.expr): + self._expr(sl, ns, loc, depth) + + def _attribute(self, e, ns, loc, depth): + # reading an attribute is fine (incl. `.weight`, `.output`, `.shape`); reading a + # dunder/introspection attribute is an escape. + if e.attr in _INTROSPECTION or (e.attr.startswith("__") and e.attr.endswith("__")): + raise _Stop(REJECT, f"attribute `.{e.attr}` is an introspection escape — rejected") + if e.attr == "weight" or e.attr == "bias": + self.v.touches_host_weights = True + self._expr(e.value, ns, loc, depth) + + def _call(self, e, ns, loc, depth): + # walk args first (they may themselves disqualify) + for a in e.args: + self._expr(a, ns, loc, depth) + for kw in e.keywords: + self._expr(kw.value, ns, loc, depth) + + func = e.func + if isinstance(func, ast.Attribute): + self._call_attribute(func, ns, loc, depth) + elif isinstance(func, ast.Name): + self._call_name(func.id, ns, loc, depth) + else: + # a computed callable: `(a or b)(x)`, `factory()(x)` — unconfirmable target + self._expr(func, ns, loc, depth) + raise _Stop(ISOLATE, "call to a computed/dynamic target — isolated") + + def _call_attribute(self, func: ast.Attribute, ns, loc, depth): + attr = func.attr + if attr in _INTROSPECTION or (attr.startswith("__") and attr.endswith("__")): + raise _Stop(REJECT, f"method `.{attr}()` is an introspection escape — rejected") + # is this a banned qualified op like torch.load / torch.save / torch.hub.*? + qual = _attr_chain(func) + if qual is not None: + for banned in _BANNED_QUALIFIED: + if qual == banned or qual.startswith(banned + "."): + raise _Stop(ISOLATE, f"`{qual}` touches fs/net/JIT — isolated") + # if it resolves into a safe module, fine; if into a banned module, isolate + head = qual.split(".", 1)[0] + obj = _lookup(head, ns, loc) + if inspect.ismodule(obj) and _module_is_safe(obj): + self._expr(func.value, ns, loc, depth) + return + # otherwise a method call on a receiver (tensor/envoy/host object). Method-name + # is not an escape (checked above); the receiver is walked. Allowed under + # trust=local — `.clone()`, `.save()`, `.sum()`, `.backward()`, `.to(...)` etc. + self._expr(func.value, ns, loc, depth) + + def _call_name(self, name: str, ns, loc, depth): + if name in _INTROSPECTION: + raise _Stop(REJECT, f"`{name}()` is an introspection escape — rejected") + if name in _SAFE_BUILTINS: + return + if name in loc: + return # a local/param: a host object or fn passed in at a walked call site + obj = _lookup(name, ns, None) + if obj is _MISSING: + raise _Stop(ISOLATE, f"call to unresolved name `{name}` — unknown authority, isolated") + kind = _classify_obj(obj) + if kind == "op" or kind == "nnsight" or kind == "host" or kind == "builtin_ok": + if kind == "host": + self.v.touches_host_weights = self.v.touches_host_weights # host call is fine + return + if kind == "banned": + raise _Stop(REJECT, f"`{name}()` is an introspection escape — rejected") + if kind == "userfn": + self._recurse_fn(obj, depth) + return + # unknown object type bound to a resolvable name — be conservative + raise _Stop(ISOLATE, f"call to `{name}` ({type(obj).__name__}) is not confirmable — isolated") + + +# --------------------------------------------------------------------------- # +# object / name resolution helpers # +# --------------------------------------------------------------------------- # +_MISSING = object() + + +def _lookup(name, ns: dict, loc): + if loc is not None and name in loc: + return _MISSING # a local: caller treats specially + if name in ns: + return ns[name] + import builtins + if hasattr(builtins, name): + return getattr(builtins, name) + return _MISSING + + +def _module_is_safe(mod) -> bool: + modname = getattr(mod, "__name__", "") + return any(modname == p or modname.startswith(p + ".") for p in _SAFE_MODULE_PREFIXES) + + +def _classify_obj(obj) -> str: + import builtins + if obj is _MISSING: + return "unknown" + # introspection builtins resolved as objects + if getattr(obj, "__name__", None) in _INTROSPECTION and inspect.isbuiltin(obj): + return "banned" + mod = getattr(obj, "__module__", None) or "" + if mod == "builtins": + nm = getattr(obj, "__name__", None) + if nm in _INTROSPECTION: + return "banned" + if nm in _SAFE_BUILTINS: + return "builtin_ok" + return "unknown" # an unlisted builtin (open/input/…) → isolate, not reject + # pure-compute libraries + if any(mod == p or mod.startswith(p + ".") for p in _SAFE_MODULE_PREFIXES): + return "op" + # nnsight primitives (Envoy, eproperty, tracer methods, save, ...) + if mod == "nnsight" or mod.startswith("nnsight."): + return "nnsight" + # a host object: Envoy / nn.Module / Tensor / Backend — calling it runs the real model + if _is_host_object(obj): + return "host" + # any other user-defined callable → recurse into its source + if inspect.isfunction(obj) or inspect.ismethod(obj) or isinstance(obj, type(_lookup)): + return "userfn" + if inspect.isfunction(getattr(obj, "__call__", None)): + return "userfn" + return "unknown" + + +def _is_host_object(obj) -> bool: + try: + import torch.nn as nn + import torch + if isinstance(obj, (nn.Module, torch.Tensor)): + return True + except Exception: # noqa: BLE001 + pass + cls = type(obj) + chain = {c.__name__ for c in cls.__mro__} + return bool(chain & {"Envoy", "Backend"}) + + +def _attr_chain(node): + """`torch.nn.functional.linear` → that dotted string; None if not a pure attr chain.""" + parts = [] + cur = node + while isinstance(cur, ast.Attribute): + parts.append(cur.attr) + cur = cur.value + if isinstance(cur, ast.Name): + parts.append(cur.id) + return ".".join(reversed(parts)) + return None + + +def _arg_names(args: ast.arguments) -> set: + names = set() + for group in (args.posonlyargs, args.args, args.kwonlyargs): + for a in group: + names.add(a.arg) + if args.vararg: + names.add(args.vararg.arg) + if args.kwarg: + names.add(args.kwarg.arg) + return names + + +def _target_names(target) -> set: + names = set() + for n in ast.walk(target): + if isinstance(n, ast.Name): + names.add(n.id) + return names + + +def _assigned_names(n) -> set: + targets = n.targets if isinstance(n, ast.Assign) else [n.target] + names = set() + for t in targets: + for x in ast.walk(t): + if isinstance(x, ast.Name): + names.add(x.id) + return names diff --git a/src/nnsight/intervention/interleaver.py b/src/nnsight/intervention/interleaver.py index 6eccf8569..1006eb19b 100755 --- a/src/nnsight/intervention/interleaver.py +++ b/src/nnsight/intervention/interleaver.py @@ -992,6 +992,11 @@ def __init__( self._iso_backward = False self._iso_grad_reals: dict = {} + # Fast lane: the static safety verdict (set in start when fast_lane is on) and the + # wall-clock watchdog for a confirmed-safe in-process run. + self._fastlane_verdict = None + self._fastlane_watchdog = None + # True only inside an isolated worker process (set by _run_one_job). Lets # cross-process-aware logic (e.g. Barrier) know it can't count locally. self._isolated_worker = False @@ -1088,31 +1093,44 @@ def start(self, interleaver: Interleaver): from .isolation import isolation_state if isolation_state()["on"]: - # Isolated path: run the intervention in an isolated GPU worker process - # (a warm pooled worker when pool_size>0, else a cold one-shot worker). - # Sets self.channel (host end), self.worker (the process), self._iso. - from .isolation import acquire_isolated_worker - - acquire_isolated_worker(self) - self.interleaver.current = self - else: - _intervention = self.intervention - _args = (self, self.info, *self.args) - - def _worker_target(): - if _caller_stream is not None: - torch.cuda.set_stream(_caller_stream) - _intervention(*_args) - - # Start the worker thread. - self.worker = Thread( - target=_worker_target, - daemon=True, - name=self.name, + # Three execution tiers under isolation. The fast lane classifies each + # mediator once: a CONFIRMED-safe intervention runs IN-PROCESS (full model + + # weights, no worker, no per-hook channel) — the ONLY tier that can run the + # weight-reading interp majority, since the worker holds weightless dummy + # modules — while the unconfirmable remainder isolates and an introspection + # escape is rejected. The classifier is a footgun selector, never a malice + # boundary (see fastlane.py); it only moves a mediator isolate -> in-process. + from .fastlane import FAST, REJECT, FastLaneRejected, Watchdog + from .isolation import ( + acquire_isolated_worker, + classify_for_fast_lane, + fast_lane_enabled, ) - self.interleaver.current = self - self.worker.start() + tier = "isolate" + if fast_lane_enabled(): + verdict = classify_for_fast_lane(self) + self._fastlane_verdict = verdict + if verdict.tier == REJECT: + raise FastLaneRejected( + f"intervention cannot be confirmed for the in-process fast lane: " + f"{verdict.reason}. Rewrite without introspection, or run without " + f"isolate_mediators()." + ) + tier = verdict.tier + + if tier == FAST: + deadline = isolation_state()["opts"].fast_lane_timeout + self._fastlane_watchdog = Watchdog(deadline) + self._run_in_process(_caller_stream, watchdog=self._fastlane_watchdog) + else: + # Isolated path: run the intervention in an isolated GPU worker process + # (warm pooled worker when pool_size>0, else a cold one-shot worker). + # Sets self.channel (host end), self.worker (the process), self._iso. + acquire_isolated_worker(self) + self.interleaver.current = self + else: + self._run_in_process(_caller_stream) self.channel.wait_event() @@ -1124,6 +1142,31 @@ def _worker_target(): self.interleaver.current = None + def _run_in_process(self, caller_stream, watchdog=None): + """Launch the intervention in a daemon thread (the in-process path, shared by the + isolation-off default and the confirmed-safe fast lane). ``watchdog`` (fast lane + only) bounds a runaway pure-Python loop: it is armed on the thread after start and + disarmed when the body returns (the ``finally``) — and again at ``cancel`` as a + backstop. The body's own ``try/except`` (invoker.compile) routes an injected + :class:`fastlane.FastLaneTimeout` through ``mediator.exception`` like any error.""" + _intervention = self.intervention + _args = (self, self.info, *self.args) + + def _worker_target(): + try: + if caller_stream is not None: + torch.cuda.set_stream(caller_stream) + _intervention(*_args) + finally: + if watchdog is not None: + watchdog.disarm() + + self.worker = Thread(target=_worker_target, daemon=True, name=self.name) + self.interleaver.current = self + self.worker.start() + if watchdog is not None: + watchdog.arm(self.worker.ident) + ### Provider Methods ### def cancel(self): @@ -1136,6 +1179,11 @@ def cancel(self): # Retained on-graph activations (isolated backward) pin the autograd graph; # drop them at trace end rather than waiting for the mediator to be GC'd. self._iso_grad_reals = {} + # Disarm the fast-lane watchdog (backstop to the thread-target finally) so a + # generous deadline can't fire into an unrelated later computation. + if self._fastlane_watchdog is not None: + self._fastlane_watchdog.disarm() + self._fastlane_watchdog = None # If the worker is still mid-protocol, unwind it with a Cancelation. For the # isolated channel the get_event/put_response calls are host-local — they do NOT @@ -1824,6 +1872,8 @@ def __setstate__(self, state): self._iso_backward = False self._iso_grad_reals = {} self._iso_caches = {} + self._fastlane_verdict = None + self._fastlane_watchdog = None self.interleaver = None self.history = set() self.user_cache: "Cache" = list() diff --git a/src/nnsight/intervention/isolation.py b/src/nnsight/intervention/isolation.py index f1b933e28..ad7048f63 100644 --- a/src/nnsight/intervention/isolation.py +++ b/src/nnsight/intervention/isolation.py @@ -96,6 +96,12 @@ class IsoOptions: gpu_mem_fraction: float = 0.3 lockdown: bool = False # functional-first; seccomp lockdown enabled separately timeout: float = 60.0 # per-step wall-clock cap on user code (hang containment) + # Fast lane: confirmed-safe interventions run IN-PROCESS (full model access, no + # worker, no per-hook channel) instead of in the GPU worker. Without it, isolation + # cannot run the weight-reading interp majority at all (the worker is weightless). + fast_lane: bool = True + trust: str = "local" # only "local" provenance is fast-lane-eligible + fast_lane_timeout: float = 120.0 # whole-intervention watchdog bound for the fast lane @property def pool_key(self) -> tuple: @@ -133,6 +139,9 @@ def isolate_mediators( timeout: float = 60.0, lockdown: bool = False, pool_size: int = 0, + fast_lane: bool = True, + trust: str = "local", + fast_lane_timeout: float = 120.0, ): """Run interventions inside ``with model.trace(...)`` in an isolated GPU worker. @@ -151,6 +160,19 @@ def isolate_mediators( Under ``lockdown=True`` a pooled worker locks its import set at warm time, so a job whose user code triggers a NEW import fails (consistently across the pool) — stricter than the cold path, which deserializes before lockdown. + fast_lane: if True (default), a per-mediator static classifier + (:mod:`nnsight.intervention.fastlane`) confirms interventions that use only + whitelisted ops / host-model access / nnsight primitives and runs THOSE + in-process (full model + weights, no worker, no per-hook channel), isolating + only the unconfirmable remainder. This is what lets isolation run the + weight-reading interp majority at all — the worker holds weightless dummy + modules. The override only ever moves a mediator from isolate to in-process; + set ``fast_lane=False`` to force pure isolation. + trust: only ``"local"`` provenance is fast-lane-eligible; the static gate is a + footgun selector, not a malice boundary, so any other value disables the fast + lane wholesale (everything isolates). + fast_lane_timeout: whole-intervention wall-clock bound for a fast-laned thread + (a best-effort watchdog restoring loop-containment in-process). """ prev = dict(_STATE) _STATE.update( @@ -162,6 +184,9 @@ def isolate_mediators( gpu_mem_fraction=gpu_mem_fraction, lockdown=lockdown, timeout=timeout, + fast_lane=fast_lane, + trust=trust, + fast_lane_timeout=fast_lane_timeout, ), ) try: @@ -170,6 +195,42 @@ def isolate_mediators( _STATE.update(prev) +def fast_lane_enabled() -> bool: + """True if confirmed-safe mediators should run in-process. Gated on the context + option, the ``trust="local"`` provenance cordon, and the global config flag (a server + can force pure isolation without code changes).""" + opts = _STATE["opts"] + if not opts.fast_lane or opts.trust != "local": + return False + try: + from .. import CONFIG + + return bool(getattr(CONFIG.APP, "FAST_LANE", True)) + except Exception: # noqa: BLE001 — config unavailable => default-on + return True + + +def classify_for_fast_lane(mediator): + """Run the static classifier on ``mediator`` (host-side, before any serialization). + Returns a :class:`fastlane.Verdict`. Cached on the intervention code object so a + re-run of the same trace pays the walk once.""" + from . import fastlane + + code = getattr(mediator.intervention, "__code__", None) + cache = _FASTLANE_VERDICT_CACHE + if code is not None and code in cache: + return cache[code] + verdict = fastlane.classify(mediator) + if code is not None: + cache[code] = verdict + return verdict + + +# Verdict cache keyed by the intervention code object identity (per-trace-shape, stable +# across re-runs of the same trace); bounded implicitly by the number of distinct traces. +_FASTLANE_VERDICT_CACHE: Dict[Any, Any] = {} + + # --------------------------------------------------------------------------- # # Host side — worker handle + pool # # --------------------------------------------------------------------------- # @@ -432,13 +493,24 @@ def _build_job(mediator) -> tuple: "cross_invoker": bool(mediator.cross_invoker), # `with tensor.backward()` detection — the single decision point for BOTH # sides: the host gates real-activation retention, the worker gates - # delivered-clone tagging. (The substring can false-positive, e.g. in a - # comment, which only costs needless tagging; tighten it here when needed.) - "backward_active": ".backward(" in mediator.intervention.__source__, + # delivered-clone tagging. Prefer the fast-lane classifier's closure-aware flag + # (it resolves through build()/capture() closures the substring is blind to); + # fall back to the source substring when the classifier did not run. + "backward_active": _backward_active(mediator), } return payload, extras, worker_opts +def _backward_active(mediator) -> bool: + """Closure-aware `with tensor.backward()` detection for the isolated job, preferring + the fast-lane classifier's verdict (which walks through user closures) over the + source substring (blind to a backward hidden in a build()/capture() closure).""" + verdict = getattr(mediator, "_fastlane_verdict", None) + if verdict is not None: + return verdict.differentiate + return ".backward(" in mediator.intervention.__source__ + + def _wire_host_channel(mediator, iso: _PooledWorker, worker_opts: dict) -> None: """Point the (possibly recycled) worker's host channel at THIS mediator.""" chan = iso.channel diff --git a/src/nnsight/schema/config.py b/src/nnsight/schema/config.py index c7f6a00f1..ad16f5303 100755 --- a/src/nnsight/schema/config.py +++ b/src/nnsight/schema/config.py @@ -29,6 +29,7 @@ class AppConfigModel(BaseModel): CACHE_DIR: str = "~/.cache/nnsight/" CROSS_INVOKER: bool = True TRACE_CACHING: bool = False + FAST_LANE: bool = True def __setattr__(self, name, value): if name == "TRACE_CACHING" and value is True: From bd598ce4fa16c1ffb28c53680a302a008a7bae2b Mon Sep 17 00:00:00 2001 From: khaiwang Date: Sat, 13 Jun 2026 00:50:07 -0400 Subject: [PATCH 16/30] docs(intervention): fast-lane design + integration-doc support-matrix note docs/developing/fast-lane.md: why isolation could not run the weight- reading interp majority (weightless worker), the three-tier design, the classifier rules + threat-model contract, the watchdog, the prior art it borrows from (Cloudflare Workers / RestrictedPython / SES / fx+JAX / gVisor / Firecracker / the pysandbox negative result), the designed part-2 declarative primitives (next increment), deferred items, and the verification matrix. Cross-linked from the integration doc's support matrix, which now notes the FAST/ISOLATE/REJECT tiering. Co-Authored-By: Claude Fable 5 --- docs/developing/fast-lane.md | 158 ++++++++++++++++++ .../mediator-gpu-trace-integration.md | 9 + 2 files changed, 167 insertions(+) create mode 100644 docs/developing/fast-lane.md diff --git a/docs/developing/fast-lane.md b/docs/developing/fast-lane.md new file mode 100644 index 000000000..af7e547b3 --- /dev/null +++ b/docs/developing/fast-lane.md @@ -0,0 +1,158 @@ +# The in-process fast lane — design + +**Status:** Implemented (first slice) · **Date:** 2026-06-13 · **Branch:** `worktree-mediator-sandbox` +**Builds on:** [mediator-gpu-trace-integration.md](mediator-gpu-trace-integration.md) (the GPU-worker isolation harness this adds a tier to). + +## 1. Why — isolation can't run the interp majority + +The GPU-worker isolation contains footguns by running each user intervention in a spawned process +that holds a **path-only, weightless mirror** of the model: `isolation._WorkerPersistent` synthesizes a +bare `nn.Module()` (only a `__path__`, no parameters, no `forward`) for every `Module:`. That is +fine for reading/swapping/saving *delivered activations*, but the actual interpretability workloads — +checked against the `interp-serve-bench` taxonomy — do more than that **inside the trace body**: + +```python +normed = model.transformer.ln_f(hidden) # call the host's real final-norm module +logits = F.linear(normed, model.lm_head.weight) # read the host's real unembed WEIGHTS +``` + +Every logit-lens, steering, ablation, activation-patching, and attribution cell reads host weights +and/or calls host modules. In the worker, `head.weight` → `AttributeError` and `head(x)` → +`NotImplementedError` on the weightless dummy. So **as it stood, isolation could not run any of the +weight-reading interp majority** — the workflow's compatibility pass found this the single largest gap +(9 of the cataloged workloads, all blocked on the same weight/host-module surface). + +The fix is not to ship weights into the worker (a per-trace cost, and a partial answer). It is to run +the confirmed-safe interventions **in-process**, where the real model and its weights already live, and +isolate only the code that genuinely needs containing. That is the fast lane. + +## 2. The three tiers + +`isolate_mediators()` now classifies each mediator once, at fork in `Mediator.start`: + +| tier | when | how it runs | +|---|---|---| +| **FAST** | the effective code is all whitelisted ops / host-model access / nnsight primitives | in-process daemon thread (`_iso=None`) — full model + weights, no worker, no per-hook channel; a watchdog bounds runaway loops | +| **ISOLATE** | anything unconfirmable (unresolved global call, import, `while`, unrecoverable closure, unknown node) | the existing GPU worker | +| **REJECT** | an introspection escape (`__globals__` / `getattr` / `eval` / …) | raise `FastLaneRejected` | + +The conservative default is **ISOLATE**: absence of proof is not proof of safety; only explicitly +whitelisted code reaches FAST. Default behavior is preserved — isolation **off** never consults the +gate; isolation **on** only ever moves a mediator *isolate → in-process*. + +## 3. Detect & confirm — the classifier (`fastlane.py`) + +A **fail-closed, default-deny** static walk over the **effective code**: the trace body PLUS every user +closure it calls, resolved through the capturing frame, the function's `__globals__`, and its closure +cells. This closure resolution is load-bearing — the harness wraps real compute in `build()` / +`capture()` / `patch()` closures, so a walk of the `with` block alone sees only `build().save()`, an +opaque call. Confirming on the with-block while the footgun lives in an unresolved closure would be the +exact false-safe the gate must never produce; an unrecoverable closure fails closed to ISOLATE. + +The walk (the rules, in `_Walker`): + +- **Node allowlist.** Permit assignment / expr / return / `if` / `for` / `with` / `raise` / `pass` and + pure expression nodes. `while` → ISOLATE (cannot be statically bounded). `import` → ISOLATE (ambient + authority). `try` → ISOLATE. `global`/`nonlocal` → REJECT. Unknown node type → ISOLATE. +- **Call targets.** Each call must resolve to: a pure-compute op (`torch` / `torch.nn.functional` / + `math` / `operator` / `numpy`), an nnsight primitive, a host `Envoy`/`nn.Module`/`Tensor` (calling it + runs the real model — the fast lane's whole point), a recursively confirmed user function, a safe + builtin, or a local/parameter name (a host object bound at a walked call site). An unresolved global + call → ISOLATE (unknown authority); `torch.load`/`save`/`hub`/`jit` → ISOLATE (fs/net/JIT). +- **Introspection ban (REJECT).** `getattr`/`setattr`/`eval`/`exec`/`compile`/`__import__`/`globals`/…, + any dunder attribute (`x.__class__`), and a dunder/private subscript key (`d["__builtins__"]`). +- **Host-state writes (ISOLATE).** An attribute store other than the nnsight boundary writes + (`.output`/`.input`/`.grad`) mutates state visible to sibling mediators. In-place + (`hidden[:] = …`) and replacement boundary writes are *allowed* (and in-place is **safer** here than + under isolation, where clone-on-receive silently no-ops it). +- **Backward detection** is closure-aware (`with x.backward():` found through closures), replacing the + old `.backward(` source substring that was blind to a backward hidden in a closure — this flag now + feeds the *isolated* job's gradient-retention too. + +Verdicts are cached by intervention code-object identity, so re-running the same trace pays the walk +once. The classifier is GPU-free and unit-tested (`test_fastlane_classifier.py`) on the real workload +shapes and on **renamed** module structures (`decoder_blocks`/`final_norm`/`output_projection`) so +nothing is keyed to GPT-2 naming. + +### Threat model (the contract) + +This is **not** an adversarial sandbox — a determined author can defeat any in-process restriction +(the pysandbox negative result). Under the harness's relaxed "contain footguns, not adversaries" model +the gate confirms the effective code: introduces no ambient authority (no import, no introspection, no +unresolved global call), has no unbounded loop, writes no host state, and is composed only of +whitelisted ops / host access / confirmed user code. It is cordoned to `trust="local"` provenance and a +`CONFIG.APP.FAST_LANE` flag; anything deserialized/remote, or with the flag off, isolates wholesale. +OOM and device-side asserts in pure tensor math are knowingly traded to the in-process tier (the same +risk a user running without isolation already accepts); a deployment that cannot tolerate them sets +`fast_lane=False`. + +## 4. The watchdog + +The static gate bans `while` and confirms loops are over bounded iterables, but a *huge* bounded +`range(10**12)` passes statically yet would hang the host — the loop-containment that turning isolation +on implies. A best-effort wall-clock `Watchdog` injects a `FastLaneTimeout` into the fast-lane thread at +its next bytecode if it overruns `fast_lane_timeout`. Because the intervention body is already wrapped +in `try/except → mediator.exception` (invoker.compile), the injected exception routes through the normal +path and the host re-raises it cleanly — no channel hang. It cannot preempt a wedged native/CUDA call +(only the worker-process kill can); the bounded-loop static rule is the primary defense, this is the +backstop. Armed after thread start, disarmed in the thread's `finally` and again at `cancel`. + +## 5. Prior art it is built on + +| system | tier | borrowed idea | +|---|---|---| +| Cloudflare Workers (V8 isolates) | fast-in-process | many tenants in one process; safety from a capability-restricted runtime, not a process boundary — the fast lane's model | +| RestrictedPython (Zope) | static-confirm | a `RestrictingNodeTransformer`-style default-deny AST pass as the gate | +| SES / Hardened JS (Endo) | fast-in-process | confirmed code runs against an explicitly granted namespace; no ambient authority | +| torch.fx / JAX tracing | static-confirm | a leaf/atomic-op allowlist; "prove the body is a pure function over provided tensors" — fail, don't guess | +| gVisor | hybrid | cost is the boundary *crossing*, not the work — minimize crossings (the fast lane removes them entirely) | +| AWS Lambda + Firecracker / SnapStart | heavy-isolated | what the existing GPU worker *is* — the slow lane you keep for unconfirmable code; warm-restore to amortize spawn (the warm pool) | +| pysandbox (negative result) | — | the contract: in-process restriction is a footgun selector, never an adversarial boundary → the `trust="local"` cordon | + +The unifying pattern across all of them: **confirm once up front, then run free** behind a small, +explicit, enumerable fallback to the heavy tier. + +## 6. Part 2 — declarative primitives (designed; next increment) + +The fast lane already runs the weight-reading workloads in-process, so they need no new APIs *for the +fast lane*. The declarative primitives' value is (a) making the workloads runnable on the **isolated** +tier too, and (b) collapsing common raw-compute patterns into named calls the gate whitelists trivially: + +| primitive | signature | taxonomy primitive | role | +|---|---|---|---| +| `tracer.unembed` | `(residual, norm, head) → logits` | host-weight read + module call | the projection every logit-lens / steering-direction / attribution metric does | +| `tracer.steer` | `(envoy, direction, alpha)` | boundary write (injection) | always a replacement swap — fixes in-place's silent no-op under isolation | +| `tracer.patch` | `(envoy, value)` | boundary write (transplant) | whole-tuple replacement | +| `tracer.ablate` | `(envoy, mode)` | boundary write (injection) | zero/mean knockout | +| `tracer.capture` | `(value) → handle` | read + run↔run transfer | cross-trace handoff; non-transmittable → clean fail, not silent drop | + +Each mirrors the existing `tracer.cache()` shape: in-process it resolves the real envoys and runs +directly (the fast-lane execution); isolated it ships a spec via a new event whose host handler runs the +real op host-side — which **also** closes the standing weight-read blocker for the heavy lane. Building +those host event handlers is the increment; the protocol slot pattern already exists (`Events.CACHE` / +`handle_cache_event`). + +## 7. What was deliberately deferred + +- **The process-global `sys.addaudithook` backstop.** Its own failure mode (a leaked thread-local flag + arms it during the model's *own* forward → server-wide outage) makes it net-negative when the static + default-deny gate already makes imports / `open` / `exec` / `socket` statically impossible in + fast-laned code. Documented as future hardening; the static gate is the confirmation. +- **A frozen-namespace `Compartment`** (SES-style) for fast-lane execution. The first slice relies on + the static pass + the `trust` cordon; namespace shadowing is a later refinement. +- **The five primitives' isolated event handlers** (§6). + +## 8. Verification + +- **Classifier units** (`test_fastlane_classifier.py`, GPU-free) — 17/17: logit-lens / steering / + patching / attribution shapes and **renamed** structures classify FAST; imports / `while` / + unresolved-call / `open` → ISOLATE; introspection → REJECT; the `differentiate`/`in_place`/ + `touches_host_weights` flags are set correctly. +- **Fast-lane end-to-end** (`test_fast_lane.py`, gpt2 + a renamed model) — 6/6: the weight-reading + logit lens is bit-identical on the fast lane (`max|Δ|=0`) **and raises under forced isolation** + (`fast_lane=False`) — proving the fast lane is the enabling tier; in-place steering bit-identical; + renamed-model lens bit-identical; a footgun routes off the fast lane and the host survives; an + introspection escape is rejected; a runaway loop is killed by the watchdog and the host survives. +- **Existing isolated WORKER path** — 9/9 still bit-identical, pinned with `fast_lane=False` so they + keep exercising the worker (otherwise the simple read/swap/save cells would now fast-lane). +- **In-process core** — 51 passed (the default in-process path is untouched). diff --git a/docs/developing/mediator-gpu-trace-integration.md b/docs/developing/mediator-gpu-trace-integration.md index 065922dff..0b579a39e 100644 --- a/docs/developing/mediator-gpu-trace-integration.md +++ b/docs/developing/mediator-gpu-trace-integration.md @@ -235,6 +235,15 @@ benign CudaIPC release warning. ## 8. Support matrix (what works under `isolate_mediators()` today, and how) +> **Fast lane (2026-06-13, [fast-lane.md](fast-lane.md)).** `isolate_mediators()` now runs each +> mediator on one of three tiers: a static classifier confirms safe interventions and runs them +> **in-process** (FAST — full model + weights), isolates the unconfirmable remainder in the GPU worker +> (ISOLATE), and rejects introspection escapes (REJECT). This is what lets the **weight-reading interp +> majority** (logit lens, steering, ablation, activation patching, attribution) run under isolation at +> all — the worker holds weightless dummy modules, so those workloads can only run on the fast lane. +> The matrix below describes the **ISOLATE (worker) tier**; a row marked weight-reading/host-module is +> served by the FAST tier. Default-on; `fast_lane=False` forces the worker tier. + | Feature | Cross-process mechanism | Status | |---|---|---| | read / swap (`=`) / `.save()` (tensors) / skip / exception | six events over the channel; host-side hook registration; worker→host saves transmission at END | ✅ bit-identical | From 948f4b725de107fed5430b3dcb9ac15bbe0b5a65 Mon Sep 17 00:00:00 2001 From: khaiwang Date: Sun, 14 Jun 2026 01:49:07 -0400 Subject: [PATCH 17/30] feat(intervention): host-routed tracer.unembed for the isolated tier MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The weight-reading interp readout — F.linear(norm(residual), head.weight), done by every logit-lens / steering-direction / attribution-metric cell — cannot run in the isolated worker because its dummy modules are weightless. tracer.unembed closes that on the isolated tier without putting weights in the worker: the worker ships the residual VALUE plus the norm/head module PATHS via a new Events.UNEMBED request; the host's handle_unembed_event resolves the real envoys, runs the real norm + unembed on the real weights, and ships back only the logits (bounce-buffer round trip, clone-on-receive). Weights never cross the boundary — so this neither binds the generic warm worker to a model nor places host weight memory in the less-trusted worker (the two costs that ruled out shipping/sharing weights). In-process / on the fast lane it just runs the real modules directly. Shaped exactly like Events.CACHE / handle_cache_event. This is the first of the part-2 declarative primitives; it also means a deployment that forces pure isolation (fast_lane=False, e.g. for OOM containment) can still run weight-reading workloads if they are written with tracer.unembed. Verified (test_isolated_unembed.py, all under forced isolation, gpt2 + renamed model): single-layer / 3-layer-interleaved / formulation="module" / norm=None / renamed-model readouts isolated-vs-in-process max|Δ|=0; tracer.unembed == the manual F.linear it replaces. Isolated trace/cache, fast-lane e2e, in-process core (31), and classifier units (17) unchanged. Co-Authored-By: Claude Fable 5 --- docs/developing/fast-lane.md | 49 ++++-- .../gpu_sandbox/test_isolated_unembed.py | 164 ++++++++++++++++++ src/nnsight/intervention/interleaver.py | 43 +++++ src/nnsight/intervention/tracing/tracer.py | 49 ++++++ 4 files changed, 293 insertions(+), 12 deletions(-) create mode 100644 prototypes/mediator-sandbox/gpu_sandbox/test_isolated_unembed.py diff --git a/docs/developing/fast-lane.md b/docs/developing/fast-lane.md index af7e547b3..fe06d0c51 100644 --- a/docs/developing/fast-lane.md +++ b/docs/developing/fast-lane.md @@ -112,25 +112,49 @@ backstop. Armed after thread start, disarmed in the thread's `finally` and again The unifying pattern across all of them: **confirm once up front, then run free** behind a small, explicit, enumerable fallback to the heavy tier. -## 6. Part 2 — declarative primitives (designed; next increment) +## 6. Part 2 — declarative primitives The fast lane already runs the weight-reading workloads in-process, so they need no new APIs *for the fast lane*. The declarative primitives' value is (a) making the workloads runnable on the **isolated** tier too, and (b) collapsing common raw-compute patterns into named calls the gate whitelists trivially: -| primitive | signature | taxonomy primitive | role | -|---|---|---|---| -| `tracer.unembed` | `(residual, norm, head) → logits` | host-weight read + module call | the projection every logit-lens / steering-direction / attribution metric does | -| `tracer.steer` | `(envoy, direction, alpha)` | boundary write (injection) | always a replacement swap — fixes in-place's silent no-op under isolation | -| `tracer.patch` | `(envoy, value)` | boundary write (transplant) | whole-tuple replacement | -| `tracer.ablate` | `(envoy, mode)` | boundary write (injection) | zero/mean knockout | -| `tracer.capture` | `(value) → handle` | read + run↔run transfer | cross-trace handoff; non-transmittable → clean fail, not silent drop | +| primitive | signature | taxonomy primitive | role | status | +|---|---|---|---|---| +| `tracer.unembed` | `(residual, norm, head, formulation="weight") → logits` | host-weight read + module call | the projection every logit-lens / steering-direction / attribution metric does | **built** | +| `tracer.steer` | `(envoy, direction, alpha)` | boundary write (injection) | always a replacement swap — fixes in-place's silent no-op under isolation | designed | +| `tracer.patch` | `(envoy, value)` | boundary write (transplant) | whole-tuple replacement | designed | +| `tracer.ablate` | `(envoy, mode)` | boundary write (injection) | zero/mean knockout | designed | +| `tracer.capture` | `(value) → handle` | read + run↔run transfer | cross-trace handoff; non-transmittable → clean fail, not silent drop | designed | Each mirrors the existing `tracer.cache()` shape: in-process it resolves the real envoys and runs directly (the fast-lane execution); isolated it ships a spec via a new event whose host handler runs the -real op host-side — which **also** closes the standing weight-read blocker for the heavy lane. Building -those host event handlers is the increment; the protocol slot pattern already exists (`Events.CACHE` / -`handle_cache_event`). +real op host-side. + +### `tracer.unembed` — host-routed readout (built, 2026-06-14) + +The first primitive, and the one that closes the standing weight-read blocker on the **isolated** tier. +`tracer.unembed(residual, norm, head)` projects a residual through the final norm + unembed: + +- **In-process / fast lane** (`mediator._isolated_worker` is False): runs the real modules directly — + `F.linear(norm(residual), head.weight)` (or `head(norm(residual))` with `formulation="module"`). + This is the same compute the workloads write by hand; `tracer.unembed` just names it. +- **Isolated worker** (weightless dummies): ships the residual VALUE plus the module **paths** + (`{norm_path, head_path, formulation}`) via a new `Events.UNEMBED` request. The host's + `handle_unembed_event` resolves the real envoys (`path_to_envoy`), runs the real norm + unembed on the + **real weights**, and ships back only the logits over the bounce buffer (clone-on-receive both ways, + like any VALUE/BACKWARD round-trip). **Weights never cross the boundary** — so the readout works on + the isolated tier without binding the generic warm worker to a model or placing host weight memory in + the less-trusted worker (the two costs that ruled out shipping/sharing weights — see §1, §7). Paths + resolve through renames host-side (the wire path is always the real path), so renamed models work. + +Touch points: `Events.UNEMBED`; `handle()` dispatch + `handle_unembed_event` (interleaver.py); the +`tracer.unembed` method with the isolated/in-process branch (tracer.py). Shaped exactly like +`Events.CACHE` / `handle_cache_event`. + +**Verified (`test_isolated_unembed.py`, all under forced isolation `fast_lane=False`, gpt2 + renamed +model):** single-layer, 3-layer-interleaved, `formulation="module"`, `norm=None`, and renamed-model +readouts all isolated-vs-in-process `max|Δ|=0`; `tracer.unembed` equals the manual +`F.linear(norm(x), head.weight)` it replaces. ## 7. What was deliberately deferred @@ -140,7 +164,8 @@ those host event handlers is the increment; the protocol slot pattern already ex fast-laned code. Documented as future hardening; the static gate is the confirmation. - **A frozen-namespace `Compartment`** (SES-style) for fast-lane execution. The first slice relies on the static pass + the `trust` cordon; namespace shadowing is a later refinement. -- **The five primitives' isolated event handlers** (§6). +- **The remaining primitives' isolated event handlers** (§6) — `steer`/`patch`/`ablate`/`capture`; + `unembed` is built. ## 8. Verification diff --git a/prototypes/mediator-sandbox/gpu_sandbox/test_isolated_unembed.py b/prototypes/mediator-sandbox/gpu_sandbox/test_isolated_unembed.py new file mode 100644 index 000000000..eea68e963 --- /dev/null +++ b/prototypes/mediator-sandbox/gpu_sandbox/test_isolated_unembed.py @@ -0,0 +1,164 @@ +#!/usr/bin/env python3 +"""Host-routed unembed on the ISOLATED tier: tracer.unembed runs the real final-norm + +unembed on the host's real weights even when the intervention runs in the weightless +worker — so the weight-reading interp readout works under forced isolation, not only on +the fast lane. + +All cases force isolation (fast_lane=False) so the cell genuinely runs in the worker and +exercises the worker -> Events.UNEMBED -> host handler -> logits-back round trip. + + single — one layer's residual projected via tracer.unembed, isolated == + in-process, bit-identical. + multi — three layers (forward order) interleaved with unembed, bit-identical. + module_form — formulation="module" (head(normed)) bit-identical. + no_norm — norm=None (skip normalization) bit-identical. + renamed — renamed model (final_norm / output_projection / decoder_blocks): + paths resolve host-side, bit-identical (no hardcoded names). + matches_raw — tracer.unembed equals the manual F.linear(norm(x), head.weight). + +Run: + CUDA_VISIBLE_DEVICES=5 PYTHONPATH=src \ + /disk/u/zikai/anaconda3/envs/hf-serve/bin/python -u \ + prototypes/mediator-sandbox/gpu_sandbox/test_isolated_unembed.py +""" +import sys + +import torch +import torch.nn.functional as F + +from nnsight import LanguageModel +from nnsight.intervention.isolation import isolate_mediators + +PROMPT = "The Eiffel Tower is in the city of" + + +def _both(build): + """Run build() in-process and under forced isolation; return (ref, got).""" + ref = build() + with isolate_mediators(fast_lane=False, timeout=30): + got = build() + return ref, got + + +def _eq(ref, got): + return (torch.is_tensor(ref) and torch.is_tensor(got) + and ref.shape == got.shape and torch.equal(ref, got)) + + +def _delta(ref, got): + return (ref - got).abs().max().item() if _eq(ref, got) else float("nan") + + +def test_single(model): + def build(): + with model.trace(PROMPT) as tracer: + with torch.no_grad(): + resid = model.transformer.h[6].output + logits = tracer.unembed(resid, model.transformer.ln_f, model.lm_head) + row = logits[:, -1, :].save() + return row + ref, got = _both(build) + ok = _eq(ref, got) + print(f"[single] isolated unembed bit-identical={ok} (max|Δ|={_delta(ref, got)})", flush=True) + return ok + + +def test_multi(model): + def build(): + with model.trace(PROMPT) as tracer: + with torch.no_grad(): + rows = [] + for i in [0, 4, 8]: + resid = model.transformer.h[i].output + rows.append(tracer.unembed(resid, model.transformer.ln_f, + model.lm_head)[:, -1, :]) + g = torch.stack(rows, dim=0).save() + return g + ref, got = _both(build) + ok = _eq(ref, got) + print(f"[multi] 3-layer isolated unembed bit-identical={ok} (max|Δ|={_delta(ref, got)})", flush=True) + return ok + + +def test_module_form(model): + def build(): + with model.trace(PROMPT) as tracer: + with torch.no_grad(): + resid = model.transformer.h[6].output + logits = tracer.unembed(resid, model.transformer.ln_f, model.lm_head, + formulation="module") + row = logits[:, -1, :].save() + return row + ref, got = _both(build) + ok = _eq(ref, got) + print(f"[module_form] formulation='module' bit-identical={ok} (max|Δ|={_delta(ref, got)})", flush=True) + return ok + + +def test_no_norm(model): + def build(): + with model.trace(PROMPT) as tracer: + with torch.no_grad(): + resid = model.transformer.h[6].output + logits = tracer.unembed(resid, None, model.lm_head) # skip norm + row = logits[:, -1, :].save() + return row + ref, got = _both(build) + ok = _eq(ref, got) + print(f"[no_norm] norm=None bit-identical={ok} (max|Δ|={_delta(ref, got)})", flush=True) + return ok + + +def test_renamed(): + rename = {"transformer.ln_f": "final_norm", "lm_head": "output_projection", + "transformer.h": "decoder_blocks"} + model = LanguageModel("gpt2", device_map="cuda", dispatch=True, rename=rename) + + def build(): + with model.trace(PROMPT) as tracer: + with torch.no_grad(): + resid = model.decoder_blocks[5].output + logits = tracer.unembed(resid, model.final_norm, model.output_projection) + row = logits[:, -1, :].save() + return row + ref, got = _both(build) + ok = _eq(ref, got) + print(f"[renamed] renamed-model isolated unembed bit-identical={ok} (max|Δ|={_delta(ref, got)})", flush=True) + return ok + + +def test_matches_raw(model): + # tracer.unembed (in-process) must equal the manual readout it replaces. + with model.trace(PROMPT) as tracer: + with torch.no_grad(): + resid = model.transformer.h[6].output + via_api = tracer.unembed(resid, model.transformer.ln_f, model.lm_head)[:, -1, :].save() + with model.trace(PROMPT): + with torch.no_grad(): + out = model.transformer.h[6].output + hidden = out[0] if isinstance(out, tuple) else out + manual = F.linear(model.transformer.ln_f(hidden), model.lm_head.weight)[:, -1, :].save() + ok = _eq(via_api, manual) + print(f"[matches_raw] tracer.unembed == manual F.linear bit-identical={ok} " + f"(max|Δ|={_delta(via_api, manual)})", flush=True) + return ok + + +def main(): + assert torch.cuda.is_available() + model = LanguageModel("gpt2", device_map="cuda", dispatch=True) + results = { + "single": test_single(model), + "multi": test_multi(model), + "module_form": test_module_form(model), + "no_norm": test_no_norm(model), + "renamed": test_renamed(), + "matches_raw": test_matches_raw(model), + } + print("=" * 72, flush=True) + print(f"ISOLATED UNEMBED: {results}", flush=True) + sys.exit(0 if all(results.values()) else 1) + + +if __name__ == "__main__": + main() diff --git a/src/nnsight/intervention/interleaver.py b/src/nnsight/intervention/interleaver.py index 1006eb19b..8e6f9a67d 100755 --- a/src/nnsight/intervention/interleaver.py +++ b/src/nnsight/intervention/interleaver.py @@ -357,6 +357,7 @@ class Events(Enum): BARRIER = "barrier" # Signal that a barrier should be set CACHE = "cache" # Register a tracer.cache() on the host's real modules (isolation) BACKWARD = "backward" # Run the backward pass on the host's real graph (isolation) + UNEMBED = "unembed" # Run norm+unembed on the host's real weights (isolation) class Cancelation(Exception): @@ -1278,6 +1279,8 @@ def handle(self, provider: Optional[str] = None, value: Optional[Any] = None): process = self.handle_cache_event(data) elif event == Events.BACKWARD: process = self.handle_backward_event(data) + elif event == Events.UNEMBED: + process = self.handle_unembed_event(data) elif event == Events.END: process = self.handle_end_event(data) @@ -1579,6 +1582,46 @@ def handle_backward_event(self, seed: dict): self.respond(result) # ack -> worker's send() returns the grad dict return True + def handle_unembed_event(self, data): + """Run a residual through the host's REAL final-norm + unembed (isolation). + + The worker holds weightless dummy modules, so a ``F.linear(normed, head.weight)`` + readout — what every logit-lens / steering-direction / attribution-metric cell + does — cannot run there. The worker instead ships the residual VALUE plus the + module PATHS via ``Events.UNEMBED``; the host resolves the real envoys, runs + ``head(norm(residual))`` on the real weights, and ships back only the logits. + Weights never cross the boundary (no containment regression, no model-binding of + the generic worker); only the result does. ``spec`` = + ``{"norm_path", "head_path", "formulation"}`` (``formulation`` ∈ weight|module). + """ + import torch.nn.functional as F + + from .isolation import path_to_envoy + + residual, spec = data + p2e = path_to_envoy(self) + head = p2e.get(spec["head_path"]) + if head is None: + self.respond( + KeyError(f"unembed: head path {spec['head_path']!r} not found on the model") + ) + return True + norm = p2e.get(spec["norm_path"]) if spec.get("norm_path") else None + if spec.get("norm_path") and norm is None: + self.respond( + KeyError(f"unembed: norm path {spec['norm_path']!r} not found on the model") + ) + return True + + normed = norm._module(residual) if norm is not None else residual + if spec["formulation"] == "weight": + logits = F.linear(normed, head._module.weight) + else: + logits = head._module(normed) + + self.respond(logits) # ack -> worker's send() returns the logits + return True + def handle_end_event(self, saves: Optional[Any] = None): """ Handle an end event by stopping the mediator. diff --git a/src/nnsight/intervention/tracing/tracer.py b/src/nnsight/intervention/tracing/tracer.py index be5ba8466..d5bcf53d3 100755 --- a/src/nnsight/intervention/tracing/tracer.py +++ b/src/nnsight/intervention/tracing/tracer.py @@ -652,6 +652,55 @@ def cache( return cache_obj.cache + def unembed(self, residual, norm, head, formulation: str = "weight"): + """Project ``residual`` through the final norm + unembed → logits, on the host's + REAL weights even when the intervention is isolated. + + The readout every logit-lens / steering-direction / attribution-metric cell does + — ``F.linear(norm(residual), head.weight)`` — reads the host model's real + weights. In-process (incl. the fast lane) that just runs the real modules. In an + isolated worker the modules are weightless dummies, so this ships the residual + VALUE plus the module PATHS to the host (``Events.UNEMBED``); the host runs the + real norm + unembed and ships back only the logits. Weights never cross the + boundary — so this works on the isolated tier without binding the generic worker + to a model or placing host weight memory in the (less-trusted) worker. + + Args: + residual: the residual-stream tensor to project (a block ``.output``; a tuple + is untupled to ``[0]``). + norm: the final-norm Envoy (e.g. ``model.transformer.ln_f``), or ``None`` to + skip normalization. + head: the unembed Envoy (e.g. ``model.lm_head``). + formulation: ``"weight"`` → ``F.linear(normed, head.weight)`` (portable; the + form the workloads use); ``"module"`` → ``head(normed)``. + + Returns: + The logits tensor. + """ + import torch.nn.functional as F + + if isinstance(residual, tuple): + residual = residual[0] + + mediator = self.model.interleaver.current + + if mediator._isolated_worker: + # the worker's modules are dummies — route the real compute to the host. + from ..interleaver import Events + + spec = { + "norm_path": norm.path if norm is not None else None, + "head_path": head.path, + "formulation": formulation, + } + return mediator.send(Events.UNEMBED, (residual, spec)) + + # in-process / fast lane: the real modules + weights are right here. + normed = norm(residual) if norm is not None else residual + if formulation == "weight": + return F.linear(normed, head.weight) + return head(normed) + def barrier(self, n_participants: int): """ nnsight barrier: A synchronization primitive for coordinating multiple concurrent invocations in nnsight. From db9652c12d3cb4070b785c34b1ebcdb2146cbf2e Mon Sep 17 00:00:00 2001 From: khaiwang Date: Mon, 15 Jun 2026 01:01:18 -0400 Subject: [PATCH 18/30] =?UTF-8?q?feat(intervention):=20tracer.steer=20?= =?UTF-8?q?=E2=80=94=20replacement-swap=20injection=20for=20the=20isolated?= =?UTF-8?q?=20tier?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Add tracer.steer(envoy, direction, alpha), the next part-2 declarative primitive after tracer.unembed. It adds alpha*direction to a module's output residual via a *replacement* boundary write (assign envoy.output), which routes through the eproperty setter -> Events.SWAP and ships the steered value back on either tier. Steering touches no host weights — only the delivered activation — so unlike unembed it needs no host round-trip and no isolated/in-process branch: the same method is correct in-process, on the fast lane, and in the isolated worker. The point is the replacement swap. The hand-written additive form is in-place (block.output[:, -1, :] += direction); under isolation that mutates only the worker's delivered clone, no SWAP fires, the host's real activation is untouched, and the steering silently no-ops. tracer.steer makes it cross the boundary by construction. Tuple outputs (attention modules) are replaced whole, steering element [0] and carrying the tail (incl. a None) through pack_cuda untouched. The classifier already treats tracer.steer as a trusted nnsight primitive (its __module__ is nnsight.*), so no gate change is needed. Verified (test_isolated_steer.py, gpt2 + a renamed model), all max|Δ|=0: steering a block, an attention tuple output, and three blocks at once are isolated-vs-in-process bit-identical and propagate through later layers; tracer.steer equals the manual whole-tuple replacement; and the crux — under forced isolation the in-place form leaves the downstream residual at the unsteered baseline (silent no-op) while tracer.steer takes effect and matches the in-process result. Classifier units still 17/17. Doc: docs/developing/fast-lane.md §6 (steer marked built + subsection), §7, §8. Co-Authored-By: Claude Opus 4.8 (1M context) --- docs/developing/fast-lane.md | 57 ++++- .../gpu_sandbox/test_isolated_steer.py | 223 ++++++++++++++++++ src/nnsight/intervention/tracing/tracer.py | 45 ++++ 3 files changed, 319 insertions(+), 6 deletions(-) create mode 100644 prototypes/mediator-sandbox/gpu_sandbox/test_isolated_steer.py diff --git a/docs/developing/fast-lane.md b/docs/developing/fast-lane.md index fe06d0c51..262451edf 100644 --- a/docs/developing/fast-lane.md +++ b/docs/developing/fast-lane.md @@ -121,14 +121,17 @@ tier too, and (b) collapsing common raw-compute patterns into named calls the ga | primitive | signature | taxonomy primitive | role | status | |---|---|---|---|---| | `tracer.unembed` | `(residual, norm, head, formulation="weight") → logits` | host-weight read + module call | the projection every logit-lens / steering-direction / attribution metric does | **built** | -| `tracer.steer` | `(envoy, direction, alpha)` | boundary write (injection) | always a replacement swap — fixes in-place's silent no-op under isolation | designed | +| `tracer.steer` | `(envoy, direction, alpha=1.0)` | boundary write (injection) | always a replacement swap — fixes in-place's silent no-op under isolation | **built** | | `tracer.patch` | `(envoy, value)` | boundary write (transplant) | whole-tuple replacement | designed | | `tracer.ablate` | `(envoy, mode)` | boundary write (injection) | zero/mean knockout | designed | | `tracer.capture` | `(value) → handle` | read + run↔run transfer | cross-trace handoff; non-transmittable → clean fail, not silent drop | designed | -Each mirrors the existing `tracer.cache()` shape: in-process it resolves the real envoys and runs -directly (the fast-lane execution); isolated it ships a spec via a new event whose host handler runs the -real op host-side. +Most mirror the existing `tracer.cache()` shape: in-process they resolve the real envoys and run +directly (the fast-lane execution); isolated they ship a spec via a new event whose host handler runs the +real op host-side. The exception is `tracer.steer` (and the boundary-write `patch`/`ablate`): steering +touches **no host weights** — only the *delivered* activation — so it needs no host round-trip and no new +event. It rides the existing `Events.SWAP`: a *replacement* write (assign `envoy.output`) ships the +steered value back, which is what makes it cross the boundary where an in-place `[:] =` silently no-ops. ### `tracer.unembed` — host-routed readout (built, 2026-06-14) @@ -156,6 +159,43 @@ model):** single-layer, 3-layer-interleaved, `formulation="module"`, `norm=None` readouts all isolated-vs-in-process `max|Δ|=0`; `tracer.unembed` equals the manual `F.linear(norm(x), head.weight)` it replaces. +### `tracer.steer` — replacement-swap injection (built, 2026-06-15) + +`tracer.steer(envoy, direction, alpha=1.0)` adds `alpha * direction` to a module's output residual — +the activation-steering / injection every steering cell does. It is the simplest of the part-2 +primitives and the structural opposite of `unembed`: it touches **no host weights**, only the delivered +activation, so it needs **no new event and no isolated/in-process branch**. The method just performs a +*replacement* boundary write: + +```python +out = envoy.output +hidden = out[0] if isinstance(out, tuple) else out +steered = hidden + alpha * direction.to(dtype=hidden.dtype, device=hidden.device) +envoy.output = (steered, *out[1:]) if isinstance(out, tuple) else steered +``` + +The eproperty setter routes that assignment through `mediator.swap` → `Events.SWAP` on **either** tier: +in-process it swaps into the batcher directly; isolated, the worker ships the steered value back over the +existing SWAP path (`pack_cuda` walks the tuple, carrying a `None` tail through untouched). The same code +is therefore correct in-process, on the fast lane, and in the isolated worker. + +The point is the **replacement** swap. The hand-written additive form is in-place +(`block.output[:, -1, :] += direction` — the canonical nnsight steering); under isolation that mutates +only the worker's *delivered clone*, no SWAP fires, the host's real activation is untouched, and the +steering silently no-ops (the save of the steered residual even looks right — only the downstream forward +reveals nothing changed). `tracer.steer` makes the steering cross the boundary by construction. Tuple +outputs (most attention modules — `(tensor, None)` here) are replaced whole, steering element `[0]`. + +Touch points: just the `tracer.steer` method (tracer.py). No event, no host handler — it reuses +`Events.SWAP` and the eproperty setter. + +**Verified (`test_isolated_steer.py`, gpt2 + renamed model):** steering one block, an attention tuple +output, and three blocks at once are all isolated-vs-in-process `max|Δ|=0` and propagate through later +layers; `tracer.steer` equals the manual untuple + whole-tuple replacement. The crux case proves the +motivation under forced isolation: the in-place form leaves the downstream residual == the unsteered +baseline (silent no-op) while `tracer.steer` changes it (steering took effect) to exactly the in-process +result. + ## 7. What was deliberately deferred - **The process-global `sys.addaudithook` backstop.** Its own failure mode (a leaked thread-local flag @@ -164,8 +204,9 @@ readouts all isolated-vs-in-process `max|Δ|=0`; `tracer.unembed` equals the man fast-laned code. Documented as future hardening; the static gate is the confirmation. - **A frozen-namespace `Compartment`** (SES-style) for fast-lane execution. The first slice relies on the static pass + the `trust` cordon; namespace shadowing is a later refinement. -- **The remaining primitives' isolated event handlers** (§6) — `steer`/`patch`/`ablate`/`capture`; - `unembed` is built. +- **The remaining primitives** (§6) — `patch`/`ablate`/`capture`. `unembed` and `steer` are built + (`steer` rides `Events.SWAP`, so it needed no new handler; `patch`/`ablate` will too, `capture` needs + a run↔run handoff). ## 8. Verification @@ -178,6 +219,10 @@ readouts all isolated-vs-in-process `max|Δ|=0`; `tracer.unembed` equals the man (`fast_lane=False`) — proving the fast lane is the enabling tier; in-place steering bit-identical; renamed-model lens bit-identical; a footgun routes off the fast lane and the host survives; an introspection escape is rejected; a runaway loop is killed by the watchdog and the host survives. +- **Isolated steer** (`test_isolated_steer.py`, gpt2 + a renamed model) — 6/6: steering a block, an + attention tuple output, and three blocks at once are isolated-vs-in-process `max|Δ|=0`; `tracer.steer` + equals the manual whole-tuple replacement; and the crux — under forced isolation the in-place form is a + no-op (downstream == unsteered baseline) while `tracer.steer` takes effect and matches in-process. - **Existing isolated WORKER path** — 9/9 still bit-identical, pinned with `fast_lane=False` so they keep exercising the worker (otherwise the simple read/swap/save cells would now fast-lane). - **In-process core** — 51 passed (the default in-process path is untouched). diff --git a/prototypes/mediator-sandbox/gpu_sandbox/test_isolated_steer.py b/prototypes/mediator-sandbox/gpu_sandbox/test_isolated_steer.py new file mode 100644 index 000000000..712ac9a5b --- /dev/null +++ b/prototypes/mediator-sandbox/gpu_sandbox/test_isolated_steer.py @@ -0,0 +1,223 @@ +#!/usr/bin/env python3 +"""Activation steering on the ISOLATED tier: tracer.steer injects a direction into a +module's output via a REPLACEMENT swap, so the steering crosses the isolation boundary and +propagates through the host's real forward — where the hand-written in-place form +(``hidden[:] = …``) silently no-ops (the worker mutates its delivered clone, no SWAP fires, +the host's activation is untouched). + +Unlike tracer.unembed, steering touches no host weights — only the delivered activation — +so it needs no host round-trip: the eproperty setter routes the SWAP on either tier and the +SAME method is correct in-process, on the fast lane, and in the isolated worker. The +steering DIRECTION is precomputed outside the trace (a host-weight read inside a forced- +isolation trace would hit the weightless worker — the unembed problem); that mirrors real +usage, where steering vectors are precomputed. + + single — steer one block (single-tensor output), forced-isolation downstream + residual == in-process, bit-identical: the replacement swap crossed the + boundary AND propagated through later layers. + replacement — THE crux: under forced isolation the in-place form is a no-op (downstream + == unsteered baseline) while tracer.steer actually steers (downstream != + baseline) and equals the in-process steered result. + tuple_output — steer an attention output (a tuple `(tensor, None)`): the whole-tuple + replacement branch, forced-isolation == in-process. + multi — steer three blocks (forward order), forced-isolation == in-process. + renamed — renamed model (decoder_blocks): forced-isolation == in-process (no names + hardcoded; the wire path is the real path). + matches_manual— tracer.steer (in-process) equals the manual untuple+replacement it names. + +Run: + CUDA_VISIBLE_DEVICES=5 PYTHONPATH=src \ + /disk/u/zikai/anaconda3/envs/hf-serve/bin/python -u \ + prototypes/mediator-sandbox/gpu_sandbox/test_isolated_steer.py +""" +import sys + +import torch +import torch.nn.functional as F + +from nnsight import LanguageModel +from nnsight.intervention.isolation import isolate_mediators + +PROMPT = "The Eiffel Tower is in the city of" +ALPHA = 8.0 + + +def _direction(head, token_id=5000): + """A precomputed unit steering vector, on CPU and detached — captured into the trace + body so the isolated worker never reads a host weight (it has none).""" + return F.normalize(head.weight[token_id].float(), dim=0).detach().cpu() + + +def _both(build): + """Run build() in-process and under forced isolation; return (ref, got).""" + ref = build() + with isolate_mediators(fast_lane=False, timeout=30): + got = build() + return ref, got + + +def _eq(ref, got): + return (torch.is_tensor(ref) and torch.is_tensor(got) + and ref.shape == got.shape and torch.equal(ref, got)) + + +def _delta(ref, got): + return (ref - got).abs().max().item() if _eq(ref, got) else float("nan") + + +def test_single(model): + direction = _direction(model.lm_head) + + def build(): + with model.trace(PROMPT) as tracer: + with torch.no_grad(): + tracer.steer(model.transformer.h[6], direction, ALPHA) + o = model.transformer.h[10].output # downstream residual + down = (o[0] if isinstance(o, tuple) else o)[:, -1, :].save() + return down + ref, got = _both(build) + ok = _eq(ref, got) + print(f"[single] isolated steer downstream bit-identical={ok} (max|Δ|={_delta(ref, got)})", flush=True) + return ok + + +def test_replacement(model): + """The crux: replacement swap crosses the boundary; in-place is a silent no-op.""" + direction = _direction(model.lm_head) + + def read_downstream(): # unsteered baseline + with model.trace(PROMPT): + with torch.no_grad(): + o = model.transformer.h[10].output + down = (o[0] if isinstance(o, tuple) else o)[:, -1, :].save() + return down + + def steer_downstream(): # tracer.steer (replacement swap) + with model.trace(PROMPT) as tracer: + with torch.no_grad(): + tracer.steer(model.transformer.h[6], direction, ALPHA) + o = model.transformer.h[10].output + down = (o[0] if isinstance(o, tuple) else o)[:, -1, :].save() + return down + + def inplace_downstream(): # hand-written in-place steering + with model.trace(PROMPT): + with torch.no_grad(): + out = model.transformer.h[6].output + hidden = out[0] if isinstance(out, tuple) else out + hidden[:] = hidden + ALPHA * direction.to(dtype=hidden.dtype, device=hidden.device) + o = model.transformer.h[10].output + down = (o[0] if isinstance(o, tuple) else o)[:, -1, :].save() + return down + + base = read_downstream() + steer_ip = steer_downstream() # in-process steer + with isolate_mediators(fast_lane=False, timeout=30): + steer_iso = steer_downstream() # isolated steer (replacement swap) + inplace_iso = inplace_downstream() # isolated in-place (no-op) + + inplace_is_noop = _eq(inplace_iso, base) # in-place never crossed the boundary + steer_took_effect = not _eq(steer_iso, base) # replacement swap changed the host forward + steer_correct = _eq(steer_iso, steer_ip) # ... to exactly the in-process result + ok = inplace_is_noop and steer_took_effect and steer_correct + print(f"[replacement] isolated in-place is a no-op={inplace_is_noop}; " + f"isolated steer took effect={steer_took_effect}; " + f"steer iso==in-process={steer_correct} (max|Δ|={_delta(steer_iso, steer_ip)})", flush=True) + return ok + + +def test_tuple_output(model): + # An attention output is a tuple (tensor, None) — exercises the whole-tuple replacement + # branch (steered element [0], None tail carried through). + direction = _direction(model.lm_head) + + def build(): + with model.trace(PROMPT) as tracer: + with torch.no_grad(): + tracer.steer(model.transformer.h[6].attn, direction, ALPHA) + o = model.transformer.h[10].output + down = (o[0] if isinstance(o, tuple) else o)[:, -1, :].save() + return down + ref, got = _both(build) + ok = _eq(ref, got) + print(f"[tuple_output] tuple (attn) steer bit-identical={ok} (max|Δ|={_delta(ref, got)})", flush=True) + return ok + + +def test_multi(model): + direction = _direction(model.lm_head) + + def build(): + with model.trace(PROMPT) as tracer: + with torch.no_grad(): + for i in [2, 5, 8]: # forward order + tracer.steer(model.transformer.h[i], direction, ALPHA) + o = model.transformer.h[11].output + down = (o[0] if isinstance(o, tuple) else o)[:, -1, :].save() + return down + ref, got = _both(build) + ok = _eq(ref, got) + print(f"[multi] 3-block isolated steer bit-identical={ok} (max|Δ|={_delta(ref, got)})", flush=True) + return ok + + +def test_renamed(): + rename = {"transformer.ln_f": "final_norm", "lm_head": "output_projection", + "transformer.h": "decoder_blocks"} + model = LanguageModel("gpt2", device_map="cuda", dispatch=True, rename=rename) + direction = _direction(model.output_projection) + + def build(): + with model.trace(PROMPT) as tracer: + with torch.no_grad(): + tracer.steer(model.decoder_blocks[3], direction, ALPHA) + o = model.decoder_blocks[9].output + down = (o[0] if isinstance(o, tuple) else o)[:, -1, :].save() + return down + ref, got = _both(build) + ok = _eq(ref, got) + print(f"[renamed] renamed-model isolated steer bit-identical={ok} (max|Δ|={_delta(ref, got)})", flush=True) + return ok + + +def test_matches_manual(model): + # tracer.steer (in-process) must equal the manual untuple + replacement it names. + direction = _direction(model.lm_head) + with model.trace(PROMPT) as tracer: + with torch.no_grad(): + tracer.steer(model.transformer.h[6], direction, ALPHA) + o = model.transformer.h[10].output + via_api = (o[0] if isinstance(o, tuple) else o)[:, -1, :].save() + with model.trace(PROMPT): + with torch.no_grad(): + out = model.transformer.h[6].output + is_tuple = isinstance(out, tuple) + hidden = out[0] if is_tuple else out + steered = hidden + ALPHA * direction.to(dtype=hidden.dtype, device=hidden.device) + model.transformer.h[6].output = (steered, *out[1:]) if is_tuple else steered + o = model.transformer.h[10].output + via_manual = (o[0] if isinstance(o, tuple) else o)[:, -1, :].save() + ok = _eq(via_api, via_manual) + print(f"[matches_manual] tracer.steer == manual replacement bit-identical={ok} " + f"(max|Δ|={_delta(via_api, via_manual)})", flush=True) + return ok + + +def main(): + assert torch.cuda.is_available() + model = LanguageModel("gpt2", device_map="cuda", dispatch=True) + results = { + "single": test_single(model), + "replacement": test_replacement(model), + "tuple_output": test_tuple_output(model), + "multi": test_multi(model), + "renamed": test_renamed(), + "matches_manual": test_matches_manual(model), + } + print("=" * 72, flush=True) + print(f"ISOLATED STEER: {results}", flush=True) + sys.exit(0 if all(results.values()) else 1) + + +if __name__ == "__main__": + main() diff --git a/src/nnsight/intervention/tracing/tracer.py b/src/nnsight/intervention/tracing/tracer.py index d5bcf53d3..1add36624 100755 --- a/src/nnsight/intervention/tracing/tracer.py +++ b/src/nnsight/intervention/tracing/tracer.py @@ -701,6 +701,51 @@ def unembed(self, residual, norm, head, formulation: str = "weight"): return F.linear(normed, head.weight) return head(normed) + def steer(self, envoy, direction, alpha: float = 1.0): + """Add ``alpha * direction`` to ``envoy``'s output residual via a **replacement** + boundary write — the activation-steering / injection every steering cell does, made + to work on the isolated tier where an in-place ``hidden[:] = …`` silently no-ops. + + Steering reads the delivered activation and writes a modified one back. Done in + place (``hidden[:] = hidden + …``) that mutation never crosses the isolation + boundary: the worker mutates its *delivered clone*, no SWAP fires, and the host's + real activation is untouched — a silent no-op (the save of the steered residual + looks right, but nothing downstream changes). ``tracer.steer`` always performs a + *replacement* swap (assign ``envoy.output``), which ships the steered value back + over the existing ``Events.SWAP`` path, so it is correct in-process, on the fast + lane, AND in the isolated worker. Unlike :meth:`unembed`, steering touches no host + weights — only the delivered activation — so it needs no host round-trip and no + isolated/in-process branch: the eproperty setter routes the swap on either tier. + + Tuple outputs (attention modules, and transformer blocks on transformers <5) are + replaced whole — element ``[0]`` is steered and the rest of the tuple rides through + unchanged. + + Args: + envoy: the module whose output residual to steer (e.g. + ``model.transformer.h[6]``). + direction: the steering vector, broadcast over the residual and cast to its + dtype/device (so a float32 direction steers a float16 residual cleanly). + alpha: the steering coefficient — a scalar or any tensor broadcastable over the + residual. Defaults to ``1.0``. + + Returns: + The steered residual tensor (element ``[0]`` for tuple outputs). + """ + out = envoy.output + is_tuple = isinstance(out, tuple) + hidden = out[0] if is_tuple else out + + direction = direction.to(dtype=hidden.dtype, device=hidden.device) + steered = hidden + alpha * direction + + if is_tuple: + envoy.output = (steered, *out[1:]) + else: + envoy.output = steered + + return steered + def barrier(self, n_participants: int): """ nnsight barrier: A synchronization primitive for coordinating multiple concurrent invocations in nnsight. From 8f8898605bfd281d9d59d2f03b8b082cbefcfd79 Mon Sep 17 00:00:00 2001 From: khaiwang Date: Wed, 17 Jun 2026 10:47:44 -0500 Subject: [PATCH 19/30] feat(isolation): warm-time preimport for whitelist parity + triton-model coverage MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Add a `preimport=` option to isolate_mediators()/warm_worker_pool() that loads modules at worker warm time, BEFORE seccomp lockdown freezes new file opens (import == open()). This brings user-facing import capability under lockdown to parity with an in-process module whitelist without weakening containment: the model's own kernels (incl. triton) run host-side and are unaffected. - thread `preimport` through _STATE -> _base_opts -> the pool key (now per device/arena/mem/lockdown/preimport signature) and consume it in _pool_worker_main warmup, before lock_down(). - test_isolated_triton_model.py: a @triton.jit-kernel model traced under isolation+lockdown is bit-identical to in-process (host compiles triton while the worker is fully locked down); worker-side triton in the intervention is blocked. Requires GPU + triton. - docs §16: the triton deployment motivation, the strictly-better-than-upstream module-restriction comparison, and the verified timeout-directionality and cold-vs-pool lockdown-ordering facts. Fix the stale docstring claiming the cold path deserializes before lockdown (the unified worker locks down before both). Co-Authored-By: Claude Opus 4.8 --- .../mediator-gpu-trace-integration.md | 66 +++++++ .../gpu_sandbox/test_isolated_triton_model.py | 164 ++++++++++++++++++ src/nnsight/intervention/isolation.py | 53 ++++-- 3 files changed, 272 insertions(+), 11 deletions(-) create mode 100644 prototypes/mediator-sandbox/gpu_sandbox/test_isolated_triton_model.py diff --git a/docs/developing/mediator-gpu-trace-integration.md b/docs/developing/mediator-gpu-trace-integration.md index 722002544..0a6b71f94 100644 --- a/docs/developing/mediator-gpu-trace-integration.md +++ b/docs/developing/mediator-gpu-trace-integration.md @@ -260,6 +260,8 @@ benign CudaIPC release warning. | `tracer.cache()` (`modules=`, `include_inputs=`) | ✅ bit-identical — CACHE event → host registers the real cache hooks; the forward fills the host CacheDict in-place (§15) | | `.source` operation-level access (`...attn.split_1.output`) | 🔜 not yet (op paths aren't in `model.modules()`) | | in-place `[:]=` | ⛔ use explicit `=` (clone semantics, §4) | +| Triton-kernel models (MoE / SSM / `torch.compile`) | ✅ host-side forward compiles Triton unrestricted — §16 | +| user-code Triton (kernel inside the intervention) | ⛔ blocked under lockdown by design — §16 | Not-yet-supported features fail **cleanly** (missed-provider error or the per-step timeout), not as a silent deadlock — the lifecycle (timeout + `finally: cancel()`) is the safety net until each feature lands. @@ -487,3 +489,67 @@ regression (trace/acceptance/multitoken/cross-invoke/pool) unchanged. **Not covered:** a cache placeholder nested inside a container save (`got = [t.cache().save()]`) — the swap matches a top-level saved `CacheDict`; nesting would need a recursive walk. `cross_invoker` + cache is untested. `modules=None` (cache *all* modules) registers a hook per module on the host — correct but heavy. + +--- + +## 16. Triton-kernel models — the deployment motivation (and why this beats the in-process whitelist) + +**Why this is urgent.** Frontier GPU model execution increasingly runs **Triton JIT kernels**, by three +independent routes: (1) architectures whose core op has no fast eager form ship Triton kernels — fused MoE +(vLLM `fused_moe`: Mixtral, DeepSeek-V2/V3/R1, Qwen-MoE, Llama-4, …) and SSM selective-scan (HF `mamba2` +imports `mamba_ssm.ops.triton.selective_state_update`; Jamba/Bamba/FalconMamba/Zamba); (2) HF's `kernels` +library + `kernelize(model, mode=inference)` pulls Triton kernels from the Hub as drop-in +norm/activation/attention replacements (Liger fuses RMSNorm/RoPE/SwiGLU/CE); (3) `torch.compile` → +TorchInductor, whose GPU codegen target *is* Triton. The triton-free path is essentially limited to a plain +dense transformer in eager mode with SDPA / precompiled-CUDA FlashAttention (which is an AOT library call, +not runtime codegen). + +**Why the in-process whitelist can't serve them.** NDIF's sandbox +(`ndif:src/ndif/services/ray/nn/security/`) is a Python import whitelist that wraps the ENTIRE +`tracer.execute(model)` — the forward runs *inside* the Protector (unavoidable in-process: user +interventions interleave with the forward in one call stack). Triton's first-use compilation needs +`subprocess`+ptxas, `tempfile`, `os`, `importlib`, `open()` — none whitelisted, several hard-blocked by the +audit hook. So the model's own Triton kernels, firing during the sandboxed forward, are denied. A +pre-warmed kernel cache doesn't save it: per-request shapes recompile, and even a cache *hit* still needs +`os`/`open`/`importlib` to load the cached `.so`. + +**Why this backend fixes it — the trust boundary aligns with the capability boundary.** The forward runs on +the TRUSTED host (unrestricted → Triton compiles normally); only UNTRUSTED user intervention code runs in +the worker. The split lands on the seam the whitelist couldn't find in-process: the thing needing dangerous +capability (runtime codegen) is the *model* (trusted); the thing we distrust is *user code*. The worker's +lockdown level is therefore **orthogonal** to whether Triton models run — worker security can be cranked to +the max without touching model Triton. + +**Strictly better than upstream on the module-restriction axis.** Against the right baseline (upstream +*sandboxed* NDIF, not un-sandboxed nnsight): + +| | Upstream whitelist | This backend (lockdown on) | +|---|---|---| +| user-code imports | restricted (allowlist) | restricted (seccomp import-freeze) — match-or-stricter | +| model serving stack | **also restricted** → Triton broken | **unrestricted** (host-side) → Triton works | +| user escape past Python layer | reaches host / tenants (whitelist README concedes C-level escapes; suite: 10/10 escapes succeed) | contained in a refless, killable process | + +The seccomp **import-freeze** (under lockdown, only warm-time-loaded modules are importable, since +import == `open()`, which seccomp then blocks) is therefore *not a regression* — upstream restricts user +imports at least as hard. The one parity item is *which* modules a legit user may reach; close it with +`isolate_mediators(..., preimport=(...))` (and the matching `warm_worker_pool(..., preimport=(...))`), which +loads the deployment's allowed-module set at warm time, before lockdown. User-code Triton is **doubly +blocked** (import wall if not pre-warmed + compile wall: ptxas execve / cache open) — intended containment, +documented as unsupported. + +**Verified properties (code, not just design):** +- *Timeout directionality is correct for slow host compiles.* The worker waits on the host with **no + timeout** (`CudaIpcWorkerChannel.wait_response` — a blocking `recv`); the timeout lives only on the host + (`CudaIpcHostChannel.wait_event`) and measures worker think-time. Host-side Triton compilation happens + during the host's *forward execution* — never while the host is in `wait_event` — so a multi-second cold + MoE autotune never false-trips the worker's hang-detector, in either direction. +- *Lockdown ordering / cold-vs-pool.* `lock_down()` runs once after warm-up, before the job loop, in the + unified `_pool_worker_main` (the only worker entrypoint, used for cold via `poolable=False` and pooled via + `poolable=True`). So the import set is frozen at warm time for **both** paths — there is no + "cold deserializes before lockdown" advantage; the cold-vs-pool difference is recycle-vs-retire. (The + earlier §7 note describing a cold deserialize-before-lockdown window predates the warm-pool unification.) + +**Coverage:** `prototypes/mediator-sandbox/gpu_sandbox/test_isolated_triton_model.py` — `host_compiles` +(isolated + `lockdown=True` Triton-kernel model bit-identical to in-process, with a cold `TRITON_CACHE_DIR` +shown to populate during the isolated run) and `user_contained` (worker-side Triton blocked under +lockdown, exercising `preimport=`). Requires a GPU + a Triton install. diff --git a/prototypes/mediator-sandbox/gpu_sandbox/test_isolated_triton_model.py b/prototypes/mediator-sandbox/gpu_sandbox/test_isolated_triton_model.py new file mode 100644 index 000000000..814140268 --- /dev/null +++ b/prototypes/mediator-sandbox/gpu_sandbox/test_isolated_triton_model.py @@ -0,0 +1,164 @@ +#!/usr/bin/env python3 +"""Triton-kernel model under isolation — the headline motivation for this backend. + +A model whose forward executes a Triton JIT kernel cannot be served under NDIF's +in-process import-whitelist sandbox: Triton's first-use compilation needs subprocess +(ptxas), tempfile, open() and importlib — exactly what the whitelist denies — and the +whitelist necessarily wraps the forward pass too. This backend moves the forward to the +TRUSTED host (unrestricted, so Triton compiles normally) and contains only the UNTRUSTED +intervention in the worker. These tests pin that property: + + host_compiles — an isolated trace with ``lockdown=True`` through a Triton-kernel model + is bit-identical to in-process. The worker is fully seccomp'd, so the + kernel MUST have compiled+run on the host => the host/worker split lets + Triton models run under isolation. (A fresh TRITON_CACHE_DIR is shown + to populate during the isolated run as corroboration.) + user_contained — a Triton kernel invoked from WITHIN the intervention (worker-side) under + lockdown fails cleanly (import wall and/or compile wall), demonstrating + that the capability preserved for the model is still denied to user code. + Uses ``preimport=("triton", "triton.language")`` to get the worker past + the import wall so the *compile* wall is what's exercised. + +Run (needs a GPU + a Triton install): + CUDA_VISIBLE_DEVICES=0 PYTHONPATH=src \ + /disk/u/zikai/anaconda3/envs/hf-serve/bin/python \ + prototypes/mediator-sandbox/gpu_sandbox/test_isolated_triton_model.py +""" +import os +import sys +import tempfile + +# Point Triton at a fresh cache BEFORE importing it, so a populated cache after the +# isolated run is evidence the host compiled during that run (not a pre-warmed hit). +_CACHE_DIR = tempfile.mkdtemp(prefix="nnsight_triton_cache_") +os.environ["TRITON_CACHE_DIR"] = _CACHE_DIR + +import torch +import torch.nn as nn + +import nnsight +from nnsight import NNsight +from nnsight.intervention.isolation import isolate_mediators + +try: + import triton + import triton.language as tl + + _HAVE_TRITON = True +except Exception: # noqa: BLE001 + _HAVE_TRITON = False + + +if _HAVE_TRITON: + + @triton.jit + def _scale_kernel(x_ptr, out_ptr, scale, n, BLOCK: tl.constexpr): + pid = tl.program_id(0) + offs = pid * BLOCK + tl.arange(0, BLOCK) + mask = offs < n + x = tl.load(x_ptr + offs, mask=mask) + tl.store(out_ptr + offs, x * scale, mask=mask) + + def triton_scale(x, scale): + """Elementwise x*scale via a Triton kernel (forces a Triton JIT at first use).""" + x = x.contiguous() + out = torch.empty_like(x) + n = x.numel() + grid = lambda meta: (triton.cdiv(n, meta["BLOCK"]),) # noqa: E731 + _scale_kernel[grid](x, out, scale, n, BLOCK=1024) + return out + + +class TritonMLP(nn.Module): + """fc1 -> (Triton kernel) -> fc2. The Triton kernel lives ON the forward path, so a + trace of this model only succeeds if that kernel compiled+ran wherever the forward + executed (the host, under isolation).""" + + def __init__(self, d=64): + super().__init__() + self.fc1 = nn.Linear(d, d) + self.fc2 = nn.Linear(d, d) + + def forward(self, x): + h = self.fc1(x) + h = triton_scale(h, 2.0) + return self.fc2(h) + + +def _build(): + torch.manual_seed(0) + net = TritonMLP().cuda().eval() + return NNsight(net), torch.randn(1, 8, 64).cuda() + + +def test_host_compiles(): + """Isolated + lockdown trace through a Triton-kernel model == in-process.""" + model, x = _build() + + # Isolated FIRST, with a cold Triton cache and the worker fully locked down. + with isolate_mediators(lockdown=True): + with model.trace(x): + got = model.fc2.output.save() + + cache_populated = os.path.isdir(_CACHE_DIR) and any(os.scandir(_CACHE_DIR)) + + # In-process reference (cache now warm; comparison is what matters). + with model.trace(x): + ref = model.fc2.output.save() + + ok = torch.equal(ref, got) + print( + f"[host_compiles] isolated(lockdown) Triton model == in-process: {ok} " + f"(max|Δ|={(ref - got).abs().max().item():.2e}) | host triton cache populated " + f"during isolated run: {cache_populated}" + ) + # Bit-identity is the load-bearing assertion; cache population is corroboration only + # (cache layout/location is version-dependent), so it does not gate the result. + return ok + + +def test_user_contained(): + """A Triton kernel called from inside the intervention (worker) is blocked under + lockdown — the model keeps Triton, untrusted user code does not.""" + model, x = _build() + raised = None + try: + # preimport gets the worker past the import wall; the compile wall (ptxas execve + + # cache open, both seccomp-blocked) is then what must stop it. + with isolate_mediators(lockdown=True, preimport=("triton", "triton.language")): + with model.trace(x): + act = model.fc1.output # real tensor in the worker + # Worker-side Triton use: the JIT (ptxas execve + cache open) must be + # blocked under lockdown. This call is what should raise. + out = triton_scale(act, 3.0) + nnsight.save(out) + except Exception as e: # noqa: BLE001 + raised = e + ok = raised is not None + print( + f"[user_contained] worker-side Triton blocked under lockdown: {ok} " + f"(raised={type(raised).__name__ if raised else None})" + ) + return ok + + +def main(): + if not torch.cuda.is_available(): + print("SKIP: no CUDA device") + sys.exit(0) + if not _HAVE_TRITON: + print("SKIP: triton not installed") + sys.exit(0) + + results = { + "host_compiles": test_host_compiles(), + "user_contained": test_user_contained(), + } + ok = all(results.values()) + print("=" * 72) + print(f"ISOLATED TRITON MODEL: {'PASS' if ok else 'FAIL'} — {results}") + sys.exit(0 if ok else 1) + + +if __name__ == "__main__": + main() diff --git a/src/nnsight/intervention/isolation.py b/src/nnsight/intervention/isolation.py index de738613c..f45514b9a 100644 --- a/src/nnsight/intervention/isolation.py +++ b/src/nnsight/intervention/isolation.py @@ -84,6 +84,7 @@ def _transmittable(v) -> bool: "timeout": 60.0, # per-step wall-clock cap on user code (hang containment) "lockdown": False, # functional-first; seccomp lockdown enabled separately "pool_size": 0, # 0 => cold one-shot worker per trace; >0 => warm pool cap + "preimport": (), # modules to load at warm time, before lockdown freezes imports } @@ -99,6 +100,7 @@ def isolate_mediators( timeout: float = 60.0, lockdown: bool = False, pool_size: int = 0, + preimport: tuple = (), ): """Run interventions inside ``with model.trace(...)`` in an isolated GPU worker. @@ -110,13 +112,19 @@ def isolate_mediators( timeout: per-step wall-clock cap on user code; a worker that produces no event within ``timeout`` is presumed hung and killed (the host survives). pool_size: if > 0, draw workers from a process-global warm pool capped at - ``pool_size`` per (device, arena_bytes, gpu_mem_fraction, lockdown) - signature (auto-grown lazily, persists across traces, falls back to a cold - one-shot worker past the cap). 0 (default) spawns a cold worker per trace — - the original behavior. Use :func:`warm_worker_pool` to pre-warm at startup. - Under ``lockdown=True`` a pooled worker locks its import set at warm time, so - a job whose user code triggers a NEW import fails (consistently across the - pool) — stricter than the cold path, which deserializes before lockdown. + ``pool_size`` per (device, arena_bytes, gpu_mem_fraction, lockdown, + preimport) signature (auto-grown lazily, persists across traces, falls back + to a cold one-shot worker past the cap). 0 (default) spawns a cold worker per + trace — the original behavior. Use :func:`warm_worker_pool` to pre-warm at + startup. + preimport: module names to import at worker warm time, before seccomp lockdown + freezes new file opens. Under ``lockdown=True`` the import set is frozen at + warm time (import == ``open()``, which seccomp then blocks), so a job whose + user code triggers a NEW import fails; list here the modules interventions + are allowed to use, to bring user-facing import capability to parity with an + in-process whitelist. The freeze holds for the cold path too (cold and pooled + share the unified worker, which locks down before any job deserializes) — the + cold-vs-pool difference is recycle-vs-retire, not lockdown timing. """ prev = dict(_STATE) _STATE.update( @@ -127,6 +135,7 @@ def isolate_mediators( timeout=timeout, lockdown=lockdown, pool_size=pool_size, + preimport=tuple(preimport), ) try: yield @@ -142,6 +151,7 @@ def _base_opts() -> Dict[str, Any]: "gpu_mem_fraction": _STATE["gpu_mem_fraction"], "lockdown": _STATE["lockdown"], "timeout": _STATE["timeout"], + "preimport": _STATE["preimport"], } @@ -211,9 +221,10 @@ class _WorkerPool: def __init__(self): # Keyed by base-opts signature: workers are interchangeable ONLY within the - # same (device, arena_bytes, gpu_mem_fraction, lockdown) — the bounce buffer is - # device- and size-specific, so reusing a worker across devices would copy into - # the wrong-device buffer (silent corruption). + # same (device, arena_bytes, gpu_mem_fraction, lockdown, preimport) — the bounce + # buffer is device- and size-specific, so reusing a worker across devices would + # copy into the wrong-device buffer (silent corruption); and a worker's frozen + # import set must match the requested preimport list. self._idle: Dict[tuple, deque] = defaultdict(deque) self._all: Dict[tuple, set] = defaultdict(set) self._lock = threading.Lock() @@ -226,6 +237,7 @@ def _key(base_opts: dict) -> tuple: int(base_opts["arena_bytes"]), float(base_opts["gpu_mem_fraction"]), bool(base_opts.get("lockdown", False)), + tuple(sorted(base_opts.get("preimport", ()))), ) def warm(self, n: int, base_opts: dict) -> None: @@ -321,12 +333,15 @@ def warm_worker_pool( gpu_mem_fraction: float = 0.3, lockdown: bool = False, timeout: float = 60.0, + preimport: tuple = (), ) -> None: """Pre-warm ``size`` generic workers (blocks until each is ready). Call once at server startup so the first request pays no spawn cost. The base options here fix the pool's per-worker configuration. Each worker costs ~0.55 GiB - GPU per GPU it touches — size the pool as a GPU-memory budget. + GPU per GPU it touches — size the pool as a GPU-memory budget. ``preimport`` and + ``lockdown`` must match the values later passed to :func:`isolate_mediators` or the + pre-warmed workers won't match its pool signature and fresh ones are spawned. """ _POOL.warm( size, @@ -336,6 +351,7 @@ def warm_worker_pool( "gpu_mem_fraction": gpu_mem_fraction, "lockdown": lockdown, "timeout": timeout, + "preimport": tuple(preimport), }, ) @@ -725,6 +741,21 @@ def _pool_worker_main(conn, buf, base_opts): _ensure_mounted() # install Object.save so `.save()` resolves in the worker + # User-configurable warm-time pre-imports: load modules interventions may need + # BEFORE seccomp freezes new file opens (import == open()). Under lockdown a module + # not loaded here is unimportable in any job, so pre-warming the deployment's + # allowed-module set brings user-facing import capability to parity with an + # in-process whitelist. A failed pre-import is non-fatal (warn + skip). + if base_opts.get("preimport"): + import importlib + import warnings + + for _mod in base_opts["preimport"]: + try: + importlib.import_module(_mod) + except Exception as _e: # noqa: BLE001 + warnings.warn(f"isolated worker pre-import of {_mod!r} failed: {_e!r}") + if base_opts.get("gpu_mem_fraction") and torch.cuda.is_available(): torch.cuda.set_per_process_memory_fraction(base_opts["gpu_mem_fraction"]) From f72a39739110aecf707fc10673857be8ccc71147 Mon Sep 17 00:00:00 2001 From: khaiwang Date: Wed, 17 Jun 2026 23:06:43 -0500 Subject: [PATCH 20/30] harden(intervention): restricted-unpickler codec for worker->host frames MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The host read worker frames with mp.Pipe.recv() (pickle) — and the worker runs UNTRUSTED user code, so a crafted __reduce__ gadget over the control plane was a host-side RCE that would bypass every other isolation layer (seccomp, namespaces, row-bounding). The design hardened the inbound user payload (unpickled inside the worker) but not the outbound results. Close it: - worker->host is now tensor-free (tensors already ride the GPU buffer / safetensors) and the small remaining structure is decoded with a RESTRICTED unpickler (transport._RestrictedUnpickler / _safe_loads): find_class allows ONLY torch dtype/device and refuses every other class/function. find_class resolves a global before the REDUCE that would call it, so a gadget is refused before it can execute. - the event rides as its string .value (no enum class) and exceptions as a (type-name, message) sentinel (no class), so the allowlist stays {torch dtype, device}. host->worker stays normal pickle (host-authored, trusted). - this also fixes a real correctness gap a prior hand-rolled JSON codec had: the Events.CACHE spec carries torch.dtype/torch.device, which the tagged codec rejected; pickle handles all plain nested structures natively (no per-type enumeration), and anything un-allowlisted fails loud at decode with the class name. - capability narrowing: .save() of an arbitrary object / numpy / framework type (e.g. ModelOutput) is no longer transmittable from a worker — save a tensor instead. test_isolated_codec_security.py: fidelity for VALUE/SWAP/END/CACHE(dtype,device)/ EXCEPTION/push, and a genuine __reduce__ gadget refused at decode without executing. CPU is enough (needs torch). Legacy AF_UNIX socket channels are unused by isolate_mediators and still plain-unpickle — noted in the module header. Co-Authored-By: Claude Opus 4.8 --- .../test_isolated_codec_security.py | 134 +++++++++++++++++ src/nnsight/intervention/transport.py | 140 ++++++++++++++++-- 2 files changed, 264 insertions(+), 10 deletions(-) create mode 100644 prototypes/mediator-sandbox/gpu_sandbox/test_isolated_codec_security.py diff --git a/prototypes/mediator-sandbox/gpu_sandbox/test_isolated_codec_security.py b/prototypes/mediator-sandbox/gpu_sandbox/test_isolated_codec_security.py new file mode 100644 index 000000000..23c1b5dcc --- /dev/null +++ b/prototypes/mediator-sandbox/gpu_sandbox/test_isolated_codec_security.py @@ -0,0 +1,134 @@ +#!/usr/bin/env python3 +"""Worker->host frame codec — security + fidelity (restricted unpickler). + +The isolated worker runs UNTRUSTED user code; the host must not *plain* ``pickle.loads`` +its frames (a ``__reduce__`` gadget would execute on the trusted host = RCE). Worker->host +frames are tensor-free (tensors ride the GPU buffer / safetensors) and the rest is decoded +with ``transport._RestrictedUnpickler`` (``find_class`` allows ONLY torch dtype/device). + + fidelity - VALUE / SWAP / END / EXCEPTION / cross_invoker-push AND a tracer.cache()-style + spec carrying ``torch.dtype`` + ``torch.device`` all round-trip exactly. + security - a frame whose data is a ``__reduce__`` gadget is REFUSED at decode + (``find_class`` rejects ``os.system`` before the REDUCE could call it) — the + gadget never executes on the host. A non-allowlisted class is also refused. + +Needs torch; **CPU is enough — no CUDA required** (the bounce buffer is just bytes here). +Run: + PYTHONPATH=src python \ + prototypes/mediator-sandbox/gpu_sandbox/test_isolated_codec_security.py +""" +import os +import pickle +import sys + +import torch + +from nnsight.intervention import transport as T +from nnsight.intervention.interleaver import Events + + +def _roundtrip(event, data, push=None, nbytes=1 << 20): + """Encode a worker->host frame and decode it host-side, via a CPU 'bounce buffer'.""" + buf = torch.empty(nbytes, dtype=torch.uint8) # CPU stand-in for the GPU buffer + frame, _had = T._encode_worker_frame(event, data, push, buf) + return T._decode_worker_frame(frame, buf) + + +def test_fidelity(): + ok = True + + # VALUE: the requester string. + ev, out, _ = _roundtrip(Events.VALUE, "transformer.h.6.output.i0") + ok &= ev is Events.VALUE and out == "transformer.h.6.output.i0" + + # SWAP: (requester, value) carrying a tensor. + t = torch.randn(2, 3) + ev, out, _ = _roundtrip(Events.SWAP, ("h.6.output", t)) + ok &= ev is Events.SWAP and out[0] == "h.6.output" and torch.equal(out[1], t) + + # END: saved dict — tensor + scalars + nested list/tuple + tuple-with-None. + saved = { + "x": torch.arange(6).float(), + "n": 3, + "f": 1.5, + "lst": [1, (2, 3)], + "tup": (torch.ones(2), None), + } + ev, out, _ = _roundtrip(Events.END, saved) + ok &= ev is Events.END + ok &= torch.equal(out["x"], saved["x"]) and out["n"] == 3 and out["f"] == 1.5 + ok &= out["lst"] == [1, (2, 3)] and type(out["lst"][1]) is tuple + ok &= type(out["tup"]) is tuple and torch.equal(out["tup"][0], torch.ones(2)) + ok &= out["tup"][1] is None + + # CACHE: spec carrying torch.dtype + torch.device (the previously-missed case). + spec = (12345, ["transformer.h.0"], torch.device("cpu"), torch.float16, + True, True, False, {}, {}) + ev, out, _ = _roundtrip(Events.CACHE, spec) + ok &= ev is Events.CACHE + ok &= out[2] == torch.device("cpu") and out[3] is torch.float16 + ok &= out[0] == 12345 and out[1] == ["transformer.h.0"] + + # EXCEPTION: rebuilt host-side from (type-name, message); no object crosses. + ev, out, _ = _roundtrip(Events.EXCEPTION, ValueError("boom")) + ok &= ev is Events.EXCEPTION and isinstance(out, ValueError) and "boom" in str(out) + + # cross_invoker push: CPU tensors + scalars ride safetensors, not pickle. + push = {"shared": torch.randn(4), "k": 7} + ev, _out, pout = _roundtrip(Events.END, {"y": torch.zeros(1)}, push=push) + ok &= pout is not None and torch.equal(pout["shared"], push["shared"]) and pout["k"] == 7 + + print(f"[fidelity] VALUE/SWAP/END/CACHE(dtype,device)/EXCEPTION/push exact: {ok}") + return ok + + +def test_security(): + probe = "/tmp/nnsight_codec_pwn_probe" + if os.path.exists(probe): + os.remove(probe) + + class Bomb: + def __reduce__(self): + return (os.system, (f"echo pwned > {probe}",)) + + buf = torch.empty(1 << 16, dtype=torch.uint8) + + # The worker CAN pickle a gadget (encoding is safe — nothing runs). The host must + # REFUSE it at decode (find_class rejects os.system) before the gadget executes. + frame, _had = T._encode_worker_frame(Events.END, {"evil": Bomb()}, None, buf) + refused = False + try: + T._decode_worker_frame(frame, buf) + except pickle.UnpicklingError: + refused = True + no_pwn = not os.path.exists(probe) + if os.path.exists(probe): + os.remove(probe) + + # A non-allowlisted (but harmless) class is also refused — the allowlist is tight. + class Plain: + pass + + frame2, _ = T._encode_worker_frame(Events.END, {"obj": Plain()}, None, buf) + refused2 = False + try: + T._decode_worker_frame(frame2, buf) + except pickle.UnpicklingError: + refused2 = True + + ok = refused and no_pwn and refused2 + print(f"[security] gadget refused at decode + never executed: {refused and no_pwn} | " + f"non-allowlisted class refused: {refused2}") + return ok + + +def main(): + results = {"fidelity": test_fidelity(), "security": test_security()} + ok = all(results.values()) + print("=" * 72) + print(f"ISOLATED CODEC SECURITY: {'PASS' if ok else 'FAIL'} — {results}") + sys.exit(0 if ok else 1) + + +if __name__ == "__main__": + main() diff --git a/src/nnsight/intervention/transport.py b/src/nnsight/intervention/transport.py index 34eac4737..6e774a1bb 100644 --- a/src/nnsight/intervention/transport.py +++ b/src/nnsight/intervention/transport.py @@ -23,9 +23,17 @@ Frame codec: length-prefixed ``pickle``. This socket path forks the worker, so both ends are mutually trusted and pickle-both-ways is fine. Two follow-ups (see the plan): -1. **Security:** once the worker is untrusted, the *jail->host* direction MUST NOT - ``pickle.loads`` arbitrary objects — restrict the host-side decoder to - tensors/known frames. (jail<-host is host-authored, so the jail trusting it is ok.) +1. **Security:** the *worker->host* direction MUST NOT plain-``pickle.loads`` untrusted + worker bytes (a ``__reduce__`` gadget would be a host-side RCE). **Done for the active + GPU path** (:class:`CudaIpcHostChannel` / :class:`CudaIpcWorkerChannel`): worker->host + frames are tensor-free (tensors ride the GPU buffer / safetensors) and the small + remaining structure is decoded with the **restricted unpickler** + (:class:`_RestrictedUnpickler` / :func:`_safe_loads`) whose ``find_class`` allows ONLY + torch dtype/device; host->worker stays a normal pickle (host-authored, trusted). The + legacy AF_UNIX socket channels below (:class:`SocketHostChannel`, + :class:`ShmSocketHostChannel`) are NOT used by ``isolate_mediators`` and still + plain-``pickle.loads`` — route them through :func:`_safe_loads` before wiring them to an + untrusted worker. 2. **Performance (measured):** ``pickle`` of a torch tensor is the per-hook bottleneck — ``dumps``+``loads`` ~22 ms per direction at 16.8 MB and **superlinear**; @@ -37,6 +45,8 @@ from __future__ import annotations +import builtins +import io import mmap import os import pickle @@ -46,7 +56,7 @@ import torch -from .interleaver import MediatorChannel +from .interleaver import Events, MediatorChannel try: from safetensors.torch import load as _st_load @@ -343,6 +353,111 @@ def unpack_cuda(skel: Any, table: dict, buf: torch.Tensor) -> Any: return _merge_tensors(skel, tensors) +# --------------------------------------------------------------------------- # +# Safe worker->host frame codec — restricted unpickler, tight allowlist # +# --------------------------------------------------------------------------- # +# The worker runs UNTRUSTED user code, so the host MUST NOT do a plain +# ``pickle.loads`` of its frames (a ``__reduce__`` gadget = host RCE). Tensors +# travel out-of-band (the GPU bounce buffer / safetensors), so a worker->host +# frame is TENSOR-FREE and carries only plain data plus, at most, a torch +# ``dtype`` / ``device`` (the ``tracer.cache()`` spec). We decode it with a +# restricted Unpickler whose ``find_class`` allows ONLY torch dtype/device and +# refuses every other class/function. ``find_class`` is consulted to resolve a +# global BEFORE the ``REDUCE`` opcode could call it, so a gadget (``os.system`` +# etc.) is refused before it can execute. The event crosses as its string +# ``.value`` (no enum class) and exceptions as a (type-name, message) sentinel +# (no class), so the allowlist stays just {torch dtype, torch device}. +# (host->worker stays a normal pickle: that direction is host-authored, trusted.) + +_EXC_TAG = "__nnsight_iso_exc__" + + +class _RestrictedUnpickler(pickle.Unpickler): + """Unpickler for UNTRUSTED worker->host frames. Reconstructs only torch + dtype/device globals; every other class/function is refused — so no gadget + callable is ever resolved and the ``REDUCE`` that would call it never runs.""" + + def find_class(self, module: str, name: str): + if module == "torch": + obj = getattr(torch, name, None) + if obj is torch.device or isinstance(obj, torch.dtype): + return obj + raise pickle.UnpicklingError( + f"refusing to unpickle {module}.{name} from an isolated worker — only " + f"tensors (out-of-band) + basic data types cross the isolation boundary " + f"(a torch dtype/device is allowed; a custom object / numpy / framework " + f"type is not transmittable from a worker)" + ) + + +def _safe_loads(data: bytes) -> Any: + """Unpickle UNTRUSTED worker bytes with the restricted (allowlist) unpickler.""" + return _RestrictedUnpickler(io.BytesIO(data)).load() + + +def _rebuild_exc(name: str, msg: str) -> BaseException: + """Rebuild an exception from a (type-name, message) pair, host side. Resolves + ONLY builtin exception types (never an arbitrary callable); else RuntimeError.""" + cls = getattr(builtins, name, None) + if isinstance(cls, type) and issubclass(cls, BaseException): + try: + return cls(msg) + except Exception: # noqa: BLE001 — some exceptions have non-(str,) __init__ + pass + return RuntimeError(f"{name}: {msg}") + + +def _encode_worker_frame(event: Events, data: Any, push: Any, buf: torch.Tensor) -> tuple: + """Build the worker->host frame: tensors into ``buf`` (D2D) / safetensors (push + CPU tensors) so the pickled part is TENSOR-FREE, the rest pickled (decoded host- + side with the restricted Unpickler). Returns ``(frame_bytes, had_tensors)``; the + caller must ``cuda.synchronize()`` before sending if ``had_tensors`` (the host + clones from ``buf`` on a separate context).""" + if event is Events.EXCEPTION: + # The exception object may be an arbitrary class; ship only (type-name, + # message) so the restricted unpickler never has to resolve its class. + data = {_EXC_TAG: [type(data).__name__, str(data)]} + skel, table = pack_cuda(data, buf) + if push is not None: + pstore: dict = {} + pskel = _split_tensors(push, pstore) + if pstore and not _HAS_SAFETENSORS: + raise RuntimeError( + "safetensors is required to transmit cross_invoker tensors from an " + "isolated worker" + ) + pblob = _st_save(pstore) if pstore else b"" + else: + pskel, pblob = None, b"" + # The event rides as its string ``.value`` (not the enum) so the allowlist need + # not include the Events class. Encoding with pickle is safe — only *decoding* + # untrusted bytes is dangerous, and that is what _safe_loads restricts. + payload = pickle.dumps( + (event.value, skel, table, pskel), protocol=pickle.HIGHEST_PROTOCOL + ) + return _HEADER.pack(len(payload)) + payload + pblob, bool(table) + + +def _decode_worker_frame(raw: bytes, buf: torch.Tensor) -> tuple: + """Reverse of :func:`_encode_worker_frame`, host side. The pickled part is decoded + with the RESTRICTED unpickler (untrusted). Returns ``(event, data, push)``.""" + (plen,) = _HEADER.unpack(raw[: _HEADER.size]) + off = _HEADER.size + event_value, skel, table, pskel = _safe_loads(raw[off : off + plen]) + pblob = raw[off + plen :] + event = Events(event_value) + data = unpack_cuda(skel, table, buf) + if event is Events.EXCEPTION and isinstance(data, dict) and _EXC_TAG in data: + name, msg = data[_EXC_TAG] + data = _rebuild_exc(name, msg) + if pskel is None: + push = None + else: + ptensors = _st_load(pblob) if pblob else {} + push = _merge_tensors(pskel, ptensors) + return event, data, push + + class CudaIpcHostChannel(MediatorChannel): """Host (main-thread) end of the GPU-bounce-buffer channel. @@ -389,7 +504,7 @@ def wait_event(self) -> None: f"— worker presumed hung (e.g. an infinite loop in user code)" ) try: - event, skel, table, push = self._conn.recv() + raw = self._conn.recv_bytes() except (EOFError, OSError) as e: # The pipe broke mid-protocol => the worker died (e.g. a segfault # in user C-code, or the GPU process was OOM-killed). Surface a @@ -398,10 +513,13 @@ def wait_event(self) -> None: raise RuntimeError( "sandboxed intervention worker died during execution" ) from e + # SECURITY: the worker is untrusted — decode with the no-pickle safe + # codec (a pickle.loads here would be a host-side RCE sink). + event, data, push = _decode_worker_frame(raw, self._buf) if push is not None and self.on_push is not None: self.on_push(push) # cross_invoker: merge into the host var store self._started = True - self._pending = (event, unpack_cuda(skel, table, self._buf)) + self._pending = (event, data) self._has = True @property @@ -473,13 +591,15 @@ def __init__(self, conn: Any, buf: torch.Tensor): # --- worker -> main --- def put_event(self, item: Any) -> None: event, data = item - skel, table = pack_cuda(data, self._buf) - if table: + push = self.push_provider() if self.push_provider is not None else None + # SECURITY: encode WITHOUT pickle — the host must never unpickle untrusted + # worker data. Tensors ride the shared GPU buffer; only tagged JSON crosses. + frame, had_tensors = _encode_worker_frame(event, data, push, self._buf) + if had_tensors: # Async D2D copies must finish before the HOST (separate CUDA context) # clones them out of the buffer. See CudaIpcHostChannel.put_response. torch.cuda.synchronize() - push = self.push_provider() if self.push_provider is not None else None - self._conn.send((event, skel, table, push)) + self._conn.send_bytes(frame) # --- main -> worker --- def wait_response(self) -> None: From 151e5a5b83c90d28ecd2c45c5fae6e04184d5069 Mon Sep 17 00:00:00 2001 From: khaiwang Date: Wed, 17 Jun 2026 23:06:43 -0500 Subject: [PATCH 21/30] docs(isolation): threat-model + security-hardening decision artifact Capture the security analysis behind the GPU-worker backend: the attacker model, the asset list (host integrity / fs / net / cross-tenant host+GPU memory / DoS / deser-RCE), the R0-R4 configuration ladder with the cost coupling (closing a deeper threat forces a slower data path; the cliff is R2->R3, i.e. leaving the GPU), which sandbox controls are compatible vs incompatible with the shared-GPU CUDA-IPC design, the co-batch tenant-isolation invariant (the empty-invoke full-batch hole; the Batcher/Interleaver = tenant boundary), and the worker->host restricted-unpickler fix. Co-Authored-By: Claude Opus 4.8 --- docs/developing/mediator-threat-models.md | 193 ++++++++++++++++++++++ 1 file changed, 193 insertions(+) create mode 100644 docs/developing/mediator-threat-models.md diff --git a/docs/developing/mediator-threat-models.md b/docs/developing/mediator-threat-models.md new file mode 100644 index 000000000..2e231fad0 --- /dev/null +++ b/docs/developing/mediator-threat-models.md @@ -0,0 +1,193 @@ +# Mediator Isolation — Threat Models & Security Hardening + +**Status:** Design notes / decision artifact · **Branch:** `worktree-mediator-sandbox` +**Builds on:** [mediator-gpu-trace-integration.md](mediator-gpu-trace-integration.md) (the GPU-worker +backend + §16 triton motivation), [mediator-isolation-sandbox.md](mediator-isolation-sandbox.md) (the +earlier CPU-only threat model + AWS deployment probe). + +This page is the security companion to the GPU-worker backend: *what* adversary each isolation +configuration faces, what it can and cannot contain, and the resulting **cost coupling** that makes +the "how much isolation" question concrete. It also records the one security fix landed in-branch — the +worker→host **restricted-unpickler** codec. + +--- + +## 1. Why isolation, and why now + +`model.trace()` runs **arbitrary user intervention code**. NDIF runs that on shared GPU infrastructure, +today behind an **in-process Python import whitelist** that wraps the *entire forward pass*. That +boundary is architecturally unfixable in-process (the security suite shows 10/10 in-process escapes and +8/8 cross-tenant capability leaks succeed), **and** it cannot serve **triton-kernel models** — frontier +MoE/SSM models and anything under `torch.compile` JIT-compile triton, which needs +`subprocess`/ptxas/`tempfile`/`open` exactly when the whitelist denies them (see §16 of +[mediator-gpu-trace-integration.md](mediator-gpu-trace-integration.md)). The GPU-worker backend fixes +both at once: the **trusted forward runs host-side, unrestricted** (triton works); only the **untrusted +intervention** runs in a separate, sandboxable worker. + +--- + +## 2. Attacker model + +The attacker submits **serialized Python (cloudpickle)** that the server deserializes and executes +against a loaded model. They control: (i) the **pickle payload** (gadgets can fire at deserialize time) +and (ii) the **intervention function** — arbitrary Python at hook points. They do *not* directly control +the host process — only what they reach through the channels we expose. Intent tiers: + +- **Footgun** — accident: infinite loop, OOM, bad index/device assert, accidental host mutation. +- **Casual adversary** — `import os`, open files, sockets, object-graph walks. +- **Determined adversary** — pickle gadgets, C-extension escapes, GPU-driver/kernel exploits, + cross-tenant GPU reads, side channels. + +--- + +## 3. Assets + +| | Asset | Notes | +|---|---|---| +| **A** | Host-process integrity (no RCE on the model server) | the crown jewel | +| **B** | Host filesystem | model weights, other tenants' cached data, **secrets/IAM** | +| **C** | Network | exfiltration, internal services, **IMDS credential theft** | +| **D1** | Cross-tenant **host-memory** isolation | only under co-batching | +| **D2** | Cross-tenant **GPU-memory** isolation | the hard one | +| **E** | Availability / DoS | resource exhaustion, crash blast-radius | +| **F** | Deserialization RCE | attacker pickle firing **on the host** | + +D1/D2 only bite when **multiple tenants share a forward pass** (co-batching). + +--- + +## 4. Configuration ladder + the cost coupling + +The deep insight: **the protection level and the data-path cost are coupled** — closing a deeper threat +*forces* the activation onto a slower path. You cannot pick "max isolation, min latency." + +| Ladder | Threats newly closed | Data path | Per-hook cost | +|---|---|---|---| +| **R0 in-process** | — | tensor never moves | ~0 (≈12 ms whole gpt2 trace) | +| **R1 worker, shared GPU** | A host-process, **D1**, **F**, footgun-DoS | stays on GPU (CUDA-IPC) | **~0.6 ms, size-independent** + spawn (warm ~0.22 s) | +| **R2 + OS jail** (seccomp / net+mount+user ns + cgroups) | B fs, C net/IMDS, privilege, RAM-DoS | **same GPU buffer** | **~same ~0.6 ms** (jail is warm-time only) | +| **R3 + no GPU in worker** | **D2**, GPU-driver/kernel surface | **GPU→CPU→GPU (PCIe)** | **~ms, size-*dependent*** (1–285 ms) | +| **R4 + CPU-mem isolation** (gVisor/microVM) | kernel-proof, mem side-channels | **socket copy** (no shared mem) | R3 + serialize + pipe (slowest) | + +Two consequences: +1. **The cost cliff is R2→R3, not R1→R2.** Adding the OS jail (closes fs/net/IMDS/privilege/RAM-DoS + against even a determined adversary) rides the *same* GPU buffer and is ~free per hook. The latency + explosion only happens when you leave the GPU to close **D2**. +2. **The decision reduces to one bit: do you require D2 (cross-tenant GPU-memory isolation)?** If no → + **R2** is the target (full host-OS isolation, fast path, accept the GPU-layer residual). If yes → + you must pay **R3** for the co-batched path. + +--- + +## 5. Security controls vs the shared-GPU design + +The GPU-worker's defining feature — host and worker **share a GPU + a CUDA-IPC memory buffer** — is +itself a cross-trust channel and requires exposing the NVIDIA driver surface to the worker. Controls +split by whether they touch that path. (Lifecycle rule for all of Group A: **map device fds + init CUDA ++ map the IPC buffer + warm, THEN jail** — fds/buffer survive; mirrors the existing +lock-down-after-warm.) + +**Group A — compatible (the free wins, take these):** +- **Network namespace** (no NIC) — zero GPU impact; kills exfil + IMDS. Take unconditionally. +- **Mount namespace** (ro allowlist; bind only `/dev/nvidia*` + CUDA libs) — removes weights/secrets/ + other-tenant disk from view. Highest value. +- **User namespace** (non-root, drop caps); **cgroups** (CPU/mem → OOM/forkbomb cap); **seccomp** (have it). + +**Group B — in tension / incompatible (forces the CPU-transport fallback or hardware partition):** +- **PID/IPC namespace** — cross-namespace CUDA-IPC is documented to fail; `CLONE_NEWPID` must be entered + before CUDA init, colliding with map-first. Likely cannot PID-isolate while sharing the buffer. +- **gVisor** — `nvproxy` implements a narrow NVIDIA ioctl subset; CUDA-IPC unconfirmed → likely drops to + CPU transport. +- **Firecracker microVM** — GPU only via exclusive passthrough; can't share one GPU. +- **MIG** — the only real GPU-memory isolation, but breaks IPC sharing by construction (and unavailable + on the T4 prod GPU). + +**Residual that no OS control closes:** even with all of Group A, the worker holds a CUDA context on a +GPU also hosting the model + (under batching) other tenants' KV-cache → GPU/driver-layer cross-tenant +surface (driver bugs, IPC-handle reachability, memory-residency side channels). Only **R3 (no GPU in +worker)** or **MIG** removes it. + +--- + +## 6. Co-batch tenant isolation + +With co-batching in scope, "can a co-batched user reach another's data / entire workflow?" resolves as +follows (verified in code). + +**Row-bounding is host-authoritative.** `batch_group` is computed host-side at invoke construction +(`tracing/invoker.py:81`); the host narrows reads to `self.batch_group` (`interleaver.py:1245`) and +bounds writes (`:1259`/`:1294`). The worker only names a `(path, kind, iteration)` string — it cannot +widen. Each user's intervention runs in its own worker **process** (separate address space). + +**The one data-plane hole.** An **empty invoke** gets `batch_group=None`, and `narrow(None)` returns the +**entire batch** while `swap(None,…)` overwrites it (`batching.py:214/217/244/248` — `None`, `[-1]`, and +`not needs_batching` are all full-batch sentinels). So a `batch_group=None` mediator is a full-batch +read+write primitive across all co-batched tenants — the in-process suite's "`swap(None)` poisons all +tenants" leak. + +**Enforce, don't guard.** `batch_group` is 100% host-computed (no user-supplied field), so this is a +**system invariant**, not input validation — but the full-batch path is reachable via the user's +*program structure* (an empty invoke is a legit single-tenant feature). The correct enforcement is +architectural: **the Batcher/Interleaver instance is the tenant boundary** — co-batch *below* the nnsight +Batcher (at the inference engine, e.g. vLLM continuous batching), never *inside* one shared Batcher. +Under that invariant all three sentinels are tenant-local and safe by construction, and +`cross_invoker`/`barrier` (per-Interleaver `_xinvoke_store`/`_barrier_acc`) close with it. The vLLM path +gets this for free (independent per-request interleavers); an HF-merged-batch path must enforce it. A +runtime reject-`None` check is defense-in-depth against our own merge bug, not the primary control. + +**What a co-batched user CAN / CANNOT do.** *Cannot* (with the empty-invoke invariant + normal CUDA): +read another tenant's rows/activations (narrowed out, never sent), their prompt/code/output (separate +process), or their **live** GPU memory by pointer (separate CUDA contexts are VA/page-table isolated). So +reconstructing another user's **entire workflow is structurally closed**. *Can*, on the shared GPU +(R1/R2): scoop **uninitialized/freed GPU memory** (untargeted *fragments*, not workflows) + coarse +timing side channels — reducible on the fast path via **zero-on-free memory scrubbing** + alloc caps + +driver capability segmentation, without going to R3. + +--- + +## 7. Worker→host RCE — finding and fix (landed in-branch) + +**Finding.** The host read worker frames with `mp.Pipe.recv()` (`transport.py`), which **unpickles**. +The worker runs attacker code, so it could send a crafted pickle and get **RCE on the trusted host** — +which bypasses *every* other layer (seccomp, namespaces, row-bounding). The design fixed deserialization +for the *inbound* user payload (unpickled inside the worker) but not the *outbound* results. This is a +precondition for the isolation guarantee holding against a determined adversary. + +**Fix.** Worker→host now never plain-`pickle.loads`. Tensors already travel out-of-band (GPU bounce +buffer / safetensors), so a frame is **tensor-free**; the remaining small structure is decoded with a +**restricted unpickler** (`transport._RestrictedUnpickler` / `_safe_loads`) whose `find_class` allows +**only torch `dtype`/`device`** and refuses every other class/function. `find_class` resolves a global +*before* the `REDUCE` that would call it, so a gadget (`os.system`, …) is refused before it can execute. +Supporting choices keep the allowlist minimal: the event rides as its string `.value` (no enum class), +and exceptions ride as a `(type-name, message)` sentinel (no class). **Host→worker stays normal pickle** +(host-authored, trusted). + +- **Why a restricted unpickler over a hand-rolled codec:** pickle handles all plain nested Python + structures natively (no per-type enumeration — an earlier hand-rolled JSON codec silently *missed* the + `Events.CACHE` event's `torch.dtype`/`torch.device` payload). Anything un-allowlisted now fails **loud + at decode with the exact refused class name**, easy to extend if legitimately safe. +- **Cost:** ~plain-pickle speed (µs for the tiny tensor-free frame); `find_class` fires only on the + handful of globals per frame. Negligible vs the ~0.6 ms/hook round-trip. +- **Why not `torch.load(weights_only=True)`:** it had an RCE bypass (CVE-2025-32434, fixed only in 2.6); + our unpickler is *tighter* — it enables no tensor-rebuild path at all (tensors aren't in the pickle). +- **Capability narrowing (documented):** `.save()` of an arbitrary object / numpy array / framework type + (e.g. `ModelOutput`) is no longer transmittable from a worker — save a tensor (or basic data) instead. + +Coverage: `prototypes/mediator-sandbox/gpu_sandbox/test_isolated_codec_security.py` — fidelity for +VALUE/SWAP/END/**CACHE (dtype+device)**/EXCEPTION/push, and a real `__reduce__` gadget refused at decode +without executing. (CPU is enough; needs torch.) The legacy AF_UNIX socket channels +(`SocketHostChannel`/`ShmSocketHostChannel`) are unused by `isolate_mediators` and still plain-unpickle — +route them through `_safe_loads` before wiring them to an untrusted worker. + +--- + +## 8. Open items / next steps + +- **Group-A namespace jail** behind a pluggable `Jailer` (applied post-warm/pre-`ready`) — the + highest-value, design-preserving hardening; R2 on the fast path. +- **Co-batch invariant** in the (future) HF-merged-batch integration: keep the nnsight Batcher + per-tenant; add a multi-tenant tripwire rejecting unbounded `batch_group`. (Single-trace=single-Batcher + already holds it today; an unconditional guard would break legit single-user empty invokes.) +- **GPU memory scrubbing** (zero-on-free) for the shared-GPU residual. +- **gVisor + CUDA-IPC empirical test** — decides whether any R3-strong path keeps GPU sharing. +- **Run the codec + isolated suite** on a torch/GPU box to confirm fidelity + bit-identical regression. From 043196c5ce259273fa68b5982ad72474ddc5c73e Mon Sep 17 00:00:00 2001 From: khaiwang Date: Fri, 19 Jun 2026 23:13:13 -0400 Subject: [PATCH 22/30] feat(intervention): replace worker->host pickle with a closed value-algebra codec MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The worker->host (untrusted) direction no longer runs pickle's VM on worker bytes. Instead the boundary transmits a CLOSED VALUE ALGEBRA — never live objects — so there is no opcode that can call a function in the first place: Value = None | bool | int | float | str | bytes | list | tuple | dict | set of Value | torch.dtype | torch.device | Array(...) (torch tensors AND numpy arrays, OUT-OF-BAND) This makes the boundary value-semantic, and makes "safe + correct" a property of the type rather than a bet on a restricted unpickler: - SAFE by construction: _codec_loads is pure data assembly (no globals, no find_class, no REDUCE), plus a size cap + bounds-checked reads for decode-bomb DoS. The previous restricted unpickler is removed entirely. - FAITHFUL: pack_cuda's Array leaf is generalized from torch.is_tensor to also cover numpy.ndarray (bridged through torch, re-materialized host-side as an ndarray), so numpy `.save()`s now cross — they were silently refused before. - HONEST contract: a value outside the algebra (custom object / framework type) is rejected at the WORKER, at ENCODE, with a clear BoundaryValueError naming it at its source — not an encode-ok / decode-refuse split. tracer.cache() shipped a live CacheDict placeholder, which is not a value; the worker now ships its token as a `{_ISO_CACHE_TAG: token}` marker (same shape as the EXCEPTION sentinel) and the host swaps in its forward-filled cache by token. (This also fixes cache under isolation, which the merged restricted unpickler had broken — it refused CacheDict the same way.) host->worker stays plain pickle (host-authored, trusted). Verified (nnsight-tf: py3.11 / torch 2.11 / transformers 5.12): - codec unit (test_isolated_codec_security.py): fidelity over the algebra incl. numpy + dtype/device; a __reduce__ gadget and a custom object rejected at encode before any __reduce__ runs; malformed/oversized/unknown-tag raise cleanly. - full isolated GPU suite bit-identical max|Δ|=0: trace, unembed, steer, cache, backward, multitoken-iter, cross-invoke, pool, acceptance. Co-Authored-By: Claude Opus 4.8 (1M context) --- .../test_isolated_codec_security.py | 138 +++++--- src/nnsight/intervention/interleaver.py | 24 +- src/nnsight/intervention/isolation.py | 9 +- src/nnsight/intervention/transport.py | 305 ++++++++++++++---- 4 files changed, 356 insertions(+), 120 deletions(-) diff --git a/prototypes/mediator-sandbox/gpu_sandbox/test_isolated_codec_security.py b/prototypes/mediator-sandbox/gpu_sandbox/test_isolated_codec_security.py index ba0c133d9..690fb157c 100644 --- a/prototypes/mediator-sandbox/gpu_sandbox/test_isolated_codec_security.py +++ b/prototypes/mediator-sandbox/gpu_sandbox/test_isolated_codec_security.py @@ -1,36 +1,36 @@ #!/usr/bin/env python3 -"""Worker->host frame codec — security + fidelity (restricted unpickler). - -The isolated worker runs UNTRUSTED user code; the host must not *plain* ``pickle.loads`` -its frames (a ``__reduce__`` gadget would execute on the trusted host = RCE). Worker->host -frames are tensor-free (tensors ride the GPU buffer / safetensors) and the rest is decoded -with ``transport._RestrictedUnpickler`` (``find_class`` allows ONLY torch dtype/device). - - fidelity - VALUE / SWAP / END / EXCEPTION / cross_invoker-push AND a tracer.cache()-style - spec carrying ``torch.dtype`` + ``torch.device`` all round-trip exactly. - security - a frame whose data is a ``__reduce__`` gadget is REFUSED at decode - (``find_class`` rejects ``os.system`` before the REDUCE could call it) — the - gadget never executes on the host. A non-allowlisted class is also refused. - -Needs torch; **CPU is enough — no CUDA required** (the bounce buffer is just bytes here). +"""Worker->host frame codec — fidelity + security (closed value algebra, no pickle VM). + +The isolated worker runs UNTRUSTED user code, so the host must never run pickle's VM on +its frames (a ``__reduce__`` gadget would execute on the trusted host = RCE). Instead the +boundary transmits a CLOSED VALUE ALGEBRA (``transport._codec_dumps`` / ``_codec_loads``): +None/bool/int/float/str/bytes, list/tuple/dict/set, torch dtype/device, and out-of-band +ARRAY leaves (torch tensors AND numpy arrays ride the GPU buffer / safetensors). There is no +opcode that can call a function, so decoding is pure data assembly. + + fidelity - VALUE / SWAP / END / EXCEPTION / cross_invoker-push, a tracer.cache()-style + spec carrying torch.dtype + torch.device, AND a numpy array all round-trip + exactly (numpy re-materializes as ndarray). + security - a value OUTSIDE the algebra (a __reduce__ gadget, a custom class) is rejected + at ENCODE, in the worker, before any pickle/__reduce__ runs — never an + encode-ok / decode-refuse split. + robustness - malformed / oversized / unknown-tag bytes raise BoundaryDecodeError host-side + (no crash, no hang), not a silent or VM-level failure. + +Needs torch + numpy; **CPU is enough — no CUDA required** (the bounce buffer is just bytes). Run: - PYTHONPATH=src python \ + PYTHONPATH=src /disk/u/zikai/anaconda3/envs/nnsight-tf/bin/python \ prototypes/mediator-sandbox/gpu_sandbox/test_isolated_codec_security.py """ import os -import pickle import sys +import numpy as np import torch from nnsight.intervention import transport as T from nnsight.intervention.interleaver import Events - - -# Module-level (a function-local class can't be pickled even to ENCODE the frame, which -# would mask the decode-time refusal we mean to test). Harmless, but not on the allowlist. -class _NonAllowlisted: - pass +from nnsight.intervention.transport import BoundaryDecodeError, BoundaryValueError def _roundtrip(event, data, push=None, nbytes=1 << 20): @@ -47,22 +47,23 @@ def test_fidelity(): ev, out, _ = _roundtrip(Events.VALUE, "transformer.h.6.output.i0") ok &= ev is Events.VALUE and out == "transformer.h.6.output.i0" - # SWAP: (requester, value) carrying a tensor. + # SWAP: (requester, value) carrying a tensor + a None tail (the attn-output shape). t = torch.randn(2, 3) - ev, out, _ = _roundtrip(Events.SWAP, ("h.6.output", t)) - ok &= ev is Events.SWAP and out[0] == "h.6.output" and torch.equal(out[1], t) + ev, out, _ = _roundtrip(Events.SWAP, ("h.6.output", (t, None))) + ok &= ev is Events.SWAP and out[0] == "h.6.output" + ok &= torch.equal(out[1][0], t) and out[1][1] is None - # END: saved dict — tensor + scalars + nested list/tuple + tuple-with-None. + # END: saved dict — tensor + scalars + bytes + set + nested list/tuple. saved = { "x": torch.arange(6).float(), - "n": 3, - "f": 1.5, + "n": 3, "f": 1.5, "b": b"hi", "set": {1, 2, 3}, "lst": [1, (2, 3)], "tup": (torch.ones(2), None), } ev, out, _ = _roundtrip(Events.END, saved) ok &= ev is Events.END ok &= torch.equal(out["x"], saved["x"]) and out["n"] == 3 and out["f"] == 1.5 + ok &= out["b"] == b"hi" and out["set"] == {1, 2, 3} ok &= out["lst"] == [1, (2, 3)] and type(out["lst"][1]) is tuple ok &= type(out["tup"]) is tuple and torch.equal(out["tup"][0], torch.ones(2)) ok &= out["tup"][1] is None @@ -75,16 +76,23 @@ def test_fidelity(): ok &= out[2] == torch.device("cpu") and out[3] is torch.float16 ok &= out[0] == 12345 and out[1] == ["transformer.h.0"] + # NUMPY: a saved ndarray crosses out-of-band (the Array leaf) and comes back ndarray. + # (A numpy *scalar* is not an ndarray, so it's not in the algebra — save a python value.) + arr = np.arange(12, dtype=np.float32).reshape(3, 4) + ev, out, _ = _roundtrip(Events.END, {"np": arr}) + ok &= isinstance(out["np"], np.ndarray) and out["np"].dtype == np.float32 + ok &= np.array_equal(out["np"], arr) + # EXCEPTION: rebuilt host-side from (type-name, message); no object crosses. ev, out, _ = _roundtrip(Events.EXCEPTION, ValueError("boom")) ok &= ev is Events.EXCEPTION and isinstance(out, ValueError) and "boom" in str(out) - # cross_invoker push: CPU tensors + scalars ride safetensors, not pickle. + # cross_invoker push: CPU tensors + scalars ride safetensors, not the value codec. push = {"shared": torch.randn(4), "k": 7} ev, _out, pout = _roundtrip(Events.END, {"y": torch.zeros(1)}, push=push) ok &= pout is not None and torch.equal(pout["shared"], push["shared"]) and pout["k"] == 7 - print(f"[fidelity] VALUE/SWAP/END/CACHE(dtype,device)/EXCEPTION/push exact: {ok}") + print(f"[fidelity] VALUE/SWAP/END/CACHE(dtype,device)/numpy/EXCEPTION/push exact: {ok}") return ok @@ -93,43 +101,75 @@ def test_security(): if os.path.exists(probe): os.remove(probe) - class Bomb: + class Bomb: # a classic __reduce__ gadget def __reduce__(self): return (os.system, (f"echo pwned > {probe}",)) + class Plain: # a harmless but non-algebra custom object + pass + buf = torch.empty(1 << 16, dtype=torch.uint8) - # The worker CAN pickle a gadget (encoding is safe — nothing runs). The host must - # REFUSE it at decode (find_class rejects os.system) before the gadget executes. - frame, _had = T._encode_worker_frame(Events.END, {"evil": Bomb()}, None, buf) - refused = False + # The gadget is OUTSIDE the value algebra, so the codec refuses it at ENCODE, in the + # worker — __reduce__ is never even consulted, nothing reaches the host. + gadget_rejected = False try: - T._decode_worker_frame(frame, buf) - except pickle.UnpicklingError: - refused = True + T._encode_worker_frame(Events.END, {"evil": Bomb()}, None, buf) + except BoundaryValueError: + gadget_rejected = True no_pwn = not os.path.exists(probe) if os.path.exists(probe): os.remove(probe) - # A non-allowlisted (but harmless) class is also refused — the allowlist is tight. - frame2, _ = T._encode_worker_frame(Events.END, {"obj": _NonAllowlisted()}, None, buf) - refused2 = False + # A harmless custom object is refused the same way — the algebra is closed, not a denylist. + plain_rejected = False try: - T._decode_worker_frame(frame2, buf) - except pickle.UnpicklingError: - refused2 = True + T._encode_worker_frame(Events.END, {"obj": Plain()}, None, buf) + except BoundaryValueError: + plain_rejected = True - ok = refused and no_pwn and refused2 - print(f"[security] gadget refused at decode + never executed: {refused and no_pwn} | " - f"non-allowlisted class refused: {refused2}") + ok = gadget_rejected and no_pwn and plain_rejected + print(f"[security] gadget rejected at encode (never ran): {gadget_rejected and no_pwn} | " + f"custom object rejected at encode: {plain_rejected}") + return ok + + +def test_robustness(): + # The host decodes UNTRUSTED bytes: malformed / oversized / unknown-tag must raise a + # clean BoundaryDecodeError, never crash or hang. + buf = torch.empty(1 << 16, dtype=torch.uint8) + from nnsight.intervention.transport import _codec_loads, _MAX_CODEC_BYTES + + cases = { + "truncated": bytes([T._T_LIST, 0xFF, 0xFF, 0x01]), # claims items, no data + "unknown_tag": bytes([200]), # tag not in the algebra + "trailing": bytes([T._T_NONE, T._T_NONE]), # two values, one expected + "bad_varint": bytes([T._T_STR] + [0x80] * 12), # runaway length varint + "oversized": b"\x00" * (_MAX_CODEC_BYTES + 1), # past the size cap + } + ok = True + for name, payload in cases.items(): + raised = False + try: + _codec_loads(payload) + except BoundaryDecodeError: + raised = True + except Exception as e: # any other exception = not a clean rejection + print(f" [{name}] WRONG exception {type(e).__name__}") + ok &= raised + print(f"[robustness] malformed/oversized/unknown-tag all raise BoundaryDecodeError: {ok}") return ok def main(): - results = {"fidelity": test_fidelity(), "security": test_security()} + results = { + "fidelity": test_fidelity(), + "security": test_security(), + "robustness": test_robustness(), + } ok = all(results.values()) print("=" * 72) - print(f"ISOLATED CODEC SECURITY: {'PASS' if ok else 'FAIL'} — {results}") + print(f"ISOLATED CODEC (value algebra): {'PASS' if ok else 'FAIL'} — {results}") sys.exit(0 if ok else 1) diff --git a/src/nnsight/intervention/interleaver.py b/src/nnsight/intervention/interleaver.py index 8e6f9a67d..8401c88c3 100755 --- a/src/nnsight/intervention/interleaver.py +++ b/src/nnsight/intervention/interleaver.py @@ -360,6 +360,14 @@ class Events(Enum): UNEMBED = "unembed" # Run norm+unembed on the host's real weights (isolation) +# A ``tracer.cache()`` placeholder is a custom ``CacheDict`` object that is NOT in the +# worker->host value algebra (only data crosses, never live objects). The worker ships its +# token as a plain-dict marker ``{_ISO_CACHE_TAG: token}`` instead; the host swaps in its +# own forward-filled CacheDict by token in ``handle_end_event``. (Same shape as the +# EXCEPTION ``(type, message)`` sentinel — a value standing in for an object.) +_ISO_CACHE_TAG = "__nnsight_iso_cache__" + + class Cancelation(Exception): """Exception raised when a request is canceled.""" @@ -1640,18 +1648,14 @@ def handle_end_event(self, saves: Optional[Any] = None): self._iso.clean = True if saves: if self._iso_caches: - # tracer.cache(): the worker shipped an EMPTY placeholder CacheDict; - # swap in the HOST cache (matched by token), which the forward fills - # in-place after this injection. The user's variable then IS the - # forward-filled host cache. - from .tracing.tracer import Cache - + # tracer.cache(): the worker shipped a ``{_ISO_CACHE_TAG: token}`` + # marker in place of its placeholder CacheDict; swap in the HOST cache + # (matched by token), which the forward fills in-place after this + # injection. The user's variable then IS the forward-filled host cache. saves = dict(saves) for name, val in list(saves.items()): - if isinstance(val, Cache.CacheDict): - host_cache = self._iso_caches.get( - getattr(val, "_iso_cache_token", None) - ) + if isinstance(val, dict) and _ISO_CACHE_TAG in val: + host_cache = self._iso_caches.get(val[_ISO_CACHE_TAG]) if host_cache is not None: saves[name] = host_cache.cache tracer = self.interleaver.tracer diff --git a/src/nnsight/intervention/isolation.py b/src/nnsight/intervention/isolation.py index d222206a7..cbab2dfbb 100644 --- a/src/nnsight/intervention/isolation.py +++ b/src/nnsight/intervention/isolation.py @@ -797,12 +797,19 @@ def end(self): # Worker→host saves transmission: bundle .save()'d values into the END event. # The intervention's compiled body calls ``end()`` on success; push() populates # the SerializedFrame's f_locals, which we filter by Globals.saves. - from .interleaver import Events + from .interleaver import Events, _ISO_CACHE_TAG from .tracing.globals import Globals + from .tracing.tracer import Cache self.push() flocals = self.info.frame.f_locals saved = {k: v for k, v in flocals.items() if id(v) in Globals.saves} + # A tracer.cache() placeholder CacheDict is a live object, not a value — ship its + # token as a plain-dict marker so the value codec accepts it; the host swaps in its + # own forward-filled cache by token (top-level saves only, as the host handler is). + for k, v in list(saved.items()): + if isinstance(v, Cache.CacheDict): + saved[k] = {_ISO_CACHE_TAG: getattr(v, "_iso_cache_token", None)} self.channel.put_event((Events.END, saved)) def exception(self, exception: Exception): diff --git a/src/nnsight/intervention/transport.py b/src/nnsight/intervention/transport.py index 6e774a1bb..dc291496d 100644 --- a/src/nnsight/intervention/transport.py +++ b/src/nnsight/intervention/transport.py @@ -26,14 +26,15 @@ 1. **Security:** the *worker->host* direction MUST NOT plain-``pickle.loads`` untrusted worker bytes (a ``__reduce__`` gadget would be a host-side RCE). **Done for the active GPU path** (:class:`CudaIpcHostChannel` / :class:`CudaIpcWorkerChannel`): worker->host - frames are tensor-free (tensors ride the GPU buffer / safetensors) and the small - remaining structure is decoded with the **restricted unpickler** - (:class:`_RestrictedUnpickler` / :func:`_safe_loads`) whose ``find_class`` allows ONLY - torch dtype/device; host->worker stays a normal pickle (host-authored, trusted). The - legacy AF_UNIX socket channels below (:class:`SocketHostChannel`, - :class:`ShmSocketHostChannel`) are NOT used by ``isolate_mediators`` and still - plain-``pickle.loads`` — route them through :func:`_safe_loads` before wiring them to an - untrusted worker. + frames carry array data out-of-band (the GPU buffer / safetensors) and the small + remaining structure is a **closed value algebra** decoded WITHOUT a pickle VM + (:func:`_codec_dumps` / :func:`_codec_loads`) — only primitives / containers / torch + dtype-device / out-of-band array headers cross; anything else is rejected at the worker + at encode (:class:`BoundaryValueError`). host->worker stays a normal pickle + (host-authored, trusted). The legacy AF_UNIX socket channels below + (:class:`SocketHostChannel`, :class:`ShmSocketHostChannel`) are NOT used by + ``isolate_mediators`` and still plain-``pickle.loads`` — route them through the value + codec before wiring them to an untrusted worker. 2. **Performance (measured):** ``pickle`` of a torch tensor is the per-hook bottleneck — ``dumps``+``loads`` ~22 ms per direction at 16.8 MB and **superlinear**; @@ -46,7 +47,6 @@ from __future__ import annotations import builtins -import io import mmap import os import pickle @@ -66,6 +66,14 @@ except Exception: # pragma: no cover _HAS_SAFETENSORS = False +try: + import numpy as _np + + _HAS_NUMPY = True +except Exception: # pragma: no cover + _np = None + _HAS_NUMPY = False + _HEADER = struct.Struct("!I") # 4-byte big-endian length prefix @@ -313,9 +321,24 @@ def pack_cuda(value: Any, buf: torch.Tensor) -> tuple: state = {"offset": 0} def walk(obj: Any) -> Any: + # The "Array" leaf of the boundary value algebra: any contiguous typed buffer — + # a torch.Tensor OR a numpy.ndarray — travels OUT-OF-BAND in ``buf``; only a small + # ``(offset, nbytes, shape, dtype, kind)`` header rides in the skeleton. numpy is + # bridged through torch (zero-copy view); the host re-materializes the right kind. + t = kind = None if torch.is_tensor(obj): + t, kind = obj.detach().contiguous(), "torch" + elif _HAS_NUMPY and isinstance(obj, _np.ndarray): + try: + t = torch.from_numpy(_np.ascontiguousarray(obj)) + except (TypeError, ValueError) as e: + raise BoundaryValueError( + f"a numpy array of dtype {obj.dtype} is not transmittable across the " + f"isolation boundary (torch has no matching dtype): {e}" + ) + kind = "numpy" + if t is not None: i = len(table) - t = obj.detach().contiguous() flat = t.reshape(-1).view(torch.uint8) n = int(flat.numel()) offset = (state["offset"] + 15) & ~15 # 16-byte align @@ -326,7 +349,7 @@ def walk(obj: Any) -> Any: ) if n: buf[offset : offset + n].copy_(flat) - table[str(i)] = (offset, n, tuple(t.shape), str(t.dtype)) + table[str(i)] = (offset, n, tuple(t.shape), str(t.dtype), kind) state["offset"] = offset + n return {_TENSOR_TAG: i} if type(obj) is tuple: @@ -346,53 +369,215 @@ def unpack_cuda(skel: Any, table: dict, buf: torch.Tensor) -> Any: later reuse of the single buffer can't corrupt it — the clone-on-receive rule) and re-injected into the skeleton.""" tensors: dict = {} - for k, (offset, n, shape, dtype_str) in table.items(): + for k, meta in table.items(): + offset, n, shape, dtype_str = meta[0], meta[1], meta[2], meta[3] + kind = meta[4] if len(meta) > 4 else "torch" dtype = _dtype_from_str(dtype_str) view = buf[offset : offset + n].view(dtype).reshape(shape) - tensors[k] = view.clone() + # CLONE out of the shared single-slot buffer (so a later reuse can't corrupt it); + # a numpy "Array" leaf is re-materialized on the CPU as an ndarray. + tensors[k] = view.clone().cpu().numpy() if kind == "numpy" else view.clone() return _merge_tensors(skel, tensors) # --------------------------------------------------------------------------- # -# Safe worker->host frame codec — restricted unpickler, tight allowlist # +# Worker->host frame codec — closed value algebra, no pickle VM # # --------------------------------------------------------------------------- # -# The worker runs UNTRUSTED user code, so the host MUST NOT do a plain -# ``pickle.loads`` of its frames (a ``__reduce__`` gadget = host RCE). Tensors -# travel out-of-band (the GPU bounce buffer / safetensors), so a worker->host -# frame is TENSOR-FREE and carries only plain data plus, at most, a torch -# ``dtype`` / ``device`` (the ``tracer.cache()`` spec). We decode it with a -# restricted Unpickler whose ``find_class`` allows ONLY torch dtype/device and -# refuses every other class/function. ``find_class`` is consulted to resolve a -# global BEFORE the ``REDUCE`` opcode could call it, so a gadget (``os.system`` -# etc.) is refused before it can execute. The event crosses as its string -# ``.value`` (no enum class) and exceptions as a (type-name, message) sentinel -# (no class), so the allowlist stays just {torch dtype, torch device}. -# (host->worker stays a normal pickle: that direction is host-authored, trusted.) +# The worker runs UNTRUSTED user code, so the host must never run pickle's VM on +# its frames (a ``__reduce__`` gadget = host RCE). Instead the boundary transmits +# a CLOSED VALUE ALGEBRA — never live objects — so there is no opcode that can +# call a function in the first place: +# +# Value = None | bool | int | float | str | bytes (primitives) +# | list | tuple | dict | set of Value (containers) +# | torch.dtype | torch.device (value leaves) +# | Array(...) -> {_TENSOR_TAG: i} (bulk, OUT-OF-BAND) +# +# Array leaves (torch tensors, numpy arrays) are already pulled out-of-band by +# ``pack_cuda``/``_split_tensors`` into the GPU buffer / safetensors, leaving only +# a ``{_TENSOR_TAG: i}`` header in the skeleton — itself plain data. So this codec +# only has to (de)serialize the algebra above; anything outside it is rejected at +# the WORKER, at ENCODE, with a clear message (``BoundaryValueError``) — never an +# encode-ok / decode-refuse split. Decoding is pure data assembly: no globals, no +# ``find_class``, no ``REDUCE``. A size cap + bounds-checked reads contain a +# decode bomb (DoS). (host->worker stays a normal pickle: host-authored, trusted.) _EXC_TAG = "__nnsight_iso_exc__" +# Max bytes for the codec part of a worker->host frame (the structured envelope; +# bulk array data is out-of-band, so this stays tiny). Caps decode-time allocation. +_MAX_CODEC_BYTES = 64 << 20 + +# Value-algebra tags (one byte each). +_T_NONE, _T_FALSE, _T_TRUE, _T_INT, _T_FLOAT, _T_STR, _T_BYTES = range(7) +_T_LIST, _T_TUPLE, _T_DICT, _T_SET, _T_DTYPE, _T_DEVICE = range(7, 13) + + +class BoundaryValueError(TypeError): + """A value outside the boundary algebra was handed to the worker->host codec + (e.g. a custom object / framework type ``.save()``-d in the worker). Raised at + ENCODE, in the worker, so the failure names the offending value at its source.""" + + +class BoundaryDecodeError(ValueError): + """A worker->host frame was malformed, oversized, or carried an unknown tag. + Raised at DECODE, host-side, instead of trusting attacker-shaped bytes.""" + + +def _w_uvarint(out: bytearray, n: int) -> None: + while True: + b = n & 0x7F + n >>= 7 + out.append(b | 0x80 if n else b) + if not n: + return + + +def _codec_dumps(value: Any) -> bytes: + out = bytearray() + + def enc(v: Any) -> None: + # bool BEFORE int (bool is a subclass of int); identity check is exact. + if v is None: + out.append(_T_NONE) + elif v is True: + out.append(_T_TRUE) + elif v is False: + out.append(_T_FALSE) + elif type(v) is int: + out.append(_T_INT) + b = v.to_bytes((v.bit_length() + 8) // 8 or 1, "little", signed=True) + _w_uvarint(out, len(b)) + out.extend(b) + elif type(v) is float: + out.append(_T_FLOAT) + out.extend(struct.pack("host frames. Reconstructs only torch - dtype/device globals; every other class/function is refused — so no gadget - callable is ever resolved and the ``REDUCE`` that would call it never runs.""" - - def find_class(self, module: str, name: str): - if module == "torch": - obj = getattr(torch, name, None) - if obj is torch.device or isinstance(obj, torch.dtype): - return obj - raise pickle.UnpicklingError( - f"refusing to unpickle {module}.{name} from an isolated worker — only " - f"tensors (out-of-band) + basic data types cross the isolation boundary " - f"(a torch dtype/device is allowed; a custom object / numpy / framework " - f"type is not transmittable from a worker)" - ) + enc(value) + return bytes(out) -def _safe_loads(data: bytes) -> Any: - """Unpickle UNTRUSTED worker bytes with the restricted (allowlist) unpickler.""" - return _RestrictedUnpickler(io.BytesIO(data)).load() +def _codec_loads(data: bytes) -> Any: + if len(data) > _MAX_CODEC_BYTES: + raise BoundaryDecodeError( + f"worker->host frame {len(data)} B exceeds the {_MAX_CODEC_BYTES} B cap" + ) + pos = 0 + n = len(data) + + def need(k: int) -> int: + nonlocal pos + if pos + k > n: + raise BoundaryDecodeError("truncated worker->host frame") + start = pos + pos += k + return start + + def ruvarint() -> int: + nonlocal pos + result = shift = 0 + while True: + if pos >= n: + raise BoundaryDecodeError("truncated varint") + b = data[pos] + pos += 1 + result |= (b & 0x7F) << shift + if not b & 0x80: + return result + shift += 7 + if shift > 63: # lengths/counts fit in 64 bits — reject a runaway varint + raise BoundaryDecodeError("varint too long") + + def dec() -> Any: + nonlocal pos + if pos >= n: + raise BoundaryDecodeError("truncated worker->host frame") + tag = data[pos] + pos += 1 + if tag == _T_NONE: + return None + if tag == _T_TRUE: + return True + if tag == _T_FALSE: + return False + if tag == _T_INT: + k = ruvarint() + return int.from_bytes(data[need(k) : pos], "little", signed=True) + if tag == _T_FLOAT: + return struct.unpack("host frame") + return value def _rebuild_exc(name: str, msg: str) -> BaseException: @@ -408,14 +593,15 @@ def _rebuild_exc(name: str, msg: str) -> BaseException: def _encode_worker_frame(event: Events, data: Any, push: Any, buf: torch.Tensor) -> tuple: - """Build the worker->host frame: tensors into ``buf`` (D2D) / safetensors (push - CPU tensors) so the pickled part is TENSOR-FREE, the rest pickled (decoded host- - side with the restricted Unpickler). Returns ``(frame_bytes, had_tensors)``; the - caller must ``cuda.synchronize()`` before sending if ``had_tensors`` (the host - clones from ``buf`` on a separate context).""" + """Build the worker->host frame: array data (torch/numpy) into ``buf`` (D2D) / + safetensors (push CPU tensors) so the structured part is ARRAY-FREE, the rest + encoded with the closed value-algebra codec (decoded host-side without a pickle + VM). Returns ``(frame_bytes, had_tensors)``; the caller must ``cuda.synchronize()`` + before sending if ``had_tensors`` (the host clones from ``buf`` on a separate + context).""" if event is Events.EXCEPTION: # The exception object may be an arbitrary class; ship only (type-name, - # message) so the restricted unpickler never has to resolve its class. + # message) — both strings, in the value algebra — so no class crosses. data = {_EXC_TAG: [type(data).__name__, str(data)]} skel, table = pack_cuda(data, buf) if push is not None: @@ -429,21 +615,20 @@ def _encode_worker_frame(event: Events, data: Any, push: Any, buf: torch.Tensor) pblob = _st_save(pstore) if pstore else b"" else: pskel, pblob = None, b"" - # The event rides as its string ``.value`` (not the enum) so the allowlist need - # not include the Events class. Encoding with pickle is safe — only *decoding* - # untrusted bytes is dangerous, and that is what _safe_loads restricts. - payload = pickle.dumps( - (event.value, skel, table, pskel), protocol=pickle.HIGHEST_PROTOCOL - ) + # The event rides as its string ``.value`` (not the enum). The codec only emits the + # value algebra; a value outside it raises ``BoundaryValueError`` HERE, in the + # worker, naming the offending object at its source. + payload = _codec_dumps((event.value, skel, table, pskel)) return _HEADER.pack(len(payload)) + payload + pblob, bool(table) def _decode_worker_frame(raw: bytes, buf: torch.Tensor) -> tuple: - """Reverse of :func:`_encode_worker_frame`, host side. The pickled part is decoded - with the RESTRICTED unpickler (untrusted). Returns ``(event, data, push)``.""" + """Reverse of :func:`_encode_worker_frame`, host side. The structured part is decoded + with the value-algebra codec (no pickle VM on untrusted bytes). Returns + ``(event, data, push)``.""" (plen,) = _HEADER.unpack(raw[: _HEADER.size]) off = _HEADER.size - event_value, skel, table, pskel = _safe_loads(raw[off : off + plen]) + event_value, skel, table, pskel = _codec_loads(raw[off : off + plen]) pblob = raw[off + plen :] event = Events(event_value) data = unpack_cuda(skel, table, buf) From 582304153bc72afe42e10d16d7b96bf22eababb6 Mon Sep 17 00:00:00 2001 From: khaiwang Date: Sat, 20 Jun 2026 13:25:45 -0400 Subject: [PATCH 23/30] =?UTF-8?q?feat(intervention):=20tracer.patch=20+=20?= =?UTF-8?q?tracer.ablate=20=E2=80=94=20boundary-write=20primitives=20for?= =?UTF-8?q?=20the=20isolated=20tier?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The two remaining single-write part-2 primitives, structural twins of tracer.steer: both are replacement boundary writes that touch no host weights, so they ride the existing Events.SWAP with no new event, no host handler, and no isolated/in-process branch. Done in place under isolation each silently no-ops (the worker mutates its delivered clone, no SWAP fires); the replacement swap makes them cross the boundary by construction. - tracer.patch(envoy, value): transplant a precomputed value into a module's output residual (activation patching / resampling). Cast to the residual's dtype/device so a value precomputed on CPU — the isolation case — transplants cleanly. Whole-tuple replacement (element [0]). - tracer.ablate(envoy, mode="zero"): zero/mean knockout. mode="mean" is the self-contained within-sequence mean; reference-distribution (dataset) mean ablation is a precomputed value transplanted via tracer.patch — not derivable from a single forward, so kept distinct to avoid silent wrong-mean semantics. Unknown mode raises ValueError. Verified bit-identical (max|Δ|=0) under forced isolation vs in-process on gpt2 + a renamed model: test_isolated_patch.py 6/6, test_isolated_ablate.py 7/7, including the crux that the in-place form is a silent no-op under isolation while the primitive takes effect. No regression in steer/unembed/trace/acceptance/cache/cross_invoke/backward. Docs: docs/developing/fast-lane.md §6 (built), new subsections, §7/§8 updated. Co-Authored-By: Claude Opus 4.8 (1M context) Claude-Session: https://claude.ai/code/session_014ZUUF44B2tfuKBNFhDFedR --- docs/developing/fast-lane.md | 92 ++++++- .../gpu_sandbox/test_isolated_ablate.py | 250 ++++++++++++++++++ .../gpu_sandbox/test_isolated_patch.py | 228 ++++++++++++++++ src/nnsight/intervention/tracing/tracer.py | 102 +++++++ 4 files changed, 667 insertions(+), 5 deletions(-) create mode 100644 prototypes/mediator-sandbox/gpu_sandbox/test_isolated_ablate.py create mode 100644 prototypes/mediator-sandbox/gpu_sandbox/test_isolated_patch.py diff --git a/docs/developing/fast-lane.md b/docs/developing/fast-lane.md index 262451edf..23f81ee5e 100644 --- a/docs/developing/fast-lane.md +++ b/docs/developing/fast-lane.md @@ -122,8 +122,8 @@ tier too, and (b) collapsing common raw-compute patterns into named calls the ga |---|---|---|---|---| | `tracer.unembed` | `(residual, norm, head, formulation="weight") → logits` | host-weight read + module call | the projection every logit-lens / steering-direction / attribution metric does | **built** | | `tracer.steer` | `(envoy, direction, alpha=1.0)` | boundary write (injection) | always a replacement swap — fixes in-place's silent no-op under isolation | **built** | -| `tracer.patch` | `(envoy, value)` | boundary write (transplant) | whole-tuple replacement | designed | -| `tracer.ablate` | `(envoy, mode)` | boundary write (injection) | zero/mean knockout | designed | +| `tracer.patch` | `(envoy, value)` | boundary write (transplant) | whole-tuple replacement | **built** | +| `tracer.ablate` | `(envoy, mode="zero")` | boundary write (injection) | zero/mean knockout | **built** | | `tracer.capture` | `(value) → handle` | read + run↔run transfer | cross-trace handoff; non-transmittable → clean fail, not silent drop | designed | Most mirror the existing `tracer.cache()` shape: in-process they resolve the real envoys and run @@ -196,6 +196,78 @@ motivation under forced isolation: the in-place form leaves the downstream resid baseline (silent no-op) while `tracer.steer` changes it (steering took effect) to exactly the in-process result. +### `tracer.patch` — replacement-swap transplant (built, 2026-06-20) + +`tracer.patch(envoy, value)` replaces a module's output residual with a precomputed `value` — the +activation-patching / resampling transplant every patching cell does. It is the structural twin of +`tracer.steer`: a boundary write that touches **no host weights** (only the *delivered* activation is +replaced), so it needs **no new event and no isolated/in-process branch** and rides the existing +`Events.SWAP`. The only difference from `steer` is the source of the new value — `steer` computes it from +the delivered activation (`hidden + alpha*direction`), `patch` takes it from the caller: + +```python +out = envoy.output +hidden = out[0] if isinstance(out, tuple) else out +value = value.to(dtype=hidden.dtype, device=hidden.device) +envoy.output = (value, *out[1:]) if isinstance(out, tuple) else value +``` + +The point, as with `steer`, is the **replacement** swap. The hand-written form is in-place +(`block.output[0][:] = clean_act`); under isolation that mutates only the worker's *delivered clone*, no +SWAP fires, and the host's real activation is untouched — the transplant silently no-ops. `tracer.patch` +ships the value back over SWAP, so it crosses the boundary by construction. The value is cast to the +residual's dtype/device, so a value precomputed **on CPU** — the isolation-relevant case, where the +clean/source activation is captured in a prior run outside the trace — transplants cleanly. `value` +replaces the residual whole (element `[0]` for tuple outputs); partial patches construct a full-shape +value (clone + edit a slice), the standard nnsight idiom. + +Touch points: just the `tracer.patch` method (tracer.py). No event, no host handler — it reuses +`Events.SWAP` and the eproperty setter. + +**Verified (`test_isolated_patch.py`, gpt2 + renamed model):** transplanting into one block, an attention +tuple output, and three blocks at once are all isolated-vs-in-process `max|Δ|=0` and propagate through +later layers; `tracer.patch` equals the manual untuple + whole-tuple replacement (with the same cast). The +crux case proves the motivation under forced isolation: the in-place form leaves the downstream residual +== the unpatched baseline (silent no-op) while `tracer.patch` changes it (transplant took effect) to +exactly the in-process result. + +### `tracer.ablate` — replacement-swap knockout (built, 2026-06-20) + +`tracer.ablate(envoy, mode="zero")` replaces a module's output residual with a baseline — the lesion-study +knockout every ablation cell does. Same shape as `patch`/`steer`: a boundary write riding `Events.SWAP`, +no new event, no isolated/in-process branch. Two self-contained modes: + +```python +out = envoy.output +hidden = out[0] if isinstance(out, tuple) else out +if mode == "zero": + ablated = torch.zeros_like(hidden) +elif mode == "mean": # within-sequence mean + ablated = hidden.mean(dim=-2, keepdim=True).expand_as(hidden).contiguous() +else: + raise ValueError(...) # no silent wrong-ablation +envoy.output = (ablated, *out[1:]) if isinstance(out, tuple) else ablated +``` + +`mode="mean"` is the **within-sequence** mean (each position → the per-example mean over the token +dimension), the only mean derivable from a single forward. The reduction decision that §6 flagged resolves +here: **reference-distribution** mean ablation (the mean activation over a *dataset*, per +`docs/patterns/ablation.md`) is not a single-forward quantity, so it is precomputed and transplanted via +`tracer.patch(envoy, mean_act)` — not this mode. Keeping the two distinct avoids a silent-semantics trap +(a user expecting dataset-mean getting sequence-mean). An unknown mode raises `ValueError` rather than +silently picking a baseline. Under isolation the worker computes the mean from its delivered clone (== +the host's real activation) and ships the result back over SWAP, so isolated == in-process bit-identically. + +Touch points: just the `tracer.ablate` method (tracer.py). No event, no host handler. + +**Verified (`test_isolated_ablate.py`, gpt2 + renamed model):** zero- and mean-ablating a block, an +attention tuple output, and the renamed model are all isolated-vs-in-process `max|Δ|=0` and change the +downstream forward vs the un-ablated baseline (the knockout took effect across the boundary); +`tracer.ablate` equals the manual `zeros_like` / mean-over-seq replacement; an unknown mode raises +`ValueError`. The crux case proves the motivation under forced isolation: the in-place zero leaves the +downstream residual == the un-ablated baseline (silent no-op) while `tracer.ablate` changes it to exactly +the in-process result. + ## 7. What was deliberately deferred - **The process-global `sys.addaudithook` backstop.** Its own failure mode (a leaked thread-local flag @@ -204,9 +276,10 @@ result. fast-laned code. Documented as future hardening; the static gate is the confirmation. - **A frozen-namespace `Compartment`** (SES-style) for fast-lane execution. The first slice relies on the static pass + the `trust` cordon; namespace shadowing is a later refinement. -- **The remaining primitives** (§6) — `patch`/`ablate`/`capture`. `unembed` and `steer` are built - (`steer` rides `Events.SWAP`, so it needed no new handler; `patch`/`ablate` will too, `capture` needs - a run↔run handoff). +- **The last primitive** (§6) — `capture`. `unembed`, `steer`, `patch`, and `ablate` are built; the + boundary-write trio (`steer`/`patch`/`ablate`) all ride `Events.SWAP` with no new handler. `capture` + remains because it needs a run↔run handoff (not a single boundary write) and would collide with the + existing `Tracer.capture(frame)` AST method — both a new mechanism and a naming decision. ## 8. Verification @@ -223,6 +296,15 @@ result. attention tuple output, and three blocks at once are isolated-vs-in-process `max|Δ|=0`; `tracer.steer` equals the manual whole-tuple replacement; and the crux — under forced isolation the in-place form is a no-op (downstream == unsteered baseline) while `tracer.steer` takes effect and matches in-process. +- **Isolated patch** (`test_isolated_patch.py`, gpt2 + a renamed model) — 6/6: transplanting into a block, + an attention tuple output, and three blocks at once are isolated-vs-in-process `max|Δ|=0`; `tracer.patch` + equals the manual whole-tuple replacement (with the dtype/device cast); and the crux — the in-place + transplant is a no-op under isolation while `tracer.patch` takes effect and matches in-process. +- **Isolated ablate** (`test_isolated_ablate.py`, gpt2 + a renamed model) — 7/7: zero/mean knockout of a + block, an attention tuple output, and the renamed model are isolated-vs-in-process `max|Δ|=0` and change + the downstream forward vs the un-ablated baseline; `tracer.ablate` equals the manual `zeros_like` / + mean-over-seq replacement; an unknown mode raises `ValueError`; and the crux — the in-place zero is a + no-op under isolation while `tracer.ablate` takes effect and matches in-process. - **Existing isolated WORKER path** — 9/9 still bit-identical, pinned with `fast_lane=False` so they keep exercising the worker (otherwise the simple read/swap/save cells would now fast-lane). - **In-process core** — 51 passed (the default in-process path is untouched). diff --git a/prototypes/mediator-sandbox/gpu_sandbox/test_isolated_ablate.py b/prototypes/mediator-sandbox/gpu_sandbox/test_isolated_ablate.py new file mode 100644 index 000000000..d586b9d46 --- /dev/null +++ b/prototypes/mediator-sandbox/gpu_sandbox/test_isolated_ablate.py @@ -0,0 +1,250 @@ +#!/usr/bin/env python3 +"""Ablation (zero / mean knockout) on the ISOLATED tier: tracer.ablate replaces a module's +output with a baseline (zeros, or the within-sequence mean) via a REPLACEMENT swap, so the +knockout crosses the isolation boundary and propagates through the host's real forward — +where the hand-written in-place form (``hidden[:] = 0``) silently no-ops (the worker mutates +its delivered clone, no SWAP fires, the host's activation is untouched). + +Like tracer.steer/patch (and unlike tracer.unembed), ablation touches no host weights — the +baseline is derived from (or independent of) the delivered activation — so it rides the +existing Events.SWAP with no host round-trip and no isolated/in-process branch: the SAME +method is correct in-process, on the fast lane, and in the isolated worker. + +``mode="mean"`` here is the SELF-CONTAINED within-sequence mean (each position → the +per-example mean over the token dimension). Reference-distribution mean ablation (the mean +activation over a dataset, per docs/patterns/ablation.md) is a precomputed value transplanted +via tracer.patch — not this mode. + + zero_single — zero-ablate one block; forced-isolation == in-process AND downstream != + baseline (the ablation took effect across the boundary). + mean_single — mean-ablate one block; forced-isolation == in-process AND != baseline. + crux — THE crux: under forced isolation the in-place zero is a no-op (downstream + == un-ablated baseline) while tracer.ablate actually ablates (downstream != + baseline) and equals the in-process zero-ablated result. + tuple_output — zero-ablate an attention output (a tuple `(tensor, None)`): the whole-tuple + replacement branch, forced-isolation == in-process. + renamed — renamed model (decoder_blocks): forced-isolation == in-process. + matches_manual— tracer.ablate (in-process) equals the manual zeros_like / mean-over-seq + replacement it names. + bad_mode — an unknown mode raises ValueError (no silent wrong-ablation). + +Run: + CUDA_VISIBLE_DEVICES=7 PYTHONPATH=src \ + /disk/u/zikai/anaconda3/envs/nnsight-tf/bin/python -u \ + prototypes/mediator-sandbox/gpu_sandbox/test_isolated_ablate.py +""" +import sys + +import torch + +from nnsight import LanguageModel +from nnsight.intervention.isolation import isolate_mediators + +PROMPT = "The Eiffel Tower is in the city of" + + +def _both(build): + """Run build() in-process and under forced isolation; return (ref, got).""" + ref = build() + with isolate_mediators(fast_lane=False, timeout=30): + got = build() + return ref, got + + +def _eq(ref, got): + return (torch.is_tensor(ref) and torch.is_tensor(got) + and ref.shape == got.shape and torch.equal(ref, got)) + + +def _delta(ref, got): + return (ref - got).abs().max().item() if _eq(ref, got) else float("nan") + + +def _baseline_downstream(model, read): + with model.trace(PROMPT): + with torch.no_grad(): + o = read.output + down = (o[0] if isinstance(o, tuple) else o)[:, -1, :].save() + return down + + +def test_zero_single(model): + base = _baseline_downstream(model, model.transformer.h[10]) + + def build(): + with model.trace(PROMPT) as tracer: + with torch.no_grad(): + tracer.ablate(model.transformer.h[6], "zero") + o = model.transformer.h[10].output + down = (o[0] if isinstance(o, tuple) else o)[:, -1, :].save() + return down + ref, got = _both(build) + ok = _eq(ref, got) and not _eq(ref, base) + print(f"[zero_single] isolated zero-ablate bit-identical={_eq(ref, got)} " + f"took_effect={not _eq(ref, base)} (max|Δ|={_delta(ref, got)})", flush=True) + return ok + + +def test_mean_single(model): + base = _baseline_downstream(model, model.transformer.h[10]) + + def build(): + with model.trace(PROMPT) as tracer: + with torch.no_grad(): + tracer.ablate(model.transformer.h[6], "mean") + o = model.transformer.h[10].output + down = (o[0] if isinstance(o, tuple) else o)[:, -1, :].save() + return down + ref, got = _both(build) + ok = _eq(ref, got) and not _eq(ref, base) + print(f"[mean_single] isolated mean-ablate bit-identical={_eq(ref, got)} " + f"took_effect={not _eq(ref, base)} (max|Δ|={_delta(ref, got)})", flush=True) + return ok + + +def test_crux(model): + """The crux: replacement swap crosses the boundary; in-place zero is a silent no-op.""" + base = _baseline_downstream(model, model.transformer.h[10]) + + def ablate_downstream(): # tracer.ablate (replacement swap) + with model.trace(PROMPT) as tracer: + with torch.no_grad(): + tracer.ablate(model.transformer.h[6], "zero") + o = model.transformer.h[10].output + down = (o[0] if isinstance(o, tuple) else o)[:, -1, :].save() + return down + + def inplace_downstream(): # hand-written in-place zero + with model.trace(PROMPT): + with torch.no_grad(): + out = model.transformer.h[6].output + hidden = out[0] if isinstance(out, tuple) else out + hidden[:] = 0 + o = model.transformer.h[10].output + down = (o[0] if isinstance(o, tuple) else o)[:, -1, :].save() + return down + + ablate_ip = ablate_downstream() # in-process ablate + with isolate_mediators(fast_lane=False, timeout=30): + ablate_iso = ablate_downstream() # isolated ablate (replacement swap) + inplace_iso = inplace_downstream() # isolated in-place (no-op) + + inplace_is_noop = _eq(inplace_iso, base) # in-place never crossed the boundary + ablate_took_effect = not _eq(ablate_iso, base) # replacement swap changed the host forward + ablate_correct = _eq(ablate_iso, ablate_ip) # ... to exactly the in-process result + ok = inplace_is_noop and ablate_took_effect and ablate_correct + print(f"[crux] isolated in-place zero is a no-op={inplace_is_noop}; " + f"isolated ablate took effect={ablate_took_effect}; " + f"ablate iso==in-process={ablate_correct} (max|Δ|={_delta(ablate_iso, ablate_ip)})", flush=True) + return ok + + +def test_tuple_output(model): + base = _baseline_downstream(model, model.transformer.h[10]) + + def build(): + with model.trace(PROMPT) as tracer: + with torch.no_grad(): + tracer.ablate(model.transformer.h[6].attn, "zero") + o = model.transformer.h[10].output + down = (o[0] if isinstance(o, tuple) else o)[:, -1, :].save() + return down + ref, got = _both(build) + ok = _eq(ref, got) and not _eq(ref, base) + print(f"[tuple_output] tuple (attn) zero-ablate bit-identical={_eq(ref, got)} " + f"took_effect={not _eq(ref, base)} (max|Δ|={_delta(ref, got)})", flush=True) + return ok + + +def test_renamed(): + rename = {"transformer.ln_f": "final_norm", "lm_head": "output_projection", + "transformer.h": "decoder_blocks"} + model = LanguageModel("gpt2", device_map="cuda", dispatch=True, rename=rename) + + def build(): + with model.trace(PROMPT) as tracer: + with torch.no_grad(): + tracer.ablate(model.decoder_blocks[3], "mean") + o = model.decoder_blocks[9].output + down = (o[0] if isinstance(o, tuple) else o)[:, -1, :].save() + return down + ref, got = _both(build) + ok = _eq(ref, got) + print(f"[renamed] renamed-model isolated ablate bit-identical={ok} (max|Δ|={_delta(ref, got)})", flush=True) + return ok + + +def test_matches_manual(model): + # tracer.ablate (in-process) must equal the manual zeros_like / mean-over-seq replacement. + with model.trace(PROMPT) as tracer: + with torch.no_grad(): + tracer.ablate(model.transformer.h[6], "zero") + z_api = (model.transformer.h[10].output[0] + if isinstance(model.transformer.h[10].output, tuple) + else model.transformer.h[10].output)[:, -1, :].save() + with model.trace(PROMPT) as tracer: + with torch.no_grad(): + tracer.ablate(model.transformer.h[6], "mean") + o = model.transformer.h[10].output + m_api = (o[0] if isinstance(o, tuple) else o)[:, -1, :].save() + + with model.trace(PROMPT): + with torch.no_grad(): + out = model.transformer.h[6].output + is_tuple = isinstance(out, tuple) + hidden = out[0] if is_tuple else out + zeroed = torch.zeros_like(hidden) + model.transformer.h[6].output = (zeroed, *out[1:]) if is_tuple else zeroed + o = model.transformer.h[10].output + z_manual = (o[0] if isinstance(o, tuple) else o)[:, -1, :].save() + with model.trace(PROMPT): + with torch.no_grad(): + out = model.transformer.h[6].output + is_tuple = isinstance(out, tuple) + hidden = out[0] if is_tuple else out + meaned = hidden.mean(dim=-2, keepdim=True).expand_as(hidden).contiguous() + model.transformer.h[6].output = (meaned, *out[1:]) if is_tuple else meaned + o = model.transformer.h[10].output + m_manual = (o[0] if isinstance(o, tuple) else o)[:, -1, :].save() + + ok = _eq(z_api, z_manual) and _eq(m_api, m_manual) + print(f"[matches_manual] zero == manual zeros_like={_eq(z_api, z_manual)} " + f"(max|Δ|={_delta(z_api, z_manual)}); mean == manual mean-over-seq={_eq(m_api, m_manual)} " + f"(max|Δ|={_delta(m_api, m_manual)})", flush=True) + return ok + + +def test_bad_mode(model): + raised = False + try: + with model.trace(PROMPT) as tracer: + with torch.no_grad(): + tracer.ablate(model.transformer.h[6], "median") + model.transformer.h[10].output # force the body to run + except ValueError: + raised = True + except Exception as e: # any other surfaced error is not the contract + print(f"[bad_mode] wrong exception type: {type(e).__name__}: {e}", flush=True) + print(f"[bad_mode] unknown mode raises ValueError={raised}", flush=True) + return raised + + +def main(): + assert torch.cuda.is_available() + model = LanguageModel("gpt2", device_map="cuda", dispatch=True) + results = { + "zero_single": test_zero_single(model), + "mean_single": test_mean_single(model), + "crux": test_crux(model), + "tuple_output": test_tuple_output(model), + "renamed": test_renamed(), + "matches_manual": test_matches_manual(model), + "bad_mode": test_bad_mode(model), + } + print("=" * 72, flush=True) + print(f"ISOLATED ABLATE: {results}", flush=True) + sys.exit(0 if all(results.values()) else 1) + + +if __name__ == "__main__": + main() diff --git a/prototypes/mediator-sandbox/gpu_sandbox/test_isolated_patch.py b/prototypes/mediator-sandbox/gpu_sandbox/test_isolated_patch.py new file mode 100644 index 000000000..be8150fa9 --- /dev/null +++ b/prototypes/mediator-sandbox/gpu_sandbox/test_isolated_patch.py @@ -0,0 +1,228 @@ +#!/usr/bin/env python3 +"""Activation patching (transplant) on the ISOLATED tier: tracer.patch replaces a module's +output with a precomputed value via a REPLACEMENT swap, so the transplant crosses the +isolation boundary and propagates through the host's real forward — where the hand-written +in-place form (``hidden[:] = value``) silently no-ops (the worker mutates its delivered +clone, no SWAP fires, the host's activation is untouched). + +Like tracer.steer (and unlike tracer.unembed), patching touches no host weights — only the +delivered activation is replaced — so it rides the existing Events.SWAP with no host +round-trip and no isolated/in-process branch: the SAME method is correct in-process, on the +fast lane, and in the isolated worker. The patch VALUE is precomputed outside the trace (on +CPU) — it crosses the boundary as a transplanted value, mirroring real usage where the +clean/source activation is captured in a prior run. + + single — patch one block (single-tensor output); forced-isolation downstream + residual == in-process, bit-identical: the replacement swap crossed the + boundary AND propagated through later layers. + transplant — THE crux: under forced isolation the in-place form is a no-op (downstream + == unpatched baseline) while tracer.patch actually transplants (downstream + != baseline) and equals the in-process patched result. + tuple_output — patch an attention output (a tuple `(tensor, None)`): the whole-tuple + replacement branch, forced-isolation == in-process. + multi — patch three blocks (forward order), forced-isolation == in-process. + renamed — renamed model (decoder_blocks): forced-isolation == in-process (no names + hardcoded; the wire path is the real path). + matches_manual— tracer.patch (in-process) equals the manual untuple+replacement (with the + same dtype/device cast) it names. + +Run: + CUDA_VISIBLE_DEVICES=7 PYTHONPATH=src \ + /disk/u/zikai/anaconda3/envs/nnsight-tf/bin/python -u \ + prototypes/mediator-sandbox/gpu_sandbox/test_isolated_patch.py +""" +import sys + +import torch + +from nnsight import LanguageModel +from nnsight.intervention.isolation import isolate_mediators + +PROMPT = "The Eiffel Tower is in the city of" + + +def _patch_value(envoy, model): + """A full-shape replacement activation, precomputed OUTSIDE the trace (on CPU) so the + isolated worker transplants a value it never had to compute. Derived from the site's own + baseline residual (so shapes match by construction) and clearly perturbed (scale+shift) + so the transplant provably changes the downstream forward.""" + with model.trace(PROMPT): + with torch.no_grad(): + o = envoy.output + r = ((o[0] if isinstance(o, tuple) else o) * 1.5 + 0.1).save() + return r.detach().cpu() + + +def _both(build): + """Run build() in-process and under forced isolation; return (ref, got).""" + ref = build() + with isolate_mediators(fast_lane=False, timeout=30): + got = build() + return ref, got + + +def _eq(ref, got): + return (torch.is_tensor(ref) and torch.is_tensor(got) + and ref.shape == got.shape and torch.equal(ref, got)) + + +def _delta(ref, got): + return (ref - got).abs().max().item() if _eq(ref, got) else float("nan") + + +def test_single(model): + value = _patch_value(model.transformer.h[6], model) + + def build(): + with model.trace(PROMPT) as tracer: + with torch.no_grad(): + tracer.patch(model.transformer.h[6], value) + o = model.transformer.h[10].output # downstream residual + down = (o[0] if isinstance(o, tuple) else o)[:, -1, :].save() + return down + ref, got = _both(build) + ok = _eq(ref, got) + print(f"[single] isolated patch downstream bit-identical={ok} (max|Δ|={_delta(ref, got)})", flush=True) + return ok + + +def test_transplant(model): + """The crux: replacement swap crosses the boundary; in-place is a silent no-op.""" + value = _patch_value(model.transformer.h[6], model) + + def read_downstream(): # unpatched baseline + with model.trace(PROMPT): + with torch.no_grad(): + o = model.transformer.h[10].output + down = (o[0] if isinstance(o, tuple) else o)[:, -1, :].save() + return down + + def patch_downstream(): # tracer.patch (replacement swap) + with model.trace(PROMPT) as tracer: + with torch.no_grad(): + tracer.patch(model.transformer.h[6], value) + o = model.transformer.h[10].output + down = (o[0] if isinstance(o, tuple) else o)[:, -1, :].save() + return down + + def inplace_downstream(): # hand-written in-place transplant + with model.trace(PROMPT): + with torch.no_grad(): + out = model.transformer.h[6].output + hidden = out[0] if isinstance(out, tuple) else out + hidden[:] = value.to(dtype=hidden.dtype, device=hidden.device) + o = model.transformer.h[10].output + down = (o[0] if isinstance(o, tuple) else o)[:, -1, :].save() + return down + + base = read_downstream() + patch_ip = patch_downstream() # in-process patch + with isolate_mediators(fast_lane=False, timeout=30): + patch_iso = patch_downstream() # isolated patch (replacement swap) + inplace_iso = inplace_downstream() # isolated in-place (no-op) + + inplace_is_noop = _eq(inplace_iso, base) # in-place never crossed the boundary + patch_took_effect = not _eq(patch_iso, base) # replacement swap changed the host forward + patch_correct = _eq(patch_iso, patch_ip) # ... to exactly the in-process result + ok = inplace_is_noop and patch_took_effect and patch_correct + print(f"[transplant] isolated in-place is a no-op={inplace_is_noop}; " + f"isolated patch took effect={patch_took_effect}; " + f"patch iso==in-process={patch_correct} (max|Δ|={_delta(patch_iso, patch_ip)})", flush=True) + return ok + + +def test_tuple_output(model): + # An attention output is a tuple (tensor, None) — exercises the whole-tuple replacement + # branch (transplanted element [0], None tail carried through). + value = _patch_value(model.transformer.h[6].attn, model) + + def build(): + with model.trace(PROMPT) as tracer: + with torch.no_grad(): + tracer.patch(model.transformer.h[6].attn, value) + o = model.transformer.h[10].output + down = (o[0] if isinstance(o, tuple) else o)[:, -1, :].save() + return down + ref, got = _both(build) + ok = _eq(ref, got) + print(f"[tuple_output] tuple (attn) patch bit-identical={ok} (max|Δ|={_delta(ref, got)})", flush=True) + return ok + + +def test_multi(model): + vals = {i: _patch_value(model.transformer.h[i], model) for i in [2, 5, 8]} + + def build(): + with model.trace(PROMPT) as tracer: + with torch.no_grad(): + for i in [2, 5, 8]: # forward order + tracer.patch(model.transformer.h[i], vals[i]) + o = model.transformer.h[11].output + down = (o[0] if isinstance(o, tuple) else o)[:, -1, :].save() + return down + ref, got = _both(build) + ok = _eq(ref, got) + print(f"[multi] 3-block isolated patch bit-identical={ok} (max|Δ|={_delta(ref, got)})", flush=True) + return ok + + +def test_renamed(): + rename = {"transformer.ln_f": "final_norm", "lm_head": "output_projection", + "transformer.h": "decoder_blocks"} + model = LanguageModel("gpt2", device_map="cuda", dispatch=True, rename=rename) + value = _patch_value(model.decoder_blocks[3], model) + + def build(): + with model.trace(PROMPT) as tracer: + with torch.no_grad(): + tracer.patch(model.decoder_blocks[3], value) + o = model.decoder_blocks[9].output + down = (o[0] if isinstance(o, tuple) else o)[:, -1, :].save() + return down + ref, got = _both(build) + ok = _eq(ref, got) + print(f"[renamed] renamed-model isolated patch bit-identical={ok} (max|Δ|={_delta(ref, got)})", flush=True) + return ok + + +def test_matches_manual(model): + # tracer.patch (in-process) must equal the manual untuple + replacement (with cast) it names. + value = _patch_value(model.transformer.h[6], model) + with model.trace(PROMPT) as tracer: + with torch.no_grad(): + tracer.patch(model.transformer.h[6], value) + o = model.transformer.h[10].output + via_api = (o[0] if isinstance(o, tuple) else o)[:, -1, :].save() + with model.trace(PROMPT): + with torch.no_grad(): + out = model.transformer.h[6].output + is_tuple = isinstance(out, tuple) + hidden = out[0] if is_tuple else out + v = value.to(dtype=hidden.dtype, device=hidden.device) + model.transformer.h[6].output = (v, *out[1:]) if is_tuple else v + o = model.transformer.h[10].output + via_manual = (o[0] if isinstance(o, tuple) else o)[:, -1, :].save() + ok = _eq(via_api, via_manual) + print(f"[matches_manual] tracer.patch == manual replacement bit-identical={ok} " + f"(max|Δ|={_delta(via_api, via_manual)})", flush=True) + return ok + + +def main(): + assert torch.cuda.is_available() + model = LanguageModel("gpt2", device_map="cuda", dispatch=True) + results = { + "single": test_single(model), + "transplant": test_transplant(model), + "tuple_output": test_tuple_output(model), + "multi": test_multi(model), + "renamed": test_renamed(), + "matches_manual": test_matches_manual(model), + } + print("=" * 72, flush=True) + print(f"ISOLATED PATCH: {results}", flush=True) + sys.exit(0 if all(results.values()) else 1) + + +if __name__ == "__main__": + main() diff --git a/src/nnsight/intervention/tracing/tracer.py b/src/nnsight/intervention/tracing/tracer.py index 1add36624..cd2a4fa9b 100755 --- a/src/nnsight/intervention/tracing/tracer.py +++ b/src/nnsight/intervention/tracing/tracer.py @@ -746,6 +746,108 @@ def steer(self, envoy, direction, alpha: float = 1.0): return steered + def patch(self, envoy, value): + """Transplant ``value`` into ``envoy``'s output residual via a **replacement** + boundary write — the activation-patching / resampling transplant every patching cell + does, made to work on the isolated tier where an in-place ``hidden[:] = value`` + silently no-ops. + + Patching replaces a module's delivered activation with one captured elsewhere (a clean + run, a different prompt, a precomputed mean). Done in place that write never crosses the + isolation boundary: the worker mutates its *delivered clone*, no SWAP fires, and the + host's real activation is untouched — a silent no-op. ``tracer.patch`` always performs a + *replacement* swap (assign ``envoy.output``), which ships the transplanted value back + over the existing ``Events.SWAP`` path, so it is correct in-process, on the fast lane, + AND in the isolated worker. Like :meth:`steer`, it touches no host weights — only the + delivered activation — so it needs no host round-trip and no isolated/in-process branch. + + ``value`` replaces the residual whole (element ``[0]`` for tuple outputs); construct a + full-shape value (e.g. clone the residual and edit a slice) for partial patches, the + standard nnsight idiom. The value is cast to the residual's dtype/device, so a value + precomputed on CPU (the isolation case) transplants cleanly. + + Tuple outputs (attention modules, and transformer blocks on transformers <5) are + replaced whole — element ``[0]`` is transplanted and the rest of the tuple rides + through unchanged. + + Args: + envoy: the module whose output residual to replace (e.g. + ``model.transformer.h[6]``). + value: the replacement activation, broadcast/cast to the residual's dtype/device. + Must match the residual's shape (element ``[0]`` for tuple outputs). + + Returns: + The transplanted residual tensor (element ``[0]`` for tuple outputs). + """ + out = envoy.output + is_tuple = isinstance(out, tuple) + hidden = out[0] if is_tuple else out + + value = value.to(dtype=hidden.dtype, device=hidden.device) + + if is_tuple: + envoy.output = (value, *out[1:]) + else: + envoy.output = value + + return value + + def ablate(self, envoy, mode: str = "zero"): + """Knock out ``envoy``'s output residual via a **replacement** boundary write — the + lesion-study ablation every ablation cell does, made to work on the isolated tier where + an in-place ``hidden[:] = 0`` silently no-ops. + + Ablation replaces a component's output with a baseline to measure how the prediction + degrades. Done in place that write never crosses the isolation boundary: the worker + mutates its *delivered clone*, no SWAP fires, and the host's real activation is + untouched — a silent no-op (the model never sees the lesion). ``tracer.ablate`` always + performs a *replacement* swap, riding the existing ``Events.SWAP`` path, so it is + correct in-process, on the fast lane, AND in the isolated worker. Like :meth:`steer`, it + touches no host weights, so it needs no host round-trip and no isolated/in-process + branch. + + Modes: + ``"zero"`` — replace with zeros (zero ablation; pushes the residual off-distribution + but is the simplest knockout). + ``"mean"`` — replace with the within-sequence mean: each position becomes the + per-example mean over the token dimension, keeping the average magnitude while + removing position-specific deviation. This is the *self-contained* mean. For + *reference-distribution* mean ablation (the mean activation over a dataset, per + docs/patterns/ablation.md), precompute that mean and transplant it via + :meth:`patch` — that mean is not derivable from a single forward. + + Tuple outputs (attention modules, and transformer blocks on transformers <5) are + replaced whole — element ``[0]`` is ablated and the rest of the tuple rides through + unchanged. + + Args: + envoy: the module whose output residual to ablate (e.g. + ``model.transformer.h[6]``). + mode: ``"zero"`` (default) or ``"mean"``. + + Returns: + The ablated residual tensor (element ``[0]`` for tuple outputs). + """ + out = envoy.output + is_tuple = isinstance(out, tuple) + hidden = out[0] if is_tuple else out + + if mode == "zero": + ablated = torch.zeros_like(hidden) + elif mode == "mean": + ablated = hidden.mean(dim=-2, keepdim=True).expand_as(hidden).contiguous() + else: + raise ValueError( + f"tracer.ablate mode must be 'zero' or 'mean', got {mode!r}" + ) + + if is_tuple: + envoy.output = (ablated, *out[1:]) + else: + envoy.output = ablated + + return ablated + def barrier(self, n_participants: int): """ nnsight barrier: A synchronization primitive for coordinating multiple concurrent invocations in nnsight. From 39e62de27066d2a3c94c4b097eedbeac29debd0f Mon Sep 17 00:00:00 2001 From: khaiwang Date: Sun, 21 Jun 2026 18:18:52 -0400 Subject: [PATCH 24/30] feat(intervention): fix session cross-trace handoff under isolation + .carry() primitive MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit model.session() carries values across its inner traces in-process (each inner trace pushes its locals up to the session frame; the session's exit-push surfaces only saves). Under isolation each inner trace runs in a worker that shipped only its .save()'d locals home, so the two-hop push was broken two ways: - a SAVED value used cross-trace was written into the session frame but its host id was never re-registered in Globals.saves, so the session's exit-push dropped it (UnboundLocalError); - a NON-saved value used cross-trace was never shipped at all (NameError). The realized form of the last part-2 primitive (the run<->run handoff originally specced as tracer.capture, which collided with the existing Tracer.capture(frame) AST method): - Saved-case fix: when the isolated END target is a nested/session frame, the host writes the worker's values into it AND re-registers the saved values' host ids in Globals.saves, so the session's exit-push keeps them. Root (single-trace) writeback is unchanged. Makes the documented `hs = x.save()` -> use `hs` session pattern work under isolation. - .carry() (universal value method, like .save(); plus nnsight.carry(x)): hand a value to a later trace in the session WITHOUT surfacing it as an output. The worker end() now ships saved-union-carried locals as (values, saved_names); the host writes all to the session frame (next trace sees them) but registers only the saved ones, so carried values drop at session exit — exactly in-process non-saved semantics, made explicit. With no .carry() in use the payload equals the prior saved-only one, so the single-trace path is unchanged. .carry() is portable: harmless in-process, load-bearing under isolation. Root cause confirmed by host-side instrumentation (saved value reaches the session frame but host Globals.saves stays empty -> session root-push drops it). Verified (nnsight-tf, GPU7, gpt2 + renamed model): test_isolated_session_handoff.py 6/6 all max|Δ|=0 (saved + carried handoff isolated==in-process; nnsight.carry==method; carried value not surfaced to caller while saved is; .carry() in-process==isolated). No regression across the isolated suite (trace/cache/backward/multitoken-iter/cross-invoke/acceptance/steer/patch/ ablate/unembed/pool) and in-process core test_lm.py 75/75. Co-Authored-By: Claude Opus 4.8 (1M context) Claude-Session: https://claude.ai/code/session_014ZUUF44B2tfuKBNFhDFedR --- docs/developing/fast-lane.md | 60 ++++- .../test_isolated_session_handoff.py | 226 ++++++++++++++++++ src/nnsight/__init__.py | 2 +- src/nnsight/intervention/interleaver.py | 48 +++- src/nnsight/intervention/isolation.py | 23 +- src/nnsight/intervention/tracing/base.py | 3 + src/nnsight/intervention/tracing/globals.py | 53 +++- 7 files changed, 385 insertions(+), 30 deletions(-) create mode 100644 prototypes/mediator-sandbox/gpu_sandbox/test_isolated_session_handoff.py diff --git a/docs/developing/fast-lane.md b/docs/developing/fast-lane.md index 23f81ee5e..9599f870b 100644 --- a/docs/developing/fast-lane.md +++ b/docs/developing/fast-lane.md @@ -124,7 +124,7 @@ tier too, and (b) collapsing common raw-compute patterns into named calls the ga | `tracer.steer` | `(envoy, direction, alpha=1.0)` | boundary write (injection) | always a replacement swap — fixes in-place's silent no-op under isolation | **built** | | `tracer.patch` | `(envoy, value)` | boundary write (transplant) | whole-tuple replacement | **built** | | `tracer.ablate` | `(envoy, mode="zero")` | boundary write (injection) | zero/mean knockout | **built** | -| `tracer.capture` | `(value) → handle` | read + run↔run transfer | cross-trace handoff; non-transmittable → clean fail, not silent drop | designed | +| `x.carry()` / `nnsight.carry(x)` | `(value) → value` | read + run↔run transfer | session cross-trace handoff under isolation (and `.save()` now crosses too) | **built** | Most mirror the existing `tracer.cache()` shape: in-process they resolve the real envoys and run directly (the fast-lane execution); isolated they ship a spec via a new event whose host handler runs the @@ -268,6 +268,50 @@ downstream forward vs the un-ablated baseline (the knockout took effect across t downstream residual == the un-ablated baseline (silent no-op) while `tracer.ablate` changes it to exactly the in-process result. +### `.carry()` — session cross-trace handoff (built, 2026-06-21) + +The last part-2 primitive started as a designed `tracer.capture(value) → handle`, but investigating the +actual gap reframed it. Cross-run handoff across **separate** `model.trace()` calls already works under +isolation — `.save()` ships the value home and the next trace re-receives it (the basis of +`tracer.patch`). The only broken capability was handoff **inside `model.session()`**, which two isolation +bugs killed: + +- **Saved value used cross-trace** → the host wrote it into the session frame but never re-registered its + host-side id in `Globals.saves`, so the session's exit-push (which filters `id(v) in Globals.saves`) + dropped it (`UnboundLocalError`). +- **Non-saved value used cross-trace** → the isolated worker's `end()` ships only `Globals.saves`-filtered + locals, so a non-saved var never reached the host session frame (`NameError`). + +In-process both work, because each inner trace's `push` is `is_root=False` and pushes *all* its locals up +to the session frame (so the next trace's `pull` sees them), while only the session's final root-push +filters to saves. The isolation END-writeback implemented the root semantics only. + +The fix restores the in-process two-hop semantics on the isolated path, plus an explicit primitive for the +non-saved case: + +- **Saved-case fix:** when the END target is a nested/session frame (`__nnsight_tracing_info__` present), + the host writes the worker's values into it AND re-registers the saved values' host ids in + `Globals.saves`, so the session's exit-push keeps them. Root (single-trace) writeback is unchanged. This + makes the documented `hs = x.save()` → use `hs` session pattern work under isolation. +- **`.carry()`** (mounted universally like `.save()`, plus `nnsight.carry(x)`): marks a value to cross to a + later trace in the session **without** surfacing it as an output. The worker's `end()` now ships saved ∪ + carried locals as `(values, saved_names)`; the host writes all of them to the session frame (so the next + trace sees them) but registers only the saved ones, so carried values drop at session exit — exactly the + in-process non-saved semantics, made explicit. With no `.carry()` in use the payload is exactly the prior + saved-only one, so the single-trace path is unchanged. `.carry()` is **portable**: harmless in-process + (non-saved vars already cross there), load-bearing under isolation. + +Touch points: `Globals.shared` + `carry()` + `Object.carry` + mount (globals.py); `nnsight.carry` export; +`Globals.shared` cleared at the root push (base.py) and per-job worker reset (isolation.py); the worker +`end()` payload (isolation.py); `handle_end_event` nested-registration + carried handling (interleaver.py). + +**Verified (`test_isolated_session_handoff.py`, gpt2 + renamed model):** the saved (`.save()`) and carried +(`.carry()`) session handoffs are isolated-vs-in-process `max|Δ|=0` and change the downstream forward; +`nnsight.carry(x)` equals the method form; a carried value does **not** surface to the caller frame while a +saved one does; `.carry()` is in-process == isolated. No regression across the isolated suite (trace, +cache, backward, multitoken-iter, cross-invoke, acceptance, steer, patch, ablate) — the END-payload change +is transparent to the single-trace path — and the in-process core (`test_lm.py`) is 75/75. + ## 7. What was deliberately deferred - **The process-global `sys.addaudithook` backstop.** Its own failure mode (a leaked thread-local flag @@ -276,10 +320,10 @@ the in-process result. fast-laned code. Documented as future hardening; the static gate is the confirmation. - **A frozen-namespace `Compartment`** (SES-style) for fast-lane execution. The first slice relies on the static pass + the `trust` cordon; namespace shadowing is a later refinement. -- **The last primitive** (§6) — `capture`. `unembed`, `steer`, `patch`, and `ablate` are built; the - boundary-write trio (`steer`/`patch`/`ablate`) all ride `Events.SWAP` with no new handler. `capture` - remains because it needs a run↔run handoff (not a single boundary write) and would collide with the - existing `Tracer.capture(frame)` AST method — both a new mechanism and a naming decision. +- *(none — the part-2 primitive set is complete.)* `unembed`, `steer`, `patch`, `ablate` are built, and + the run↔run handoff shipped as `.carry()` + the session saved-case fix (the original `tracer.capture` + name was dropped — it collided with the existing `Tracer.capture(frame)` AST method, and the realized + primitive is a value method parallel to `.save()`, not a tracer-level handle). ## 8. Verification @@ -305,6 +349,10 @@ the in-process result. the downstream forward vs the un-ablated baseline; `tracer.ablate` equals the manual `zeros_like` / mean-over-seq replacement; an unknown mode raises `ValueError`; and the crux — the in-place zero is a no-op under isolation while `tracer.ablate` takes effect and matches in-process. +- **Isolated session handoff** (`test_isolated_session_handoff.py`, gpt2 + a renamed model) — 6/6: the + `.save()` and `.carry()` session cross-trace handoffs are isolated-vs-in-process `max|Δ|=0` and change + the downstream forward; `nnsight.carry(x)` equals the method form; a carried value does not surface to + the caller frame while a saved one does; `.carry()` is in-process == isolated. - **Existing isolated WORKER path** — 9/9 still bit-identical, pinned with `fast_lane=False` so they keep exercising the worker (otherwise the simple read/swap/save cells would now fast-lane). -- **In-process core** — 51 passed (the default in-process path is untouched). +- **In-process core** — `test_lm.py` 75 passed (the default in-process path is untouched). diff --git a/prototypes/mediator-sandbox/gpu_sandbox/test_isolated_session_handoff.py b/prototypes/mediator-sandbox/gpu_sandbox/test_isolated_session_handoff.py new file mode 100644 index 000000000..8dd781406 --- /dev/null +++ b/prototypes/mediator-sandbox/gpu_sandbox/test_isolated_session_handoff.py @@ -0,0 +1,226 @@ +#!/usr/bin/env python3 +"""Cross-trace (run<->run) handoff inside model.session() on the ISOLATED tier. + +In-process, a session carries values across its inner traces: a value produced in trace 1 +is visible in trace 2 (saved or not), because each inner trace pushes its locals up to the +session frame. Under isolation each inner trace runs in a worker that ships only its +``.save()``'d values home, so cross-trace handoff broke two ways: + - SAVED value used cross-trace -> the host wrote it to the session frame but never + re-registered its host id in Globals.saves, so the session's exit-push dropped it + (UnboundLocalError). + - NON-saved value used cross-trace -> the worker never shipped it (NameError). + +This verifies both are fixed: + - the documented ``hs = x.save()`` -> use ``hs`` session pattern works under isolation; + - ``x.carry()`` explicitly hands a value to a later trace WITHOUT surfacing it as a saved + output, and is portable (in-process == isolated). + + saved_handoff — session, trace1 ``hs = x.save()``, trace2 patches ``hs*1.5`` -> downstream + bit-identical isolated vs in-process, and changed vs the unpatched baseline. + carry_handoff — same but ``hs = x.carry()`` (not saved): bit-identical isolated vs in-process. + carry_func — ``nnsight.carry(x)`` functional form equals the method form. + not_surfaced — after an isolated session, the carried var is NOT in the caller frame + (it was not saved) while the saved output IS. + portable — the same ``.carry()`` code gives identical results in-process and isolated. + renamed — renamed model (decoder_blocks): carry handoff isolated == in-process. + +Run: + CUDA_VISIBLE_DEVICES=7 PYTHONPATH=src \ + /disk/u/zikai/anaconda3/envs/nnsight-tf/bin/python -u \ + prototypes/mediator-sandbox/gpu_sandbox/test_isolated_session_handoff.py +""" +import sys + +import torch + +import nnsight +from nnsight import LanguageModel +from nnsight.intervention.isolation import isolate_mediators + +P = "The Eiffel Tower is in the city of" + + +def U(o): + return o[0] if isinstance(o, tuple) else o + + +def _eq(a, b): + return (torch.is_tensor(a) and torch.is_tensor(b) + and a.shape == b.shape and torch.equal(a, b)) + + +def _delta(a, b): + return (a - b).abs().max().item() if _eq(a, b) else float("nan") + + +def _baseline(model, block): + with model.trace(P): + with torch.no_grad(): + d = U(block.output)[:, -1, :].save() + return d + + +def _ref(model, src, mid, dst): + """In-process reference: session, save src.output, patch dst-input with src*1.5, + read mid downstream. (Works in-process today via the session var-flow.)""" + with model.session(): + with model.trace(P): + with torch.no_grad(): + hs = U(src.output).save() + with model.trace(P) as tracer: + with torch.no_grad(): + tracer.patch(mid, hs * 1.5) + d = U(dst.output)[:, -1, :].save() + return d + + +def test_saved_handoff(model): + base = _baseline(model, model.transformer.h[10]) + ref = _ref(model, model.transformer.h[6], model.transformer.h[6], model.transformer.h[10]) + + def build(): + with model.session(): + with model.trace(P): + with torch.no_grad(): + hs = U(model.transformer.h[6].output).save() + with model.trace(P) as tracer: + with torch.no_grad(): + tracer.patch(model.transformer.h[6], hs * 1.5) + d = U(model.transformer.h[10].output)[:, -1, :].save() + return d + + with isolate_mediators(fast_lane=False, timeout=30): + got = build() + ok = _eq(ref, got) and not _eq(got, base) + print(f"[saved_handoff] isolated == in-process={_eq(ref, got)} " + f"changed_vs_baseline={not _eq(got, base)} (max|Δ|={_delta(ref, got)})", flush=True) + return ok + + +def test_carry_handoff(model): + base = _baseline(model, model.transformer.h[10]) + ref = _ref(model, model.transformer.h[6], model.transformer.h[6], model.transformer.h[10]) + + def build(): + with model.session(): + with model.trace(P): + with torch.no_grad(): + hs = U(model.transformer.h[6].output).carry() # NOT saved + with model.trace(P) as tracer: + with torch.no_grad(): + tracer.patch(model.transformer.h[6], hs * 1.5) + d = U(model.transformer.h[10].output)[:, -1, :].save() + return d + + with isolate_mediators(fast_lane=False, timeout=30): + got = build() + ok = _eq(ref, got) and not _eq(got, base) + print(f"[carry_handoff] isolated == in-process={_eq(ref, got)} " + f"changed_vs_baseline={not _eq(got, base)} (max|Δ|={_delta(ref, got)})", flush=True) + return ok + + +def test_carry_func(model): + ref = _ref(model, model.transformer.h[6], model.transformer.h[6], model.transformer.h[10]) + + def build(): + with model.session(): + with model.trace(P): + with torch.no_grad(): + hs = nnsight.carry(U(model.transformer.h[6].output)) # functional form + with model.trace(P) as tracer: + with torch.no_grad(): + tracer.patch(model.transformer.h[6], hs * 1.5) + d = U(model.transformer.h[10].output)[:, -1, :].save() + return d + + with isolate_mediators(fast_lane=False, timeout=30): + got = build() + ok = _eq(ref, got) + print(f"[carry_func] nnsight.carry(x) isolated == in-process={ok} (max|Δ|={_delta(ref, got)})", flush=True) + return ok + + +def test_not_surfaced(model): + """A carried (non-saved) value must NOT leak to the caller frame; a saved one must.""" + with isolate_mediators(fast_lane=False, timeout=30): + with model.session(): + with model.trace(P): + with torch.no_grad(): + hs = U(model.transformer.h[6].output).carry() + with model.trace(P) as tracer: + with torch.no_grad(): + tracer.patch(model.transformer.h[6], hs * 1.5) + down = U(model.transformer.h[10].output)[:, -1, :].save() + loc = locals() + down_surfaced = torch.is_tensor(loc.get("down")) + hs_surfaced = "hs" in loc + ok = down_surfaced and not hs_surfaced + print(f"[not_surfaced] saved 'down' surfaced={down_surfaced} " + f"carried 'hs' surfaced={hs_surfaced} (want True/False)", flush=True) + return ok + + +def test_portable(model): + """The same .carry() code must give identical results in-process and isolated.""" + def build(): + with model.session(): + with model.trace(P): + with torch.no_grad(): + hs = U(model.transformer.h[6].output).carry() + with model.trace(P) as tracer: + with torch.no_grad(): + tracer.patch(model.transformer.h[6], hs * 1.5) + d = U(model.transformer.h[10].output)[:, -1, :].save() + return d + + ip = build() + with isolate_mediators(fast_lane=False, timeout=30): + iso = build() + ok = _eq(ip, iso) + print(f"[portable] .carry() in-process == isolated={ok} (max|Δ|={_delta(ip, iso)})", flush=True) + return ok + + +def test_renamed(): + rename = {"transformer.ln_f": "final_norm", "lm_head": "output_projection", + "transformer.h": "decoder_blocks"} + model = LanguageModel("gpt2", device_map="cuda", dispatch=True, rename=rename) + ref = _ref(model, model.decoder_blocks[3], model.decoder_blocks[3], model.decoder_blocks[9]) + + def build(): + with model.session(): + with model.trace(P): + with torch.no_grad(): + hs = U(model.decoder_blocks[3].output).carry() + with model.trace(P) as tracer: + with torch.no_grad(): + tracer.patch(model.decoder_blocks[3], hs * 1.5) + d = U(model.decoder_blocks[9].output)[:, -1, :].save() + return d + + with isolate_mediators(fast_lane=False, timeout=30): + got = build() + ok = _eq(ref, got) + print(f"[renamed] renamed-model carry handoff isolated == in-process={ok} (max|Δ|={_delta(ref, got)})", flush=True) + return ok + + +def main(): + assert torch.cuda.is_available() + model = LanguageModel("gpt2", device_map="cuda", dispatch=True) + results = { + "saved_handoff": test_saved_handoff(model), + "carry_handoff": test_carry_handoff(model), + "carry_func": test_carry_func(model), + "not_surfaced": test_not_surfaced(model), + "portable": test_portable(model), + "renamed": test_renamed(), + } + print("=" * 72, flush=True) + print(f"ISOLATED SESSION HANDOFF: {results}", flush=True) + sys.exit(0 if all(results.values()) else 1) + + +if __name__ == "__main__": + main() diff --git a/src/nnsight/__init__.py b/src/nnsight/__init__.py index c6e3f3c60..7fc1eecd2 100755 --- a/src/nnsight/__init__.py +++ b/src/nnsight/__init__.py @@ -54,7 +54,7 @@ except ImportError: __version__ = "unknown version" -from .intervention.tracing.globals import save +from .intervention.tracing.globals import carry, save from .ndif import * from IPython import get_ipython diff --git a/src/nnsight/intervention/interleaver.py b/src/nnsight/intervention/interleaver.py index 8401c88c3..31704376b 100755 --- a/src/nnsight/intervention/interleaver.py +++ b/src/nnsight/intervention/interleaver.py @@ -1630,34 +1630,45 @@ def handle_unembed_event(self, data): self.respond(logits) # ack -> worker's send() returns the logits return True - def handle_end_event(self, saves: Optional[Any] = None): + def handle_end_event(self, payload: Optional[Any] = None): """ Handle an end event by stopping the mediator. - Worker→host saves transmission: in the isolated path the worker bundles its ``.save()``'d values - (already filtered by ``Globals.saves``) into the END event, since the - worker's frame + ``Globals.saves`` live in another process. The host - injects them directly into the real **user** frame — the tracer's - ``info.frame`` (where the in-process two-hop push ultimately lands saved - vars). The worker already applied the ``Globals.saves`` filter, so no - second filtering hop is needed here. + Worker→host saves transmission: in the isolated path the worker bundles its + ``.save()``'d and ``.carry()``'d values into the END event as + ``(values, saved_names)``, since the worker's frame + ``Globals`` live in another + process. ``values`` is the union of saved + carried locals; ``saved_names`` is the + ``.save()``'d subset. + + The host injects these into the tracer's ``info.frame`` — the **user** frame for a + root trace, or the **session** frame for an inner trace within ``model.session()``. + For an inner trace (the target carries ``__nnsight_tracing_info__``), it replicates + the in-process two-hop push: write ALL values so the next trace sees them, and + re-register the SAVED values' host ids in ``Globals.saves`` so the session's exit-push + keeps them (the worker's ids live in another process; without this the session's + root filter would drop everything). Carried values are written for the next trace but + not registered, so they drop at session exit — matching in-process non-saved + semantics. For a root trace, only the saved subset surfaces to the user frame. """ + from .tracing.globals import Globals + if self._iso is not None: # A consumed END means the worker ended cleanly => recyclable by the pool. # cancel() (below) reads this to recycle vs retire the worker. self._iso.clean = True - if saves: + values, saved_names = (payload if payload else ({}, [])) + if values: if self._iso_caches: # tracer.cache(): the worker shipped a ``{_ISO_CACHE_TAG: token}`` # marker in place of its placeholder CacheDict; swap in the HOST cache # (matched by token), which the forward fills in-place after this # injection. The user's variable then IS the forward-filled host cache. - saves = dict(saves) - for name, val in list(saves.items()): + values = dict(values) + for name, val in list(values.items()): if isinstance(val, dict) and _ISO_CACHE_TAG in val: host_cache = self._iso_caches.get(val[_ISO_CACHE_TAG]) if host_cache is not None: - saves[name] = host_cache.cache + values[name] = host_cache.cache tracer = self.interleaver.tracer user_frame = ( tracer.info.frame @@ -1665,7 +1676,18 @@ def handle_end_event(self, saves: Optional[Any] = None): else self.info.frame ) if user_frame is not None: - push_variables(user_frame, saves) + nested = "__nnsight_tracing_info__" in user_frame.f_locals + if nested: + # inner trace in a session: all values cross to the next trace; + # only saved values are registered so they survive the session push. + push_variables(user_frame, values) + for name in saved_names: + if name in values: + Globals.saves.add(id(values[name])) + else: + # root trace: only saved values surface to the user frame. + saved_only = {k: values[k] for k in saved_names if k in values} + push_variables(user_frame, saved_only) self.cancel() diff --git a/src/nnsight/intervention/isolation.py b/src/nnsight/intervention/isolation.py index cbab2dfbb..eb7e7ff27 100644 --- a/src/nnsight/intervention/isolation.py +++ b/src/nnsight/intervention/isolation.py @@ -794,9 +794,9 @@ def _tag(t): apply(value, _tag, torch.Tensor) def end(self): - # Worker→host saves transmission: bundle .save()'d values into the END event. - # The intervention's compiled body calls ``end()`` on success; push() populates - # the SerializedFrame's f_locals, which we filter by Globals.saves. + # Worker→host saves transmission: bundle .save()'d (and .carry()'d) values into the + # END event. The intervention's compiled body calls ``end()`` on success; push() + # populates the SerializedFrame's f_locals, which we filter by Globals.saves/shared. from .interleaver import Events, _ISO_CACHE_TAG from .tracing.globals import Globals from .tracing.tracer import Cache @@ -804,13 +804,23 @@ def end(self): self.push() flocals = self.info.frame.f_locals saved = {k: v for k, v in flocals.items() if id(v) in Globals.saves} + # Carried (.carry()) values: cross-trace handoffs within a session. Ship them too so + # the host can write them to the session frame for the next trace, but tag which are + # SAVED (saved_names) so the host surfaces only those to the user frame. With no + # .carry() in play this is exactly the prior payload (saved only) — no regression for + # the single-trace path. + carried = { + k: v for k, v in flocals.items() + if id(v) in Globals.shared and id(v) not in Globals.saves + } + values = {**saved, **carried} # A tracer.cache() placeholder CacheDict is a live object, not a value — ship its # token as a plain-dict marker so the value codec accepts it; the host swaps in its # own forward-filled cache by token (top-level saves only, as the host handler is). - for k, v in list(saved.items()): + for k, v in list(values.items()): if isinstance(v, Cache.CacheDict): - saved[k] = {_ISO_CACHE_TAG: getattr(v, "_iso_cache_token", None)} - self.channel.put_event((Events.END, saved)) + values[k] = {_ISO_CACHE_TAG: getattr(v, "_iso_cache_token", None)} + self.channel.put_event((Events.END, (values, list(saved.keys())))) def exception(self, exception: Exception): super().exception(_transmissible_exc(exception)) @@ -884,6 +894,7 @@ def _run_one_job(channel, payload, extras, opts, device) -> None: global _WORKER_CURRENT try: Globals.saves.clear() # per-job reset (the only worker-side global state) + Globals.shared.clear() interleaver = _WorkerInterleaver(default_all=opts.get("default_all")) mediator = serialization.loads(payload, _WorkerPersistent(interleaver, extras)) diff --git a/src/nnsight/intervention/tracing/base.py b/src/nnsight/intervention/tracing/base.py index f5fa48dbe..bca7cad35 100755 --- a/src/nnsight/intervention/tracing/base.py +++ b/src/nnsight/intervention/tracing/base.py @@ -566,7 +566,10 @@ def push(self, state: Dict = None) -> Dict: filtered_state = { k: v for k, v in filtered_state.items() if id(v) in Globals.saves } + # Carried (.carry()) values are cross-trace handoffs, never surfaced to the + # user frame; clear them at the root/session exit alongside saves. Globals.saves.clear() + Globals.shared.clear() if target_frame is not None: # Push the filtered variables back to the original frame diff --git a/src/nnsight/intervention/tracing/globals.py b/src/nnsight/intervention/tracing/globals.py index 2594cb8ea..30c7bb4fa 100755 --- a/src/nnsight/intervention/tracing/globals.py +++ b/src/nnsight/intervention/tracing/globals.py @@ -10,14 +10,15 @@ def _ensure_mounted(): - """Mount Object.save as the universal `.save` method. + """Mount Object.save / Object.carry as the universal `.save` / `.carry` methods. - Lazy one-time setup. Called from .save() / nnsight.save() so we only - pay the C-level mount cost once, on first use. + Lazy one-time setup, run at trace setup (``_setup_interleaver``) and from + ``.save()`` / ``.carry()`` so we only pay the C-level mount cost once. """ global _mounted if CONFIG.APP.PYMOUNT and not _mounted: mount(Object.save, "save") + mount(Object.carry, "carry") _mounted = True @@ -28,6 +29,25 @@ def save(object: Any): return object +def carry(object: Any): + """Mark ``object`` to be handed to a later trace in the same ``model.session()`` + WITHOUT surfacing it as a saved output. + + The portable counterpart to relying on an inner trace's locals flowing implicitly: + in-process a non-saved value already crosses to the next trace, but under isolation + each inner trace runs in a worker that only ships its outputs home, so a non-saved + value would vanish. ``.carry()`` explicitly registers the value to cross the boundary + — so the same code is correct in-process AND isolated. Unlike :func:`save`, a carried + value is dropped at session exit (it is not in ``Globals.saves``), so it never appears + in the caller's frame; use it for cross-trace handoffs (an activation to patch into a + later run) that are not themselves results. + """ + + Globals.shared.add(id(object)) + + return object + + class Object(torch.Tensor): def save(self, _=0): @@ -46,6 +66,24 @@ def save(self, _=0): return self + def carry(self, _=0): + """Hand this value to a later trace in the same ``model.session()`` without saving + it as an output. See :func:`carry`. + + Examples: + + >>> with model.session(): + ... with model.trace("clean prompt"): + ... act = model.transformer.h[6].output.carry() # not a result + ... with model.trace("corrupt prompt") as tracer: + ... tracer.patch(model.transformer.h[6], act) # transplanted in + ... logits = model.lm_head.output.save() + """ + + carry(self) + + return self + def __getattr__(self, name: str) -> Self: return super().__getattr__(name) @@ -100,10 +138,14 @@ def clear(self): class Globals: """Process-wide tracing state. - Holds two pieces of true global state: + Holds these pieces of true global state: - ``saves``: set of ``id()`` for objects marked via ``.save()``. The root tracer's ``push()`` filters its frame locals against this set so only saved values propagate out of the trace. + - ``shared``: set of ``id()`` for objects marked via ``.carry()`` — + cross-trace handoffs within a ``model.session()`` that are NOT surfaced + as outputs. Consulted by the isolated worker's ``end()`` to ship carried + values across the boundary; dropped at session exit (not in ``saves``). - ``cache``: source/AST/code-object memoization across traces. Root-vs-inner detection lives on the tracer itself — see @@ -113,9 +155,12 @@ class Globals: saves = set() + shared = set() + cache = TracingCache() @staticmethod def clear(): Globals.saves.clear() + Globals.shared.clear() Globals.cache.clear() From f7d134d531f91685b7305ca528b391c20ddd0bff Mon Sep 17 00:00:00 2001 From: khaiwang Date: Sun, 21 Jun 2026 23:51:37 -0400 Subject: [PATCH 25/30] fix(isolation): defer seccomp lockdown to after the first job's deserialize MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The isolated worker installed its seccomp filter before the job loop (warm time), so the very first conn.recv() — which unpickles the host's job message with standard pickle — triggered a lazy transformers submodule import (transformers loads modeling submodules only at unpickle time, so preimport=("transformers",) did NOT help) → open() → EPERM under seccomp. EPERM is an OSError, which the recv loop's `except (EOFError, OSError): break` swallowed → os._exit(0) → the host saw a pipe EOF and reported "worker died during execution." This broke every lockdown=True trace (root cause confirmed by worker-side instrumentation: death at conn.recv, importing transformers/models/gpt2/__init__.py). The job message and mediator payload are host-authored, TRUSTED data; only the user intervention code is untrusted. So lockdown belongs after deserialization, before user code — exactly what _sandbox.lock_down's own docstring already stated. Move lock_down() out of _pool_worker_main's warm section into _run_one_job, installed once (guarded by a worker global) after the first job's payload is deserialized and before its intervention runs: - the first conn.recv runs unlocked, so a fresh worker's first job needs no preimport=; - one-way + once, so a warm pool's later jobs run under the first job's lockdown — a homogeneous model needs nothing (already imported), a different model needs preimport=; - cold (pool_size=0) and pooled share the path, so both deserialize their first job first. Containment is unchanged: user-code open/socket/exec under lockdown are still blocked and now surface as a clean NNsightException (shipped via the EXCEPTION path) rather than a silent death. Verified (nnsight-tf, GPU7, gpt2): test_isolated_lockdown_safety.py 4/4 — read under lockdown max|Δ|=0, fs/net blocked, and a NEW warm-pool case (3 traces on one pooled worker under lockdown, all bit-identical). No regression: trace/cache/pool/session_handoff pass (the lockdown=False default path is untouched). Docs: mediator-gpu-trace-integration.md support matrix + lockdown-ordering notes updated (break -> fixed). Co-Authored-By: Claude Opus 4.8 (1M context) Claude-Session: https://claude.ai/code/session_014ZUUF44B2tfuKBNFhDFedR --- .../mediator-gpu-trace-integration.md | 24 ++++++---- .../test_isolated_lockdown_safety.py | 24 ++++++++++ src/nnsight/intervention/isolation.py | 47 ++++++++++++++----- 3 files changed, 72 insertions(+), 23 deletions(-) diff --git a/docs/developing/mediator-gpu-trace-integration.md b/docs/developing/mediator-gpu-trace-integration.md index 0f25342be..e103e1cbd 100644 --- a/docs/developing/mediator-gpu-trace-integration.md +++ b/docs/developing/mediator-gpu-trace-integration.md @@ -249,7 +249,7 @@ benign CudaIPC release warning. | read / swap (`=`) / `.save()` (tensors) / skip / exception | six events over the channel; host-side hook registration; worker→host saves transmission at END | ✅ bit-identical | | multi-invoke + batch narrowing | per-invoke worker + host mediator; Batcher stays host-side | ✅ bit-identical | | single-forward `generate(...)` (no iter) | same as trace | ✅ verified | -| seccomp lockdown (fs/net/exec) | `_sandbox.lock_down` after warm-up | ⚠️ broken since the warm-pool unification (§14): the worker locks down BEFORE receiving its first job, and unpickling the job's extras (Tokenizer) triggers a new `transformers` submodule import that seccomp blocks → the worker dies at job-recv. Pre-dates the 2026-06-10 refactors (reproduced on the pre-refactor commit). Mitigation now available: list the needed modules in `preimport=` so they load before lockdown ([threat-models](mediator-threat-models.md) §8); a default fix (auto-warm the model's transformers modules) is still open. | +| seccomp lockdown (fs/net/exec) | `_sandbox.lock_down` after the first job's deserialize, in `_run_one_job` | ✓ a read under `lockdown=True` is `max|Δ|=0`, user-code `open`/`socket`/`exec` are blocked, and a warm pool keeps serving (`test_isolated_lockdown_safety.py` 4/4). The break — the worker locked down BEFORE receiving its first job, so unpickling the job message (its `transformers` lazy submodules load only at unpickle time) hit a seccomp-blocked `open` (an `OSError` the recv loop swallowed → silent death) — is fixed by deferring lockdown to after the (host-authored, trusted) first job is deserialized, before its user code. One-way and installed once: in a warm pool, a later job needing a *new* model's modules still needs `preimport=` (a homogeneous model needs nothing). | | `iter`/`all`/`next` (multi-token) | step stamped in the requester; host iter-hooks bump the tracker; live `default_all` piggyback (§9) | ✅ bit-identical (`iter[N]`, `iter[:]`, per-step swap) | | `tracer.barrier()` | worker sends the target count; host accumulates participants + runs the coordination loop (§10) | ✅ | | `cross_invoker` variable sharing | host variable store; worker pushes data locals, pulls the merged store; transmittable data only (§10) | ✅ | @@ -431,10 +431,13 @@ model-weight-independent, linear in worker count) — and **MPS does not reduce an 80 GB A100 but ~55% of a 16 GB T4) — the cap must be deliberate, with the cold-spawn fallback past it. (`probe_pool_gpu_footprint.py`.) -**Lockdown + pool.** Seccomp lockdown happens once after warm-up (before the job loop), so a pooled worker -locks the import set at warm time — a job whose user code triggers a *new* import fails (deterministically, -since every worker shares the same warm-time import set). This is **stricter than the cold path** -(`pool_size=0`), which deserializes the mediator *before* lockdown and so allows deserialize-time imports. +**Lockdown + pool.** Seccomp lockdown is installed once, in `_run_one_job`, after the **first** job's +(host-authored, trusted) payload is deserialized and before its user code runs — so deserialization's own +imports (e.g. transformers' lazy modeling submodules, loaded only at unpickle time) succeed. It is one-way: +in a warm pool, later jobs run under the first job's lockdown, so a later job whose deserialize needs a +*new* import fails — true only across *different* models (a homogeneous model is already imported); pre-load +the deployment's model set via `preimport=` to serve heterogeneous models under lockdown. The cold path +(`pool_size=0`) is the same code, so it behaves identically (deserialize the one job, then lock down). Lockdown defaults off. **Hardening (independent review, 2026-06-08).** Both passes confirmed no Critical issue — the cross-request @@ -603,11 +606,12 @@ documented as unsupported. (`CudaIpcHostChannel.wait_event`) and measures worker think-time. Host-side Triton compilation happens during the host's *forward execution* — never while the host is in `wait_event` — so a multi-second cold MoE autotune never false-trips the worker's hang-detector, in either direction. -- *Lockdown ordering / cold-vs-pool.* `lock_down()` runs once after warm-up, before the job loop, in the - unified `_pool_worker_main` (the only worker entrypoint, used for cold via `poolable=False` and pooled via - `poolable=True`). So the import set is frozen at warm time for **both** paths — there is no - "cold deserializes before lockdown" advantage; the cold-vs-pool difference is recycle-vs-retire. (The - earlier §7 note describing a cold deserialize-before-lockdown window predates the warm-pool unification.) +- *Lockdown ordering / cold-vs-pool.* `lock_down()` runs once in `_run_one_job`, after the first job's + payload is deserialized and before its user code runs — in the unified `_pool_worker_main` (the only + worker entrypoint, cold via `poolable=False`, pooled via `poolable=True`). So **both** paths deserialize + their first job before lockdown; the cold-vs-pool difference is recycle-vs-retire. The first `conn.recv` + (which unpickles the job message) runs unlocked, so a fresh worker's first job needs no `preimport=`; a + warm pool's *later* jobs run under the first job's lockdown (so a different model would need `preimport=`). **Coverage:** `prototypes/mediator-sandbox/gpu_sandbox/test_isolated_triton_model.py` — `host_compiles` (isolated + `lockdown=True` Triton-kernel model bit-identical to in-process, with a cold `TRITON_CACHE_DIR` diff --git a/prototypes/mediator-sandbox/gpu_sandbox/test_isolated_lockdown_safety.py b/prototypes/mediator-sandbox/gpu_sandbox/test_isolated_lockdown_safety.py index 36feb68d2..8f10b8882 100644 --- a/prototypes/mediator-sandbox/gpu_sandbox/test_isolated_lockdown_safety.py +++ b/prototypes/mediator-sandbox/gpu_sandbox/test_isolated_lockdown_safety.py @@ -73,6 +73,29 @@ def test_net_blocked(model): return ok +def test_warm_pool_under_lockdown(model): + """Lockdown is one-way and installed after the FIRST job's deserialize, so a warm pool + must keep serving later jobs under that lockdown. Three traces on one pooled worker: + job 1 deserializes (imports the model's modules) then locks down; jobs 2-3 reuse the + locked worker — their deserialize needs no NEW import (same model), so they work.""" + from nnsight.intervention.isolation import shutdown_worker_pool + + with model.trace(PROMPT): + ref = model.transformer.h[6].output[0].save() + got = [] + try: + with isolate_mediators(fast_lane=False, lockdown=True, pool_size=1): + for _ in range(3): + with model.trace(PROMPT): + g = model.transformer.h[6].output[0].save() + got.append(g) + finally: + shutdown_worker_pool() + ok = len(got) == 3 and all(torch.equal(ref, g) for g in got) + print(f"[pool] 3 traces on one pooled worker under lockdown all bit-identical: {ok}") + return ok + + def main(): assert torch.cuda.is_available() model = LanguageModel("gpt2", device_map="cuda", dispatch=True) @@ -80,6 +103,7 @@ def main(): "func": test_functional_under_lockdown(model), "fs": test_fs_blocked(model), "net": test_net_blocked(model), + "pool": test_warm_pool_under_lockdown(model), } ok = all(results.values()) print("=" * 72) diff --git a/src/nnsight/intervention/isolation.py b/src/nnsight/intervention/isolation.py index eb7e7ff27..b2ec6029f 100644 --- a/src/nnsight/intervention/isolation.py +++ b/src/nnsight/intervention/isolation.py @@ -869,6 +869,12 @@ def push_locals(self) -> Optional[dict]: # ``worker_backward_context`` so BackwardsTracer can find the ambient mediator. _WORKER_CURRENT: Optional[WorkerMediator] = None +# Whether this worker has installed its (one-way) seccomp lockdown yet. Lockdown is +# deferred to the first job — after the job's host-authored payload is deserialized, before +# its user code runs — so the deserialize's imports (e.g. transformers' lazy submodules, +# loaded only at unpickle time) are not blocked. See ``_run_one_job``. +_LOCKED_DOWN: bool = False + def worker_backward_context() -> Optional[WorkerMediator]: """Return this worker's mediator if its trace contains a backward block, else None. @@ -883,15 +889,21 @@ def worker_backward_context() -> Optional[WorkerMediator]: return None -def _run_one_job(channel, payload, extras, opts, device) -> None: +def _run_one_job(channel, payload, extras, opts, device, lockdown: bool = False) -> None: """Deserialize one mediator against fresh dummies, adopt it as this job's :class:`WorkerMediator`, run its intervention, and ship saves at END. Any failure (including a bad payload) is reported as an EXCEPTION event so the host never - waits on a worker that won't speak.""" + waits on a worker that won't speak. + + Under ``lockdown`` the seccomp filter is installed HERE — after the (host-authored, + trusted) payload is deserialized, before the (untrusted) user intervention runs — so + deserialization's own imports succeed. It is one-way and installed once; in a warm pool + later jobs run under the first job's lockdown, so their deserialize must not need a NEW + import (true for a homogeneous model; otherwise pre-load via ``preimport=``).""" from .interleaver import Events from .tracing.globals import Globals - global _WORKER_CURRENT + global _WORKER_CURRENT, _LOCKED_DOWN try: Globals.saves.clear() # per-job reset (the only worker-side global state) Globals.shared.clear() @@ -903,6 +915,15 @@ def _run_one_job(channel, payload, extras, opts, device) -> None: channel.on_meta = mediator.apply_meta channel.push_provider = mediator.push_locals + # Footgun containment: seccomp-block new fs/net/exec syscalls NOW — the trusted + # payload is deserialized (its imports done), the untrusted user code is next. CUDA + # + the control Pipe + the IPC buffer use already-open fds, so they keep working. + if lockdown and not _LOCKED_DOWN: + from ._sandbox import lock_down + + lock_down() + _LOCKED_DOWN = True + mediator.intervention(mediator, mediator.info, *mediator.args) except BaseException as e: # noqa: BLE001 — contain the footgun; report it try: @@ -960,15 +981,15 @@ def _pool_worker_main(conn, buf, worker_iso_opts: IsoOptions): channel = CudaIpcWorkerChannel(conn, buf) # persistent; rebinds handlers per job - # Footgun containment: after CUDA is warm and base imports are done (both may open - # files), seccomp-block new fs/net/exec syscalls. Under the pool this locks the - # import set for ALL jobs (jobs whose user code triggers a NEW import will fail) — - # lockdown defaults off; document the trade-off. CUDA + the control Pipe + the IPC - # buffer use already-open fds. - if worker_iso_opts.lockdown: - from ._sandbox import lock_down - - lock_down() + # Lockdown is NOT installed here: the job message (received + unpickled by ``conn.recv`` + # below) and the mediator payload are host-authored, trusted data whose deserialization + # may need new imports (transformers loads its modeling submodules lazily, only at + # unpickle time). Installing seccomp before the first ``conn.recv`` blocks those opens + # and the worker dies (EPERM is an ``OSError`` → swallowed by the recv loop's + # ``except``). Instead ``_run_one_job`` installs lockdown after the first job's payload + # is deserialized, before its user code runs. The first ``conn.recv`` therefore runs + # unlocked; a warm pool's later recvs run under the first job's lockdown, so a NEW model + # type would need ``preimport=`` (a homogeneous model needs nothing — already imported). # One-time ready ack: the spawner consumes this before the channel reads protocol. conn.send("ready") @@ -983,7 +1004,7 @@ def _pool_worker_main(conn, buf, worker_iso_opts: IsoOptions): if not (isinstance(msg, tuple) and msg and msg[0] == "job"): continue # ignore stray control messages _, payload, extras, opts = msg - _run_one_job(channel, payload, extras, opts, device) + _run_one_job(channel, payload, extras, opts, device, worker_iso_opts.lockdown) # Job done; the worker is idle and recyclable. Loop for the next job/stop. # Skip interpreter atexit handlers: under seccomp lockdown, tempfile's atexit From 0852817141a486fd2c8e8946ec3bd0edf8e359d8 Mon Sep 17 00:00:00 2001 From: khaiwang Date: Mon, 22 Jun 2026 00:43:03 -0400 Subject: [PATCH 26/30] docs(isolation): sync integration support matrix with part-2 primitives, session handoff, backward state MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Bring docs/developing/mediator-gpu-trace-integration.md §8 (and the §10 cross-trace note) up to date with the features landed this session (the detail already lived in fast-lane.md §6): - support matrix: the part-2 primitive row now lists all of unembed/steer/patch/ablate; a new row documents session cross-trace handoff (.save() used in a later trace, and .carry()/nnsight.carry). - backward row: reflects the current state — read-path bit-identical; multi-token backward is a clean-fail (in-process doesn't support it either); grad-through-a-swap cleanly errors (the swapped value is a host-side leaf, severing the host graph at the seam) — the next backward increment. - §10 cross-trace note: the per-job reset clears Globals.shared too, and clarifies that the no-leak property is about UNRELATED traces — intentional in-session handoff is a separate supported path. Docs-only; no code change. Co-Authored-By: Claude Opus 4.8 (1M context) Claude-Session: https://claude.ai/code/session_014ZUUF44B2tfuKBNFhDFedR --- docs/developing/mediator-gpu-trace-integration.md | 12 ++++++++---- 1 file changed, 8 insertions(+), 4 deletions(-) diff --git a/docs/developing/mediator-gpu-trace-integration.md b/docs/developing/mediator-gpu-trace-integration.md index e103e1cbd..91462941c 100644 --- a/docs/developing/mediator-gpu-trace-integration.md +++ b/docs/developing/mediator-gpu-trace-integration.md @@ -254,9 +254,10 @@ benign CudaIPC release warning. | `tracer.barrier()` | worker sends the target count; host accumulates participants + runs the coordination loop (§10) | ✅ | | `cross_invoker` variable sharing | host variable store; worker pushes data locals, pulls the merged store; transmittable data only (§10) | ✅ | | warm worker pool (`pool_size=N`, `warm_worker_pool`) | generic workers receive serialized mediators as jobs; clean-END recycle (§14) | ✅ ~21× faster per request once warm | -| `with tensor.backward()` / `.grad` | BACKWARD event: worker seeds `dL/d(delivered clone)`, host continues `torch.autograd.grad` on the real graph, `.grad` by provenance path (§16) | ✅ read-path bit-identical (scalar loss; single invoke; no swap-then-backward; `.grad` editing raises) | +| `with tensor.backward()` / `.grad` | BACKWARD event: worker seeds `dL/d(delivered clone)`, host continues `torch.autograd.grad` on the real graph, `.grad` by provenance path (§16) | ✅ read-path bit-identical (scalar loss; single invoke); multi-token backward is a clean-fail (in-process doesn't support it either); **grad through a swap** cleanly errors (the swapped value is a host-side leaf, so the host graph is severed at the seam) — §16 | | `tracer.cache()` (`modules=`, `include_inputs=`) | CACHE event → host registers the real cache hooks; host CacheDict swapped in at END, filled in-place by the forward (§15) | ✅ bit-identical | -| `tracer.unembed` / `tracer.steer` | host-routed weight read (UNEMBED event) / replacement-swap injection (rides SWAP) — [fast-lane.md](fast-lane.md) §6 | ✅ bit-identical (isolated and in-process) | +| part-2 primitives: `tracer.unembed` / `tracer.steer` / `tracer.patch` / `tracer.ablate` | host-routed weight read (UNEMBED event); replacement-swap injection/transplant/knockout (ride SWAP, no new event) — [fast-lane.md](fast-lane.md) §6 | ✅ bit-identical (isolated and in-process) | +| session cross-trace handoff (`.save()` used in a later trace; `.carry()` / `nnsight.carry(x)`) | inner-trace END writeback to the session frame: saved values re-registered host-side so the session exit-push keeps them; carried (non-saved) values written for the next trace only — [fast-lane.md](fast-lane.md) §6 | ✅ bit-identical (`.carry()` is portable: harmless in-process, load-bearing under isolation) | | Triton-kernel models (MoE / SSM / `torch.compile`) | host-side forward compiles/runs Triton unrestricted (the worker holds only the intervention) | ✅ — §17 | | user-code Triton (kernel inside the intervention) | — (compiling a kernel needs `open`/`subprocess`/ptxas, which lockdown blocks by design) | ⛔ under lockdown — §17 | | `.source` operation-level access (`...attn.split_1.output`) | — (op paths aren't in `model.modules()`) | 🔜 not yet | @@ -324,8 +325,11 @@ regression PASS (4-tuple event). **Findings:** the variable-sharing push must filter to transmittable data (else the `Barrier` local pulls in the whole model) AND move tensors to CPU (CUDA-IPC tensors can't be re-shared by the host). The `_xinvoke_store` is per-interleaver (reset each trace) so no cross-trace leak today; the warm pool (§14) -rebuilds the interleaver + dummy modules per job and clears `Globals.saves`, so there is no cross-trace -leak under reuse either. Known coverage gaps: no tests yet for +rebuilds the interleaver + dummy modules per job and clears `Globals.saves`/`Globals.shared`, so there is +no cross-trace leak between *unrelated* traces under reuse either. (Intentional cross-trace handoff +*within a `model.session()`* — a `.save()`d value used in a later trace, or `.carry()` — is a separate, +supported path: the worker ships those values to the session frame at END; see the §8 support matrix and +[fast-lane.md](fast-lane.md) §6.) Known coverage gaps: no tests yet for multi-barrier-in-one-trace, 3+ participants, multi-token + barrier, or variable sharing without a barrier; the store grows monotonically per trace and ships all shared tensors CPU-serialized on every response (a perf cliff for large cross-invoke tensors). From 7a4c819343041fb000e08e8231c30c80896c9593 Mon Sep 17 00:00:00 2001 From: khaiwang Date: Tue, 23 Jun 2026 00:07:22 -0400 Subject: [PATCH 27/30] =?UTF-8?q?feat(isolation):=20gradient=20through=20a?= =?UTF-8?q?=20swap=20under=20isolation=20=E2=80=94=20stitch=20the=20backwa?= =?UTF-8?q?rd=20seam?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit An isolated SWAP installs the worker-computed value on the host as a detached leaf (clone-on-receive strips grad_fn), severing the host autograd graph at the swap point. So a downstream loss differentiated w.r.t. an UPSTREAM activation dead-ended at the swap ("no gradient available ... off the backward path"), while in-process gradients flow through swaps. The read-path backward split the chain rule once at the worker→host seam; a swap adds a second seam that splits it the other way (host downstream → worker swap tape → host pre-swap). Fix: iterate the existing Events.BACKWARD exchange to a fixpoint over swap seams. - Host (interleaver.py): handle_swap_event, under _iso_backward, makes the swap leaf requires_grad_(True) and retains it (_iso_grad_swaps) so the downstream forward tracks it and it is a backward target; handle_backward_event adds swap leaves to its targets and returns dL/d(swap leaf) under a reserved key (kept separate from reals so a read-then- swapped module sharing one requester path doesn't collide). - Worker (isolation.py): WorkerMediator.swap keeps the worker-tape swap value (with grad_fn); reset alongside the other _bwd state. - Worker backward (backwards.py): loop — send seeds, receive dL/d(swap leaf), backprop it through the swap tape to dL/d(delivered clone), re-seed the pre-swap graph, repeat; a read reached both directly and through a swap SUMS its gradient across rounds. With no swaps the loop is exactly the prior single exchange. Verified (nnsight-tf, GPU7, gpt2 + renamed): test_isolated_grad_through_swap.py 5/5 all max|Δ|=0 — grad through h*2, h+vec, tracer.steer, TWO chained swaps (loop fixpoint), and the renamed model, isolated == in-process. No regression across the isolated suite 13/13 (read-path backward, multi-token-backward clean-fail parity, trace/steer/patch/ablate swaps without backward, multitoken-iter/cross-invoke/session-handoff/cache/lockdown/acceptance/pool). Docs: mediator-gpu-trace-integration.md §16 + support matrix (grad-through-swap DONE). Co-Authored-By: Claude Opus 4.8 (1M context) Claude-Session: https://claude.ai/code/session_014ZUUF44B2tfuKBNFhDFedR --- .../mediator-gpu-trace-integration.md | 22 ++- .../test_isolated_grad_through_swap.py | 159 ++++++++++++++++++ src/nnsight/intervention/interleaver.py | 53 ++++-- src/nnsight/intervention/isolation.py | 12 ++ src/nnsight/intervention/tracing/backwards.py | 35 +++- 5 files changed, 260 insertions(+), 21 deletions(-) create mode 100644 prototypes/mediator-sandbox/gpu_sandbox/test_isolated_grad_through_swap.py diff --git a/docs/developing/mediator-gpu-trace-integration.md b/docs/developing/mediator-gpu-trace-integration.md index 91462941c..ed8ec0539 100644 --- a/docs/developing/mediator-gpu-trace-integration.md +++ b/docs/developing/mediator-gpu-trace-integration.md @@ -254,7 +254,7 @@ benign CudaIPC release warning. | `tracer.barrier()` | worker sends the target count; host accumulates participants + runs the coordination loop (§10) | ✅ | | `cross_invoker` variable sharing | host variable store; worker pushes data locals, pulls the merged store; transmittable data only (§10) | ✅ | | warm worker pool (`pool_size=N`, `warm_worker_pool`) | generic workers receive serialized mediators as jobs; clean-END recycle (§14) | ✅ ~21× faster per request once warm | -| `with tensor.backward()` / `.grad` | BACKWARD event: worker seeds `dL/d(delivered clone)`, host continues `torch.autograd.grad` on the real graph, `.grad` by provenance path (§16) | ✅ read-path bit-identical (scalar loss; single invoke); multi-token backward is a clean-fail (in-process doesn't support it either); **grad through a swap** cleanly errors (the swapped value is a host-side leaf, so the host graph is severed at the seam) — §16 | +| `with tensor.backward()` / `.grad` | BACKWARD event: worker seeds `dL/d(delivered clone)`, host continues `torch.autograd.grad` on the real graph, `.grad` by provenance path; **grad-through-swap** iterates the exchange — host returns `dL/d(swap leaf)`, worker backprops it through its swap tape, re-seeds the pre-swap graph (§16) | ✅ read-path AND grad-through-swap bit-identical (scalar loss; single invoke); multi-token backward is a clean-fail (in-process doesn't support it either) | | `tracer.cache()` (`modules=`, `include_inputs=`) | CACHE event → host registers the real cache hooks; host CacheDict swapped in at END, filled in-place by the forward (§15) | ✅ bit-identical | | part-2 primitives: `tracer.unembed` / `tracer.steer` / `tracer.patch` / `tracer.ablate` | host-routed weight read (UNEMBED event); replacement-swap injection/transplant/knockout (ride SWAP, no new event) — [fast-lane.md](fast-lane.md) §6 | ✅ bit-identical (isolated and in-process) | | session cross-trace handoff (`.save()` used in a later trace; `.carry()` / `nnsight.carry(x)`) | inner-trace END writeback to the session frame: saved values re-registered host-side so the session exit-push keeps them; carried (non-saved) values written for the next trace only — [fast-lane.md](fast-lane.md) §6 | ✅ bit-identical (`.carry()` is portable: harmless in-process, load-bearing under isolation) | @@ -501,7 +501,7 @@ untested. `modules=None` (cache *all* modules) registers a hook per module on th --- -## 16. `with tensor.backward()` — read-path DONE (2026-06-10) +## 16. `with tensor.backward()` — read-path DONE (2026-06-10), grad-through-swap DONE (2026-06-23) **The gap (§11).** Clone-on-receive strips `grad_fn` — the worker's delivered activations are detached clones, the autograd graph lives only on the host, and `.grad` providers were keyed by `id(tensor)` @@ -538,12 +538,22 @@ renamed model (`final_norm`/`output_projection`) `max|Δ|=0`; user-derived-tenso Independent review (7 finder angles, dedup + verify): **no silent-wrong in the in-scope path** (single invoke, on-path tensor-output target, scalar loss, no swaps). +**Gradient-through-swap — DONE (2026-06-23).** An isolated SWAP installs a worker-computed value as a host +*leaf* (clone-on-receive strips `grad_fn`), so the host backward used to dead-end at the swap — while +in-process gradients flow through swaps. Now the seam is stitched by iterating the existing `Events.BACKWARD` +exchange to a fixpoint. Host: `handle_swap_event` (under `_iso_backward`) makes the swap leaf +`requires_grad_(True)` and retains it (`_iso_grad_swaps`), so the downstream forward tracks it and it is a +backward target; `handle_backward_event` adds swap leaves to its targets and returns `dL/d(swap leaf)` under +a reserved key. Worker: `WorkerMediator.swap` keeps the worker-tape swap value (with `grad_fn`); the backward +block loops — send seeds, receive `dL/d(swap leaf)`, backprop it through the swap tape to `dL/d(delivered +clone)`, re-seed the pre-swap graph, repeat — accumulating each read's gradient across rounds (a clone +reached both directly and through a swap sums both paths). With no swaps the loop is the original single +exchange. Chained swaps converge in N rounds. **Verified (`test_isolated_grad_through_swap.py`, gpt2 + +renamed):** grad through `h*2`, `h+vec`, `tracer.steer`, and TWO chained swaps, plus the renamed model, all +isolated-vs-in-process `max|Δ|=0`. + **Limits / open:** - Scalar loss only — `loss.backward(gradient=...)` is not honored (scalar-only error). -- **Gradient-through-swap unsupported** (next increment): an isolated SWAP splices a graph *leaf* into - the host forward (worker-computed values carry no host `grad_fn`), so the host backward dead-ends at - the swap — while in-process gradients DO flow through swaps. Fix = recursive seam stitch (host ships - `dL/d(swap)` to the worker, the worker tape backprops to its leaves, ships back). - Batched traces error with a cryptic shape mismatch (host retains the full-batch tensor, worker seeds the narrowed clone) — needs a clear error or narrowed retention. - Multi-token backward: **not supported in-process either** (characterized 2026-06-10, diff --git a/prototypes/mediator-sandbox/gpu_sandbox/test_isolated_grad_through_swap.py b/prototypes/mediator-sandbox/gpu_sandbox/test_isolated_grad_through_swap.py new file mode 100644 index 000000000..4d27dbd57 --- /dev/null +++ b/prototypes/mediator-sandbox/gpu_sandbox/test_isolated_grad_through_swap.py @@ -0,0 +1,159 @@ +#!/usr/bin/env python3 +"""Gradient THROUGH a swap under isolation == in-process, bit-identical. + +A SWAP (replacement write) installs a worker-computed value on the host as a detached +*leaf*, so the host autograd graph is severed at the swap seam: a downstream loss +differentiated w.r.t. an UPSTREAM activation dead-ends at the swap (it returns no gradient), +while in-process the gradient flows through the swap. The fix stitches the seam: the host +returns dL/d(swap leaf), the worker backprops through its swap tape to dL/d(delivered clone), +and the host continues the pre-swap backward — iterated to a fixpoint, grads summed across +rounds (a clone reached both directly and through a swap contributes via both paths). + + mul — swap h[6].output := 2*h[6].output; grad of an upstream h[3].output matches + in-process (the *2 factor propagates back through the seam). + add — swap := h[6].output + vec (additive, like steering); grad of h[3] matches. + steer — tracer.steer at h[6], then backward; grad of h[3] matches (steer is a swap). + two_swaps — swaps at h[4] AND h[7]; grad of h[2] flows through BOTH seams (loop fixpoint). + renamed — renamed model (decoder_blocks): grad through a swap matches in-process. + +Run: + CUDA_VISIBLE_DEVICES=7 PYTHONPATH=src \ + /disk/u/zikai/anaconda3/envs/nnsight-tf/bin/python -u \ + prototypes/mediator-sandbox/gpu_sandbox/test_isolated_grad_through_swap.py +""" +import sys + +import torch + +from nnsight import LanguageModel +from nnsight.intervention.isolation import isolate_mediators + +P = "The Eiffel Tower is in the city of" + + +def U(o): + return o[0] if isinstance(o, tuple) else o + + +def _eq(a, b): + return (torch.is_tensor(a) and torch.is_tensor(b) + and a.shape == b.shape and torch.equal(a, b)) + + +def _delta(a, b): + return (a - b).abs().max().item() if _eq(a, b) else float("nan") + + +def _both(build): + ref = build() + with isolate_mediators(fast_lane=False, timeout=30): + got = build() + return ref, got + + +def test_mul(model): + def build(): + with model.trace(P): + up = U(model.transformer.h[3].output) + mid = U(model.transformer.h[6].output) + model.transformer.h[6].output = mid * 2.0 + loss = model.lm_head.output.sum() + with loss.backward(): + g = up.grad.save() + return g + ref, got = _both(build) + ok = _eq(ref, got) + print(f"[mul] grad through h6:=2*h6, wrt h3, isolated==in-process={ok} (max|Δ|={_delta(ref, got)})", flush=True) + return ok + + +def test_add(model): + vec = torch.randn(768) + + def build(): + with model.trace(P): + up = U(model.transformer.h[3].output) + mid = U(model.transformer.h[6].output) + model.transformer.h[6].output = mid + vec.to(dtype=mid.dtype, device=mid.device) + loss = model.lm_head.output.sum() + with loss.backward(): + g = up.grad.save() + return g + ref, got = _both(build) + ok = _eq(ref, got) + print(f"[add] grad through h6:=h6+vec, wrt h3, isolated==in-process={ok} (max|Δ|={_delta(ref, got)})", flush=True) + return ok + + +def test_steer(model): + direction = torch.randn(768) + + def build(): + with model.trace(P) as tracer: + up = U(model.transformer.h[3].output) + tracer.steer(model.transformer.h[6], direction, 4.0) + loss = model.lm_head.output.sum() + with loss.backward(): + g = up.grad.save() + return g + ref, got = _both(build) + ok = _eq(ref, got) + print(f"[steer] grad through tracer.steer(h6), wrt h3, isolated==in-process={ok} (max|Δ|={_delta(ref, got)})", flush=True) + return ok + + +def test_two_swaps(model): + def build(): + with model.trace(P): + up = U(model.transformer.h[2].output) + m4 = U(model.transformer.h[4].output) + model.transformer.h[4].output = m4 * 1.5 + m7 = U(model.transformer.h[7].output) + model.transformer.h[7].output = m7 * 0.7 + loss = model.lm_head.output.sum() + with loss.backward(): + g = up.grad.save() + return g + ref, got = _both(build) + ok = _eq(ref, got) + print(f"[two_swaps] grad through h4 & h7 swaps, wrt h2, isolated==in-process={ok} (max|Δ|={_delta(ref, got)})", flush=True) + return ok + + +def test_renamed(): + rename = {"transformer.ln_f": "final_norm", "lm_head": "output_projection", + "transformer.h": "decoder_blocks"} + model = LanguageModel("gpt2", device_map="cuda", dispatch=True, rename=rename) + + def build(): + with model.trace(P): + up = U(model.decoder_blocks[3].output) + mid = U(model.decoder_blocks[6].output) + model.decoder_blocks[6].output = mid * 2.0 + loss = model.output_projection.output.sum() + with loss.backward(): + g = up.grad.save() + return g + ref, got = _both(build) + ok = _eq(ref, got) + print(f"[renamed] grad through swap (renamed), isolated==in-process={ok} (max|Δ|={_delta(ref, got)})", flush=True) + return ok + + +def main(): + assert torch.cuda.is_available() + model = LanguageModel("gpt2", device_map="cuda", dispatch=True) + results = { + "mul": test_mul(model), + "add": test_add(model), + "steer": test_steer(model), + "two_swaps": test_two_swaps(model), + "renamed": test_renamed(), + } + print("=" * 72, flush=True) + print(f"ISOLATED GRAD-THROUGH-SWAP: {results}", flush=True) + sys.exit(0 if all(results.values()) else 1) + + +if __name__ == "__main__": + main() diff --git a/src/nnsight/intervention/interleaver.py b/src/nnsight/intervention/interleaver.py index 31704376b..a2a2402eb 100755 --- a/src/nnsight/intervention/interleaver.py +++ b/src/nnsight/intervention/interleaver.py @@ -1000,6 +1000,12 @@ def __init__( # non-backward isolated trace pays nothing. self._iso_backward = False self._iso_grad_reals: dict = {} + # Grad-through-swap: each isolated SWAP installs a worker-computed value as a host + # *leaf* (clone-on-receive strips grad_fn), severing the host graph at the seam. + # Keep the (tagged, requires_grad) swap leaf per requester so handle_backward_event + # can return dL/d(swap leaf); the worker backprops that through its swap tape and + # re-seeds the pre-swap graph (a second BACKWARD round). Gated on ``_iso_backward``. + self._iso_grad_swaps: dict = {} # Fast lane: the static safety verdict (set in start when fast_lane is on) and the # wall-clock watchdog for a confirmed-safe in-process run. @@ -1188,6 +1194,7 @@ def cancel(self): # Retained on-graph activations (isolated backward) pin the autograd graph; # drop them at trace end rather than waiting for the mediator to be GC'd. self._iso_grad_reals = {} + self._iso_grad_swaps = {} # Disarm the fast-lane watchdog (backstop to the thread-target finally) so a # generous deadline can't fire into an unrelated later computation. if self._fastlane_watchdog is not None: @@ -1370,6 +1377,19 @@ def handle_swap_event(self, provider: Any, requester: Any, swap_value: Any): """ # If fulfilled by this processor, swap the value and respond with the value and continue processing events. if provider == requester: + # Grad-through-swap: the swap value arrived clone-on-receive (a detached host + # leaf), so the downstream host graph would dead-end here. Make the seam tensor + # require grad BEFORE the forward consumes it (so downstream tracks it and it is + # a backward target) and retain it; handle_backward_event returns dL/d(this leaf) + # for the worker to backprop through its swap tape. Tuple outputs: the residual is + # element [0]. + if self._iso is not None and self._iso_backward: + seam = swap_value[0] if isinstance(swap_value, tuple) else swap_value + if (torch.is_tensor(seam) and seam.is_leaf + and seam.is_floating_point() and not seam.requires_grad): + seam.requires_grad_(True) + self._iso_grad_swaps[requester] = seam + # Swap the value in the batcher. Might only replace a slice of the value if this mediator is part of a batch group. self.interleaver.batcher.swap(self.batch_group, swap_value) @@ -1551,8 +1571,11 @@ def handle_backward_event(self, seed: dict): ``torch.autograd.grad`` on the host's graph with the worker-supplied seeds. """ reals = self._iso_grad_reals + swaps = self._iso_grad_swaps - # The seeds: (real activation tensor, dL/d(activation) from the worker). + # The seeds: (real activation tensor, dL/d(activation) from the worker). Seeds key + # into the pre-swap reals — round 1 seeds the loss's delivered reads, round 2 (after + # the worker backprops a swap seam) seeds the pre-swap activation that was swapped. out_tensors, out_grads = [], [] for path, grad in seed.items(): real = reals.get(path) @@ -1560,14 +1583,16 @@ def handle_backward_event(self, seed: dict): out_tensors.append(real) out_grads.append(grad) - # The targets: every delivered activation still on the host graph. The worker - # reads a subset of these via ``.grad``; computing all keeps it source-agnostic. - target_paths = [ - p for p, r in reals.items() if torch.is_tensor(r) and r.requires_grad - ] - target_tensors = [reals[p] for p in target_paths] + # The targets: every delivered activation AND every swap leaf still on the host + # graph. Reals serve `.grad` reads (and round-2 re-seeding); swap leaves let the + # worker continue the chain rule through a swap seam (grad-through-swap). Computing + # all keeps it source-agnostic. Kept as separate lists so a real and a swap sharing + # one requester path (read then swapped) don't collide. + real_paths = [p for p, r in reals.items() if torch.is_tensor(r) and r.requires_grad] + swap_paths = [p for p, s in swaps.items() if torch.is_tensor(s) and s.requires_grad] + target_tensors = [reals[p] for p in real_paths] + [swaps[p] for p in swap_paths] - if reals and not target_tensors: + if (reals or swaps) and not target_tensors: # Activations were delivered but NONE is on a graph — the forward ran # without gradient tracking (e.g. generate() runs grad-less). Tell the # worker so its .grad reads blame the real cause, not "off the path". @@ -1583,9 +1608,16 @@ def handle_backward_event(self, seed: dict): allow_unused=True, retain_graph=True, ) - result = { - p: g for p, g in zip(target_paths, grads) if torch.is_tensor(g) + n = len(real_paths) + result = {p: g for p, g in zip(real_paths, grads[:n]) if torch.is_tensor(g)} + swap_grads = { + p: g for p, g in zip(swap_paths, grads[n:]) if torch.is_tensor(g) } + if swap_grads: + # The worker backprops these through its swap tape to dL/d(delivered clone) + # and re-seeds (a further BACKWARD round); kept under a reserved key so they + # are not mistaken for `.grad`-readable activation gradients. + result["__nnsight_swap_grads__"] = swap_grads self.respond(result) # ack -> worker's send() returns the grad dict return True @@ -1940,6 +1972,7 @@ def __setstate__(self, state): self._isolated_worker = False self._iso_backward = False self._iso_grad_reals = {} + self._iso_grad_swaps = {} self._iso_caches = {} self._fastlane_verdict = None self._fastlane_watchdog = None diff --git a/src/nnsight/intervention/isolation.py b/src/nnsight/intervention/isolation.py index b2ec6029f..352dd4aa4 100644 --- a/src/nnsight/intervention/isolation.py +++ b/src/nnsight/intervention/isolation.py @@ -771,6 +771,7 @@ def adopt(cls, mediator, channel, interleaver, opts: dict, device) -> "WorkerMed mediator._bwd_active = opts["backward_active"] mediator._bwd_prov = {} # id(delivered clone) -> requester string mediator._bwd_tagged = [] # delivered clones made to require grad + mediator._bwd_swaps = {} # requester string -> worker-tape swap value (grad-through-swap) interleaver.current = mediator return mediator @@ -780,6 +781,17 @@ def request(self, requester: str): self._tag_delivered(value, requester) return value + def swap(self, requester: str, value): + # Grad-through-swap: keep the worker-TAPE swap value (with grad_fn back to the + # delivered clones) before send() ships its detached copy. The backward block uses + # it to backprop the host-returned dL/d(swap leaf) into dL/d(delivered clone). Tuple + # outputs: the residual is element [0]. + if self._bwd_active: + seam = value[0] if isinstance(value, tuple) else value + if torch.is_tensor(seam) and seam.requires_grad: + self._bwd_swaps[requester] = seam + return super().swap(requester, value) + def _tag_delivered(self, value, requester: str) -> None: """Tag each delivered activation tensor with its requester provenance and make it require grad, so worker-side ops on it build the worker's half of the graph.""" diff --git a/src/nnsight/intervention/tracing/backwards.py b/src/nnsight/intervention/tracing/backwards.py index 2a5b03c62..a0fed97c0 100755 --- a/src/nnsight/intervention/tracing/backwards.py +++ b/src/nnsight/intervention/tracing/backwards.py @@ -139,6 +139,7 @@ def _execute_isolated(self, fn: Callable, worker_mediator): forward_mediator = worker_mediator provenance = worker_mediator._bwd_prov tagged = [t for t in worker_mediator._bwd_tagged if t.requires_grad] + swaps = worker_mediator._bwd_swaps # requester -> worker-tape swap value # Worker half of the chain rule: seed = dL/d(delivered leaf) for leaves the loss # actually depends on (allow_unused drops the rest). @@ -151,11 +152,35 @@ def _execute_isolated(self, fn: Callable, worker_mediator): if grad is not None: seed[provenance[id(leaf)]] = grad - # Host runs its half and returns dL/d(activation) keyed by requester string. - worker_grads = forward_mediator.send(Events.BACKWARD, seed) or {} - # The host signals "no graph at all" (forward ran without gradient tracking, - # e.g. generate()) distinctly from "this particular read is off the path". - no_graph = bool(worker_grads.pop("__nnsight_backward_no_graph__", False)) + # Stitch the chain across the boundary, iterating over swap seams. Each round the + # host returns dL/d(activation) for reads AND dL/d(swap leaf) (a swap installs a host + # leaf, severing the host graph at the seam). The worker backprops each swap-leaf + # grad through its swap tape to dL/d(delivered clone) and re-seeds the pre-swap graph + # (another BACKWARD round). A read reached both directly and through a swap sums its + # contributions across rounds. With no swaps this is the original single exchange. + worker_grads: dict = {} + no_graph = False + while seed: + resp = forward_mediator.send(Events.BACKWARD, seed) or {} + no_graph = no_graph or bool(resp.pop("__nnsight_backward_no_graph__", False)) + swap_grads = resp.pop("__nnsight_swap_grads__", {}) + for path, g in resp.items(): + if torch.is_tensor(g): + worker_grads[path] = g if path not in worker_grads else worker_grads[path] + g + # Next round's seed: backprop each returned swap-leaf grad through the worker's + # swap tape to its delivered-clone leaves. + seed = {} + for swap_path, sg in swap_grads.items(): + swapped = swaps.get(swap_path) + if not (tagged and torch.is_tensor(swapped) and torch.is_tensor(sg)): + continue + grads = torch.autograd.grad( + swapped, tagged, grad_outputs=sg, allow_unused=True, retain_graph=True + ) + for leaf, g in zip(tagged, grads): + if g is not None: + p = provenance[id(leaf)] + seed[p] = g if p not in seed else seed[p] + g mediator = BackwardsMediator(fn, self.info) interleaver = Interleaver([mediator], self) From 40d20c08ac711b9dfe75efa724341c2d8e21740b Mon Sep 17 00:00:00 2001 From: khaiwang Date: Wed, 24 Jun 2026 17:42:15 -0400 Subject: [PATCH 28/30] docs(isolation): reconcile mediator-sandbox docs with shipped code + downgrade two overclaims MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit A deep audit found doc-vs-code drift across the isolation docs. Code is correct; the docs lagged. Doc-only changes (verified against transport.py / isolation.py / _sandbox.py / fastlane.py): - threat-models §7 + line 11: the worker->host fix is the shipped closed value-algebra codec (transport._codec_dumps / _codec_loads), not a restricted unpickler. The documented _RestrictedUnpickler / _safe_loads do not exist (replaced in 043196c). Rewrote the Fix paragraph to the codec (closed algebra, pure data assembly, no find_class/REDUCE/pickle VM, size cap + bounds checks, BoundaryValueError at encode); deleted the "restricted-unpickler vs hand-rolled codec" subsection and replaced it with why a closed codec is stronger (no opcode can call anything, so the restricted-unpickler bypass class, e.g. CVE-2025-32434, cannot exist). Legacy-socket remediation pointer _safe_loads -> _codec_loads. - integration §14: pool_key is the 5-tuple (device, arena_bytes, gpu_mem_fraction, lockdown, preimport), not 4 (preimport was missing). - integration §3/§7/§9: back-patch renamed symbols ensure_provider -> ensure_isolated_provider, spawn_isolated_worker -> _spawn_worker, _worker_main -> _pool_worker_main (kept the one "previously _worker_main, now _pool_worker_main" history line intact). - integration: clarify set_per_process_memory_fraction caps the allocator pool (the 20 GB footgun), distinct from the ~0.55 GiB CUDA-context cost it does not reduce. - mediator-isolation-sandbox.md + gpu-sandbox.md: SUPERSEDED / pre-integration banners pointing to the authoritative integration + threat-model docs; dropped topk from the op list and noted "capture" shipped as .carry(). Posture claims downgraded to match the shipped footgun-containment model: - threat-models §4/§5: the shipped seccomp is a default-ALLOW denylist of 7 fs/net/exec syscalls (plus GPU mem-fraction), i.e. R1 + footgun containment. Full R2 (allowlist-default seccomp + ptrace/clone/fork + namespaces + cgroups) is designed, not built; R2's "determined adversary" closure is the roadmap target, matching gpu-sandbox.md. - fast-lane §7: the static gate rejects import/open/exec/socket AST nodes, but the fast lane runs in-process and allowlists numpy/torch calls, which is a footgun selector, not an adversarial boundary. Safety rests on the trust="local" cordon (unreachable for non-author code), made explicit. Co-Authored-By: Claude Opus 4.8 (1M context) Claude-Session: https://claude.ai/code/session_014ZUUF44B2tfuKBNFhDFedR --- docs/developing/fast-lane.md | 9 ++- docs/developing/gpu-sandbox.md | 8 +++ .../mediator-gpu-trace-integration.md | 20 +++--- docs/developing/mediator-isolation-sandbox.md | 15 +++- docs/developing/mediator-threat-models.md | 70 +++++++++++-------- 5 files changed, 81 insertions(+), 41 deletions(-) diff --git a/docs/developing/fast-lane.md b/docs/developing/fast-lane.md index 9599f870b..9c62f8d77 100644 --- a/docs/developing/fast-lane.md +++ b/docs/developing/fast-lane.md @@ -316,8 +316,13 @@ is transparent to the single-trace path — and the in-process core (`test_lm.py - **The process-global `sys.addaudithook` backstop.** Its own failure mode (a leaked thread-local flag arms it during the model's *own* forward → server-wide outage) makes it net-negative when the static - default-deny gate already makes imports / `open` / `exec` / `socket` statically impossible in - fast-laned code. Documented as future hardening; the static gate is the confirmation. + default-deny gate already rejects `import` / `open` / `exec` / `socket` *nodes* outright in fast-laned + code (those AST forms route to ISOLATE, so they cannot appear in the in-process body). The remaining + boundary is **provenance, not the gate**: the fast lane is reachable only for `trust="local"` + (author-local) interventions, never remote/NDIF-submitted code, and the numpy/torch call-allowlist is a + **footgun selector, not an adversarial boundary** (a torch or numpy call runs arbitrary compute and both + have escape hatches). So fast-lane safety rests on the `trust="local"` cordon; the static gate only keeps + honest code honest. Documented as future hardening. - **A frozen-namespace `Compartment`** (SES-style) for fast-lane execution. The first slice relies on the static pass + the `trust` cordon; namespace shadowing is a later refinement. - *(none — the part-2 primitive set is complete.)* `unembed`, `steer`, `patch`, `ablate` are built, and diff --git a/docs/developing/gpu-sandbox.md b/docs/developing/gpu-sandbox.md index f4c9f202c..ab017abac 100644 --- a/docs/developing/gpu-sandbox.md +++ b/docs/developing/gpu-sandbox.md @@ -3,6 +3,14 @@ **Status:** Built + tested (functional + safety pass) · **Date:** 2026-06-06 **Supersedes** the CPU-only transport work (harness-plan Phases 5/5b/6) for the GPU path. +> **Pre-integration prototype.** This page documents the standalone prototype under +> `prototypes/mediator-sandbox/gpu_sandbox/` (`sandbox.py` seccomp, `gpu_worker.py`, `gpu_sandbox.py`). The +> design has since landed in `src/nnsight/intervention/`: the prototype `sandbox.py` became `_sandbox.py`, +> wired through `isolation.py` + `transport.py` + `interleaver.py`. For the shipped behavior and current +> support matrix see [mediator-gpu-trace-integration.md](mediator-gpu-trace-integration.md) (authoritative) +> and [mediator-threat-models.md](mediator-threat-models.md) (security posture). The "why this design" +> rationale below still holds. + ## Why this design The threat model was relaxed to **contain footguns, not defeat a determined adversary** — stop a careless diff --git a/docs/developing/mediator-gpu-trace-integration.md b/docs/developing/mediator-gpu-trace-integration.md index ed8ec0539..ac99417b4 100644 --- a/docs/developing/mediator-gpu-trace-integration.md +++ b/docs/developing/mediator-gpu-trace-integration.md @@ -101,7 +101,7 @@ isolation the worker can't (no real module), so the host registers it. - **Where:** in the host `handle` loop, when a `VALUE`/`SWAP`/`SKIP` arrives for a requester `R` whose provider is not yet set up (the existing `handle_value_event` else-branch, `interleaver.py:1203`, that today does `history.add`/`restore_event`/`return False`). -- **What:** `ensure_provider(R)` — parse `R` (`"..i"`) → `iteration = N`, +- **What:** `ensure_isolated_provider(R)`: parse `R` (`"..i"`) → `iteration = N`, `kind`, `path`; resolve `envoy = root_envoy.get(path)` (`envoy.py:586`); call the **existing** `output_hook`/`input_hook(host_mediator, envoy._module, f"{path}.{kind}")` (`hooks.py:224`/`154`). The host-side mediator is passed, so the hook closure delivers over the channel. @@ -113,7 +113,7 @@ isolation the worker can't (no real module), so the host registers it. across traces and need explicit per-trace cleanup. Keeping `eproperty`/`hooks.py` unchanged was the reason to register on a dead dummy rather than guard `_hook` with an `interleaver.isolated` flag; revisit if dummy-module accumulation shows up in profiling. -- **Idempotency:** `ensure_provider` registers a given `R` once per mediator (tracked in a set on the +- **Idempotency:** `ensure_isolated_provider` registers a given `R` once per mediator (tracked in a set on the host mediator), mirroring the `current_provider`-skip logic in `requires_output`. For a single forward pass the iteration is always `0`, so `R` is `"..i0"`. Multi-token @@ -188,7 +188,7 @@ Single forward pass; single and multiple invokes; read (`.output`/`.input`/`.inp 2. **Isolated `Mediator.start`** (`interleaver.py`): a branch that spawns the worker (CUDA → `spawn`), ships `self` via source-serialization with the model→path-only-envoy persistent-object map, and a worker bootstrap that builds the interleaver stub + path-only envoy mirror + runs the intervention. -3. **Host-side on-demand hook registration** (`interleaver.py`): `ensure_provider(R)` in `handle`. +3. **Host-side on-demand hook registration** (`interleaver.py`): `ensure_isolated_provider(R)` in `handle`. 4. **Saves transmission + lifecycle** (`interleaver.py` + bootstrap): worker ships `Globals.saves`- filtered frame locals at END; host injects; `cancel()` kills the process. 5. **Opt-in** (§12): `CONFIG.APP.ISOLATE_MEDIATORS` + `nnsight.isolate_mediators()` context. @@ -206,7 +206,7 @@ Single forward pass; single and multiple invokes; read (`.output`/`.input`/`.inp ### Status: DONE (2026-06-06) Built and verified (TDD; harnesses in `prototypes/mediator-sandbox/gpu_sandbox/`): - **Code:** `transport.py` (`pack_cuda`/`unpack_cuda` + the two channel ends, clone-on-receive, per-wait - timeout); `isolation.py` (`isolate_mediators` + `spawn_isolated_worker` + `_worker_main` + + timeout); `isolation.py` (`isolate_mediators` + `_spawn_worker` + `_pool_worker_main` + `ensure_isolated_provider` + `_WorkerInterleaver`/`_WorkerPersistent`); `_sandbox.py` (seccomp `lock_down`, relocated from the prototype); `interleaver.py` seam (`_iso` field, isolated `Mediator.start` branch, the on-demand hook-registration call in `handle`, saves injection in @@ -277,7 +277,7 @@ gained three pieces: - **Per-step hook registration:** `ensure_isolated_provider` parses the step `N` from the requester and passes `iteration=N` to `output_hook`/`input_hook` (new optional param), so the hook fires on step N — not the host mediator's iteration. -- **Host iter-hooks:** `spawn_isolated_worker` calls `register_iter_hooks(host_mediator, real_model)` so +- **Host iter-hooks:** `_spawn_worker` calls `register_iter_hooks(host_mediator, real_model)` so the host `iteration_tracker` advances per forward. - **Live host→worker piggyback:** `default_all` (= `generate(max_new_tokens)`) is set *after* the worker spawns, so it's piggybacked on each response frame (`CudaIpcHostChannel.meta_provider` → @@ -303,7 +303,7 @@ Both isolated == in-process. can't count cross-invoke. `Barrier.__call__` (isolated) sends the TARGET count; the host accumulates participant names in `Interleaver._barrier_acc` and, once all arrive, runs the existing coordination loop (`handle_barrier_event` iterates the host mediators, `respond()`+`handle()` each over its own - channel). `Mediator._isolated_worker` (set in `_worker_main`) gates the worker-side behavior. + channel). `Mediator._isolated_worker` (set in `_pool_worker_main`) gates the worker-side behavior. - **Variable sharing (host store):** worker frames aren't shared across processes, so each worker pushes its *data* locals to `Interleaver._xinvoke_store` (a 4th `push` field on the event frame) and pulls the merged store back (piggybacked on the response, reusing the multi-token channel). Only **transmittable @@ -422,7 +422,8 @@ hook-registration state; the worker rebuilds its interleaver + dummy modules per **Opt-in.** `isolate_mediators(..., pool_size=N)` routes through the pool (`pool_size=0`, the default, is the unchanged cold-spawn path). `warm_worker_pool(N, ...)` pre-warms at startup (blocks until N ack ready); `shutdown_worker_pool()` tears it down. Workers are pooled **per (device, arena_bytes, gpu_mem_fraction, -lockdown) signature** — a worker is reused only for a matching signature, so a process hosting models on +lockdown, preimport) signature** (the `IsoOptions.pool_key` 5-tuple; `preimport` and `lockdown` are +warm-time, so they partition the pool); a worker is reused only for a matching signature, so a process hosting models on different GPUs gets a per-device sub-pool (NOT a shared pool whose bounce buffer is fixed to the first model's device — that would copy into the wrong-device buffer). Per-trace options (`default_all`, `cross_invoker`, `timeout`) ride each job. @@ -433,7 +434,10 @@ Each warm worker costs **~0.55 GiB GPU per GPU it touches** (CUDA context + cuBL model-weight-independent, linear in worker count) — and **MPS does not reduce this** (Ampere MPS shares the *scheduler*, not context memory; measured identical under MPS). So at batch-16/single-GPU ≈ 8.7 GB (11% of an 80 GB A100 but ~55% of a 16 GB T4) — the cap must be deliberate, with the cold-spawn fallback past it. -(`probe_pool_gpu_footprint.py`.) +(`probe_pool_gpu_footprint.py`.) Note this ~0.55 GiB is the fixed CUDA-context cost and is distinct from +`set_per_process_memory_fraction`: that knob caps the worker's **allocator pool** (so a runaway allocation +can't exhaust the device, the 20 GB-symptom footgun), not the context, so it does not reduce the per-worker +0.55 GiB. **Lockdown + pool.** Seccomp lockdown is installed once, in `_run_one_job`, after the **first** job's (host-authored, trusted) payload is deserialized and before its user code runs — so deserialization's own diff --git a/docs/developing/mediator-isolation-sandbox.md b/docs/developing/mediator-isolation-sandbox.md index dd7f0ee35..02eb394ed 100644 --- a/docs/developing/mediator-isolation-sandbox.md +++ b/docs/developing/mediator-isolation-sandbox.md @@ -1,9 +1,17 @@ # Mediator Isolation Sandbox — Design -**Status:** Draft / design (prototype scope) · **Date:** 2026-06-05 · **Author:** zikai +**Status:** SUPERSEDED (historical) · **Date:** 2026-06-05 · **Author:** zikai **Related:** `ndif` security regression suite (`src/services/ray/tests/security/`), `NDIF.md` §6–7, nnsight `src/nnsight/intervention/interleaver.py` (Mediator / event protocol). +> **Superseded.** This is the earlier CPU-only, two-tier (fork-jail) design from 2026-06-05. On 2026-06-06 +> the project chose the GPU-sandbox approach (contain footguns, not a determined adversary), built and +> documented in [gpu-sandbox.md](gpu-sandbox.md) and landed via the integration in +> [mediator-gpu-trace-integration.md](mediator-gpu-trace-integration.md), which is the authoritative +> current reference (plus [mediator-threat-models.md](mediator-threat-models.md) for the security posture). +> Kept for historical context; the CPU-transport and two-tier op-interpreter details below do not reflect +> the shipped code. + ## 1. Motivation NDIF executes **arbitrary user-submitted Python** on shared GPU infrastructure. Today that code runs @@ -283,8 +291,9 @@ pool. That caps pool size and the per-hook D2H/H2D budget (the per-hook D2H/H2D ## 7. Future optimization (out of scope now) Where do user tensor ops execute? The prototype picks **per-hook D2H/H2D in the jail, on CPU** (D2H/H2D per hook). -The destination is **the two-tier approach**: run the common tiny tensor algebra (read/write/project/ablate/steer/ -patch/topk/cache) on the host via a *validated op interpreter* (data never leaves the GPU, no jail in the +The destination is **the two-tier approach**: run the common tiny tensor algebra (read/swap/save/unembed/steer/ +patch/ablate/cache, plus the `.carry()` cross-trace handoff; the realized primitive set, "capture" shipped as +`.carry()`) on the host via a *validated op interpreter* (data never leaves the GPU, no jail in the hot path), and route only genuinely-arbitrary Python to the CPU jail. (B) GPU-in-jail via CUDA IPC is rejected (surface + breaks the pool); (C) fully symbolic host execution is a breaking change to nnsight's "real tensors" contract. The two-tier approach attacks the root — *arbitrary code next to the data* — rather than paying to diff --git a/docs/developing/mediator-threat-models.md b/docs/developing/mediator-threat-models.md index 3d2eff6e9..ae7f17bb3 100644 --- a/docs/developing/mediator-threat-models.md +++ b/docs/developing/mediator-threat-models.md @@ -7,8 +7,8 @@ earlier CPU-only threat model + AWS deployment probe). This page is the security companion to the GPU-worker backend: *what* adversary each isolation configuration faces, what it can and cannot contain, and the resulting **cost coupling** that makes -the "how much isolation" question concrete. It also records the one security fix landed in-branch — the -worker→host **restricted-unpickler** codec. +the "how much isolation" question concrete. It also records the one security fix landed in-branch: the +worker→host **closed value-algebra codec**. --- @@ -77,6 +77,16 @@ Two consequences: **R2** is the target (full host-OS isolation, fast path, accept the GPU-layer residual). If yes → you must pay **R3** for the co-batched path. +**Shipped vs designed (R2 is the target, not the current state).** The code today is **R1 plus +footgun-containment**: the spawned worker (shared GPU, CUDA-IPC), a seccomp **denylist** that EPERMs new +filesystem (`open`/`openat`/`openat2`), network (`socket`/`connect`), and exec (`execve`/`execveat`) +syscalls, and a GPU allocator memory-fraction cap. That contains careless or buggy interventions, matching +the relaxed model in [gpu-sandbox.md](gpu-sandbox.md) ("contain footguns, not a determined adversary"). +The full **R2** jail in the ladder above (an allowlist-default seccomp filter that also blocks +`ptrace`/`clone`/`fork`, plus net/mount/user namespaces and cgroups) is **designed, not yet built**, so +R2's "closes ... against even a determined adversary" line is a roadmap target, not what ships today. +Raising the posture to R2 is the §8 hardening work. + --- ## 5. Security controls vs the shared-GPU design @@ -91,7 +101,8 @@ lock-down-after-warm.) - **Network namespace** (no NIC) — zero GPU impact; kills exfil + IMDS. Take unconditionally. - **Mount namespace** (ro allowlist; bind only `/dev/nvidia*` + CUDA libs) — removes weights/secrets/ other-tenant disk from view. Highest value. -- **User namespace** (non-root, drop caps); **cgroups** (CPU/mem → OOM/forkbomb cap); **seccomp** (have it). +- **User namespace** (non-root, drop caps); **cgroups** (CPU/mem → OOM/forkbomb cap); **seccomp** (today a + footgun **denylist** of fs/net/exec syscalls; the R2 target is allowlist-default plus `ptrace`/`clone`/`fork`). **Group B — in tension / incompatible (forces the CPU-transport fallback or hardware partition):** - **PID/IPC namespace** — cross-namespace CUDA-IPC is documented to fail; `CLONE_NEWPID` must be entered @@ -153,31 +164,34 @@ which bypasses *every* other layer (seccomp, namespaces, row-bounding). The desi for the *inbound* user payload (unpickled inside the worker) but not the *outbound* results. This is a precondition for the isolation guarantee holding against a determined adversary. -**Fix.** Worker→host now never plain-`pickle.loads`. Tensors already travel out-of-band (GPU bounce -buffer / safetensors), so a frame is **tensor-free**; the remaining small structure is decoded with a -**restricted unpickler** (`transport._RestrictedUnpickler` / `_safe_loads`) whose `find_class` allows -**only torch `dtype`/`device`** and refuses every other class/function. `find_class` resolves a global -*before* the `REDUCE` that would call it, so a gadget (`os.system`, …) is refused before it can execute. -Supporting choices keep the allowlist minimal: the event rides as its string `.value` (no enum class), -and exceptions ride as a `(type-name, message)` sentinel (no class). **Host→worker stays normal pickle** -(host-authored, trusted). - -- **Why a restricted unpickler over a hand-rolled codec:** pickle handles all plain nested Python - structures natively (no per-type enumeration — an earlier hand-rolled JSON codec silently *missed* the - `Events.CACHE` event's `torch.dtype`/`torch.device` payload). Anything un-allowlisted now fails **loud - at decode with the exact refused class name**, easy to extend if legitimately safe. -- **Cost:** ~plain-pickle speed (µs for the tiny tensor-free frame); `find_class` fires only on the - handful of globals per frame. Negligible vs the ~0.6 ms/hook round-trip. -- **Why not `torch.load(weights_only=True)`:** it had an RCE bypass (CVE-2025-32434, fixed only in 2.6); - our unpickler is *tighter* — it enables no tensor-rebuild path at all (tensors aren't in the pickle). -- **Capability narrowing (documented):** `.save()` of an arbitrary object / numpy array / framework type - (e.g. `ModelOutput`) is no longer transmittable from a worker — save a tensor (or basic data) instead. - -Coverage: `prototypes/mediator-sandbox/gpu_sandbox/test_isolated_codec_security.py` — fidelity for -VALUE/SWAP/END/**CACHE (dtype+device)**/EXCEPTION/push, and a real `__reduce__` gadget refused at decode -without executing. (CPU is enough; needs torch.) The legacy AF_UNIX socket channels -(`SocketHostChannel`/`ShmSocketHostChannel`) are unused by `isolate_mediators` and still plain-unpickle — -route them through `_safe_loads` before wiring them to an untrusted worker. +**Fix.** Worker→host now never `pickle.loads` at all. Tensors already travel out-of-band (GPU bounce +buffer / safetensors), so a frame is **tensor-free**; the remaining small structure crosses as a **closed +value algebra** encoded/decoded by a hand-written codec (`transport._codec_dumps` / `_codec_loads`). The +algebra is exactly: `None | bool | int | float | str | bytes | list | tuple | dict | set | torch.dtype | +torch.device | Array (out-of-band)`. Decode is **pure data assembly**: no `find_class`, no `REDUCE`, no +global resolution, so there is no opcode that can construct or call an arbitrary object. It is bounded +(a `_MAX_CODEC_BYTES` cap and bounds-checked reads guard against decode-bomb / OOB), and a value outside +the algebra fails **loud at encode in the worker** (`BoundaryValueError`, naming the offending object), +not at decode. The event rides as its string `.value`, and exceptions ride as a `(type-name, message)` +sentinel, so neither pulls in a class. **Host→worker stays normal pickle** (host-authored, trusted). + +- **Why a closed codec rather than a restricted unpickler:** a restricted unpickler still runs the pickle + VM and gates classes at `find_class`; that gate has had RCE bypasses (e.g. `torch.load(weights_only=True)`, + CVE-2025-32434). Because the worker is the less-trusted side, running *any* pickle VM on its bytes is the + wrong shape. The codec removes the class of bug outright: there is no instruction that can call anything. + Per-type enumeration is the price, and it is small (the frame is tensor-free and the algebra is fixed). +- **Cost:** measured ~1.25–1.6x the old restricted-unpickler frame round-trip, swamped by the ~0.6 ms/hook + GPU context-switch, so negligible end to end. +- **Capability scope:** `.save()` / `.carry()` of a tensor, a numpy array, or basic scalars/containers + crosses; an arbitrary live object / framework type (e.g. `ModelOutput`) does not, and is refused at + encode with a clear error. + +Coverage: `prototypes/mediator-sandbox/gpu_sandbox/test_isolated_codec_security.py` covers fidelity for +VALUE/SWAP/END/**CACHE (dtype+device)**/EXCEPTION/push (including numpy + dtype/device), a `__reduce__` +gadget and a custom object both refused at **encode**, and malformed/oversized frames raising at decode. +(CPU is enough; needs torch.) The legacy AF_UNIX socket channels +(`SocketHostChannel`/`ShmSocketHostChannel`) are unused by `isolate_mediators` and still plain-unpickle: +route them through `_codec_loads` before wiring them to an untrusted worker. --- From 745575c270c3cf2de49c10ef8b4c98018a173e34 Mon Sep 17 00:00:00 2001 From: khaiwang Date: Sun, 28 Jun 2026 22:43:20 -0400 Subject: [PATCH 29/30] test(isolation): lock in batched backward (list input); correct the multi-invoke framing MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Investigating the "batched backward" open item showed it splits two ways, and checking the in-process baseline first reframed it: - Batched backward via LIST input (model.trace([A, B, ...])) already works under isolation, bit-identical. It is one mediator over the padded batch (batch_group=None, no per-invoke narrowing), so the worker's delivered clone and the host's retained real are both full-batch and shapes match; the read-path and grad-through-swap seam-stitch run unchanged on a (batch, seq, hidden) tensor. Added test_isolated_batched_backward.py (2-row, 3-row, upstream-block, batched grad-through-swap, renamed): all isolated-vs-in-process max|Δ|=0. - Backward inside MULTIPLE tracer.invoke(...) contexts raises MissedProviderError IN-PROCESS too (the .grad provider is never registered across invoke contexts), so it is a core nnsight limitation, not an isolation gap. The prior doc framing ("cryptic shape mismatch / needs narrowed retention") predated checking the in-process baseline; same category as multi-token backward (parity, not a gap). Doc-only + new test; no code change. mediator-gpu-trace-integration.md §16 + support-matrix backward row corrected. Co-Authored-By: Claude Opus 4.8 (1M context) Claude-Session: https://claude.ai/code/session_014ZUUF44B2tfuKBNFhDFedR --- .../mediator-gpu-trace-integration.md | 16 +- .../test_isolated_batched_backward.py | 146 ++++++++++++++++++ 2 files changed, 158 insertions(+), 4 deletions(-) create mode 100644 prototypes/mediator-sandbox/gpu_sandbox/test_isolated_batched_backward.py diff --git a/docs/developing/mediator-gpu-trace-integration.md b/docs/developing/mediator-gpu-trace-integration.md index ac99417b4..f7fe3a852 100644 --- a/docs/developing/mediator-gpu-trace-integration.md +++ b/docs/developing/mediator-gpu-trace-integration.md @@ -254,7 +254,7 @@ benign CudaIPC release warning. | `tracer.barrier()` | worker sends the target count; host accumulates participants + runs the coordination loop (§10) | ✅ | | `cross_invoker` variable sharing | host variable store; worker pushes data locals, pulls the merged store; transmittable data only (§10) | ✅ | | warm worker pool (`pool_size=N`, `warm_worker_pool`) | generic workers receive serialized mediators as jobs; clean-END recycle (§14) | ✅ ~21× faster per request once warm | -| `with tensor.backward()` / `.grad` | BACKWARD event: worker seeds `dL/d(delivered clone)`, host continues `torch.autograd.grad` on the real graph, `.grad` by provenance path; **grad-through-swap** iterates the exchange — host returns `dL/d(swap leaf)`, worker backprops it through its swap tape, re-seeds the pre-swap graph (§16) | ✅ read-path AND grad-through-swap bit-identical (scalar loss; single invoke); multi-token backward is a clean-fail (in-process doesn't support it either) | +| `with tensor.backward()` / `.grad` | BACKWARD event: worker seeds `dL/d(delivered clone)`, host continues `torch.autograd.grad` on the real graph, `.grad` by provenance path; **grad-through-swap** iterates the exchange (host returns `dL/d(swap leaf)`, worker backprops it through its swap tape, re-seeds the pre-swap graph) (§16) | ✅ read-path AND grad-through-swap bit-identical (scalar loss; single mediator, including list-batched `trace([A,B,...])`); multi-invoke-context and multi-token backward are clean-fail parity (in-process does not support them either) | | `tracer.cache()` (`modules=`, `include_inputs=`) | CACHE event → host registers the real cache hooks; host CacheDict swapped in at END, filled in-place by the forward (§15) | ✅ bit-identical | | part-2 primitives: `tracer.unembed` / `tracer.steer` / `tracer.patch` / `tracer.ablate` | host-routed weight read (UNEMBED event); replacement-swap injection/transplant/knockout (ride SWAP, no new event) — [fast-lane.md](fast-lane.md) §6 | ✅ bit-identical (isolated and in-process) | | session cross-trace handoff (`.save()` used in a later trace; `.carry()` / `nnsight.carry(x)`) | inner-trace END writeback to the session frame: saved values re-registered host-side so the session exit-push keeps them; carried (non-saved) values written for the next trace only — [fast-lane.md](fast-lane.md) §6 | ✅ bit-identical (`.carry()` is portable: harmless in-process, load-bearing under isolation) | @@ -505,7 +505,7 @@ untested. `modules=None` (cache *all* modules) registers a hook per module on th --- -## 16. `with tensor.backward()` — read-path DONE (2026-06-10), grad-through-swap DONE (2026-06-23) +## 16. `with tensor.backward()`: read-path DONE (2026-06-10), grad-through-swap DONE (2026-06-23), batched (list) verified (2026-06-28) **The gap (§11).** Clone-on-receive strips `grad_fn` — the worker's delivered activations are detached clones, the autograd graph lives only on the host, and `.grad` providers were keyed by `id(tensor)` @@ -558,8 +558,16 @@ isolated-vs-in-process `max|Δ|=0`. **Limits / open:** - Scalar loss only — `loss.backward(gradient=...)` is not honored (scalar-only error). -- Batched traces error with a cryptic shape mismatch (host retains the full-batch tensor, worker seeds - the narrowed clone) — needs a clear error or narrowed retention. +- **Batched backward (list input) works.** `model.trace([A, B, ...])` is one mediator over the padded + batch (no per-invoke narrowing, `batch_group=None`), so the worker's clone and the host's retained real + are both full-batch and their shapes match; the read-path and grad-through-swap seam-stitch run unchanged + on a `(batch, seq, hidden)` tensor. Verified bit-identical isolated-vs-in-process in + `test_isolated_batched_backward.py` (2-row, 3-row, upstream-block, batched grad-through-swap, renamed). +- Backward inside MULTIPLE `tracer.invoke(...)` contexts is **not supported in-process either**: it raises + `MissedProviderError` (the `.grad` provider is never registered across invoke contexts), so it is a core + nnsight limitation, not an isolation gap. (The earlier "shape mismatch / narrowed retention" framing + predated checking the in-process baseline.) The isolated path also fails; making that failure as clean as + the in-process error is minor parity hygiene, not a capability. - Multi-token backward: **not supported in-process either** (characterized 2026-06-10, `test_isolated_multitoken_backward.py`) — `generate()` runs the forward without gradient tracking, so the first `.grad` read fails in-process ("cannot register a hook on a tensor that doesn't require diff --git a/prototypes/mediator-sandbox/gpu_sandbox/test_isolated_batched_backward.py b/prototypes/mediator-sandbox/gpu_sandbox/test_isolated_batched_backward.py new file mode 100644 index 000000000..0e960cda6 --- /dev/null +++ b/prototypes/mediator-sandbox/gpu_sandbox/test_isolated_batched_backward.py @@ -0,0 +1,146 @@ +#!/usr/bin/env python3 +"""Batched backward under isolation == in-process, bit-identical. + +"Batched" here is a LIST input (`model.trace([A, B, ...])`): one mediator, one forward over +the padded batch, one backward, per-row gradients. There is no per-invoke narrowing (the +mediator's batch_group is None), so the worker's delivered clone and the host's retained real +are both full-batch and their shapes match: the read-path and grad-through-swap seam-stitch +work unchanged, just on a (batch, seq, hidden) tensor. + +(Backward inside MULTIPLE `tracer.invoke(...)` contexts is a separate, unsupported structure: +it raises `MissedProviderError` IN-PROCESS too, so it is a core nnsight limitation, not an +isolation gap; not covered here.) + + two_rows — grad of ln_f.output for a 2-prompt batch, isolated == in-process, shape (2,·,768). + three_rows — same for a 3-prompt batch. + upstream — grad of an upstream block (h[3].output) in a batched trace. + swap — batched + a swap at h[6]; grad of upstream h[3] flows through the seam (batched + grad-through-swap), isolated == in-process. + renamed — renamed model (decoder_blocks), batched backward. + +Run: + CUDA_VISIBLE_DEVICES=7 PYTHONPATH=src \ + /disk/u/zikai/anaconda3/envs/nnsight-tf/bin/python -u \ + prototypes/mediator-sandbox/gpu_sandbox/test_isolated_batched_backward.py +""" +import sys + +import torch + +from nnsight import LanguageModel +from nnsight.intervention.isolation import isolate_mediators + +A = "The Eiffel Tower is in the city of" +B = "A red bicycle was left near the river of" +C = "She quietly closed the heavy wooden front" + + +def U(o): + return o[0] if isinstance(o, tuple) else o + + +def _eq(a, b): + return (torch.is_tensor(a) and torch.is_tensor(b) + and a.shape == b.shape and torch.equal(a, b)) + + +def _delta(a, b): + return (a - b).abs().max().item() if _eq(a, b) else float("nan") + + +def _both(build): + ref = build() + with isolate_mediators(fast_lane=False, timeout=30): + got = build() + return ref, got + + +def test_two_rows(model): + def build(): + with model.trace([A, B]): + hs = U(model.transformer.ln_f.output) + with model.lm_head.output.sum().backward(): + g = hs.grad.save() + return g + ref, got = _both(build) + ok = _eq(ref, got) and ref.shape[0] == 2 + print(f"[two_rows] batched grad bit-identical={_eq(ref, got)} shape={tuple(ref.shape)} (max|Δ|={_delta(ref, got)})", flush=True) + return ok + + +def test_three_rows(model): + def build(): + with model.trace([A, B, C]): + hs = U(model.transformer.ln_f.output) + with model.lm_head.output.sum().backward(): + g = hs.grad.save() + return g + ref, got = _both(build) + ok = _eq(ref, got) and ref.shape[0] == 3 + print(f"[three_rows] batched grad bit-identical={_eq(ref, got)} shape={tuple(ref.shape)} (max|Δ|={_delta(ref, got)})", flush=True) + return ok + + +def test_upstream(model): + def build(): + with model.trace([A, B]): + up = U(model.transformer.h[3].output) + with model.lm_head.output.sum().backward(): + g = up.grad.save() + return g + ref, got = _both(build) + ok = _eq(ref, got) + print(f"[upstream] batched grad of h[3] bit-identical={ok} shape={tuple(ref.shape)} (max|Δ|={_delta(ref, got)})", flush=True) + return ok + + +def test_swap(model): + def build(): + with model.trace([A, B]): + up = U(model.transformer.h[3].output) + mid = U(model.transformer.h[6].output) + model.transformer.h[6].output = mid * 2.0 + loss = model.lm_head.output.sum() + with loss.backward(): + g = up.grad.save() + return g + ref, got = _both(build) + ok = _eq(ref, got) + print(f"[swap] batched grad-through-swap bit-identical={ok} shape={tuple(ref.shape)} (max|Δ|={_delta(ref, got)})", flush=True) + return ok + + +def test_renamed(): + rename = {"transformer.ln_f": "final_norm", "lm_head": "output_projection", + "transformer.h": "decoder_blocks"} + model = LanguageModel("gpt2", device_map="cuda", dispatch=True, rename=rename) + + def build(): + with model.trace([A, B]): + hs = U(model.final_norm.output) + with model.output_projection.output.sum().backward(): + g = hs.grad.save() + return g + ref, got = _both(build) + ok = _eq(ref, got) and ref.shape[0] == 2 + print(f"[renamed] batched grad (renamed) bit-identical={_eq(ref, got)} shape={tuple(ref.shape)} (max|Δ|={_delta(ref, got)})", flush=True) + return ok + + +def main(): + assert torch.cuda.is_available() + model = LanguageModel("gpt2", device_map="cuda", dispatch=True) + results = { + "two_rows": test_two_rows(model), + "three_rows": test_three_rows(model), + "upstream": test_upstream(model), + "swap": test_swap(model), + "renamed": test_renamed(), + } + print("=" * 72, flush=True) + print(f"ISOLATED BATCHED BACKWARD: {results}", flush=True) + sys.exit(0 if all(results.values()) else 1) + + +if __name__ == "__main__": + main() From b67e363ff6c4b188202c883d6f52cda43ae70834 Mon Sep 17 00:00:00 2001 From: khaiwang Date: Mon, 20 Jul 2026 13:22:24 -0400 Subject: [PATCH 30/30] refactor(isolation): host owns the iteration tag for in-loop requesters Converge the isolation seam with ndif2's process sandbox on its one structural improvement: a single iteration authority. Previously an isolated worker tagged every requester itself (from its tracer.iter pin; its per-path step counter is frozen at 0 because its modules are dummies) while the host separately advanced its own step counter to gate hook firing, so two counters had to stay in agreement across the boundary. Now a requester issued inside a tracer.iter scope crosses UNTAGGED as a marker tuple (base path, pin, pin-dirty flag). The host resolves the tag with the exact in-process rule against its own pin/step counter, applies the worker's pin only when the loop assigned it since the last event (so fire-time pin relaxation on the host is not clobbered), registers the hook from the resolved tag, and ships the tag back on the response's meta piggyback. The worker keys backward provenance (delivered-clone tags, grad-through-swap seams) off that returned tag, keeping it aligned with the host's retention maps. Out-of-loop requesters keep local resolution: the pin lives with the worker's code and is authoritative there, and the worker's frozen step counter is never consulted. Mediator gains iter_depth (incremented by IteratorTracer around each loop) so the worker's interleaver stub knows which mode applies; in-process nothing reads it. Verified on gpt2/cuda: isolated suite 17/17 (multi-token iteration, backward, grad-through-swap, batched backward all bit-identical, max|delta|=0); in-process tests/test_lm.py + tests/test_iter_edge_cases.py 95/95. Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_01LxKMBwcSwnpwShDYqjkxMz --- src/nnsight/intervention/interleaver.py | 39 +++++++- src/nnsight/intervention/isolation.py | 99 ++++++++++++++++++-- src/nnsight/intervention/tracing/iterator.py | 8 ++ 3 files changed, 131 insertions(+), 15 deletions(-) diff --git a/src/nnsight/intervention/interleaver.py b/src/nnsight/intervention/interleaver.py index a2a2402eb..2356a446b 100755 --- a/src/nnsight/intervention/interleaver.py +++ b/src/nnsight/intervention/interleaver.py @@ -1016,6 +1016,11 @@ def __init__( # cross-process-aware logic (e.g. Barrier) know it can't count locally. self._isolated_worker = False + # Host side of an isolated mediator: the final tagged requester of the + # event currently being handled, shipped back to the worker on the + # response's meta piggyback (the worker keys backward provenance off it). + self._iso_last_tag = None + self.skip_container = None self.history = set() @@ -1023,6 +1028,12 @@ def __init__( self.hooks: List[Any] = list() self.iteration_tracker = defaultdict(int) self.iteration = 0 + # Number of `tracer.iter[...]` scopes currently open (IteratorTracer + # increments/decrements around each loop). Read by the isolated worker's + # interleaver stub to decide whether a requester's iteration must be + # resolved on the host (in-loop: the host owns the step counter and the + # pin-relaxation state) or locally (out-of-loop: the pin is authoritative). + self.iter_depth = 0 self.all_stop: Optional[int] = stop self.args = list() self.cross_invoker = None @@ -1190,6 +1201,7 @@ def cancel(self): self.history = set() self.iteration_tracker = defaultdict(int) self.iteration = 0 + self.iter_depth = 0 self.worker = None # Retained on-graph activations (isolated backward) pin the autograd graph; # drop them at trace end rather than waiting for the mediator to be GC'd. @@ -1267,17 +1279,32 @@ def handle(self, provider: Optional[str] = None, value: Optional[Any] = None): event, data = self.channel.get_event() - # Host-side hook registration: in the isolated path the worker has no real module, so the - # host registers the one-shot hook on demand from the requester string - # the worker just sent. + # Isolated path: the worker has no real module, so the host both + # resolves the requester's iteration tag and registers the one-shot + # hook on demand. An in-loop requester crosses UNTAGGED (a marker + # tuple carrying the base path + the worker's live `tracer.iter` pin): + # the host is the single iteration authority, so the tag is computed + # here from the host mediator's own pin/step counter (the same state + # the in-process resolution reads). Out-of-loop requesters arrive + # already tagged and pass through. The resolved string is stashed on + # `_iso_last_tag` so the response's meta piggyback ships it back to + # the worker (backward provenance keys off it). Events restored by + # `handle_value_event` re-enter this loop already resolved, so the + # tag is never recomputed at a later (wrong) step. if self._iso is not None and event in ( Events.VALUE, Events.SWAP, Events.SKIP, ): - from .isolation import ensure_isolated_provider + from .isolation import ensure_isolated_provider, resolve_iso_requester - requester = data if event == Events.VALUE else data[0] + if event == Events.VALUE: + data = resolve_iso_requester(self, data) + requester = data + else: + requester = resolve_iso_requester(self, data[0]) + data = (requester, *data[1:]) + self._iso_last_tag = requester ensure_isolated_provider(self, requester) if event == Events.VALUE: @@ -1981,6 +2008,8 @@ def __setstate__(self, state): self.user_cache: "Cache" = list() self.hooks: List[Any] = list() self.iteration = 0 + self.iter_depth = 0 + self._iso_last_tag = None self.args = list() self.original_globals = {} self.cross_invoker = None diff --git a/src/nnsight/intervention/isolation.py b/src/nnsight/intervention/isolation.py index 352dd4aa4..e6036060f 100644 --- a/src/nnsight/intervention/isolation.py +++ b/src/nnsight/intervention/isolation.py @@ -538,9 +538,14 @@ def _wire_host_channel(mediator, iso: _PooledWorker, worker_opts: dict) -> None: # ._execute), so a snapshot is stale. Piggyback the LIVE value + the cross_invoker # var store on each response; the worker reads default_all before bounding its # iter[:] loop and pulls the store before each access. + mediator._iso_last_tag = None chan.meta_provider = lambda: { "default_all": mediator.interleaver.default_all, "xinvoke_store": mediator.interleaver._xinvoke_store, + # Host-resolved tag of the event this response answers (in-loop + # requesters cross untagged; the worker keys backward provenance + # off the tag the host actually used). + "resolved": mediator._iso_last_tag, } chan.on_push = mediator.interleaver._xinvoke_store.update mediator.channel = chan @@ -623,6 +628,36 @@ def path_to_envoy(mediator) -> dict: return iso.path2envoy +# First element of the wire tuple an isolated worker sends for an in-loop +# (`tracer.iter`) requester: ("", base_path, pin, pin_dirty). The host +# resolves the iteration tag; out-of-loop requesters cross as plain tagged strings. +ISO_ITER_REQ = "__nnsight_iso_iter_req__" + + +def resolve_iso_requester(mediator, wire) -> str: + """Resolve a worker-sent requester into its final tagged string, host-side. + + A plain string (out-of-loop, already tagged by the worker) passes through. + A marker tuple (in-loop) is resolved against the HOST mediator's state with + the exact in-process rule (`Interleaver.iterate_requester`): the pin if set, + else the per-path step counter. The worker's pin rides the tuple and is + applied only when `pin_dirty` says the worker's iter loop assigned it since + the last event; otherwise the host's own pin state (including fire-time + relaxation done by the one-shot hooks) is authoritative. + """ + if isinstance(wire, str): + return wire + _, base, pin, pin_dirty = wire + if pin_dirty: + mediator.iteration = pin + iteration = ( + mediator.iteration + if mediator.iteration is not None + else mediator.iteration_tracker[base] + ) + return f"{base}.i{iteration}" + + def ensure_isolated_provider(mediator, requester: str) -> None: """Host-side hook registration: register the one-shot hook for ``requester`` on the *real* module. @@ -684,8 +719,18 @@ def __init__(self, default_all=None): # to run; set by generate(max_new_tokens=N) on the host and shipped over. self.default_all = default_all - def iterate_requester(self, requester: str) -> str: + def iterate_requester(self, requester: str): med = self.current + # Inside a `tracer.iter[...]` scope the step counter lives on the HOST + # (the worker's tracker never advances: its modules are dummies), so the + # requester crosses untagged and the host resolves it. The pin rides + # along, flagged dirty only if the loop assigned it since the last event, + # so host-side pin relaxation is not clobbered by a stale worker pin. + if med.iter_depth: + dirty, med._pin_dirty = med._pin_dirty, False + return (ISO_ITER_REQ, requester, med.iteration, dirty) + # Out of loop the pin is authoritative and lives here; resolve locally + # (same rule as in-process). iteration = ( med.iteration if med.iteration is not None else med.iteration_tracker[requester] ) @@ -753,10 +798,31 @@ class WorkerMediator(Mediator): channel's ``on_meta`` / ``push_provider`` per job. """ + # The `tracer.iter` pin, instrumented so `iterate_requester` can tell the host + # whether the worker's loop assigned it since the last event (dirty). The class + # swap in `adopt` makes this data descriptor shadow the deserialized instance + # attribute, which `adopt` migrates into `_iteration`. + @property + def iteration(self): + return self._iteration + + @iteration.setter + def iteration(self, value): + self._iteration = value + self._pin_dirty = True + @classmethod def adopt(cls, mediator, channel, interleaver, opts: dict, device) -> "WorkerMediator": """Turn a freshly-deserialized mediator into this job's worker mediator.""" mediator.__class__ = cls + # Migrate the pin under the property; not dirty: the host mediator starts + # from the identical deserialized state, so there is nothing to push. + mediator._iteration = mediator.__dict__.pop("iteration") + mediator._pin_dirty = False + # Final tag of the last-delivered value, resolved on the host and shipped + # back on the response meta (None for out-of-loop requesters, which the + # worker tagged itself). + mediator._last_resolved = None mediator.channel = channel mediator.interleaver = interleaver mediator.idx = 0 @@ -775,22 +841,33 @@ def adopt(cls, mediator, channel, interleaver, opts: dict, device) -> "WorkerMed interleaver.current = mediator return mediator - def request(self, requester: str): + def _resolved_tag(self, requester) -> str: + """The final tagged requester of the exchange that just completed: an + out-of-loop requester was tagged locally; an in-loop one crossed as a + marker tuple and its tag came back on the response meta.""" + return requester if isinstance(requester, str) else self._last_resolved + + def request(self, requester): value = super().request(requester) if self._bwd_active: - self._tag_delivered(value, requester) + self._tag_delivered(value, self._resolved_tag(requester)) return value - def swap(self, requester: str, value): + def swap(self, requester, value): # Grad-through-swap: keep the worker-TAPE swap value (with grad_fn back to the # delivered clones) before send() ships its detached copy. The backward block uses # it to backprop the host-returned dL/d(swap leaf) into dL/d(delivered clone). Tuple - # outputs: the residual is element [0]. + # outputs: the residual is element [0]. Keyed by the HOST-resolved tag (known + # only after the exchange), matching the host's `_iso_grad_swaps` keys. + seam = None if self._bwd_active: seam = value[0] if isinstance(value, tuple) else value - if torch.is_tensor(seam) and seam.requires_grad: - self._bwd_swaps[requester] = seam - return super().swap(requester, value) + if not (torch.is_tensor(seam) and seam.requires_grad): + seam = None + result = super().swap(requester, value) + if seam is not None: + self._bwd_swaps[self._resolved_tag(requester)] = seam + return result def _tag_delivered(self, value, requester: str) -> None: """Tag each delivered activation tensor with its requester provenance and make @@ -838,8 +915,10 @@ def exception(self, exception: Exception): super().exception(_transmissible_exc(exception)) def apply_meta(self, m: dict) -> None: - # Live host state piggybacked on each response: the iter[:] bound and the - # cross_invoker var store (pulled into the frame so push()/pull() see it). + # Live host state piggybacked on each response: the iter[:] bound, the + # cross_invoker var store (pulled into the frame so push()/pull() see it), + # and the host-resolved tag of the event this response answers. + self._last_resolved = m["resolved"] self.interleaver.default_all = m.get( "default_all", self.interleaver.default_all ) diff --git a/src/nnsight/intervention/tracing/iterator.py b/src/nnsight/intervention/tracing/iterator.py index 6361ca435..3f1cf85fd 100755 --- a/src/nnsight/intervention/tracing/iterator.py +++ b/src/nnsight/intervention/tracing/iterator.py @@ -230,6 +230,11 @@ def __iter__(self): # in the finally block below so they don't leak outside the loop. iter_handles = register_iter_hooks(mediator, self.model) + # Mark the iter scope open. In an isolated worker this routes requester + # tagging to the host (the worker's tracker never advances on its dummy + # modules); in-process nothing reads it. + mediator.iter_depth += 1 + try: if isinstance(self.iteration, slice): @@ -284,6 +289,7 @@ def __iter__(self): yield self.iteration finally: + mediator.iter_depth -= 1 mediator.iteration = original_iteration # Remove the iteration-tracking hooks. @@ -331,6 +337,7 @@ def execute(self, fn: Callable): mediator.push() iter_handles = register_iter_hooks(mediator, self.model) + mediator.iter_depth += 1 def do_iteration(iter: int): @@ -381,6 +388,7 @@ def do_iteration(iter: int): do_iteration(self.iteration) finally: + mediator.iter_depth -= 1 mediator.iteration = original_iteration for handle in iter_handles: