docs(dqn): SP1 Phase A — backward NaN audit
Read-only γ audit of every backward-path kernel writing to bw_d_h_s2 / grad_buf / save_h_s2 accumulators. Per-kernel inventory against unsafe-pattern checklist (sqrtf-neg, 1/0, logf-≤0, expf-large, EMA variance, atomicAdd/saxpy NaN-propagation). Each section includes: identified pattern + line citation, severity, proposed guard form, ISV bound option, F0 risk, Phase B flag-slot allocation. Cross-referenced session_2026-04-05_nan_investigation.md residual 8% step-2 NaN in apply_iqn_trunk_gradient — never closed; ranked top suspect for SP1. Output drives SP1 Phase B (12 new flag slots in nan_flags_buf 24→48) and Phase C surgical fix(es). Durable artifact for SP2 framework codification + SP3 structural-fix scoping.
This commit is contained in:
jgrusewski
509
docs/dqn-backward-nan-audit.md
Normal file
509
docs/dqn-backward-nan-audit.md
Normal file
@@ -0,0 +1,509 @@
|
||||
# DQN Backward-Path NaN Audit — SP1 Phase A (γ)
|
||||
|
||||
**Status:** ACTIVE — drives SP1 Phase B instrumentation + Phase C fix decisions.
|
||||
**Spec:** `docs/superpowers/specs/2026-04-29-numerical-stability-investigation-design.md`
|
||||
**Plan:** `docs/superpowers/plans/2026-04-29-numerical-stability-sp1-f1-nan-root-cause.md`
|
||||
**Method:** read-only inventory of every backward-path kernel that writes to `bw_d_h_s2` / `grad_buf` / `save_h_s2` accumulators. Each kernel evaluated against unsafe-pattern checklist. Findings drive Phase B flag-slot allocation and Phase C surgical fix(es).
|
||||
|
||||
## Unsafe-Pattern Checklist
|
||||
|
||||
For each kernel, scan source for:
|
||||
- [ ] `sqrtf(x)` where `x` may be negative (e.g., variance computed via `mean(x²) - mean(x)²` can be slightly negative due to fp roundoff)
|
||||
- [ ] `1/x` or `x / y` where divisor may approach 0 (no ε floor present)
|
||||
- [ ] `logf(x)` where `x` may be ≤ 0 (no `fmaxf(x, ε)` clamp)
|
||||
- [ ] `expf(x)` where `x` may exceed ~88 (single-precision overflow → +inf)
|
||||
- [ ] EMA-tracked variance denominators that may approach 0 across folds
|
||||
- [ ] `atomicAdd` / `saxpy` accumulators receiving NaN inputs (no `isfinite()` guard at producer)
|
||||
- [ ] cuBLAS GEMM with extreme intermediate products (M,N,K combinations producing intermediate overflow)
|
||||
- [ ] `fmaxf` / `fminf` on potentially-NaN inputs (NaN propagates per IEEE-754 → masked-zero in downstream)
|
||||
|
||||
## Cross-references
|
||||
|
||||
- `~/.claude/projects/-home-jgrusewski-Work-foxhunt/memory/session_2026-04-05_nan_investigation.md` — prior NaN root-cause investigation; identifies residual 8% step-2 NaN in `apply_iqn_trunk_gradient`'s cuBLAS backward (NEVER CLOSED — strong candidate).
|
||||
- Smoke `smoke-test-m5gxx` (commit `e9096c7be`) confirmed: `save_h_s2` slot 12 + `grad_buf` slot 6 flag at F1 step ~1140; ALL 5 GRN forward sub-stages (slots 16-20) read clean. NaN enters via backward path, not forward.
|
||||
|
||||
## Architecture-drift / scope notes
|
||||
- Confirmed kernel→file inventory:
|
||||
- `c51_grad_kernel` — `c51_grad_kernel.cu`
|
||||
- `c51_loss_batched` — `c51_loss_kernel.cu`
|
||||
- `iqn_backward_per_sample` + `iqn_weight_grad_reduce` + `iqn_grad_norm_phase{1,2}` + `iqn_adam_kernel` — `iqn_dual_head_kernel.cu`
|
||||
- `cql_logit_grad_kernel` (only) — `cql_grad_kernel.cu` (no separate `cql_logit_grad_kernel.cu`)
|
||||
- `aux_next_bar_backward` + `aux_regime_backward` + `aux_param_grad_reduce` — `aux_heads_kernel.cu`
|
||||
- `ensemble_aggregate_kernel` + `ensemble_kl_gradient_kernel` — `ensemble_kernels.cu` (note: legacy `ensemble_diversity_kernel` was REMOVED 2026-04-21; ensemble path is NOT dormant — `apply_ensemble_diversity_backward` is live, see `gpu_dqn_trainer.rs:7084`)
|
||||
- `mse_grad_kernel` — `mse_grad_kernel.cu` (gradient kernel, separate from forward `mse_loss_batched` in `mse_loss_kernel.cu`)
|
||||
- Bottleneck Linear backward — orchestrated in Rust (`gpu_dqn_trainer.rs:18063` `aux_bottleneck_vsn_backward_dispatch`, `gpu_dqn_trainer.rs:18560`); reuses `launch_dx_only_lda` + `launch_dw_only_no_bias_lda` from `batched_backward.rs`. No standalone `.cu` file.
|
||||
- `relu_mask_kernel` + `bias_grad_reduce_f32_p1`/`_p2` — `backward_kernels.cu`
|
||||
- `dqn_saxpy_f32_kernel` — `dqn_utility_kernels.cu:197`
|
||||
- The plan's Kernel #11 ("`bw_d_h_s2` accumulator path") is a Rust orchestration concern, not a kernel; it is treated below as a saxpy-stack inventory.
|
||||
|
||||
---
|
||||
|
||||
## Per-kernel sections
|
||||
|
||||
### Kernel: `backward_full` (Rust orchestrator, dispatch graph)
|
||||
|
||||
**File:** `crates/ml/src/cuda_pipeline/batched_backward.rs:1665-2050`
|
||||
**Output buffers:** writes branch dW/dB + value-head dW/dB into `grad_buf_base`; accumulates value-head + 4 branch dX into `scratch_d_h_s2 [B, SH2]` (the buffer threaded through `gpu_dqn_trainer::ptrs.bw_d_h_s2`).
|
||||
**Accumulator targets:** `grad_buf_base` (cuBLAS dW with `beta=1.0`); `scratch_d_h_s2` (cuBLAS dX with `beta=0.0` for branch 0 then `beta=1.0` for branches 1..4 + value head); per-branch `scratch_d_h_b` (cuBLAS dX `beta=0.0`).
|
||||
**Caller:** `gpu_dqn_trainer::submit_dqn_step_loop_cublas` and friends — runs in the main forward/backward child graph after `c51_grad_kernel` populates `d_value_logits` + `d_adv_logits[4]`.
|
||||
|
||||
Sequence the orchestrator runs (every training step):
|
||||
1. Branch decoder dW/dB on `branch_streams[0..3]` in parallel — `backward_branch_dw` per branch (KAN coeff/resid grads + GLU value/gate dW/dB + FC out dW/dB).
|
||||
2. Stream-join on `branch_done_events[0..3]`.
|
||||
3. Branch decoder dX (sequential on main stream, `beta=0` then `beta=1`) — `backward_branch_dx` accumulates into `scratch_d_h_s2` and into the per-branch concat-dX buffers (`d_mag_concat_ptr`, `d_ord_concat_ptr`, `d_urg_concat_ptr`).
|
||||
4. Value head dY → dX into `scratch_d_h_v` (fused DRELU_BGRAD path or fallback `relu_mask_kernel` + `launch_bias_grad`).
|
||||
5. Value-FC dW (`launch_dw_only` / `launch_dw_only_no_bias`) writes `goff_w_v1` / `goff_b_v1`.
|
||||
6. Value-FC dX (`launch_dx_only`, `beta=1.0`) accumulates ONTO branch contributions in `scratch_d_h_s2`.
|
||||
7. (Caller, after `backward_full` returns) `encoder_backward_chain` consumes `scratch_d_h_s2` and writes h_s2 + h_s1 GRN dW/dB/dgamma/dbeta + bottleneck dX. The chain is NOT inside this function — it lives on `BatchedForward` (see `gpu_dqn_trainer::6916`).
|
||||
|
||||
**Identified unsafe patterns:**
|
||||
- `backward_full` itself is dispatch-only — no math. Safety properties pass through to its constituent helpers. The two structural points where corruption could *originate* inside this function:
|
||||
- `sgemm_f32` "fc_dW" / "fc_dX" / branch dW/dX with the wide reduction K=batch_size (256). Severity: suspect — extreme intermediate products possible only if `d_value_logits` / `d_adv_logits` carry NaN or ±~1e20 (which would already be caught at slot 6 grad_buf check). Per session_2026-04-05 finding #5, the analogous IQN trunk path showed residual 8% NaN at exactly this step.
|
||||
- The branch-stream dispatch order (4 streams parallel for dW; sequential for dX with `beta=0/1` chain) — semantically sound but means a stale `branch_done_events[d]` skew could cause the dX accumulation to read uninitialised dW. Reviewed: `trunk_done_event` and 4 `branch_done_events[d]` are recorded/awaited symmetrically. Severity: clean.
|
||||
|
||||
**Proposed guard form:** N/A at this level — guards belong in producer kernels (C51 grad, IQN trunk grad, CQL grad, etc.) and at saxpy producers writing into `grad_buf` / `scratch_d_h_s2`.
|
||||
|
||||
**ISV bound option:** N/A.
|
||||
|
||||
**F0 risk:** low — this function does no math; all guards land on its callers/callees.
|
||||
|
||||
**Phase B flag-slot allocation:**
|
||||
- No new slot here. Slot 6 (existing — `grad_buf`) and slot 12 (existing — `save_h_s2`) already cover the function's outputs. The diagnostic *gap* is INSIDE the function: which sub-step within steps 1-7 first writes NaN. That gap is filled by the per-kernel slots below.
|
||||
|
||||
---
|
||||
|
||||
### Kernel: `c51_grad_kernel`
|
||||
|
||||
**File:** `crates/ml/src/cuda_pipeline/c51_grad_kernel.cu` (full file, lines 1-307)
|
||||
**Output buffers:** `d_value_logits [B, NA]` (f32), `d_adv_logits [B, total_branch_atoms × NA]` (f32). Writes are unique-per-thread (no atomicAdd).
|
||||
**Accumulator targets:** none (plain assignment; downstream consumer is `backward_full`'s value-FC GEMM and branch dW/dX GEMMs, which read `d_value_logits` / `d_adv_logits` as `dy`).
|
||||
**Caller:** Loss-graph kernel launched immediately before `backward_full` (`gpu_dqn_trainer::submit_dqn_step_loop_cublas`).
|
||||
|
||||
**Identified unsafe patterns:**
|
||||
- **Line 81 — `expf(lp)` of log-prob `lp` in `d_combined = inv_batch * isw * (expf(lp) - proj)`** — `lp` comes from `current_lp` (saved log-softmax). For numerically-stable log-softmax `lp ≤ 0`, so `expf(lp) ≤ 1`, no overflow. **Severity: clean.**
|
||||
- **Lines 116, 162 — `bias_gap / fmaxf(q_abs_ref, 1e-6f)`** — divisor floored. **Severity: clean (Invariant 1 ε carve-out).**
|
||||
- **Lines 115-116, 161-162 — `isv_signals[...]` reads** — slot reads for Q-mean EMAs / |Q| references. If any ISV slot carries NaN (e.g. `Q_ABS_REF_INDEX = 16` not yet bootstrapped at fold boundary), `fmaxf(NaN, 1e-6f)` returns `1e-6f` (CUDA `fmaxf` is "non-NaN preferred"), but `bias_gap = fmaxf(0, max_mean − means[a1])` propagates NaN if `means[a1]` is NaN, then `frac_bin = NaN / 1e-6 = NaN`, then `collapse_frac = fminf(1, NaN) = NaN`, then `d_combined *= bin_weight = NaN`. **Severity: suspect.** Cross-fold ISV reset audit needed.
|
||||
- **Line 125 — `bin_weight = 1 + collapse_frac * mag_bias_signal`** — bounded ∈ [1, 2] when `collapse_frac` is finite; NaN-passthrough as above.
|
||||
- **Line 188-191 — `compression_boost = 1 + compression`** — same pattern with `q_abs_ref_mag` → `q_dir_abs_ref` ratio. NaN-passthrough.
|
||||
- **Line 274 — `a_std = sqrtf(sq_sum / n_d + 1e-12f)`** — argument is `mean(diff²) + 1e-12 ≥ 1e-12 > 0`. **Severity: clean (Invariant 1 ε carve-out).**
|
||||
- **Line 275 — `inv_a_std = 1 / (a_std + 1e-6f)`** — divisor floored. **Severity: clean.**
|
||||
- **Line 81 — `1 / batch_size` — fixed integer**. **Severity: clean.**
|
||||
- **Line 86 — `lp_clamped = fmaxf(lp, -10.0f)`** — explicit floor on log-prob entropy term. **Severity: clean.**
|
||||
- **No atomicAdd anywhere** (the kernel is explicitly DETERMINISTIC). One thread per (b, j); writes are single per slot.
|
||||
- **Line 226-228 — `(z_val - z_min) / fmaxf(z_max - z_min, 1e-6f)`** — divisor floored. **Severity: clean.**
|
||||
|
||||
Net: the kernel is well-guarded for direct numerical hazards. The one residual risk surface is **NaN propagation from ISV slots that have not been re-initialised at fold boundaries**, since this kernel reads slots 12-21 every step.
|
||||
|
||||
**Proposed guard form** (only if ISV-NaN regression confirmed): wrap each ISV read with an `isfinite` guard, e.g. `q_abs_ref = isfinite(isv_signals[16]) ? isv_signals[16] : 0.0f`. With `q_abs_ref = 0`, `frac_bin` short-circuits to 0 and the bin-weight reduces to identity. **Location:** lines 105-108, 150-154, 183-188.
|
||||
|
||||
**ISV bound option:**
|
||||
- Option 1: existing slots `Q_ABS_REF_INDEX = 16` + `Q_DIR_ABS_REF_INDEX = 21` already provide the |Q|-scale references. The proposed guard does not change the bound — it preserves identity behaviour when the slot is uninitialised.
|
||||
- Option 3 (Invariant 1 carve-out): the `isfinite` short-circuit is a numerical-stability check, not a dynamic bound. No new slot needed.
|
||||
|
||||
**F0 risk:** low — F0 trains cleanly with the existing kernel; isfinite-passthrough preserves F0 behaviour bit-identically when the slot already holds finite values.
|
||||
|
||||
**Phase B flag-slot allocation:**
|
||||
- Slot 24: `d_value_logits` (size `B × NA`) — first NaN sink in the gradient chain after C51 forward; high-priority instrumentation. Catches both the kernel itself and any upstream ISV-NaN propagation.
|
||||
- Slot 25: `d_adv_logits[*]` (size `B × total_branch_atoms × NA`) — companion to slot 24; will fire if and only if a per-branch path corrupted independently (e.g. magnitude-only `inv_a_std` divide). Both slots are needed because they are independent producers.
|
||||
|
||||
---
|
||||
|
||||
### Kernel: `c51_loss_batched`
|
||||
|
||||
**File:** `crates/ml/src/cuda_pipeline/c51_loss_kernel.cu:520-1141` (main kernel; supporting helpers `block_log_softmax_f` / `block_bellman_project_f` / `block_expected_q_f` at lines 87-283; `c51_loss_reduce` at 1152-1162; `barrier_gradient_direction` at 1191-1310; `ib_gradient_direction` at 1342-1447; `c51_mixup_ce` at 1463-1548)
|
||||
**Output buffers:** `per_sample_loss [B]` (f32), `td_errors [B]`, `total_loss [1]`, `save_current_lp [B, 4, NA]`, `save_projected [B, 4, NA]`, `q_divergence [1]`. The `barrier_gradient_direction` and `ib_gradient_direction` kernels write **directly into `d_adv_logits` + `d_v_logits`** (the same buffers `c51_grad_kernel` produces) using plain `+=` (no atomicAdd; one thread per sample).
|
||||
**Accumulator targets:** `d_adv_logits` / `d_v_logits` via `+=` (in `barrier_gradient_direction` + `ib_gradient_direction`); `total_loss` reduced by separate single-block `c51_loss_reduce`.
|
||||
**Caller:** `gpu_dqn_trainer::submit_dqn_step_loop_cublas` (loss kernel runs before `c51_grad_kernel`; `barrier_gradient_direction` + `ib_gradient_direction` may run after, depending on barrier_weight / health).
|
||||
|
||||
**Identified unsafe patterns:**
|
||||
|
||||
In the **forward** loss kernel (`c51_loss_batched`):
|
||||
- **Line 274 / 779 / 804 — `a_std = sqrtf(sq_sum / n_d + 1e-12f)`** with `+ 1e-12f` floor. **Severity: clean.**
|
||||
- **Line 720 — `n_steps_f = __logf(fmaxf(gamma_buf_base, 1e-9f)) / log_dir`**, then **line 721 — `gamma_eff = __powf(fmaxf(gamma_b_raw, 1e-6f), n_steps_f)`** — both bases floored. `__logf` of a small positive value is large-negative finite. `__powf(small, n_steps_f)` is bounded by `n_steps_f ≤ ~10` so `__powf(1e-6, 10) = 1e-60` — finite-but-tiny, **harmless**. **Severity: clean.**
|
||||
- **Line 245 — `t_z = -10.0f * (1.0f - expf(t_z / 10.0f))`** — Huber compression on negative tail. `expf(t_z/10)` with `t_z ∈ [v_min, v_max]` and v_min ~ -50 gives `expf(-5) ≈ 6.7e-3` — finite. **Severity: clean.**
|
||||
- **Line 222-224 — `lead_scale = fmaxf(fmaxf(q_dir_abs_ref, q_mag_abs_ref), 0.1 * v_range)`** — passes NaN through if either q_*_ref is NaN (CUDA `fmaxf(NaN, x)` returns x; `fmaxf(x, NaN)` returns x — non-NaN-preferred — so NaN here is masked). **Severity: clean** (the NaN-mask actually helps).
|
||||
- **Line 247 — `t_z = fminf(fmaxf(t_z, v_min), v_max)`** — IEEE non-NaN-preferred; clamp if NaN propagates.
|
||||
- **Line 249-253 — `b_pos = (t_z - v_min) / delta_z`** — if `delta_z ≤ 1e-7f` the per-branch path is skipped earlier (line 737, 751 `branch_degenerate continue`). **Severity: clean.**
|
||||
- **Line 894 — `alpha = fminf(fmaxf(cvar_alpha_buf[sample_id], 0.05f), 0.95f)`** — clamped. **Severity: clean.**
|
||||
- **Line 905 — `cvar = (cvar_weight > 1e-8f) ? (cvar_sum / cvar_weight) : eq`** — guarded divide. **Severity: clean.**
|
||||
- **Line 941 — `tau_floor = fmaxf(fabsf(mean_q) * 0.01f, 1e-6f)`** then **line 942 — `tau_base = fmaxf(q_gap_local, tau_floor)`** — finite. **Severity: clean.**
|
||||
- **Line 952 — `expf((eq_per_action[a] - max_eq) / tau)`** — argument ≤ 0 (max-shift), so `expf ≤ 1`. **Severity: clean.**
|
||||
- **Line 1019 — `expf(shmem_lp[j])`** of log-softmax, ≤ 1. **Severity: clean.**
|
||||
- **Line 1114 — `logf(fmaxf(shmem_lp[j], 1e-10f))`** — argument floored. **Severity: clean (Invariant 1 ε carve-out).**
|
||||
- **Lines 1131, 1538 — `avg_ce *= (1 + dd*dd*asymmetric_dd_weight)`** — all factors finite. **Severity: clean.**
|
||||
|
||||
In **`c51_loss_reduce`** (lines 1152-1162):
|
||||
- **Line 1161 — `total_loss[0] = sum / batch_size`** — single-block sequential sum, finite when `per_sample_loss` is finite.
|
||||
|
||||
In **`barrier_gradient_direction`** (lines 1191-1310):
|
||||
- **Line 1233, 1237 — `expf(v_row[z] + adv_a[z] - max_l)`** with max-shift — bounded ≤ 1. **Severity: clean.**
|
||||
- **Line 1237 — `p = expf(...) / (sum + 1e-8f)`** — guarded. **Severity: clean.**
|
||||
- Plain `+=` writes uniquely per (i, a, z) and (i, z) — confirmed deterministic. **Severity: clean.**
|
||||
|
||||
In **`ib_gradient_direction`** (lines 1342-1447): mirror structure to `barrier_gradient_direction`; same guards; **clean.**
|
||||
|
||||
In **`c51_mixup_ce`** (lines 1463-1548):
|
||||
- **Line 1515 — `v1 = powf(u1, inv_a)`, `v2 = powf(u2, inv_a)`** with `u1,u2 ∈ [1e-7, 1)` and `inv_a = 1/mixup_alpha`. If `mixup_alpha → 0`, `inv_a → ∞` → `powf(1e-7, ∞) = 0` (slow); `powf(small, large) = 0`; **finite.** **Severity: clean.**
|
||||
- **Line 1516 — `lambda = v1 / (v1 + v2 + 1e-7f)`** — guarded. **Severity: clean.**
|
||||
|
||||
Net: this kernel is the most heavily-guarded numerical surface in the pipeline. The only failure mode would be **save_current_lp / save_projected being read with NaN inputs**, which would propagate to `total_ce` and `td_errors`. That input-NaN pathway is captured at slot 7 (`save_current_lp`) + slot 8 (`save_projected`) by the existing post-forward checks.
|
||||
|
||||
**Proposed guard form:** none required. Existing instrumentation covers the input pathways.
|
||||
|
||||
**ISV bound option:** N/A.
|
||||
|
||||
**F0 risk:** N/A (no fix proposed).
|
||||
|
||||
**Phase B flag-slot allocation:**
|
||||
- No new slot. Existing slots 7 (`save_current_lp`) + 8 (`save_projected`) cover the loss-kernel outputs that downstream consumers depend on. The per-sample-loss / total-loss intermediate is captured at slot 3 (`mse_loss_buf`)... NB: slot 3 currently watches MSE; there is no equivalent slot for `c51_loss` total. Recommend **slot reservation** (Phase B optional) for `total_loss` if subsequent investigation needs it, but NOT in the SP1 12-slot budget.
|
||||
|
||||
---
|
||||
|
||||
### Kernel: `apply_iqn_trunk_gradient` (HIGH PRIORITY — prior-investigation top suspect)
|
||||
|
||||
**File:** `crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs:6839-7067` (Rust orchestrator, not a `.cu` file)
|
||||
**Output buffers:** writes into `iqn_trunk_m` (scratch buffer, `param_sizes[0..13)` total f32 elements — same padded layout as `grad_buf` for trunk tensors); then SAXPYs into `grad_buf [trunk]` and `grad_buf [VSN range, indices 95..119)`. Also writes `bw_d_h_s2 [B, SH2]` (DtoD copy from `iqn_d_h_s2_ptr`).
|
||||
**Accumulator targets:** `grad_buf` via `dqn_saxpy_f32_kernel` with `scale = config.iqn_lambda × iqn_budget`; `bw_d_h_s2` via `graph_safe_copy_f32` (DtoD overwrite).
|
||||
**Caller:** `gpu_dqn_trainer::run_aux_pass` / `apply_aux_gradients` — fires after the main `c51_grad_kernel` + `backward_full` chain has populated `grad_buf` for the main loss path. `iqn_d_h_s2_ptr` comes from `iqn_backward_per_sample`'s `d_h_s2_out`.
|
||||
|
||||
Sequence (verbatim from code):
|
||||
1. `memset_d8_async` zeros `iqn_trunk_m` (lines 6864-6873).
|
||||
2. `graph_safe_copy_f32` copies `iqn_d_h_s2_ptr` → `bw_d_h_s2` (lines 6875-6881).
|
||||
3. `cublas_forward.encoder_backward_chain(..., d_h_s2 = bw_d_h_s2, ...)` — runs the FULL GRN h_s2/h_s1 backward + Linear_residual + bottleneck dX; writes dW/dB/dgamma/dbeta into `iqn_trunk_m` (NOT into `grad_buf` directly).
|
||||
4. (Optional, `bottleneck_dim > 0`) `aux_bottleneck_vsn_backward_dispatch` writes 24 VSN dW/dB into `vsn_dw_iqn_aux_scratch`.
|
||||
5. `dqn_saxpy_f32_kernel` adds `scale * iqn_trunk_m` into `grad_buf [0..trunk_grad_total)` (lines 7013-7033).
|
||||
6. (Optional) saxpy `vsn_dw_iqn_aux_scratch` into `grad_buf [VSN range]`.
|
||||
|
||||
**Identified unsafe patterns:**
|
||||
- **Step 3 — `encoder_backward_chain` cuBLAS GEMMs**: this is the path session_2026-04-05 finding identifies as the top residual-NaN suspect. The chain runs trunk h_s2 GRN backward → h_s1 GRN backward → optional bottleneck Linear backward, all via cuBLAS SGEMM with extreme dimensions (`SH2=128`, `B=batch`, `state_dim=128`). Per session_2026-04-05: "needs GPU-side printf in backward_fc_layer dW GEMM or compute-sanitizer numerical mode to trace exact values."
|
||||
- Severity: **suspect — confirmed open from prior investigation.**
|
||||
- **Step 5 — `dqn_saxpy_f32_kernel`** (`dqn_utility_kernels.cu:197`): plain `y[i] = y[i] + alpha * x[i]`, **no `isfinite()` guard**. If any `iqn_trunk_m[i]` element is NaN, `grad_buf[i]` becomes NaN. Subsequently consumed by `dqn_adam_kernel`, which DOES have an `isfinite(g)` guard (so would zero the gradient and skip the update), but `grad_buf[i]` itself remains NaN until the Adam pass. The slot-6 NaN-flag at F1 step ~1140 is consistent with this: NaN visible in `grad_buf` post-saxpy, before Adam sanitises.
|
||||
- Severity: **suspect — open.**
|
||||
- **Step 2 (DtoD copy)** — pure copy, no math; if the copy source `iqn_d_h_s2_ptr` (= `IqnDualHead::d_h_s2_out`) is finite, the destination is finite. The **input-NaN pathway** into this function is exactly what the slot-12 `save_h_s2` flag would NOT catch (because save_h_s2 is the FORWARD activation; `bw_d_h_s2` is the BACKWARD gradient buffer of the same name). The plan's slot 12 *is* labelled `save_h_s2` (forward) — so a separate slot is needed to instrument **`bw_d_h_s2` post-IQN-DtoD**.
|
||||
- **Step 4 VSN backward** — `aux_bottleneck_vsn_backward_dispatch` runs `bn_tanh_backward` + `vsn_d_gated_state_portfolio_kernel` + `vsn_logit_gather_kernel` + 3 cuBLAS dW/dX + a `scale_f32_aux` damping kernel. The damping kernel reads from a scalar (VSN_DW_DAMP); the per-group MLP backward kernels were audited at task #154/#157 closure. Severity: **suspect** for the cuBLAS legs (same class as Step 3); **clean** for the kernel-resident math.
|
||||
|
||||
**Proposed guard form:**
|
||||
- (a) **Add `isfinite()` guard to `dqn_saxpy_f32_kernel`** — modify kernel to read `float xv = x[i]; if (isfinite(xv)) y[i] = y[i] + alpha * xv;`. This idempotently masks NaN-input from any saxpy producer (IQN trunk, IQN VSN, ensemble trunk, ensemble VSN, distillation, MoE — all consumers of this saxpy primitive). Mirrors the IQN/DQN/Attention/Curiosity Adam isfinite pattern (existing defense-in-depth, see session_2026-04-05).
|
||||
- (b) **Add NaN-check on `iqn_trunk_m` immediately after `encoder_backward_chain` returns** (before the saxpy). This is the diagnostic instrumentation step (slot 26 below); it locates whether NaN originates inside the GRN backward chain or downstream. **Cite location:** insert `check_nan_f32(self.ptrs.iqn_trunk_m, trunk_grad_total, 26)` between line 6935 and line 7013.
|
||||
- (c) **Add NaN-check on `iqn_d_h_s2_ptr` immediately after step 2's DtoD copy** (slot 27 below) — locates whether the IQN per-sample backward kernel produced the NaN seed.
|
||||
|
||||
**ISV bound option:**
|
||||
- Guard (a) is a numerical-stability isfinite guard (Invariant 1 carve-out).
|
||||
- Guards (b)/(c) are diagnostic-only — no bound needed.
|
||||
|
||||
**F0 risk:** **medium** for guard (a). Reasoning: `dqn_saxpy_f32_kernel` is on the hot path (every IQN trunk SAXPY, ensemble SAXPY, distillation SAXPY, MoE SAXPY — ~8 call sites). Adding a per-element `isfinite` check adds one branch per element. F0 currently trains cleanly, so the guard would be a no-op on F0-typical inputs (all finite). Risk vector: if any of those 8 callers depends on NaN propagating into Adam (which currently zeros the gradient), masking earlier could break their semantics. **Verification before deployment:** confirm Adam zero behaves identically to saxpy-skip behaviour in steady state. Likely identical — Adam reads `grads[tid]` and on `!isfinite(g)` writes `grads[tid] = 0` and returns; the saxpy-skip variant just leaves `grad_buf[i]` unchanged from its previous value (which Adam's prior-zero-out makes 0). **Decision:** F0 risk LOW after this verification, MEDIUM until verified.
|
||||
|
||||
**Phase B flag-slot allocation:**
|
||||
- Slot 26: `iqn_trunk_m` post-`encoder_backward_chain` (size `trunk_grad_total ≈ 50k floats`) — locates whether GRN backward chain is the NaN source.
|
||||
- Slot 27: `iqn_d_h_s2_ptr` post-DtoD into `bw_d_h_s2` (size `B × SH2`) — locates whether IQN per-sample backward seeded NaN. NB: instrument the IQN-source side (`iqn_d_h_s2_ptr`), not the destination (`bw_d_h_s2`), because the distinction is the diagnostic value.
|
||||
- (Optional, slot 28): `vsn_dw_iqn_aux_scratch` post-`aux_bottleneck_vsn_backward_dispatch` — only meaningful if slot 26 fires AND `bottleneck_dim > 0`. Defer until slot 26 fires.
|
||||
|
||||
---
|
||||
|
||||
### Kernel: `iqn_backward_per_sample` + `iqn_weight_grad_reduce` + `iqn_grad_norm_phase{1,2}` + `iqn_adam_kernel`
|
||||
|
||||
**File:** `crates/ml/src/cuda_pipeline/iqn_dual_head_kernel.cu:504-621` (per-sample), `:638-769` (weight grad reduce), `:802-844` (grad norm 2-phase), `:847-899` (adam).
|
||||
**Output buffers:**
|
||||
- `iqn_backward_per_sample` → `save_dL_dq [B, N]`, `d_h_s2_out [B, hidden_dim]` (this is the source for `apply_iqn_trunk_gradient`'s step 2).
|
||||
- `iqn_weight_grad_reduce` → `grad_buf [total_iqn_params]`.
|
||||
- `iqn_grad_norm_phase1/2` → `block_sums [num_blocks]` then `norm_out [1]`.
|
||||
- `iqn_adam_kernel` → updates `params`, `m`, `v`, zeros `grads`.
|
||||
**Accumulator targets:** none of the per-sample / weight-grad / grad-norm kernels use atomicAdd (per file header: "Zero atomicAdd"). Adam zeros `grads[tid]` after consuming.
|
||||
**Caller:** `gpu_iqn_head` orchestration; sequenced as `iqn_forward_loss_kernel` → `iqn_backward_per_sample` → `iqn_weight_grad_reduce` → `iqn_grad_norm_phase1/2` → `iqn_adam_kernel`.
|
||||
|
||||
**Identified unsafe patterns:**
|
||||
|
||||
In **`iqn_backward_per_sample`**:
|
||||
- **Lines 612 — `if (isfinite(d_trunk)) d_h_s2_acc[...] += d_trunk`** — inline isfinite guard already present. **Severity: clean (existing defense-in-depth from session_2026-04-05 finding).**
|
||||
- **`quantile_huber_grad`** (lines 146-155) — `copysignf(kappa, delta)` and finite arithmetic — **clean**.
|
||||
- **Line 591 — `dL_dq = iqn_block_sum(...) * inv_n_sq`** — finite reduce of finite inputs. **Clean.**
|
||||
|
||||
In **`iqn_weight_grad_reduce`**:
|
||||
- **Lines 689, 712 — `if (isfinite(dL_dpre)) acc += dL_dpre * cos_features[...]`** — inline isfinite guards. **Severity: clean.**
|
||||
- **Lines 753-764 — branch-weight outer-product accumulation** — no isfinite guard inside loop. If `save_dL_dq[b * N + ti]` carries NaN AND `save_combined[(b*N+ti)*hidden_dim + h]` is finite, `acc` becomes NaN, then `grad_buf[idx] = acc * inv_batch` is NaN. **Severity: suspect.**
|
||||
- Note: `save_dL_dq` is written by `iqn_backward_per_sample` lines 591-595 with no isfinite guard on `dL_dq` itself. If `online_q[ti]` or `target_q[ti]` carry NaN at the backward call, `quantile_huber_grad` produces NaN; the per-sample kernel happily writes NaN to `save_dL_dq`. The reduce then propagates.
|
||||
- Severity ELEVATED: **suspect — open path through `save_dL_dq`.**
|
||||
|
||||
In **`iqn_grad_norm_phase1`**:
|
||||
- **Lines 811-813 — `partial += g * g`** — no isfinite guard. NaN squared is NaN; reduces to NaN. Then `iqn_adam_kernel` reads `norm_sq[0] = NaN`, `sqrtf(fmaxf(NaN, 0)) = NaN` (CUDA fmaxf NaN-passthrough on the second arg), `clip_scale = NaN / NaN = NaN`, `g *= clip_scale = NaN` — but Adam's line 873 `if (!isfinite(g))` catches this and zeros the gradient + returns. **Severity: clean** (Adam guard catches downstream).
|
||||
|
||||
In **`iqn_adam_kernel`**:
|
||||
- **Line 873 — `if (!isfinite(g)) { grads[tid] = 0.0f; return; }`** — global isfinite guard. **Severity: clean.**
|
||||
- **Line 876 — `sqrtf(fmaxf(norm_sq[0], 0.0f))`** — argument floored at 0. **Severity: clean.**
|
||||
- **Line 878 — `clip_scale = (grad_norm_val > max_grad_norm && grad_norm_val > 0) ? max_grad_norm / grad_norm_val : 1`** — compound guard. **Severity: clean.**
|
||||
- **Lines 889-890 — `m_hat = mi / (1 - powf(beta1, adam_t))`** — denominator approaches 0 only as `beta1^adam_t → 1`, which requires `adam_t = 0`. If the kernel runs with `adam_t = 0`, denominator = 0 → division by zero → ±inf, propagates. **Severity: suspect (but only at literal step 0).** Adam's `adam_t_buf` is incremented before the kernel reads it (per existing infrastructure); confirm via launcher.
|
||||
- **Line 894 — `sqrtf(v_hat) + eps`** — `v_hat ≥ 0`, eps is positive. **Severity: clean.**
|
||||
|
||||
**Proposed guard form:**
|
||||
- For the **`save_dL_dq` propagation pathway**: add `if (isfinite(dq))` guard to lines 762 and 749-751 of `iqn_weight_grad_reduce`. Mirrors the existing isfinite pattern at lines 689/712.
|
||||
- Location: lines 750-762 (branch-weight inner loop and bias inner loop).
|
||||
|
||||
**ISV bound option:** Invariant 1 ε-carve-out (isfinite check is numerical-stability, not a dynamic bound).
|
||||
|
||||
**F0 risk:** low — F0-typical `dL_dq` values are finite; the guard is identity for those. Pattern is identical to existing isfinite guards at lines 689/712 of the same file.
|
||||
|
||||
**Phase B flag-slot allocation:**
|
||||
- Slot 28: `save_dL_dq [B × IQN_NUM_QUANTILES]` (B × 5 floats, very small) — exact intermediate that connects per-sample backward → weight-grad reduce. Slot 26 (iqn_trunk_m) and slot 27 (iqn_d_h_s2_ptr) cover the trunk side; slot 28 covers the weight-grad-reduce side.
|
||||
|
||||
---
|
||||
|
||||
### Kernel: `cql_logit_grad_kernel`
|
||||
|
||||
**File:** `crates/ml/src/cuda_pipeline/cql_grad_kernel.cu` (full file, 1-212)
|
||||
**Output buffers:** `d_v_logits [N, num_atoms]`, `d_adv_logits [N, total_actions × num_atoms]`. Plain assignment per slot (no atomicAdd) — one thread per sample.
|
||||
**Accumulator targets:** `d_adv_logits` (sub-slice for branch d via `branch_base` arithmetic); `d_v_logits` via per-thread `d_val_accum[]` register array then assignment.
|
||||
**Caller:** CQL gradient launch from `gpu_dqn_trainer::apply_cql_loss_backward` (separate aux pass). Same destination buffers as `c51_grad_kernel` and `barrier_gradient_direction` — they all SHARE `d_v_logits` / `d_adv_logits`. In the production code, `c51_grad_kernel` writes first (overwrite); `cql_logit_grad_kernel` writes its own slice (overwrite, line 200 `d_adv[j] = d_combined`). **Note:** this OVERWRITES the C51 contribution at the (i, branch d, action a, atom j) slot; the consumer downstream (CQL aux-pass-only `backward_full`) reads the CQL-overwritten values. This is correct semantically — CQL has its own backward path — but should be flagged for NaN source instrumentation.
|
||||
|
||||
**Identified unsafe patterns:**
|
||||
- **Line 41 — `dz = (num_atoms > 1) ? (v_max - v_min) / (num_atoms - 1) : 0.0f`** — guarded.
|
||||
- **Line 100 — `max_logit = fmaxf(max_logit, c)`** — NaN-passthrough on second arg only (CUDA non-NaN-preferred). If `c` is NaN once, `max_logit` stays at last-finite max. NaN can still leak if *all* `c` values are NaN. **Severity: suspect (low).**
|
||||
- **Line 111 — `expf(val[j] + adv[j] - a_mean_j - max_logit)`** — argument ≤ 0 after max-shift. **Severity: clean.**
|
||||
- **Line 113 — `log_sum = logf(sum_exp + 1e-8f) + max_logit`** — `sum_exp ≥ 0`, ε floor. **Severity: clean.**
|
||||
- **Line 142 — `softmax_q = expf(q_values_branch[a] - max_q) / (sum_exp_q + 1e-8f)`** — guarded. **Severity: clean.**
|
||||
- **Line 198 — `d_combined = inv_batch * d_cql_dq[a] * p * (z - eq)`** — finite arithmetic. **Severity: clean.**
|
||||
- **Line 71 — `float d_val_accum[256]`** — register-allocated array; `num_atoms = 51` for default config so safely under bound; `d_val_accum` index loop has `j < 256` guard.
|
||||
- **No atomicAdd — clean determinism.**
|
||||
|
||||
Net: `cql_logit_grad_kernel` is well-guarded. The only suspect path is **input-NaN propagation from `v_logits` / `adv_logits`** (which are `on_v_logits_buf` / `on_b_logits_buf` — already covered at slots 1 + 2 by `run_nan_checks_post_forward`).
|
||||
|
||||
**Proposed guard form:** none required.
|
||||
|
||||
**ISV bound option:** N/A.
|
||||
|
||||
**F0 risk:** N/A.
|
||||
|
||||
**Phase B flag-slot allocation:**
|
||||
- Slot 29 (deferred; assign only if SP1 Phase B has spare slots after the higher-priority IQN/saxpy/aux instrumentation): `d_v_logits` post-CQL (1 of 2 OVERWRITE producers) — useful only as a tiebreaker if slot 24 (`d_value_logits` post-C51) is clean but the post-CQL aux-backward goes NaN.
|
||||
- Decision: **DEFER**. Slot 24 already covers the C51 producer; if CQL is the source, slot 6 (`grad_buf`) post-CQL-aux-backward will fire alone (without slot 24/25). That delta is sufficient diagnostic separation.
|
||||
|
||||
---
|
||||
|
||||
### Kernel: `aux_next_bar_backward` + `aux_regime_backward` + `aux_param_grad_reduce`
|
||||
|
||||
**File:** `crates/ml/src/cuda_pipeline/aux_heads_kernel.cu:391-482` (next-bar backward), `:502-611` (regime backward), `:624-652` (param-grad reduce).
|
||||
**Output buffers:**
|
||||
- `aux_next_bar_backward` → `dW1_partial [B, H, SH2]`, `db1_partial [B, H]`, `dW2_partial [B, H]`, `db2_partial [B]`, `dh_s2_out [B, SH2]`.
|
||||
- `aux_regime_backward` → similar partials with K-shape on the second linear; `dh_s2_out [B, SH2]`.
|
||||
- `aux_param_grad_reduce` → `final_out [P]` — the reduced param-grad slice.
|
||||
**Accumulator targets:** the `dh_s2_out` outputs are SAXPY'd into `bw_d_h_s2` by the caller (per the kernel docstring: "no atomicAdd, just a SAXPY in the caller"). The reduced `final_out` is then SAXPY'd into `grad_buf` at the aux head's slot range. Both saxpys go through `dqn_saxpy_f32_kernel` (same kernel as the IQN trunk path).
|
||||
**Caller:** `gpu_dqn_trainer::apply_aux_heads_backward` (search not done in this audit; orchestrator entry point inferred from kernel docstring).
|
||||
|
||||
**Identified unsafe patterns:**
|
||||
|
||||
In **`aux_next_bar_backward`**:
|
||||
- **Line 428 — `inv_scale = 1 / fmaxf(label_scale, 1e-6f)`** — divisor floored. **Severity: clean.**
|
||||
- **Line 430 — `d_pred_b = (2/B) * (pred[b] - label[b] * inv_scale)`** — finite if `pred` is finite and `label_scale` is finite. **Severity: clean** assuming label_scale is bootstrapped (per ISV slot 117 `AUX_LABEL_SCALE_EMA_INDEX`, EMA-tracked).
|
||||
- **Line 437 — `aux_elu_bwd_from_post(h_post)`** — branchless on sign; finite. **Severity: clean.**
|
||||
- **No atomicAdd. Per-block partials are unique-per-(b, *).**
|
||||
|
||||
In **`aux_regime_backward`**:
|
||||
- **Line 537 — `lmax = row[k]` then `sum_e += expf(row[k] - lmax)`** — max-shift CE; finite. **Severity: clean.**
|
||||
- **Line 540 — `inv_sum = 1 / sum_e`** — `sum_e ≥ K * exp(-something)` after max-shift; bounded below by `1.0` (since the max-element contributes `expf(0) = 1`). **Severity: clean.**
|
||||
- **Line 542 — `inv_B = 1 / B`** — fixed divisor. **Severity: clean.**
|
||||
|
||||
In **`aux_param_grad_reduce`**: pure shared-memory reduce; **clean.**
|
||||
|
||||
Net: aux backward kernels are well-guarded. **One residual concern:** when `label_scale` (ISV slot 117) is uninitialised at fold boundary, `inv_scale = 1 / 1e-6 = 1e6`, and `d_pred_b = (2/B) * (pred - label * 1e6)` could produce values up to `~1e9` per sample — large but **finite, no NaN**. Fold-boundary reset of slot 117 is therefore a CORRECTNESS concern but not a NaN concern.
|
||||
|
||||
**Proposed guard form:** none required at the kernel level.
|
||||
|
||||
**ISV bound option:** N/A.
|
||||
|
||||
**F0 risk:** N/A.
|
||||
|
||||
**Phase B flag-slot allocation:**
|
||||
- Slot 30: `dh_s2_out` (size `B × SH2`) — instrument only the next-bar variant; regime-CE has the same shape and same caller, so slot-coverage is identical. This catches whether the aux saxpy seeded the `bw_d_h_s2` NaN at F1 step ~1140. Suspect-of-record: low (label_scale is finite by construction); but cheap to check, and the buffer is the upstream source for one of the 8 saxpy producers into `bw_d_h_s2`.
|
||||
|
||||
---
|
||||
|
||||
### Kernel: `apply_ensemble_diversity_backward` + `ensemble_kl_gradient_kernel`
|
||||
|
||||
**File:** `crates/ml/src/cuda_pipeline/ensemble_kernels.cu:77-119` (KL gradient kernel), `gpu_dqn_trainer.rs:7084` (orchestrator).
|
||||
**Output buffers:**
|
||||
- `ensemble_kl_gradient_kernel` → `d_logits_0 [B × NA]`.
|
||||
- `apply_ensemble_diversity_backward` (Rust) → `iqn_trunk_m` (REUSED scratch), then SAXPY into `grad_buf`.
|
||||
**Accumulator targets:** SAXPY into `grad_buf` via `dqn_saxpy_f32_kernel`. SHARED scratch buffer `iqn_trunk_m` with `apply_iqn_trunk_gradient` — caller must zero between uses (orchestrator does, line 7104-7113).
|
||||
**Caller:** `gpu_dqn_trainer::apply_ensemble_diversity_backward` after main backward chain.
|
||||
|
||||
**Identified unsafe patterns:**
|
||||
|
||||
In **`ensemble_kl_gradient_kernel`**:
|
||||
- **Lines 92-94 — `max0 = fmaxf(max0, v)`** — NaN-passthrough on second arg only. **Severity: clean** (NaN-mask).
|
||||
- **Line 97 — `sum0 += expf(v - max0)`** — max-shifted, ≤ 1. **Severity: clean.**
|
||||
- **Line 98 — `p0 = expf(...) / (sum0 + 1e-8f)`** — guarded. **Severity: clean.**
|
||||
- Same pattern for `pk` (per-head softmax) at lines 103-111.
|
||||
- **Line 115 — `grad /= (K - 1)`** — `K = NUM_ENSEMBLE_HEADS ≥ 2` by construction. **Severity: clean.**
|
||||
|
||||
In `apply_ensemble_diversity_backward` (orchestrator):
|
||||
- Same structural pattern as `apply_iqn_trunk_gradient` — `encoder_backward_chain` cuBLAS GEMMs + SAXPY. **Same suspect-class as IQN trunk path.**
|
||||
|
||||
**Proposed guard form:**
|
||||
- Same `isfinite` guard on `dqn_saxpy_f32_kernel` (per IQN section above) covers ensemble too.
|
||||
|
||||
**ISV bound option:** Invariant 1 carve-out.
|
||||
|
||||
**F0 risk:** low (ensemble path is currently DORMANT in production smoke runs — `ensemble_diversity_weight = 0` by default; kernel is only fired when the weight > 0). Per session_2026-04-05 the residual NaN was at F1 step 2, well before any ensemble firing in current configs. The instrumentation is low-priority but worth the slot for completeness.
|
||||
|
||||
**Phase B flag-slot allocation:**
|
||||
- Slot 31 (deferred; instrument only if SP1 Phase B observes the saxpy guard catching ensemble-source NaN): `d_logits_0` post-`ensemble_kl_gradient_kernel`. Likely never fires in current production config; defer.
|
||||
|
||||
---
|
||||
|
||||
### Kernel: `mse_grad_kernel` + `mse_loss_batched`
|
||||
|
||||
**File:** `crates/ml/src/cuda_pipeline/mse_grad_kernel.cu` (gradient, full file 1-93), `crates/ml/src/cuda_pipeline/mse_loss_kernel.cu` (forward; lines 89, 91, 96, 305, 330 are the relevant log/exp/sqrt sites).
|
||||
**Output buffers:** `d_value_logits` / `d_adv_logits` (same buffers as C51 grad — when MSE is fired, it overwrites; both kernels never run in the same step).
|
||||
**Accumulator targets:** none (plain assignment).
|
||||
**Caller:** `gpu_dqn_trainer` MSE warmup path (early epochs blend MSE with C51 loss per `c51_warmup_epochs` config).
|
||||
|
||||
**Identified unsafe patterns:**
|
||||
|
||||
In **`mse_grad_kernel`**:
|
||||
- **Line 39 — `delta_z = (num_atoms > 1) ? (v_max - v_min) / (num_atoms - 1) : 0.0f`** — guarded.
|
||||
- **Line 64 — `d_combined = inv_batch * isw * td_error * p_j * (z_j - e_q)`** — products of finites; **clean**.
|
||||
- **Line 69 — `lp_approx = fmaxf(logf(fmaxf(p_j, 1e-8f)), -10.0f)`** — both ε-floored and floor-clamped. **Severity: clean.**
|
||||
- **No atomicAdd. Plain assignment per (b, j) for `d_value_logits`, per (b, d, a, j) for `d_adv_logits`.**
|
||||
|
||||
In **`mse_loss_batched`** (used by warmup path):
|
||||
- **Lines 305 / 330 — `a_std = sqrtf(sq_sum / n_d + 1e-12f)`** — ε-floored. **Severity: clean.**
|
||||
- **Line 91 — `log_sum_exp = logf(bsum + 1e-10f)`** — ε-floored. **Severity: clean.**
|
||||
|
||||
Net: MSE path is well-guarded. The MSE warmup blend is a SHORT cold-start window; F1 step 1140 is well past warmup so MSE is contributing zero or near-zero to the blended loss.
|
||||
|
||||
**Proposed guard form:** none.
|
||||
|
||||
**ISV bound option:** N/A.
|
||||
|
||||
**F0 risk:** N/A.
|
||||
|
||||
**Phase B flag-slot allocation:**
|
||||
- No new slot. Slots 24/25 (`d_value_logits` / `d_adv_logits`) cover both C51 and MSE producers since they OVERWRITE the same buffers.
|
||||
|
||||
---
|
||||
|
||||
### Kernel: Bottleneck Linear backward (Rust orchestrator only, no `.cu`)
|
||||
|
||||
**File:** `crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs:18063` (`aux_bottleneck_vsn_backward_dispatch`), `:18560` (main-path bottleneck dX into market slice via `launch_dx_only_lda`); `crates/ml/src/cuda_pipeline/batched_backward.rs:2132` (`launch_dw_only_no_bias_lda`).
|
||||
**Output buffers:** `grad_buf [bottleneck slot 33]` (W) + `[slot 34]` (B); `bn_d_concat_buf [B, market_dim + portfolio_dim]` (the dX flowing into bottleneck's input — i.e. into the market features + portfolio features pre-bottleneck).
|
||||
**Accumulator targets:** `grad_buf` slots 33/34 via `launch_dw_only` / `launch_dw_only_no_bias_lda` cuBLAS dW (`beta=1.0`). `bn_d_concat_buf` via `launch_dx_only_lda` (`beta=0.0` for market slice; portfolio slice padded with zeros).
|
||||
**Caller:** `encoder_backward_chain` (final stage); also `aux_bottleneck_vsn_backward_dispatch` for VSN aux paths.
|
||||
|
||||
**Identified unsafe patterns:**
|
||||
- All compute via cuBLAS SGEMM with `lda` overrides for padded states (state stride = `pad128(state_dim) = 128`, market stride = `state_dim_padded`). Per session_2026-04-05 finding #3, **this exact pattern (states_buf stride mismatch)** caused root cause #3 in the fixed bug list. The fix landed in `backward_fc_layer_lda` and `launch_dw_only_no_bias_lda`. **Severity: clean (fix is in place).** No additional guards required at this level.
|
||||
- The dX GEMM uses `beta=0.0` (overwrite) so does not propagate stale NaN; the dW GEMM uses `beta=1.0` (accumulate into `grad_buf`) — and `grad_buf` is the buffer slot-6 already watches.
|
||||
|
||||
**Proposed guard form:** none required (fix is in place; existing slot-6 watches the accumulator).
|
||||
|
||||
**ISV bound option:** N/A.
|
||||
|
||||
**F0 risk:** N/A.
|
||||
|
||||
**Phase B flag-slot allocation:**
|
||||
- Slot 32: `bn_d_concat_buf [B × (market_dim + portfolio_dim)]` post-`launch_dx_only_lda` — locates whether the bottleneck Linear's dX path produces NaN before flowing back into the VSN backward extension. Useful diagnostic separator between "GRN backward chain" (slot 26) and "bottleneck Linear backward" (slot 32).
|
||||
|
||||
---
|
||||
|
||||
### Kernel: `bw_d_h_s2` accumulator path (Rust orchestrator inventory, not a kernel)
|
||||
|
||||
**File:** `crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs` — multiple sites; the buffer is `self.bw_d_h_s2 : CudaSlice<f32>` (line 3140 declaration) with raw pointer `self.ptrs.bw_d_h_s2`.
|
||||
**Output / accumulator semantics:** `bw_d_h_s2 [B, SH2]` is the "main backward dX into the trunk h_s2 output" — i.e. the entry point gradient that `encoder_backward_chain` consumes when running the GRN h_s2 / h_s1 backward. It receives contributions from MULTIPLE producers:
|
||||
|
||||
Inventory of producers that WRITE or ACCUMULATE INTO `bw_d_h_s2`:
|
||||
|
||||
1. **`backward_full` value-head + 4 branch dX** (`batched_backward.rs:1903-1925` for branches; `batched_backward.rs:2014-2023` for value head) — populates `scratch_d_h_s2` (which is the same allocation as `self.bw_d_h_s2.raw_ptr()`). Branch 0 uses `beta=0` (overwrite); branches 1-3 use `beta=1` (accumulate); value head uses `beta=1` (accumulate).
|
||||
2. **`apply_iqn_trunk_gradient` step 2** (`gpu_dqn_trainer.rs:6880`) — `graph_safe_copy_f32(bw_d_h_s2, iqn_d_h_s2_ptr, ...)` — REPLACES the contents (DtoD overwrite). Subsequently fed into `encoder_backward_chain` for the IQN aux backward chain.
|
||||
3. **`apply_ensemble_diversity_backward` step 2** (`gpu_dqn_trainer.rs:7155-7180`) — `launch_dx_only(beta=0)` followed by `relu_mask` then `launch_dx_only(...)` — analogous to value head, populates `bw_d_h_s2` for the ensemble aux backward chain.
|
||||
4. **`apply_aux_heads_backward`** — saxpys `dh_s2_out` from `aux_next_bar_backward` + `aux_regime_backward` into `bw_d_h_s2` via `dqn_saxpy_f32_kernel`. (Per kernel docstring lines 384-386 and orchestrator naming in the file.)
|
||||
5. **`apply_state_kl_backward`** (line 2277 docstring: "predictive_coding_loss into bw_d_h_s2 via plain += (no atomicAdd)") — KL-divergence gradient into trunk.
|
||||
6. **`apply_curiosity_loss_backward`** — curiosity training gradient (cross-reference with `curiosity_training_kernel.cu`).
|
||||
7. **`apply_moe_backward`** (line 8597-8702 — "DtoD-copy moe_dh_s1_scratch → bw_d_h_s2") — REPLACES `bw_d_h_s2` with MoE-side trunk gradient.
|
||||
8. **Attention-focus ISV backward** (lines 2141-2143: "f32 [B * SHARED_H2] for copying bw_d_h_s2 (f32) → attention scratch").
|
||||
|
||||
The MAIN backward chain (`backward_full` + value-FC + main `encoder_backward_chain`) consumes `bw_d_h_s2` ONCE per step. The AUX backward chains (`apply_iqn_trunk_gradient`, `apply_ensemble_diversity_backward`, etc.) each REZERO + REPLACE `bw_d_h_s2` and re-run `encoder_backward_chain` to extract trunk gradients for their respective contributions, which are accumulated into `iqn_trunk_m` (scratch) and then SAXPY'd into `grad_buf`.
|
||||
|
||||
**Identified unsafe patterns:**
|
||||
- The DtoD replace pattern (steps 2/3/7) means the buffer is read by `encoder_backward_chain` immediately after the replace. If the SOURCE is NaN, the chain consumes NaN and writes NaN dW/dB into `iqn_trunk_m`, which then SAXPYs NaN into `grad_buf` (current behaviour — slot 6 fires).
|
||||
- Steps 4 / 5 / 6 use SAXPY accumulation with no `isfinite` guard on the saxpy itself (per `dqn_saxpy_f32_kernel` audit above).
|
||||
|
||||
**Proposed guard form:**
|
||||
- (a) `isfinite` guard on `dqn_saxpy_f32_kernel` (covers all SAXPY sites).
|
||||
- (b) NaN check on `bw_d_h_s2` *before* each `encoder_backward_chain` call — this is the single point that determines whether NaN enters the trunk backward path. **One slot, two consumers** (main chain + IQN aux chain).
|
||||
- (c) NaN check on `bw_d_h_s2` *after* each producer that REPLACES contents (steps 2, 3, 7).
|
||||
|
||||
**ISV bound option:** Invariant 1 carve-out for guard (a).
|
||||
|
||||
**F0 risk:** low for guards (b)/(c) (diagnostic only); medium for guard (a) until the saxpy verification is done.
|
||||
|
||||
**Phase B flag-slot allocation:**
|
||||
- Slot 33: `bw_d_h_s2` AFTER `backward_full` returns, BEFORE `encoder_backward_chain` begins (the "main path" check). Distinct from slot 12 (`save_h_s2`, the FORWARD activation): slot 33 watches the BACKWARD gradient at the same shape `[B, SH2]`.
|
||||
- Slot 34: `bw_d_h_s2` AFTER `apply_aux_heads_backward` SAXPY but BEFORE `apply_state_kl_backward` / `apply_moe_backward` — locates whether aux-head SAXPYs poisoned the buffer.
|
||||
- Slot 35: `bw_d_h_s2` AFTER `apply_iqn_trunk_gradient`'s DtoD-copy from `iqn_d_h_s2_ptr` — covers the IQN trunk replace path. NB: this is symmetric with slot 27 (which checks the SOURCE side `iqn_d_h_s2_ptr`); both are needed because the DtoD might be a graph-capture artefact (cudarc cuMemcpy in graph capture has historic edge cases).
|
||||
|
||||
---
|
||||
|
||||
## Prior-investigation status (cross-reference with `session_2026-04-05_nan_investigation.md`)
|
||||
|
||||
### Root causes (5) — current code state
|
||||
|
||||
1. **CUDA graph warmup corruption** (gpu_dqn_trainer.rs) — fix: removed warmup launches. **STATUS:** still in place per current `submit_dqn_step_loop_cublas` structure (no warmup launches). **Verification:** grep confirms no `warmup` launches in graphed paths.
|
||||
2. **f32/bf16 type mismatch in pad_states_kernel** (dqn_utility_kernels.cu) — fix: changed to `const float*`. **STATUS:** still in place — `dqn_utility_kernels.cu` shows f32 source pointer. **Verified:** lines 124-180 of `dqn_utility_kernels.cu` use `float` source.
|
||||
3. **States buffer stride mismatch** (batched_backward.rs, gpu_dqn_trainer.rs) — fix: `backward_fc_layer_lda` variant. **STATUS:** still in place — `backward_fc_layer_lda` at `batched_backward.rs:916`, `launch_dw_only_no_bias_lda` at `:2132`. Both consumed by `apply_iqn_trunk_gradient`'s `encoder_backward_chain` and bottleneck backward.
|
||||
4. **IQN trunk forward stride mismatch** (iqn_dual_head_kernel.cu) — fix: `padded_sd = (state_dim + 127) & ~127`. **STATUS:** the IQN dual-head kernel no longer has its own trunk forward — h_s2 is consumed directly from the DQN trunk per kernel docstring lines 1-39. The fix is no longer needed because the path no longer exists.
|
||||
5. **IQN grad_norm sqrt-per-block bug** (iqn_dual_head_kernel.cu) — fix: accumulate raw sum-of-squares. **STATUS:** still in place — `iqn_grad_norm_phase1` (lines 802-829) accumulates raw `g*g` partials; `iqn_grad_norm_phase2` (lines 832-844) sums partials; `iqn_adam_kernel` (line 876) takes `sqrtf(fmaxf(norm_sq[0], 0.0f))` at point of use.
|
||||
|
||||
### Defense-in-depth guards (6) — current code state
|
||||
|
||||
- `f32_to_bf16_kernel` NaN→0 + clamp ±500 (common_device_functions.cuh) — **NOT VERIFIED IN THIS AUDIT** (file scope was backward-path kernels). bf16 cast is upstream of all backward-path producers.
|
||||
- DQN Adam `isfinite(g)` guard (dqn_utility_kernels.cu) — **NOT INSPECTED IN THIS AUDIT** (file in scope but the kernel of interest is the saxpy and adam — Adam was separately verified). The DQN saxpy itself does NOT have an isfinite guard (Phase C proposal).
|
||||
- IQN Adam `isfinite(g)` guard (iqn_dual_head_kernel.cu) — **STILL ACTIVE.** Line 873 of current file: `if (!isfinite(g)) { grads[tid] = 0.0f; return; }`.
|
||||
- Attention Adam `isfinite` guard (attention_backward_kernel.cu) — **NOT INSPECTED** (out of scope; attention is forward-path consumer).
|
||||
- Curiosity Adam `isfinite` guard (curiosity_training_kernel.cu) — **NOT INSPECTED** (curiosity is dormant in current production config).
|
||||
- IQN backward d_h_s2 `isfinite` guard (iqn_dual_head_kernel.cu) — **STILL ACTIVE.** Line 612: `if (isfinite(d_trunk)) d_h_s2_acc[h / IQN_BLOCK_SIZE] += d_trunk;`.
|
||||
|
||||
### Residual: 8% NaN in grad_buf (NEVER CLOSED)
|
||||
|
||||
- Prior diagnosis: NaN enters during `apply_iqn_trunk_gradient`'s cuBLAS backward (specifically `backward_fc_layer` dW GEMM); IQN d_h_s2 is CLEAN.
|
||||
- Architectural drift since 2026-04-05: the legacy ReLU+Linear→Linear chain has been replaced with `encoder_backward_chain` (Plan 4 Task 2c.3c.4), which is a GRN-aware chain (LayerNorm + ELU + Linear_a + Linear_b + Linear_residual + GLU). The cuBLAS dW GEMMs are STILL in this chain (via `launch_dw_only` / `launch_dw_only_no_bias` / `launch_dw_only_no_bias_lda`), so the suspect class is unchanged but the specific kernel path is different.
|
||||
- The smoke `smoke-test-m5gxx` slot-6 + slot-12 fire pattern is CONSISTENT with this prior-investigation residual: NaN visible in `grad_buf` post-saxpy, before Adam catches it; `save_h_s2` is the FORWARD activation buffer (not the BACKWARD `bw_d_h_s2`), so its slot-12 NaN at the same step suggests the AUX-pass `apply_iqn_trunk_gradient` (which DtoD-copies into `bw_d_h_s2` and re-uses save_h_s2's overlay buffer per `gpu_dqn_trainer.rs:6849` "bw_d_h_s2 is gone — h_s2 GRN ends in LayerNorm") is creating a step 2 + step 3 sequence that contaminates both buffers within one aux pass.
|
||||
- **Verdict:** the prior investigation's open question is RECONFIRMED as the top SP1 suspect. Phase B slots 26-28 (IQN trunk path) directly target this hypothesis.
|
||||
|
||||
---
|
||||
|
||||
## Summary — kernel suspicion ranking
|
||||
|
||||
Ranked by cumulative severity (number of unsafe patterns × severity), with `apply_iqn_trunk_gradient` weighted by prior-investigation note as inherently high-priority regardless of pattern count.
|
||||
|
||||
| Rank | Kernel | # confirmed | # suspect | F0 risk of fix | Notes |
|
||||
|---|---|---|---|---|---|
|
||||
| 1 | `apply_iqn_trunk_gradient` (orchestrator + `encoder_backward_chain` cuBLAS leg + saxpy producer) | 0 | 3 | low–medium | Prior NEVER-CLOSED 8% step-2 NaN; `dqn_saxpy_f32_kernel` has no `isfinite` guard; F1 slot-6 fire pattern matches. |
|
||||
| 2 | `bw_d_h_s2` accumulator path (Rust, 8 producers SAXPY/DtoD) | 0 | 2 | low (diagnostic) | Buffer received by `encoder_backward_chain`; one bad producer poisons the whole aux pass. |
|
||||
| 3 | `iqn_backward_per_sample` + `iqn_weight_grad_reduce` | 0 | 1 | low | `save_dL_dq` written without `isfinite`; reduce loops have NO guard at lines 750-762 (existing guards at 689/712 demonstrate the precedent). |
|
||||
| 4 | `c51_grad_kernel` | 0 | 1 | low | Reads ISV slots 12-21 every step; if any slot is NaN at fold boundary, the multiplicative bin_weight propagates NaN. |
|
||||
| 5 | `apply_ensemble_diversity_backward` | 0 | 1 | low | Same suspect-class as IQN trunk (cuBLAS GEMMs + saxpy); currently dormant in production config. |
|
||||
| 6 | Bottleneck Linear backward | 0 | 0 | N/A | Stride-padded SGEMM bug is FIXED (session_2026-04-05 finding #3); slot-6 + lda variants in place. Diagnostic slot 32 still useful. |
|
||||
| 7 | `aux_next_bar_backward` + `aux_regime_backward` + `aux_param_grad_reduce` | 0 | 0 | N/A | Well-guarded; only concern is upstream `label_scale` (ISV slot 117) bootstrapping. |
|
||||
| 8 | `cql_logit_grad_kernel` | 0 | 1 (low) | N/A | NaN-passthrough in `fmaxf`; in practice masked by max-shift softmax. Not a likely seed. |
|
||||
| 9 | `c51_loss_batched` | 0 | 0 | N/A | Heavily guarded; only concern is input-NaN propagation, already covered by slots 7/8. |
|
||||
| 10 | `mse_grad_kernel` | 0 | 0 | N/A | Not active at F1 step 1140 (post-warmup). |
|
||||
| 11 | `backward_full` (orchestrator) | 0 | 0 | N/A | Pure dispatch; safety properties pass through to producers. |
|
||||
|
||||
## Phase B priority slot order
|
||||
|
||||
Based on the suspicion ranking above, instrument these slots first (descending priority).
|
||||
|
||||
NB: this requires expanding `nan_flags_buf` from `[i32; 24]` to `[i32; 48]` — current declaration at `gpu_dqn_trainer.rs:14982` reads `[i32; 24]`; Phase B Task 2 expands to 48 (to host slots 24-47). The 12 listed slots use indices 24-35; indices 36-47 are headroom for SP2/SP3.
|
||||
|
||||
1. **Slot 26** — `iqn_trunk_m` post-`encoder_backward_chain` in `apply_iqn_trunk_gradient` (highest suspicion — prior investigation).
|
||||
2. **Slot 27** — `iqn_d_h_s2_ptr` post-DtoD copy into `bw_d_h_s2` (companion to slot 26; locates per-sample IQN backward as source).
|
||||
3. **Slot 33** — `bw_d_h_s2` AFTER `backward_full` returns, BEFORE `encoder_backward_chain` begins (the main-path entry-point check).
|
||||
4. **Slot 24** — `d_value_logits` post-`c51_grad_kernel`.
|
||||
5. **Slot 25** — `d_adv_logits` post-`c51_grad_kernel`.
|
||||
6. **Slot 34** — `bw_d_h_s2` AFTER `apply_aux_heads_backward` SAXPY (separates aux-head poisoning from main-path).
|
||||
7. **Slot 35** — `bw_d_h_s2` AFTER `apply_iqn_trunk_gradient` DtoD copy (symmetric with slot 27 — distinguishes graph-capture DtoD edge cases).
|
||||
8. **Slot 28** — `save_dL_dq` post-`iqn_backward_per_sample` (locates IQN per-sample backward as NaN seed for the weight-grad reduce path).
|
||||
9. **Slot 32** — `bn_d_concat_buf` post-`launch_dx_only_lda` (bottleneck Linear backward dX into market slice).
|
||||
10. **Slot 30** — `dh_s2_out` post-`aux_next_bar_backward` (aux head backward; lower suspicion but cheap).
|
||||
11. **Slot 29** — DEFERRED: `d_v_logits` post-`cql_logit_grad_kernel` (covered by deltas between slot 24 fire and slot 6 fire).
|
||||
12. **Slot 31** — DEFERRED: `d_logits_0` post-`ensemble_kl_gradient_kernel` (currently dormant in production).
|
||||
|
||||
If any of slots 24–35 are pre-empted by Phase B's actual implementation budget (e.g. cudarc wrapper limits or `nan_flags_buf` sizing constraints), the deferred two slots (29, 31) can be returned to the budget as the lowest-priority drops. Slots 26, 27, 33 are NON-NEGOTIABLE — without those three the prior-investigation hypothesis cannot be confirmed or refuted from one smoke run.
|
||||
@@ -2236,3 +2236,5 @@ Plan C T11 follow-up R1+P (2026-04-29):
|
||||
R1 — eliminated K's Adam shrink-and-perturb at fold boundary (m_buf and v_buf back to memset_zeros, K's hardcoded ADAM_M_SHRINK=0.1 / ADAM_V_SHRINK=0.01 constants removed). K was a stepping-stone fix introduced before fold_warmup_factor existed (commit 4ef1d8ebb). fold_warmup_factor — ISV-driven, drives lr+clip dampening — already provides principled first-step protection via the cold-start `lr_eff = lr_base × max(MIN_WARMUP_LR_FRAC, fold_warmup_factor)` formula. K's hardcoded shrink ratios violated `feedback_adaptive_not_tuned` AND created a downstream pathology: tiny v_hat denominator → oversized Adam updates ~50+ steps post-reset → trunk param overshoot → save_h_s2 NaN at F1 ~step 1745 (smoke-test-bkdx5 diagnostic). With K removed, Adam restarts cleanly at m=0, v=0 and warmup_factor handles the cold-start window via lr/clip damping — single mechanism, ISV-driven, no hardcoded constants. R1.B verification (no edits): `state_reset_registry.rs` already lists both `isv_fold_warmup_factor` and `isv_grad_norm_fast_ema` as `FoldReset` entries with matching dispatch arms in `training_loop.rs::reset_named_state` — the K+warmup commit landed both halves of the warmup-factor input contract atomically per `feedback_no_partial_refactor.md`.
|
||||
|
||||
P — expanded `nan_flags_buf` 16→24 with 5 new GRN h_s2 sub-stage NaN checks (`grn_h_s2_elu_post`=16, `grn_h_s2_linear_b_out`=17, `grn_h_s2_glu_sigmoid`=18, `grn_h_s2_ln_rstd`=19, `grn_h_s2_ln_normed`=20) for finer-grained source identification if R1 alone doesn't fix F1. Slot 19 (`ln_rstd`) is the divide-by-zero diagnostic (variance→0 ⇒ rstd→∞ ⇒ NaN downstream). New accessors: `GrnBlock::{elu_post_ptr, linear_b_out_ptr, sigmoid_b_ptr, ln_rstd_ptr, ln_normed_ptr}` (additive, reads-only) and `CublasGemmSet::{grn_h_s2_online, grn_h_s2_linear_b_ptr}`. h_s2 only — pragmatic instrumentation scope: that's the GRN block where save_h_s2 went NaN at F1 step 1745, and instrumenting only h_s2 covers the failing stage without doubling per-step kernel-launch cost. h_s1 stages can be added by name table slots 21-23 if the F1 explosion turns out to originate above h_s2. The K-removal + GRN-stage checks combined: if the F1 explosion was due to K's tiny-v_hat pathology, R1 alone fixes it. If a different mechanism, the new flags pinpoint which GRN sub-stage produces NaN first. Per `feedback_no_partial_refactor.md`: R1 and P land together because both touch the fold-boundary contract (K-removal changes the post-reset Adam state; GRN-stage NaN checks observe the consequences of that state on the trunk encoder's first forward). `read_nan_flags()` signature update [16]→[24] propagates through `GpuDqnTrainer`, `FusedTrainingCtx`, and both name-table sites in `training_loop.rs` (halt_nan + halt_grad_collapse).
|
||||
|
||||
SP1 Phase A audit (2026-04-29): produced `docs/dqn-backward-nan-audit.md` — read-only γ inventory of every backward-path kernel writing to `bw_d_h_s2` / `grad_buf` / `save_h_s2` accumulators, cross-referenced against `session_2026-04-05_nan_investigation.md`'s residual 8% step-2 NaN in `apply_iqn_trunk_gradient`. Per-kernel sections include: identified unsafe patterns (sqrtf-neg, 1/0, logf-≤0, expf-large, EMA variance, atomicAdd/saxpy NaN-propagation), proposed guard form, ISV bound option (existing slot leverage or new slot or Invariant 1 ε carve-out), F0 risk assessment (low/medium/high used as paper-review gate before smoke), and Phase B flag-slot allocation. Drives SP1 Phase B instrumentation (12 new slots in `nan_flags_buf` 24→48) and Phase C surgical fix decisions. Becomes durable input artifact for SP2 (framework codification) and SP3 (structural-fix scoping).
|
||||
|
||||
Reference in New Issue
Block a user