diff options
Diffstat (limited to 'ep_run/test_bp_free.py')
| -rw-r--r-- | ep_run/test_bp_free.py | 572 |
1 files changed, 572 insertions, 0 deletions
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" + ) |