From 60e96bf55c1bf69685daa5f56c758efa813fcfd2 Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Fri, 29 May 2026 08:45:24 +0200 Subject: [PATCH] diag(ml-alpha): add nan_scan kernel + trainer helper for step-4 NaN hunt MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Single-block warp-shuffle scanner that printfs the first non-finite index/value in a labeled [N] float buffer. Loaded unconditionally so the module handle is always live; per-launch gating uses the new `nan_scan_enabled` field, set from `FOXHUNT_NAN_SCAN=1` at trainer construction. Production training pays zero cost when the env var is unset. label_id table (kernel + host helper kept in lockstep): 0 h_t, 1 v_pred, 2 q_logits, 3 pi_logits, 4 advantages, 5 ss_pi_l_pi_per_batch. Diagnostic context: residual ~33% NaN at step 4 (seed 16962) after the atomicAdd/V-envelope/0×Inf fixes (10d4614fb, b4aadff75, a6acc25ec). Adding printf to v_head_fwd_bwd or running under nsys cured it, so the remaining cause is microarchitectural (timing/memory-ordering) or localized to a kernel not yet on the trail. nan_scan pinpoints the first stage whose buffer goes bad when the NaN fires; instrumentation of the per-step pipeline lands in a follow-up commit. Co-Authored-By: Claude Opus 4.7 --- crates/ml-alpha/build.rs | 1 + crates/ml-alpha/cuda/nan_scan.cu | 101 ++++++++++++++++++++++ crates/ml-alpha/src/trainer/integrated.rs | 84 ++++++++++++++++++ 3 files changed, 186 insertions(+) create mode 100644 crates/ml-alpha/cuda/nan_scan.cu diff --git a/crates/ml-alpha/build.rs b/crates/ml-alpha/build.rs index b1fe3a42f..0f686f256 100644 --- a/crates/ml-alpha/build.rs +++ b/crates/ml-alpha/build.rs @@ -137,6 +137,7 @@ const KERNELS: &[&str] = &[ "snapshot_aos_to_soa", // AoS→SoA scatter: one thread per snapshot reads contiguous Mbp10RawInput, writes into 10 SoA device buffers; replaces host nested loops + 10 DtoD copies "gpu_sample_and_gather", // GPU-resident batch sampler: random file+anchor sampling + AoS→SoA gather from pre-uploaded dataset; eliminates ALL per-step CPU data loading "rl_v_target_envelope_update", // Phase 2.0 (2026-05-28): done-gated EMA of trade magnitudes → V regression target clamp envelope (slots 593-596); pairs with v_head_bwd clamp to bound (v_pred - target)² regardless of fat-tail trade-close spikes + "nan_scan", // 2026-05-29: diagnostic non-finite scanner for step-4 intermittent NaN hunt; opt-in via FOXHUNT_NAN_SCAN=1; printfs first non-finite (idx, val) per labeled buffer ]; // Cache bust v31 — five new reduce / derive kernels populate the input diff --git a/crates/ml-alpha/cuda/nan_scan.cu b/crates/ml-alpha/cuda/nan_scan.cu new file mode 100644 index 000000000..b963219db --- /dev/null +++ b/crates/ml-alpha/cuda/nan_scan.cu @@ -0,0 +1,101 @@ +// nan_scan.cu — diagnostic non-finite scanner for the alpha-rl pipeline. +// +// Single-block, single-warp scanner over an `[N] float` buffer. Reports +// (via CUDA printf) the first batch/element index whose value is not +// finite (NaN, +Inf, -Inf, denormal sentinel). Used to pinpoint which +// stage of the training step first produces a non-finite buffer when the +// step-4 intermittent NaN fires. +// +// Bug context (session 2026-05-29 follow-up): +// - Intermittent NaN at step 4 (~33% rate, seed 16962). +// - Heisenbug: adding printf to v_head_fwd_bwd or running under nsys +// profiling makes it disappear, ruling out a deterministic math bug. +// - Prior atomicAdd / V envelope / 0×Inf fixes (10d4614fb, b4aadff75, +// a6acc25ec) closed the cleanly-reproducible failure modes; the +// residual is microarchitectural (timing / memory ordering) or +// localized to a kernel not yet on the trail. +// +// The scanner runs OPT-IN via `FOXHUNT_NAN_SCAN=1`; the trainer holds a +// boolean derived from that env var at construction and skips the kernel +// launch when disabled. When enabled, every key buffer in the per-step +// pipeline is scanned immediately after its producer kernel returns; the +// first scan that reports a non-finite hit identifies the producing +// stage. +// +// label_id values (kept here so the host helper and kernel stay in sync; +// adding a new label requires updating BOTH sides in the same commit per +// `feedback_no_partial_refactor`): +// +// 0 — h_t (encoder output, post forward_encoder) +// 1 — v_pred (V-head forward) +// 2 — q_logits (DQN distributional Q forward, post-flatten) +// 3 — pi_logits (policy head forward) +// 4 — advantages (compute_advantage_return output) +// 5 — ss_pi_l_pi_per_batch (PPO surrogate per-batch loss) +// +// Format printed on hit (one line per offending block): +// "NAN_SCAN step= label= first_bad_idx= val=\n" +// +// Single block + single warp = no cross-block sync, no shared mem. +// Warp-shuffle reduction picks the lowest offending index. Lane 0 +// printfs the result; if no offender is found, nothing is printed +// (clean-run noise suppression). + +#include +#include +#include + +extern "C" __global__ void nan_scan( + const float* __restrict__ buf, // [N] + int n, + int label_id, + int step +) { + if (blockIdx.x != 0) return; // single-block contract + const int tid = threadIdx.x; + const int nt = blockDim.x; + + // Each lane scans its strided slice for the first non-finite value. + // INT_MAX sentinel means "no hit on this lane". + int my_idx = 0x7fffffff; + float my_val = 0.0f; + for (int i = tid; i < n; i += nt) { + const float v = buf[i]; + if (!isfinite(v)) { + if (i < my_idx) { + my_idx = i; + my_val = v; + } + } + } + + // Warp-shuffle min on (idx, val) — lane with smallest idx wins. + // val rides along with idx via parallel shuffle so the winning val + // matches the winning idx. + #pragma unroll + for (int off = 16; off > 0; off >>= 1) { + const int other_idx = __shfl_xor_sync(0xffffffff, my_idx, off); + const float other_val = __shfl_xor_sync(0xffffffff, my_val, off); + if (other_idx < my_idx) { + my_idx = other_idx; + my_val = other_val; + } + } + + // Lane 0 reports. Hardcoded label names — kept in lockstep with the + // host-side enum. + if (tid == 0 && my_idx != 0x7fffffff) { + const char* label = "unknown"; + switch (label_id) { + case 0: label = "h_t"; break; + case 1: label = "v_pred"; break; + case 2: label = "q_logits"; break; + case 3: label = "pi_logits"; break; + case 4: label = "advantages"; break; + case 5: label = "ss_pi_l_pi_per_batch"; break; + default: break; + } + printf("NAN_SCAN step=%d label=%s first_bad_idx=%d val=%f n=%d\n", + step, label, my_idx, my_val, n); + } +} diff --git a/crates/ml-alpha/src/trainer/integrated.rs b/crates/ml-alpha/src/trainer/integrated.rs index 43de6059e..710b01cd8 100644 --- a/crates/ml-alpha/src/trainer/integrated.rs +++ b/crates/ml-alpha/src/trainer/integrated.rs @@ -202,6 +202,14 @@ const RL_ATOM_SUPPORT_UPDATE_CUBIN: &[u8] = // §3 Sub-phase 2.0. const RL_V_TARGET_ENVELOPE_UPDATE_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/rl_v_target_envelope_update.cubin")); +// NaN scanner — 2026-05-29 diagnostic for step-4 intermittent NaN hunt. +// Single-block warp-shuffle min over an [N] float buffer; on hit, lane 0 +// printfs "NAN_SCAN step= label= first_bad_idx= val=". +// Opt-in via `FOXHUNT_NAN_SCAN=1`; the trainer caches the env var as a +// boolean at construction and gates `nan_scan` launches behind it so +// production training pays zero cost. +const NAN_SCAN_CUBIN: &[u8] = + include_bytes!(concat!(env!("OUT_DIR"), "/nan_scan.cubin")); // Q→π distillation gradient — audit 2026-05-24 vj5f6 followup. ADDS // λ × (π_new - π_target) to pi_grad_logits after the PPO surrogate // backward, where π_target = softmax(E_Q[s,*] / τ). Couples Q's @@ -1263,6 +1271,19 @@ pub struct IntegratedTrainer { pub ss_iqn_grad_b_out_d: CudaSlice, pub ss_iqn_grad_w_embed_d: CudaSlice, pub ss_iqn_grad_b_embed_d: CudaSlice, + + // ── NaN scanner (2026-05-29 step-4 NaN diagnostic) ────────────────── + // Loaded unconditionally so the module/function handles are always + // available; per-launch gating uses `nan_scan_enabled` so production + // training pays only a single `if` per insertion point. + _nan_scan_module: Arc, + nan_scan_fn: CudaFunction, + /// `FOXHUNT_NAN_SCAN=1` at trainer construction → true. When true, + /// `nan_scan` is launched after each key forward in the per-step + /// pipeline (encoder h_t, V forward, Q forward, π forward, + /// compute_advantage_return, PPO surrogate). When false, all scan + /// launches are skipped — zero per-step cost. + nan_scan_enabled: bool, } impl IntegratedTrainer { @@ -1690,6 +1711,23 @@ impl IntegratedTrainer { .load_function("abs_copy") .context("load abs_copy")?; + // NaN scanner (step-4 NaN diagnostic, 2026-05-29). Loaded + // unconditionally so the handles are always live; per-launch + // gating uses `nan_scan_enabled` so production training pays + // zero per-step cost when the env var is unset. + let nan_scan_module = ctx + .load_cubin(NAN_SCAN_CUBIN.to_vec()) + .context("load nan_scan cubin")?; + let nan_scan_fn = nan_scan_module + .load_function("nan_scan") + .context("load nan_scan")?; + let nan_scan_enabled = std::env::var("FOXHUNT_NAN_SCAN") + .map(|v| v == "1" || v.eq_ignore_ascii_case("true")) + .unwrap_or(false); + if nan_scan_enabled { + eprintln!(" ml-alpha: FOXHUNT_NAN_SCAN=1 — NaN scanner active on key buffers"); + } + // Streaming kernels for the variance-ratio + kurtosis EMAs. // Per the streaming refactor: these now fold ACROSS STEPS (Welford-EMA on // per-batch means in ISV state slots) instead of across batch @@ -2806,6 +2844,9 @@ impl IntegratedTrainer { ss_iqn_grad_b_out_d, ss_iqn_grad_w_embed_d, ss_iqn_grad_b_embed_d, + _nan_scan_module: nan_scan_module, + nan_scan_fn, + nan_scan_enabled, hot, } .with_controllers_bootstrapped()?) @@ -3839,6 +3880,49 @@ impl IntegratedTrainer { Ok(()) } + /// Diagnostic NaN scanner (2026-05-29). Launches `nan_scan` on a + /// device buffer; on hit, lane 0 of the kernel printfs: + /// `NAN_SCAN step= label= first_bad_idx= val= n=` + /// + /// No-ops when `FOXHUNT_NAN_SCAN` was unset at trainer construction. + /// `label_id` maps to the hardcoded names in `nan_scan.cu`: + /// 0 h_t, 1 v_pred, 2 q_logits, 3 pi_logits, + /// 4 advantages, 5 ss_pi_l_pi_per_batch. + /// Adding a label requires updating BOTH this enum AND the switch in + /// the kernel in the same commit per `feedback_no_partial_refactor`. + /// + /// Single block, 32 threads — no smem, no atomics. Cheap enough that + /// instrumenting six insertion points adds ~10-20 μs/step worst case, + /// but production keeps the env var unset and pays nothing. + pub fn nan_scan( + &self, + label_id: i32, + step: u32, + buf_ptr: u64, + n: usize, + ) -> Result<()> { + if !self.nan_scan_enabled { + return Ok(()); + } + let n_i = n as i32; + let step_i = step as i32; + let mut args = RawArgs::new(); + args.push_ptr(buf_ptr); + args.push_i32(n_i); + args.push_i32(label_id); + args.push_i32(step_i); + let mut ptrs = args.build_arg_ptrs(); + unsafe { + raw_launch( + self.nan_scan_fn.cu_function(), + (1, 1, 1), (32, 1, 1), 0, + self.raw_stream, + &mut ptrs[..args.len()], + ).map_err(|e| anyhow::anyhow!("nan_scan(label={}): {:?}", label_id, e))?; + } + Ok(()) + } + /// Phase R7a: launch `abs_copy` — element-wise `dst[b] = fabsf(src[b])`. /// Feeds the `|reward|` signal into `ema_update_on_done` for the /// `MEAN_ABS_PNL_EMA` slot without mutating the signed `rewards_d`.