spec(dqn): SP2 + SP3 combined — F0 regression fix + Q-learning structural stability

Designs the fixes for the two SP1 leftover issues identified at SP1 closure
(commit 047f175fb):

1. SP2 — F0 regression fix (instrumentation overhead optimization):
   Consolidate the 11 per-step check_nan_f32 launches in
   run_nan_checks_post_backward into a single fused-reduce kernel
   (dqn_nan_check_fused_f32_kernel). 1 launch instead of 11; expected
   F0 return to ~55 baseline (was regressing to ~35 across 3 SP1 smokes).

2. SP3 — Q-learning structural stability (5 mechanisms):
   - M1: Target-Q clipping at single-point source, ISV-driven
     (10 × isv[Q_ABS_REF].max(1.0))
   - M2: C51 atom-position growth bounds, same ISV bound
   - M3: Hard target sync at fold boundary (DQN + IQN + ensemble target nets)
   - M4: Adam EMA reset at fold transitions (comprehensive — all Adams)
   - M5: Embedded Adam-centric diagnostic instrumentation (slots 36-47):
     Adam m/v max-abs (slots 36-43), weight max-abs (44-45),
     target_q post-clip (46), atom span (47). Sticky-bit threshold
     checks; SP2's fused kernel is extended (NOT replaced) to handle
     all 24 slots in 24 blocks.

Architecture: SP2 lands first; Gate 1 must pass (F0 ≥ 53.08 + slot
24-35 firing topology unchanged) before SP3 starts. SP3's 5 mechanisms
ship in ONE commit per feedback_no_partial_refactor (related fixes
for the same root pathology — Q-target inflation drives Adam EMA
saturation drives weight drift drives cuBLAS overflow). Gate 2
evaluates the full SP1 7-criterion bar.

Operating principles (mandatory):
- ISV-driven design for every dynamic bound (no hardcoded clip values)
- ε floor on ISV multiplier per SP1 pearl (`isv.max(1.0)`)
- Combined RELATED commits per feedback_no_partial_refactor
- Permanent diagnostic infrastructure (always-on, no debug flags)
- No deferrals — anomalies during impl get fixed within scope

Validation cost: SP2 ~1-2 days + 1 smoke; SP3 ~3-5 days + 1-2 smokes.
Combined ~1000 LOC across 2 commits.

Spec covers: architecture, components (SP2 fused kernel + SP3 5
mechanisms), data flow, ISV slot bindings, validation gates, error
handling failure modes, operating principles, anti-patterns,
sequencing constraints, effort estimate, handoff conditions.
This commit is contained in:
jgrusewski
2026-04-30 08:36:32 +02:00
parent 047f175fbf
commit ed727b51c0

View File

