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; } }