BP-free formal audit (5.6-sol): EP update verified local (108 bitwise zero-influence checks); fixes — gate_every<=0 = true off switch, governor reaction now opt-in (--gate_govern, default observe-only), dFdtheta restored as self-sealing invariant-test surface; test_bp_free.py 4/4 green

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_014FAPDWQ49M5Ye3NpTndTpn

2 files changed, 588 insertions, 8 deletions

diff --git a/ep_run/casc_eq_train.py b/ep_run/casc_eq_train.py
index 4eba267..f1f26a4 100644
--- a/ep_run/casc_eq_train.py
+++ b/ep_run/casc_eq_train.py
@@ -27,7 +27,8 @@ ap.add_argument('--tok_init', type=float, default=0.0)  # >0: init tok/pos with
 ap.add_argument('--compile', action='store_true')       # torch.compile each block (free speed where supported)
 ap.add_argument('--sig_every', type=int, default=25)    # tok-sigma refresh interval (amortized)
 ap.add_argument('--dtop_every', type=int, default=1)    # 1 = exact (DEFAULT, BP-parity); 2 = fast mode (~20% cheaper, ~4% CE tax at high lr)
-ap.add_argument('--gate_every', type=int, default=200)  # in-training cos(EP,BP) telemetry
+ap.add_argument('--gate_every', type=int, default=200)  # in-training cos(EP,BP) telemetry; <=0 = fully BP-free (no bp_gate at all)
+ap.add_argument('--gate_govern', action='store_true')   # let gate cos adjust K/bscale (default: observe-only => training control is BP-free)
 args = ap.parse_args()
 torch.manual_seed(args.seed)
 dev = 'cuda' if torch.cuda.is_available() else 'cpu'
@@ -129,14 +130,20 @@ def relax(z0, zs, ins, outs, y, beta, K, x):
     return zs, outs
 
 def dFdtheta(zs, x, y, beta):
-    """dF/dtheta at fixed relaxed states (z0 rebuilt WITH graph so emb gets its E-path grad)."""
+    """theta-readout at FIXED states. Not used by the training loop (relax reuses its own
+    graphs); kept as the INVARIANT-TEST surface for test_bp_free.py. Self-sealing: inputs
+    are detached here so the local-graph property holds for any caller."""
+    zs = [z.detach() for z in zs]
     prev = tok(x) + pos(torch.arange(args.T, device=dev))[None]
     E = 0.0
-    for z, b in zip(zs, blocks): E = E + 0.5 * ((z - b(prev, mask)) ** 2).sum(); prev = z
+    for z, b in zip(zs, blocks):
+        E = E + 0.5 * ((z - b(prev, mask)) ** 2).sum()
+        prev = z   # zs detached at entry => blocks l>0 get detached inputs; block 0 gets the graphed emb
     obj = E / NBT + beta * F.cross_entropy(readout(zs[-1]).reshape(-1, vocab), y.reshape(-1))
     gs = torch.autograd.grad(obj, all_params, allow_unused=True)
     return [g if g is not None else None for g in gs]
 
+
 SIG0 = None
 GOV = {'K': None, 'bscale': 1.0, 'gema': None, 'drift': 0.0, 'gn': 0.0, 'sig': 0.0}
 def ep_step(x, y):
@@ -215,17 +222,18 @@ for step in range(args.steps + 1):
     ce, beta_t, rounds, ok = ep_step(x, y)
     if not ok: skips += 1
     gcos = float('nan')
-    if step % args.gate_every == 0 and ok:
+    if args.gate_every > 0 and step % args.gate_every == 0 and ok:
         gbp = bp_gate(x, y)
         num = den1 = den2 = 0.0
         for p, g in zip(all_params, gbp):
             if p.grad is None or g is None: continue
             num += float((p.grad * g).sum()); den1 += float((p.grad ** 2).sum()); den2 += float((g ** 2).sum())
         gcos = num / max((den1 ** 0.5) * (den2 ** 0.5), 1e-12)
-        if gcos < 0.97:                                   # estimator governor: spend more
-            GOV['K'] = min(GOV['K'] + 2, args.kmax); GOV['bscale'] = max(GOV['bscale'] * 0.7, 0.05)
-        elif gcos > 0.995 and GOV['K'] > args.K:          # relax back when quality is abundant
-            GOV['K'] -= 1; GOV['bscale'] = min(GOV['bscale'] * 1.05, 1.0)
+        if args.gate_govern:                              # opt-in: BP-informed control flow
+            if gcos < 0.97:
+                GOV['K'] = min(GOV['K'] + 2, args.kmax); GOV['bscale'] = max(GOV['bscale'] * 0.7, 0.05)
+            elif gcos > 0.995 and GOV['K'] > args.K:
+                GOV['K'] -= 1; GOV['bscale'] = min(GOV['bscale'] * 1.05, 1.0)
     torch.nn.utils.clip_grad_norm_(all_params, 1.0)
     opt.step(); sched.step(); opt.zero_grad(set_to_none=True)
     if step % args.log == 0:
