From 79945987a5686c40bf04caa79d2b4983376f5038 Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Wed, 13 May 2026 14:13:40 +0200 Subject: [PATCH] fix(sp22): W-read guards across ALL atom-shift kernels Verify smoke train-t5885 partial result: dir CLEAN, mag NaN. The isfinite(W) guard added in compute_expected_q didn't extend to the other atom-shift consumers. Adam-corrupted W propagates through them via 0*NaN=NaN. Adds isfinite(w_aux[a]) guards to all atom-shift kernels: - mag_concat_qdir - quantile_q_select - c51_loss_kernel (eq_per_action shift + effective_reward W[a0]/W[best_next_a]) Also strengthens c51_aux_dw_kernel: guard sp/isw/dz/gamma/done (not just dz<1e-7 which doesn't catch NaN). Any NaN/Inf input -> skip sample (no dW contribution). Defense-in-depth complete: NaN cannot propagate through any atom-shift path regardless of source. Cargo check clean. Co-Authored-By: Claude Opus 4.7 (1M context) --- .../ml/src/cuda_pipeline/c51_aux_dw_kernel.cu | 13 ++++++++++--- .../ml/src/cuda_pipeline/c51_loss_kernel.cu | 13 ++++++++++--- .../src/cuda_pipeline/experience_kernels.cu | 17 +++++++++++++---- docs/dqn-wire-up-audit.md | 19 +++++++++++++++++++ 4 files changed, 52 insertions(+), 10 deletions(-) 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 b6dab9bfc..f052faa14 100644 --- a/crates/ml/src/cuda_pipeline/c51_aux_dw_kernel.cu +++ b/crates/ml/src/cuda_pipeline/c51_aux_dw_kernel.cu @@ -93,11 +93,18 @@ extern "C" __global__ void c51_aux_dw_kernel( float sp = aux_proj_logdiff_dir[b]; float isw = fminf(is_weights[b], 10.0f); float dz = per_sample_support[b * 12 + 0 * 3 + 2]; - if (dz < 1e-7f) continue; /* degenerate dir support: no gradient (matches forward skip) */ - - float dL_dDelta_online = sp / dz; + /* DEFENSIVE NaN guards (2026-05-13): the dz check `dz < 1e-7` does + * NOT catch NaN (NaN comparison is always false), so NaN dz would + * slip through and produce NaN dW. Bisect cut #1 proved this + * kernel + Adam are the NaN source. Guard ALL upstream-derived + * inputs: sp / is_weights / dz / gamma / done. Any NaN/Inf in + * any input → skip sample (no contribution to dW). */ + if (!isfinite(sp) || !isfinite(isw) || !isfinite(dz) || dz < 1e-7f) continue; float gamma_eff = gamma_buf[b]; float done = dones[b]; + if (!isfinite(gamma_eff) || !isfinite(done)) continue; + + float dL_dDelta_online = sp / dz; float dL_dDelta_target = -gamma_eff * (1.0f - done) * dL_dDelta_online; if (a == a_d) { diff --git a/crates/ml/src/cuda_pipeline/c51_loss_kernel.cu b/crates/ml/src/cuda_pipeline/c51_loss_kernel.cu index 0bb299c85..6e9ee0452 100644 --- a/crates/ml/src/cuda_pipeline/c51_loss_kernel.cu +++ b/crates/ml/src/cuda_pipeline/c51_loss_kernel.cu @@ -1088,7 +1088,11 @@ extern "C" __global__ void c51_loss_batched( * the same amount (linearity of conditional expectation). */ if (d == 0 && aux_shift_active) { - cvar += w_aux[a] * next_state_121; + /* DEFENSIVE NaN guard (2026-05-13): W can carry NaN from + * Adam corruption; 0 * NaN = NaN propagates. */ + float w_a = w_aux[a]; + if (!isfinite(w_a)) w_a = 0.0f; + cvar += w_a * next_state_121; } eq_per_action[a] = cvar; } @@ -1233,8 +1237,11 @@ extern "C" __global__ void c51_loss_batched( * branches and when aux_shift_active == false. */ float effective_reward = reward; if (d == 0 && aux_shift_active) { - float aux_shift_online = w_aux[a0] * state_121; - float aux_shift_target = w_aux[best_next_a] * next_state_121; + /* DEFENSIVE NaN guard (2026-05-13): W can carry NaN. */ + float w_ad = w_aux[a0]; if (!isfinite(w_ad)) w_ad = 0.0f; + float w_star = w_aux[best_next_a]; if (!isfinite(w_star)) w_star = 0.0f; + float aux_shift_online = w_ad * state_121; + float aux_shift_target = w_star * next_state_121; effective_reward = reward + gamma_eff * aux_shift_target * (1.0f - done) - aux_shift_online; diff --git a/crates/ml/src/cuda_pipeline/experience_kernels.cu b/crates/ml/src/cuda_pipeline/experience_kernels.cu index 4244ab6e5..e6a1003af 100644 --- a/crates/ml/src/cuda_pipeline/experience_kernels.cu +++ b/crates/ml/src/cuda_pipeline/experience_kernels.cu @@ -5660,9 +5660,14 @@ extern "C" __global__ void mag_concat_qdir( /* SP22 H6 Phase 3 α atom-shift: per-action shift = W[a]*state_121. * Effective atom positions for this action: z_n + W[a]*state_121. * Since Σ_n prob_n = 1, the expected Q shifts by exactly the - * scalar W[a]*state_121 (linearity of expectation). */ + * scalar W[a]*state_121 (linearity of expectation). + * DEFENSIVE NaN guard (2026-05-13): same pattern as + * compute_expected_q — W can carry NaN from Adam corruption, and + * 0 * NaN = NaN propagates. Verify finite. */ if (aux_shift_active) { - eq += w_aux[a] * state_121; + float w_a = w_aux[a]; + if (!isfinite(w_a)) w_a = 0.0f; + eq += w_a * state_121; } eq_local[a] = eq; } @@ -6311,9 +6316,13 @@ extern "C" __global__ void quantile_q_select( /* SP22 H6 Phase 3 α atom-shift: dir branch (d==0) only. * Per-action shift = W[a]*state_121 uniformly shifts q10/q50/q90 * (they're all positions on the shifted support), so the blend - * shifts by exactly W[a]*state_121. */ + * shifts by exactly W[a]*state_121. + * DEFENSIVE NaN guard (2026-05-13): same pattern as + * compute_expected_q + mag_concat_qdir. */ if (d == 0 && aux_shift_active) { - q_blended += w_aux[a] * state_121; + float w_a = w_aux[a]; + if (!isfinite(w_a)) w_a = 0.0f; + q_blended += w_a * state_121; } q_select[(long long)i * total_actions + action_idx] = q_blended; diff --git a/docs/dqn-wire-up-audit.md b/docs/dqn-wire-up-audit.md index 426fa7a0f..c49f02db1 100644 --- a/docs/dqn-wire-up-audit.md +++ b/docs/dqn-wire-up-audit.md @@ -17428,3 +17428,22 @@ Together these provide belt-and-suspenders against NaN propagation: 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. + +#### Comprehensive W-read guards across all atom-shift kernels (2026-05-13) + +**Verification smoke `train-t5885` @ commit `57cf5c1e0` partial result**: +- `q_var_per_branch [dir=0 mag=NaN ord=0 urg=0]` — dir CLEAN, mag NaN +- `q_by_action [hold=-0 long=-0 short=-0 flat=-0]` — dir CLEAN +- `v_a_means [v=-0.11 a_dir=-0.003 a_mag=NaN a_ord=0.004 a_urg=0.003]` — a_mag NaN + +**Diagnostic**: defensive W-guard in `compute_expected_q` fixed dir branch. But `mag_concat_qdir` (which feeds mag SGEMM input via the dir-Q concat) does NOT have the W guard yet. NaN W there propagates: W=NaN * state_121=0 = NaN → eq_local NaN → mag_concat_buf NaN → mag SGEMM forward NaN → a_mag NaN. + +**This commit**: applies the SAME isfinite guard on `w_aux[a]` reads across ALL atom-shift consumer kernels: +- `mag_concat_qdir` (experience_kernels.cu): guard W in `eq += W[a]*state_121` +- `quantile_q_select` (experience_kernels.cu): guard W in `q_blended += W[a]*state_121` +- `c51_loss_kernel`: guard W in `cvar += W[a]*next_state_121` AND in `effective_reward` substitution (both W[a0] and W[best_next_a]) +- `c51_aux_dw_kernel`: extended guards to cover sp/isw/dz/gamma/done — any NaN/Inf in any input → skip sample (the dz < 1e-7 check did NOT catch NaN since NaN comparisons are always false) + +Together with the previously-added guards on `state_121` and on `total` (dW write), Adam inputs (`g, lr, m, v`), and `W` in `compute_expected_q`, this provides defense-in-depth across the entire atom-shift path. No NaN can propagate through any of the 4 forward kernels OR the dW/Adam backward. + +Cargo check clean.