From 95e96aa976c8c4c30733a67cdbb8cebad2b32458 Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Sat, 4 Apr 2026 23:00:16 +0200 Subject: [PATCH] docs: spec for Q-value explosion / NaN investigation Documented all findings from the training step performance session: - Shared memory race was real but not the NaN source (racecheck: 0 hazards) - NaN is non-deterministic numerical instability, NOT a concurrency bug - Suspected math bug in gradient budget allocation during MSE warmup - Investigation plan: GPU NaN detection kernels, gradient audit, git bisect Co-Authored-By: Claude Opus 4.6 (1M context) --- ...6-04-04-q-value-explosion-investigation.md | 115 ++++++++++++++++++ 1 file changed, 115 insertions(+) create mode 100644 docs/superpowers/specs/2026-04-04-q-value-explosion-investigation.md diff --git a/docs/superpowers/specs/2026-04-04-q-value-explosion-investigation.md b/docs/superpowers/specs/2026-04-04-q-value-explosion-investigation.md new file mode 100644 index 000000000..a462085a7 --- /dev/null +++ b/docs/superpowers/specs/2026-04-04-q-value-explosion-investigation.md @@ -0,0 +1,115 @@ +# Q-Value Explosion / NaN Loss — Root Cause Investigation + +**Status**: Open — code bug suspected in loss/gradient math +**Priority**: Critical — blocks reliable smoketest +**Date**: 2026-04-04 + +--- + +## Symptom + +Training produces `NaN/Inf at step N: loss=NaN, grad=0` non-deterministically. +- Step N varies: 2, 22, 62, 92 — no fixed pattern +- Fold 1 usually passes, folds 2-3 more likely to fail +- Same code sometimes passes, sometimes fails (GPU FP non-determinism) +- When loss explodes, it jumps from ~3-4 to ~53+ in one epoch, with grad collapsing from ~0.9 to ~0.02 + +## What we ruled out + +1. **Shared memory race**: compute-sanitizer `--tool racecheck` with `--force-blocking-launches` reports **0 hazards**. The `__syncthreads()` fix in block_reduce (commit 137480c88) addressed a real race but did NOT eliminate the NaN. + +2. **GPU printf deadlock**: Removed printf from C51 kernels (was filling CUDA graph printf buffer → hang). Not the NaN source. + +3. **Buffer size**: Tested 256, 1024, 2048 — NaN persists at all sizes. Not a PER IS-weight variance issue. + +4. **Dual graph forward**: NaN occurs with AND without the MSE-only graph. Not caused by the dual graph optimization. + +5. **NaN-safe scale kernel**: `0 * NaN → 0` guard in dqn_scale_f32_kernel. Helps prevent propagation but doesn't fix the source. + +6. **Huber loss + NaN guards in MSE kernel**: Made it WORSE (0/5 pass → masking the problem caused silent divergence). + +## What we know + +- `loss=NaN, grad=0` — the training guard reads `mse_loss_buf` (during MSE warmup) and it contains NaN. Grad is 0 (or effectively 0). +- The NaN is **non-deterministic** — depends on GPU floating-point reduction ordering. +- The racecheck with `--force-blocking-launches` (fully serialized kernels) STILL produces NaN — this is a numerical bug, not a concurrency bug. +- The **original single mega-graph** (before the 3-graph split into graph_forward + graph_aux + graph_adam) had stable training with <1s epochs. The instability appeared after the graph split refactor. +- The MSE loss kernel computes `0.5 * td^2` where `td = online_E[Q] - target_E[Q]`. Both E[Q] are bounded by support [-24, 24]. The theoretical max MSE is `0.5 * 48^2 = 1152` — large but not NaN. +- The MSE gradient is `isw * td * p_j * (z_j - E[Q])` — proportional to td, which can be up to 48. + +## Key suspicion: math bug introduced during graph split refactor + +Before the 3-graph architecture: +- All ops (forward + loss + grad + aux + adam) were in a single CUDA graph +- The kernel execution order was fixed and captured atomically +- Training was stable with consistent ~<1s epochs + +After the 3-graph split: +- `graph_forward`: forward + loss + backward +- `graph_aux`: HER + clip + EMA + IQL + IQN + CQL +- `graph_adam`: grad_norm + adam + unflatten +- Between these graphs, ungraphed operations run (spectral norm, HER donors, pruning, grad_norm_outside_graph) +- The **gradient accumulation** across graphs may have a math error: + - C51 gradient → `d_value_logits_buf` / `d_adv_logits_buf` (graph_forward) + - gradient budget clip (graph_aux) + - IQN trunk gradient injection (graph_aux) + - CQL gradient SAXPY (graph_aux) + - grad_norm computation (ungraphed, between aux and adam) + - Adam update (graph_adam) + +## Investigation approach + +### Phase 1: Isolate which kernel produces the NaN + +Add a GPU-side NaN detection buffer (`nan_flag_buf[8]`), one flag per kernel stage. After each stage in `submit_forward_ops_mse_only()`, launch a tiny check kernel: +```cuda +extern "C" __global__ void nan_check(const float* buf, int n, int* flag, int flag_idx) { + for (int i = threadIdx.x; i < n; i += blockDim.x) + if (!isfinite(buf[i])) { flag[flag_idx] = 1; return; } +} +``` +Insert between: (1) cuBLAS forward, (2) MSE loss, (3) MSE grad, (4) bf16 cast, (5) cuBLAS backward. +Read back `nan_flag_buf` at epoch boundary to identify exactly which stage produces the first NaN. + +**Important**: These check kernels must NOT be inside the CUDA graph. Run them ungraphed on the first N steps only. + +### Phase 2: Compare single-graph vs 3-graph gradient accumulation + +The original mega-graph had all gradient contributions (C51, IQN, CQL, ensemble) computed and accumulated within a single graph launch. The 3-graph split introduced inter-graph gradient accumulation: +- `graph_forward` writes C51/MSE gradients to `grad_buf` +- `graph_aux` modifies `grad_buf` (clip, IQN injection, CQL SAXPY) +- `compute_grad_norm_outside_graph()` reads `grad_buf` between aux and adam + +Check: is `grad_buf` correctly zeroed before graph_forward? Is the gradient budget clip using the right norm? Are the SAXPY coefficients correct for the blended MSE+C51 case? + +### Phase 3: Verify the gradient budget math with c51_alpha=0 + +During MSE warmup (c51_alpha=0): +- MSE grad writes to `d_value_logits_buf` (main buffers) via the MSE-only graph +- OR MSE grad writes to scratch, then SAXPY blends: `main = 0 * C51 + 1 * MSE` +- The gradient budget clip in graph_aux applies: `c51_budget = max_grad_norm * c51_frac` +- With c51_frac computed from enabled sub-trainers, NOT from c51_alpha + +**Potential bug**: The gradient budget allocates 70% to C51 regardless of c51_alpha. During warmup with alpha=0, the C51 gradient is all zeros (from the NaN-safe scale), so 70% of the budget is "spent" on zeros. The remaining 30% is split between CQL (15%), IQN (10%), Ensemble (5%). The MSE gradient gets NO explicit budget — it's whatever remains after the C51 clip. + +This means the effective MSE gradient is potentially unconstrained or incorrectly constrained. + +### Phase 4: Binary search the refactor + +If phases 1-3 don't find it, use git bisect between the last stable single-graph commit and the current 3-graph architecture to find exactly which commit introduced the instability. + +## Files involved + +| File | Role | +|------|------| +| `crates/ml/src/cuda_pipeline/mse_loss_kernel.cu` | MSE loss computation | +| `crates/ml/src/cuda_pipeline/mse_grad_kernel.cu` | MSE gradient computation | +| `crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs` | Graph capture, gradient accumulation, Adam | +| `crates/ml/src/trainers/dqn/fused_training.rs` | Orchestration: graph_forward → aux → adam | +| `crates/ml/src/trainers/dqn/trainer/training_loop.rs` | Training guard, epoch boundary | + +## Success criteria + +- `test_dqn_trains_on_es_fut` passes 10/10 times deterministically +- No NaN/Inf at any step in any fold +- Epoch times remain <5s on RTX 3050