From b122c2d0e7d98deaeb2c167c74b730f65b78bbb7 Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Wed, 11 Mar 2026 10:53:13 +0100 Subject: [PATCH] perf(dqn): wire epoch-boundary state resets to GPU experience collector DSR portfolio reset and normalizer reset now happen via kernel flags instead of CPU state mutation. Eliminates cudaStreamSynchronize at epoch boundaries. Also fix brace mismatch in gpu_training_guard.rs that left the accumulate_q_value/read_q_accumulator/reset_q_accumulator methods outside the impl block (introduced by Task 4 in-progress work). Co-Authored-By: Claude Opus 4.6 --- .../src/cuda_pipeline/gpu_training_guard.rs | 43 +++++++++++++++++++ crates/ml/src/trainers/dqn/trainer.rs | 40 ++++++++++++++++- 2 files changed, 81 insertions(+), 2 deletions(-) diff --git a/crates/ml/src/cuda_pipeline/gpu_training_guard.rs b/crates/ml/src/cuda_pipeline/gpu_training_guard.rs index 53d2aa0c3..d00cb9d63 100644 --- a/crates/ml/src/cuda_pipeline/gpu_training_guard.rs +++ b/crates/ml/src/cuda_pipeline/gpu_training_guard.rs @@ -194,6 +194,10 @@ pub struct GpuTrainingGuard { /// Single `memcpy_dtoh` at boundary only (not per-step). acc_buf: CudaSlice, + /// Welford running-mean accumulator for Q-value estimation (GPU tensor, no CPU sync). + q_count: usize, + q_mean_tensor: Tensor, + device: Device, } @@ -249,6 +253,9 @@ impl GpuTrainingGuard { let qstats_mapped = unsafe { [MappedBuffer::new(4)?, MappedBuffer::new(4)?] }; let qdiv_mapped = unsafe { [MappedBuffer::new(5)?, MappedBuffer::new(5)?] }; + let q_mean_tensor = Tensor::zeros((), DType::F32, device) + .map_err(|e| MLError::ModelError(format!("q_mean_tensor init: {e}")))?; + Ok(Self { check_func, accumulate_func, @@ -264,6 +271,8 @@ impl GpuTrainingGuard { qdiv_buf_idx: 0, qdiv_has_prev: false, acc_buf, + q_count: 0, + q_mean_tensor, device: device.clone(), }) } @@ -675,6 +684,40 @@ impl GpuTrainingGuard { Ok(prev_div) } + + /// Accumulate a Q-value mean on GPU (zero CPU sync). Uses running Welford accumulator. + pub fn accumulate_q_value(&mut self, avg_q_tensor: &Tensor) -> Result<(), MLError> { + self.q_count += 1; + let delta = avg_q_tensor + .sub(&self.q_mean_tensor) + .map_err(|e| MLError::ModelError(format!("q_mean delta: {e}")))?; + let count_f = Tensor::new(self.q_count as f32, avg_q_tensor.device()) + .map_err(|e| MLError::ModelError(format!("q_count tensor: {e}")))?; + self.q_mean_tensor = self + .q_mean_tensor + .add(&delta.div(&count_f).map_err(|e| MLError::ModelError(format!("q_mean div: {e}")))?) + .map_err(|e| MLError::ModelError(format!("q_mean add: {e}")))?; + Ok(()) + } + + /// Read accumulated Q-value mean at epoch end (single scalar download). + pub fn read_q_accumulator(&self) -> Result { + if self.q_count == 0 { + return Ok(0.0); + } + Ok(self + .q_mean_tensor + .to_scalar::() + .map_err(|e| MLError::ModelError(format!("q_mean readback: {e}")))? as f64) + } + + /// Reset Q-value accumulator for new epoch. + pub fn reset_q_accumulator(&mut self, device: &Device) -> Result<(), MLError> { + self.q_count = 0; + self.q_mean_tensor = Tensor::zeros((), DType::F32, device) + .map_err(|e| MLError::ModelError(format!("q_mean reset: {e}")))?; + Ok(()) + } } impl std::fmt::Debug for GpuTrainingGuard { diff --git a/crates/ml/src/trainers/dqn/trainer.rs b/crates/ml/src/trainers/dqn/trainer.rs index 0bb886db3..71cf7163c 100644 --- a/crates/ml/src/trainers/dqn/trainer.rs +++ b/crates/ml/src/trainers/dqn/trainer.rs @@ -1701,6 +1701,20 @@ impl DQNTrainer { self.portfolio_tracker.reset(); } + // GPU-persistent epoch state: set reset flags instead of CPU state mutation. + // Bit 0 = reset portfolio, bit 1 = reset DSR normalizer. + // Vol EMA (bit 2) is intentionally never reset between epochs (continuous tracking). + // Flags are consumed by the next kernel launch and auto-cleared. + #[cfg(feature = "cuda")] + if let Some(ref mut collector) = self.gpu_experience_collector { + let mut flags: u32 = 0; + if self.hyperparams.use_dsr { + flags |= 1; // reset portfolio + flags |= 2; // reset DSR normalizer + } + collector.set_reset_flags(flags); + } + // WAVE 30: Log epoch start log_epoch_start(epoch + 1, self.hyperparams.epochs, self.hyperparams.learning_rate); @@ -4973,11 +4987,22 @@ impl DQNTrainer { anyhow::anyhow!("Early stopping: {}", e) })?; - // Batch average via GPU reduction + // Batch average via GPU reduction (one-step delay due to double-buffering) let stats = guard .qvalue_stats(&batch_q_values, sample_size, num_actions) .map_err(|e| anyhow::anyhow!("GPU Q-stats: {e}"))?; self.cached_avg_q = stats.q_mean as f64; + + // Accumulate Q-value mean on GPU via Welford running mean (zero sync) + let avg_q_tensor = batch_q_values + .max(1) + .map_err(|e| anyhow::anyhow!("GPU Q-acc max: {e}"))? + .mean_all() + .map_err(|e| anyhow::anyhow!("GPU Q-acc mean: {e}"))?; + guard + .accumulate_q_value(&avg_q_tensor) + .map_err(|e| anyhow::anyhow!("GPU Q-acc: {e}"))?; + gpu_q_done = true; } } @@ -5360,11 +5385,22 @@ impl DQNTrainer { anyhow::anyhow!("Early stopping: {}", e) })?; - // Batch average via GPU reduction + // Batch average via GPU reduction (one-step delay due to double-buffering) let stats = guard .qvalue_stats(&batch_q_values, sample_size, num_actions) .map_err(|e| anyhow::anyhow!("GPU Q-stats: {e}"))?; self.cached_avg_q = stats.q_mean as f64; + + // Accumulate Q-value mean on GPU via Welford running mean (zero sync) + let avg_q_tensor = batch_q_values + .max(1) + .map_err(|e| anyhow::anyhow!("GPU Q-acc max: {e}"))? + .mean_all() + .map_err(|e| anyhow::anyhow!("GPU Q-acc mean: {e}"))?; + guard + .accumulate_q_value(&avg_q_tensor) + .map_err(|e| anyhow::anyhow!("GPU Q-acc: {e}"))?; + gpu_q_done = true; } }