foxhunt/scripts/compare-mamba2-state.py

#!/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 <run_a_dir> <run_b_dir>

Each dir contains `step_{0,1,2,3}_<name>.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())