diff --git a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs index 591ffd392..d16b578aa 100644 --- a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs +++ b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs @@ -696,10 +696,15 @@ pub struct GpuDqnTrainer { // ── Forward output buffers ────────────────────────────────────── per_sample_loss_buf: CudaSlice, // [B] td_errors_buf: CudaSlice, // [B] - pub(crate) total_loss_buf: CudaSlice, // [1] float accumulator (deterministic reduce) — C51 loss - pub(crate) mse_loss_buf: CudaSlice, // [1] MSE loss accumulator (separate from C51) - /// [1] Mean squared Q-divergence between online and target networks (atomicAdd). - pub(crate) q_divergence_buf: CudaSlice, + /// C51 loss accumulator [1] — pinned device-mapped. GPU writes via dev_ptr, CPU reads via host ptr. + total_loss_pinned: *mut f32, + total_loss_dev_ptr: u64, + /// MSE loss accumulator [1] — pinned device-mapped. GPU writes via dev_ptr, CPU reads via host ptr. + mse_loss_pinned: *mut f32, + mse_loss_dev_ptr: u64, + /// [1] Mean squared Q-divergence — pinned device-mapped. GPU writes via dev_ptr, CPU reads via host ptr. + q_divergence_pinned: *mut f32, + q_divergence_dev_ptr: u64, // ── Forward-only Q-value output ───────────────────────────────── q_out_buf: CudaSlice, // [B, TOTAL_ACTIONS(11)] @@ -714,7 +719,9 @@ pub struct GpuDqnTrainer { target_params_buf: CudaSlice, // [TOTAL_PARAMS] f32 master target parameters (EMA operates here) m_buf: CudaSlice, // [TOTAL_PARAMS] Adam first moment (f32 for precision) v_buf: CudaSlice, // [TOTAL_PARAMS] Adam second moment (f32 for precision) - pub(crate) grad_norm_buf: CudaSlice, // [1] f32 L2 norm (written by finalize) + /// Gradient L2 norm [1] — pinned device-mapped. Finalize kernel writes via dev_ptr, CPU reads via host ptr. + grad_norm_pinned: *mut f32, + grad_norm_dev_ptr: u64, grad_norm_partials: CudaSlice, // [grad_norm_blocks] per-block partial sums grad_norm_blocks: usize, // number of blocks for grad_norm kernel cql_grad_scratch: CudaSlice, // [TOTAL_PARAMS] f32 CQL gradient isolation buffer @@ -732,7 +739,8 @@ pub struct GpuDqnTrainer { /// Per-branch gradient scales [B, 4] — pointer set by IQL trainer. branch_scales_ptr: u64, /// Fixed wide v_range [2] for compute_expected_q argmax (not C51 loss). - eval_v_range_buf: CudaSlice, + /// Pinned device-mapped: CPU writes v_min/v_max, GPU reads via dev_ptr. No HtoD copy. + eval_v_range_pinned: *mut f32, eval_v_range_ptr: u64, /// EMA-smoothed Q-stats for stable eval v_range updates. pub(crate) eval_q_mean_ema: f32, @@ -1020,6 +1028,9 @@ pub struct GpuDqnTrainer { q_stats_buf: CudaSlice, /// GPU buffer for atom utilization accumulation [2 floats: sum_entropy, sum_utilized] atom_stats_buf: CudaSlice, + /// Pinned readback buffer for q_stats [7 floats] + q_out sample 0 [12 floats] = 19 floats. + /// DtoH uses pinned DMA-capable destination for faster async transfer. + q_readback_pinned: *mut f32, // ── CQL (Conservative Q-Learning) penalty kernel ───────────────── /// Computes CQL logit gradients: dCQL/d_value_logits and dCQL/d_adv_logits. @@ -1092,8 +1103,9 @@ pub struct GpuDqnTrainer { qlstm_n: CudaSlice, /// [8] output context vector (written each qlstm_step call) qlstm_context: CudaSlice, - /// [4] per-branch Q-gaps device buffer (written in reduce_current_q_stats) - per_branch_q_gaps_buf: CudaSlice, + /// [4] per-branch Q-gaps — pinned device-mapped. CPU writes, GPU reads via dev_ptr. No HtoD copy. + per_branch_q_gaps_pinned: *mut f32, + per_branch_q_gaps_dev_ptr: u64, /// qlstm_step kernel handle qlstm_step_kernel: CudaFunction, } @@ -1134,8 +1146,11 @@ impl GpuDqnTrainer { // Reset Q-divergence EMA — fold 2's divergence baseline differs from fold 1. self.q_div_ema = 0.0; - self.stream.memset_zeros(&mut self.q_divergence_buf) - .map_err(|e| MLError::ModelError(format!("reset q_divergence: {e}")))?; + unsafe { + cudarc::driver::sys::cuMemsetD8Async( + self.q_divergence_dev_ptr, 0, std::mem::size_of::(), self.stream.cu_stream(), + ); + } tracing::info!("Adam optimizer + PopArt + grad_clip state reset for new fold"); Ok(()) @@ -1214,6 +1229,27 @@ impl Drop for GpuDqnTrainer { if !self.tau_pinned.is_null() { let _ = unsafe { cudarc::driver::result::free_host(self.tau_pinned.cast()) }; } + if !self.total_loss_pinned.is_null() { + let _ = unsafe { cudarc::driver::result::free_host(self.total_loss_pinned.cast()) }; + } + if !self.mse_loss_pinned.is_null() { + let _ = unsafe { cudarc::driver::result::free_host(self.mse_loss_pinned.cast()) }; + } + if !self.grad_norm_pinned.is_null() { + let _ = unsafe { cudarc::driver::result::free_host(self.grad_norm_pinned.cast()) }; + } + if !self.eval_v_range_pinned.is_null() { + let _ = unsafe { cudarc::driver::result::free_host(self.eval_v_range_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()) }; + } + if !self.q_readback_pinned.is_null() { + let _ = unsafe { cudarc::driver::result::free_host(self.q_readback_pinned.cast()) }; + } + if !self.q_divergence_pinned.is_null() { + let _ = unsafe { cudarc::driver::result::free_host(self.q_divergence_pinned.cast()) }; + } } } @@ -1225,13 +1261,20 @@ impl GpuDqnTrainer { pub fn eval_v_range_ptr(&self) -> u64 { self.eval_v_range_ptr } + /// Device pointer to total_loss (pinned device-mapped). + pub fn total_loss_dev_ptr(&self) -> u64 { self.total_loss_dev_ptr } + /// Device pointer to mse_loss (pinned device-mapped). + pub fn mse_loss_dev_ptr(&self) -> u64 { self.mse_loss_dev_ptr } + /// Device pointer to grad_norm (pinned device-mapped). + pub fn grad_norm_dev_ptr(&self) -> u64 { self.grad_norm_dev_ptr } + /// Update per-branch liquid tau modulation — GPU kernel (RK4/Euler adaptive ODE). /// - /// Reads per_branch_q_gaps_buf and qlstm_context from device. + /// Reads per_branch_q_gaps (pinned device-mapped) and qlstm_context from device. /// Updates per_branch_q_gap_ema_buf and liquid_mod_buf on device. /// Zero CPU compute. pub fn update_liquid_tau(&mut self, delta_z_approx: f32) -> Result<(), MLError> { - let gaps_ptr = self.per_branch_q_gaps_buf.raw_ptr(); + let gaps_ptr = self.per_branch_q_gaps_dev_ptr; let ema_ptr = self.per_branch_q_gap_ema_buf.raw_ptr(); let mod_ptr = self.liquid_mod_buf.raw_ptr(); let ctx_ptr = self.qlstm_context.raw_ptr(); @@ -1276,8 +1319,8 @@ impl GpuDqnTrainer { self.eval_q_std_ema = (1.0 - alpha_std) * self.eval_q_std_ema + alpha_std * q_std; } - // Keep last_per_branch_q_gaps in sync (per_branch_q_gaps_buf already written in reduce_current_q_stats). - let _ = per_branch_q_gaps; // already on device via per_branch_q_gaps_buf + // Keep last_per_branch_q_gaps in sync (per_branch_q_gaps pinned already written in reduce_current_q_stats). + let _ = per_branch_q_gaps; // already on device via per_branch_q_gaps_dev_ptr // GPU kernel: liquid tau RK4/Euler adaptive ODE (zero CPU compute). let delta_z = self.eval_q_std_ema.max(0.01); @@ -1293,13 +1336,10 @@ impl GpuDqnTrainer { let half = gap_width; let v_min = self.eval_q_mean_ema - half; let v_max = self.eval_q_mean_ema + half; + // Pinned device-mapped: CPU writes directly, GPU reads via dev_ptr. No HtoD copy. unsafe { - cudarc::driver::sys::cuMemcpyHtoDAsync_v2( - self.eval_v_range_ptr, - [v_min, v_max].as_ptr().cast(), - 2 * std::mem::size_of::(), - self.stream.cu_stream(), - ); + *self.eval_v_range_pinned = v_min; + *self.eval_v_range_pinned.add(1) = v_max; } } pub fn set_per_sample_support_ptr(&mut self, ptr: u64) { @@ -1596,9 +1636,9 @@ impl GpuDqnTrainer { /// Read Q-divergence from pinned readback buffer (offset [12]). /// Returns the mean squared Q-divergence between online and target networks. - /// Valid after `replay_adam_and_readback()` has been called (async copy landed). + /// Pinned device-mapped: reads directly from host pointer (no DtoH copy). pub(crate) fn q_divergence_readback(&self) -> f32 { - unsafe { *self.readback_pinned.add(12) } + unsafe { *self.q_divergence_pinned } } /// Current IQN readiness scalar [0, 1]. @@ -1612,7 +1652,7 @@ impl GpuDqnTrainer { /// When divergence is high (networks drifted apart), increase tau to close the gap faster. /// Uses EMA-smoothed divergence as the baseline. pub fn compute_adaptive_tau(&mut self, base_tau: f32) -> f32 { - let raw_div = unsafe { *self.readback_pinned.add(12) }; + let raw_div = unsafe { *self.q_divergence_pinned }; // EMA smooth the divergence signal const DIV_BETA: f32 = 0.95; if self.q_div_ema <= 0.0 { @@ -2594,16 +2634,8 @@ impl GpuDqnTrainer { unsafe { cudarc::driver::sys::cuStreamSynchronize(self.stream.cu_stream()); } - // grad_norm_buf[0] now contains the L2 norm directly (sqrt applied by finalize) - let mut norm_val = [0.0_f32; 1]; - unsafe { - cudarc::driver::sys::cuMemcpyDtoH_v2( - norm_val.as_mut_ptr().cast(), - self.grad_norm_buf.raw_ptr(), - std::mem::size_of::(), - ); - } - Ok(norm_val[0]) + // grad_norm is pinned device-mapped — read directly from host pointer after sync. + Ok(unsafe { *self.grad_norm_pinned }) } /// Compute and return the current grad_buf L2 norm using the STANDALONE @@ -2622,16 +2654,8 @@ impl GpuDqnTrainer { unsafe { cudarc::driver::sys::cuStreamSynchronize(self.stream.cu_stream()); } - let mut norm_f32 = [0.0_f32; 1]; - unsafe { - cudarc::driver::sys::cuMemcpyDtoH_v2( - norm_f32.as_mut_ptr().cast(), - self.grad_norm_buf.raw_ptr(), - std::mem::size_of::(), - ); - } - // grad_norm_buf holds L2 norm directly (finalize computes sqrt) - Ok(norm_f32[0]) + // grad_norm is pinned device-mapped — read directly from host pointer after sync. + Ok(unsafe { *self.grad_norm_pinned }) } /// Apply GPU multi-head feature attention to `save_h_s2` (post-graph). @@ -2833,12 +2857,44 @@ impl GpuDqnTrainer { // ── Allocate forward output buffers ───────────────────────── let per_sample_loss_buf = alloc_f32(&stream, b, "per_sample_loss")?; let td_errors_buf = alloc_f32(&stream, b, "td_errors")?; - let total_loss_buf = stream.alloc_zeros::(1) - .map_err(|e| MLError::ModelError(format!("alloc total_loss_f32: {e}")))?; - let mse_loss_buf = stream.alloc_zeros::(1) - .map_err(|e| MLError::ModelError(format!("alloc mse_loss_f32: {e}")))?; - let q_divergence_buf = stream.alloc_zeros::(1) - .map_err(|e| MLError::ModelError(format!("alloc q_divergence_f32: {e}")))?; + // total_loss + mse_loss — pinned device-mapped (zero-copy readback). + let total_loss_pinned: *mut f32 = unsafe { + let flags = cudarc::driver::sys::CU_MEMHOSTALLOC_DEVICEMAP; + cudarc::driver::result::malloc_host(std::mem::size_of::(), flags) + .map_err(|e| MLError::ModelError(format!("pinned total_loss alloc: {e}")))? + as *mut f32 + }; + unsafe { *total_loss_pinned = 0.0; } + let total_loss_dev_ptr = unsafe { + let mut dp = 0u64; + cudarc::driver::sys::cuMemHostGetDevicePointer_v2(&mut dp as *mut u64, total_loss_pinned.cast(), 0); + dp + }; + let mse_loss_pinned: *mut f32 = unsafe { + let flags = cudarc::driver::sys::CU_MEMHOSTALLOC_DEVICEMAP; + cudarc::driver::result::malloc_host(std::mem::size_of::(), flags) + .map_err(|e| MLError::ModelError(format!("pinned mse_loss alloc: {e}")))? + as *mut f32 + }; + unsafe { *mse_loss_pinned = 0.0; } + let mse_loss_dev_ptr = unsafe { + let mut dp = 0u64; + cudarc::driver::sys::cuMemHostGetDevicePointer_v2(&mut dp as *mut u64, mse_loss_pinned.cast(), 0); + dp + }; + // q_divergence — pinned device-mapped (zero-copy readback). + let q_divergence_pinned: *mut f32 = unsafe { + let flags = cudarc::driver::sys::CU_MEMHOSTALLOC_DEVICEMAP; + cudarc::driver::result::malloc_host(std::mem::size_of::(), flags) + .map_err(|e| MLError::ModelError(format!("pinned q_divergence alloc: {e}")))? + as *mut f32 + }; + unsafe { *q_divergence_pinned = 0.0; } + let q_divergence_dev_ptr = unsafe { + let mut dp = 0u64; + cudarc::driver::sys::cuMemHostGetDevicePointer_v2(&mut dp as *mut u64, q_divergence_pinned.cast(), 0); + dp + }; let total_actions = config.branch_0_size + config.branch_1_size + config.branch_2_size + config.branch_3_size; let q_out_buf = alloc_f32(&stream, b * total_actions, "q_out")?; let eval_td_snapshot = alloc_f32(&stream, b, "eval_td_snapshot")?; @@ -2864,7 +2920,19 @@ impl GpuDqnTrainer { .map_err(|e| MLError::ModelError(format!("alloc adam_m f32: {e}")))?; let v_buf = stream.alloc_zeros::(total_params + cutlass_tile_pad) .map_err(|e| MLError::ModelError(format!("alloc adam_v f32: {e}")))?; - let grad_norm_buf = alloc_f32(&stream, 1, "grad_norm")?; + // grad_norm — pinned device-mapped (zero-copy readback). + let grad_norm_pinned: *mut f32 = unsafe { + let flags = cudarc::driver::sys::CU_MEMHOSTALLOC_DEVICEMAP; + cudarc::driver::result::malloc_host(std::mem::size_of::(), flags) + .map_err(|e| MLError::ModelError(format!("pinned grad_norm alloc: {e}")))? + as *mut f32 + }; + unsafe { *grad_norm_pinned = 0.0; } + let grad_norm_dev_ptr = unsafe { + let mut dp = 0u64; + cudarc::driver::sys::cuMemHostGetDevicePointer_v2(&mut dp as *mut u64, grad_norm_pinned.cast(), 0); + dp + }; // Size for worst case: d_logits clipping uses val_blocks + adv_blocks // which can exceed total_params blocks at large batch sizes. let tba = config.branch_0_size + config.branch_1_size @@ -2878,14 +2946,22 @@ impl GpuDqnTrainer { let cql_grad_scratch = stream.alloc_zeros::(total_params + cutlass_tile_pad) .map_err(|e| MLError::ModelError(format!("alloc cql_grad_scratch f32: {e}")))?; - // Fixed wide v_range for compute_expected_q (argmax, not C51 loss). - let mut eval_v_range_buf = stream.alloc_zeros::(2) - .map_err(|e| MLError::ModelError(format!("alloc eval_v_range: {e}")))?; - // Eval v_range: initial from config, updated per-epoch from Q-stats. - let eval_range = [config.v_min, config.v_max]; - stream.memcpy_htod(&eval_range, &mut eval_v_range_buf) - .map_err(|e| MLError::ModelError(format!("eval_v_range HtoD: {e}")))?; - let eval_v_range_ptr = eval_v_range_buf.raw_ptr(); + // 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; + cudarc::driver::result::malloc_host(2 * std::mem::size_of::(), flags) + .map_err(|e| MLError::ModelError(format!("pinned eval_v_range alloc: {e}")))? + as *mut f32 + }; + unsafe { + *eval_v_range_pinned = config.v_min; + *eval_v_range_pinned.add(1) = config.v_max; + } + let eval_v_range_ptr = unsafe { + let mut dp = 0u64; + cudarc::driver::sys::cuMemHostGetDevicePointer_v2(&mut dp as *mut u64, eval_v_range_pinned.cast(), 0); + dp + }; // Adam step counter — pinned device-mapped (no HtoD copies). let t_pinned: *mut i32 = unsafe { let flags = cudarc::driver::sys::CU_MEMHOSTALLOC_DEVICEMAP; @@ -3021,6 +3097,14 @@ impl GpuDqnTrainer { .map_err(|e| MLError::ModelError(format!("alloc q_stats_f32: {e}")))?; let atom_stats_buf = stream.alloc_zeros::(2) .map_err(|e| MLError::ModelError(format!("alloc atom_stats: {e}")))?; + // q_readback — pinned host buffer for DMA-capable DtoH of q_stats[7] + q_out[12] + let q_readback_pinned: *mut f32 = unsafe { + let flags = cudarc::driver::sys::CU_MEMHOSTALLOC_DEVICEMAP; + cudarc::driver::result::malloc_host(19 * std::mem::size_of::(), flags) + .map_err(|e| MLError::ModelError(format!("pinned q_readback alloc: {e}")))? + as *mut f32 + }; + unsafe { std::ptr::write_bytes(q_readback_pinned, 0, 19); } info!("GpuDqnTrainer: expected_q + q_stats kernels compiled"); // ── Load mag_concat and strided_accumulate from experience_kernels cubin ── @@ -3470,13 +3554,13 @@ impl GpuDqnTrainer { params_ptr: params_buf.raw_ptr(), target_ptr: target_params_buf.raw_ptr(), grad_buf: grad_buf.raw_ptr(), - grad_norm_buf: grad_norm_buf.raw_ptr(), + grad_norm_buf: grad_norm_dev_ptr, m_buf: m_buf.raw_ptr(), v_buf: v_buf.raw_ptr(), t_buf: t_dev_ptr, adaptive_clip_buf: adaptive_clip_dev_ptr, - total_loss_buf: total_loss_buf.raw_ptr(), - mse_loss_buf: mse_loss_buf.raw_ptr(), + total_loss_buf: total_loss_dev_ptr, + mse_loss_buf: mse_loss_dev_ptr, cql_grad_scratch: cql_grad_scratch.raw_ptr(), states_buf: states_buf.raw_ptr(), next_states_buf: next_states_buf.raw_ptr(), @@ -3694,8 +3778,21 @@ impl GpuDqnTrainer { .map_err(|e| MLError::ModelError(format!("alloc qlstm_n: {e}")))?; let qlstm_context = stream.alloc_zeros::(QLSTM_HEAD_DIM) // [8] .map_err(|e| MLError::ModelError(format!("alloc qlstm_context: {e}")))?; - let per_branch_q_gaps_buf = stream.alloc_zeros::(4) // [4] - .map_err(|e| MLError::ModelError(format!("alloc per_branch_q_gaps_buf: {e}")))?; + // per_branch_q_gaps — pinned device-mapped (zero-copy write from CPU). + let per_branch_q_gaps_pinned: *mut f32 = unsafe { + let flags = cudarc::driver::sys::CU_MEMHOSTALLOC_DEVICEMAP; + cudarc::driver::result::malloc_host(4 * std::mem::size_of::(), flags) + .map_err(|e| MLError::ModelError(format!("pinned per_branch_q_gaps alloc: {e}")))? + as *mut f32 + }; + unsafe { + std::ptr::write_bytes(per_branch_q_gaps_pinned, 0, 4); + } + let per_branch_q_gaps_dev_ptr = unsafe { + let mut dp = 0u64; + cudarc::driver::sys::cuMemHostGetDevicePointer_v2(&mut dp as *mut u64, per_branch_q_gaps_pinned.cast(), 0); + dp + }; let cpbi_module_qlstm = stream.context().load_cubin(EXPECTED_Q_CUBIN.to_vec()) .map_err(|e| MLError::ModelError(format!("cpbi cubin (qlstm): {e}")))?; let qlstm_step_kernel = cpbi_module_qlstm.load_function("qlstm_step") @@ -3813,9 +3910,12 @@ impl GpuDqnTrainer { save_projected, per_sample_loss_buf, td_errors_buf, - total_loss_buf, - mse_loss_buf, - q_divergence_buf, + total_loss_pinned, + total_loss_dev_ptr, + mse_loss_pinned, + mse_loss_dev_ptr, + q_divergence_pinned, + q_divergence_dev_ptr, q_out_buf, eval_td_snapshot, eval_loss_snapshot, @@ -3824,7 +3924,8 @@ impl GpuDqnTrainer { target_params_buf, m_buf, v_buf, - grad_norm_buf, + grad_norm_pinned, + grad_norm_dev_ptr, grad_norm_partials, grad_norm_blocks, cql_grad_scratch, @@ -3834,7 +3935,7 @@ impl GpuDqnTrainer { tau_dev_ptr, per_sample_support_ptr: 0, branch_scales_ptr: 0, - eval_v_range_buf, + eval_v_range_pinned, eval_v_range_ptr, eval_q_mean_ema: 0.0, eval_q_std_ema: 0.0, @@ -3918,6 +4019,7 @@ impl GpuDqnTrainer { q_stats_kernel, q_stats_buf, atom_stats_buf, + q_readback_pinned, cql_logit_grad_kernel, cql_d_value_logits, cql_d_adv_logits, @@ -3996,7 +4098,8 @@ impl GpuDqnTrainer { qlstm_c, qlstm_n, qlstm_context, - per_branch_q_gaps_buf, + per_branch_q_gaps_pinned, + per_branch_q_gaps_dev_ptr, qlstm_step_kernel, }) } @@ -4165,68 +4268,26 @@ impl GpuDqnTrainer { } pub fn replay_adam_and_readback(&mut self) -> Result { - // 1. Collect previous step's scalars (if pending) - // Non-blocking: read from pinned buffer without synchronizing. - // These values are monitoring-only (the training guard reads GPU buffers - // directly via loss_gpu_buf()/grad_norm_gpu_buf()). If the async DtoH - // from the previous step hasn't landed yet, we read stale values — acceptable - // for logging scalars. In practice the 12-byte copy completes in <1 µs - // while the GPU does thousands of kernel launches between steps. - let prev_scalars = if self.readback_pending { - self.readback_pending = false; - // Read from pinned buffer: [0]=loss, [1]=mse_loss, [2]=grad_norm - let (loss, mse_loss, grad_norm) = unsafe { - (*self.readback_pinned, *self.readback_pinned.add(1), *self.readback_pinned.add(2)) - }; - let blended = Self::blend_loss(self.c51_alpha, mse_loss, loss); - FusedTrainScalars { - total_loss: blended, - grad_norm, - } - } else { - // First step — no previous data - FusedTrainScalars { total_loss: 0.0, grad_norm: 0.0 } + // 1. Read previous step's scalars from pinned device-mapped memory. + // These buffers were written by the GPU during the PREVIOUS step's forward + // pass. The current step's forward has not yet overwritten them (it runs + // below). No sync needed — thousands of kernel launches have elapsed since + // the previous forward, guaranteeing the GPU writes have landed in host DRAM. + let (loss, mse_loss, grad_norm) = unsafe { + (*self.total_loss_pinned, *self.mse_loss_pinned, *self.grad_norm_pinned) + }; + let blended = Self::blend_loss(self.c51_alpha, mse_loss, loss); + let prev_scalars = FusedTrainScalars { + total_loss: blended, + grad_norm, }; // 2. Launch current step: grad_norm (two-phase, no atomicAdd) + adam self.compute_grad_norm_outside_graph()?; self.replay_adam()?; - // 3. Async DtoH into pinned host buffer (truly async — no CPU blocking) - unsafe { - // [0] = loss - cudarc::driver::sys::cuMemcpyDtoHAsync_v2( - self.readback_pinned.cast(), - self.total_loss_buf.raw_ptr(), - std::mem::size_of::(), - self.stream.cu_stream(), - ); - // [1] = mse_loss - cudarc::driver::sys::cuMemcpyDtoHAsync_v2( - self.readback_pinned.add(1).cast(), - self.mse_loss_buf.raw_ptr(), - std::mem::size_of::(), - self.stream.cu_stream(), - ); - // [2] = grad_norm (L2 norm, finalize computes sqrt) - cudarc::driver::sys::cuMemcpyDtoHAsync_v2( - self.readback_pinned.add(2).cast(), - self.grad_norm_buf.raw_ptr(), - std::mem::size_of::(), - self.stream.cu_stream(), - ); - // [12] = q_divergence (mean squared Q-divergence between online and target) - cudarc::driver::sys::cuMemcpyDtoHAsync_v2( - self.readback_pinned.add(12).cast(), - self.q_divergence_buf.raw_ptr(), - std::mem::size_of::(), - self.stream.cu_stream(), - ); - } - - // No event recording needed — flush_readback reads pinned buffer - // directly without synchronizing (async copy completes in <1 µs). - self.readback_pending = true; + // All per-step scalars are now pinned device-mapped — zero cuMemcpy. + // q_divergence is also pinned: GPU writes via dev_ptr, CPU reads via host ptr. // Return previous step's values self.scalars_readback_host = [prev_scalars.total_loss, prev_scalars.grad_norm]; @@ -4242,14 +4303,12 @@ impl GpuDqnTrainer { /// to have landed. The training guard reads loss/grad_norm from GPU /// buffers directly — these CPU scalars are monitoring-only. pub fn flush_readback(&mut self) -> Result { - if !self.readback_pending { - return Ok(FusedTrainScalars { total_loss: 0.0, grad_norm: 0.0 }); - } - // Non-blocking: no event.synchronize(). The async DtoH enqueued by - // replay_adam_and_readback() has completed long before epoch end. self.readback_pending = false; + // Read directly from pinned device-mapped host memory — no DtoH needed. + // The last forward pass wrote these values to host DRAM via PCIe; by epoch + // end, thousands of kernel launches have elapsed, guaranteeing visibility. let (loss, mse_loss, grad_norm) = unsafe { - (*self.readback_pinned, *self.readback_pinned.add(1), *self.readback_pinned.add(2)) + (*self.total_loss_pinned, *self.mse_loss_pinned, *self.grad_norm_pinned) }; let blended = Self::blend_loss(self.c51_alpha, mse_loss, loss); Ok(FusedTrainScalars { @@ -4260,10 +4319,10 @@ impl GpuDqnTrainer { /// Phase 2: reduce block partials → L2 norm. /// Single block, 256 threads — reduces grad_norm_partials[0..grad_norm_blocks]. - /// Writes L2 norm (with sqrt) into grad_norm_buf. + /// Writes L2 norm (with sqrt) into grad_norm pinned device-mapped buffer. fn launch_grad_norm_finalize(&self) -> Result<(), MLError> { let partials_ptr = self.grad_norm_partials.raw_ptr(); - let buf_ptr = self.grad_norm_buf.raw_ptr(); + let buf_ptr = self.grad_norm_dev_ptr; let nb = self.grad_norm_blocks as i32; unsafe { self.stream @@ -4289,27 +4348,14 @@ impl GpuDqnTrainer { unsafe { cudarc::driver::sys::cuStreamSynchronize(self.stream.cu_stream()); } - let mut c51_loss_f32 = [0.0_f32; 1]; - let mut mse_loss_f32 = [0.0_f32; 1]; - let mut norm_val = [0.0_f32; 1]; - unsafe { - cudarc::driver::sys::cuMemcpyDtoH_v2( - c51_loss_f32.as_mut_ptr().cast(), - self.ptrs.total_loss_buf, std::mem::size_of::(), - ); - cudarc::driver::sys::cuMemcpyDtoH_v2( - mse_loss_f32.as_mut_ptr().cast(), - self.ptrs.mse_loss_buf, std::mem::size_of::(), - ); - cudarc::driver::sys::cuMemcpyDtoH_v2( - norm_val.as_mut_ptr().cast(), - self.ptrs.grad_norm_buf, std::mem::size_of::(), - ); - } - let blended = Self::blend_loss(self.c51_alpha, mse_loss_f32[0], c51_loss_f32[0]); + // After sync, pinned device-mapped host memory is coherent — read directly. + let c51_loss = unsafe { *self.total_loss_pinned }; + let mse_loss = unsafe { *self.mse_loss_pinned }; + let norm_val = unsafe { *self.grad_norm_pinned }; + let blended = Self::blend_loss(self.c51_alpha, mse_loss, c51_loss); Ok(FusedTrainScalars { total_loss: blended, - grad_norm: norm_val[0], + grad_norm: norm_val, }) } @@ -4344,38 +4390,25 @@ impl GpuDqnTrainer { self.compute_grad_norm_outside_graph()?; self.replay_adam()?; - // ── Consolidate readback: gather 4 GPU buffers → 1 readback buf ── - // Layout: [total_loss(1) | mse_loss(1) | grad_norm(1) | td_errors(B)] - // D2D copies are async on the same stream — no host sync until the - // single memcpy_dtoh at the end. + // ── Readback: td_errors from GPU, scalars from pinned host memory ── + // loss/mse_loss/grad_norm are pinned device-mapped — GPU already wrote + // to host DRAM. Only td_errors[B] needs DtoH transfer. let readback_base = self.readback_buf.raw_ptr(); let f32_bytes = std::mem::size_of::(); - // total_loss (C51) → readback[0] - let loss_src = self.total_loss_buf.raw_ptr(); - dtod_copy(readback_base, loss_src, f32_bytes, &self.stream, 0, "readback_gather")?; - - // mse_loss → readback[1] - let mse_src = self.mse_loss_buf.raw_ptr(); - dtod_copy(readback_base + f32_bytes as u64, mse_src, f32_bytes, &self.stream, 1, "readback_gather")?; - - // grad_norm → readback[2] - let norm_src = self.grad_norm_buf.raw_ptr(); - dtod_copy(readback_base + (2 * f32_bytes) as u64, norm_src, f32_bytes, &self.stream, 2, "readback_gather")?; - - // td_errors → readback[3..3+B] + // td_errors → readback[3..3+B] (only td_errors needs DtoD → DtoH) let td_src = self.td_errors_buf.raw_ptr(); let td_bytes = b * f32_bytes; dtod_copy(readback_base + (3 * f32_bytes) as u64, td_src, td_bytes, &self.stream, 3, "readback_gather")?; - // ── Single DtoH transfer ────────────────────────────────────── + // ── Single DtoH transfer (only td_errors portion matters) ───── super::dtoh_f32(&self.stream, &self.readback_buf, &mut self.readback_host)?; - // ── Unpack on CPU ───────────────────────────────────────────── - let c51_loss = self.readback_host[0]; - let mse_loss = self.readback_host[1]; - let grad_norm_sq = self.readback_host[2]; - let td_errors = self.readback_host[3..3 + b].to_vec(); // cpu-side host-buf slice + // ── Unpack: scalars from pinned host, td_errors from readback ─ + let (c51_loss, mse_loss, grad_norm_sq) = unsafe { + (*self.total_loss_pinned, *self.mse_loss_pinned, *self.grad_norm_pinned) + }; + let td_errors = self.readback_host[3..3 + b].to_vec(); let blended = Self::blend_loss(self.c51_alpha, mse_loss, c51_loss); Ok(FusedTrainResult { @@ -4649,10 +4682,10 @@ impl GpuDqnTrainer { self.stream.memset_zeros(&mut self.d_adv_logits_mse) .map_err(|e| MLError::ModelError(format!("vaccine zero mse2: {e}")))?; self.launch_mse_loss()?; - self.launch_loss_reduce(&self.mse_loss_buf)?; + self.launch_loss_reduce(self.mse_loss_dev_ptr)?; self.launch_mse_grad_to_scratch()?; self.launch_c51_loss()?; - self.launch_loss_reduce(&self.total_loss_buf)?; + self.launch_loss_reduce(self.total_loss_dev_ptr)?; self.launch_c51_grad()?; // Backward writes g_val to scratch (not grad_buf) self.launch_cublas_backward_to(g_val_ptr)?; @@ -4784,29 +4817,15 @@ impl GpuDqnTrainer { self.compute_grad_norm_outside_graph()?; self.replay_adam()?; - // ── Raw sync + readback (bypasses stale events from graph capture) ── + // ── Sync then read pinned device-mapped host memory (no DtoH copy) ── unsafe { cudarc::driver::sys::cuStreamSynchronize(self.stream.cu_stream()); } - let mut c51_loss_f32 = [0.0_f32; 1]; - let mut mse_loss_f32 = [0.0_f32; 1]; - let mut norm_val = [0.0_f32; 1]; - unsafe { - cudarc::driver::sys::cuMemcpyDtoH_v2( - c51_loss_f32.as_mut_ptr().cast(), - self.total_loss_buf.raw_ptr(), std::mem::size_of::(), - ); - cudarc::driver::sys::cuMemcpyDtoH_v2( - mse_loss_f32.as_mut_ptr().cast(), - self.mse_loss_buf.raw_ptr(), std::mem::size_of::(), - ); - cudarc::driver::sys::cuMemcpyDtoH_v2( - norm_val.as_mut_ptr().cast(), - self.grad_norm_buf.raw_ptr(), std::mem::size_of::(), - ); - } - let blended = Self::blend_loss(self.c51_alpha, mse_loss_f32[0], c51_loss_f32[0]); - self.scalars_readback_host = [blended, norm_val[0]]; + let c51_loss = unsafe { *self.total_loss_pinned }; + let mse_loss = unsafe { *self.mse_loss_pinned }; + let norm_val = unsafe { *self.grad_norm_pinned }; + let blended = Self::blend_loss(self.c51_alpha, mse_loss, c51_loss); + self.scalars_readback_host = [blended, norm_val]; Ok(FusedTrainScalars { total_loss: self.scalars_readback_host[0], @@ -4857,7 +4876,7 @@ impl GpuDqnTrainer { // Flag 2: on_b_logits (f32 cuBLAS forward output — branch advantage logits) self.check_nan_f32(self.on_b_logits_buf.raw_ptr(), b * (b0 + b1 + b2 + b3) * na, 2)?; // Flag 3: mse_loss_buf [1] (MSE loss scalar) - self.check_nan_f32(self.mse_loss_buf.raw_ptr(), 1, 3)?; + self.check_nan_f32(self.mse_loss_dev_ptr, 1, 3)?; // Flag 6: grad_buf (cuBLAS backward output) — use ptrs for consistency self.check_nan_f32(self.ptrs.grad_buf, self.total_params, 6)?; // Flag 7: save_current_lp (softmax probs from MSE loss — f32) @@ -5299,26 +5318,36 @@ impl GpuDqnTrainer { // action selection. Stale values cause 2-state oscillation where // the system flips between cached results from different steps. // Cost: ~5µs sync + 28B + 48B DtoH. Runs every 50 training steps. - unsafe { cudarc::driver::sys::cuStreamSynchronize(self.stream.cu_stream()); } - let mut host = [0.0_f32; 7]; + // DtoHAsync to pinned readback buffer (DMA-capable destination), then sync. unsafe { - cudarc::driver::sys::cuMemcpyDtoH_v2( - host.as_mut_ptr().cast(), + cudarc::driver::sys::cuMemcpyDtoHAsync_v2( + self.q_readback_pinned.cast(), self.q_stats_buf.raw_ptr(), 7 * std::mem::size_of::(), + self.stream.cu_stream(), ); - } - - // Per-branch Q-gap: read first sample's 12 Q-values from q_out_buf. - // Stream already synced. Branch layout: [dir(3), mag(3), ord(3), urg(3)]. - let mut q12 = [0.0_f32; 12]; - unsafe { - cudarc::driver::sys::cuMemcpyDtoH_v2( - q12.as_mut_ptr().cast(), + cudarc::driver::sys::cuMemcpyDtoHAsync_v2( + self.q_readback_pinned.add(7).cast(), self.q_out_buf.raw_ptr(), 12 * std::mem::size_of::(), + self.stream.cu_stream(), ); + cudarc::driver::sys::cuStreamSynchronize(self.stream.cu_stream()); } + // Read from pinned buffer — host[0..7] = q_stats, host[7..19] = q_out sample 0 + let host: [f32; 7] = unsafe { + let mut h = [0.0_f32; 7]; + std::ptr::copy_nonoverlapping(self.q_readback_pinned, h.as_mut_ptr(), 7); + h + }; + let q12: [f32; 12] = unsafe { + let mut q = [0.0_f32; 12]; + std::ptr::copy_nonoverlapping(self.q_readback_pinned.add(7), q.as_mut_ptr(), 12); + q + }; + + // Per-branch Q-gap: first sample's 12 Q-values. + // Branch layout: [dir(3), mag(3), ord(3), urg(3)]. let branch_sizes = [ self.config.branch_0_size, self.config.branch_1_size, self.config.branch_2_size, self.config.branch_3_size, @@ -5333,12 +5362,13 @@ impl GpuDqnTrainer { offset += bs; } - // Upload per-branch Q-gaps to device buffer (16 bytes HtoD — once per 50 steps). + // Write per-branch Q-gaps to pinned device-mapped buffer — no HtoD copy. + // GPU reads via dev_ptr (qlstm_step + liquid_tau_rk4_step kernels). unsafe { - cudarc::driver::sys::cuMemcpyHtoD_v2( - self.per_branch_q_gaps_buf.raw_ptr(), - self.last_per_branch_q_gaps.as_ptr().cast(), - 4 * std::mem::size_of::(), + std::ptr::copy_nonoverlapping( + self.last_per_branch_q_gaps.as_ptr(), + self.per_branch_q_gaps_pinned, + 4, ); } @@ -5368,13 +5398,13 @@ impl GpuDqnTrainer { Ok(result) } - /// Launch the xLSTM mLSTM step: reads q_stats_buf[7] + per_branch_q_gaps_buf[4], + /// Launch the xLSTM mLSTM step: reads q_stats_buf[7] + per_branch_q_gaps[4] (pinned), /// updates persistent matrix memory C[8,8] and normalizer n[8], /// writes context[8] output. Runs in a single thread (grid=1, block=1). /// Called every 50 training steps from reduce_current_q_stats. pub(crate) fn launch_qlstm_step(&self) -> Result<(), MLError> { let q_stats_ptr = self.q_stats_buf.raw_ptr(); - let gaps_ptr = self.per_branch_q_gaps_buf.raw_ptr(); + let gaps_ptr = self.per_branch_q_gaps_dev_ptr; let c_ptr = self.qlstm_c.raw_ptr(); let n_ptr = self.qlstm_n.raw_ptr(); let ctx_ptr = self.qlstm_context.raw_ptr(); @@ -5557,14 +5587,15 @@ impl GpuDqnTrainer { /// Pass 3 is submitted separately via `submit_forward_ops_ddqn()`. pub(crate) fn submit_forward_ops_main(&mut self) -> Result<(), MLError> { // ── Zero accumulators (all REQUIRED — deterministic reduce / beta=1.0 accumulation) ─ - // total_loss_buf: c51_loss_reduce writes this scalar deterministically - self.stream - .memset_zeros(&mut self.total_loss_buf) - .map_err(|e| MLError::ModelError(format!("zero total_loss: {e}")))?; - // mse_loss_buf: c51_loss_reduce writes this scalar deterministically - self.stream - .memset_zeros(&mut self.mse_loss_buf) - .map_err(|e| MLError::ModelError(format!("zero mse_loss: {e}")))?; + // total_loss + mse_loss are pinned device-mapped — zero via cuMemsetD8Async on dev_ptr. + unsafe { + cudarc::driver::sys::cuMemsetD8Async( + self.total_loss_dev_ptr, 0, std::mem::size_of::(), self.stream.cu_stream(), + ); + cudarc::driver::sys::cuMemsetD8Async( + self.mse_loss_dev_ptr, 0, std::mem::size_of::(), self.stream.cu_stream(), + ); + } // grad_buf: backward_full uses beta=1.0 GEMM accumulation — zero via ptrs unsafe { cudarc::driver::sys::cuMemsetD8Async( @@ -5594,14 +5625,17 @@ impl GpuDqnTrainer { self.stream.memset_zeros(&mut self.d_adv_logits_mse) .map_err(|e| MLError::ModelError(format!("zero d_adv_mse: {e}")))?; self.launch_mse_loss()?; - self.launch_loss_reduce(&self.mse_loss_buf)?; + self.launch_loss_reduce(self.mse_loss_dev_ptr)?; self.launch_mse_grad_to_scratch()?; // C51 path → main buffers (already zeroed above) - self.stream.memset_zeros(&mut self.q_divergence_buf) - .map_err(|e| MLError::ModelError(format!("zero q_divergence: {e}")))?; + unsafe { + cudarc::driver::sys::cuMemsetD8Async( + self.q_divergence_dev_ptr, 0, std::mem::size_of::(), self.stream.cu_stream(), + ); + } self.launch_c51_loss()?; - self.launch_loss_reduce(&self.total_loss_buf)?; + self.launch_loss_reduce(self.total_loss_dev_ptr)?; self.launch_c51_grad()?; // Blend: main = α * C51 + (1-α) * MSE @@ -5713,11 +5747,15 @@ impl GpuDqnTrainer { /// Submit loss computation + gradient ops (everything between forward and backward). /// Extracted from submit_forward_ops_main for sub-graph timing. pub(crate) fn submit_loss_and_grad_ops(&mut self) -> Result<(), MLError> { - // Zero accumulators - self.stream.memset_zeros(&mut self.total_loss_buf) - .map_err(|e| MLError::ModelError(format!("zero total_loss: {e}")))?; - self.stream.memset_zeros(&mut self.mse_loss_buf) - .map_err(|e| MLError::ModelError(format!("zero mse_loss: {e}")))?; + // Zero accumulators — pinned device-mapped, use cuMemsetD8Async on dev_ptr. + unsafe { + cudarc::driver::sys::cuMemsetD8Async( + self.total_loss_dev_ptr, 0, std::mem::size_of::(), self.stream.cu_stream(), + ); + cudarc::driver::sys::cuMemsetD8Async( + self.mse_loss_dev_ptr, 0, std::mem::size_of::(), self.stream.cu_stream(), + ); + } unsafe { cudarc::driver::sys::cuMemsetD8Async( self.ptrs.grad_buf, @@ -5739,14 +5777,17 @@ impl GpuDqnTrainer { self.stream.memset_zeros(&mut self.d_adv_logits_mse) .map_err(|e| MLError::ModelError(format!("zero d_adv_mse: {e}")))?; self.launch_mse_loss()?; - self.launch_loss_reduce(&self.mse_loss_buf)?; + self.launch_loss_reduce(self.mse_loss_dev_ptr)?; self.launch_mse_grad_to_scratch()?; // C51 path - self.stream.memset_zeros(&mut self.q_divergence_buf) - .map_err(|e| MLError::ModelError(format!("zero q_divergence: {e}")))?; + unsafe { + cudarc::driver::sys::cuMemsetD8Async( + self.q_divergence_dev_ptr, 0, std::mem::size_of::(), self.stream.cu_stream(), + ); + } self.launch_c51_loss()?; - self.launch_loss_reduce(&self.total_loss_buf)?; + self.launch_loss_reduce(self.total_loss_dev_ptr)?; self.launch_c51_grad()?; // Blend: main = α * C51 + (1-α) * MSE @@ -6262,7 +6303,7 @@ impl GpuDqnTrainer { // ── Outputs (3) ── .arg(&self.per_sample_loss_buf) .arg(&self.td_errors_buf) - .arg(&self.total_loss_buf) + .arg(&self.total_loss_dev_ptr) // ── Saved for backward (2) ── .arg(&self.save_current_lp) .arg(&self.save_projected) @@ -6287,7 +6328,7 @@ impl GpuDqnTrainer { // ── Spectral decoupling (1) ── .arg(&self.config.spectral_decoupling_lambda) // ── Q-divergence accumulator (1) ── - .arg(&self.q_divergence_buf) + .arg(&self.q_divergence_dev_ptr) // ── Adam step counter for stochastic Expected SARSA ── .arg(&self.ptrs.t_buf) .launch(LaunchConfig { @@ -6303,13 +6344,13 @@ impl GpuDqnTrainer { /// Deterministic loss reduction: sequential sum of per_sample_loss → total_loss / batch_size. /// Grid=(1,1,1), Block=(1,1,1). Replaces atomicAdd for fully deterministic training. - fn launch_loss_reduce(&self, total_loss_buf: &CudaSlice) -> Result<(), MLError> { + fn launch_loss_reduce(&self, total_loss_dev_ptr: u64) -> Result<(), MLError> { let b = self.config.batch_size as i32; unsafe { self.stream .launch_builder(&self.c51_loss_reduce_kernel) .arg(&self.per_sample_loss_buf) - .arg(total_loss_buf) + .arg(&total_loss_dev_ptr) .arg(&b) .launch(LaunchConfig { grid_dim: (1, 1, 1), @@ -6464,7 +6505,7 @@ impl GpuDqnTrainer { // ── Outputs (3) ── .arg(&self.per_sample_loss_buf) .arg(&self.td_errors_buf) - .arg(&self.mse_loss_buf) // MSE writes to separate accumulator (not total_loss_buf) + .arg(&self.mse_loss_dev_ptr) // MSE writes to separate accumulator (not total_loss) // ── Saved for backward (2) ── repurposed: save_current_lp = softmax probs, // save_projected = per-branch E[Q] values (4 floats per sample per branch) .arg(&self.save_current_lp) diff --git a/crates/ml/src/trainers/dqn/fused_training.rs b/crates/ml/src/trainers/dqn/fused_training.rs index 481cc702e..65e4afc96 100644 --- a/crates/ml/src/trainers/dqn/fused_training.rs +++ b/crates/ml/src/trainers/dqn/fused_training.rs @@ -2268,18 +2268,20 @@ impl FusedTrainingCtx { Ok(()) } - pub(crate) fn loss_gpu_buf(&self) -> &cudarc::driver::CudaSlice { + /// Raw device pointer to the active loss scalar (pinned device-mapped). + /// Returns MSE dev_ptr when c51_alpha ≈ 0, otherwise C51 total_loss dev_ptr. + pub(crate) fn loss_gpu_ptr(&self) -> u64 { if self.trainer.c51_alpha() < 1e-6 { - &self.trainer.mse_loss_buf + self.trainer.mse_loss_dev_ptr() } else { - &self.trainer.total_loss_buf + self.trainer.total_loss_dev_ptr() } } - /// Return a reference to the GPU-resident grad_norm scalar (f32). + /// Raw device pointer to grad_norm scalar (pinned device-mapped). /// Written by the CUDA graph's grad_norm kernel — valid after `replay_adam()`. - pub(crate) fn grad_norm_gpu_buf(&self) -> &cudarc::driver::CudaSlice { - &self.trainer.grad_norm_buf + pub(crate) fn grad_norm_gpu_ptr(&self) -> u64 { + self.trainer.grad_norm_dev_ptr() } /// Read NaN detection flags (synchronizes stream). Used by training guard on NaN halt. diff --git a/crates/ml/src/trainers/dqn/trainer/training_loop.rs b/crates/ml/src/trainers/dqn/trainer/training_loop.rs index 847d87c85..0627ea596 100644 --- a/crates/ml/src/trainers/dqn/trainer/training_loop.rs +++ b/crates/ml/src/trainers/dqn/trainer/training_loop.rs @@ -1376,8 +1376,8 @@ impl DQNTrainer { if train_step_count % 5 == 0 { if let Some(ref mut guard) = self.training_guard { let fused = self.fused_ctx.as_ref().unwrap(); - let loss_raw = fused.loss_gpu_buf().raw_ptr(); - let grad_raw = fused.grad_norm_gpu_buf().raw_ptr(); + let loss_raw = fused.loss_gpu_ptr(); + let grad_raw = fused.grad_norm_gpu_ptr(); let gr = guard.check_and_accumulate( loss_raw, grad_raw,