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

fix(sp22): defensive NaN guards on Step 8/11 dW + Adam path

Browse Source 
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) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-05-13 11:12:28 +02:00
parent a1945de031
commit 57cf5c1e08
5 changed files with 56 additions and 19 deletions
Download Patch File Download Diff File Expand all files Collapse all files

7
crates/ml/src/cuda_pipeline/adam_w_aux_kernel.cu
View File

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

9
crates/ml/src/cuda_pipeline/c51_aux_dw_kernel.cu
View File

@@ -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;
}
}

18
crates/ml/src/cuda_pipeline/experience_kernels.cu
View File

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

18
crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs
View File

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