From b2e695fe085c27d76ff18f25e1b63d95cfc224f4 Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Mon, 13 Apr 2026 15:18:12 +0200 Subject: [PATCH] feat(c51): per-sample support + branch scales in kernel launches, delete v_range Replace global v_range infrastructure with per-sample C51 support and per-branch gradient scales. The IQL trainer sets these pointers after construction, enabling sample-level distributional RL customization. - Add per_sample_support_ptr and branch_scales_ptr fields + setters - Update all kernel launches (C51 loss, C51 grad, CQL, MSE, expected_q) to pass per_sample_support_ptr instead of v_range_dev_ptr - Add branch_scales_ptr arg to c51_grad_kernel launch - Delete adapt_v_range_full, adapt_v_range, adapt_v_range_with_gap, v_range getter, v_range_buf_ptr getter - Delete v_range_pinned/v_range_dev_ptr/reward_std_ema fields + alloc/free Co-Authored-By: Claude Opus 4.6 (1M context) --- .../ml/src/cuda_pipeline/gpu_dqn_trainer.rs | 200 ++++++------------ 1 file changed, 69 insertions(+), 131 deletions(-) diff --git a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs index ba0aea880..a5c3c3701 100644 --- a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs +++ b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs @@ -25,7 +25,7 @@ //! (IQN trunk, attention, ensemble) into `grad_buf` via SAXPY. The Adam graph //! then sees the combined C51 + auxiliary gradients. //! -//! Per-step scalars (adam_step, tau, v_range, adaptive_clip) use pinned +//! Per-step scalars (adam_step, tau, adaptive_clip) use pinned //! device-mapped memory — GPU reads via pointer, host writes directly, //! no HtoD copies. Only batch input data needs explicit upload before replay. //! @@ -626,9 +626,10 @@ pub struct GpuDqnTrainer { tau_pinned: *mut f32, tau_dev_ptr: u64, - /// Adaptive C51 z-support [v_min, v_max] — pinned device-mapped, no HtoD copy. - v_range_pinned: *mut f32, - v_range_dev_ptr: u64, + /// Per-sample C51 support [B, 3] — pointer set by IQL trainer. + per_sample_support_ptr: u64, + /// Per-branch gradient scales [B, 4] — pointer set by IQL trainer. + branch_scales_ptr: u64, /// Adaptive gradient clip norm — pinned device-mapped host memory. /// GPU reads via device pointer (graph-safe), host writes directly (no HtoD copy). @@ -639,8 +640,6 @@ pub struct GpuDqnTrainer { grad_norm_ema: f32, /// EMA of Q-divergence for adaptive tau computation. q_div_ema: f32, - /// EMA of observed reward std for adaptive v_range floor. - reward_std_ema: f32, // ── Training state ────────────────────────────────────────────── pub(crate) adam_step: i32, @@ -963,9 +962,6 @@ impl GpuDqnTrainer { self.stream.memset_zeros(&mut self.popart_count) .map_err(|e| MLError::ModelError(format!("reset popart_count: {e}")))?; - // Reset adaptive C51 z-support — fold 2's Q-scale may differ from fold 1. - unsafe { *self.v_range_pinned = -0.005; *self.v_range_pinned.add(1) = 0.005; } - // Keep adaptive gradient clip EMA across folds — gradient scale is a property // of the model architecture and loss function, not the data window. // Resetting it would leave the first fold-2 epochs unprotected. @@ -975,7 +971,7 @@ impl GpuDqnTrainer { self.stream.memset_zeros(&mut self.q_divergence_buf) .map_err(|e| MLError::ModelError(format!("reset q_divergence: {e}")))?; - tracing::info!("Adam optimizer + PopArt + v_range + grad_clip state reset for new fold"); + tracing::info!("Adam optimizer + PopArt + grad_clip state reset for new fold"); Ok(()) } } @@ -1000,9 +996,6 @@ impl Drop for GpuDqnTrainer { if !self.tau_pinned.is_null() { let _ = unsafe { cudarc::driver::result::free_host(self.tau_pinned.cast()) }; } - if !self.v_range_pinned.is_null() { - let _ = unsafe { cudarc::driver::result::free_host(self.v_range_pinned.cast()) }; - } } } @@ -1012,6 +1005,13 @@ impl GpuDqnTrainer { self.config.batch_size } + pub fn set_per_sample_support_ptr(&mut self, ptr: u64) { + self.per_sample_support_ptr = ptr; + } + pub fn set_branch_scales_ptr(&mut self, ptr: u64) { + self.branch_scales_ptr = ptr; + } + /// Raw pointer to the pinned host readback buffer [16 × f32]. /// Used by `FusedTrainingCtx` for async diversity loss readback at offset 9. pub fn readback_pinned_ptr(&self) -> *mut f32 { @@ -1727,7 +1727,7 @@ impl GpuDqnTrainer { .arg(&b1_i32) .arg(&b2_i32) .arg(&b3_i32) - .arg(&self.v_range_dev_ptr) + .arg(&self.per_sample_support_ptr) .launch(LaunchConfig { grid_dim: (blocks, 1, 1), block_dim: (256, 1, 1), @@ -2454,24 +2454,6 @@ impl GpuDqnTrainer { dev_ptr }; - // Adaptive C51 z-support — pinned device-mapped, no HtoD copy. - let v_range_pinned: *mut f32 = unsafe { - let flags = cudarc::driver::sys::CU_MEMHOSTALLOC_DEVICEMAP; - cudarc::driver::result::malloc_host(2 * std::mem::size_of::(), flags) - .map_err(|e| MLError::ModelError(format!("pinned v_range alloc: {e}")))? - as *mut f32 - }; - unsafe { *v_range_pinned = -0.005_f32; *v_range_pinned.add(1) = 0.005_f32; } - let v_range_dev_ptr = unsafe { - let mut dev_ptr: u64 = 0; - cudarc::driver::sys::cuMemHostGetDevicePointer_v2( - &mut dev_ptr as *mut u64, - v_range_pinned.cast(), - 0, - ); - dev_ptr - }; - // Adaptive gradient clip norm — pinned device-mapped host memory. // GPU reads via device pointer (zero-copy, no HtoD), host writes directly. let adaptive_clip_init = config.max_grad_norm; @@ -2982,7 +2964,7 @@ impl GpuDqnTrainer { .map_err(|e| MLError::ModelError(format!("alloc popart_count: {e}")))?; // v8: Pessimistic Q-value initialization — shift value head bias to -0.1 - // Pessimistic Q-init REMOVED — incompatible with adaptive v_range. + // Pessimistic Q-init REMOVED — incompatible with per-sample support. // Xavier init gives near-zero value head output, correct for adaptive C51. Ok(Self { @@ -3066,13 +3048,12 @@ impl GpuDqnTrainer { t_dev_ptr, tau_pinned, tau_dev_ptr, - v_range_pinned, - v_range_dev_ptr, + per_sample_support_ptr: 0, + branch_scales_ptr: 0, adaptive_clip_pinned, adaptive_clip_dev_ptr, grad_norm_ema: 0.0, q_div_ema: 0.0, - reward_std_ema: 0.0, adam_step: 0, total_params, params_initialized: false, @@ -4131,13 +4112,55 @@ impl GpuDqnTrainer { self.config.branch_0_size + self.config.branch_1_size + self.config.branch_2_size + self.config.branch_3_size } - /// Reference to the Q-value output buffer from the last `forward_only_q()` call. + /// Reference to the Q-value output buffer. /// - /// Shape: `[B, total_actions]`. Only valid after `forward_only_q()`. + /// Shape: `[B, total_actions]`. Valid after `populate_q_out()` or `replay_forward_for_q_values()`. pub fn q_out_buf(&self) -> &CudaSlice { &self.q_out_buf } + /// Convert current logits → Q-values in q_out_buf. + /// + /// Runs compute_expected_q on the logits already in on_v_logits_buf / on_b_logits_buf + /// (populated by graph_forward replay). Does NOT replay any graph — just the + /// lightweight expected_q kernel. + /// + /// Used by IQL to get Q(s, a_taken) for expectile regression target. + pub fn populate_q_out(&self, batch_size: usize) -> Result<&CudaSlice, MLError> { + let n = batch_size as i32; + let na = self.config.num_atoms as i32; + let b0 = self.config.branch_0_size as i32; + let b1 = self.config.branch_1_size as i32; + let b2 = self.config.branch_2_size as i32; + let b3 = self.config.branch_3_size as i32; + let block_dim = 256_u32; + let grid_dim = ((batch_size as u32 + block_dim - 1) / block_dim).max(1); + let q_out_ptr = self.q_out_buf.raw_ptr(); + let null_atom_stats = 0u64; + unsafe { + self.stream + .launch_builder(&self.expected_q_kernel) + .arg(&self.ptrs.on_v_logits_buf) + .arg(&self.ptrs.on_b_logits_buf) + .arg(&q_out_ptr) + .arg(&n) + .arg(&na) + .arg(&b0) + .arg(&b1) + .arg(&b2) + .arg(&b3) + .arg(&self.per_sample_support_ptr) + .arg(&null_atom_stats) + .launch(LaunchConfig { + grid_dim: (grid_dim, 1, 1), + block_dim: (block_dim, 1, 1), + shared_mem_bytes: 0, + }) + .map_err(|e| MLError::ModelError(format!("populate_q_out expected_q: {e}")))?; + } + Ok(&self.q_out_buf) + } + /// Run the cuBLAS forward pass without graph (used during graph capture). fn replay_forward_ungraphed(&self) -> Result<(), MLError> { let param_sizes = compute_param_sizes(&self.config); @@ -4233,7 +4256,7 @@ impl GpuDqnTrainer { .arg(&b1) .arg(&b2) .arg(&b3) - .arg(&self.v_range_dev_ptr) + .arg(&self.per_sample_support_ptr) .arg(&null_atom_stats) .launch(LaunchConfig { grid_dim: (grid_dim, 1, 1), @@ -4385,7 +4408,7 @@ impl GpuDqnTrainer { .arg(&b1) .arg(&b2) .arg(&b3) - .arg(&self.v_range_dev_ptr) + .arg(&self.per_sample_support_ptr) .arg(&atom_stats_ptr) .launch(LaunchConfig { grid_dim: (eq_blocks, 1, 1), @@ -5334,11 +5357,11 @@ impl GpuDqnTrainer { // ── Curiosity Q-penalty (2) ── .arg(&self.curiosity_error_buf) .arg(&self.config.curiosity_q_penalty_lambda) - // ── Config (8 — v_range_buf replaces v_min+v_max) ── + // ── Config (8 — per_sample_support replaces v_min+v_max) ── .arg(&gamma) .arg(&batch_i32) .arg(&na_i32) - .arg(&self.v_range_dev_ptr) + .arg(&self.per_sample_support_ptr) .arg(&b0_i32) .arg(&b1_i32) .arg(&b2_i32) @@ -5478,6 +5501,7 @@ impl GpuDqnTrainer { .arg(&b3_i32) .arg(&total_branch_atoms_i32) .arg(&entropy_coeff) + .arg(&self.branch_scales_ptr) .launch(LaunchConfig { grid_dim: (blocks, 1, 1), block_dim: (256, 1, 1), @@ -5577,11 +5601,11 @@ impl GpuDqnTrainer { // ── Curiosity Q-penalty (2) ── .arg(&self.curiosity_error_buf) .arg(&self.config.curiosity_q_penalty_lambda) - // ── Config (8 — v_range_buf replaces v_min+v_max) ── + // ── Config (8 — per_sample_support replaces v_min+v_max) ── .arg(&gamma) .arg(&batch_i32) .arg(&na_i32) - .arg(&self.v_range_dev_ptr) + .arg(&self.per_sample_support_ptr) .arg(&b0_i32) .arg(&b1_i32) .arg(&b2_i32) @@ -5666,7 +5690,7 @@ impl GpuDqnTrainer { .arg(&b2_i32) .arg(&b3_i32) .arg(&total_branch_atoms_i32) - .arg(&self.v_range_dev_ptr) + .arg(&self.per_sample_support_ptr) .launch(LaunchConfig { grid_dim: (blocks, 1, 1), block_dim: (256, 1, 1), @@ -6359,92 +6383,6 @@ impl GpuDqnTrainer { Ok(()) } - // ── Adaptive C51 z-support ───────────────────────────────────── - - /// Q-stats-driven v_range: center at q_mean, width = 3σ + Bellman headroom. - /// - /// Replaces the old heuristic expand/contract with a single data-driven formula. - /// - 3σ coverage captures 99.7% of Q-distribution - /// - Bellman headroom (γ * 3σ) prevents projection clipping on reward shift - /// - MIN_RANGE floor prevents atom collapse when Q-variance is near-zero - /// - Mean-centered: no wasted atoms on impossible Q-values - pub fn adapt_v_range(&mut self, q_mean: f32, q_variance: f32) -> bool { - self.adapt_v_range_full(q_mean, q_variance, 0.0, 0.0) - } - - pub fn adapt_v_range_with_gap(&mut self, q_mean: f32, q_variance: f32, q_gap: f32) -> bool { - self.adapt_v_range_full(q_mean, q_variance, q_gap, 0.0) - } - - /// Fully adaptive v_range from Q-stats + reward scale. - /// - /// The reward_std parameter comes from the experience collector's observed - /// reward statistics. It varies 940× between smoketest (0.007) and production - /// (6.57) — hardcoding it was fundamentally wrong. - /// - /// The v_range floor = reward_std (single Bellman step shift magnitude). - /// This ensures the C51 projection can shift atoms by at least one reward - /// unit, breaking the stable fixed point at any scale. - pub fn adapt_v_range_full(&mut self, q_mean: f32, q_variance: f32, q_gap: f32, reward_std: f32) -> bool { - const SIGMA_COVERAGE: f32 = 3.0; - const GAP_ATOM_TARGET: f32 = 10.0; - - let q_std = (q_variance.max(0.0) + 1e-10).sqrt(); - let half_width = SIGMA_COVERAGE * q_std; - let bellman_headroom = self.config.gamma * half_width; - let total_half = half_width + bellman_headroom; - - let target_min = q_mean - total_half; - let target_max = q_mean + total_half; - - // Reward-scale floor: derived from ACTUAL observed reward_std. - // Smoketest: reward_std ≈ 0.007 → floor = 0.007 - // Production: reward_std ≈ 6.57 → floor = 6.57 - // Falls back to EMA of observed reward_std (self.reward_std_ema) - // when per-step reward_std is not available. - let effective_reward_std = if reward_std > 1e-8 { - // Update EMA with fresh observation - const R_BETA: f32 = 0.99; - if self.reward_std_ema <= 1e-8 { - self.reward_std_ema = reward_std; - } else { - self.reward_std_ema = R_BETA * self.reward_std_ema + (1.0 - R_BETA) * reward_std; - } - self.reward_std_ema - } else if self.reward_std_ema > 1e-8 { - self.reward_std_ema // use cached EMA - } else { - 0.01 // absolute fallback before first observation - }; - - // Decay floor as Q-values mature (same exp decay as before, but - // using the ACTUAL reward scale instead of hardcoded 0.01) - let maturity = (q_mean.abs() / effective_reward_std.max(1e-8)).min(5.0); - let num_atoms = self.config.num_atoms.max(2) as f32; - let gap_based_range = q_gap * (num_atoms / GAP_ATOM_TARGET); - let decayed_floor = effective_reward_std * (-maturity).exp(); - let adaptive_min = decayed_floor - .max(gap_based_range) - .max(target_max - target_min); - let (final_min, final_max) = if (target_max - target_min) < adaptive_min { - (q_mean - adaptive_min * 0.5, q_mean + adaptive_min * 0.5) - } else { - (target_min, target_max) - }; - - let cur = self.v_range(); - let changed = (final_min - cur[0]).abs() > 1e-6 || (final_max - cur[1]).abs() > 1e-6; - if changed { - unsafe { *self.v_range_pinned = final_min; *self.v_range_pinned.add(1) = final_max; } - } - changed - } - - pub fn v_range(&self) -> [f32; 2] { - unsafe { [*self.v_range_pinned, *self.v_range_pinned.add(1)] } - } - pub fn v_range_buf_ptr(&self) -> u64 { self.v_range_dev_ptr } - // ── Adaptive gradient clipping ────────────────────────────────── /// Update EMA-based adaptive gradient clip norm.