foxhunt/scripts/compare-backward-state.py

#!/usr/bin/env python3
"""Compare raw backward-pipeline state dumps between two determinism runs.

Reads the dumps produced by
`IntegratedTrainer::dump_backward_state_for_debug` (env-gated on
`FOXHUNT_DETERMINISM_DEBUG_BACKWARD=1`) and prints a per-step
exact-bytewise-equality verdict for each of the 4 buffer groups
identified by Phase 2.5's residual:

  Group G — PER replay sample hidden states + scalars:
    - sampled_h_t, sampled_h_tp1
    - sampled_actions, sampled_rewards, sampled_dones, sampled_log_pi_old

  Group H — Head forward outputs (post-fwd, pre-bwd):
    - q_logits, pi_logits, v_pred
    - iqn_q_values, dueling_q_composed

  Group I — Pre-reduce per-batch / per-row param-grad scratches:
    - cfc_grad_w_in_scratch, cfc_grad_w_rec_scratch
    - cfc_grad_b_scratch, cfc_grad_tau_scratch
    - vsn_grad_w_scratch, vsn_grad_b_scratch
    - attn_grad_q_scratch
    - grad_ln_a_gain_per_row, grad_ln_a_bias_per_row
    - grad_ln_b_gain_per_row, grad_ln_b_bias_per_row

  Group J — Per-head grad_h_t outputs + combined accumulator:
    - grad_h_t_q, grad_h_t_pi, grad_h_t_v
    - grad_h_t_frd, grad_h_t_outcome
    - grad_h_t_combined

Decision tree (per Phase 2.6 dispatch):
  (1) Group G divergent → step-2 PER K-loop sample is non-det (aliased
      pearl-twin to Phase 2's PER tree rebuild bug)
  (2) Group H divergent → head forward kernel (likely NoisyLinear noise
      resampling between forward and backward) is non-det
  (3) Group I EQUAL but later (final-grad) DIVERGE → `reduce_axis0` (or
      `layer_norm_reduce_param_grads`) is the bug; canonical
      Phase-2-PER-rebuild fix applies (single block + __syncthreads +
      fixed thread-ID order)
  (4) Group J divergent → `grad_h_accumulate_scaled` or a specific
      head backward kernel writing grad_h_t is non-det

The script reports the FIRST step that diverges and which group(s)
diverged there. If Group I is the only divergent group at the first
divergent step → confirms hypothesis (3). If both Group I and Group J
diverge → upstream head bwd is also wrong (because Group J feeds
encoder backward via grad_h_accumulate; if its inputs differ, encoder
gradients differ). If Group H diverges but I/J don't → head forward is
non-det but somehow downstream gradients normalize (unlikely).

Exit codes:
  0  both runs identical for all dumped buffers at all dumped steps
  1  divergence found; verdict printed naming the candidate group
  2  inputs missing or malformed

Usage:
  python3 scripts/compare-backward-state.py <run_a_dir> <run_b_dir>

Each dir contains `step_{0,1,2,3}_g_<name>.bin`,
`step_{0,1,2,3}_h_<name>.bin`, `step_{0,1,2,3}_i_<name>.bin`,
`step_{0,1,2,3}_j_<name>.bin` (prefix discriminates buffer group).

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


# Each entry: (logical name, on-disk basename, element struct format,
# element size in bytes). The groups partition the candidate space
# so the verdict can name the culprit class precisely.

# Naming convention: files are written as
# `step_<N>_<group_letter>_<basename>.bin` where group_letter ∈ {g,h,i,j}.
GROUP_G_PER_SAMPLE = [
    ("sampled_h_t",          "g_sampled_h_t",           "f", 4),
    ("sampled_h_tp1",        "g_sampled_h_tp1",         "f", 4),
    ("sampled_actions",      "g_sampled_actions",       "i", 4),
    ("sampled_rewards",      "g_sampled_rewards",       "f", 4),
    ("sampled_dones",        "g_sampled_dones",         "f", 4),
    ("sampled_log_pi_old",   "g_sampled_log_pi_old",    "f", 4),
]
GROUP_H_HEAD_OUTPUTS = [
    ("q_logits",             "h_q_logits",              "f", 4),
    ("pi_logits",            "h_pi_logits",             "f", 4),
    ("v_pred",               "h_v_pred",                "f", 4),
    ("iqn_q_values",         "h_iqn_q_values",          "f", 4),
    ("dueling_q_composed",   "h_dueling_q_composed",    "f", 4),
]
GROUP_I_GRAD_SCRATCH = [
    ("cfc_grad_w_in_scratch",  "i_cfc_grad_w_in_scratch",  "f", 4),
    ("cfc_grad_w_rec_scratch", "i_cfc_grad_w_rec_scratch", "f", 4),
    ("cfc_grad_b_scratch",     "i_cfc_grad_b_scratch",     "f", 4),
    ("cfc_grad_tau_scratch",   "i_cfc_grad_tau_scratch",   "f", 4),
    ("vsn_grad_w_scratch",     "i_vsn_grad_w_scratch",     "f", 4),
    ("vsn_grad_b_scratch",     "i_vsn_grad_b_scratch",     "f", 4),
    ("attn_grad_q_scratch",    "i_attn_grad_q_scratch",    "f", 4),
    ("grad_ln_a_gain_per_row", "i_grad_ln_a_gain_per_row", "f", 4),
    ("grad_ln_a_bias_per_row", "i_grad_ln_a_bias_per_row", "f", 4),
    ("grad_ln_b_gain_per_row", "i_grad_ln_b_gain_per_row", "f", 4),
    ("grad_ln_b_bias_per_row", "i_grad_ln_b_bias_per_row", "f", 4),
]
GROUP_J_GRAD_H_T = [
    ("grad_h_t_q",         "j_grad_h_t_q",         "f", 4),
    ("grad_h_t_pi",        "j_grad_h_t_pi",        "f", 4),
    ("grad_h_t_v",         "j_grad_h_t_v",         "f", 4),
    ("grad_h_t_frd",       "j_grad_h_t_frd",       "f", 4),
    ("grad_h_t_outcome",   "j_grad_h_t_outcome",   "f", 4),
    ("grad_h_t_combined",  "j_grad_h_t_combined",  "f", 4),
]
# Phase 2.6 sub-2 (after first run surfaced grad_h_t_outcome
# divergence at step 2 with Δ ~ 0.003): outcome head's inputs
# (labels + w_d/b_d) are now dumped as Group K to localise the
# upstream-input that drifted (kernel is sole-writer so can't
# introduce divergence from equal inputs).
GROUP_K_OUTCOME_INPUTS = [
    ("outcome_labels",    "k_outcome_labels", "i", 4),
    ("outcome_w",         "k_outcome_w",      "f", 4),
    ("outcome_b",         "k_outcome_b",      "f", 4),
]

ALL_GROUPS = [
    ("G (PER replay sample)",     GROUP_G_PER_SAMPLE),
    ("H (head fwd outputs)",      GROUP_H_HEAD_OUTPUTS),
    ("I (pre-reduce grad scratch)", GROUP_I_GRAD_SCRATCH),
    ("J (per-head grad_h_t)",     GROUP_J_GRAD_H_T),
    ("K (outcome head inputs)",   GROUP_K_OUTCOME_INPUTS),
]


def read_typed(path: Path, fmt: str, elem_size: int) -> list:
    """Read a raw little-endian binary file as `fmt`-typed elements."""
    data = path.read_bytes()
    if len(data) % elem_size != 0:
        raise ValueError(
            f"{path}: size {len(data)} not a multiple of {elem_size}"
        )
    n = len(data) // elem_size
    return list(struct.unpack(f"<{n}{fmt}", data))


def compare_lists(a: list, b: list, label: str) -> tuple[bool, int, str]:
    """Bytewise exact equality. Returns (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] = []
    max_abs_delta = 0.0
    for i, (x, y) in enumerate(zip(a, b)):
        if x != y:
            diffs.append(i)
            try:
                ad = abs(float(x) - float(y))
                if ad > max_abs_delta:
                    max_abs_delta = ad
            except (TypeError, ValueError):
                pass
        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"max|Δ|={max_abs_delta:.6g} "
        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

    overall_any_diverge = False
    per_step_group_results: dict[tuple[int, str], list] = {}
    first_diverge_step: int | None = None
    first_diverge_groups: set[str] = set()

    for step in (0, 1, 2, 3):
        print(f"=== step {step} ===")
        any_present = False
        step_had_diverge = False
        step_diverged_groups: set[str] = set()

        for group_label, members in ALL_GROUPS:
            print(f"  -- group {group_label} --")
            group_results: list[tuple[str, bool, str]] = []
            group_any_present = False
            for name, basename, fmt, elem_size in members:
                pa = args.run_a / f"step_{step}_{basename}.bin"
                pb = args.run_b / f"step_{step}_{basename}.bin"
                if not (pa.is_file() and pb.is_file()):
                    continue
                group_any_present = True
                any_present = True
                try:
                    a = read_typed(pa, fmt, elem_size)
                    b = read_typed(pb, fmt, elem_size)
                except ValueError as e:
                    print(f"    ERROR reading {name}: {e}")
                    return 2
                equal, _, summary = compare_lists(a, b, name)
                print(f"    {summary}")
                group_results.append((name, equal, summary))
                if not equal:
                    overall_any_diverge = True
                    step_had_diverge = True
                    step_diverged_groups.add(group_label)

            if not group_any_present:
                print("    SKIP — no buffers dumped for this group/step")
            per_step_group_results[(step, group_label)] = group_results

        if not any_present:
            print("  SKIP — no buffers dumped for this step")
            continue

        if step_had_diverge and first_diverge_step is None:
            first_diverge_step = step
            first_diverge_groups = step_diverged_groups

    print()
    if not overall_any_diverge:
        print(
            "VERDICT: all dumped backward-pipeline buffers match across\n"
            "all dumped steps. The bug is OUTSIDE Groups G/H/I/J. Per\n"
            "the Phase 2.6 dispatch STOP rule: this is an UNEXPECTED\n"
            "finding (the 4 groups were chosen to enumerate the entire\n"
            "step-2 backward pipeline). Possibilities:\n"
            "  - Adam optimizer step is non-deterministic (parallel\n"
            "    reduction in moment update)\n"
            "  - Bellman target projection has order-dependent writes\n"
            "  - A reducer kernel writes scratch BEFORE Group I dump\n"
            "    captures it (race between dump and last reducer launch)\n"
            "Phase 2.7 dispatch needed — STOP per\n"
            "feedback_systematic_debugging."
        )
        return 0

    assert first_diverge_step is not None
    print(f"VERDICT (first divergent step = {first_diverge_step}):")
    print(f"  diverged group(s): {sorted(first_diverge_groups)}")
    print()
    print("Per-group culprit table:")
    for group_label, _ in ALL_GROUPS:
        results = per_step_group_results.get((first_diverge_step, group_label), [])
        any_div = any(not eq for _, eq, _ in results)
        if any_div:
            divnames = [name for name, eq, _ in results if not eq]
            print(f"  Group {group_label}: DIVERGE in {divnames}")
        else:
            present = len(results) > 0
            tag = "EQUAL" if present else "(no dump)"
            print(f"  Group {group_label}: {tag}")
    print()
    print("Next-step interpretation (per Phase 2.6 decision tree):")

    # Subcase (1): only Group G diverges.
    if first_diverge_groups == {"G (PER replay sample)"}:
        print(
            "  Group G (PER replay sample) is the ONLY divergent group.\n"
            "  → step-2 PER K-loop sample IS the bug, despite Phase 2.5's\n"
            "    end-of-step PER state EQUAL. Most likely: the K-loop's\n"
            "    rl_per_sample is reading from `replay_h_t_d` /\n"
            "    `replay_h_tp1_d` which were written by a non-deterministic\n"
            "    `rl_per_push_ring`/flush kernel earlier in step 2. Inspect\n"
            "    `rl_per_push_ring.cu` and `rl_per_push_flush.cu` for\n"
            "    parallel-reduction / atomicAdd patterns. (Aliased pearl-\n"
            "    twin to Phase 2's PER tree rebuild bug.)"
        )
        return 1

    # Subcase (2): only Group H diverges (or H is the most upstream).
    if "H (head fwd outputs)" in first_diverge_groups and \
       "G (PER replay sample)" not in first_diverge_groups:
        print(
            "  Group H (head fwd outputs) diverges with EQUAL Group G.\n"
            "  → head forward kernel is non-det. Most likely candidate:\n"
            "    NoisyLinear noise resampling (factored noise) between\n"
            "    online forward and target forward, especially if the\n"
            "    resampling kernel uses parallel write to shared state\n"
            "    or if cuBLAS PEDANTIC isn't actually active (sanity-\n"
            "    check `[cublas_determinism]` line in run logs). Less\n"
            "    likely but possible: IQN tau sampling order across the\n"
            "    K-loop, or NoisyLinear reset_noise + forward overlap.\n"
            "  → check rl_iqn_forward / rl_noisy_linear_forward for\n"
            "    parallel-reduction or atomic patterns; verify cuBLAS\n"
            "    handle's math mode is CUBLAS_PEDANTIC_MATH at every\n"
            "    site (mamba2_block, dqn, iqn — 3 sites)."
        )
        return 1

    # Subcase (3): Group I diverges but Group J doesn't (or upstream of J).
    # NB: if Group I diverges, Group J will almost always diverge too
    # because J is downstream of the same head backward kernels. But
    # the KEY distinction is: do the pre-reduce scratches (I) diverge,
    # or do they match while the post-reduce final grads (dumped by
    # dump_mamba2_state_for_debug under "grad_vsn_w" etc.) diverge?
    # That comparison happens IN the mamba2 dump verdict; we report it
    # via interpretation hint here.
    if "I (pre-reduce grad scratch)" in first_diverge_groups:
        print(
            "  Group I (pre-reduce per-batch/per-row grad scratches)\n"
            "  diverges. → upstream per-batch backward kernel is non-det.\n"
            "  Candidates by scratch:\n"
            "    cfc_grad_*_scratch  → cfc_step_backward_batched\n"
            "    vsn_grad_*_scratch  → variable_selection_bwd\n"
            "    attn_grad_q_scratch → attention_pool_bwd\n"
            "    grad_ln_*_per_row   → layer_norm_bwd\n"
            "  Inspect which scratch's first divergence is the most\n"
            "  upstream (smallest k_iter or smallest n_rows index).\n"
            "  Look for parallel reductions with __threadfence (not\n"
            "  __syncthreads), atomicAdd, or warp-shuffle ordering\n"
            "  issues. Canonical Phase-2 fix: single block + __syncthreads\n"
            "  + fixed thread-ID accumulation order."
        )
    elif "J (per-head grad_h_t)" in first_diverge_groups:
        print(
            "  Group J (per-head grad_h_t) diverges with EQUAL Group I.\n"
            "  → `grad_h_accumulate_scaled` or a head backward writing\n"
            "    grad_h_t (without scratch reduction) is the culprit.\n"
            "  → check grad_h_accumulate.cu: should be sole-writer-per-i\n"
            "    (element-wise `+=`). If λ is host-passed scalar, host-\n"
            "    side ISV mirror read should be order-stable. Also check\n"
            "    head bwd kernels (dqn_grad_w_b_h_t / pi_head_bwd / v_head_\n"
            "    bwd / iqn_bwd_relu_hadamard) — they all write grad_h_t\n"
            "    as OVERWRITE per (b, c) sole writer; if any uses `+=`\n"
            "    with parallel writes, that's the bug."
        )

    # If post-reduce final grads (dumped by mamba2 dumper) diverge while
    # pre-reduce scratches (Group I) are EQUAL, that's the canonical
    # `reduce_axis0` bug. Surface that diagnostic in the message even
    # if we can't directly check the mamba2 dump here.
    if "I (pre-reduce grad scratch)" not in first_diverge_groups and \
       "J (per-head grad_h_t)" not in first_diverge_groups and \
       first_diverge_groups:
        print(
            "\n  WARNING: divergent group(s) not in expected G/H/I/J set\n"
            "  → unexpected finding, STOP per dispatch rule."
        )
        return 1

    # Cross-check: if BOTH I and J are EQUAL but the mamba2 dump
    # shows trunk grads DIVERGE (per Phase 2.5's finding), the only
    # remaining place is `reduce_axis0` / `layer_norm_reduce_param_grads`.
    print(
        "\n  Cross-check: if mamba2 dump shows final trunk grads (post-\n"
        "  reduce) DIVERGE but Group I scratches above are EQUAL → the\n"
        "  reducer kernel itself (`reduce_axis0` or\n"
        "  `layer_norm_reduce_param_grads`) is the bug. Apply canonical\n"
        "  Phase-2-PER-rebuild fix: single block + __syncthreads +\n"
        "  fixed thread-ID accumulation order, no __threadfence."
    )
    return 1


if __name__ == "__main__":
    sys.exit(main())