#!/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 Each dir contains `step_{0,1,2,3}_g_.bin`, `step_{0,1,2,3}_h_.bin`, `step_{0,1,2,3}_i_.bin`, `step_{0,1,2,3}_j_.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___.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())