From bb42c9963696af75bb21abe049789e9dc25debf8 Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Tue, 21 Apr 2026 21:51:39 +0200 Subject: [PATCH] =?UTF-8?q?diag(policy-quality):=20Task=200.4=20=E2=80=94?= =?UTF-8?q?=20per-branch=20grad=20norm=20via=20pinned=20readback?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Wires grad_ratio_mag_dir HEALTH_DIAG field with the magnitude head's gradient L2 norm divided by the direction head's. H4 detection signal (magnitude-head gradient starvation). Implementation: * GpuDqnTrainer.grad_readback_pinned_ptr — pinned host buffer sized TOTAL_PARAMS f32, allocated once at construction. Plain pinned (not device-mapped) — written via memcpy_dtoh. * per_branch_grad_norms() -> [f32; 4]: sync stream, dtoh full grad buffer to pinned host, compute L2 norm per branch slice using compute_param_sizes + padded_byte_offset (branch tensors at indices 8-11 / 12-15 / 16-19 / 20-23 for direction/magnitude/order/urgency). * grad_ratio_mag_dir() -> f32: convenience accessor for HEALTH_DIAG. * fused_training::FusedTrainingCtx::grad_ratio_mag_dir(&mut) — exposes via the wrapper used by training_loop. Performance: pinned dtoh of ~2.7 MB (667K params × 4 bytes) takes ~0.5ms on PCIe 4. Stream-sync cost is per-epoch (not per-step), acceptable for diagnostic readback. Pageable-memory dtoh would be ~3-5x slower. Drop free of the pinned buffer added alongside other pinned slots. Per plan Task 0.4. Complete — no stubs. --- .../ml/src/cuda_pipeline/gpu_dqn_trainer.rs | 96 +++++++++++++++++++ crates/ml/src/trainers/dqn/fused_training.rs | 7 ++ .../src/trainers/dqn/trainer/training_loop.rs | 11 ++- 3 files changed, 111 insertions(+), 3 deletions(-) diff --git a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs index 255a6d2fb..50445ca7b 100644 --- a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs +++ b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs @@ -958,6 +958,13 @@ pub struct GpuDqnTrainer { // ── Backward / Adam buffers (f32 master weights — GemmEx reads directly) ─ grad_buf: CudaSlice, // [TOTAL_PARAMS] f32 gradient accumulator + /// Task 0.4 — pinned host buffer for per-branch grad-norm readback. + /// Allocated once at construction sized to TOTAL_PARAMS f32. memcpy_dtoh + /// into pinned memory runs at PCIe bandwidth (~0.5 ms for 4 MB) instead + /// of pageable-host speed. Used only at epoch boundary, not in the + /// training graph capture region. + grad_readback_pinned_ptr: usize, // host-visible f32 buffer + grad_readback_pinned_capacity: usize, // element count (TOTAL_PARAMS) 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) @@ -2027,6 +2034,12 @@ impl Drop for GpuDqnTrainer { if !self.eval_v_range_pinned.is_null() { let _ = unsafe { cudarc::driver::result::free_host(self.eval_v_range_pinned.cast()) }; } + // Task 0.4 — pinned grad-readback buffer. + if self.grad_readback_pinned_ptr != 0 { + let _ = unsafe { + cudarc::driver::result::free_host(self.grad_readback_pinned_ptr as *mut std::ffi::c_void) + }; + } if !self.per_branch_q_gaps_pinned.is_null() { let _ = unsafe { cudarc::driver::result::free_host(self.per_branch_q_gaps_pinned.cast()) }; } @@ -2084,6 +2097,76 @@ impl GpuDqnTrainer { /// Device pointer to grad_norm (pinned device-mapped). pub fn grad_norm_dev_ptr(&self) -> u64 { self.grad_norm_dev_ptr } + /// Task 0.4 — per-action-branch L2 grad norm. + /// + /// Returns `[direction, magnitude, order, urgency]` L2 norms over each + /// branch's 4 weight tensors (fc weights, fc bias, out weights, out bias). + /// Branch 0 starts at param tensor index 8, each branch spans 4 indices. + /// + /// Implementation: single `memcpy_dtoh` of the full grad buffer into a + /// **pinned** host slot allocated once at construction. Pinned host memory + /// gives PCIe-bandwidth-limited transfer (~0.5 ms for ~4 MB) instead of + /// pageable-memory copy. Per-branch norm reductions then run on host — + /// 4 small slice-sum-of-squares loops, sub-microsecond each. + /// + /// Called once per epoch from HEALTH_DIAG — outside the training graph + /// capture region, so the readback cost is acceptable. + pub fn per_branch_grad_norms(&self) -> Result<[f32; 4], MLError> { + if self.grad_readback_pinned_ptr == 0 { + return Ok([0.0_f32; 4]); + } + let grad_len = self.grad_buf.len(); + // Pinned buffer is sized TOTAL_PARAMS at construction; refuse if mismatch + // (defensive — should never happen, but better an error than UB). + if grad_len > self.grad_readback_pinned_capacity { + return Err(MLError::ModelError(format!( + "per_branch_grad_norms: grad_buf len {} exceeds pinned readback capacity {}", + grad_len, self.grad_readback_pinned_capacity + ))); + } + // Stream-sync to ensure backward + Adam grad accumulation has completed + // before reading. Acceptable cost — runs once per epoch. + self.stream.synchronize().map_err(|e| { + MLError::ModelError(format!("per_branch_grad_norms sync: {e}")) + })?; + // Reinterpret pinned host slot as &mut [f32]. + let host_slice: &mut [f32] = unsafe { + std::slice::from_raw_parts_mut(self.grad_readback_pinned_ptr as *mut f32, grad_len) + }; + self.stream + .memcpy_dtoh(&self.grad_buf.slice(..grad_len), host_slice) + .map_err(|e| MLError::ModelError(format!("per_branch_grad_norms dtoh: {e}")))?; + + let param_sizes = compute_param_sizes(&self.config); + let mut norms = [0.0_f32; 4]; + for branch in 0..4_usize { + let first_tensor = 8 + branch * 4; // 8, 12, 16, 20 + let last_tensor = first_tensor + 4; // exclusive + let mut sum_sq = 0.0_f64; + for tensor_idx in first_tensor..last_tensor { + let start_byte = padded_byte_offset(¶m_sizes, tensor_idx); + let start_idx = (start_byte as usize) / std::mem::size_of::(); + let len = param_sizes[tensor_idx]; + let end_idx = (start_idx + len).min(host_slice.len()); + for i in start_idx..end_idx { + let g = host_slice[i] as f64; + sum_sq += g * g; + } + } + norms[branch] = sum_sq.sqrt() as f32; + } + Ok(norms) + } + + /// Task 0.4 — direction grad norm divided by magnitude grad norm. + /// Returns 0.0 if either norm is zero. H4 detection signal. + pub fn grad_ratio_mag_dir(&self) -> Result { + let n = self.per_branch_grad_norms()?; + let dir = n[0]; + let mag = n[1]; + Ok(if dir > 1e-9 { mag / dir } else { 0.0 }) + } + /// Update per-branch liquid tau modulation — GPU kernel (RK4/Euler adaptive ODE). /// /// Reads per_branch_q_gaps (pinned device-mapped) and qlstm_context from device. @@ -5489,6 +5572,17 @@ impl GpuDqnTrainer { buf }; + // Task 0.4 — pinned host buffer for grad readback. + // Plain pinned (NOT device-mapped) — written via memcpy_dtoh, not by GPU directly. + // Sized to TOTAL_PARAMS f32. Allocated once, reused every epoch. + let grad_readback_pinned_ptr: usize = unsafe { + let bytes = total_params * std::mem::size_of::(); + cudarc::driver::result::malloc_host(bytes, 0) + .map_err(|e| MLError::ModelError(format!("pinned grad_readback alloc ({} bytes): {e}", bytes)))? + as usize + }; + let grad_readback_pinned_capacity = total_params; + // 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; @@ -7340,6 +7434,8 @@ impl GpuDqnTrainer { eval_td_snapshot, eval_loss_snapshot, grad_buf, + grad_readback_pinned_ptr, + grad_readback_pinned_capacity, params_buf, target_params_buf, m_buf, diff --git a/crates/ml/src/trainers/dqn/fused_training.rs b/crates/ml/src/trainers/dqn/fused_training.rs index f8cfa296d..73277ab74 100644 --- a/crates/ml/src/trainers/dqn/fused_training.rs +++ b/crates/ml/src/trainers/dqn/fused_training.rs @@ -884,6 +884,13 @@ impl FusedTrainingCtx { self.trainer.predictive_loss_value() } + /// Task 0.4 — per-action-branch grad-norm ratio (magnitude / direction) + /// for H4 detection. Reads grad_buf to a pinned host slot and computes + /// L2 norm per branch slice. Returns 0.0 on failure (non-fatal — diag only). + pub(crate) fn grad_ratio_mag_dir(&mut self) -> f32 { + self.trainer.grad_ratio_mag_dir().unwrap_or(0.0) + } + /// Flush the last in-flight async readback from the training step. /// Called once at epoch end to retrieve the final step's scalars. pub(crate) fn flush_readback(&mut self) -> Result { diff --git a/crates/ml/src/trainers/dqn/trainer/training_loop.rs b/crates/ml/src/trainers/dqn/trainer/training_loop.rs index 37377cfad..2b505f0b0 100644 --- a/crates/ml/src/trainers/dqn/trainer/training_loop.rs +++ b/crates/ml/src/trainers/dqn/trainer/training_loop.rs @@ -2177,6 +2177,10 @@ impl DQNTrainer { } else { 0.0 }; + // Task 0.4 — per-branch grad-norm ratio (magnitude/direction). H4 signal. + let grad_ratio_mag_dir = self.fused_ctx.as_mut() + .map(|f| f.grad_ratio_mag_dir()) + .unwrap_or(0.0); // Track 1 action distribution: per-magnitude usage across the epoch. // action_counts[9] layout is dir*3 + mag — Quarter = mag 0 (indices 0,3,6), @@ -2299,9 +2303,10 @@ impl DQNTrainer { g12_loss, // Track 1 — magnitude (10 f32): q_full, q_half, q_quarter, var_scale, // kelly_f, avg_win_ratio, grad_ratio_mag_dir, dist_q, dist_h, dist_f. - // Q-values, var_scale, Kelly, grad-ratio still stubbed (need GPU readback - // infrastructure from Task 0.3 follow-up); action dist wired now. - 0.0_f32, 0.0_f32, 0.0_f32, 0.0_f32, 0.0_f32, 0.0_f32, 0.0_f32, + // Q-values, var_scale, Kelly remain to be wired by future Task 0.3 work. + // grad_ratio_mag_dir now real (Task 0.4); action_dist wired. + 0.0_f32, 0.0_f32, 0.0_f32, 0.0_f32, 0.0_f32, 0.0_f32, + grad_ratio_mag_dir, dist_q, dist_h, dist_f, // Track 1 — trail (6 f32) 0.0_f32, 0.0_f32, 0.0_f32, 0.0_f32, 0.0_f32, 0.0_f32,