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

fix(dqn): wire GPU PER into hot path, fix IS-weight ordering and IQN CVaR sort

Browse Source 
- Wire GpuReplayBuffer into DQN constructor with OOM fallback to CPU PER
- Fix P0: GPU PER priorities never updated in single-batch train_step()
  (result.indices was empty CPU Vec; GPU tensors td_errors_gpu/indices_gpu
  were silently dropped)
- Fix IS-weight ordering: apply weights AFTER Huber loss, not before
  (weighting before nonlinear Huber shifts quadratic/linear regime boundary)
- Fix IQN CVaR: add .contiguous() before sort_last_dim (341/341 tests pass)
- Defer loss scalar readback to after backward pass (piggyback on grad flush)
- Add next_states to ExperienceBatch with episode-aware shift computation
- Add insert_batch_tensors() for direct GPU tensor insertion into replay buffer
- Make Q-value estimation periodic (every 50 steps) to reduce forward passes
- Fix CPU fallback path types (u8 action, i32 fixed-point reward, timestamp)

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-03-08 21:50:26 +01:00
parent 7e5af20373
commit e35ead5f3e
5 changed files with 237 additions and 41 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -740,8 +740,9 @@ pub struct GradientResult {
struct ComputeLossResult {
/// The loss tensor (still in the computation graph for backward pass).
loss_tensor: Tensor,
/// Scalar loss value extracted from the tensor.
loss_value: f32,
/// Loss tensor cast to F32 for deferred scalar readback.
/// Read via `to_scalar::<f32>()` AFTER backward pass to avoid premature GPU flush.
loss_f32_tensor: Tensor,
/// TD errors for PER priority updates.
td_errors: Vec<f32>,
/// Replay buffer indices for PER priority updates.
@@ -1204,14 +1205,41 @@ impl DQN {
// Create experience replay buffer (uniform or prioritized based on config)
let memory = if config.use_per {
// Create Prioritized Experience Replay buffer
super::replay_buffer_type::ReplayBufferType::new_prioritized(
config.replay_buffer_capacity,
config.per_alpha,
config.per_beta_start,
config.per_beta_max,
config.per_beta_annealing_steps,
)?
// GPU PER: on CUDA, allocate GPU-resident ring buffer with OOM fallback to CPU PER.
// This activates the GpuBatch fast path in compute_loss_internal() and
// GPU TD error retention — eliminating 5 of 6 CPU↔GPU roundtrips per train step.
#[cfg(feature = "cuda")]
{
if device.is_cuda() {
super::replay_buffer_type::ReplayBufferType::try_gpu_prioritized_with_fallback(
config.replay_buffer_capacity,
config.state_dim,
config.per_alpha,
config.per_beta_start,
config.per_beta_max,
config.per_beta_annealing_steps,
&device,
)?
} else {
super::replay_buffer_type::ReplayBufferType::new_prioritized(
config.replay_buffer_capacity,
config.per_alpha,
config.per_beta_start,
config.per_beta_max,
config.per_beta_annealing_steps,
)?
}
}
#[cfg(not(feature = "cuda"))]
{
super::replay_buffer_type::ReplayBufferType::new_prioritized(
config.replay_buffer_capacity,
config.per_alpha,
config.per_beta_start,
config.per_beta_max,
config.per_beta_annealing_steps,
)?
}
} else {
// Create uniform replay buffer (standard DQN)
// P2 FIX: Initialize at BASE_CAPACITY (10K) for adaptive growth
@@ -2639,14 +2667,14 @@ impl DQN {
.map_err(|e| MLError::TrainingError(format!("Failed to cast weights tensor: {}", e)))?
.detach()
};
let weighted_diff = (&diff * &weights_tensor)?;
if self.config.use_huber_loss {
// Huber loss: L(x) = 0.5 * x^2 if |x| <= delta, else delta * (|x| - 0.5*delta)
// Huber loss on RAW diff first, then weight by IS weights.
// L_i = Huber(diff_i), loss = mean(w_i * L_i)
// IS weights must be applied AFTER the nonlinear Huber function,
// not before — otherwise the quadratic/linear regime boundaries shift.
let delta = self.config.huber_delta;
let abs_diff = weighted_diff.abs()?;
let abs_diff = diff.abs()?;
// Element-wise Huber loss (use weighted_diff for loss computation)
let half = Tensor::from_vec(vec![0.5_f32; batch_size], batch_size, device)
.map_err(|e| {
MLError::TrainingError(format!("Failed to create half tensor: {}", e))
@@ -2655,9 +2683,8 @@ impl DQN {
.map_err(|e| {
MLError::TrainingError(format!("Failed to cast half tensor: {}", e))
})?;
let squared_loss = (&weighted_diff * &weighted_diff)?.broadcast_mul(&half)?; // 0.5 * x^2
let squared_loss = (&diff * &diff)?.broadcast_mul(&half)?; // 0.5 * x^2
// Create delta tensor for operations
let delta_tensor = Tensor::from_vec(vec![delta; batch_size], batch_size, device)
.map_err(|e| {
MLError::TrainingError(format!("Failed to create delta tensor: {}", e))
@@ -2681,16 +2708,18 @@ impl DQN {
.map_err(|e| {
MLError::TrainingError(format!("Failed to cast linear term tensor: {}", e))
})?;
let linear_loss = (linear_loss_term1 - &linear_loss_term2_tensor)?; // delta * (|x| - 0.5*delta)
let linear_loss = (linear_loss_term1 - &linear_loss_term2_tensor)?;
// Condition: use squared if |x| <= delta, else linear
let mask = abs_diff.le(delta)?.to_dtype(dtype)?; // 1.0 if |x| <= delta, 0.0 otherwise
let mask = abs_diff.le(delta)?.to_dtype(dtype)?;
let one_minus_mask = (Tensor::ones(mask.shape(), dtype, device)? - &mask)?;
let huber_loss = ((&squared_loss * &mask)? + (&linear_loss * &one_minus_mask)?)?;
huber_loss.mean_all()?
// NOW apply IS weights to the per-sample Huber losses
(huber_loss * weights_tensor)?.mean_all()?
} else {
// MSE fallback (use weighted_diff)
(&weighted_diff * &weighted_diff)?.mean_all()?
// MSE: L_i = diff_i^2, loss = mean(w_i * L_i)
((&diff * &diff)? * weights_tensor)?.mean_all()?
}
};
@@ -2740,13 +2769,12 @@ impl DQN {
loss_with_entropy
};
// Extract scalar loss value (forces GPU→CPU pipeline flush)
// Cast to F32 for scalar extraction — loss_tensor may be BF16 on Ampere+ GPUs
let loss_value = loss_tensor
// Defer loss scalar readback: keep tensor on GPU to avoid premature pipeline flush.
// to_scalar() is called AFTER backward pass in train_step() / compute_gradients(),
// piggybacking on the grad norm flush (zero additional roundtrip).
let loss_f32_tensor = loss_tensor
.to_dtype(DType::F32)
.map_err(|e| MLError::TrainingError(format!("Failed to cast loss to F32: {}", e)))?
.to_scalar::<f32>()
.map_err(|e| MLError::TrainingError(format!("Failed to extract loss: {}", e)))?;
.map_err(|e| MLError::TrainingError(format!("Failed to cast loss to F32: {}", e)))?;
// BUG #41 FIX: Detach diff for PER priority updates (values only, no gradients).
// GPU SATURATION: When GpuPrioritized replay is active, keep TD errors on GPU
@@ -2792,7 +2820,7 @@ impl DQN {
Ok(ComputeLossResult {
loss_tensor,
loss_value,
loss_f32_tensor: loss_f32_tensor,
td_errors: td_errors_vec,
indices,
states_tensor,
@@ -2835,8 +2863,21 @@ impl DQN {
self.training_steps += 1;
// Update priorities for PER (if using prioritized replay)
#[cfg(feature = "cuda")]
{
if let (Some(ref td_gpu), Some(ref idx_gpu)) =
(&result.td_errors_gpu, &result.indices_gpu)
{
self.memory.update_priorities_gpu(idx_gpu, td_gpu)?;
} else if !result.indices.is_empty() {
self.memory
.update_priorities(&result.indices, &result.td_errors)?;
}
}
#[cfg(not(feature = "cuda"))]
if !result.indices.is_empty() {
self.memory.update_priorities(&result.indices, &result.td_errors)?;
self.memory
.update_priorities(&result.indices, &result.td_errors)?;
}
// Step beta annealing for PER
@@ -2854,7 +2895,13 @@ impl DQN {
// Update target network with cosine-annealed EMA or hard updates
self.update_target_networks()?;
Ok((result.loss_value, grad_norm))
// Deferred loss readback: AFTER backward pass, the GPU pipeline is already
// flushed by grad_norm readback — this to_scalar() is near-free.
let loss_value = result.loss_f32_tensor
.to_scalar::<f32>()
.map_err(|e| MLError::TrainingError(format!("Deferred loss readback: {}", e)))?;
Ok((loss_value, grad_norm))
}
/// Compute clipped gradients for a batch WITHOUT applying an optimizer step.
@@ -2888,8 +2935,13 @@ impl DQN {
));
};
// Deferred loss readback: AFTER backward_and_clip, GPU pipeline is flushed.
let loss_value = result.loss_f32_tensor
.to_scalar::<f32>()
.map_err(|e| MLError::TrainingError(format!("Deferred loss readback: {}", e)))?;
Ok(GradientResult {
loss: result.loss_value,
loss: loss_value,
grad_norm: grad_norm as f32,
grads,
td_errors: result.td_errors,

2
crates/ml-dqn/src/quantile_regression.rs
View File

@@ -295,7 +295,7 @@ impl QuantileNetwork {
// QR-DQN has fixed uniform taus so quantiles are already ordered, but IQN
// uses random taus and the quantile outputs may not be monotonic. Sorting
// guarantees correct CVaR (mean of the worst-alpha fraction) in both cases.
let sorted_quantiles = quantiles.sort_last_dim(true)?.0; // ascending
let sorted_quantiles = quantiles.contiguous()?.sort_last_dim(true)?.0; // ascending
let tail = sorted_quantiles.narrow(2, 0, num_tail)?;
tail.mean(2)
}

34
crates/ml/src/cuda_pipeline/gpu_experience_collector.rs
View File

@@ -157,6 +157,9 @@ impl ExperienceCollectorConfig {
pub struct ExperienceBatch {
/// Flattened state vectors `[N * L * STATE_DIM]`
pub states: Vec<f32>,
/// Flattened next-state vectors `[N * L * STATE_DIM]`
/// `next_states[t] = states[t+1]` within each episode; terminal uses current state.
pub next_states: Vec<f32>,
/// Selected action indices `[N * L]`
pub actions: Vec<i32>,
/// Combined rewards `[N * L]`
@@ -687,9 +690,38 @@ impl GpuExperienceCollector {
"Experience collection complete"
);
// ---- Step 7: Return batch ----
// ---- Step 7: Build next_states from states (episode-aware shift) ----
// next_states[t] = states[t+1] within same episode; terminal/last uses current state.
let sd = self.state_dim;
let mut next_states = Vec::with_capacity(total * sd);
for ep in 0..n_episodes {
for t in 0..timesteps {
let idx = ep * timesteps + t;
let is_done = done_flags.get(idx).copied().unwrap_or(0) != 0;
if !is_done && t + 1 < timesteps {
let next_idx = ep * timesteps + t + 1;
let ns_start = next_idx * sd;
let ns_end = ns_start + sd;
if let Some(ns) = states.get(ns_start..ns_end) {
next_states.extend_from_slice(ns);
} else {
// Out of bounds — copy current state as fallback
let s_start = idx * sd;
next_states.extend_from_slice(&states[s_start..s_start + sd]);
}
} else {
// Terminal or last timestep — next_state = current state
// (masked by done flag in Bellman equation)
let s_start = idx * sd;
next_states.extend_from_slice(&states[s_start..s_start + sd]);
}
}
}
// ---- Step 8: Return batch ----
Ok(ExperienceBatch {
states,
next_states,
actions,
rewards,
done_flags,

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

@@ -281,6 +281,65 @@ impl DQNAgentType {
}
}
/// Insert a batch of experience tensors directly into the GPU replay buffer.
///
/// When the replay buffer is `GpuPrioritized`, this feeds tensors straight to
/// `GpuReplayBuffer::insert_batch()` — zero CPU Experience intermediaries.
/// For CPU buffers, converts tensors back to `Vec<Experience>` (fallback path).
#[cfg(feature = "cuda")]
pub fn insert_batch_tensors(
&self,
states: &candle_core::Tensor,
next_states: &candle_core::Tensor,
actions: &candle_core::Tensor,
rewards: &candle_core::Tensor,
dones: &candle_core::Tensor,
) -> Result<(), MLError> {
match self {
Self::Standard(agent) => {
if let Some(mut gpu_buf) = agent.memory.as_gpu_buffer() {
gpu_buf.insert_batch(states, next_states, actions, rewards, dones)
} else {
// CPU PER fallback: convert tensors back to experiences
let batch_size = states.dim(0).map_err(|e| {
MLError::TrainingError(format!("states dim: {e}"))
})?;
let state_dim = states.dim(1).map_err(|e| {
MLError::TrainingError(format!("states dim1: {e}"))
})?;
let s_flat: Vec<f32> = states.to_dtype(candle_core::DType::F32)?.flatten_all()?.to_vec1()
.map_err(|e| MLError::TrainingError(format!("states to_vec1: {e}")))?;
let ns_flat: Vec<f32> = next_states.to_dtype(candle_core::DType::F32)?.flatten_all()?.to_vec1()
.map_err(|e| MLError::TrainingError(format!("next_states to_vec1: {e}")))?;
let a_vec: Vec<u32> = actions.to_dtype(candle_core::DType::U32)?.to_vec1()
.map_err(|e| MLError::TrainingError(format!("actions to_vec1: {e}")))?;
let r_vec: Vec<f32> = rewards.to_dtype(candle_core::DType::F32)?.to_vec1()
.map_err(|e| MLError::TrainingError(format!("rewards to_vec1: {e}")))?;
let d_vec: Vec<f32> = dones.to_dtype(candle_core::DType::F32)?.to_vec1()
.map_err(|e| MLError::TrainingError(format!("dones to_vec1: {e}")))?;
let mut experiences = Vec::with_capacity(batch_size);
for i in 0..batch_size {
let start = i * state_dim;
let end = start + state_dim;
experiences.push(crate::dqn::Experience {
state: s_flat.get(start..end).unwrap_or(&[]).to_vec(),
action: a_vec.get(i).copied().unwrap_or(0) as u8,
reward: (r_vec.get(i).copied().unwrap_or(0.0) * 1_000_000.0) as i32,
next_state: ns_flat.get(start..end).unwrap_or(&[]).to_vec(),
done: d_vec.get(i).map_or(false, |&v| v > 0.5),
timestamp: 0, // fallback path — timestamp not meaningful for replay
});
}
agent.memory.add_batch(experiences)
}
}
Self::RegimeConditional(_) => {
// Regime-conditional agent doesn't support GPU PER (each head has its own buffer)
Ok(())
}
}
}
/// Check if agent can train (has enough replay buffer samples)
pub fn can_train(&self) -> bool {
match self {

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

@@ -209,6 +209,11 @@ pub struct DQNTrainer {
/// Current effective batch size (may be reduced by OOM recovery)
current_batch_size: usize,
/// Cached Q-value estimate for periodic monitoring (avoids extra forward pass every step)
cached_avg_q: f64,
/// Counter for Q-value estimation frequency (estimate every N training steps)
q_estimation_counter: u64,
/// Pre-uploaded GPU training data (set once, reused across epochs)
gpu_data: Option<DqnGpuData>,
@@ -864,6 +869,10 @@ impl DQNTrainer {
// OOM recovery: track effective batch size
current_batch_size: initial_batch_size,
// Q-value estimation: periodic (every 50 steps) instead of every step
cached_avg_q: 0.0,
q_estimation_counter: 0,
// GPU pipeline: pre-uploaded training data (initialized lazily at first epoch)
gpu_data: None,
#[cfg(feature = "cuda")]
@@ -1926,11 +1935,46 @@ impl DQNTrainer {
// (56 with OFI, 48 without) — must match kernel_dims.0.
let raw_dim = if self.hyperparams.mbp10_data_dir.is_some() { 51 } else { 43 };
let aligned_dim = crate::dqn::mixed_precision::align_dim_for_tensor_cores(raw_dim, &self.device);
let experiences = gpu_batch_to_experiences(&batch, aligned_dim);
let count = experiences.len();
// GPU PER fast path: create tensors from ExperienceBatch and
// insert directly into GpuReplayBuffer — no CPU Experience structs.
let count = batch.states.len() / aligned_dim;
info!("GPU collected {} experiences ({} episodes × {} timesteps)",
count, batch.n_episodes, batch.timesteps);
self.store_experiences_batch(experiences).await?;
{
let agent = self.agent.read().await;
if agent.memory().is_gpu_prioritized() && count > 0 {
// Build tensors on GPU directly from ExperienceBatch vecs.
// ExperienceBatch provides both states and next_states
// (episode-aware shift computed at collection time).
let states_t = candle_core::Tensor::from_vec(
batch.states.clone(), (count, aligned_dim), &self.device
).map_err(|e| anyhow::anyhow!("GPU PER states tensor: {e}"))?;
let next_states_t = candle_core::Tensor::from_vec(
batch.next_states.clone(), (count, aligned_dim), &self.device
).map_err(|e| anyhow::anyhow!("GPU PER next_states tensor: {e}"))?;
let actions_t = candle_core::Tensor::from_vec(
batch.actions.iter().map(|&a| a as u32).collect::<Vec<_>>(),
count, &self.device
).map_err(|e| anyhow::anyhow!("GPU PER actions tensor: {e}"))?;
let rewards_t = candle_core::Tensor::from_vec(
batch.rewards.clone(), count, &self.device
).map_err(|e| anyhow::anyhow!("GPU PER rewards tensor: {e}"))?;
let dones_t = candle_core::Tensor::from_vec(
batch.done_flags.iter().map(|&d| d as f32).collect::<Vec<_>>(),
count, &self.device
).map_err(|e| anyhow::anyhow!("GPU PER dones tensor: {e}"))?;
agent.insert_batch_tensors(
&states_t, &next_states_t, &actions_t, &rewards_t, &dones_t
).map_err(|e| anyhow::anyhow!("GPU PER insert_batch: {e}"))?;
} else {
// CPU PER fallback
drop(agent); // release read lock before store_experiences_batch
let experiences = gpu_batch_to_experiences(&batch, aligned_dim);
self.store_experiences_batch(experiences).await?;
}
}
true
}
Err(e) => {
@@ -3883,8 +3927,13 @@ impl DQNTrainer {
loss_f32
};
// Get Q-values from a sample state for monitoring (also performs Q-value divergence check)
let avg_q_value = self.estimate_avg_q_value_with_early_stopping(&mut agent).await?;
// Q-value estimation: periodic (every 50 steps) to avoid extra forward pass overhead.
// Q-value divergence check still runs at 50-step intervals — sufficient for early stopping.
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?;
}
let avg_q_value = self.cached_avg_q;
// Convert to f64 for monitoring
let grad_norm = grad_norm_f32 as f64;
@@ -4092,8 +4141,12 @@ impl DQNTrainer {
avg_loss
};
// Get Q-values for monitoring
let avg_q_value = self.estimate_avg_q_value_with_early_stopping(&mut agent).await?;
// Q-value estimation: periodic (every 50 steps) to avoid extra forward pass overhead.
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?;
}
let avg_q_value = self.cached_avg_q;
Ok((loss_clipped, avg_q_value, final_grad_norm as f64))
}