diff --git a/ep_run/test_bp_free.py b/ep_run/test_bp_free.py
new file mode 100644
index 0000000..886a7a5
--- /dev/null
+++ b/ep_run/test_bp_free.py
@@ -0,0 +1,572 @@
+import ast
+import builtins
+from contextlib import ExitStack
+import io
+import inspect
+import pickle
+import sys
+import types
+from pathlib import Path
+from types import SimpleNamespace
+from unittest import mock
+
+import numpy as np
+import torch
+import torch.nn.functional as F
+
+
+ROOT = Path(__file__).resolve().parent
+TRAINER = ROOT / "casc_eq_train.py"
+PROBE = ROOT / "cascade_probe.py"
+_REAL_OPEN = builtins.open
+_META = pickle.dumps({"vocab_size": 17})
+
+
+def _fake_open(file, *args, **kwargs):
+    normalized = str(file).replace("\\", "/")
+    if normalized.endswith("/data/tinystories_bpe/meta.pkl"):
+        return io.BytesIO(_META)
+    return _REAL_OPEN(file, *args, **kwargs)
+
+
+def _load_prefix(path, stop_marker, argv, fake_memmap=False):
+    source = path.read_text()
+    assert stop_marker in source, f"loader marker moved in {path.name}"
+    source = source.split(stop_marker, 1)[0]
+
+    module = types.ModuleType(f"_bp_free_audit_{path.stem}")
+    module.__file__ = str(path)
+
+    with ExitStack() as stack:
+        stack.enter_context(mock.patch.object(sys, "argv", [str(path), *argv]))
+        stack.enter_context(mock.patch.object(builtins, "open", _fake_open))
+        stack.enter_context(
+            mock.patch.object(torch.cuda, "is_available", return_value=False)
+        )
+        if fake_memmap:
+            dummy = np.arange(256, dtype=np.uint16)
+            stack.enter_context(
+                mock.patch.object(np, "memmap", return_value=dummy)
+            )
+        exec(compile(source, str(path), "exec"), module.__dict__)
+
+    assert module.dev == "cpu"
+    return module
+
+
+def _load_trainer():
+    return _load_prefix(
+        TRAINER,
+        "\nwb = None",
+        [
+            "--L", "4",
+            "--C", "8",
+            "--H", "2",
+            "--T", "4",
+            "--B", "2",
+            "--K", "2",
+            "--steps", "-1",
+            "--noguard",
+            "--gate_every", "0",
+        ],
+    )
+
+
+def _load_probe():
+    return _load_prefix(
+        PROBE,
+        "\nflat = lambda",
+        [
+            "--L", "4",
+            "--C", "8",
+            "--H", "2",
+            "--T", "4",
+            "--B", "2",
+            "--K", "2",
+            "--scheme", "fb",
+            "--include_io",
+            "--batches", "1",
+        ],
+        fake_memmap=True,
+    )
+
+
+class _BlockMarker(torch.autograd.Function):
+    @staticmethod
+    def forward(ctx, x, block_idx):
+        ctx.block_idx = int(block_idx)
+        return x
+
+    @staticmethod
+    def backward(ctx, grad):
+        return grad, None
+
+
+def _tag_blocks(module):
+    for idx, block in enumerate(module.blocks):
+        original = block.forward
+
+        def tagged(self, z, mask, _original=original, _idx=idx):
+            output = _original(z, mask)
+            return _BlockMarker.apply(output, _idx)
+
+        block.forward = types.MethodType(tagged, block)
+
+
+def _tensors(value):
+    if isinstance(value, torch.Tensor):
+        yield value
+    elif isinstance(value, (tuple, list)):
+        for item in value:
+            yield from _tensors(item)
+
+
+def _roots(outputs):
+    return [
+        tensor.grad_fn
+        for tensor in _tensors(outputs)
+        if tensor.grad_fn is not None
+    ]
+
+
+def _graph_facts(outputs, param_owner):
+    """
+    Walk every root-to-leaf dependency path.
+
+    A summed theta objective may contain markers from all blocks, but no
+    individual path may pass through more than one marker.
+    """
+    stack = [(root, ()) for root in _roots(outputs)]
+    seen = set()
+    markers = set()
+    max_chain = 0
+    foreign_param_paths = []
+    unmarked_block_params = []
+
+    while stack:
+        node, path = stack.pop()
+
+        block_idx = getattr(node, "block_idx", None)
+        if block_idx is not None:
+            path = path + ((id(node), int(block_idx)),)
+            markers.add(int(block_idx))
+            max_chain = max(max_chain, len(path))
+
+        key = (node, tuple(marker_id for marker_id, _ in path))
+        if key in seen:
+            continue
+        seen.add(key)
+
+        variable = getattr(node, "variable", None)
+        owner = param_owner.get(id(variable))
+        if owner is not None:
+            if not path:
+                unmarked_block_params.append(owner)
+            elif path[-1][1] != owner:
+                foreign_param_paths.append(
+                    (tuple(idx for _, idx in path), owner)
+                )
+
+        for next_node, _ in node.next_functions:
+            if next_node is not None:
+                stack.append((next_node, path))
+
+    return {
+        "markers": markers,
+        "max_chain": max_chain,
+        "foreign": foreign_param_paths,
+        "unmarked": unmarked_block_params,
+    }
+
+
+def _assert_local(outputs, param_owner, label):
+    facts = _graph_facts(outputs, param_owner)
+
+    assert facts["max_chain"] <= 1, (
+        f"{label}: GLOBAL transformer chain detected: {facts}"
+    )
+    assert not facts["foreign"], (
+        f"{label}: block marker reached another block's parameters: {facts}"
+    )
+    assert not facts["unmarked"], (
+        f"{label}: an unmarked path reached block parameters: {facts}"
+    )
+    return facts
+
+
+def _checked_grad(real_grad, param_owner, records):
+    def checked(outputs, inputs, *args, **kwargs):
+        frame = inspect.currentframe().f_back
+        label = (
+            f"{Path(frame.f_code.co_filename).name}:"
+            f"{frame.f_code.co_name}:{frame.f_lineno}"
+        )
+
+        facts = _assert_local(outputs, param_owner, label)
+
+        grad_outputs = kwargs.get(
+            "grad_outputs", args[0] if args else None
+        )
+        assert all(
+            tensor.grad_fn is None and not tensor.requires_grad
+            for tensor in _tensors(grad_outputs)
+        ), f"{label}: VJP seed carries an autograd graph"
+
+        assert all(
+            tensor.is_leaf for tensor in _tensors(inputs)
+        ), f"{label}: differentiation target is not a leaf"
+
+        result = real_grad(outputs, inputs, *args, **kwargs)
+
+        assert all(
+            grad is None
+            or (grad.grad_fn is None and not grad.requires_grad)
+            for grad in result
+        ), f"{label}: autograd.grad returned a higher-order graph"
+
+        records.append((frame.f_code.co_name, frame.f_lineno, facts))
+        return result
+
+    return checked
+
+
+def _dotted_name(node):
+    if isinstance(node, ast.Name):
+        return node.id
+    if isinstance(node, ast.Attribute):
+        prefix = _dotted_name(node.value)
+        return f"{prefix}.{node.attr}" if prefix else node.attr
+    return ""
+
+
+def _grad_sites(path):
+    tree = ast.parse(path.read_text(), filename=str(path))
+    sites = []
+
+    class Visitor(ast.NodeVisitor):
+        def __init__(self):
+            self.functions = []
+
+        def visit_FunctionDef(self, node):
+            self.functions.append(node.name)
+            self.generic_visit(node)
+            self.functions.pop()
+
+        visit_AsyncFunctionDef = visit_FunctionDef
+
+        def visit_Call(self, node):
+            if _dotted_name(node.func) == "torch.autograd.grad":
+                owner = self.functions[-1] if self.functions else "<module>"
+                sites.append((owner, node.lineno))
+            self.generic_visit(node)
+
+    Visitor().visit(tree)
+    return set(sites)
+
+
+def test_casc_eq_training_graphs_are_local_and_gate_is_read_only():
+    module = _load_trainer()
+    _tag_blocks(module)
+
+    param_owner = {
+        id(param): idx
+        for idx, block in enumerate(module.blocks)
+        for param in block.parameters()
+    }
+
+    generator = torch.Generator().manual_seed(123)
+    x = torch.randint(
+        module.vocab,
+        (module.args.B, module.args.T),
+        generator=generator,
+    )
+    y = torch.randint(
+        module.vocab,
+        (module.args.B, module.args.T),
+        generator=generator,
+    )
+
+    # Free-state firewalls.
+    z0, states, ins, outs = module.free_states_graphed(x)
+    assert z0.grad_fn is None and not z0.requires_grad
+    assert all(
+        state.grad_fn is None and not state.requires_grad
+        for state in states
+    )
+    assert all(item.is_leaf and item.requires_grad for item in ins)
+
+    for idx, output in enumerate(outs):
+        facts = _assert_local(output, param_owner, f"free output {idx}")
+        assert facts["markers"] == {idx}
+
+    real_grad = torch.autograd.grad
+    records = []
+    checked = _checked_grad(real_grad, param_owner, records)
+
+    with mock.patch.object(torch.autograd, "grad", checked):
+        relaxed, last_outs = module.relax(
+            z0,
+            states,
+            ins,
+            outs,
+            y,
+            module.args.beta,
+            module.args.K,
+            x,
+        )
+
+        assert all(
+            state.grad_fn is None and not state.requires_grad
+            for state in relaxed
+        )
+
+        for idx, output in enumerate(last_outs):
+            facts = _assert_local(
+                output, param_owner, f"last-round output {idx}"
+            )
+            assert facts["markers"] == {idx}
+
+        # Exercise the otherwise-unused helper call site too.
+        module.dFdtheta(relaxed, x, y, module.args.beta)
+
+        _, _, _, ok = module.ep_step(x, y)
+        assert ok
+
+    # Every non-gate autograd.grad call currently in the source must have
+    # been exercised by the instrumented tiny run.
+    sites = _grad_sites(TRAINER)
+    local_sites = {
+        site for site in sites if site[0] != "bp_gate"
+    }
+    observed = {(name, line) for name, line, _ in records}
+    assert local_sites <= observed
+
+    # The theta objective should contain all block markers, but as a
+    # disconnected forest with one marker maximum per path.
+    assert any(
+        facts["markers"] == set(range(module.args.L))
+        and facts["max_chain"] == 1
+        for _, _, facts in records
+    )
+
+    # Known global exception: prove both that it is truly global and that
+    # it cannot overwrite the EP gradients stored in p.grad.
+    for idx, param in enumerate(module.all_params):
+        param.grad = torch.full_like(param, (idx + 1) / 1000)
+
+    before = [param.grad.clone() for param in module.all_params]
+    gate_facts = []
+
+    def gate_grad(outputs, inputs, *args, **kwargs):
+        gate_facts.append(_graph_facts(outputs, param_owner))
+        return real_grad(outputs, inputs, *args, **kwargs)
+
+    with mock.patch.object(torch.autograd, "grad", gate_grad):
+        gbp = module.bp_gate(x, y)
+
+    assert len(gate_facts) == 1
+    assert gate_facts[0]["max_chain"] == module.args.L
+    assert all(
+        torch.equal(old, param.grad)
+        for old, param in zip(before, module.all_params)
+    ), "bp_gate populated or changed p.grad"
+    assert all(
+        grad is None
+        or (grad.grad_fn is None and not grad.requires_grad)
+        for grad in gbp
+    )
+
+    # The trainer should not acquire any Tensor.backward/autograd.backward
+    # calls later.
+    tree = ast.parse(TRAINER.read_text())
+    backward_lines = [
+        node.lineno
+        for node in ast.walk(tree)
+        if isinstance(node, ast.Call)
+        and _dotted_name(node.func).endswith(".backward")
+    ]
+    assert not backward_lines, (
+        f"new .backward() calls at lines {backward_lines}"
+    )
+
+
+def test_theta_readout_has_zero_distant_block_influence():
+    module = _load_trainer()
+    generator = torch.Generator().manual_seed(456)
+
+    x = torch.randint(
+        module.vocab,
+        (module.args.B, module.args.T),
+        generator=generator,
+    )
+    y = torch.randint(
+        module.vocab,
+        (module.args.B, module.args.T),
+        generator=generator,
+    )
+    states = [
+        torch.randn(
+            module.args.B,
+            module.args.T,
+            module.args.C,
+            generator=generator,
+        )
+        for _ in range(module.args.L)
+    ]
+
+    baseline = module.dFdtheta(states, x, y, module.args.beta)
+    param_index = {
+        id(param): idx
+        for idx, param in enumerate(module.all_params)
+    }
+
+    for changed_state in range(module.args.L):
+        altered = [state.clone() for state in states]
+        altered[changed_state].add_(
+            0.125
+            * torch.randn(
+                altered[changed_state].shape,
+                generator=generator,
+            )
+        )
+        got = module.dFdtheta(altered, x, y, module.args.beta)
+
+        for block_idx, block in enumerate(module.blocks):
+            # E_i may depend numerically only on z_i and z_{i-1}.
+            if changed_state in {block_idx - 1, block_idx}:
+                continue
+
+            for param in block.parameters():
+                idx = param_index[id(param)]
+                assert torch.equal(baseline[idx], got[idx]), (
+                    f"block {block_idx} gradient changed after distant "
+                    f"state {changed_state} changed"
+                )
+
+
+def test_cascade_probe_fb_ep_side_is_local():
+    module = _load_probe()
+    _tag_blocks(module)
+
+    param_owner = {
+        id(param): idx
+        for idx, block in enumerate(module.blocks)
+        for param in block.parameters()
+    }
+
+    generator = torch.Generator().manual_seed(789)
+    x = torch.randint(
+        module.vocab,
+        (module.args.B, module.args.T),
+        generator=generator,
+    )
+    y = torch.randint(
+        module.vocab,
+        (module.args.B, module.args.T),
+        generator=generator,
+    )
+    positions = torch.arange(module.args.T)
+    z0 = (module.tok(x) + module.pos(positions)[None]).detach()
+
+    real_grad = torch.autograd.grad
+    records = []
+    checked = _checked_grad(real_grad, param_owner, records)
+
+    with (
+        mock.patch.object(torch.autograd, "grad", checked),
+        mock.patch.object(
+            torch.Tensor,
+            "backward",
+            side_effect=AssertionError("fb called backward()"),
+        ),
+    ):
+        relaxed = module.relax(z0, y, module.args.beta)
+        module.dFdtheta(
+            relaxed,
+            z0,
+            y,
+            module.args.beta,
+            module.gate_params,
+            x,
+        )
+
+    assert all(
+        state.grad_fn is None and not state.requires_grad
+        for state in relaxed
+    )
+
+    # Per round: one top-force call plus L-1 feedback calls.
+    # Then one disconnected theta-readout forest.
+    assert len(records) == module.args.K * module.args.L + 1
+    assert sum(
+        not facts["markers"] for _, _, facts in records
+    ) == module.args.K
+    assert any(
+        facts["markers"] == set(range(module.args.L))
+        and facts["max_chain"] == 1
+        for _, _, facts in records
+    )
+
+    # Positive control corresponding to the probe's intended BP reference.
+    for idx, param in enumerate(module.gate_params):
+        param.grad = torch.full_like(param, (idx + 1) / 1000)
+    before = [param.grad.clone() for param in module.gate_params]
+
+    z = module.tok(x) + module.pos(positions)[None]
+    for block in module.blocks:
+        z = block(z, module.mask)
+    ce = F.cross_entropy(
+        module.readout(z).reshape(-1, module.vocab),
+        y.reshape(-1),
+    )
+
+    facts = _graph_facts(ce, param_owner)
+    assert facts["max_chain"] == module.args.L
+
+    gbp = real_grad(ce, module.gate_params, allow_unused=True)
+    assert all(
+        torch.equal(old, param.grad)
+        for old, param in zip(before, module.gate_params)
+    )
+    assert all(grad is None or not grad.requires_grad for grad in gbp)
+
+
+def test_gate_every_zero_is_a_true_off_switch():
+    """
+    This intentionally fails on the audited source. It passes after changing
+    the loop condition to guard gate_every > 0 before taking the modulo.
+    """
+    tree = ast.parse(TRAINER.read_text(), filename=str(TRAINER))
+    gate_ifs = []
+
+    for node in ast.walk(tree):
+        if isinstance(node, ast.If) and any(
+            isinstance(child, ast.Call)
+            and _dotted_name(child.func) == "bp_gate"
+            for child in ast.walk(node)
+        ):
+            gate_ifs.append(node)
+
+    assert len(gate_ifs) == 1
+    expression = ast.Expression(gate_ifs[0].test)
+    ast.fix_missing_locations(expression)
+
+    try:
+        enabled = eval(
+            compile(expression, str(TRAINER), "eval"),
+            {},
+            {
+                "args": SimpleNamespace(gate_every=0),
+                "step": 0,
+                "ok": True,
+            },
+        )
+    except ZeroDivisionError:
+        raise AssertionError(
+            "--gate_every 0 crashes; guard modulo with "
+            "args.gate_every > 0"
+        ) from None
+
+    assert enabled is False, (
+        "--gate_every 0 must prevent the global bp_gate call"
+    )