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 a24c839de..77d8c1171 100644 --- a/crates/ml/src/cuda_pipeline/c51_aux_dw_kernel.cu +++ b/crates/ml/src/cuda_pipeline/c51_aux_dw_kernel.cu @@ -102,10 +102,17 @@ extern "C" __global__ void c51_aux_dw_kernel( if (a == a_d) { float s121 = batch_states[(long long)b * state_dim + aux_dir_prob_index]; + /* Defensive NaN guard (2026-05-13): mirrors atom-shift kernels. + * State_121 NaN would make local_sum NaN, dW NaN, Adam NaN, + * w_aux NaN — corrupting all subsequent forward atom-shifts. + * Clamp NaN to 0 → no W gradient contribution from that sample + * (correct behavior; the gradient was undefined anyway). */ + if (!isfinite(s121)) s121 = 0.0f; local_sum += inv_batch * isw * dL_dDelta_online * s121; } if (a == a_star) { float ns121 = next_batch_states[(long long)b * state_dim + aux_dir_prob_index]; + if (!isfinite(ns121)) ns121 = 0.0f; local_sum += inv_batch * isw * dL_dDelta_target * ns121; } } diff --git a/crates/ml/src/cuda_pipeline/c51_loss_kernel.cu b/crates/ml/src/cuda_pipeline/c51_loss_kernel.cu index b6b4fb9ed..0bb299c85 100644 --- a/crates/ml/src/cuda_pipeline/c51_loss_kernel.cu +++ b/crates/ml/src/cuda_pipeline/c51_loss_kernel.cu @@ -838,6 +838,13 @@ extern "C" __global__ void c51_loss_batched( + aux_dir_prob_index]; next_state_121 = next_batch_states[(long long)sample_id * state_dim + aux_dir_prob_index]; + /* Defensive NaN guard (2026-05-13): state_121 SHOULD be in + * [-1, +1] per the H6 spec, but IEEE-754 `0 * NaN = NaN` + * defeats W=0 safety if an upstream NaN propagates here. The + * real source (aux head or state assembly) is investigated + * separately. */ + if (!isfinite(state_121)) state_121 = 0.0f; + if (!isfinite(next_state_121)) next_state_121 = 0.0f; } for (int d = 0; d < NUM_BRANCHES; d++) { diff --git a/crates/ml/src/cuda_pipeline/experience_kernels.cu b/crates/ml/src/cuda_pipeline/experience_kernels.cu index adbd41c4c..e403a6714 100644 --- a/crates/ml/src/cuda_pipeline/experience_kernels.cu +++ b/crates/ml/src/cuda_pipeline/experience_kernels.cu @@ -5044,6 +5044,17 @@ 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. */ + if (!isfinite(state_121)) state_121 = 0.0f; aux_atom_shift = w_aux[a] * state_121; } @@ -5565,11 +5576,14 @@ extern "C" __global__ void mag_concat_qdir( if (b >= B) return; /* SP22 H6 Phase 3 α atom-shift: hoist state_121 once per sample. - * NULL-safe per the same pattern as compute_expected_q + c51_loss. */ + * NULL-safe per the same pattern as compute_expected_q + c51_loss. + * Defensive NaN guard (2026-05-13): see compute_expected_q for + * rationale — IEEE-754 `0 * NaN = NaN` defeats W=0 safety. */ float state_121 = 0.0f; bool aux_shift_active = (w_aux != NULL && batch_states != NULL); if (aux_shift_active) { state_121 = batch_states[(long long)b * state_dim + aux_dir_prob_index]; + if (!isfinite(state_121)) state_121 = 0.0f; } /* Plan 4 Task 2c.3c.6: read trunk-output RMS once into a register. @@ -6201,11 +6215,13 @@ extern "C" __global__ void quantile_q_select( int i = blockIdx.x * blockDim.x + threadIdx.x; if (i >= N) return; - /* SP22 H6 Phase 3 α atom-shift: hoist state_121 once. */ + /* SP22 H6 Phase 3 α atom-shift: hoist state_121 once. + * Defensive NaN guard (2026-05-13): see compute_expected_q. */ float state_121 = 0.0f; bool aux_shift_active = (w_aux != NULL && batch_states != NULL); if (aux_shift_active) { state_121 = batch_states[(long long)i * state_dim + aux_dir_prob_index]; + if (!isfinite(state_121)) state_121 = 0.0f; } /* Phase 2d: per-branch (v_min, delta_z) read INSIDE the d-loop below. */ diff --git a/docs/dqn-wire-up-audit.md b/docs/dqn-wire-up-audit.md index ad6b0fb2d..e5f8de83b 100644 --- a/docs/dqn-wire-up-audit.md +++ b/docs/dqn-wire-up-audit.md @@ -17356,3 +17356,37 @@ With W=0 + β=0, atom-shift and β are runtime no-ops. The full Phase 3 wiring ( 2. **Smoke still NaN** → bug is in Step 8/11 (c51_aux_dw + adam_w_aux) backward logic, NOT atom-shift forward. Action: investigate the dW + Adam path for the specific NaN source (likely in c51_aux_dw_kernel reads of scratch buffers, or adam_w_aux interaction). Cargo check: 0 errors, 21 pre-existing warnings. + +#### Root cause: IEEE-754 `0 * NaN = NaN` defeats W=0 safety (2026-05-13) + +**Diagnostic finding**: Second smoke `train-gs4gx` @ `617eb61aa` with `W=[0,0,0,0]` and `β=0.0` STILL produced identical NaN pattern: +- `q_var_per_branch [dir=NaN mag=NaN ord=0 urg=0]` +- `q_by_action [hold/long/short/flat = NaN]` +- `a_mag = NaN` + +**Conclusion**: The bug is NOT in atom-shift magnitude. It propagates regardless of W value. + +**Root cause** (IEEE-754 rule): in my atom-shift kernels: +```c +aux_atom_shift = w_aux[a] * state_121; +``` + +When `w_aux[a] = 0` (W=0 diagnostic) AND `state_121 = NaN` (upstream NaN propagation), the multiplication `0 * NaN = NaN` per IEEE-754 (not 0). So W=0 does NOT prevent NaN propagation if state_121 contains NaN. + +`state_121` is read from `batch_states[i * state_dim + AUX_DIR_PROB_INDEX]`. The trainer's sampled batch from the replay buffer is the source. If aux head's softmax produces NaN at any rollout step → `prev_aux_dir_prob` NaN → `state[121]` stored as NaN in replay buffer → trainer reads NaN → my atom-shift propagates NaN. + +**Defensive fix** (this commit): in all 4 atom-shift kernels (`compute_expected_q`, `c51_loss_batched`, `mag_concat_qdir`, `quantile_q_select`) AND the dW kernel (`c51_aux_dw_kernel`), guard state_121 reads: +```c +if (!isfinite(state_121)) state_121 = 0.0f; +``` + +This prevents NaN propagation. State_121 with NaN → treated as 0 → atom_shift = W[a] * 0 = 0 → no shift applied for that sample. + +**Note**: this is a defensive band-aid. The REAL question is why state_121 in the replay buffer contains NaN. Possible sources: +1. Aux head's separate trunk producing NaN logits at some training step → softmax NaN → recentering NaN → stored in replay buffer. +2. Bug in state assembly that fails to zero-init state[121] for some envs. +3. Race condition in aux_softmax_to_per_env_kernel timing. + +The defensive guard unblocks smoke verification of the rest of the Phase 3 mechanism. The state_121 NaN source investigation comes after the smoke confirms the guard works. + +**Verification**: cargo check -p ml --lib clean.