@@ -0,0 +1,312 @@
# Numerical Stability SP2 + SP3 — Combined Design Spec
**Status:** APPROVED 2026-04-29 — ready for `superpowers:writing-plans` to produce implementation plan.
**Branch:** `plan-c-phase-2-thompson` (continuation; latest commit `047f175fb` SP1 closure)
**Scope:** combined spec covering SP2 (instrumentation overhead) + SP3 (Q-learning structural stability)
**Predecessor:** `docs/superpowers/specs/2026-04-29-numerical-stability-investigation-design.md` (the original 3-sub-project decomposition; SP1 closed at commit `047f175fb`)
**Audit doc (durable artifact):** `docs/dqn-backward-nan-audit.md`
## 1. Background and motivation
SP1 (F1 NaN root-cause investigation) closed at commit `047f175fb` with surgical-fix-scope deliverables:
- Permanent diagnostic infrastructure (`nan_flags_buf` 24→48, slots 24-35 verified working across 3 L40S smokes)
- F1 NaN entry kernels pinpointed: `apply_iqn_trunk_gradient` cuBLAS sgemm (slot 26) + Bottleneck Linear backward sgemm (slot 32)
- Cold-start clamp pathology fixed via ε-floor pearl: `1e6 × isv.max(1.0)` (validated F1 NaN moved step 240 → 3540, 15× later)
Two open issues remain that SP1's surgical-fix scope cannot address:
1. **F0 regression** to ~35 (vs baseline 55.87). Three consecutive smokes (smoke-test-xvzgk, dr2bn, ckgv8) show F0 ≈ 35 across post-Phase-B commits. Phase C surgical fix did not change F0 — confirming the regression is structural to Phase B instrumentation. **SP2 territory.**
2. **F1 step-3540 structural NaN.** Same flagged signature `[6, 12, 26, 32]` recurs at step 3540 with clean inputs (slots 27, 35). cuBLAS GEMM produces NaN/Inf because weights themselves accumulate NaN/Inf via Adam updates over many training steps. Surgical clamps on gradient inputs cannot prevent this; root cause is structural Q-learning instability. **SP3 territory.**
This combined spec designs fixes for both, sequenced so SP2 lands first (its F0 fix unblocks SP3 evaluation).
## 2. Success criteria
The combined spec must satisfy SP1's original 7 pass criteria on a single multi-fold L40S smoke:
1. Zero `NaN-CLAMPED-TO-ZERO` log lines
2. Zero `NaN/Inf at step` errors
3. F0 Best Sharpe ≥ 53.08 (95% of 55.87 baseline)
4. F1 last-epoch Best Sharpe ≥ first-epoch Best Sharpe
5. F2 last-epoch Best Sharpe ≥ first-epoch Best Sharpe
6. Zero flagged NaN/threshold slots (slots 24-47 all remain at zero)
7. All 3 folds complete 5 epochs each
Criterion 6 must hold across the new 24-47 slot range (slots 36-47 are SP3's threshold-trigger diagnostics; their staying at zero confirms SP3 mechanisms are no-ops on healthy training).
## 3. Architecture
```
SP2 (foundation) SP3 (structural)
───────────────── ─────────────────
Fused NaN check kernel Mech 1: Target-Q clipping
scans 12 backward slots → Mech 2: C51 atom-position growth bounds
in parallel, 1 launch Mech 3: Hard target sync at fold boundary
vs current 11 launches Mech 4: Adam EMA reset at fold transitions
Mech 5: Adam-centric diagnostic (slots 36-47)
Extends SP2's fused kernel — no second kernel
```
**Sequencing constraint:** SP2 lands and passes Gate 1 BEFORE SP3 starts. SP3 evaluation (full SP1 7 criteria) is confounded by F0 regression.
**ISV-driven principle (mandatory per `feedback_isv_for_adaptive_bounds`):** every dynamic bound in SP3 derives from existing ISV slots. ε floors go on the ISV multiplier (`isv.max(1.0)`), not on the bound itself (per the SP1 ε-floor pearl).
**Data flow per training step (post SP2+SP3):**
```
forward → 1 fused fwd-NaN check (existing) → loss kernels
target Q → CLIP (Mech 1) → loss → backward
backward kernels → 1 fused bwd-NaN+threshold check (NEW; SP2 + SP3 diag)
SAXPY → grad_buf → Adam (with EMA scan via diag) → weights
fold boundary? → hard target sync (Mech 3)
→ Adam EMA reset (Mech 4)
→ atom EMA reset (Mech 2 init)
```
## 4. SP2 — Fused NaN check kernel
### Component A: `dqn_nan_check_fused_f32_kernel` (new CUDA kernel)
Replaces 8 individual `check_nan_f32` launches in `GpuDqnTrainer::run_nan_checks_post_backward`. Slots 33/34/35 (`bw_d_h_s2` multi-point snapshots) stay as inline checks because they fire at distinct backward-orchestration phases for localization.
**Kernel signature:**
```cuda
extern "C" __global__ void dqn_nan_check_fused_f32_kernel(
const float* const* buf_ptrs, // [N_SLOTS] device array of buffer pointers (or null = skip)
const int* buf_lens, // [N_SLOTS] device array of buffer lengths
int base_flag_idx, // first slot index in nan_flags_buf (24)
int* nan_flags_buf
);
```
**Grid layout:** 1 block per slot, 256 threads per block. Block N processes slot `base_flag_idx + N`. Block-local grid-strided scan reduces "any NaN" via `__syncthreads_or`.
**File:** appended to `crates/ml/src/cuda_pipeline/dqn_utility_kernels.cu` (alongside existing `dqn_nan_check_f32`).
### Component B: `nan_check_meta` device buffer
New field on `GpuDqnTrainer`:
```rust
/// SP2 fused NaN check metadata: 12 (ptr, len) pairs for slots 24-35.
/// Populated ONCE at construction; reused every step (slot pointers
/// are stable). Slot 31 (ensemble_d_logits_buf) entries are null/0
/// when gpu_iqn = None — kernel skips that block.
nan_check_meta: CudaSlice<u8>,
```
Sized `12 × (sizeof::<u64>() + sizeof::<i32>())`. Populated at construction with the (ptr, len) tuples for slots 24-30, 31 (deferred = null), 32. No per-step host→device update.
### Component C: `launch_nan_check_fused_f32` (Rust wrapper)
Replaces 8 sequential `check_nan_f32` calls in `run_nan_checks_post_backward`. Single launch, `grid_dim=(12, 1, 1)`, `block_dim=(256, 1, 1)`.
### Component D: SP2 Gate 1 validation
After SP2 lands, single L40S smoke must show:
- F0 Best Sharpe ≥ 53.08 (returns to ~55 baseline expected)
- All 12 backward-path slots still detectable — confirmed by F1 step-3540 NaN STILL FIRING with the existing ε-floor fix (the live regression test embedded in the codebase)
- No new errors / warnings
If F0 doesn't recover, the diagnosis (Phase B kernel-launch overhead) was wrong; halt before SP3 starts.
## 5. SP3 — Q-learning structural stability (5 mechanisms)
### Mechanism 1: Target-Q clipping (always-on, single-point)
**Where:** at the target Q computation site — clip the `target_q` output buffer BEFORE it enters the loss kernels (C51 / IQN / MSE all consume the same buffer).
**Algorithm:**
```rust
let q_abs_ref = self.read_isv_signal_at(Q_ABS_REF_INDEX); // existing slot 16
let max_abs_target_q = 10.0_f32 * q_abs_ref.max(1.0_f32); // ε on multiplier
self.launch_clamp_finite_f32(target_q_ptr, target_q_len, max_abs_target_q)?;
```
**Reuses** `dqn_clamp_finite_f32_kernel` from SP1. No new kernel.
**ISV bound:** `Q_ABS_REF_INDEX = 16` (existing) — tracks typical Q magnitude. Multiplier `10×` allows full Q distribution coverage; ε-floor `1.0` on multiplier guarantees `max_abs ≥ 10` even with cold-start ISV.
**F0 risk: low.** F0-typical |target_q| bounded by reward × horizon ≤ ~10. With Q_ABS_REF EMA ≈ 1-5 at F0, max_abs = 10-50. Guard is no-op for F0. F1 inflation (thousands+) gets clamped.
**Why single-point not per-loss-kernel:** all three losses consume the same target_q buffer. Single source covers all consumers in one launch. Per-loss-kernel clipping would create 3 sites with risk of divergent semantics.
### Mechanism 2: C51 atom-position growth bounds
**Where:** in the existing C51 atom-update kernel where atoms grow via EMA toward the observed Q distribution.
**Algorithm:** during atom-position EMA update, clamp the new atom values:
```cuda
const float max_atom_abs = 10.0f * fmaxf(isv_signals[Q_ABS_REF_INDEX], 1.0f);
new_pos = fminf(fmaxf(new_pos, -max_atom_abs), max_atom_abs);
```
(~5 lines added to existing kernel.)
**ISV bound:** same as Mech 1. Atoms span the Q range, so their bound matches.
**F0 risk: low.** Atoms naturally span typical Q range; capping at 10× allows full distribution coverage without extreme-tail growth.
### Mechanism 3: Hard target sync at fold boundary
**Current state:** target net updates via Polyak-EMA every step (`target ← τ·online + (1-τ)·target`). IQN side already has partial coverage per memory file `project_fold_boundary_q_drift_resolved.md` ("IQN sync_target_from_online + 7 aux Adam resets").
**Proposal:** at fold boundary, replace the gradual Polyak with a HARD sync (`target ← online`) for ALL target nets:
- DQN main target net
- IQN target (extend existing partial coverage)
- Ensemble target nets (if active)
**Where:** in the fold-transition handler invoked by `StateResetRegistry` from `training_loop.rs`. Extends the existing IQN-side mechanism to cover all target nets uniformly.
**F0 risk: not applicable.** F0 is the first fold; no prior fold to inherit from. F1+F2 benefit only.
### Mechanism 4: Adam EMA reset at fold transitions
**Current state:** Adam (m, v) state persists across fold boundaries. Memory file mentions "7 aux Adam resets" — partial coverage already exists for some Adams.
**Proposal:** at fold boundary, zero ALL Adam (m, v) state for: trunk Adam, value Adam, branch Adam, IQN Adam, aux head Adams, ensemble Adams. Comprehensive — no Adam state survives a fold transition.
**Where:** same fold-transition handler as Mechanism 3. Single rich commit per `feedback_no_partial_refactor` since both are fold-boundary-only mechanisms addressing the same target staleness pathway.
**F0 risk: not applicable.** F0 starts with zero Adam state by construction.
### Mechanism 5: Embedded Adam-centric diagnostic (slots 36-47)
The 12 reserved headroom slots from SP1 are populated with sticky-bit threshold checks. Allocation:
| Slot | Signal | Threshold | Source buffer |
|---|---|---|---|
| 36 | Adam m max-abs (trunk) | ≥ 100 × isv[16].max(1) | trunk Adam m |
| 37 | Adam m max-abs (value) | same | value Adam m |
| 38 | Adam m max-abs (branch) | same | branch Adam m |
| 39 | Adam m max-abs (IQN) | same | IQN Adam m |
| 40 | Adam v max-abs (trunk) | ≥ 1e6 × isv[16].max(1)² | trunk Adam v |
| 41 | Adam v max-abs (value) | same | value Adam v |
| 42 | Adam v max-abs (branch) | same | branch Adam v |
| 43 | Adam v max-abs (IQN) | same | IQN Adam v |
| 44 | Weight max-abs (trunk) | ≥ 1000 × isv[16].max(1) | params_buf trunk slice |
| 45 | Weight max-abs (heads) | same | params_buf heads slice |
| 46 | target_q max-abs (post-clip) | ≥ 9.5 × max_abs_target_q | target_q buffer |
| 47 | atom_span max | ≥ 9.5 × max_atom_abs × 2 | atom_positions buffer |
**Implementation:** the SP2 fused kernel is **extended** (NOT replaced) to include these threshold checks. Each kernel block now handles either:
- A NaN-check buffer (slots 24-35; sets flag if `!isfinite(v)`)
- A threshold-check buffer (slots 36-47; sets flag if `fabsf(v) >= threshold`)
Single fused kernel handles all 24 slots in 24 blocks (plus null-skip for deferred slot 31). Slots 36-47's thresholds passed via a `[24]` device array of f32 thresholds (NaN-check entries hold sentinel `+inf` so the check is no-op for them).
**Sticky-flag semantics preserved.** Slots 36-47 fire ONCE when threshold first exceeded; never clear. Same readback infrastructure as slots 24-35 (consumes via existing `read_nan_flags()` returning `[i32; 48]`).
**Diagnostic interpretation if SP3 fix fails:**
- Slots 38-43 fire (Adam EMAs saturated despite reset) → Mech 4 insufficient; need more frequent resets or post-Adam weight clamping
- Slots 44-45 fire but Adam clean → direct weight drift via gradient (not Adam mediation); add post-Adam weight clamp
- Slot 46 (target_q post-clip) NEVER fires across F1+F2 → clip is unnecessary defensive; no action (overshooting protection still useful as sentinel)
- Slot 46 fires often → clip bound too generous; tighten multiplier 10× → 5×
- Slot 47 (atom span) fires often → atom bounds engaging; mechanism doing real work
- Slots 24-35 fire post-fix → SP3 missed a class of pathology; new investigation thread
## 6. ISV slot bindings (summary)
| Mechanism | ISV slot | Index | Bound formula |
|---|---|---|---|
| Mech 1: Target-Q clip | Q_ABS_REF | 16 | `10 × isv.max(1.0)` |
| Mech 2: Atom bounds | Q_ABS_REF | 16 | `10 × isv.max(1.0)` |
| Slots 36-39 (Adam m thresholds) | Q_ABS_REF | 16 | `100 × isv.max(1.0)` |
| Slots 40-43 (Adam v thresholds) | Q_ABS_REF | 16 | `1e6 × isv.max(1.0)²` |
| Slots 44-45 (weight max thresholds) | Q_ABS_REF | 16 | `1000 × isv.max(1.0)` |
| Slot 46 (target_q post-clip) | (derived from Mech 1 bound) | n/a | `9.5 × max_abs_target_q` |
| Slot 47 (atom span) | (derived from Mech 2 bound) | n/a | `9.5 × max_atom_abs × 2` |
**No new ISV slots.** All bounds use existing `Q_ABS_REF_INDEX = 16` per audit's no-new-slot principle.
## 7. Validation gates and error handling
### Gate 1 — SP2 intermediate (mandatory)
After SP2 lands, before SP3 starts. Single L40S smoke at SP2 commit.
**Pass:**
- F0 Best Sharpe ≥ 53.08
- Slot 24-35 firing topology unchanged (slot 26 + 32 still fire at F1 step ~3540 — proves the fused kernel correctly observes NaN entry)
- Cargo check clean, no new errors
**Fail (F0 doesn't recover):** halt before SP3. Re-investigate F0 cause. Possible alternates:
- Buffer-allocation alignment effects from `nan_flags_buf` 24→48 expansion (test: revert size, observe F0)
- Graph-capture topology effects (test: profile graph node count vs F0)
- Stochastic variance (test: 5+ baseline F0 runs at e9096c7be)
### Gate 2 — SP3 full validation
After all 5 SP3 mechanisms land. Full SP1 7-criterion evaluation.
| Failure mode | Diagnostic signal | Response |
|---|---|---|
| F1 still NaNs at slots 26+32 | Sentinel slots 24-35 fire as before | Slot 36-47 diagnostic identifies which stage saturated. Iterate on that mechanism. |
| Slots 38-43 fire (Adam EMAs saturated) | despite Mech 4 fold reset | Increase Adam reset frequency (every-N-steps), OR add post-Adam weight clamp |
| Slots 44-45 fire but Adam clean | direct weight drift from gradient | Add post-Adam weight clamp `1000 × isv[16].max(1)` |
| Slot 46 NEVER fires | target_q clip overshooting | Note in audit doc; clip stays as defense-in-depth; no action |
| Slot 47 fires often | atom bounds engaging frequently | Confirms Mech 2 doing real work; no action |
| F0 regresses after SP3 | F0 drops vs SP2 baseline | ISV bound too tight; widen multiplier (10× → 30×). Iterate. |
| New NaN topology | slots not previously flagged | New pathology class; document, decide whether SP3 needs another mechanism or SP4 territory |
## 8. Operating principles (mandatory)
1. **No deferrals.** If investigation reveals additional anomalies during SP2 or SP3 implementation, fix them within scope, not punted.
2. **Combined RELATED fixes — rich commits.** Per `feedback_no_partial_refactor`. SP2 = one commit. SP3's 5 mechanisms = one commit (related to same root pathology).
3. **Mandatory ISV-driven design.** Every dynamic bound MUST derive from existing ISV slot. ε floors go on the ISV multiplier per the SP1 pearl.
4. **F0 paper-review BEFORE smoke.** Each ISV bound documented with low/medium/high F0 risk + reasoning. High-risk bounds inspect F0-typical input ranges before consuming an L40S smoke.
5. **Permanent diagnostic infrastructure.** Slots 36-47 stay populated as production-grade regression sentinel. Never gated behind a debug flag.
## 9. Anti-patterns explicitly rejected
- **Hardcoded clip bounds.** No `clamp(target_q, -100.0, +100.0)` literals. ISV-driven only.
- **Per-loss-kernel target-Q clipping.** Single-source clip at the target_q buffer is mandatory.
- **Pessimistic ensemble.** Spec lists it as alternative to target-Q clipping but it requires ensemble of target nets — heavy infra. Skip; defer to SP4 if Mech 1 insufficient.
- **Partial mechanism landing.** Per `feedback_no_partial_refactor`, SP3's 5 mechanisms ship in ONE commit. Splitting creates intermediate states with partial structural coverage strictly worse than original.
- **Gating diagnostic behind debug flags.** Slots 36-47 always populate (production-grade sentinel). Never `#[cfg(debug)]` or runtime-gated.
- **New ISV slots without justification.** Reuse `Q_ABS_REF_INDEX = 16` for all SP3 bounds. Adding a new slot requires a per-slot rationale in the audit doc.
## 10. Sequencing constraints
1. SP2 lands first. Gate 1 must pass before SP3 work begins.
2. SP3's 5 mechanisms ship in one commit (per `feedback_no_partial_refactor`).
3. Gate 2 evaluates the full SP1 7-criterion bar on SP3's commit.
4. If Gate 2 fails, slot 36-47 diagnostic readout drives the next iteration. Per no-deferrals, iterate within SP3, do not punt.
5. Branch `plan-c-phase-2-thompson` continues; both SP2 and SP3 commits land on same branch.
## 11. Estimated effort and validation cost
- **SP2:** ~200 LOC, 1 new kernel, 1 Rust wrapper, 1 metadata buffer, 1 L40S smoke. ~1-2 days.
- **SP3:** ~600-800 LOC, 0 new kernels (extends SP2's fused kernel), 4-5 fold-transition wire-ups, 1 atom-update kernel modification, 1 target-Q clip site, 1 L40S smoke. ~3-5 days. Validation: 1-2 smokes (best case 1 if Gate 2 passes; worst case 3-4 if iteration needed via diagnostic).
**Combined budget:** ~1000 LOC across 2 commits, 2-4 L40S smokes, 4-7 days.
## 12. Handoff conditions
After Gate 2 passes:
- Memory entry `project_sp2_sp3_numerical_stability_resolved.md` documenting:
- Final SP2 commit + diagnostic of fix effectiveness (F0 Sharpe restored)
- Final SP3 commit + per-mechanism contribution analysis (which mechanisms triggered slot 36-47 firing — confirms which were necessary)
- Open follow-ups (e.g., periodic Adam reset cadence calibration, atom bound sensitivity)
- Audit doc `docs/dqn-backward-nan-audit.md` updated with Phase E (SP2) + Phase F (SP3) closure sections
- Wire-up audit `docs/dqn-wire-up-audit.md` Invariant 7 entries for both commits
- SP1 closure memory entry `project_sp1_f1_nan_root_cause_resolved.md` updated to mark SP2/SP3 handoff complete
If Gate 2 doesn't pass after 3 SP3 iterations, escalate: pause SP3, re-spec with new pathology class identified by diagnostic readout. New design doc per `superpowers:brainstorming` cycle.
## References
- Predecessor spec: `docs/superpowers/specs/2026-04-29-numerical-stability-investigation-design.md`
- SP1 plan: `docs/superpowers/plans/2026-04-29-numerical-stability-sp1-f1-nan-root-cause.md`
- SP1 closure memory: `~/.claude/projects/-home-jgrusewski-Work-foxhunt/memory/project_sp1_f1_nan_root_cause_resolved.md`
- Audit doc: `docs/dqn-backward-nan-audit.md`
- Wire-up audit: `docs/dqn-wire-up-audit.md`
- Pearls (documented in `project_sp1_f1_nan_root_cause_resolved.md`): (a) ε-floor on ISV multiplier, not bound; (b) NaN diagnostic ordering — inline check before each clamp to preserve sentinel