diff --git a/docs/superpowers/specs/2026-04-29-numerical-stability-sp2-sp3-design.md b/docs/superpowers/specs/2026-04-29-numerical-stability-sp2-sp3-design.md new file mode 100644 index 000000000..d1d0e1e2f --- /dev/null +++ b/docs/superpowers/specs/2026-04-29-numerical-stability-sp2-sp3-design.md @@ -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, +``` + +Sized `12 × (sizeof::() + sizeof::())`. 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