diff --git a/crates/ml-alpha/cuda/nan_scan.cu b/crates/ml-alpha/cuda/nan_scan.cu index f419e698d..fbbe1e582 100644 --- a/crates/ml-alpha/cuda/nan_scan.cu +++ b/crates/ml-alpha/cuda/nan_scan.cu @@ -98,6 +98,37 @@ // H2 verdict: {36, 37} CLEAN at save time but {20, 21} NaN at bwd entry → // shared-scratch overwrite by a second forward pass. // +// Step-entry temporal-localization scans (added 2026-05-29 follow-up — find +// the FIRST step where any persistent weight or step-input first goes NaN). +// Reading these AT THE TOP OF the per-step dispatch — BEFORE any forward +// kernel runs — disambiguates between: +// (a) the corruption originating in step N-1's backward / AdamW (label +// 39-42 fire at step N entry, label 38 clean → AdamW NaN-passthrough) +// (b) the corruption originating in the data path upstream of perception +// (label 38 NaN at step N entry → snap_assemble / SoA source). +// Per pearl_atomicadd_masks_v_instability the bug is in step ≤4; these +// scans give the earliest step number on which a NaN is observable. +// +// 38 — step_entry_window_tensor (window_tensor_d at top-of-step entry — +// STALE from prior step before snap_assemble +// overwrites; NaN here means prior step's +// snap_assemble produced NaN that was never +// cleared, i.e. upstream of perception) +// 39 — step_entry_vsn_W (trunk.vsn_w_d at top-of-step entry — +// persistent across steps; NaN here means +// step N-1's VSN backward + AdamW path +// produced NaN) +// 40 — step_entry_vsn_b (trunk.vsn_b_d at top-of-step entry — +// same diagnostic as 39 for the bias) +// 41 — step_entry_mamba2_l1_w_a (trunk.mamba2_l1().w_a.weight — Mamba2 +// L1 A-projection weight at step entry; +// NaN here means step N-1's Mamba2 L1 +// backward + AdamW path produced NaN) +// 42 — step_entry_mamba2_l2_w_a (trunk.mamba2_l2().w_a.weight — Mamba2 +// L2 A-projection weight at step entry; +// comparison probe to 41 — if 42 NaN but +// 41 clean, corruption is L2-side first) +// // Format printed on hit (one line per offending block): // "NAN_SCAN step= label= first_bad_idx= val=\n" // @@ -199,6 +230,11 @@ extern "C" __global__ void nan_scan( case 35: label = "ln_a_bwd_out_grad_in"; break; case 36: label = "mamba2_l1_fwd_a_proj_at_save"; break; case 37: label = "mamba2_l1_fwd_x_at_save"; break; + case 38: label = "step_entry_window_tensor"; break; + case 39: label = "step_entry_vsn_W"; break; + case 40: label = "step_entry_vsn_b"; break; + case 41: label = "step_entry_mamba2_l1_w_a"; break; + case 42: label = "step_entry_mamba2_l2_w_a"; break; default: break; } const int step = (int)isv[RL_STEP_COUNTER_ISV_INDEX]; diff --git a/crates/ml-alpha/src/trainer/perception.rs b/crates/ml-alpha/src/trainer/perception.rs index ce3e776d8..913476b80 100644 --- a/crates/ml-alpha/src/trainer/perception.rs +++ b/crates/ml-alpha/src/trainer/perception.rs @@ -5173,6 +5173,39 @@ impl PerceptionTrainer { } } + // Step-4 NaN temporal-localization scans (labels 38-42, added + // 2026-05-29 follow-up to pearl_atomicadd_masks_v_instability). + // Read AT TOP OF STEP, BEFORE snap_assemble / forward kernels — + // tells us the EARLIEST step on which any persistent weight or + // step-input first observes a NaN. Persistent params (39-42) + // survive across steps, so a NaN here means step N-1's + // backward + AdamW produced NaN (AdamW NaN-passthrough class). + // Label 38 reads window_tensor_d BEFORE snap_assemble overwrites + // it — STALE from prior step; if NaN here, prior step's + // snap_feature_assemble_batched wrote NaN that survived. + // + // Step counter is read from ISV[548] by the kernel; the + // perception trainer is called from IntegratedTrainer AFTER + // rl_increment_step bumps the counter, so the printed `step=N` + // matches the step about to execute. + self.nan_scan( + 38, + self.window_tensor_d.cuda_data().raw_ptr(), + self.window_tensor_d.numel(), + )?; + self.nan_scan(39, self.trunk.vsn_w_d.raw_ptr(), self.trunk.vsn_w_d.len())?; + self.nan_scan(40, self.trunk.vsn_b_d.raw_ptr(), self.trunk.vsn_b_d.len())?; + let (l1_w_a_ptr, l1_w_a_len) = { + let w = &self.trunk.mamba2_l1().w_a.weight; + (w.raw_ptr(), w.len()) + }; + self.nan_scan(41, l1_w_a_ptr, l1_w_a_len)?; + let (l2_w_a_ptr, l2_w_a_len) = { + let w = &self.trunk.mamba2_l2().w_a.weight; + (w.raw_ptr(), w.len()) + }; + self.nan_scan(42, l2_w_a_ptr, l2_w_a_len)?; + // AoS→SoA scatter: GPU kernel reads contiguous Mbp10RawInput from // mapped-pinned buffer, writes the 10 SoA device buffers. Skipped // when the SoA buffers are pre-filled by GPU gather kernels.