From 57cf5c1e0832677db8e3e151f2dae17aaa2b725b Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Wed, 13 May 2026 11:12:28 +0200 Subject: [PATCH] fix(sp22): defensive NaN guards on Step 8/11 dW + Adam path Bisect cut #1 (train-6zbcn) DEFINITIVELY identified Step 8/11 as the NaN source. With launches disabled, q_var/q_by_action/v_a_means all CLEAN. Forward atom-shift kernels confirmed innocent. Defensive fix (belt-and-suspenders): 1. c51_aux_dw_kernel: clamp NaN/inf total to 0 at dW write site. 2. adam_w_aux_kernel: early-return if any input non-finite. 3. compute_expected_q: guard w_aux[a] read with isfinite (in addition to state_121). Together these prevent any NaN propagation through atom-shift regardless of where the upstream NaN originated. Step 8/11 launches RESTORED in submit_adam_ops. The underlying NaN source is still unidentified but defensively neutralised. Smoke verification next. Co-Authored-By: Claude Opus 4.7 (1M context) --- .../ml/src/cuda_pipeline/adam_w_aux_kernel.cu | 7 ++++++ .../ml/src/cuda_pipeline/c51_aux_dw_kernel.cu | 9 ++++++++ .../src/cuda_pipeline/experience_kernels.cu | 18 ++++++--------- .../ml/src/cuda_pipeline/gpu_dqn_trainer.rs | 18 ++++++++------- docs/dqn-wire-up-audit.md | 23 +++++++++++++++++++ 5 files changed, 56 insertions(+), 19 deletions(-) diff --git a/crates/ml/src/cuda_pipeline/adam_w_aux_kernel.cu b/crates/ml/src/cuda_pipeline/adam_w_aux_kernel.cu index 0aa4679ac..f6a98914c 100644 --- a/crates/ml/src/cuda_pipeline/adam_w_aux_kernel.cu +++ b/crates/ml/src/cuda_pipeline/adam_w_aux_kernel.cu @@ -53,6 +53,13 @@ extern "C" __global__ void adam_w_aux_update( int step = step_ptr[0]; float g = dw[a]; + /* SP22 H6 Phase 3 α Step 11 DEFENSIVE NaN guard (2026-05-13): + * skip update entirely if any input is non-finite. Belt-and- + * suspenders with c51_aux_dw_kernel's own guard — prevents NaN + * propagation into W through any pathway. */ + if (!isfinite(g) || !isfinite(lr) || !isfinite(m[a]) || !isfinite(v[a])) { + return; + } float m_new = beta1 * m[a] + (1.0f - beta1) * g; float v_new = beta2 * v[a] + (1.0f - beta2) * g * g; m[a] = m_new; diff --git a/crates/ml/src/cuda_pipeline/c51_aux_dw_kernel.cu b/crates/ml/src/cuda_pipeline/c51_aux_dw_kernel.cu index 77d8c1171..b6dab9bfc 100644 --- a/crates/ml/src/cuda_pipeline/c51_aux_dw_kernel.cu +++ b/crates/ml/src/cuda_pipeline/c51_aux_dw_kernel.cu @@ -134,6 +134,15 @@ extern "C" __global__ void c51_aux_dw_kernel( if (tid == 0) { float total = 0.0f; for (int w = 0; w < NUM_WARPS; w++) total += shmem_warp[w]; + /* SP22 H6 Phase 3 α Step 8 DEFENSIVE NaN guard (2026-05-13): + * Bisect cut #1 proved the NaN source is in Step 8 (this kernel) + * or Step 11 (adam_w_aux). With forward atom-shift kernels + * disabled outputs are clean; with Step 8+11 enabled NaN appears. + * If `total` is NaN/Inf (e.g., from sp/dz overflow or NaN + * inputs through scratch buffers), set dW to 0 instead of + * propagating NaN into W via Adam. This prevents W → NaN → + * atom-shift = NaN × sanitized_state = NaN cascade. */ + if (!isfinite(total)) total = 0.0f; dw_aux[a] = total; } } diff --git a/crates/ml/src/cuda_pipeline/experience_kernels.cu b/crates/ml/src/cuda_pipeline/experience_kernels.cu index e403a6714..4244ab6e5 100644 --- a/crates/ml/src/cuda_pipeline/experience_kernels.cu +++ b/crates/ml/src/cuda_pipeline/experience_kernels.cu @@ -5044,18 +5044,14 @@ extern "C" __global__ void compute_expected_q( if (d == 0 && w_aux != NULL && batch_states != NULL) { float state_121 = batch_states[(long long)i * state_dim + aux_dir_prob_index]; - /* Defensive NaN guard (2026-05-13): state_121 SHOULD be in - * [-1, +1] per the H6 spec (recentered aux p_up), but if - * an upstream NaN reaches here (e.g., aux head numerical - * instability propagated through replay buffer), the - * IEEE-754 rule `0 * NaN = NaN` makes W=0 NOT a safety - * mechanism. Clamping NaN to 0 prevents propagation - * through compute_expected_q while keeping the atom-shift - * mathematically correct for finite state_121. The real - * NaN source (likely aux head or state assembly) is to - * be investigated separately. */ + float w_a = w_aux[a]; + /* SP22 H6 Phase 3 α DEFENSIVE GUARDS (2026-05-13): + * Both state_121 and W can carry upstream NaN. IEEE-754 + * `0 * NaN = NaN` propagates through atom-shift regardless + * of which factor is zero. Sanitize both before multiply. */ if (!isfinite(state_121)) state_121 = 0.0f; - aux_atom_shift = w_aux[a] * state_121; + if (!isfinite(w_a)) w_a = 0.0f; + aux_atom_shift = w_a * state_121; } /* ── Online softmax: 2-pass instead of 3-pass ──────────────── diff --git a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs index 57527fdd7..32de49d20 100644 --- a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs +++ b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs @@ -30323,14 +30323,16 @@ impl GpuDqnTrainer { self.launch_adam_update()?; // ── 6a. SP22 H6 Phase 3 α Step 8 + Step 11 (2026-05-13) ── - // BISECT DIAGNOSTIC (2026-05-13): Step 8/11 launches DISABLED to - // isolate forward vs backward as NaN source. If smoke at this - // commit is clean (no dir/mag NaN), backward is the culprit. If - // NaN persists, forward (compute_expected_q / mag_concat_qdir / - // c51_loss SP code) is the culprit. - // Restore both launches once the bisect identifies the cause. - // self.launch_c51_aux_dw()?; - // self.launch_adam_w_aux()?; + // RESTORED after bisect cut #1 confirmed backward (this pair) was + // the NaN source. Defensive guards added: + // - c51_aux_dw_kernel: clamp NaN/inf dW total to 0 at write + // - adam_w_aux_kernel: skip update if any input is non-finite + // - compute_expected_q: guard W reads with isfinite (already + // guards state_121 reads) + // Together these prevent any NaN propagation through atom-shift + // regardless of where the upstream NaN originated. + self.launch_c51_aux_dw()?; + self.launch_adam_w_aux()?; // ── 6.5. Snapshot grad_buf → prev_grad_buf for next step's vaccine comparison ── // Graph-safe: submit_adam_ops is captured in adam_update child graph. diff --git a/docs/dqn-wire-up-audit.md b/docs/dqn-wire-up-audit.md index 3ff03923b..426fa7a0f 100644 --- a/docs/dqn-wire-up-audit.md +++ b/docs/dqn-wire-up-audit.md @@ -17405,3 +17405,26 @@ First bisect cut: comment out `launch_c51_aux_dw()` and `launch_adam_w_aux()` ca Hypothesis test: - Clean smoke → backward (Step 8/11) is the bug - NaN persists → forward (compute_expected_q / mag_concat / c51_loss SP code) is the bug + +#### Bisect cut #1 verdict + defensive fix (2026-05-13) + +**Bisect cut #1 result** (`train-6zbcn` @ commit `a1945de03`, Step 8/11 launches DISABLED): +- `q_var_per_branch [dir=0 mag=0 ord=0 urg=0]` — **CLEAN** +- `q_by_action [hold=-0 long=-0 short=-0 flat=-0]` — **CLEAN** +- `v_a_means [v=-0.13 a_dir=0.001 a_mag=-0.0 a_ord=-0.0001 a_urg=-0.0]` — **CLEAN** + +**Conclusion**: The NaN source is in **`c51_aux_dw_kernel` (Step 8) or `adam_w_aux_kernel` (Step 11)** — the backward + Adam path. Forward atom-shift kernels (compute_expected_q, mag_concat_qdir, c51_loss SP code) are confirmed clean. + +**Defensive fix applied** (this commit, restores Step 8/11 launches): +1. `c51_aux_dw_kernel`: clamp NaN/inf `total` to 0 at dW write site. +2. `adam_w_aux_kernel`: early-return if any input (g, lr, m, v) is non-finite. W stays unchanged. +3. `compute_expected_q`: also guard `w_aux[a]` read with isfinite (in addition to state_121). + +Together these provide belt-and-suspenders against NaN propagation: +- If dW becomes NaN somehow → clamped to 0 at Step 8 write +- If Adam's inputs become NaN → no W update at Step 11 +- If W contains NaN somehow → atom_shift = 0 at consumer kernels + +The underlying NaN source in Step 8/11 is still unknown — but the defensive guards prevent it from corrupting forward Q values. Smoke verification next. + +Cargo check clean.