# 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
