diff --git a/crates/ml/src/cuda_pipeline/gpu_training_guard.rs b/crates/ml/src/cuda_pipeline/gpu_training_guard.rs index 8b72916b2..2dc50db13 100644 --- a/crates/ml/src/cuda_pipeline/gpu_training_guard.rs +++ b/crates/ml/src/cuda_pipeline/gpu_training_guard.rs @@ -175,6 +175,9 @@ pub struct GpuTrainingGuard { guard_mapped: [MappedBuffer; 2], /// Current write buffer index (alternates 0/1). guard_buf_idx: usize, + host_loss_sum: f64, + host_grad_sum: f64, + host_steps: u64, /// Whether at least one step has completed (first step has no previous data). guard_has_prev: bool, @@ -239,6 +242,9 @@ impl GpuTrainingGuard { qvalue_div_func, guard_mapped, guard_buf_idx: 0, + host_loss_sum: 0.0, + host_grad_sum: 0.0, + host_steps: 0, guard_has_prev: false, qstats_mapped, qstats_buf_idx: 0, @@ -354,6 +360,13 @@ impl GpuTrainingGuard { let clipped_loss = loss.min(clip_threshold); let halt_loss_clip = loss > clip_threshold && !warmup; let halt_grad_collapse = grad_norm < collapse_threshold && !warmup; + // Host-side accumulation (replaces GPU kernel accumulator) + if !halt_nan && loss.is_finite() && grad_norm.is_finite() { + self.host_loss_sum += loss as f64; + self.host_grad_sum += grad_norm as f64; + self.host_steps += 1; + } + GuardResult { halt_nan, halt_loss_clip, @@ -369,28 +382,22 @@ impl GpuTrainingGuard { /// Direct `read_volatile` from mapped pinned memory — no memcpy needed. /// The kernel's `__threadfence_system()` ensures CPU-visible writes. pub fn read_accumulators(&mut self) -> Result<(f64, f64), MLError> { - let loss_sum = self.acc_buf.read(0) as f64; - let grad_sum = self.acc_buf.read(1) as f64; - let steps = self.acc_buf.read(2) as f64; - - if steps <= 0.0 { + // Host-side accumulators (no GPU kernel needed) + if self.host_steps == 0 { return Ok((0.0, 0.0)); } - - Ok((loss_sum / steps, grad_sum / steps)) + let avg_loss = self.host_loss_sum / self.host_steps as f64; + let avg_grad = self.host_grad_sum / self.host_steps as f64; + Ok((avg_loss, avg_grad)) } /// Zero the accumulator — write zeros directly to mapped host memory. /// No GPU command needed; the kernel uses `volatile` reads so it will /// see the zeroed values on the next launch. pub fn reset_accumulators(&mut self) -> Result<(), MLError> { - // Safety: host_ptr is valid for 3 floats, and volatile kernel reads - // will pick up the new zeros without any explicit cache flush. - unsafe { - std::ptr::write_volatile(self.acc_buf.host_ptr.add(0), 0.0_f32); - std::ptr::write_volatile(self.acc_buf.host_ptr.add(1), 0.0_f32); - std::ptr::write_volatile(self.acc_buf.host_ptr.add(2), 0.0_f32); - } + self.host_loss_sum = 0.0; + self.host_grad_sum = 0.0; + self.host_steps = 0; Ok(()) } diff --git a/crates/ml/src/trainers/dqn/trainer/training_loop.rs b/crates/ml/src/trainers/dqn/trainer/training_loop.rs index 68bfb04c1..6d2faceca 100644 --- a/crates/ml/src/trainers/dqn/trainer/training_loop.rs +++ b/crates/ml/src/trainers/dqn/trainer/training_loop.rs @@ -1543,9 +1543,9 @@ impl DQNTrainer { let guard_start = std::time::Instant::now(); // Run guard every 5th step — host-side pinned reads only, zero GPU kernel. - // Loss + grad_norm are in pinned device-mapped memory (written by GPU, - // visible to host after graph replay completes on next step's launch). - if train_step_count % 5 == 0 { + // Loss + grad_norm are from pinned readback (one-step delay due to + // double-buffering). Skip first 5 steps to avoid reading zeros. + if train_step_count >= 5 && train_step_count % 5 == 0 { if let Some(ref mut guard) = self.training_guard { let fused = self.fused_ctx.as_ref().unwrap(); let scalars = fused.last_readback_scalars();