diff --git a/crates/ml/build.rs b/crates/ml/build.rs index 0876c9456..8c0637d3e 100644 --- a/crates/ml/build.rs +++ b/crates/ml/build.rs @@ -81,6 +81,7 @@ fn main() { "iqn_cvar_kernel.cu", "mamba2_temporal_kernel.cu", "graph_utility_kernels.cu", + "grad_decomp_kernel.cu", ]; // ALL kernels get common header (BF16 types + wrappers) diff --git a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs index 37c29f069..878b67fee 100644 --- a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs +++ b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs @@ -75,6 +75,7 @@ static ENSEMBLE_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/ensembl static CQL_GRAD_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/cql_grad_kernel.cubin")); static MAMBA2_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/mamba2_temporal_kernel.cubin")); pub(crate) static GRAPH_UTILITY_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/graph_utility_kernels.cubin")); +static GRAD_DECOMP_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/grad_decomp_kernel.cubin")); /// Mamba2 temporal scan configuration. const MAMBA2_HISTORY_K: usize = 8; // Rolling history length @@ -970,6 +971,49 @@ pub struct GpuDqnTrainer { /// training graph capture region. grad_readback_pinned_ptr: usize, // host-visible f32 buffer grad_readback_pinned_capacity: usize, // element count (TOTAL_PARAMS) + /// Task 2.0 — per-component magnitude-branch grad decomposition. + /// + /// Four device-side scratch buffers sized to the combined length of + /// branch 0 (direction, tensors 8..12) + branch 1 (magnitude, tensors + /// 12..16). BEFORE each component's backward fires (captured in the + /// training graph), a `copy_f32` kernel snapshots `grad_buf` over the + /// branch 0+1 byte range into these slots. AFTER the component's + /// backward, `grad_component_delta_norm` computes `‖current−snapshot‖` + /// on direction and magnitude slices and writes two f32 values into + /// the pinned result slot. + /// + /// The C51 slot is paired with `grad_zero_ref_buf` (permanent zeros) + /// because `grad_buf` is memset-zeroed at the start of every + /// `submit_forward_ops_main` — there is no "pre-C51" grad to snapshot. + /// Uniform subtraction keeps the kernel interface simple. + grad_snapshot_iqn: CudaSlice, // [branch01_elems] — snapshot before IQN trunk grad + grad_snapshot_cql: CudaSlice, // [branch01_elems] — snapshot before CQL SAXPY + grad_zero_ref_buf: CudaSlice, // [branch01_elems] — permanent zeros; C51 snapshot substitute + grad_snapshot_ens: CudaSlice, // [branch01_elems] — snapshot before ensemble diversity + /// Pinned device-mapped 8-float result slot written by the reduction + /// kernel. Layout (matches plan Task 2.0 Step 4): + /// [0]=iqn_mag [1]=iqn_dir + /// [2]=cql_mag [3]=cql_dir + /// [4]=c51_mag [5]=c51_dir + /// [6]=ens_mag [7]=ens_dir + grad_decomp_result_pinned: *mut f32, + grad_decomp_result_dev_ptr: u64, + /// Cached grad-decomp slice metadata (element indices into `grad_buf`). + /// `grad_decomp_dir_start`/`dir_len` cover branch 0 tensors 8..12, + /// `grad_decomp_mag_start`/`mag_len` cover branch 1 tensors 12..16. + /// `grad_decomp_snapshot_len` = dir_len + mag_len. + grad_decomp_dir_start: i32, + grad_decomp_dir_len: i32, + grad_decomp_mag_start: i32, + grad_decomp_mag_len: i32, + grad_decomp_snapshot_len: usize, + /// Reduction kernel (one-block 256-thread launch, no atomics). + grad_decomp_kernel: CudaFunction, + /// Host-side cached per-component norms populated at epoch boundary + /// from the pinned result slot. Index order matches `grad_mag_*_ratio` + /// accessors: `[iqn, cql, c51, ens]`. + grad_component_norms_mag: [f32; 4], + grad_component_norms_dir: [f32; 4], params_buf: CudaSlice, // [TOTAL_PARAMS] f32 master online parameters (Adam operates here) target_params_buf: CudaSlice, // [TOTAL_PARAMS] f32 master target parameters (EMA operates here) m_buf: CudaSlice, // [TOTAL_PARAMS] Adam first moment (f32 for precision) @@ -2045,6 +2089,12 @@ impl Drop for GpuDqnTrainer { cudarc::driver::result::free_host(self.grad_readback_pinned_ptr as *mut std::ffi::c_void) }; } + // Task 2.0 — pinned device-mapped grad-component result slot. + if !self.grad_decomp_result_pinned.is_null() { + let _ = unsafe { + cudarc::driver::result::free_host(self.grad_decomp_result_pinned.cast()) + }; + } if !self.per_branch_q_gaps_pinned.is_null() { let _ = unsafe { cudarc::driver::result::free_host(self.per_branch_q_gaps_pinned.cast()) }; } @@ -2172,6 +2222,136 @@ impl GpuDqnTrainer { Ok(if dir > 1e-9 { mag / dir } else { 0.0 }) } + /// Task 2.0 — in-graph DtoD snapshot of `grad_buf`'s branch 0+1 slice into + /// the caller-supplied scratch buffer. Uses the `copy_f32` kernel (graph- + /// captured DtoD) because `cuMemcpyDtoDAsync` is NOT captured by CUDA + /// Graph (see `graph_safe_copy_f32` doc comment). + /// + /// Copies element range `[grad_decomp_dir_start, grad_decomp_dir_start + + /// grad_decomp_snapshot_len)` from `grad_buf` (which spans branches 0+1 + /// contiguously — see `compute_param_sizes`) into the first + /// `grad_decomp_snapshot_len` elements of `dst`. + pub(crate) fn grad_decomp_snapshot(&self, dst: &CudaSlice, ctx: &str) + -> Result<(), MLError> + { + let elem_size = std::mem::size_of::() as u64; + let src_ptr = self.ptrs.grad_buf + (self.grad_decomp_dir_start as u64) * elem_size; + let dst_ptr = dst.raw_ptr(); + let n_bytes = self.grad_decomp_snapshot_len * elem_size as usize; + self.graph_safe_copy_f32(dst_ptr, src_ptr, n_bytes, ctx) + } + + /// Task 2.0 — launch `grad_component_delta_norm` on a per-component + /// snapshot, writing (mag_norm, dir_norm) into the pinned result slot at + /// the component's 2-float offset. + /// + /// `slot_offset_elems` is the element offset into the 8-float pinned slot + /// — 0=iqn, 2=cql, 4=c51, 6=ens (matches pinned-layout convention + /// documented on `grad_decomp_result_pinned`). + pub(crate) fn launch_grad_decomp( + &self, + snapshot: &CudaSlice, + slot_offset_elems: i32, + ctx: &str, + ) -> Result<(), MLError> { + let current_ptr = self.ptrs.grad_buf; + let snapshot_ptr = snapshot.raw_ptr(); + let result_ptr = self.grad_decomp_result_dev_ptr + + (slot_offset_elems as u64) * (std::mem::size_of::() as u64); + let cfg = LaunchConfig { + grid_dim: (1, 1, 1), + block_dim: (256, 1, 1), + shared_mem_bytes: 0, + }; + unsafe { + self.stream.launch_builder(&self.grad_decomp_kernel) + .arg(¤t_ptr) + .arg(&snapshot_ptr) + .arg(&self.grad_decomp_dir_start) + .arg(&self.grad_decomp_dir_len) + .arg(&self.grad_decomp_mag_start) + .arg(&self.grad_decomp_mag_len) + .arg(&result_ptr) + .launch(cfg) + .map_err(|e| MLError::ModelError(format!("grad_decomp {ctx}: {e}")))?; + } + Ok(()) + } + + /// Task 2.0 — convenience: snapshot `grad_buf` into the component's + /// snapshot buffer (captured DtoD via `copy_f32` kernel). Called BEFORE + /// the component's backward fires. + pub(crate) fn grad_decomp_snapshot_iqn(&self) -> Result<(), MLError> { + self.grad_decomp_snapshot(&self.grad_snapshot_iqn, "grad_snapshot_iqn") + } + pub(crate) fn grad_decomp_snapshot_cql(&self) -> Result<(), MLError> { + self.grad_decomp_snapshot(&self.grad_snapshot_cql, "grad_snapshot_cql") + } + pub(crate) fn grad_decomp_snapshot_ens(&self) -> Result<(), MLError> { + self.grad_decomp_snapshot(&self.grad_snapshot_ens, "grad_snapshot_ens") + } + + /// Task 2.0 — convenience: launch the reduction kernel for each component + /// with its pre-configured snapshot buffer and pinned-slot offset. + pub(crate) fn grad_decomp_launch_iqn(&self) -> Result<(), MLError> { + self.launch_grad_decomp(&self.grad_snapshot_iqn, 0, "iqn") + } + pub(crate) fn grad_decomp_launch_cql(&self) -> Result<(), MLError> { + self.launch_grad_decomp(&self.grad_snapshot_cql, 2, "cql") + } + /// C51 uses `grad_zero_ref_buf` (permanent zeros) because `grad_buf` is + /// zero before `submit_forward_ops_main` — there is no pre-C51 snapshot. + /// Kernel computes `‖grad_buf − 0‖` on branch 0+1 slices. + pub(crate) fn grad_decomp_launch_c51(&self) -> Result<(), MLError> { + self.launch_grad_decomp(&self.grad_zero_ref_buf, 4, "c51") + } + pub(crate) fn grad_decomp_launch_ens(&self) -> Result<(), MLError> { + self.launch_grad_decomp(&self.grad_snapshot_ens, 6, "ens") + } + + /// Task 2.0 — populate the host-side `grad_component_norms_mag/_dir` + /// caches from the pinned result slot. Called at epoch boundary, before + /// HEALTH_DIAG emission. + /// + /// Stream-syncs to ensure the final step's reduction kernel has completed + /// (same pattern as `per_branch_grad_norms`). Zero-copy: the pinned + /// memory is device-mapped, so the host read is a direct dereference of + /// the last-step kernel output. + /// + /// Pinned layout (written by `grad_component_delta_norm`, 2 floats per + /// component — mag then dir): + /// [0]=iqn_mag [1]=iqn_dir [2]=cql_mag [3]=cql_dir + /// [4]=c51_mag [5]=c51_dir [6]=ens_mag [7]=ens_dir + /// + /// Host caches use `[iqn, cql, c51, ens]` index order matching the + /// `grad_mag_{iqn,cql,c51,ens}_ratio` accessor families. + pub fn refresh_grad_component_norms(&mut self) -> Result<(), MLError> { + if self.grad_decomp_result_pinned.is_null() { + return Err(MLError::ModelError( + "refresh_grad_component_norms: pinned result slot is null".into() + )); + } + self.stream.synchronize().map_err(|e| { + MLError::ModelError(format!("grad_component_norms sync: {e}")) + })?; + let host_slice: &[f32] = unsafe { + std::slice::from_raw_parts(self.grad_decomp_result_pinned, 8) + }; + for comp in 0..4_usize { + self.grad_component_norms_mag[comp] = host_slice[2 * comp]; + self.grad_component_norms_dir[comp] = host_slice[2 * comp + 1]; + } + Ok(()) + } + + /// Task 2.0 — accessor for cached per-component magnitude norms + /// populated by `refresh_grad_component_norms`. Index order: + /// `[iqn, cql, c51, ens]`. + pub fn grad_component_norms_mag(&self) -> [f32; 4] { self.grad_component_norms_mag } + /// Task 2.0 — accessor for cached per-component direction norms + /// populated by `refresh_grad_component_norms`. + pub fn grad_component_norms_dir(&self) -> [f32; 4] { self.grad_component_norms_dir } + /// Task 0.6 — per-branch target/online parameter drift (H8 detection signal). /// /// Returns `[direction, magnitude, order, urgency]` RMS deviation between @@ -5790,6 +5970,79 @@ impl GpuDqnTrainer { }; let grad_readback_pinned_capacity = total_params; + // Task 2.0 — per-component grad-decomposition snapshots + pinned result slot. + // + // The snapshots only need to cover branch 0 (direction, tensors 8..12) and + // branch 1 (magnitude, tensors 12..16) — a single contiguous element range + // in `grad_buf`. Sizes are computed from `compute_param_sizes` exactly like + // `per_branch_grad_norms` (Task 0.4) so direction/magnitude bounds stay in + // lock-step with the flat parameter layout. + let (grad_decomp_dir_start, grad_decomp_dir_len, grad_decomp_mag_start, + grad_decomp_mag_len, grad_decomp_snapshot_len) = { + let param_sizes = compute_param_sizes(&config); + let dir_start_byte = padded_byte_offset(¶m_sizes, 8) as usize; + let mag_start_byte = padded_byte_offset(¶m_sizes, 12) as usize; + let branch2_start_byte = padded_byte_offset(¶m_sizes, 16) as usize; + let elem_size = std::mem::size_of::(); + let dir_start = dir_start_byte / elem_size; + let mag_start = mag_start_byte / elem_size; + let mag_end = branch2_start_byte / elem_size; + let dir_len = mag_start - dir_start; + let mag_len = mag_end - mag_start; + let snap_len = dir_len + mag_len; + (dir_start as i32, dir_len as i32, mag_start as i32, mag_len as i32, snap_len) + }; + let grad_snapshot_iqn = stream.alloc_zeros::(grad_decomp_snapshot_len) + .map_err(|e| MLError::ModelError(format!("alloc grad_snapshot_iqn: {e}")))?; + let grad_snapshot_cql = stream.alloc_zeros::(grad_decomp_snapshot_len) + .map_err(|e| MLError::ModelError(format!("alloc grad_snapshot_cql: {e}")))?; + // C51 fires BEFORE the aux phase and `grad_buf` is memset-zeroed at the + // start of `submit_forward_ops_main` — no "pre-C51" grad exists to + // snapshot. `grad_zero_ref_buf` is initialized to zeros and never + // written again; the reduction kernel reads (current − 0) = current on + // the C51 path. + let grad_zero_ref_buf = stream.alloc_zeros::(grad_decomp_snapshot_len) + .map_err(|e| MLError::ModelError(format!("alloc grad_zero_ref_buf: {e}")))?; + let grad_snapshot_ens = stream.alloc_zeros::(grad_decomp_snapshot_len) + .map_err(|e| MLError::ModelError(format!("alloc grad_snapshot_ens: {e}")))?; + + // Pinned device-mapped 8-float result slot. Device writes via + // `grad_decomp_result_dev_ptr` (captured in graph); host reads via + // `grad_decomp_result_pinned` at epoch boundary (zero-copy). + let grad_decomp_result_pinned: *mut f32 = unsafe { + let flags = cudarc::driver::sys::CU_MEMHOSTALLOC_DEVICEMAP; + cudarc::driver::result::malloc_host(8 * std::mem::size_of::(), flags) + .map_err(|e| MLError::ModelError(format!("pinned grad_decomp_result alloc: {e}")))? + as *mut f32 + }; + unsafe { + for i in 0..8 { *grad_decomp_result_pinned.add(i) = 0.0; } + } + let grad_decomp_result_dev_ptr = unsafe { + let mut dp = 0u64; + cudarc::driver::sys::cuMemHostGetDevicePointer_v2( + &mut dp as *mut u64, + grad_decomp_result_pinned.cast(), + 0, + ); + dp + }; + + // Reduction kernel — loaded from a dedicated CUmodule so its CUfunction + // is isolated from all other graphs (Hopper CUfunction-isolation rule). + let grad_decomp_kernel = { + let module = stream.context().load_cubin(GRAD_DECOMP_CUBIN.to_vec()) + .map_err(|e| MLError::ModelError(format!("grad_decomp cubin load: {e}")))?; + module.load_function("grad_component_delta_norm") + .map_err(|e| MLError::ModelError(format!("grad_component_delta_norm load: {e}")))? + }; + info!( + "GpuDqnTrainer: grad_decomp kernel loaded — snapshot_len={} dir=[{}..+{}] mag=[{}..+{}]", + grad_decomp_snapshot_len, + grad_decomp_dir_start, grad_decomp_dir_len, + grad_decomp_mag_start, grad_decomp_mag_len, + ); + // eval_v_range — pinned device-mapped (zero-copy write from CPU). let eval_v_range_pinned: *mut f32 = unsafe { let flags = cudarc::driver::sys::CU_MEMHOSTALLOC_DEVICEMAP; @@ -7644,6 +7897,20 @@ impl GpuDqnTrainer { grad_buf, grad_readback_pinned_ptr, grad_readback_pinned_capacity, + grad_snapshot_iqn, + grad_snapshot_cql, + grad_zero_ref_buf, + grad_snapshot_ens, + grad_decomp_result_pinned, + grad_decomp_result_dev_ptr, + grad_decomp_dir_start, + grad_decomp_dir_len, + grad_decomp_mag_start, + grad_decomp_mag_len, + grad_decomp_snapshot_len, + grad_decomp_kernel, + grad_component_norms_mag: [0.0_f32; 4], + grad_component_norms_dir: [0.0_f32; 4], params_buf, target_params_buf, m_buf, diff --git a/crates/ml/src/cuda_pipeline/grad_decomp_kernel.cu b/crates/ml/src/cuda_pipeline/grad_decomp_kernel.cu new file mode 100644 index 000000000..466872b75 --- /dev/null +++ b/crates/ml/src/cuda_pipeline/grad_decomp_kernel.cu @@ -0,0 +1,72 @@ +/** + * grad_decomp_kernel — Task 2.0 per-component gradient decomposition. + * + * Computes L2 norms of (current − snapshot) restricted to the direction + * (branch 0) and magnitude (branch 1) slices of grad_buf, writes two + * f32 results to a pinned device-mapped output slot. + * + * Launched four times per step — once per loss component (IQN/CQL/C51/Ens) + * — always AFTER the component's backward kernels finish writing into + * `grad_buf`. The matching BEFORE-backward snapshot is produced by the + * `copy_f32` kernel over the branch 0+1 byte range so that + * `current − snapshot` isolates exactly the component's additive delta. + * + * All reductions are block-internal via shared-memory tree reduction — + * NO atomicAdd (Task 0.5 / 0.8 no-atomics rule on the captured path). + * + * Layout: + * current — points to grad_buf[0] (f32, length total_params) + * snapshot — points to the per-component snapshot (f32, length + * dir_len + mag_len), laid out as: + * [0 .. dir_len) = direction slice + * [dir_len .. dir_len+mag_len) = magnitude slice + * indexed from zero; the caller is responsible for + * having copied the corresponding grad_buf range into + * this buffer BEFORE the component's backward fired. + * grad_dir_start/len, grad_mag_start/len + * — element offsets + lengths into `current` for each + * branch slice (so the kernel can read both sides of + * the subtraction from the correct places). + * result_out — [2] output: (mag_norm, dir_norm) — order chosen to + * match the grad_ratio_mag_dir convention. + * + * Launch config: one block, 256 threads. + */ +extern "C" __global__ void grad_component_delta_norm( + const float* __restrict__ current, + const float* __restrict__ snapshot, + int grad_dir_start, int dir_len, + int grad_mag_start, int mag_len, + float* __restrict__ result_out +) { + __shared__ float sum_mag[256]; + __shared__ float sum_dir[256]; + int tid = threadIdx.x; + sum_mag[tid] = 0.0f; + sum_dir[tid] = 0.0f; + + // Direction slice — snapshot offset starts at 0. + for (int i = tid; i < dir_len; i += blockDim.x) { + float d = current[grad_dir_start + i] - snapshot[i]; + sum_dir[tid] += d * d; + } + // Magnitude slice — snapshot offset starts at dir_len. + for (int i = tid; i < mag_len; i += blockDim.x) { + float d = current[grad_mag_start + i] - snapshot[dir_len + i]; + sum_mag[tid] += d * d; + } + __syncthreads(); + + // Block-level tree reduction (256 → 1). + for (int s = blockDim.x / 2; s > 0; s >>= 1) { + if (tid < s) { + sum_mag[tid] += sum_mag[tid + s]; + sum_dir[tid] += sum_dir[tid + s]; + } + __syncthreads(); + } + if (tid == 0) { + result_out[0] = sqrtf(sum_mag[0]); + result_out[1] = sqrtf(sum_dir[0]); + } +} diff --git a/crates/ml/src/trainers/dqn/fused_training.rs b/crates/ml/src/trainers/dqn/fused_training.rs index 0ba196c87..2680d1774 100644 --- a/crates/ml/src/trainers/dqn/fused_training.rs +++ b/crates/ml/src/trainers/dqn/fused_training.rs @@ -891,6 +891,43 @@ impl FusedTrainingCtx { self.trainer.grad_ratio_mag_dir().unwrap_or(0.0) } + /// Task 2.0 — refresh the cached per-component grad-norm arrays from the + /// pinned result slot populated by the in-graph reduction kernel. + /// Epoch-boundary cold path; stream-syncs once. Returns `Err` on + /// pinned-buffer absence or sync failure — the caller is expected to + /// `tracing::warn!` + fall back to stale values (same pattern as + /// `update_q_mag_means_cached`). + pub(crate) fn refresh_grad_component_norms(&mut self) -> Result<(), crate::MLError> { + self.trainer.refresh_grad_component_norms() + } + + /// Task 2.0 — IQN magnitude/direction grad-norm ratio. Index 0 in the + /// `grad_component_norms_*` caches (populated by + /// `refresh_grad_component_norms`). + pub(crate) fn grad_mag_iqn_ratio(&self) -> f32 { + let m = self.trainer.grad_component_norms_mag()[0]; + let d = self.trainer.grad_component_norms_dir()[0]; + if d > 1e-9 { m / d } else { 0.0 } + } + /// Task 2.0 — CQL magnitude/direction grad-norm ratio (index 1). + pub(crate) fn grad_mag_cql_ratio(&self) -> f32 { + let m = self.trainer.grad_component_norms_mag()[1]; + let d = self.trainer.grad_component_norms_dir()[1]; + if d > 1e-9 { m / d } else { 0.0 } + } + /// Task 2.0 — C51 magnitude/direction grad-norm ratio (index 2). + pub(crate) fn grad_mag_c51_ratio(&self) -> f32 { + let m = self.trainer.grad_component_norms_mag()[2]; + let d = self.trainer.grad_component_norms_dir()[2]; + if d > 1e-9 { m / d } else { 0.0 } + } + /// Task 2.0 — Ensemble magnitude/direction grad-norm ratio (index 3). + pub(crate) fn grad_mag_ens_ratio(&self) -> f32 { + let m = self.trainer.grad_component_norms_mag()[3]; + let d = self.trainer.grad_component_norms_dir()[3]; + if d > 1e-9 { m / d } else { 0.0 } + } + /// Task 0.6 — per-branch target/online param drift (H8 detection signal). /// Returns `[direction, magnitude, order, urgency]` RMS ‖target − online‖. /// Zeros on accessor failure (non-fatal — diag only). @@ -1170,6 +1207,18 @@ impl FusedTrainingCtx { ).map_err(|e| anyhow::anyhow!("HER in-place relabel kernel: {e}"))?; } + // Task 2.0 — measure C51 gradient contribution to branch 0 (direction) + // and branch 1 (magnitude). C51's backward ran inside + // `submit_forward_ops_main` (Phase 2); `grad_buf` was memset-zeroed + // at its start. At this point `grad_buf` holds exactly the C51/MSE + // blended backward — no IQN/CQL/Ens contributions yet. Reduction + // kernel runs BEFORE `apply_c51_budget_scale` so we capture the raw + // backward contribution rather than the budget-scaled amplitude. + // Uses `grad_zero_ref_buf` as the "snapshot" (permanent zeros), so + // kernel computes `‖grad_buf − 0‖` on branch 0+1 slices. + self.trainer.grad_decomp_launch_c51() + .map_err(|e| anyhow::anyhow!("Task 2.0 grad_decomp C51: {e}"))?; + // NOTE: Per-component gradient budget clip REMOVED (old explicit clip → constant grad norm). // B4/G5: Budgets are now applied as SCALE factors on individual SAXPY contributions, // not as norm clips. This preserves gradient magnitude variation while steering @@ -1443,6 +1492,17 @@ impl FusedTrainingCtx { // ── Post-join IQN processing (main stream) ── // trunk gradient, EMA, IQR run on self.stream (main) because they use // the trainer's cuBLAS handle which is bound to the main stream. + // + // Task 2.0 — snapshot `grad_buf` branch 0+1 slice BEFORE the IQN + // trunk gradient fires, so the post-IQN reduction sees only the + // additive delta. Snapshot runs unconditionally (even when + // `iqn_ok=false` or `gpu_iqn` is None) so the reduction below can + // also run unconditionally and keep the pinned result slot's IQN + // field (slots [0..2]) populated every step (all-zero delta → 0.0 + // ratio, which is the honest signal for a step that skipped IQN). + self.trainer.grad_decomp_snapshot_iqn() + .map_err(|e| anyhow::anyhow!("Task 2.0 grad_decomp_snapshot_iqn: {e}"))?; + if iqn_ok { if let Some(ref mut iqn) = self.gpu_iqn { let d_h_s2_ptr = iqn.d_h_s2_raw_ptr(); @@ -1469,6 +1529,13 @@ impl FusedTrainingCtx { } } + // Task 2.0 — IQN reduction. Runs unconditionally so the pinned slot + // stays populated even on steps where `iqn_ok=false`. When IQN was + // skipped the snapshot and current `grad_buf` are identical over the + // branch 0+1 range → delta norm = 0.0 (honest semantic). + self.trainer.grad_decomp_launch_iqn() + .map_err(|e| anyhow::anyhow!("Task 2.0 grad_decomp IQN: {e}"))?; + // IQN→PER loss: cast f32 per_sample_loss → f32 td_errors_buf. // IQN is now f32 native; PER system expects f32 td_errors. if let Some(ref iqn) = self.gpu_iqn { @@ -1490,6 +1557,14 @@ impl FusedTrainingCtx { let meta_q_pred = self.trainer.last_meta_q_pred(); let temporal_amp = 1.0_f32 + (meta_q_pred - 0.5).max(0.0) * 2.0; let f5_barrier_weight = 0.20_f32 * temporal_amp; + + // Task 2.0 — snapshot BEFORE the CQL gradient path. The snapshot runs + // whether or not `has_cql()` fires so the post-CQL reduction below + // can run unconditionally (keeps the pinned result slot's CQL field + // populated every step — 0.0 on steps where CQL was skipped). + self.trainer.grad_decomp_snapshot_cql() + .map_err(|e| anyhow::anyhow!("Task 2.0 grad_decomp_snapshot_cql: {e}"))?; + if self.trainer.has_cql() { match self.trainer.apply_cql_gradient(f5_barrier_weight) { Ok(true) => { @@ -1503,6 +1578,10 @@ impl FusedTrainingCtx { } } + // Task 2.0 — CQL reduction. Delta = grad_buf − snapshot over branch 0+1. + self.trainer.grad_decomp_launch_cql() + .map_err(|e| anyhow::anyhow!("Task 2.0 grad_decomp CQL: {e}"))?; + // D1/N1: Temporal self-distillation — per-step kernel pull toward the // best historical snapshot. Fully GPU-native: alpha is computed // in-kernel from ISV[LEARNING_HEALTH_INDEX] (pinned device-mapped), @@ -1547,6 +1626,13 @@ impl FusedTrainingCtx { // wiring backward adds gradient noise without value. Both calls removed // here; the kernels and Rust glue are deleted in a follow-up cleanup. + // Task 2.0 — snapshot BEFORE ensemble diversity step. Runs + // unconditionally even when `ensemble_extra_heads.is_empty()` so the + // post-ens reduction below can run unconditionally (0.0 delta on + // steps without ensemble diversity — honest semantic). + self.trainer.grad_decomp_snapshot_ens() + .map_err(|e| anyhow::anyhow!("Task 2.0 grad_decomp_snapshot_ens: {e}"))?; + // Ensemble diversity: heads 1..K-1 forward + KL gradient + trunk backward. // Runs inside aux_child graph — gradients land in grad_buf before Adam. if !self.ensemble_extra_heads.is_empty() { @@ -1554,6 +1640,10 @@ impl FusedTrainingCtx { .map_err(|e| anyhow::anyhow!("Ensemble diversity: {e}"))?; } + // Task 2.0 — ensemble reduction. Delta = grad_buf − snapshot over branch 0+1. + self.trainer.grad_decomp_launch_ens() + .map_err(|e| anyhow::anyhow!("Task 2.0 grad_decomp ens: {e}"))?; + // Regime-adaptive PER scaling. self.trainer.regime_scale_td_errors() .map_err(|e| anyhow::anyhow!("Regime PER scaling: {e}"))?; diff --git a/crates/ml/src/trainers/dqn/trainer/training_loop.rs b/crates/ml/src/trainers/dqn/trainer/training_loop.rs index de24d328d..63d9cd158 100644 --- a/crates/ml/src/trainers/dqn/trainer/training_loop.rs +++ b/crates/ml/src/trainers/dqn/trainer/training_loop.rs @@ -2228,6 +2228,28 @@ impl DQNTrainer { .map(|f| f.grad_ratio_mag_dir()) .unwrap_or(0.0); + // Task 2.0 — per-component grad decomposition (IQN/CQL/C51/Ens). + // Refresh the cached norms from the pinned result slot populated + // by the in-graph reduction kernel on the last step of this epoch, + // then compute mag/dir ratios per component. Non-fatal on failure + // (warn + stale values) — same pattern as update_q_mag_means_cached. + let (grad_mag_iqn, grad_mag_cql, grad_mag_c51, grad_mag_ens) = + if let Some(fused) = self.fused_ctx.as_mut() { + if let Err(e) = fused.refresh_grad_component_norms() { + tracing::warn!( + "HEALTH_DIAG grad_component_norms refresh failed: {e}" + ); + } + ( + fused.grad_mag_iqn_ratio(), + fused.grad_mag_cql_ratio(), + fused.grad_mag_c51_ratio(), + fused.grad_mag_ens_ratio(), + ) + } else { + (0.0_f32, 0.0_f32, 0.0_f32, 0.0_f32) + }; + // Task 0.6 — per-branch NoisyNets σ mean. H7 detection signal. // 0 = direction head, 1 = magnitude head. let (sigma_mag, sigma_dir) = { @@ -2570,7 +2592,7 @@ impl DQNTrainer { }; tracing::info!( - "HEALTH_DIAG[{}]: health={:.2} components [q_gap={:.2} q_var={:.2} atoms={:.2} grad_stable={:.2} ens_agree={:.2} grad_cos={:.2} spectral={:.2}] effective [cql_alpha={:.4} iqn_budget={:.2} cql_budget={:.2} c51_budget={:.2} tau={:.5} sarsa_tau={:.2} gamma={:.3} cf_ratio={:.2}] novels [distill={} barrier={:.3} plasticity={} ib={:.3} ensemble_collapse={:.2} contrarian={} meta_q_pred={:.2}] diag [sharpe_ema={:.3} action_entropy={:.2}] gems [g12_predictive={:.4}] mag [q_full={:.3} q_half={:.3} q_quarter={:.3} var_scale={:.3} kelly_f={:.3} avg_win_ratio={:.3} grad_ratio_mag_dir={:.4} dist_q={:.3} dist_h={:.3} dist_f={:.3}] trail [fire_q={:.3} fire_h={:.3} fire_f={:.3} hold_q={:.2} hold_h={:.2} hold_f={:.2}] noisy [vsn_mag={:.3} vsn_dir={:.3} sigma_mag={:.4} sigma_dir={:.4} drift_mag={:.3} drift_dir={:.3}] eval_dist [eq={:.3} eh={:.3} ef={:.3}] reward_contrib [popart={:.3} cf={:.3} trail_r={:.3} micro={:.3} la={:.3}] controller [anti_lr={} tau={} gamma={} clip={} cql={} cost={} fire_frac={:.2}] explore [ent_mag={:.2} ent_dir={:.2} sigma_mean={:.4}]", + "HEALTH_DIAG[{}]: health={:.2} components [q_gap={:.2} q_var={:.2} atoms={:.2} grad_stable={:.2} ens_agree={:.2} grad_cos={:.2} spectral={:.2}] effective [cql_alpha={:.4} iqn_budget={:.2} cql_budget={:.2} c51_budget={:.2} tau={:.5} sarsa_tau={:.2} gamma={:.3} cf_ratio={:.2}] novels [distill={} barrier={:.3} plasticity={} ib={:.3} ensemble_collapse={:.2} contrarian={} meta_q_pred={:.2}] diag [sharpe_ema={:.3} action_entropy={:.2}] gems [g12_predictive={:.4}] mag [q_full={:.3} q_half={:.3} q_quarter={:.3} var_scale={:.3} kelly_f={:.3} avg_win_ratio={:.3} grad_ratio_mag_dir={:.4} dist_q={:.3} dist_h={:.3} dist_f={:.3}] grad_split [iqn={:.4} cql={:.4} c51={:.4} ens={:.4}] trail [fire_q={:.3} fire_h={:.3} fire_f={:.3} hold_q={:.2} hold_h={:.2} hold_f={:.2}] noisy [vsn_mag={:.3} vsn_dir={:.3} sigma_mag={:.4} sigma_dir={:.4} drift_mag={:.3} drift_dir={:.3}] eval_dist [eq={:.3} eh={:.3} ef={:.3}] reward_contrib [popart={:.3} cf={:.3} trail_r={:.3} micro={:.3} la={:.3}] controller [anti_lr={} tau={} gamma={} clip={} cql={} cost={} fire_frac={:.2}] explore [ent_mag={:.2} ent_dir={:.2} sigma_mean={:.4}]", epoch, health_value, self.learning_health.components.q_gap_norm, @@ -2613,6 +2635,12 @@ impl DQNTrainer { avg_win_ratio, grad_ratio_mag_dir, dist_q, dist_h, dist_f, + // Task 2.0 — per-component grad decomposition (4 f32): + // iqn / cql / c51 / ens magnitude/direction grad-norm + // ratios. In-graph pinned-snapshot + reduction-kernel + // pipeline; zero-copy readback at epoch boundary. + // Feeds Task 2.1 decision tree for the H4 fix. + grad_mag_iqn, grad_mag_cql, grad_mag_c51, grad_mag_ens, // Track 1 — trail (6 f32): fire_q/h/f, hold_q/h/f (H6 — Task 0.5). // Real values from GPU experience collector's per-sample buffers. trail_rates[0], trail_rates[1], trail_rates[2],