From ce534c26a5ec6c9bb03a9afc3a428edae5bfd26d Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Tue, 10 Mar 2026 14:38:47 +0100 Subject: [PATCH] =?UTF-8?q?perf(dqn):=20wire=20GPU=20Q-value=20monitoring?= =?UTF-8?q?=20=E2=80=94=20eliminate=20to=5Fscalar/to=5Fvec2=20readbacks?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Replace CPU-bound Q-value monitoring with GPU-resident qvalue_stats and qvalue_divergence kernels from GpuTrainingGuard. The two readback sites (estimate_avg_q_value_with_early_stopping's mean_all().to_scalar() and log_q_values' to_vec2()) are now bypassed when the GPU training guard is active. CPU fallback path preserved for non-CUDA builds and guard-absent scenarios. Changes: - Add DQN::log_q_values_from_stats() accepting pre-computed GPU stats - Add DQNAgentType::log_q_values_from_stats() delegate - Replace Q-estimation in train_step_single_batch with GPU kernel path - Replace Q-estimation in train_step_with_accumulation with GPU kernel path - Import IndexOp trait for Tensor::i() in trainer.rs Co-Authored-By: Claude Opus 4.6 --- crates/ml-dqn/src/dqn.rs | 69 ++++++++++ crates/ml/src/trainers/dqn/config.rs | 21 +++ crates/ml/src/trainers/dqn/trainer.rs | 190 +++++++++++++++++++++++++- 3 files changed, 277 insertions(+), 3 deletions(-) diff --git a/crates/ml-dqn/src/dqn.rs b/crates/ml-dqn/src/dqn.rs index 17b6a2234..8bcda44f3 100644 --- a/crates/ml-dqn/src/dqn.rs +++ b/crates/ml-dqn/src/dqn.rs @@ -3701,6 +3701,75 @@ impl DQN { Ok(()) } + /// Log Q-values from pre-computed GPU statistics (zero tensor readback). + /// + /// Same divergence detection and early-stopping logic as [`log_q_values`], + /// but takes pre-computed min/max/mean/variance from the GPU reduction + /// kernel instead of downloading the full Q-matrix to CPU. + pub fn log_q_values_from_stats( + &mut self, + q_min: f32, + q_max: f32, + q_mean: f32, + q_variance: f32, + n_actions: usize, + ) -> Result<(), MLError> { + // Push to Prometheus gauges (cheap atomic set, scrapable every 15s) + training_metrics::set_training_step("dqn", "current", self.training_steps as f64); + training_metrics::set_q_value_stats("dqn", "current", q_mean as f64, q_max as f64); + + tracing::debug!( + "Step {} Q-values ({} actions): range=[{:.2}, {:.2}], mean={:.2}", + self.training_steps, n_actions, q_min, q_max, q_mean + ); + + // Q-value collapse detection + if q_mean.abs() < 0.0001 && q_variance < 0.0001 { + tracing::warn!( + "Q-VALUE COLLAPSE DETECTED at step {}: mean={:.6}, variance={:.6}, range=[{:.6}, {:.6}]", + self.training_steps, q_mean, q_variance, q_min, q_max + ); + } + + // Divergence detection (same logic as log_q_values) + const Q_VALUE_DIVERGENCE_THRESHOLD: f32 = 10000.0; + const Q_VALUE_DIVERGENCE_PATIENCE: usize = 3; + + if q_max.abs() > Q_VALUE_DIVERGENCE_THRESHOLD + || q_min.abs() > Q_VALUE_DIVERGENCE_THRESHOLD + { + self.q_value_divergence_counter += 1; + tracing::warn!( + "Q-VALUE DIVERGENCE: range=[{:.2}, {:.2}] exceeds threshold ±{:.0} at step {} (consecutive: {}/{})", + q_min, q_max, Q_VALUE_DIVERGENCE_THRESHOLD, + self.training_steps, + self.q_value_divergence_counter, Q_VALUE_DIVERGENCE_PATIENCE + ); + if self.q_value_divergence_counter >= Q_VALUE_DIVERGENCE_PATIENCE { + return Err(MLError::TrainingError(format!( + "Training stopped: Q-value divergence for {} consecutive checks (threshold: ±{:.0}, range: [{:.2}, {:.2}]). \ + Consider reducing learning rate ({:.2e}), increasing gradient clipping ({:.2}), or reducing Huber delta ({:.2})", + self.q_value_divergence_counter, + Q_VALUE_DIVERGENCE_THRESHOLD, + q_min, + q_max, + self.config.learning_rate, + self.config.gradient_clip_norm, + self.config.huber_delta + ))); + } + } else if self.q_value_divergence_counter > 0 { + tracing::info!( + "Q-values normalized: range=[{:.2}, {:.2}] within threshold ±{:.0}, resetting counter (was: {})", + q_min, q_max, Q_VALUE_DIVERGENCE_THRESHOLD, + self.q_value_divergence_counter + ); + self.q_value_divergence_counter = 0; + } + + Ok(()) + } + /// Detect dead neurons and log comprehensive diagnostics (Wave 10-A4) /// /// WAVE 23 P0 Fix #2: Now implements early stopping on consecutive gradient collapse diff --git a/crates/ml/src/trainers/dqn/config.rs b/crates/ml/src/trainers/dqn/config.rs index 81b0eea94..d9de25df1 100644 --- a/crates/ml/src/trainers/dqn/config.rs +++ b/crates/ml/src/trainers/dqn/config.rs @@ -442,6 +442,27 @@ impl DQNAgentType { } } + /// Log Q-values from pre-computed GPU statistics (zero tensor readback). + /// Returns Err if Q-value divergence detected for consecutive checks. + pub fn log_q_values_from_stats( + &mut self, + q_min: f32, + q_max: f32, + q_mean: f32, + q_variance: f32, + n_actions: usize, + ) -> Result<(), MLError> { + match self { + Self::Standard(agent) => { + agent.log_q_values_from_stats(q_min, q_max, q_mean, q_variance, n_actions) + } + Self::RegimeConditional(_agent) => { + // Regime-conditional doesn't implement Q-value divergence detection yet + Ok(()) + } + } + } + /// Get device (CPU or CUDA) pub fn device(&self) -> &Device { match self { diff --git a/crates/ml/src/trainers/dqn/trainer.rs b/crates/ml/src/trainers/dqn/trainer.rs index a9f456be8..9f72d0555 100644 --- a/crates/ml/src/trainers/dqn/trainer.rs +++ b/crates/ml/src/trainers/dqn/trainer.rs @@ -8,7 +8,7 @@ use std::sync::Arc; use std::time::Duration; use anyhow::{Context, Result}; -use candle_core::{Device, Tensor}; +use candle_core::{Device, IndexOp, Tensor}; use common::metrics::questdb_sink; use common::metrics::training_metrics; use common::CommonError; @@ -4660,7 +4660,99 @@ impl DQNTrainer { // Q-value estimation: periodic (every 50 steps) self.q_estimation_counter += 1; if self.q_estimation_counter % 50 == 1 { - self.cached_avg_q = self.estimate_avg_q_value_with_early_stopping(&mut agent).await?; + // GPU path: use qvalue_stats / qvalue_divergence kernels (zero to_scalar readback) + #[cfg(feature = "cuda")] + let mut gpu_q_done = false; + #[cfg(feature = "cuda")] + { + if let Some(ref mut guard) = self.training_guard { + let buffer = agent.memory(); + if buffer.len() > 0 { + let sample_size = buffer.len().min(10); + let batch_sample = buffer + .sample(sample_size) + .map_err(|e| anyhow::anyhow!("Q-est sample: {e}"))?; + + let state_dim = agent.get_state_dim(); + let mut batch_tensor_opt: Option = None; + + if let Some(ref gpu) = batch_sample.gpu_batch { + batch_tensor_opt = Some( + gpu.states + .to_dtype(training_dtype(agent.device())) + .map_err(|e| anyhow::anyhow!("Q-est dtype: {e}"))?, + ); + } + if batch_tensor_opt.is_none() { + let mut state_data = + Vec::with_capacity(sample_size * state_dim); + for exp in &batch_sample.experiences { + state_data.extend_from_slice(&exp.state); + } + if !state_data.is_empty() { + let tensor = Tensor::from_vec( + state_data, + (sample_size, state_dim), + &self.device, + ) + .map_err(|e| anyhow::anyhow!("Q-est tensor: {e}"))? + .to_dtype(training_dtype(&self.device)) + .map_err(|e| anyhow::anyhow!("Q-est dtype: {e}"))?; + batch_tensor_opt = Some(tensor); + } + } + + if let Some(ref batch_tensor) = batch_tensor_opt { + let batch_q_values = agent + .forward(batch_tensor) + .map_err(|e| anyhow::anyhow!("Q-est forward: {e}"))?; + let num_actions = + batch_q_values.dims().get(1).copied().unwrap_or(5); + + // Divergence check on first sample + let first_q = batch_q_values + .i(0) + .map_err(|e| anyhow::anyhow!("Q-est index: {e}"))?; + let div_result = guard + .qvalue_divergence(&first_q, num_actions, 10000.0) + .map_err(|e| anyhow::anyhow!("GPU Q-div: {e}"))?; + agent + .log_q_values_from_stats( + div_result.q_min, + div_result.q_max, + div_result.q_mean, + div_result.q_variance, + num_actions, + ) + .map_err(|e| { + tracing::info!( + "Early stopping (Q-value divergence): {}", + e + ); + anyhow::anyhow!("Early stopping: {}", e) + })?; + + // Batch average via GPU reduction + 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; + gpu_q_done = true; + } + } + } + } + // CPU fallback (no CUDA feature, or guard absent, or buffer empty) + #[cfg(feature = "cuda")] + if !gpu_q_done { + self.cached_avg_q = + self.estimate_avg_q_value_with_early_stopping(&mut agent).await?; + } + #[cfg(not(feature = "cuda"))] + { + self.cached_avg_q = + self.estimate_avg_q_value_with_early_stopping(&mut agent).await?; + } } let avg_q_value = self.cached_avg_q; @@ -4931,7 +5023,99 @@ impl DQNTrainer { // Q-value estimation: periodic (every 50 steps) self.q_estimation_counter += 1; if self.q_estimation_counter % 50 == 1 { - self.cached_avg_q = self.estimate_avg_q_value_with_early_stopping(&mut agent).await?; + // GPU path: use qvalue_stats / qvalue_divergence kernels (zero to_scalar readback) + #[cfg(feature = "cuda")] + let mut gpu_q_done = false; + #[cfg(feature = "cuda")] + { + if let Some(ref mut guard) = self.training_guard { + let buffer = agent.memory(); + if buffer.len() > 0 { + let sample_size = buffer.len().min(10); + let batch_sample = buffer + .sample(sample_size) + .map_err(|e| anyhow::anyhow!("Q-est sample: {e}"))?; + + let state_dim = agent.get_state_dim(); + let mut batch_tensor_opt: Option = None; + + if let Some(ref gpu) = batch_sample.gpu_batch { + batch_tensor_opt = Some( + gpu.states + .to_dtype(training_dtype(agent.device())) + .map_err(|e| anyhow::anyhow!("Q-est dtype: {e}"))?, + ); + } + if batch_tensor_opt.is_none() { + let mut state_data = + Vec::with_capacity(sample_size * state_dim); + for exp in &batch_sample.experiences { + state_data.extend_from_slice(&exp.state); + } + if !state_data.is_empty() { + let tensor = Tensor::from_vec( + state_data, + (sample_size, state_dim), + &self.device, + ) + .map_err(|e| anyhow::anyhow!("Q-est tensor: {e}"))? + .to_dtype(training_dtype(&self.device)) + .map_err(|e| anyhow::anyhow!("Q-est dtype: {e}"))?; + batch_tensor_opt = Some(tensor); + } + } + + if let Some(ref batch_tensor) = batch_tensor_opt { + let batch_q_values = agent + .forward(batch_tensor) + .map_err(|e| anyhow::anyhow!("Q-est forward: {e}"))?; + let num_actions = + batch_q_values.dims().get(1).copied().unwrap_or(5); + + // Divergence check on first sample + let first_q = batch_q_values + .i(0) + .map_err(|e| anyhow::anyhow!("Q-est index: {e}"))?; + let div_result = guard + .qvalue_divergence(&first_q, num_actions, 10000.0) + .map_err(|e| anyhow::anyhow!("GPU Q-div: {e}"))?; + agent + .log_q_values_from_stats( + div_result.q_min, + div_result.q_max, + div_result.q_mean, + div_result.q_variance, + num_actions, + ) + .map_err(|e| { + tracing::info!( + "Early stopping (Q-value divergence): {}", + e + ); + anyhow::anyhow!("Early stopping: {}", e) + })?; + + // Batch average via GPU reduction + 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; + gpu_q_done = true; + } + } + } + } + // CPU fallback (no CUDA feature, or guard absent, or buffer empty) + #[cfg(feature = "cuda")] + if !gpu_q_done { + self.cached_avg_q = + self.estimate_avg_q_value_with_early_stopping(&mut agent).await?; + } + #[cfg(not(feature = "cuda"))] + { + self.cached_avg_q = + self.estimate_avg_q_value_with_early_stopping(&mut agent).await?; + } } let avg_q_value = self.cached_avg_q;