#!/usr/bin/env python3 """Compare raw mamba2 state dumps between two determinism runs. Reads the dumps produced by `PerceptionTrainer::dump_mamba2_state_for_debug` (env-gated on `FOXHUNT_DETERMINISM_DEBUG_MAMBA2=1`) and prints a per-step exact-bytewise-equality verdict for each of: Per mamba2 layer (l1, l2): - w_in, b_in, w_a, b_a, w_b, b_b, w_c, w_out, b_out (weights) - h_enriched_seq (forward output) Encoder activations (single tensors, shared across layers): - vsn_out, ln_a_out, ln_b_out, attn_context, h_t Hypothesis table (spec §2.A sub-investigation, Phase 2.2): CASE 2.A.1 (reduce non-det): mamba2_grad_w_c (after-reduce) diverges, intermediates before reduce match → `mamba2_alpha_reduce_d_w_c` or `_d_proj` kernel is non-det. (only catchable via post-reduce checksum — dump is for weights/outputs) CASE 2.A.2 (forward non-det): weights match at step N, but `mamba2_l1_h_enriched_seq.bin` (forward output) diverges with identical weights → forward scan kernel is non-deterministic (very unlikely given per-(i, j) thread design). CASE 2.A.3 (upstream non-det): mamba2 L1 forward output diverges AT THE SAME STEP AS vsn_out → the non-determinism enters BEFORE mamba2 (VSN forward / backward, or earlier). CASE 2.A.4 (weights diverge): weights diverge at step N before any activation diverges → upstream optimizer / gradient accumulation (Adam iteration order, grad-reduction non-det) wrote different bits to the weights. Output is written to stdout. Exit codes: 0 both runs identical for all dumped buffers at all dumped steps 1 divergence found; verdict printed (one of 2.A.2 / 2.A.3 / 2.A.4) 2 inputs missing or malformed Usage: python3 scripts/compare-mamba2-state.py Each dir contains `step_{0,1,2,3}_.bin`. Per `feedback_cpu_is_read_only`: this script ONLY reads device dumps; it performs no compute that affects training. """ from __future__ import annotations import argparse import struct import sys from pathlib import Path # Buffers dumped at each step (defined once, used for all steps). # Order MATCHES the dump order in `dump_mamba2_state_for_debug` so the # output diff is reproducible. WEIGHTS_L1 = [ "mamba2_l1_w_in", "mamba2_l1_b_in", "mamba2_l1_w_a", "mamba2_l1_b_a", "mamba2_l1_w_b", "mamba2_l1_b_b", "mamba2_l1_w_c", "mamba2_l1_w_out", "mamba2_l1_b_out", ] WEIGHTS_L2 = [ "mamba2_l2_w_in", "mamba2_l2_b_in", "mamba2_l2_w_a", "mamba2_l2_b_a", "mamba2_l2_w_b", "mamba2_l2_b_b", "mamba2_l2_w_c", "mamba2_l2_w_out", "mamba2_l2_b_out", ] ACTIVATIONS = [ "vsn_out", "mamba2_l1_h_enriched_seq", "ln_a_out", "mamba2_l2_h_enriched_seq", "ln_b_out", "attn_context", "h_t", ] # Phase 2.4 (2026-06-02): backward scratch + reduced output, to localise # the `mamba2_grad_w_c_l1` residual that survives PEDANTIC cuBLAS. BACKWARD_SCRATCH = [ "mamba2_l1_d_w_c_per_sample", # pre-reduce, [N, sh2, state_d] "mamba2_l1_dw_c_reduced", # post-reduce, [sh2, state_d] "mamba2_l1_d_h_s2", # [N, sh2] — scan_bwd_seq output "mamba2_l2_d_w_c_per_sample", "mamba2_l2_dw_c_reduced", # Phase 2.4 follow-up: dw_in / dw_a / dw_b (cuBLAS GEMM outputs) # and d_a_proj / d_b_proj (reduce_d_proj outputs). "mamba2_l1_d_a_per_channel", # scan_bwd_seq output "mamba2_l1_d_b_per_channel", # scan_bwd_seq output "mamba2_l1_d_a_proj_2d", # reduce_d_proj output "mamba2_l1_d_b_proj_2d", # reduce_d_proj output "mamba2_l1_dw_in", # cuBLAS GEMM output "mamba2_l1_dw_a", # cuBLAS GEMM output "mamba2_l1_dw_b", # cuBLAS GEMM output "mamba2_l1_db_in", "mamba2_l1_db_a", "mamba2_l1_db_b", "mamba2_l1_d_x", "mamba2_l1_d_x_from_in", ] # Phase 2.4 follow-up #2: trunk weights + their grads to localise which # upstream-of-mamba2 weight first diverges (a divergent vsn_w_d at step 1 # would propagate to differing forward inputs to mamba2 at step 2). TRUNK_STATE = [ "vsn_w", "attn_q", "cfc_w_in", "cfc_w_rec", "ln_a_gain", "ln_b_gain", "grad_vsn_w", "grad_attn_q", "grad_cfc_w_in", "grad_cfc_w_rec", "grad_ln_a_gain", "grad_ln_b_gain", ] ALL_BUFFERS = WEIGHTS_L1 + WEIGHTS_L2 + ACTIVATIONS + BACKWARD_SCRATCH + TRUNK_STATE def read_f32(path: Path) -> list[float]: """Read a raw little-endian f32 binary file into a Python list.""" data = path.read_bytes() if len(data) % 4 != 0: raise ValueError(f"{path}: size {len(data)} not a multiple of 4") n = len(data) // 4 return list(struct.unpack(f"<{n}f", data)) def compare_lists(a: list, b: list, label: str) -> tuple[bool, int, str]: """Compare two equal-length lists; return (equal, first_diff_idx, summary).""" if len(a) != len(b): return False, -1, f"{label}: length mismatch A={len(a)} B={len(b)}" diffs: list[int] = [] for i, (x, y) in enumerate(zip(a, b)): if x != y: diffs.append(i) if len(diffs) >= 5: break if not diffs: return True, -1, f"{label}: EQUAL ({len(a)} elements)" sample = ", ".join( f"[{i}]: A={a[i]!r} B={b[i]!r}" for i in diffs[:3] ) return False, diffs[0], ( f"{label}: DIVERGE at idx {diffs[0]} " f"(showing up to 3): {sample}" ) def main() -> int: parser = argparse.ArgumentParser(description=__doc__) parser.add_argument("run_a", type=Path) parser.add_argument("run_b", type=Path) args = parser.parse_args() if not args.run_a.is_dir() or not args.run_b.is_dir(): print( f"ERROR: run dirs missing: {args.run_a} or {args.run_b}", file=sys.stderr, ) return 2 # Per-step verdict. Tracks the FIRST step that diverges and which # buffers diverge there. overall_any_diverge = False first_diverge_step: int | None = None first_diverge_state: dict[str, bool] = {} for step in (0, 1, 2, 3): print(f"=== step {step} ===") step_state: dict[str, bool] = {} any_present = False for name in ALL_BUFFERS: pa = args.run_a / f"step_{step}_{name}.bin" pb = args.run_b / f"step_{step}_{name}.bin" if not (pa.is_file() and pb.is_file()): # Buffer missing — silently skip; partial dumps okay. continue any_present = True try: a = read_f32(pa) b = read_f32(pb) except ValueError as e: print(f" ERROR reading {name}: {e}") return 2 equal, _, summary = compare_lists(a, b, name) # Compact one-line summary so the per-step block is readable # at a glance. print(f" {summary}") step_state[name] = equal if not equal: overall_any_diverge = True if not any_present: print(" SKIP — no buffers dumped for this step") continue if (not all(step_state.values())) and first_diverge_step is None: first_diverge_step = step first_diverge_state = dict(step_state) print() if not overall_any_diverge: print("VERDICT: all dumped buffers match across all dumped steps.") print("Mamba2 state itself is bit-identical between runs. The") print("downstream `checksums.encoder_output` divergence observed") print("by determinism-check.sh must come from later kernels (CfC") print("/ heads / attention) AFTER mamba2 produces its outputs.") return 0 assert first_diverge_step is not None # Categorise: which buffers diverge first? diverged = [name for name, eq in first_diverge_state.items() if not eq] l1_weights_diverged = any( name in WEIGHTS_L1 for name in diverged ) l2_weights_diverged = any( name in WEIGHTS_L2 for name in diverged ) vsn_out_diverged = "vsn_out" in diverged l1_out_diverged = "mamba2_l1_h_enriched_seq" in diverged l2_out_diverged = "mamba2_l2_h_enriched_seq" in diverged ln_a_diverged = "ln_a_out" in diverged ln_b_diverged = "ln_b_out" in diverged attn_diverged = "attn_context" in diverged h_t_diverged = "h_t" in diverged print(f"VERDICT (first divergent step = {first_diverge_step}):") print(f" diverged buffers: {diverged}") print() if l1_weights_diverged or l2_weights_diverged: print(" CASE 2.A.4 — WEIGHTS DIVERGE upstream of forward.") which = [] if l1_weights_diverged: which.append("L1") if l2_weights_diverged: which.append("L2") print(f" Mamba2 {'/'.join(which)} weights diverged at step " f"{first_diverge_step}. This means the PRIOR step's") print(" backward+Adam update wrote different bits — the") print(" non-determinism source is in:") print(" (a) mamba2 backward kernels (d_a/d_b/d_w_c per-channel") print(" writes are deterministic by construction, but") print(" their REDUCE kernels mamba2_alpha_reduce_d_proj /") print(" _d_w_c could be non-det if launch geometry races)") print(" (b) Mamba2AdamW.step_from_buffers_gpu_clip — the") print(" grad-norm phase1+phase2 tree-reduce. Phase 2") print(" appears single-block per `grad_norm.cu`, but") print(" re-audit if reduce kernels look fine.") print(" (c) Encoder upstream backward (CfC / LN / VSN /") print(" attention) writing non-det grads that feed") print(" mamba2 W_in via dx propagation through L1 / L2.") print(" Fix: identify which weight diverged FIRST and audit") print(" its gradient pipeline. Next dispatch should add") print(" checksums for the encoder backward grad outputs") print(" (grad_vsn_w_d, grad_cfc_*, grad_attn_q_d, ...) and the") print(" mamba2 reduced grad outputs (mamba2_grads_buffers.dw_*).") return 1 if vsn_out_diverged and not l1_out_diverged: print(" CASE 2.A.3 — UPSTREAM (VSN) non-determinism.") print(" vsn_out diverges before mamba2 L1 output does. The") print(" issue is in VSN forward (vsn_fwd kernel) or earlier") print(" (snap_feature_assemble_batched). Apply §2.D audit") print(" to the VSN forward kernel.") return 1 if l1_out_diverged and not l1_weights_diverged: # Forward output diverges with identical weights. print(" CASE 2.A.2 — MAMBA2 L1 FORWARD non-determinism.") print(f" At step {first_diverge_step}, mamba2 L1 weights match") print(" bit-exactly between runs, but the L1 forward output") print(" h_enriched_seq diverges. The forward scan kernel") print(" `mamba2_alpha_scan_fwd_seq` is per-thread sequential") print(" (one thread per (i, j), each thread does its own") print(" K-step scan with no atomics) — divergence here would") print(" be surprising and would require auditing the per-thread") print(" register-array initialisation + the cuBLAS GEMMs that") print(" feed a_proj / b_proj inputs (w_in / w_a / w_b forward).") print(" Likely culprit: cuBLAS split-K non-determinism in the") print(" W_in / W_a / W_b forward GEMMs (TF32 / heuristic).") return 1 if ln_a_diverged or ln_b_diverged or attn_diverged or l2_out_diverged or h_t_diverged: # Divergence enters downstream of mamba2 L1 but with L1 outputs/weights equal. print(" CASE 2.A.5 — DOWNSTREAM (post-mamba2) non-determinism.") print(" Mamba2 L1 outputs match but later layer outputs diverge.") for n in ("ln_a_out", "mamba2_l2_h_enriched_seq", "ln_b_out", "attn_context", "h_t"): if n in first_diverge_state: eq = first_diverge_state[n] print(f" {n}: {'EQUAL' if eq else 'DIVERGE'}") print(" Audit the FIRST diverged downstream buffer's kernel.") return 1 print(" UNEXPECTED state — divergence pattern does not match a") print(" defined sub-case. Review per-buffer summary above and") print(" extend the verdict logic.") return 1 if __name__ == "__main__": sys.exit(main())