gitadmin/foxhunt
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

diag(ml-alpha): add nan_scan kernel + trainer helper for step-4 NaN hunt

Browse Source 
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 <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-05-29 08:45:24 +02:00
parent a6acc25ec8
commit 60e96bf55c
3 changed files with 186 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

1
crates/ml-alpha/build.rs
View File

@@ -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

101
crates/ml-alpha/cuda/nan_scan.cu Normal file
View File

@@ -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=<step> label=<name> first_bad_idx=<i> val=<f>\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 <cstdint>
#include <cuda_runtime.h>
#include <cstdio>
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);
}
}

84
crates/ml-alpha/src/trainer/integrated.rs
View File

@@ -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=<s> label=<name> first_bad_idx=<i> val=<v>".
// 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<f32>,
pub ss_iqn_grad_w_embed_d: CudaSlice<f32>,
pub ss_iqn_grad_b_embed_d: CudaSlice<f32>,
// ── 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<CudaModule>,
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=<step> label=<name> first_bad_idx=<i> val=<v> n=<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`.