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 <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-03-11 10:53:13 +01:00
parent 4a90330a55
commit b122c2d0e7
2 changed files with 81 additions and 2 deletions

View File

@@ -194,6 +194,10 @@ pub struct GpuTrainingGuard {
/// Single `memcpy_dtoh` at boundary only (not per-step).
acc_buf: CudaSlice<f32>,
/// 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<f64, MLError> {
if self.q_count == 0 {
return Ok(0.0);
}
Ok(self
.q_mean_tensor
.to_scalar::<f32>()
.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 {

View File

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