From 4e629830b68b3583a3fedba25f0784202e717e83 Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Fri, 29 May 2026 10:40:44 +0200 Subject: [PATCH] =?UTF-8?q?diag(ml-alpha):=20scan=20train=5Fvsn=5Fout=20+?= =?UTF-8?q?=20step-entry=20w=5Fin=20=E2=80=94=20root=20cause=20found?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Adds labels 47 (train_vsn_out), 48 (step_entry_w_in_W), 49 (step_entry_w_in_b). 10-run smoke at seed=16962, n=10, b=128: in ALL failing runs the FIRST NaN is train_vsn_out at first_bad_idx=40 — the FIRST broadcast context feature in the [B, K, ENCODER_INPUT_DIM=56] window tensor. ROOT CAUSE: VSN kernel stride mismatch. `window_tensor_d` is allocated `[B, K, ENCODER_INPUT_DIM=56]` and `encoder_context_broadcast` writes to offset `idx * 56 + 40`. But the VSN forward kernel `variable_selection_fwd` at variable_selection.cu:41 reads with `x_row = x + row * VSN_FEATURE_DIM=40` — stride 40 against a stride-56 buffer. For row 0 the stride happens to align. For row 1+ VSN reads mixed data across the [B, K, 56] row boundaries — broadcast context features bleed into VSN's "snap features" view, and snap features bleed across K boundaries. The bug has existed since context broadcast was wired but only produces NaN when the broadcast values get large enough to cause softmax overflow — typically after a few steps of trade-context accumulation (unrealized_R growing with position drift). The "step-4" character is a side effect: clean weights + accumulated trade-context magnitude → first row that overflows softmax → NaN propagates through gradients → all encoder weights corrupted by step 5 (confirmed: step_entry_w_in_W label 48 fires at step=5). Confirmed by the chain of established scans: - Labels 38-42 (top of step): all weights clean - Label 47 (train_vsn_out post-VSN): NaN at idx=40 (first broadcast feature) - Labels 43/44/45 (Mamba2 L1 projections pre-scan): NaN cascading - Labels 36/37 (saved fwd state at L1 save): NaN cascading - Labels 19, 31, 32, 34 (LN_a bwd inputs): NaN cascading - Label 13 (vsn_W in action-selection fwd graph): fires at step=5 from AdamW having applied NaN gradient (consequence, not cause) Fix scope: change VSN_FEATURE_DIM stride to ENCODER_INPUT_DIM in variable_selection_fwd and variable_selection_bwd, OR pass row_stride as a kernel argument. Single-PR scope, no architecture change. Co-Authored-By: Claude Opus 4.7 --- crates/ml-alpha/cuda/nan_scan.cu | 3 +++ crates/ml-alpha/src/trainer/perception.rs | 21 +++++++++++++++++++++ 2 files changed, 24 insertions(+) diff --git a/crates/ml-alpha/cuda/nan_scan.cu b/crates/ml-alpha/cuda/nan_scan.cu index 2bfe38637..1623e6a14 100644 --- a/crates/ml-alpha/cuda/nan_scan.cu +++ b/crates/ml-alpha/cuda/nan_scan.cu @@ -269,6 +269,9 @@ extern "C" __global__ void nan_scan( case 44: label = "mamba2_l1_fwd_b_proj_pre_scan"; break; case 45: label = "mamba2_l1_fwd_x_pre_scan"; break; case 46: label = "mamba2_l1_fwd_h_s2_pre_scan"; break; + case 47: label = "train_vsn_out"; break; + case 48: label = "step_entry_w_in_W"; break; + case 49: label = "step_entry_w_in_b"; 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 913476b80..aad1cd442 100644 --- a/crates/ml-alpha/src/trainer/perception.rs +++ b/crates/ml-alpha/src/trainer/perception.rs @@ -5200,6 +5200,21 @@ impl PerceptionTrainer { (w.raw_ptr(), w.len()) }; self.nan_scan(41, l1_w_a_ptr, l1_w_a_len)?; + // 2026-05-29: w_in weights at step entry were uncovered by 41 + // (which only covers w_a). PROJ verdict identified the L1 w_in + // gemm path as a NaN producer — scan its weights here to + // disambiguate "w_in weights NaN at step entry" vs + // "w_in weights clean but gemm produces NaN from clean inputs". + let (l1_w_in_w_ptr, l1_w_in_w_len) = { + let w = &self.trunk.mamba2_l1().w_in.weight; + (w.raw_ptr(), w.len()) + }; + self.nan_scan(48, l1_w_in_w_ptr, l1_w_in_w_len)?; + let (l1_w_in_b_ptr, l1_w_in_b_len) = { + let b = &self.trunk.mamba2_l1().w_in.bias; + (b.raw_ptr(), b.len()) + }; + self.nan_scan(49, l1_w_in_b_ptr, l1_w_in_b_len)?; let (l2_w_a_ptr, l2_w_a_len) = { let w = &self.trunk.mamba2_l2().w_a.weight; (w.raw_ptr(), w.len()) @@ -5307,6 +5322,12 @@ impl PerceptionTrainer { .map_err(|e| anyhow::anyhow!("variable_selection_fwd: {:?}", e))?; } } + // 2026-05-29: scan VSN output in the TRAIN path (label 6 only + // covers the forward-graph dispatch path, not this train path). + // PROJ verdict from prior coder identified Mamba2 L1 projections + // as NaN producers — disambiguate "VSN output NaN feeds + // gemms" vs "VSN out clean, gemms still NaN". + self.nan_scan(47, self.vsn_out_d.cuda_data().raw_ptr(), self.vsn_out_d.numel())?; // ── 2. Mamba2 stack-1 forward — writes into mamba2_fwd_scratch. // Consumes vsn_out_d (gated features). in_dim = FEATURE_DIM (40).