gitadmin/foxhunt
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

fix(dqn): SP3 close-out v2 — Mech 9 to all 5 Adam kernels + IQN weight diag

Browse Source 
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) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-04-30 13:59:23 +02:00
parent 5e6706a1f9
commit 1bcc703920
16 changed files with 271 additions and 41 deletions
Download Patch File Download Diff File Expand all files Collapse all files

16
crates/ml/src/cuda_pipeline/attention_backward_kernel.cu
View File

@@ -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;

14
crates/ml/src/cuda_pipeline/curiosity_training_kernel.cu
View File

@@ -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;
}
/* ------------------------------------------------------------------ */

17
crates/ml/src/cuda_pipeline/dqn_utility_kernels.cu
View File

@@ -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;

5
crates/ml/src/cuda_pipeline/gpu_attention.rs
View File

@@ -975,6 +975,10 @@ impl GpuAttention {
max_grad_norm: f32,
override_stream: Option<&Arc<CudaStream>>,
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}")))?;
}

19
crates/ml/src/cuda_pipeline/gpu_curiosity_trainer.rs
View File

@@ -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<f32>,
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<f32>,
actions: &CudaSlice<i32>,
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)");

65
crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs
View File

@@ -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::<i32>(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::<i32>(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<usize>,
iqn_adam_v_ptr: Option<u64>,
iqn_adam_v_len: Option<usize>,
// 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<u64>,
iqn_online_params_len: Option<usize>,
) -> 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)

7
crates/ml/src/cuda_pipeline/gpu_experience_collector.rs
View File

@@ -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(())

5
crates/ml/src/cuda_pipeline/gpu_iql_trainer.rs
View File

@@ -524,6 +524,10 @@ impl GpuIqlTrainer {
pub fn train_value_step(
&mut self,
states_f32: &CudaSlice<f32>,
/* 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),

37
crates/ml/src/cuda_pipeline/gpu_iqn_head.rs
View File

@@ -803,12 +803,24 @@ impl GpuIqnHead {
dqn_actions_buf: &CudaSlice<i32>,
dqn_rewards_buf: &CudaSlice<f32>,
dqn_dones_buf: &CudaSlice<f32>,
/* 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<f32, MLError> {
// 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<CudaStream>>,
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<f32, MLError> {
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) {

13
crates/ml/src/cuda_pipeline/gpu_tlob.rs
View File

@@ -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),

12
crates/ml/src/cuda_pipeline/iql_value_kernel.cu
View File

@@ -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;

20
crates/ml/src/cuda_pipeline/iqn_dual_head_kernel.cu
View File

@@ -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 */

55
crates/ml/src/trainers/dqn/fused_training.rs
View File

@@ -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()
}

20
crates/ml/src/trainers/dqn/trainer/training_loop.rs
View File

@@ -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<String> = flags.iter().enumerate()
.filter(|(_, &f)| f != 0)

5
docs/dqn-backward-nan-audit.md
View File

@@ -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.

2
docs/dqn-wire-up-audit.md
View File

File diff suppressed because one or more lines are too long