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

diag(ml-alpha): scan persistent weights at step entry (labels 38-42)

Browse Source 
Adds top-of-step temporal-localization scans BEFORE snap_assemble /
forward kernels to find the FIRST step on which any persistent
weight first observes a NaN:

  - 38 step_entry_window_tensor (stale from prior step — upstream probe)
  - 39 step_entry_vsn_W
  - 40 step_entry_vsn_b
  - 41 step_entry_mamba2_l1_w_a
  - 42 step_entry_mamba2_l2_w_a

The integrated trainer bumps ISV[548] before calling
dispatch_train_step_inner, so the printed step=N matches the step
about to execute. Persistent weights (39-42) carry NaN across steps,
so a hit at step N means step N-1's backward + AdamW path produced
the NaN (AdamW NaN-passthrough class). A hit on label 38 (stale
window_tensor_d) points upstream to snap_feature_assemble_batched
or the SoA data path.

Per feedback_no_partial_refactor, host scan calls + kernel switch
arms land in the same commit.
This commit is contained in:
jgrusewski
2026-05-29 10:04:26 +02:00
parent 6748f6432a
commit 85531c8c3a
2 changed files with 69 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

36
crates/ml-alpha/cuda/nan_scan.cu
View File

@@ -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=<step> label=<name> first_bad_idx=<i> val=<f>\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];

33
crates/ml-alpha/src/trainer/perception.rs
View File

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