From 1bcc70392041a2bfca48697dee336847cfa2642e Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Thu, 30 Apr 2026 13:59:23 +0200 Subject: [PATCH] =?UTF-8?q?fix(dqn):=20SP3=20close-out=20v2=20=E2=80=94=20?= =?UTF-8?q?Mech=209=20to=20all=205=20Adam=20kernels=20+=20IQN=20weight=20d?= =?UTF-8?q?iag?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Real fix for the F1 step-3540 NaN. Commit 8956c2fb7 wired Mech 9 (post-Adam |p_val| clamp) into ONE Adam kernel (dqn_adam_update_kernel) out of FIVE in the codebase. The slot-26 GEMM (apply_iqn_trunk_gradient output) computes `grad_iqn @ W_iqn^T`; W_iqn lives in IQN's separate param buffer, updated by iqn_adam_kernel — never reached by Mech 9. Slots 44-45 only cover trunk + heads, leaving IQN weights without a diagnostic. The chain: IQN weights drift via iqn_adam_kernel → finite gradient × non-finite weight → slot-26 NaN. Same partial-refactor class as Mech 4's missing GpuAttention reset (caught by B5 audit), corrected the same way: every consumer of the contract migrates together. Extended Mech 9 to all four remaining Adam kernels: - iqn_adam_kernel (iqn_dual_head_kernel.cu) - iql_adam_kernel (iql_value_kernel.cu) - attn_adam_kernel (attention_backward_kernel.cu — used by GpuAttention + GpuTlob) - curiosity_adam_step (curiosity_training_kernel.cu) Same `if (weight_clamp_max_abs > 0.0f) p = fminf(fmaxf(p, -bound), bound)` pattern. Same `100 × Q_ABS_REF.max(1.0)` ISV-driven bound. Each launch site computes the bound host-side and passes it as the trailing kernel arg, mirroring the existing dqn_adam_update_kernel pattern. DT launch keeps the 0.0 disable (offline-RL, outside SP3 scope). Curiosity reads ISV from FusedTrainingCtx in training_loop and threads through the collector wrapper (collector doesn't own ISV). Added IQN-weight diagnostic slot 48: - nan_flags_buf 48 → 49 - Fused-kernel block count 24 → 25; new branch for absolute slot 48 (relative slot 24): threshold = 1e3 × q_abs_ref_eff (matches slot 44-45) - Metadata buffers (nan_check_buf_ptrs/_lens) populate slot 48 with IQN online_params pointer + length (new public accessors on GpuIqnHead) - name table (halt_grad_collapse): 49 entries with "iqn_weight_max" - Audit doc: slot 48 row in Mech 5 table; Mech 9 row updated to span all 5 Adam kernels No new ISV slots. No graph-topology change beyond the +1 block in the already-fused NaN check. cargo check -p ml --lib clean (12 pre-existing warnings, none new). Validation: deferred to one L40S smoke at this HEAD. Expected: F1 trains past step 3720 (Mech-6-only ceiling) and slots 36-43 + new slot 48 stay quiet. Co-Authored-By: Claude Opus 4.7 (1M context) --- .../attention_backward_kernel.cu | 16 ++++- .../curiosity_training_kernel.cu | 14 +++- .../src/cuda_pipeline/dqn_utility_kernels.cu | 17 ++++- crates/ml/src/cuda_pipeline/gpu_attention.rs | 5 ++ .../cuda_pipeline/gpu_curiosity_trainer.rs | 19 ++++-- .../ml/src/cuda_pipeline/gpu_dqn_trainer.rs | 65 +++++++++++++------ .../cuda_pipeline/gpu_experience_collector.rs | 7 ++ .../ml/src/cuda_pipeline/gpu_iql_trainer.rs | 5 ++ crates/ml/src/cuda_pipeline/gpu_iqn_head.rs | 37 ++++++++++- crates/ml/src/cuda_pipeline/gpu_tlob.rs | 13 +++- .../ml/src/cuda_pipeline/iql_value_kernel.cu | 12 +++- .../src/cuda_pipeline/iqn_dual_head_kernel.cu | 20 +++++- crates/ml/src/trainers/dqn/fused_training.rs | 55 ++++++++++++++-- .../src/trainers/dqn/trainer/training_loop.rs | 20 ++++++ docs/dqn-backward-nan-audit.md | 5 +- docs/dqn-wire-up-audit.md | 2 + 16 files changed, 271 insertions(+), 41 deletions(-) diff --git a/crates/ml/src/cuda_pipeline/attention_backward_kernel.cu b/crates/ml/src/cuda_pipeline/attention_backward_kernel.cu index 4ed8bef3d..d6369da3e 100644 --- a/crates/ml/src/cuda_pipeline/attention_backward_kernel.cu +++ b/crates/ml/src/cuda_pipeline/attention_backward_kernel.cu @@ -76,7 +76,12 @@ extern "C" __global__ void attn_adam_kernel( float lr, float beta1, float beta2, float eps, float weight_decay, float max_grad_norm, const int* __restrict__ adam_t_buf, - int total_params + int total_params, + /* SP3 Mech 9 (close-out v2): ISV-driven |p| bound (100 × Q_ABS_REF.max(1.0)) + * — same pattern wired into all five Adam kernels. Used by both + * GpuAttention and GpuTlob (TLOB shares this kernel via the same cubin). + * 0=disabled (debug only). */ + float weight_clamp_max_abs ) { int tid = blockIdx.x * blockDim.x + threadIdx.x; if (tid >= total_params) return; @@ -111,7 +116,14 @@ extern "C" __global__ void attn_adam_kernel( float v_hat = v_val / (bc2 + 1e-12f); /* Parameter update */ - params[tid] = params[tid] - lr * m_hat / (sqrtf(v_hat) + eps); + float p = params[tid] - lr * m_hat / (sqrtf(v_hat) + eps); + + /* SP3 Mech 9: post-Adam |p| clamp — last step before writeback. */ + if (weight_clamp_max_abs > 0.0f) { + p = fminf(fmaxf(p, -weight_clamp_max_abs), weight_clamp_max_abs); + } + + params[tid] = p; /* Zero gradient for next step */ grads[tid] = 0.0f; diff --git a/crates/ml/src/cuda_pipeline/curiosity_training_kernel.cu b/crates/ml/src/cuda_pipeline/curiosity_training_kernel.cu index ca7ef22c3..a4f78cbae 100644 --- a/crates/ml/src/cuda_pipeline/curiosity_training_kernel.cu +++ b/crates/ml/src/cuda_pipeline/curiosity_training_kernel.cu @@ -320,7 +320,10 @@ extern "C" __global__ void curiosity_adam_step( float beta1, float beta2, float eps, - int step /* 1-based step counter */ + int step, /* 1-based step counter */ + /* SP3 Mech 9 (close-out v2): ISV-driven |p| bound (100 × Q_ABS_REF.max(1.0)) + * — same pattern wired into all five Adam kernels. 0=disabled (debug only). */ + float weight_clamp_max_abs ) { int i = blockIdx.x * blockDim.x + threadIdx.x; if (i >= num_params) return; @@ -340,7 +343,14 @@ extern "C" __global__ void curiosity_adam_step( v[i] = beta2_bf * v[i] + (one_bf - beta2_bf) * g * g; float m_hat = m[i] / (1.0f - powf(beta1, (float)step)); float v_hat = v[i] / (1.0f - powf(beta2, (float)step)); - params[i] = params[i] - lr_bf * m_hat / (sqrtf(v_hat) + eps_bf); + float p = params[i] - lr_bf * m_hat / (sqrtf(v_hat) + eps_bf); + + /* SP3 Mech 9: post-Adam |p| clamp — last step before writeback. */ + if (weight_clamp_max_abs > 0.0f) { + p = fminf(fmaxf(p, -weight_clamp_max_abs), weight_clamp_max_abs); + } + + params[i] = p; } /* ------------------------------------------------------------------ */ diff --git a/crates/ml/src/cuda_pipeline/dqn_utility_kernels.cu b/crates/ml/src/cuda_pipeline/dqn_utility_kernels.cu index 4b0ba2edc..1d2a8ead6 100644 --- a/crates/ml/src/cuda_pipeline/dqn_utility_kernels.cu +++ b/crates/ml/src/cuda_pipeline/dqn_utility_kernels.cu @@ -1741,6 +1741,13 @@ extern "C" __global__ void popart_normalize_robust( // atom-positions buffers. Threshold is computed inline per-slot from a single // `q_abs_ref_eff` arg (no per-slot threshold buffer, no per-step HtoD copies). // +// SP3 close-out v2 extension: grid extended 24 → 25 blocks. Relative slot 24 +// (= absolute 48) covers IQN's separate online_params buffer — the F1 step- +// 3540 NaN was traced to this buffer drifting through `iqn_adam_kernel` (its +// own param + Adam state, never reached by the single-kernel Mech 9 in +// `dqn_adam_update_kernel`). Same `1e3 × q_abs_ref_eff` weight-slice +// threshold as slots 44-45 — IQN online_params share the trunk weight regime. +// // Slot-index → threshold mapping (relative slot ≥ 12; absolute = slot + 24): // 12-15 (abs 36-39): Adam m max-abs ≥ 100 × q_abs_ref_eff // 16-19 (abs 40-43): Adam v max-abs ≥ 1e6 × q_abs_ref_eff² @@ -1751,6 +1758,8 @@ extern "C" __global__ void popart_normalize_robust( // 23 (abs 47) : atom-span max ≥ 190 × q_abs_ref_eff // (= 9.5 × max_atom_abs × 2, max_atom_abs = 10 × // q_abs_ref_eff) +// 24 (abs 48) : IQN online_params max-abs ≥ 1e3 × q_abs_ref_eff +// (SP3 close-out v2 — slot-26 GEMM regression sentinel) // // q_abs_ref_eff = max(ISV[Q_ABS_REF=16], 1.0) computed host-side at launch, // passed by-value as a float kernel arg. ε on multiplier per SP1 pearl: cold- @@ -1763,7 +1772,7 @@ extern "C" __global__ void popart_normalize_robust( // - graph-capture safe (pure kernel launch, no host ops) // ============================================================================ extern "C" __global__ void dqn_nan_check_fused_f32_kernel( - const float* const* buf_ptrs, // [N_SLOTS] (24 entries: 12 NaN-only + 12 threshold) + const float* const* buf_ptrs, // [N_SLOTS] (25 entries: 12 NaN-only + 13 threshold) const int* buf_lens, // [N_SLOTS] float q_abs_ref_eff, // SP3 Mech 5: ISV-derived multiplier (slots ≥ 12) int base_flag_idx, // first slot index (24 for backward checks) @@ -1776,7 +1785,7 @@ extern "C" __global__ void dqn_nan_check_fused_f32_kernel( // SP3 Mech 5: per-slot threshold by relative slot index. Slots 0-11 // (= absolute 24-35) are NaN-only checks (threshold = +inf so the - // magnitude check is trivially false). Slots 12-23 (= absolute 36-47) + // magnitude check is trivially false). Slots 12-24 (= absolute 36-48) // check magnitude exceedance per the audit table above. All thresholds // derive from q_abs_ref_eff = max(isv[Q_ABS_REF=16], 1.0). float thr = INFINITY; @@ -1795,6 +1804,10 @@ extern "C" __global__ void dqn_nan_check_fused_f32_kernel( } else if (slot == 23) { // slot 47: atom-positions span (max_atom_abs = 10 × q_abs_ref_eff) thr = 190.0f * q_abs_ref_eff; + } else if (slot == 24) { + // slot 48 (SP3 close-out v2): IQN online_params — same weight-slice + // regime as slots 44-45. Captures the slot-26 GEMM blind spot. + thr = 1.0e3f * q_abs_ref_eff; } int flag_set = 0; diff --git a/crates/ml/src/cuda_pipeline/gpu_attention.rs b/crates/ml/src/cuda_pipeline/gpu_attention.rs index 6f2418010..891eb7b86 100644 --- a/crates/ml/src/cuda_pipeline/gpu_attention.rs +++ b/crates/ml/src/cuda_pipeline/gpu_attention.rs @@ -975,6 +975,10 @@ impl GpuAttention { max_grad_norm: f32, override_stream: Option<&Arc>, override_workspace: Option<(u64, usize)>, + /* SP3 Mech 9 (close-out v2): ISV-driven |p| bound for `attn_adam_kernel` + * — caller computes `100 × ISV[Q_ABS_REF].max(1.0)` host-side. + * 0=disabled (debug only). */ + weight_clamp_max_abs: f32, ) -> Result<(), MLError> { let stream = override_stream.unwrap_or(&self.stream); // override_workspace is accepted for API symmetry; adam_step does not use @@ -1044,6 +1048,7 @@ impl GpuAttention { .arg(&max_grad_norm) .arg(&t_ptr) .arg(&tp_i32) + .arg(&weight_clamp_max_abs) // SP3 Mech 9 (close-out v2): post-Adam |p| clamp .launch(adam_cfg) .map_err(|e| MLError::ModelError(format!("attention adam kernel: {e}")))?; } diff --git a/crates/ml/src/cuda_pipeline/gpu_curiosity_trainer.rs b/crates/ml/src/cuda_pipeline/gpu_curiosity_trainer.rs index 414581da4..4677f59fb 100644 --- a/crates/ml/src/cuda_pipeline/gpu_curiosity_trainer.rs +++ b/crates/ml/src/cuda_pipeline/gpu_curiosity_trainer.rs @@ -497,6 +497,11 @@ pub struct GpuCuriosityTrainer { } /// Launch Adam optimizer step for one parameter group. +/// +/// `weight_clamp_max_abs` is SP3 Mech 9 (close-out v2): ISV-driven |p| +/// bound for `curiosity_adam_step`. Caller threads `100 × Q_ABS_REF.max(1.0)` +/// from training_loop down through `train_on_collector_buffers`. 0=disabled +/// (debug only). fn launch_adam_step( stream: &CudaStream, adam_func: &CudaFunction, @@ -507,6 +512,7 @@ fn launch_adam_step( v: &mut CudaSlice, batch_size: usize, step: i32, + weight_clamp_max_abs: f32, ) -> Result<(), MLError> { let cfg = LaunchConfig { grid_dim: (((num_params as u32) + 255) / 256, 1, 1), @@ -529,6 +535,7 @@ fn launch_adam_step( .arg(&ADAM_BETA2) .arg(&ADAM_EPS) .arg(&step) + .arg(&weight_clamp_max_abs) // SP3 Mech 9 (close-out v2): post-Adam |p| clamp .launch(cfg) .map_err(|e| { MLError::ModelError(format!("curiosity_adam_step launch: {e}")) @@ -693,6 +700,10 @@ impl GpuCuriosityTrainer { states: &CudaSlice, actions: &CudaSlice, n_samples: usize, + /* SP3 Mech 9 (close-out v2): ISV-driven |p| bound for the four + * `curiosity_adam_step` launches (W1, b1, W2, b2). Caller computes + * `100 × ISV[Q_ABS_REF].max(1.0)` host-side. 0=disabled (debug only). */ + weight_clamp_max_abs: f32, ) -> Result<(), MLError> { if n_samples < 2 { return Ok(()); @@ -909,10 +920,10 @@ impl GpuCuriosityTrainer { self.step += 1; let step = self.step; - launch_adam_step(stream, &self.adam_func, &mut weights.w1, &self.grad_w1, CUR_W1_LEN, &mut self.adam_m_w1, &mut self.adam_v_w1, n_train, step)?; - launch_adam_step(stream, &self.adam_func, &mut weights.b1, &self.grad_b1, CUR_B1_LEN, &mut self.adam_m_b1, &mut self.adam_v_b1, n_train, step)?; - launch_adam_step(stream, &self.adam_func, &mut weights.w2, &self.grad_w2, CUR_W2_LEN, &mut self.adam_m_w2, &mut self.adam_v_w2, n_train, step)?; - launch_adam_step(stream, &self.adam_func, &mut weights.b2, &self.grad_b2, CUR_B2_LEN, &mut self.adam_m_b2, &mut self.adam_v_b2, n_train, step)?; + launch_adam_step(stream, &self.adam_func, &mut weights.w1, &self.grad_w1, CUR_W1_LEN, &mut self.adam_m_w1, &mut self.adam_v_w1, n_train, step, weight_clamp_max_abs)?; + launch_adam_step(stream, &self.adam_func, &mut weights.b1, &self.grad_b1, CUR_B1_LEN, &mut self.adam_m_b1, &mut self.adam_v_b1, n_train, step, weight_clamp_max_abs)?; + launch_adam_step(stream, &self.adam_func, &mut weights.w2, &self.grad_w2, CUR_W2_LEN, &mut self.adam_m_w2, &mut self.adam_v_w2, n_train, step, weight_clamp_max_abs)?; + launch_adam_step(stream, &self.adam_func, &mut weights.b2, &self.grad_b2, CUR_B2_LEN, &mut self.adam_m_b2, &mut self.adam_v_b2, n_train, step, weight_clamp_max_abs)?; debug!(step, n_train, "curiosity GPU training step complete (cuBLAS GEMM pipeline)"); diff --git a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs index 472f72f13..2008609e5 100644 --- a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs +++ b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs @@ -11476,15 +11476,19 @@ impl GpuDqnTrainer { // SP1 Phase B: expanded 24 → 48 to make room for backward-path NaN // checks (slots 24-35) plus 12 reserved headroom slots (36-47) for - // SP2 framework checker hooks + SP3 structural observers. Buffer - // size reviewable by SP2 framework codification. - let nan_flags_buf = stream.alloc_zeros::(48) + // SP2 framework checker hooks + SP3 structural observers. SP3 close- + // out v2 (slot 48): one additional slot for IQN online_params (the + // separate weight buffer that drives the slot-26 cuBLAS GEMM and was + // the F1 step-3540 NaN root). Buffer size reviewable by SP2 framework + // codification. + let nan_flags_buf = stream.alloc_zeros::(49) .map_err(|e| MLError::ModelError(format!("nan_flags alloc: {e}")))?; // SP2 + SP3 Mech 5: fused NaN/threshold-check (ptr, len) tables. - // 24 entries each — slots 0-11 (= absolute flag idx 24-35) NaN-only + // 25 entries each — slots 0-11 (= absolute flag idx 24-35) NaN-only // (SP2 Task A3) + slots 12-23 (= absolute 36-47) ISV-threshold - // (SP3 Mech 5 Task B6). Allocated as mapped-pinned + // (SP3 Mech 5 Task B6) + slot 24 (= absolute 48) IQN online_params + // (SP3 close-out v2). Allocated as mapped-pinned // (`cuMemHostAlloc(DEVICEMAP|PORTABLE)`) per // `feedback_no_htod_htoh_only_mapped_pinned` — host-side writes are // visible to the kernel through the device-mapped pointer with zero @@ -11497,9 +11501,9 @@ impl GpuDqnTrainer { // Safety: a CUDA context is active on this thread (the GpuDqnTrainer // constructor runs within `FusedTrainingCtx::new` which has already // initialised CUDA via the shared cublas handle). - let nan_check_buf_ptrs = unsafe { MappedU64Buffer::new(24) } + let nan_check_buf_ptrs = unsafe { MappedU64Buffer::new(25) } .map_err(|e| MLError::ModelError(format!("nan_check_buf_ptrs alloc: {e}")))?; - let nan_check_buf_lens = unsafe { MappedI32Buffer::new(24) } + let nan_check_buf_lens = unsafe { MappedI32Buffer::new(25) } .map_err(|e| MLError::ModelError(format!("nan_check_buf_lens alloc: {e}")))?; // v8: PopArt running statistics buffers @@ -15368,6 +15372,12 @@ impl GpuDqnTrainer { /// 22 (slot 46) target_q — post-clip ≥ 95 × q_abs_ref_eff /// 23 (slot 47) atom_positions — span max ≥ 190 × q_abs_ref_eff /// + /// SP3 close-out v2 — slot 48 (IQN online_params): + /// 24 (slot 48) iqn_online_params — weight max-abs ≥ 1e3 × q_abs_ref_eff + /// (same regime as slots 44-45; + /// captures the slot-26 GEMM blind + /// spot that hid the F1 step-3540 NaN) + /// /// Mapped-pinned host write — no HtoD copy. The kernel reads the same /// memory through the device-mapped pointer (`dev_ptr`) once the launch /// stream barrier ensures coherence. @@ -15386,6 +15396,13 @@ impl GpuDqnTrainer { iqn_adam_m_len: Option, iqn_adam_v_ptr: Option, iqn_adam_v_len: Option, + // SP3 close-out v2 addition: IQN online_params ptr + len for slot 48 + // (the diagnostic blind spot that hid the F1 step-3540 NaN — IQN's + // separate weight buffer drives the slot-26 cuBLAS GEMM in + // `apply_iqn_trunk_gradient`, never reached by slots 44-45). Same + // None=inactive convention as the slot 27/28/39/43 entries. + iqn_online_params_ptr: Option, + iqn_online_params_len: Option, ) -> Result<(), MLError> { // Slot-28 length expression: `tba × b × q` per // `run_nan_checks_post_backward` slot 28. tba = b0+b1+b2+b3. @@ -15419,7 +15436,7 @@ impl GpuDqnTrainer { let atom_positions_ptr = self.atom_positions_ptr(); let atom_positions_len = self.atom_positions_len() as i32; - let entries: [(u64, i32); 24] = [ + let entries: [(u64, i32); 25] = [ // SP2 Task A3 — slot 24 — post-c51_grad value gradient (self.d_value_logits_buf_ptr(), self.d_value_logits_buf.len() as i32), // slot 25 — post-c51_grad branch advantage @@ -15483,28 +15500,37 @@ impl GpuDqnTrainer { (target_q_ptr, target_q_len), // slot 47 — atom_positions (atom_positions_ptr, atom_positions_len), + // SP3 close-out v2 — slot 48 — IQN online_params (caller-supplied; + // 0 when IQN inactive — null-pointer guard skips block). + ( + iqn_online_params_ptr.unwrap_or(0), + iqn_online_params_len.unwrap_or(0) as i32, + ), ]; // Host-side write through the mapped-pinned pages — kernel sees the // values via dev_ptr after the next stream-sync barrier (no HtoD copy // required, per `feedback_no_htod_htoh_only_mapped_pinned`). - let ptrs_view: [u64; 24] = std::array::from_fn(|i| entries[i].0); - let lens_view: [i32; 24] = std::array::from_fn(|i| entries[i].1); + let ptrs_view: [u64; 25] = std::array::from_fn(|i| entries[i].0); + let lens_view: [i32; 25] = std::array::from_fn(|i| entries[i].1); self.nan_check_buf_ptrs.write_from_slice(&ptrs_view); self.nan_check_buf_lens.write_from_slice(&lens_view); Ok(()) } - /// SP2 + SP3 Mech 5: launch the fused NaN/threshold-check kernel. - /// Single launch processes all 24 backward-path slots (24-47) in 24 + /// SP2 + SP3 Mech 5 + SP3 close-out v2: launch the fused + /// NaN/threshold-check kernel. + /// Single launch processes all 25 backward-path slots (24-48) in 25 /// blocks, in parallel. Slots with null pointer (deferred slot 31, - /// optional IQN slots 27/28/39/43 when inactive, slots 33-35 inline + /// optional IQN slots 27/28/39/43/48 when inactive, slots 33-35 inline /// elsewhere) no-op via the kernel's null-pointer guard. Sticky-flag /// semantics preserved (kernel only writes 1, never clears). /// /// Slots 0-11 (= absolute 24-35) are NaN-only checks (threshold = +inf /// trivially false). Slots 12-23 (= absolute 36-47) are SP3 Mech 5 - /// ISV-threshold checks: threshold computed inline per-slot from a single + /// ISV-threshold checks. Slot 24 (= absolute 48) is SP3 close-out v2 — + /// IQN online_params, threshold = 1e3 × q_abs_ref_eff (same regime as + /// slots 44-45). Threshold computed inline per-slot from a single /// `q_abs_ref_eff` arg (no per-slot threshold buffer, no per-step HtoD). /// /// Replaces the 8-launch sequence in `run_nan_checks_post_backward`. @@ -15512,7 +15538,7 @@ impl GpuDqnTrainer { let nan_flags_ptr = self.nan_flags_buf.raw_ptr(); let buf_ptrs_dev = self.nan_check_buf_ptrs.dev_ptr; let buf_lens_dev = self.nan_check_buf_lens.dev_ptr; - // SP3 Mech 5: read ISV[Q_ABS_REF=16] for slot 36-47 threshold + // SP3 Mech 5: read ISV[Q_ABS_REF=16] for slot 36-48 threshold // computation. Inline ε on multiplier (`max(.., 1.0)`) per the SP1 // pearl — cold-start ISV (~0) gives q_abs_ref_eff = 1, so all // thresholds floor at their ε-floor multiplier × 1. Pass as f32 to @@ -15520,7 +15546,7 @@ impl GpuDqnTrainer { // bug pearl: passing f64 to a `float` param reads only low 4 bytes). let q_abs_ref = self.read_isv_signal_at(Q_ABS_REF_INDEX); let q_abs_ref_eff: f32 = q_abs_ref.max(1.0_f32); - const BLOCKS: u32 = 24; + const BLOCKS: u32 = 25; const THREADS: u32 = 256; const BASE_FLAG_IDX: i32 = 24; let cfg = LaunchConfig { @@ -15596,9 +15622,10 @@ impl GpuDqnTrainer { /// Read NaN flags back to CPU (synchronizes stream). Returns [48] flags. /// Index → buffer mapping is documented in `run_nan_checks_post_forward`. /// SP1 Phase B: slots 24-35 reserved for backward-path checks (wired by - /// Task 4); slots 36-47 reserved for SP2/SP3 headroom. - pub fn read_nan_flags(&self) -> Result<[i32; 48], MLError> { - let mut host = [0_i32; 48]; + /// Task 4); slots 36-47 SP3 Mech 5 ISV-threshold checks; slot 48 SP3 + /// close-out v2 IQN online_params weight diagnostic. + pub fn read_nan_flags(&self) -> Result<[i32; 49], MLError> { + let mut host = [0_i32; 49]; self.stream.memcpy_dtoh(&self.nan_flags_buf, &mut host) .map_err(|e| MLError::ModelError(format!("nan_flags read: {e}")))?; Ok(host) diff --git a/crates/ml/src/cuda_pipeline/gpu_experience_collector.rs b/crates/ml/src/cuda_pipeline/gpu_experience_collector.rs index 4eddae72f..625b82563 100644 --- a/crates/ml/src/cuda_pipeline/gpu_experience_collector.rs +++ b/crates/ml/src/cuda_pipeline/gpu_experience_collector.rs @@ -4233,10 +4233,16 @@ impl GpuExperienceCollector { } /// Train curiosity forward model directly on GPU using experience data. + /// + /// `weight_clamp_max_abs` is SP3 Mech 9 (close-out v2): ISV-driven |p| + /// bound for `curiosity_adam_step`. The collector doesn't own ISV — the + /// caller (training_loop) reads `100 × ISV[Q_ABS_REF].max(1.0)` from the + /// FusedTrainingCtx and threads it through here. 0=disabled (debug only). pub fn train_curiosity_gpu( &mut self, n_episodes: usize, timesteps: usize, + weight_clamp_max_abs: f32, ) -> Result<(), MLError> { if let Some(ref mut trainer) = self.curiosity_trainer { let total = n_episodes.saturating_mul(timesteps); @@ -4248,6 +4254,7 @@ impl GpuExperienceCollector { &self.states_out, &self.actions_out, total, + weight_clamp_max_abs, )?; } Ok(()) diff --git a/crates/ml/src/cuda_pipeline/gpu_iql_trainer.rs b/crates/ml/src/cuda_pipeline/gpu_iql_trainer.rs index ba544b41d..9e2336876 100644 --- a/crates/ml/src/cuda_pipeline/gpu_iql_trainer.rs +++ b/crates/ml/src/cuda_pipeline/gpu_iql_trainer.rs @@ -524,6 +524,10 @@ impl GpuIqlTrainer { pub fn train_value_step( &mut self, states_f32: &CudaSlice, + /* SP3 Mech 9 (close-out v2): ISV-driven |p| bound for `iql_adam_kernel` + * — caller computes `100 × ISV[Q_ABS_REF].max(1.0)` host-side. + * 0=disabled (debug only). */ + weight_clamp_max_abs: f32, ) -> Result<(), MLError> { let b = self.config.batch_size; let h = self.config.value_hidden_dim; @@ -995,6 +999,7 @@ impl GpuIqlTrainer { .arg(&mgn) .arg(&t_ptr) .arg(&total_params_i32) + .arg(&weight_clamp_max_abs) // SP3 Mech 9 (close-out v2): post-Adam |p| clamp .launch(LaunchConfig { grid_dim: (adam_blocks as u32, 1, 1), block_dim: (256, 1, 1), diff --git a/crates/ml/src/cuda_pipeline/gpu_iqn_head.rs b/crates/ml/src/cuda_pipeline/gpu_iqn_head.rs index d7ab15213..837d2c077 100644 --- a/crates/ml/src/cuda_pipeline/gpu_iqn_head.rs +++ b/crates/ml/src/cuda_pipeline/gpu_iqn_head.rs @@ -803,12 +803,24 @@ impl GpuIqnHead { dqn_actions_buf: &CudaSlice, dqn_rewards_buf: &CudaSlice, dqn_dones_buf: &CudaSlice, + /* SP3 Mech 9 (close-out v2): ISV-derived |p| bound for `iqn_adam_kernel` + * — caller computes `100 × ISV[Q_ABS_REF].max(1.0)` and passes it + * through. 0=disabled (debug only); production callers always set + * a non-zero bound. */ + weight_clamp_max_abs: f32, ) -> Result { // 1-2. Decode actions + DtoD copy rewards/dones self.prepare_buffers(dqn_actions_buf, dqn_rewards_buf, dqn_dones_buf)?; // 3-9. Execute core training pipeline - self.execute_training_pipeline(online_h_s2, next_states_buf, target_dueling, None, None) + self.execute_training_pipeline( + online_h_s2, + next_states_buf, + target_dueling, + None, + None, + weight_clamp_max_abs, + ) } /// Set a pre-computed target h_s2 buffer from the DQN trainer's Pass 2. @@ -832,6 +844,10 @@ impl GpuIqnHead { target_dueling: &DuelingWeightSet, override_stream: Option<&Arc>, override_workspace: Option<(u64, usize)>, + /* SP3 Mech 9 (close-out v2): ISV-derived |p| bound for the trailing + * `iqn_adam_kernel` arg. Caller computes `100 × ISV[Q_ABS_REF].max(1.0)` + * host-side. 0=disabled (debug only). */ + weight_clamp_max_abs: f32, ) -> Result { let effective_stream = override_stream.unwrap_or(&self.stream); let (lt_ws_ptr, lt_ws_size) = override_workspace @@ -1482,6 +1498,7 @@ impl GpuIqnHead { .arg(&b1_i32) .arg(&b2_i32) .arg(&b3_i32) + .arg(&weight_clamp_max_abs) // SP3 Mech 9 (close-out v2): post-Adam |p| clamp .launch(adam_config) .map_err(|e| MLError::ModelError(format!("IQN adam kernel: {e}")))?; } @@ -1655,6 +1672,24 @@ impl GpuIqnHead { self.v_buf.len() } + /// SP3 close-out v2 / slot 48: raw device pointer to the IQN online + /// parameter buffer. Companion to `adam_m_ptr` / `adam_v_ptr` — + /// captures the IQN-internal weights that drive the slot-26 cuBLAS + /// GEMM (`apply_iqn_trunk_gradient`). Wired into the fused + /// threshold-check kernel as the diagnostic blind spot identified in + /// the F1 step-3540 NaN trace: slots 44-45 only cover the main DQN + /// trunk + heads, leaving IQN's separate param buffer unmonitored. + pub fn online_params_ptr(&self) -> u64 { + self.online_params.raw_ptr() + } + + /// Element count of the IQN online parameter buffer. + /// Length = `total_params + cublas_pad` (last tensor can be overread + /// by 32-element cuBLAS tiles — see `init_iqn_xavier_weights`). + pub fn online_params_len(&self) -> usize { + self.online_params.len() + } + /// Write tau into the pinned+device-mapped page for graph replay. /// The EMA kernel reads via `tau_dev_ptr` which aliases the same physical page. pub fn set_tau_host(&mut self, tau: f32) { diff --git a/crates/ml/src/cuda_pipeline/gpu_tlob.rs b/crates/ml/src/cuda_pipeline/gpu_tlob.rs index 9b0617d8c..476baf1d3 100644 --- a/crates/ml/src/cuda_pipeline/gpu_tlob.rs +++ b/crates/ml/src/cuda_pipeline/gpu_tlob.rs @@ -646,7 +646,17 @@ impl GpuTlob { } /// Adam step: grad norm clip + Adam update on TLOB params. - pub(crate) fn adam_step(&mut self, lr: f32, max_grad_norm: f32) -> Result<(), MLError> { + /// + /// SP3 Mech 9 (close-out v2): `weight_clamp_max_abs` is the ISV-driven + /// |p| bound for `attn_adam_kernel` (TLOB shares this kernel via the + /// attention cubin). Caller computes `100 × ISV[Q_ABS_REF].max(1.0)`. + /// 0=disabled (debug only). + pub(crate) fn adam_step( + &mut self, + lr: f32, + max_grad_norm: f32, + weight_clamp_max_abs: f32, + ) -> Result<(), MLError> { let tp = TLOB_TOTAL_PARAMS; let tp_i32 = tp as i32; @@ -703,6 +713,7 @@ impl GpuTlob { .arg(&max_grad_norm) .arg(&t_ptr) .arg(&tp_i32) + .arg(&weight_clamp_max_abs) // SP3 Mech 9 (close-out v2): post-Adam |p| clamp .launch(LaunchConfig { grid_dim: (adam_blocks as u32, 1, 1), block_dim: (256, 1, 1), diff --git a/crates/ml/src/cuda_pipeline/iql_value_kernel.cu b/crates/ml/src/cuda_pipeline/iql_value_kernel.cu index f91727497..76c43e950 100644 --- a/crates/ml/src/cuda_pipeline/iql_value_kernel.cu +++ b/crates/ml/src/cuda_pipeline/iql_value_kernel.cu @@ -200,7 +200,10 @@ void iql_adam_kernel( float weight_decay, float max_grad_norm, const int* __restrict__ t_buf, /* [1] Adam step on device */ - int total_params + int total_params, + /* SP3 Mech 9 (close-out v2): ISV-driven |p| bound (100 × Q_ABS_REF.max(1.0)) + * — same pattern wired into all five Adam kernels. 0=disabled (debug only). */ + float weight_clamp_max_abs ) { int i = blockIdx.x * blockDim.x + threadIdx.x; @@ -227,6 +230,13 @@ void iql_adam_kernel( /* Parameter update */ p -= lr * m_hat / (sqrtf(v_hat) + eps); + /* SP3 Mech 9: post-Adam |p| clamp. Last step before writeback — + * applies after all decay/decoupled updates so the bound holds for + * the value the next forward sees. */ + if (weight_clamp_max_abs > 0.0f) { + p = fminf(fmaxf(p, -weight_clamp_max_abs), weight_clamp_max_abs); + } + params[i] = p; m[i] = m_i; v[i] = v_i; diff --git a/crates/ml/src/cuda_pipeline/iqn_dual_head_kernel.cu b/crates/ml/src/cuda_pipeline/iqn_dual_head_kernel.cu index 08a5e9aed..f47f88ed1 100644 --- a/crates/ml/src/cuda_pipeline/iqn_dual_head_kernel.cu +++ b/crates/ml/src/cuda_pipeline/iqn_dual_head_kernel.cu @@ -844,6 +844,14 @@ void iqn_grad_norm_phase2( } /* ── Adam Update Kernel ─────────────────────────────────────────────── */ +/* SP3 Mech 9 (close-out v2): trailing `weight_clamp_max_abs` arg matches the + * pattern wired into all five Adam kernels (dqn/iqn/iql/attn/curiosity). + * Bound = 100 × ISV[Q_ABS_REF].max(1.0) computed host-side. Targets the + * IQN-internal weight buffer that the slot-26 cuBLAS GEMM reads — + * `apply_iqn_trunk_gradient` saw NaN at slot 26 with finite inputs because + * `online_params` here drifted via this Adam, never reached by the + * single-kernel Mech 9 in `dqn_adam_update_kernel`. Disabled when bound==0. + */ extern "C" __global__ void iqn_adam_kernel( float* __restrict__ params, @@ -861,7 +869,8 @@ void iqn_adam_kernel( int b0_size, /* runtime: unused, for consistent interface */ int b1_size, /* runtime: unused, for consistent interface */ int b2_size, /* runtime: unused, for consistent interface */ - int b3_size /* runtime: unused, for consistent interface */ + int b3_size, /* runtime: unused, for consistent interface */ + float weight_clamp_max_abs /* SP3 Mech 9: ISV-driven |p| bound (100 × Q_ABS_REF.max(1.0)). 0=disabled. */ ) { int tid = blockIdx.x * blockDim.x + threadIdx.x; @@ -892,6 +901,15 @@ void iqn_adam_kernel( /* Parameter update with decoupled weight decay */ float p = params[tid]; p = p - lr * (m_hat / (sqrtf(v_hat) + eps) + weight_decay * p); + + /* SP3 Mech 9: post-Adam |p| clamp. Bound is host-supplied + * (100 × Q_ABS_REF.max(1.0)). One OOM below slot 48 diagnostic + * threshold (1e3 × Q_ABS_REF) so the regression sentinel keeps + * 10× firing headroom. */ + if (weight_clamp_max_abs > 0.0f) { + p = fminf(fmaxf(p, -weight_clamp_max_abs), weight_clamp_max_abs); + } + params[tid] = p; /* Zero gradient for next step */ diff --git a/crates/ml/src/trainers/dqn/fused_training.rs b/crates/ml/src/trainers/dqn/fused_training.rs index 14a5d36e9..798f0bc05 100644 --- a/crates/ml/src/trainers/dqn/fused_training.rs +++ b/crates/ml/src/trainers/dqn/fused_training.rs @@ -37,7 +37,8 @@ use crate::cuda_pipeline::mapped_pinned::{MappedF32Buffer, MappedI32Buffer}; use crate::cuda_pipeline::gpu_attention::{GpuAttention, GpuAttentionConfig}; use crate::cuda_pipeline::gpu_tlob::GpuTlob; use crate::cuda_pipeline::gpu_dqn_trainer::{ - GpuDqnTrainConfig, GpuDqnTrainer, TAU_EFF_INDEX, TLOB_REGIME_FOCUS_EMA_INDEX, + GpuDqnTrainConfig, GpuDqnTrainer, Q_ABS_REF_INDEX, TAU_EFF_INDEX, + TLOB_REGIME_FOCUS_EMA_INDEX, }; use crate::cuda_pipeline::gpu_her::{GpuHer, GpuHerConfig, parse_her_strategy}; use crate::cuda_pipeline::gpu_iql_trainer::{GpuIqlConfig, GpuIqlTrainer}; @@ -678,6 +679,11 @@ impl FusedTrainingCtx { gpu_iqn.as_ref().map(|h| h.adam_m_len()), gpu_iqn.as_ref().map(|h| h.adam_v_ptr()), gpu_iqn.as_ref().map(|h| h.adam_v_len()), + // SP3 close-out v2: IQN online_params ptr+len for slot 48 — the + // diagnostic blind spot that hid the F1 step-3540 NaN. Same + // None=inactive convention. + gpu_iqn.as_ref().map(|h| h.online_params_ptr()), + gpu_iqn.as_ref().map(|h| h.online_params_len()), ).map_err(|e| anyhow::anyhow!("populate_nan_check_meta: {e}"))?; // GPU attention — always active. 4-head self-attention on h_s2. @@ -1609,10 +1615,15 @@ impl FusedTrainingCtx { self.gpu_tlob .backward(d_concat_ref, self.batch_size, concat_dim, tlob_concat_off) .map_err(|e| anyhow::anyhow!("TLOB backward: {e}"))?; + // SP3 Mech 9 (close-out v2): TLOB shares `attn_adam_kernel` + // (cubin-shared with GpuAttention). Same ISV-derived bound. + let q_abs_ref_eff_tlob = self.trainer.read_isv_signal_at(Q_ABS_REF_INDEX).max(1.0_f32); + let tlob_weight_clamp_max_abs: f32 = 100.0_f32 * q_abs_ref_eff_tlob; self.gpu_tlob .adam_step( self.trainer.config().lr, self.trainer.config().max_grad_norm, + tlob_weight_clamp_max_abs, ) .map_err(|e| anyhow::anyhow!("TLOB Adam: {e}"))?; } @@ -1829,11 +1840,19 @@ impl FusedTrainingCtx { .map_err(|e| anyhow::anyhow!("IQL low gather: {e}"))?; // 3. Train V(s) — both tau=high and tau=low. + // + // SP3 Mech 9 (close-out v2): ISV-driven post-Adam |p| clamp for + // `iql_adam_kernel`. Same `100 × Q_ABS_REF.max(1.0)` bound as + // the main DQN Adam — IQL's separate value-net params drift via + // its own Adam state and were never reached by the single-kernel + // Mech 9 in commit 8956c2fb7. ε on multiplier per SP1 pearl. + let q_abs_ref_eff = self.trainer.read_isv_signal_at(Q_ABS_REF_INDEX).max(1.0_f32); + let weight_clamp_max_abs: f32 = 100.0_f32 * q_abs_ref_eff; let states_f32 = self.trainer.states_buf(); - self.gpu_iql.train_value_step(states_f32) + self.gpu_iql.train_value_step(states_f32, weight_clamp_max_abs) .map_err(|e| anyhow::anyhow!("IQL high train: {e}"))?; let states_f32 = self.trainer.states_buf(); - self.gpu_iql_low.train_value_step(states_f32) + self.gpu_iql_low.train_value_step(states_f32, weight_clamp_max_abs) .map_err(|e| anyhow::anyhow!("IQL low train: {e}"))?; // 4. Compute advantage weights (from high-tau V(s)). @@ -1925,9 +1944,19 @@ impl FusedTrainingCtx { let online_h_s2 = self.trainer.save_h_s2(); let next_states_buf = self.trainer.next_states_buf(); + // SP3 Mech 9 (close-out v2): ISV-driven post-Adam |p| clamp + // for `iqn_adam_kernel`. Same `100 × Q_ABS_REF.max(1.0)` + // bound as the main DQN Mech 9 launch (`launch_adam_update`) + // — IQN's separate online_params drives the slot-26 GEMM and + // was never reached by the single-kernel Mech 9 in commit + // 8956c2fb7. ε on multiplier per SP1 pearl. + let q_abs_ref_eff = self.trainer.read_isv_signal_at(Q_ABS_REF_INDEX).max(1.0_f32); + let weight_clamp_max_abs: f32 = 100.0_f32 * q_abs_ref_eff; + match iqn.execute_training_pipeline( online_h_s2, next_states_buf, &self.target_dueling, Some(&iqn_stream_ref), ws_override, + weight_clamp_max_abs, ) { Ok(_) => { iqn_ok = true; } Err(e) => { @@ -1996,7 +2025,11 @@ impl FusedTrainingCtx { // Adam on attn_stream. let lr = self.trainer.config().lr; let mgn = self.trainer.config().max_grad_norm; - attn.adam_step(lr, mgn, Some(&attn_stream_ref), ws_override) + // SP3 Mech 9 (close-out v2): ISV-driven post-Adam |p| clamp + // for `attn_adam_kernel`. Same pattern as IQN/IQL/main DQN. + let q_abs_ref_eff_attn = self.trainer.read_isv_signal_at(Q_ABS_REF_INDEX).max(1.0_f32); + let attn_weight_clamp_max_abs: f32 = 100.0_f32 * q_abs_ref_eff_attn; + attn.adam_step(lr, mgn, Some(&attn_stream_ref), ws_override, attn_weight_clamp_max_abs) .map_err(|e| anyhow::anyhow!("Attention Adam: {e}"))?; // Record completion on attn_stream. @@ -2029,7 +2062,11 @@ impl FusedTrainingCtx { .map_err(|e| anyhow::anyhow!("Attention backward: {e}"))?; let lr = self.trainer.config().lr; let mgn = self.trainer.config().max_grad_norm; - attn.adam_step(lr, mgn, None, None) + // SP3 Mech 9 (close-out v2): same ISV-derived bound as + // parallel arm. Sequential-fallback path. + let q_abs_ref_eff_attn = self.trainer.read_isv_signal_at(Q_ABS_REF_INDEX).max(1.0_f32); + let attn_weight_clamp_max_abs: f32 = 100.0_f32 * q_abs_ref_eff_attn; + attn.adam_step(lr, mgn, None, None, attn_weight_clamp_max_abs) .map_err(|e| anyhow::anyhow!("Attention Adam: {e}"))?; } @@ -2038,9 +2075,15 @@ impl FusedTrainingCtx { let online_h_s2 = self.trainer.save_h_s2(); let next_states_buf = self.trainer.next_states_buf(); + // SP3 Mech 9 (close-out v2): same ISV-derived bound as the + // parallel arm above. Mirror the dqn_adam pattern. + let q_abs_ref_eff = self.trainer.read_isv_signal_at(Q_ABS_REF_INDEX).max(1.0_f32); + let weight_clamp_max_abs: f32 = 100.0_f32 * q_abs_ref_eff; + match iqn.execute_training_pipeline( online_h_s2, next_states_buf, &self.target_dueling, None, None, + weight_clamp_max_abs, ) { Ok(_) => { iqn_ok = true; } Err(e) => { @@ -3709,7 +3752,7 @@ impl FusedTrainingCtx { /// `GpuDqnTrainer::run_nan_checks_post_forward`. /// SP1 Phase B: returns 48 flags (was 24). Slots 24-35 cover backward-path /// checks; slots 36-47 reserved for SP2/SP3 headroom. - pub(crate) fn read_nan_flags(&self) -> Result<[i32; 48], crate::MLError> { + pub(crate) fn read_nan_flags(&self) -> Result<[i32; 49], crate::MLError> { self.trainer.read_nan_flags() } diff --git a/crates/ml/src/trainers/dqn/trainer/training_loop.rs b/crates/ml/src/trainers/dqn/trainer/training_loop.rs index 2f897ccbc..f8e5110d7 100644 --- a/crates/ml/src/trainers/dqn/trainer/training_loop.rs +++ b/crates/ml/src/trainers/dqn/trainer/training_loop.rs @@ -1753,8 +1753,24 @@ impl DQNTrainer { } if count > 0 { + // SP3 Mech 9 (close-out v2): ISV-driven post-Adam |p| clamp for + // `curiosity_adam_step`. Same `100 × Q_ABS_REF.max(1.0)` bound as + // the main DQN / IQN / IQL / attention Adam launches. Curiosity + // collector doesn't own ISV — read from fctx here and thread the + // bound through. ε on multiplier per SP1 pearl. + let weight_clamp_max_abs: f32 = if let Some(ref fused) = self.fused_ctx { + let q_abs_ref_eff = fused.read_isv_signal_at(Q_ABS_REF_INDEX).max(1.0_f32); + 100.0_f32 * q_abs_ref_eff + } else { + // No fctx → no ISV → disable Mech 9. Reachable only on the + // CPU-only build path (CUDA build asserts fctx Some at + // line 442). 0=disabled keeps curiosity functional without + // unbounded clamps. + 0.0_f32 + }; if let Err(e) = collector.train_curiosity_gpu( gpu_batch.n_episodes, gpu_batch.timesteps, + weight_clamp_max_abs, ) { debug!("GPU curiosity training failed (non-fatal): {e}"); } @@ -2035,6 +2051,10 @@ impl DQNTrainer { "heads_weight_max", // 45 "target_q_post_clip", // 46 "atom_span_max", // 47 + // SP3 close-out v2 (slot 48): IQN online_params + // weight diagnostic — the slot-26 GEMM blind + // spot that hid the F1 step-3540 NaN. + "iqn_weight_max", // 48 ]; let flagged: Vec = flags.iter().enumerate() .filter(|(_, &f)| f != 0) diff --git a/docs/dqn-backward-nan-audit.md b/docs/dqn-backward-nan-audit.md index bafd2a5f2..e2ab5e2d5 100644 --- a/docs/dqn-backward-nan-audit.md +++ b/docs/dqn-backward-nan-audit.md @@ -823,11 +823,11 @@ If Gate 1 fails, halt before SP3 starts and re-investigate F0 cause (alignment e | Mech 6 | Anchored upper_bound on adaptive grad clip (100× multiplier) | `48c25d999` | gpu_dqn_trainer.rs | | ~~Mech 7~~ | ~~Per-element gradient clip in Adam~~ | reverted in `f67ede94f` | (n/a) | | ~~Mech 8~~ | ~~slow_ema fold-boundary reset~~ | reverted in `8956c2fb7` | (n/a) | -| Mech 9 | Post-Adam ISV-driven weight clamp | `8956c2fb7` | dqn_utility_kernels.cu, gpu_dqn_trainer.rs, decision_transformer.rs | +| Mech 9 | Post-Adam ISV-driven weight clamp (all 5 Adam kernels) | `8956c2fb7` (dqn) + close-out v2 (iqn/iql/attn/curiosity) | dqn_utility_kernels.cu, iqn_dual_head_kernel.cu, iql_value_kernel.cu, attention_backward_kernel.cu, curiosity_training_kernel.cu, gpu_dqn_trainer.rs, gpu_iqn_head.rs, gpu_iql_trainer.rs, gpu_attention.rs, gpu_tlob.rs, gpu_curiosity_trainer.rs, gpu_experience_collector.rs, fused_training.rs, training_loop.rs, decision_transformer.rs | **Supporting:** Task B1 (`aee11f321`) — 24 diagnostic accessors (Adam m/v + weight + target_q + atom_positions). -### Mech 5 slot allocation (slots 36-47) +### Mech 5 slot allocation (slots 36-48) | Slot | Name | Threshold | Source | |---|---|---|---| @@ -843,6 +843,7 @@ If Gate 1 fails, halt before SP3 starts and re-investigate F0 cause (alignment e | 45 | heads_weight_max | same | heads params slice | | 46 | target_q_post_clip | ≥ 95 × q_abs_ref_eff | denoise_target_q_buf | | 47 | atom_span_max | ≥ 190 × q_abs_ref_eff | per_sample_support | +| 48 | iqn_weight_max | ≥ 1e3 × q_abs_ref_eff | IQN online_params (SP3 close-out v2 — slot-26 GEMM blind spot, fold-1 step-3540 sentinel) | `q_abs_ref_eff = max(isv[Q_ABS_REF_INDEX=16], 1.0)` — ε on multiplier per SP1 pearl. Cold-start ISV (~0) → all thresholds floor at their multiplier × 1. diff --git a/docs/dqn-wire-up-audit.md b/docs/dqn-wire-up-audit.md index 374730db3..cd0869951 100644 --- a/docs/dqn-wire-up-audit.md +++ b/docs/dqn-wire-up-audit.md @@ -2,6 +2,8 @@ **Status:** Populated during Plan 1 Task 6 (A.5 orphan audit). Updated on every commit per Invariant 7. +SP3 close-out v2 — Mech 9 to all 5 Adam kernels + IQN weight diag slot 48 (2026-04-30): real fix for the F1 step-3540 NaN. Commit `8956c2fb7` wired Mech 9 (post-Adam `|p_val|` clamp) into ONE Adam kernel (`dqn_adam_update_kernel`) out of FIVE in the codebase. The slot-26 GEMM (`apply_iqn_trunk_gradient` output) computes `grad_iqn @ W_iqn^T`; `W_iqn` lives in IQN's separate `online_params` buffer, updated by `iqn_adam_kernel` — never reached by Mech 9. Slots 44-45 only cover trunk + heads weights, leaving IQN's separate param buffer without a diagnostic. The chain: IQN weights drift via `iqn_adam_kernel` → finite gradient × non-finite weight → slot-26 NaN. Same partial-refactor class as Mech 4's missing GpuAttention reset (caught by B5 audit), corrected the same way per `feedback_no_partial_refactor.md`: every consumer of the contract migrates together. Extended Mech 9 to all four remaining Adam kernels — `iqn_adam_kernel` (`iqn_dual_head_kernel.cu`), `iql_adam_kernel` (`iql_value_kernel.cu`), `attn_adam_kernel` (`attention_backward_kernel.cu`, used by both `GpuAttention` and `GpuTlob` via the shared cubin), `curiosity_adam_step` (`curiosity_training_kernel.cu`). Same `if (weight_clamp_max_abs > 0.0f) p = fminf(fmaxf(p, -bound), bound)` pattern, last step before writeback. Same `100 × Q_ABS_REF.max(1.0)` ISV-driven bound (same slot 16, ε on multiplier per SP1 pearl). Each launch site computes the bound host-side and passes it as the trailing kernel arg, mirroring the existing dqn_adam_update_kernel pattern. IQN: `train_iqn_step_gpu` + `execute_training_pipeline` take new `weight_clamp_max_abs: f32` parameters; both fork-join arms in `fused_training.rs::submit_aux_ops` compute and pass the bound. IQL: `train_value_step` takes the bound; both `gpu_iql.train_value_step` (high-tau) and `gpu_iql_low.train_value_step` (low-tau) wired. Attention: `GpuAttention::adam_step` and `GpuTlob::adam_step` both take new `weight_clamp_max_abs: f32`; all 3 call sites in fused_training (parallel attn, sequential attn, TLOB Phase 6) wired. Curiosity: `GpuCuriosityTrainer::train_on_collector_buffers` and the inner `launch_adam_step` helper take the bound; `GpuExperienceCollector::train_curiosity_gpu` wraps and forwards; `training_loop.rs::collect_step` reads ISV from `fused_ctx` (curiosity collector doesn't own ISV) and passes through. DT launch in `decision_transformer.rs` keeps the 0.0 disable (offline-RL, outside SP3 scope). Added IQN-weight diagnostic slot 48 to plug the slot-26 GEMM blind spot: `nan_flags_buf` 48 → 49 (`stream.alloc_zeros::(49)`), fused-kernel block count 24 → 25 with new branch for absolute slot 48 = relative slot 24 at threshold `1e3 × q_abs_ref_eff` (matches slot 44-45 weight regime), metadata buffers `nan_check_buf_ptrs` / `nan_check_buf_lens` extended 24 → 25 with the new slot 48 entry populated by `populate_nan_check_meta` from `gpu_iqn.online_params_ptr()` + `online_params_len()` (new public accessors on `GpuIqnHead`; existing `cfg(test)` `online_params_slice()` accessor kept). Name table in `training_loop.rs::halt_grad_collapse` extended to 49 entries with `"iqn_weight_max"` at slot 48. `read_nan_flags()` return type bumped `[i32; 48] → [i32; 49]` on both `GpuDqnTrainer` and the `FusedTrainingCtx` wrapper. Audit doc `dqn-backward-nan-audit.md` Mech 5 slot table now spans 36-48 with the new row; Mech 9 row updated to span all 5 Adam kernels (file list expanded). No new ISV slots, no new mapped-pinned allocations beyond the metadata-buffer +1, no graph-topology change beyond the +1 block in the already-fused NaN check. Touched: `cuda_pipeline/iqn_dual_head_kernel.cu` (kernel signature + clamp), `cuda_pipeline/iql_value_kernel.cu` (kernel signature + clamp), `cuda_pipeline/attention_backward_kernel.cu` (kernel signature + clamp), `cuda_pipeline/curiosity_training_kernel.cu` (kernel signature + clamp), `cuda_pipeline/dqn_utility_kernels.cu` (fused NaN check: grid 24→25, slot 24 branch), `cuda_pipeline/gpu_iqn_head.rs` (new `online_params_ptr` / `online_params_len` accessors + `train_iqn_step_gpu` + `execute_training_pipeline` signature + Adam launch arg), `cuda_pipeline/gpu_iql_trainer.rs` (`train_value_step` signature + Adam launch arg), `cuda_pipeline/gpu_attention.rs` (`adam_step` signature + launch arg), `cuda_pipeline/gpu_tlob.rs` (`adam_step` signature + launch arg), `cuda_pipeline/gpu_curiosity_trainer.rs` (`launch_adam_step` + `train_on_collector_buffers` signatures + 4 launch args), `cuda_pipeline/gpu_experience_collector.rs` (`train_curiosity_gpu` signature wrapper), `cuda_pipeline/gpu_dqn_trainer.rs` (alloc 48→49, metadata 24→25, `populate_nan_check_meta` 2 new params + slot 48 entry, fused-kernel `BLOCKS` 24→25, `read_nan_flags` 48→49), `trainers/dqn/fused_training.rs` (Q_ABS_REF_INDEX import + 4 wiring sites + read_nan_flags type), `trainers/dqn/trainer/training_loop.rs` (curiosity ISV read + name table 49 entries), `docs/dqn-backward-nan-audit.md` (slot table + Mech 9 row). cargo check -p ml --lib clean (12 pre-existing warnings, none new). Validation deferred to one L40S smoke at this HEAD; expected F1 trains past step 3720 ceiling and slots 36-43 + new slot 48 stay quiet. + SP3 Mech 9 — post-Adam ISV-driven weight clamp + Mech 8 revert (2026-04-30): coordinated close-out of SP3 Q-learning numerical-stability per `feedback_no_partial_refactor.md`. Mech 8 (slow_ema fold-boundary reset, commit `b8a7ac6f7`) reverted after smoke-test-rxhjh F1 NaN @ step 2300 — WORSE than the 3720-step ceiling Mech 6 alone produced. Persisting the α=0.001 `grad_norm_slow_ema` across folds (anchored at F0's smaller grad scale) was providing UNINTENTIONAL F1 protection by tightening Mech 6's `100 × slow_ema × isv` upper_bound during F1 ramp-up; resetting it loosened the bound during the very transient when F1 needed it tightest, accelerating Adam saturation. Removed `GpuDqnTrainer::reset_grad_norm_slow_ema` method and the `fused_training.rs::reset_for_fold` call site; kept the `grad_norm_slow_ema_pinned` field (still consumed by Mech 6) and Mech 6 multiplier at 100× (correct steady-state value, already restored from the 5× experiment in the Mech 8 commit). Mech 9 is the real fix for the cross-fold Adam pathology — root cause being targeted is cuBLAS sgemm f32-accumulator overflow at slots 26 (`iqn_trunk_m`) + 32 (`bn_d_concat_buf`) at F1 step ~3540 with INPUTS clean (slots 27, 35 unflagged). The matmul output saturates because WEIGHT tensors have drifted to extreme-but-finite magnitudes via Adam updates over thousands of steps. Mech 1 (gradient clamp) and Mech 6 (grad-norm clip) bound gradients, not parameters; neither prevents this drift. Mech 9 clamps `|p_val| ≤ 100 × Q_ABS_REF.max(1.0)` inside `dqn_adam_update_kernel` after the L1 proximal step (new trailing arg `weight_clamp_max_abs` on the kernel signature). ISV-driven: bound reads `Q_ABS_REF_INDEX = 16` host-side (same slot consumed by Mechs 1+2+5+6 — no new slot), ε on the multiplier per the SP1 pearl (`isv.max(1.0)`, never floor the bound itself or re-introduce the cold-start clamp pathology), 1 OOM below the slot 44-45 diagnostic threshold (`1e3 × Q_ABS_REF`) so the regression sentinel retains 10× firing headroom — symmetric with Mech 1's clamp/diagnostic ratio. All four Adam launch sites in `gpu_dqn_trainer.rs` wired with the bound (`launch_adam_update` at ~line 19470, `step_selectivity_adam` at ~line 20860, denoise-head Adam at ~line 17014, OFI-embed Adam at ~line 5758); the Decision Transformer Adam launch in `decision_transformer.rs` passes 0.0 to disable the clamp (DT is a separate offline-RL trainer outside SP3 scope, no ISV bus access; kernel branch `if (bound > 0.0f)` makes the disabled path zero-cost). No new ISV slots, no new kernels, no graph topology change. Touched: `cuda_pipeline/dqn_utility_kernels.cu` (kernel signature + clamp), `cuda_pipeline/gpu_dqn_trainer.rs` (4 launch sites + Mech 6 comment update + `reset_grad_norm_slow_ema` method removed), `cuda_pipeline/decision_transformer.rs` (DT launch site, disabled), `trainers/dqn/fused_training.rs` (Mech 8 call removed), `docs/dqn-backward-nan-audit.md` (mechanism table + Mech 8/9 subsection). cargo check clean. No fingerprint change — kernel signature change but param-tensor layout unchanged. SP3 Mech 8 — `grad_norm_slow_ema` fold-boundary reset (2026-04-29): closed the partial-refactor gap where the α=0.001 grad-norm slow EMA (driving Mech 6's anchored upper bound on `adaptive_clip`) persisted across fold boundaries while every other distribution-tracking signal — Adam m/v (`reset_adam_state`), main DQN target params (`sync_target_from_online`), IQN target + Adam (`iqn.sync_target_from_online` + `iqn.reset_adam_state`), TLOB / IQL / IQL-low / 4-head attention Adam (`reset_adam_state` on each), C51 atom-position EMAs, replay buffer — already reset per existing `reset_for_fold` infrastructure. The α=0.001 half-life (~693 steps) meant `grad_norm_slow_ema` lagged the new fold's grad-norm regime by hundreds of steps; Mech 6's `100 × slow_ema × isv` upper bound was anchored to the WRONG fold's scale during F1 ramp-up, producing the non-monotonic multiplier-tuning dance observed across smokes (smoke-test-fxvkk 100× F1-NaN @ 3720 / smoke-test-ftdjz 100×+Mech7 F1-collapse @ 2040 / smoke-test-d25vq 5× F1-collapse @ 2820 — the multiplier was searching the wrong dimension). Three coordinated changes per `feedback_no_partial_refactor.md`: (1) Mech 6 multiplier restored 5× → 100× in `gpu_dqn_trainer.rs::update_adaptive_clip` (the original principled setting; v2's 5× was over-clipping legitimate F0-ramp gradients while still allowing F1-anchor saturation), (2) new `GpuDqnTrainer::reset_grad_norm_slow_ema` method zeroes the existing mapped-pinned scalar (no new buffer, no new ISV slot — reuses Plan C Phase 2 follow-up K's grad_norm_slow_ema_pinned at `gpu_dqn_trainer.rs:3337`), (3) wired into `fused_training.rs::reset_for_fold` immediately after the SP3 Mech 4 attention Adam reset block (line ~1066) — same fold-boundary contract consumer chain as Mech 3+4. F0 unaffected (cold-starts at slow_ema=0 already, so Mech 8 is a no-op on F0); F1+F2 first step transient sees `upper_bound = 100 × 0 × isv → MIN_CLIP=1.0` floor for ~10-50 steps as the EMA warms, then converges to legitimate 100× headroom over CURRENT fold's grad norm. Mech 7 stays reverted (per-element clip was misdiagnosis). Touched: `cuda_pipeline/gpu_dqn_trainer.rs` (Mech 6 comment + multiplier 5.0→100.0 line ~20370; new `reset_grad_norm_slow_ema` method line ~3858), `trainers/dqn/fused_training.rs` (call wired line ~1067). cargo check clean. No fingerprint change — touches an already-allocated mapped-pinned scalar, not param-tensor layout.