gitadmin/foxhunt
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

perf(dqn): wire GPU Q-value monitoring — eliminate to_scalar/to_vec2 readbacks

Browse Source 
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 <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-03-10 14:38:47 +01:00
parent ed91f12dd2
commit ce534c26a5
3 changed files with 277 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files

69
crates/ml-dqn/src/dqn.rs
View File

@@ -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

21
crates/ml/src/trainers/dqn/config.rs
View File

@@ -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 {

190
crates/ml/src/trainers/dqn/trainer.rs
View File

@@ -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<Tensor> = 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<Tensor> = 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;