diff --git a/crates/ml/src/cuda_pipeline/gpu_iql_trainer.rs b/crates/ml/src/cuda_pipeline/gpu_iql_trainer.rs index e747b82a7..02583eeb2 100644 --- a/crates/ml/src/cuda_pipeline/gpu_iql_trainer.rs +++ b/crates/ml/src/cuda_pipeline/gpu_iql_trainer.rs @@ -3,12 +3,13 @@ //! GPU-accelerated IQL (Implicit Q-Learning) value network trainer. //! //! Implements the value network component of IQL (Kostrikov et al., 2021) -//! as a fused CUDA kernel pipeline: +//! as a fused CUDA kernel pipeline — zero atomicAdd, fully deterministic: //! -//! 1. **Forward + Expectile Loss** -- V(s) MLP forward + asymmetric loss -//! 2. **Gradient Norm** -- L2 norm for gradient clipping -//! 3. **Backward** -- backprop through SiLU MLP, atomicAdd gradients -//! 4. **Adam** -- AdamW update with gradient clipping +//! 1. **Forward + Expectile Loss** — V(s) MLP forward + asymmetric loss +//! 2. **Backward (per-sample)** — backprop into per-sample gradient buffer +//! 3. **Weight Grad Reduce** — deterministic sum across samples +//! 4. **Grad Norm Phase 1+2** — two-phase L2 norm (no atomicAdd) +//! 5. **Adam** — AdamW update with gradient clipping //! //! ## Architecture //! @@ -17,9 +18,9 @@ //! //! ## Integration //! -//! Called after the main DQN training step in `FusedTrainingCtx::run_full_step()`. -//! The value loss trains V(s) to approximate the expectile of Q(s,a), and the -//! advantage weights `exp(beta * (Q(s,a) - V(s)))` can modulate the DQN policy. +//! Called after the main DQN training step in `FusedTrainingCtx::submit_aux_ops()`. +//! The V(s) output drives advantage weights `exp(beta * (Q(s,a) - V(s)))` that +//! modulate PER priorities for advantage-weighted replay. use std::sync::Arc; @@ -57,6 +58,18 @@ pub struct GpuIqlConfig { pub weight_decay: f32, /// Maximum gradient L2 norm for clipping. pub max_grad_norm: f32, + /// Number of C51 atoms. + pub num_atoms: usize, + /// Total factored actions (b0*b1*b2*b3). + pub total_actions: usize, + /// Branch sizes [b0, b1, b2, b3] for advantage decomposition. + pub branch_sizes: [usize; 4], + /// Discount factor for Bellman headroom in per-sample support. + pub gamma: f32, + /// Staleness decay rate for PER modulation. + pub staleness_lambda: f32, + /// Staleness age normalizer (steps). + pub staleness_tau: f32, } impl Default for GpuIqlConfig { @@ -73,6 +86,12 @@ impl Default for GpuIqlConfig { epsilon: 1e-8, weight_decay: 1e-5, max_grad_norm: 1.0, + num_atoms: 51, + total_actions: 81, + branch_sizes: [3, 3, 3, 3], + gamma: 0.99, + staleness_lambda: 0.001, + staleness_tau: 10000.0, } } } @@ -94,9 +113,11 @@ impl GpuIqlConfig { /// GPU-accelerated IQL value network trainer. /// /// Owns pre-allocated GPU buffers for the V(s) network weights, Adam state, -/// activation saves, and the compiled CUDA kernels. Operates entirely on GPU -/// with no per-step CPU-GPU data transfers (states and Q-values are already -/// on device from the DQN training step). +/// per-sample gradients, and compiled CUDA kernels. Operates entirely on GPU +/// with no per-step CPU-GPU data transfers. +/// +/// Zero atomicAdd — all gradient accumulation uses per-sample buffers with +/// deterministic cross-sample reduction. #[allow(missing_debug_implementations)] pub struct GpuIqlTrainer { config: GpuIqlConfig, @@ -104,19 +125,33 @@ pub struct GpuIqlTrainer { // ── Compiled kernels ──────────────────────────────────────────── forward_loss_kernel: CudaFunction, - backward_kernel: CudaFunction, - grad_norm_kernel: CudaFunction, + backward_per_sample_kernel: CudaFunction, + weight_grad_reduce_kernel: CudaFunction, + loss_reduce_kernel: CudaFunction, + grad_norm_phase1_kernel: CudaFunction, + grad_norm_phase2_kernel: CudaFunction, adam_kernel: CudaFunction, forward_kernel: CudaFunction, + gather_q_taken_kernel: CudaFunction, + advantage_weight_kernel: CudaFunction, + modulate_td_kernel: CudaFunction, + adv_variance_kernel: CudaFunction, + per_sample_support_kernel: CudaFunction, + branch_advantage_kernel: CudaFunction, + expectile_gap_kernel: CudaFunction, + gap_mean_kernel: CudaFunction, + per_sample_epsilon_kernel: CudaFunction, // ── V network parameters (flat f32 on GPU) ───────────────────── params_buf: CudaSlice, // ── Adam optimizer state ──────────────────────────────────────── - m_buf: CudaSlice, // first moment - v_buf: CudaSlice, // second moment - grad_buf: CudaSlice, // gradient accumulator - grad_norm_buf: CudaSlice, // [1] gradient L2 norm + m_buf: CudaSlice, // first moment + v_buf: CudaSlice, // second moment + grad_buf: CudaSlice, // reduced gradient [total_params] + grads_per_sample: CudaSlice, // per-sample gradients [B * total_params] + grad_norm_buf: CudaSlice, // [1] gradient L2 norm (sum-of-squares) + grad_norm_partials: CudaSlice, // [num_blocks] partial sums for phase1 // ── Activation save buffers (forward -> backward) ─────────────── save_pre1: CudaSlice, // [B, H] pre-activation layer 1 @@ -125,22 +160,33 @@ pub struct GpuIqlTrainer { save_h2: CudaSlice, // [B, H] post-activation layer 2 // ── Output buffers ────────────────────────────────────────────── - v_out_buf: CudaSlice, // [B] V(s) predictions - loss_buf: CudaSlice, // [B] per-sample loss - total_loss_buf: CudaSlice, // [1] batch-mean loss + v_out_buf: CudaSlice, // [B] V(s) predictions + loss_buf: CudaSlice, // [B] per-sample loss + total_loss_buf: CudaSlice, // [1] batch-mean loss + q_taken_buf: CudaSlice, // [B] Q(s, a_taken) gathered from q_out + advantage_weights_buf: CudaSlice, // [B] advantage weights + + // ── New integration buffers ───────────────────────────────────── + adv_stats_buf: CudaSlice, // [2] mean, variance + adv_sigma_ema: f32, // host-side EMA + per_sample_support_buf: CudaSlice, // [B, 3] + branch_scales_buf: CudaSlice, // [B, 4] + expectile_gap_buf: CudaSlice, // [B] + gap_mean_buf: CudaSlice, // [1] + per_sample_epsilon_buf: CudaSlice, // [B] // ── Training state ────────────────────────────────────────────── adam_step: i32, t_buf: CudaSlice, total_params: usize, + grad_norm_blocks: usize, } impl GpuIqlTrainer { /// Create a new GPU IQL trainer. /// - /// Compiles the 5 CUDA kernels (forward+loss, backward, grad_norm, adam, - /// forward-only), initializes V(s) weights with Xavier/Glorot uniform, - /// and pre-allocates all GPU buffers. + /// Compiles 9 CUDA kernels, initializes V(s) weights with Xavier/Glorot, + /// and pre-allocates all GPU buffers including per-sample gradient storage. pub fn new( stream: Arc, config: GpuIqlConfig, @@ -149,9 +195,8 @@ impl GpuIqlTrainer { let b = config.batch_size; let h = config.value_hidden_dim; - // Compile all 5 kernels - let (forward_loss_kernel, backward_kernel, grad_norm_kernel, adam_kernel, forward_kernel) = - compile_iql_kernels(&stream, &config)?; + // Compile all 9 kernels + let kernels = compile_iql_kernels(&stream, &config)?; // Allocate parameter buffer and initialize with Xavier/Glorot let params_buf = init_xavier_weights(&stream, &config)?; @@ -160,7 +205,12 @@ impl GpuIqlTrainer { let m_buf = alloc_f32(&stream, total_params, "iql_m")?; let v_buf = alloc_f32(&stream, total_params, "iql_v")?; let grad_buf = alloc_f32(&stream, total_params, "iql_grad")?; + let grads_per_sample = alloc_f32(&stream, b * total_params, "iql_grads_per_sample")?; + + // Grad norm buffers (two-phase) + let grad_norm_blocks = (total_params + 255) / 256; let grad_norm_buf = alloc_f32(&stream, 1, "iql_grad_norm")?; + let grad_norm_partials = alloc_f32(&stream, grad_norm_blocks, "iql_grad_norm_partials")?; // Allocate activation save buffers let save_pre1 = alloc_f32(&stream, b * h, "iql_save_pre1")?; @@ -172,11 +222,33 @@ impl GpuIqlTrainer { let v_out_buf = alloc_f32(&stream, b, "iql_v_out")?; let loss_buf = alloc_f32(&stream, b, "iql_loss")?; let total_loss_buf = alloc_f32(&stream, 1, "iql_total_loss")?; + let q_taken_buf = alloc_f32(&stream, b, "iql_q_taken")?; + let advantage_weights_buf = alloc_f32(&stream, b, "iql_adv_weights")?; - let vram_bytes = (total_params * 4 // params + m + v + grad = 4 copies - + b * h * 4 // 4 activation buffers - + b * 2 + 2 // v_out + loss + total_loss + grad_norm - ) * 4; + // New integration buffers + let adv_stats_buf = alloc_f32(&stream, 2, "iql_adv_stats")?; + let mut per_sample_support_buf = alloc_f32(&stream, b * 3, "iql_per_sample_support")?; + let branch_scales_buf = alloc_f32(&stream, b * 4, "iql_branch_scales")?; + let expectile_gap_buf = alloc_f32(&stream, b, "iql_expectile_gap")?; + let gap_mean_buf = alloc_f32(&stream, 1, "iql_gap_mean")?; + let per_sample_epsilon_buf = alloc_f32(&stream, b, "iql_per_sample_epsilon")?; + + // Seed per_sample_support with defaults for step 0 + let na = config.num_atoms.max(2) as f32; + let default_delta_z = 2.0 / (na - 1.0); + let mut default_support = vec![0.0_f32; b * 3]; + for i in 0..b { + default_support[i * 3] = -1.0; + default_support[i * 3 + 1] = 1.0; + default_support[i * 3 + 2] = default_delta_z; + } + super::htod_f32(&stream, &default_support, &mut per_sample_support_buf)?; + + let per_sample_vram = b * total_params * 4; + let vram_bytes = total_params * 4 * 4 // params + m + v + grad + + per_sample_vram // per-sample gradients + + b * h * 4 * 4 // 4 activation buffers + + (b * 3 + 2) * 4; // v_out + loss + total_loss + adv_weights + grad_norm info!( state_dim = config.state_dim, @@ -185,8 +257,9 @@ impl GpuIqlTrainer { total_params, expectile_tau = config.expectile_tau, advantage_temperature = config.advantage_temperature, + per_sample_grad_mb = per_sample_vram / (1024 * 1024), vram_kb = vram_bytes / 1024, - "GpuIqlTrainer initialized: 5 kernels compiled, V(s) weights Xavier-initialized" + "GpuIqlTrainer initialized: 17 kernels, zero atomicAdd, V(s) Xavier-initialized" ); let t_buf = stream.alloc_zeros::(1) @@ -194,16 +267,30 @@ impl GpuIqlTrainer { Ok(Self { config, stream, - forward_loss_kernel, - backward_kernel, - grad_norm_kernel, - adam_kernel, - forward_kernel, + forward_loss_kernel: kernels.forward_loss, + backward_per_sample_kernel: kernels.backward_per_sample, + weight_grad_reduce_kernel: kernels.weight_grad_reduce, + loss_reduce_kernel: kernels.loss_reduce, + grad_norm_phase1_kernel: kernels.grad_norm_phase1, + grad_norm_phase2_kernel: kernels.grad_norm_phase2, + adam_kernel: kernels.adam, + forward_kernel: kernels.forward, + gather_q_taken_kernel: kernels.gather_q_taken, + advantage_weight_kernel: kernels.advantage_weight, + modulate_td_kernel: kernels.modulate_td, + adv_variance_kernel: kernels.adv_variance, + per_sample_support_kernel: kernels.per_sample_support, + branch_advantage_kernel: kernels.branch_advantage, + expectile_gap_kernel: kernels.expectile_gap, + gap_mean_kernel: kernels.gap_mean, + per_sample_epsilon_kernel: kernels.per_sample_epsilon, params_buf, m_buf, v_buf, grad_buf, + grads_per_sample, grad_norm_buf, + grad_norm_partials, save_pre1, save_pre2, save_h1, @@ -211,122 +298,156 @@ impl GpuIqlTrainer { v_out_buf, loss_buf, total_loss_buf, + q_taken_buf, + advantage_weights_buf, + adv_stats_buf, + adv_sigma_ema: 0.0, + per_sample_support_buf, + branch_scales_buf, + expectile_gap_buf, + gap_mean_buf, + per_sample_epsilon_buf, adam_step: 0, t_buf, total_params, + grad_norm_blocks, }) } - /// Run one IQL value network training step. + /// Run one IQL value network training step — fully deterministic. /// - /// Executes the full forward + expectile loss + backward + Adam cycle: - /// 1. Zero total_loss and grad_norm scalars - /// 2. Forward + loss kernel (8 warps per sample) - /// 3. Backward kernel (8 warps per sample, atomicAdd gradients) - /// 4. Grad norm kernel (parallel reduction) - /// 5. Adam update kernel (one thread per parameter) + /// Pipeline: forward+loss → backward_per_sample → weight_grad_reduce + /// → loss_reduce → grad_norm phase1+2 → Adam. /// - /// Returns the batch-mean expectile loss (read back from GPU). + /// Zero atomicAdd. All cross-sample reduction is deterministic. /// - /// # Arguments - /// - /// * `states_f32` - F32 CudaSlice `[B * STATE_DIM]` on CUDA device - /// * `q_values_f32` - F32 CudaSlice `[B]` on CUDA device (target Q-values from DQN) - /// - /// Takes pre-converted F32 CudaSlice buffers directly — zero Candle - /// tensor involvement. + /// `gather_q_taken()` MUST be called first to populate `q_taken_buf`. + /// The expectile regression target is Q(s, a_taken) from the DQN. pub fn train_value_step( &mut self, states_f32: &CudaSlice, - q_values_f32: &CudaSlice, - ) -> Result { + ) -> Result<(), MLError> { let b = self.config.batch_size; - - // Zero total_loss and grad_norm before kernel launches - let zero_f32 = [0.0_f32]; - super::htod_f32(&self.stream, &zero_f32, &mut self.total_loss_buf)?; - super::htod_f32(&self.stream, &zero_f32, &mut self.grad_norm_buf)?; - + let state_dim_i32 = self.config.state_dim as i32; let batch_size_i32 = b as i32; let total_params_i32 = self.total_params as i32; + let expectile_tau = self.config.expectile_tau; - // 1. Forward + loss kernel (256 threads: 8 warps per sample) - // Shared memory: warp_sums[8] in f32 for block-level reduction + // 1. Forward + loss kernel (256 threads per sample) let fwd_shmem = 8 * std::mem::size_of::(); - let fwd_config = LaunchConfig { - grid_dim: (b as u32, 1, 1), - block_dim: (256, 1, 1), - shared_mem_bytes: fwd_shmem as u32, - }; - - // Safety: states_f32 and q_values_f32 are valid F32 CudaSlice buffers - // on the same CUDA context. All internal buffers are pre-allocated. - let state_dim_i32 = self.config.state_dim as i32; unsafe { self.stream .launch_builder(&self.forward_loss_kernel) .arg(states_f32) - .arg(q_values_f32) + .arg(&self.q_taken_buf) .arg(&self.params_buf) .arg(&mut self.v_out_buf) .arg(&mut self.loss_buf) - .arg(&mut self.total_loss_buf) .arg(&mut self.save_pre1) .arg(&mut self.save_pre2) .arg(&mut self.save_h1) .arg(&mut self.save_h2) .arg(&batch_size_i32) .arg(&state_dim_i32) - .launch(fwd_config) - .map_err(|e| MLError::ModelError(format!("IQL forward+loss kernel: {e}")))?; + .arg(&expectile_tau) + .launch(LaunchConfig { + grid_dim: (b as u32, 1, 1), + block_dim: (256, 1, 1), + shared_mem_bytes: fwd_shmem as u32, + }) + .map_err(|e| MLError::ModelError(format!("IQL forward+loss: {e}")))?; } - // 2. Backward kernel (256 threads: 8 warps per sample) - let bwd_config = LaunchConfig { - grid_dim: (b as u32, 1, 1), - block_dim: (256, 1, 1), - shared_mem_bytes: 0, - }; - - // Safety: all pointers are valid GPU allocations on the same stream. + // 2. Backward per-sample (256 threads per sample, no atomicAdd) unsafe { self.stream - .launch_builder(&self.backward_kernel) + .launch_builder(&self.backward_per_sample_kernel) .arg(states_f32) - .arg(q_values_f32) + .arg(&self.q_taken_buf) .arg(&self.v_out_buf) .arg(&self.params_buf) .arg(&self.save_pre1) .arg(&self.save_pre2) .arg(&self.save_h1) .arg(&self.save_h2) - .arg(&mut self.grad_buf) + .arg(&mut self.grads_per_sample) .arg(&batch_size_i32) .arg(&state_dim_i32) - .launch(bwd_config) - .map_err(|e| MLError::ModelError(format!("IQL backward kernel: {e}")))?; + .arg(&total_params_i32) + .arg(&expectile_tau) + .launch(LaunchConfig { + grid_dim: (b as u32, 1, 1), + block_dim: (256, 1, 1), + shared_mem_bytes: 0, + }) + .map_err(|e| MLError::ModelError(format!("IQL backward_per_sample: {e}")))?; } - // 3. Grad norm kernel - let norm_blocks = (self.total_params + 255) / 256; - let norm_config = LaunchConfig { - grid_dim: (norm_blocks as u32, 1, 1), - block_dim: (256, 1, 1), - shared_mem_bytes: 0, - }; - - // Safety: grad_buf and grad_norm_buf are valid allocations. + // 3. Weight grad reduce (deterministic sum across samples) + let reduce_blocks = (self.total_params + 255) / 256; unsafe { self.stream - .launch_builder(&self.grad_norm_kernel) - .arg(&self.grad_buf) - .arg(&mut self.grad_norm_buf) + .launch_builder(&self.weight_grad_reduce_kernel) + .arg(&self.grads_per_sample) + .arg(&mut self.grad_buf) + .arg(&batch_size_i32) .arg(&total_params_i32) - .launch(norm_config) - .map_err(|e| MLError::ModelError(format!("IQL grad_norm kernel: {e}")))?; + .launch(LaunchConfig { + grid_dim: (reduce_blocks as u32, 1, 1), + block_dim: (256, 1, 1), + shared_mem_bytes: 0, + }) + .map_err(|e| MLError::ModelError(format!("IQL weight_grad_reduce: {e}")))?; } - // 4. Adam update kernel — adam_step on GPU (async HtoD) + // 4. Loss reduce (deterministic sequential sum) + unsafe { + self.stream + .launch_builder(&self.loss_reduce_kernel) + .arg(&self.loss_buf) + .arg(&mut self.total_loss_buf) + .arg(&batch_size_i32) + .launch(LaunchConfig { + grid_dim: (1, 1, 1), + block_dim: (1, 1, 1), + shared_mem_bytes: 0, + }) + .map_err(|e| MLError::ModelError(format!("IQL loss_reduce: {e}")))?; + } + + // 5. Grad norm phase 1 (per-block partial sums) + let norm_blocks = self.grad_norm_blocks; + unsafe { + self.stream + .launch_builder(&self.grad_norm_phase1_kernel) + .arg(&self.grad_buf) + .arg(&mut self.grad_norm_partials) + .arg(&total_params_i32) + .launch(LaunchConfig { + grid_dim: (norm_blocks as u32, 1, 1), + block_dim: (256, 1, 1), + shared_mem_bytes: 256 * std::mem::size_of::() as u32, + }) + .map_err(|e| MLError::ModelError(format!("IQL grad_norm_phase1: {e}")))?; + } + + // 6. Grad norm phase 2 (sequential sum of partials) + let norm_blocks_i32 = norm_blocks as i32; + unsafe { + self.stream + .launch_builder(&self.grad_norm_phase2_kernel) + .arg(&self.grad_norm_partials) + .arg(&mut self.grad_norm_buf) + .arg(&norm_blocks_i32) + .launch(LaunchConfig { + grid_dim: (1, 1, 1), + block_dim: (1, 1, 1), + shared_mem_bytes: 0, + }) + .map_err(|e| MLError::ModelError(format!("IQL grad_norm_phase2: {e}")))?; + } + + // 7. Adam update — adam_step on GPU (async HtoD) self.adam_step += 1; unsafe { cudarc::driver::sys::cuMemcpyHtoDAsync_v2( @@ -337,12 +458,6 @@ impl GpuIqlTrainer { ); } let adam_blocks = (self.total_params + 255) / 256; - let adam_config = LaunchConfig { - grid_dim: (adam_blocks as u32, 1, 1), - block_dim: (256, 1, 1), - shared_mem_bytes: 0, - }; - let lr = self.config.lr; let beta1 = self.config.beta1; let beta2 = self.config.beta2; @@ -367,13 +482,109 @@ impl GpuIqlTrainer { .arg(&mgn) .arg(&t_ptr) .arg(&total_params_i32) - .launch(adam_config) - .map_err(|e| MLError::ModelError(format!("IQL adam kernel: {e}")))?; + .launch(LaunchConfig { + grid_dim: (adam_blocks as u32, 1, 1), + block_dim: (256, 1, 1), + shared_mem_bytes: 0, + }) + .map_err(|e| MLError::ModelError(format!("IQL adam: {e}")))?; } - // Loss stays on GPU — no per-step DtoH. Caller discards the value. - // Removing the synchronous readback unblocks CUDA Graph capture. - Ok(0.0) + Ok(()) + } + + /// Gather Q(s, a_taken) from the full Q-value buffer into q_taken_buf. + /// + /// Must be called BEFORE `train_value_step` so IQL trains on Q-values, + /// not raw rewards. + pub fn gather_q_taken( + &mut self, + q_out_buf: &CudaSlice, + actions_buf: &CudaSlice, + total_actions: usize, + ) -> Result<&CudaSlice, MLError> { + let b = self.config.batch_size; + let batch_size_i32 = b as i32; + let total_actions_i32 = total_actions as i32; + let blocks = (b + 255) / 256; + + unsafe { + self.stream + .launch_builder(&self.gather_q_taken_kernel) + .arg(q_out_buf) + .arg(actions_buf) + .arg(&mut self.q_taken_buf) + .arg(&batch_size_i32) + .arg(&total_actions_i32) + .launch(LaunchConfig { + grid_dim: (blocks as u32, 1, 1), + block_dim: (256, 1, 1), + shared_mem_bytes: 0, + }) + .map_err(|e| MLError::ModelError(format!("IQL gather_q_taken: {e}")))?; + } + + Ok(&self.q_taken_buf) + } + + /// Reference to the gathered Q(s, a_taken) buffer [B]. + pub fn q_taken_buf(&self) -> &CudaSlice { + &self.q_taken_buf + } + + /// Compute IQL advantage weights: w[b] = exp(beta * (Q_taken - V(s))). + /// + /// Must be called AFTER `train_value_step` (v_out_buf is populated) + /// and AFTER `gather_q_taken` (q_taken_buf is populated). + /// + /// Advantage weights are written to `advantage_weights_buf [B]` and + /// can be used to modulate PER priorities. + pub fn compute_advantage_weights( + &mut self, + q_out_buf: &CudaSlice, + actions_buf: &CudaSlice, + total_actions: usize, + ) -> Result<(), MLError> { + let b = self.config.batch_size; + let beta = self.config.advantage_temperature; + let batch_size_i32 = b as i32; + let total_actions_i32 = total_actions as i32; + + let blocks = (b + 255) / 256; + unsafe { + self.stream + .launch_builder(&self.advantage_weight_kernel) + .arg(q_out_buf) + .arg(actions_buf) + .arg(&self.v_out_buf) + .arg(&mut self.advantage_weights_buf) + .arg(&beta) + .arg(&batch_size_i32) + .arg(&total_actions_i32) + .launch(LaunchConfig { + grid_dim: (blocks as u32, 1, 1), + block_dim: (256, 1, 1), + shared_mem_bytes: 0, + }) + .map_err(|e| MLError::ModelError(format!("IQL advantage_weights: {e}")))?; + } + + Ok(()) + } + + /// Raw device pointer to V(s) output buffer [B]. + pub fn v_out_ptr(&self) -> u64 { + self.v_out_buf.raw_ptr() + } + + /// Raw device pointer to advantage weights buffer [B]. + pub fn advantage_weights_ptr(&self) -> u64 { + self.advantage_weights_buf.raw_ptr() + } + + /// Raw device pointer to total loss buffer [1]. + pub fn total_loss_ptr(&self) -> u64 { + self.total_loss_buf.raw_ptr() } /// Current Adam step count. @@ -386,6 +597,260 @@ impl GpuIqlTrainer { pub fn increment_adam_step(&mut self) { self.adam_step += 1; } + + // ── New integration launchers ──────────────────────────────────── + + /// Modulate td_errors in-place with advantage weights and staleness decay. + pub fn modulate_td_errors( + &self, + td_errors: &mut CudaSlice, + indices_ptr: u64, + write_pos: i32, + capacity: i32, + ) -> Result<(), MLError> { + let b = self.config.batch_size; + let batch_i32 = b as i32; + let beta = self.config.advantage_temperature; + let sigma = self.adv_sigma_ema; + let lambda = self.config.staleness_lambda; + let tau = self.config.staleness_tau; + let blocks = (b + 255) / 256; + + unsafe { + self.stream + .launch_builder(&self.modulate_td_kernel) + .arg(td_errors) + .arg(&self.advantage_weights_buf) + .arg(&indices_ptr) + .arg(&sigma) + .arg(&beta) + .arg(&lambda) + .arg(&tau) + .arg(&write_pos) + .arg(&capacity) + .arg(&batch_i32) + .launch(LaunchConfig { + grid_dim: (blocks as u32, 1, 1), + block_dim: (256, 1, 1), + shared_mem_bytes: 0, + }) + .map_err(|e| MLError::ModelError(format!("IQL modulate_td_errors: {e}")))?; + } + Ok(()) + } + + /// Update the EMA of advantage standard deviation from the current batch. + pub fn update_adv_sigma(&mut self) -> Result<(), MLError> { + let batch_i32 = self.config.batch_size as i32; + + unsafe { + self.stream + .launch_builder(&self.adv_variance_kernel) + .arg(&self.advantage_weights_buf) + .arg(&mut self.adv_stats_buf) + .arg(&batch_i32) + .launch(LaunchConfig { + grid_dim: (1, 1, 1), + block_dim: (1, 1, 1), + shared_mem_bytes: 0, + }) + .map_err(|e| MLError::ModelError(format!("IQL adv_variance_reduce: {e}")))?; + } + + let mut stats = [0.0_f32; 2]; + self.stream.memcpy_dtoh(&self.adv_stats_buf, &mut stats) + .map_err(|e| MLError::ModelError(format!("IQL adv_stats DtoH: {e}")))?; + let sigma = stats[1].max(0.0).sqrt(); + const EMA_BETA: f32 = 0.99; + if self.adv_sigma_ema < 1e-8 { + self.adv_sigma_ema = sigma; + } else { + self.adv_sigma_ema = EMA_BETA * self.adv_sigma_ema + (1.0 - EMA_BETA) * sigma; + } + Ok(()) + } + + /// Compute per-sample C51 atom support centered on V(s). + pub fn compute_per_sample_support( + &mut self, + q_out_buf: &CudaSlice, + ) -> Result<(), MLError> { + let b = self.config.batch_size; + let batch_i32 = b as i32; + let total_actions_i32 = self.config.total_actions as i32; + let num_atoms_i32 = self.config.num_atoms as i32; + let gamma = self.config.gamma; + let blocks = (b + 255) / 256; + + unsafe { + self.stream + .launch_builder(&self.per_sample_support_kernel) + .arg(&self.v_out_buf) + .arg(q_out_buf) + .arg(&mut self.per_sample_support_buf) + .arg(&gamma) + .arg(&batch_i32) + .arg(&total_actions_i32) + .arg(&num_atoms_i32) + .launch(LaunchConfig { + grid_dim: (blocks as u32, 1, 1), + block_dim: (256, 1, 1), + shared_mem_bytes: 0, + }) + .map_err(|e| MLError::ModelError(format!("IQL per_sample_support: {e}")))?; + } + Ok(()) + } + + /// Compute per-branch gradient scales from Q-value marginalization. + pub fn compute_branch_scales( + &mut self, + q_out_buf: &CudaSlice, + actions_buf: &CudaSlice, + ) -> Result<(), MLError> { + let b = self.config.batch_size; + let batch_i32 = b as i32; + let ta = self.config.total_actions as i32; + let bs = self.config.branch_sizes; + let b0 = bs[0] as i32; + let b1 = bs[1] as i32; + let b2 = bs[2] as i32; + let b3 = bs[3] as i32; + let blocks = (b + 255) / 256; + + unsafe { + self.stream + .launch_builder(&self.branch_advantage_kernel) + .arg(q_out_buf) + .arg(&self.v_out_buf) + .arg(actions_buf) + .arg(&mut self.branch_scales_buf) + .arg(&batch_i32) + .arg(&ta) + .arg(&b0) + .arg(&b1) + .arg(&b2) + .arg(&b3) + .launch(LaunchConfig { + grid_dim: (blocks as u32, 1, 1), + block_dim: (256, 1, 1), + shared_mem_bytes: 0, + }) + .map_err(|e| MLError::ModelError(format!("IQL branch_advantage: {e}")))?; + } + Ok(()) + } + + /// Compute expectile gap between two V(s) estimates (high-tau minus low-tau). + pub fn compute_expectile_gap( + &mut self, + v_low_ptr: u64, + ) -> Result<(), MLError> { + let b = self.config.batch_size; + let batch_i32 = b as i32; + let blocks = (b + 255) / 256; + + unsafe { + self.stream + .launch_builder(&self.expectile_gap_kernel) + .arg(&self.v_out_buf) + .arg(&v_low_ptr) + .arg(&mut self.expectile_gap_buf) + .arg(&batch_i32) + .launch(LaunchConfig { + grid_dim: (blocks as u32, 1, 1), + block_dim: (256, 1, 1), + shared_mem_bytes: 0, + }) + .map_err(|e| MLError::ModelError(format!("IQL expectile_gap: {e}")))?; + } + + unsafe { + self.stream + .launch_builder(&self.gap_mean_kernel) + .arg(&self.expectile_gap_buf) + .arg(&mut self.gap_mean_buf) + .arg(&batch_i32) + .launch(LaunchConfig { + grid_dim: (1, 1, 1), + block_dim: (1, 1, 1), + shared_mem_bytes: 0, + }) + .map_err(|e| MLError::ModelError(format!("IQL gap_mean_reduce: {e}")))?; + } + Ok(()) + } + + /// Compute per-sample epsilon from expectile gap. + pub fn compute_per_sample_epsilon( + &mut self, + base_epsilon: f32, + ) -> Result<(), MLError> { + let b = self.config.batch_size; + let batch_i32 = b as i32; + let blocks = (b + 255) / 256; + + unsafe { + self.stream + .launch_builder(&self.per_sample_epsilon_kernel) + .arg(&self.expectile_gap_buf) + .arg(&self.gap_mean_buf) + .arg(&mut self.per_sample_epsilon_buf) + .arg(&base_epsilon) + .arg(&batch_i32) + .launch(LaunchConfig { + grid_dim: (blocks as u32, 1, 1), + block_dim: (256, 1, 1), + shared_mem_bytes: 0, + }) + .map_err(|e| MLError::ModelError(format!("IQL per_sample_epsilon: {e}")))?; + } + Ok(()) + } + + /// Raw pointer to per-sample support buffer [B, 3]. + pub fn per_sample_support_ptr(&self) -> u64 { + self.per_sample_support_buf.raw_ptr() + } + + /// Raw pointer to branch scales buffer [B, 4]. + pub fn branch_scales_ptr(&self) -> u64 { + self.branch_scales_buf.raw_ptr() + } + + /// Raw pointer to per-sample epsilon buffer [B]. + pub fn per_sample_epsilon_ptr(&self) -> u64 { + self.per_sample_epsilon_buf.raw_ptr() + } + + /// Reference to per-sample epsilon buffer. + pub fn per_sample_epsilon_buf(&self) -> &CudaSlice { + &self.per_sample_epsilon_buf + } +} + +// --------------------------------------------------------------------------- +// Compiled kernel set +// --------------------------------------------------------------------------- + +struct IqlKernels { + forward_loss: CudaFunction, + backward_per_sample: CudaFunction, + weight_grad_reduce: CudaFunction, + loss_reduce: CudaFunction, + grad_norm_phase1: CudaFunction, + grad_norm_phase2: CudaFunction, + adam: CudaFunction, + forward: CudaFunction, + gather_q_taken: CudaFunction, + advantage_weight: CudaFunction, + modulate_td: CudaFunction, + adv_variance: CudaFunction, + per_sample_support: CudaFunction, + branch_advantage: CudaFunction, + expectile_gap: CudaFunction, + gap_mean: CudaFunction, + per_sample_epsilon: CudaFunction, } // --------------------------------------------------------------------------- @@ -395,17 +860,17 @@ impl GpuIqlTrainer { /// Precompiled IQL value kernel cubin, embedded at compile time by build.rs. static IQL_VALUE_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/iql_value_kernel.cubin")); -/// Load all 5 IQL CUDA kernels from precompiled cubin. +/// Load all 17 IQL CUDA kernels from precompiled cubin. fn compile_iql_kernels( stream: &Arc, config: &GpuIqlConfig, -) -> Result<(CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction), MLError> { +) -> Result { info!( state_dim = config.state_dim, hidden_dim = config.value_hidden_dim, expectile_tau = config.expectile_tau, total_params = config.total_params(), - "GpuIqlTrainer: loading precompiled IQL kernels" + "GpuIqlTrainer: loading 17 deterministic IQL kernels (zero atomicAdd)" ); let context = stream.context(); @@ -413,24 +878,33 @@ fn compile_iql_kernels( MLError::ModelError(format!("iql_value module load: {e}")) })?; - let forward_loss = module - .load_function("iql_forward_loss_kernel") - .map_err(|e| MLError::ModelError(format!("iql_forward_loss_kernel load: {e}")))?; - let backward = module - .load_function("iql_backward_kernel") - .map_err(|e| MLError::ModelError(format!("iql_backward_kernel load: {e}")))?; - let grad_norm = module - .load_function("iql_grad_norm_kernel") - .map_err(|e| MLError::ModelError(format!("iql_grad_norm_kernel load: {e}")))?; - let adam = module - .load_function("iql_adam_kernel") - .map_err(|e| MLError::ModelError(format!("iql_adam_kernel load: {e}")))?; - let forward = module - .load_function("iql_forward_kernel") - .map_err(|e| MLError::ModelError(format!("iql_forward_kernel load: {e}")))?; + let load = |name: &str| -> Result { + module.load_function(name) + .map_err(|e| MLError::ModelError(format!("{name} load: {e}"))) + }; - info!("GpuIqlTrainer: 5 kernels compiled and loaded"); - Ok((forward_loss, backward, grad_norm, adam, forward)) + let kernels = IqlKernels { + forward_loss: load("iql_forward_loss_kernel")?, + backward_per_sample: load("iql_backward_per_sample")?, + weight_grad_reduce: load("iql_weight_grad_reduce")?, + loss_reduce: load("iql_loss_reduce")?, + grad_norm_phase1: load("iql_grad_norm_phase1")?, + grad_norm_phase2: load("iql_grad_norm_phase2")?, + adam: load("iql_adam_kernel")?, + forward: load("iql_forward_kernel")?, + gather_q_taken: load("iql_gather_q_taken")?, + advantage_weight: load("iql_compute_advantage_weights")?, + modulate_td: load("iql_modulate_td_errors")?, + adv_variance: load("iql_adv_variance_reduce")?, + per_sample_support: load("iql_compute_per_sample_support")?, + branch_advantage: load("iql_per_branch_advantage")?, + expectile_gap: load("iql_expectile_gap")?, + gap_mean: load("iql_gap_mean_reduce")?, + per_sample_epsilon: load("iql_compute_per_sample_epsilon")?, + }; + + info!("GpuIqlTrainer: 17 kernels loaded"); + Ok(kernels) } // --------------------------------------------------------------------------- @@ -503,4 +977,3 @@ fn alloc_f32( MLError::ModelError(format!("GpuIql alloc {name} ({n} f32): {e}")) }) } -