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.

23
docs/dqn-wire-up-audit.md
View File

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