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) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
@@ -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
|
||||
Reference in New Issue
Block a user