perf(dqn): eliminate GPU→CPU roundtrips from training hot path
Three optimizations targeting GPU allocation churn and lock contention: 1. In-place polyak update: Var::set() reuses existing GPU buffer instead of Var::from_tensor() which allocates a new one per call. Eliminates ~10,000 cudaMalloc/cudaFree per epoch (20 params × 500 steps). 2. Fused affine ops: Replace 13 Tensor::full()/Tensor::ones() constant tensor allocations per step with tensor.affine(mul, add) — a single fused kernel. Patterns: 1-x → x.affine(-1,1), γ*x → x.affine(γ,0), 0.5*x² → (x*x).affine(0.5,0). Applied across all 4 loss paths (branching Bellman/Huber, standard Bellman/Huber, IQN). Eliminates ~6,500 GPU allocs/epoch. 3. Batch pre-sampling (K=8): Sample 8 batches under one READ lock, train all 8 under one WRITE lock. Reduces async lock acquisitions from 2×N to 2×ceil(N/8). Priority staleness across 8 steps is negligible. Combined estimated impact: 20-35% H100 throughput improvement. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
This commit is contained in:
jgrusewski
@@ -1152,6 +1152,7 @@ pub struct DQN {
|
||||
|
||||
/// When true, skip monitoring in forward() to avoid GPU->CPU sync in training hot path
|
||||
training_forward_active: bool,
|
||||
|
||||
}
|
||||
|
||||
impl DQN {
|
||||
@@ -2589,10 +2590,9 @@ impl DQN {
|
||||
.detach().to_dtype(dtype)?
|
||||
};
|
||||
|
||||
let not_done = (Tensor::ones(&[batch_size], dtype, device)? - &dones_tensor)?;
|
||||
let gamma_tensor = Tensor::full(gamma_n, &[batch_size], device)?
|
||||
.to_dtype(dtype)?;
|
||||
let target_q_values = (&rewards_tensor + (next_state_values * not_done)?.broadcast_mul(&gamma_tensor)?)?
|
||||
// affine(-1,1) = 1-x; affine(γ,0) = γ*x — fused scalar ops, zero Tensor::full allocs
|
||||
let not_done = dones_tensor.affine(-1.0, 1.0)?;
|
||||
let target_q_values = (&rewards_tensor + (next_state_values * not_done)?.affine(gamma_n as f64, 0.0)?)?
|
||||
.detach()
|
||||
.to_dtype(state_action_values.dtype())?;
|
||||
|
||||
@@ -2602,13 +2602,11 @@ impl DQN {
|
||||
let per_sample = if self.config.use_huber_loss {
|
||||
let delta = self.config.huber_delta;
|
||||
let abs_diff = weighted_diff.abs()?;
|
||||
let half = Tensor::full(0.5_f32, &[batch_size], device)?.to_dtype(dtype)?;
|
||||
let squared = (&weighted_diff * &weighted_diff)?.broadcast_mul(&half)?;
|
||||
let delta_t = Tensor::full(delta, &[batch_size], device)?.to_dtype(dtype)?;
|
||||
let half_delta_sq = Tensor::full(delta * delta * 0.5, &[batch_size], device)?.to_dtype(dtype)?;
|
||||
let linear = ((&abs_diff * &delta_t)? - half_delta_sq)?;
|
||||
// affine(0.5,0) = 0.5*x²; affine(δ,-δ²/2) = δ|x|-δ²/2; affine(-1,1) = 1-mask
|
||||
let squared = (&weighted_diff * &weighted_diff)?.affine(0.5, 0.0)?;
|
||||
let linear = abs_diff.affine(delta as f64, -(delta * delta * 0.5) as f64)?;
|
||||
let mask = abs_diff.le(delta)?.to_dtype(dtype)?;
|
||||
let inv_mask = (Tensor::ones(mask.shape(), dtype, device)? - &mask)?;
|
||||
let inv_mask = mask.affine(-1.0, 1.0)?;
|
||||
((&squared * &mask)? + (&linear * &inv_mask)?)?
|
||||
} else {
|
||||
(&weighted_diff * &weighted_diff)?
|
||||
@@ -2780,16 +2778,9 @@ impl DQN {
|
||||
// R_t + gamma*R_{t+1} + ... + gamma^{n-1}*R_{t+n-1} in the reward.
|
||||
// The bootstrap discount must therefore be gamma^n (not gamma^1).
|
||||
let gamma_n = self.config.gamma.powi(self.config.n_steps as i32);
|
||||
// GPU-native: Tensor::full creates constant tensor on device (no CPU Vec)
|
||||
let gamma_tensor = Tensor::full(gamma_n, &[batch_size], device)
|
||||
.map_err(|e| MLError::TrainingError(format!("Failed to create gamma tensor: {}", e)))?
|
||||
.to_dtype(dtype)
|
||||
.map_err(|e| MLError::TrainingError(format!("Failed to cast gamma tensor: {}", e)))?;
|
||||
|
||||
let not_done = (Tensor::ones(&[batch_size], dtype, device)? - &dones_tensor)?;
|
||||
let gamma_next = (&gamma_tensor * &next_state_values)
|
||||
.map_err(|e| MLError::TrainingError(format!("Gamma multiplication failed: {}", e)))?;
|
||||
let discounted = (&gamma_next * ¬_done)?;
|
||||
// affine(-1,1) = 1-dones; affine(γ^n,0) = γ^n * x — fused scalar ops, zero allocs
|
||||
let not_done = dones_tensor.affine(-1.0, 1.0)?;
|
||||
let discounted = (&next_state_values * ¬_done)?.affine(gamma_n as f64, 0.0)?;
|
||||
let target_q_values = (&rewards_tensor + &discounted)?
|
||||
.detach() // Stop gradient computation
|
||||
.to_dtype(state_action_values.dtype())?; // Match forward pass dtype (F32) for sub
|
||||
@@ -2869,15 +2860,11 @@ impl DQN {
|
||||
let rewards_broadcast = rewards_tensor
|
||||
.unsqueeze(1)?
|
||||
.broadcast_as((batch_size, num_quantiles))?;
|
||||
let not_done_broadcast = (Tensor::ones(&[batch_size], dtype, device)? - &dones_tensor)?
|
||||
let not_done_broadcast = dones_tensor.affine(-1.0, 1.0)?
|
||||
.unsqueeze(1)?
|
||||
.broadcast_as((batch_size, num_quantiles))?;
|
||||
|
||||
let gamma_t = Tensor::full(gamma, &[batch_size, num_quantiles], device)
|
||||
.map_err(|e| MLError::TrainingError(format!("Failed to create gamma tensor: {}", e)))?
|
||||
.to_dtype(dtype)
|
||||
.map_err(|e| MLError::TrainingError(format!("Failed to cast gamma tensor: {}", e)))?;
|
||||
let target_quantiles = (rewards_broadcast + (next_quantiles * not_done_broadcast)?.broadcast_mul(&gamma_t)?)?
|
||||
let target_quantiles = (rewards_broadcast + (next_quantiles * not_done_broadcast)?.affine(gamma as f64, 0.0)?)?
|
||||
.detach();
|
||||
|
||||
// Per-sample quantile Huber loss (Dabney et al. 2018b, Eq. 10)
|
||||
@@ -3086,7 +3073,7 @@ impl DQN {
|
||||
|
||||
// Condition: use squared if |x| <= delta, else linear
|
||||
let mask = abs_diff.le(delta)?.to_dtype(dtype)?;
|
||||
let one_minus_mask = (Tensor::ones(mask.shape(), dtype, device)? - &mask)?;
|
||||
let one_minus_mask = mask.affine(-1.0, 1.0)?;
|
||||
let huber_loss = ((&squared_loss * &mask)? + (&linear_loss * &one_minus_mask)?)?;
|
||||
|
||||
// NOW apply IS weights to the per-sample Huber losses
|
||||
|
||||
@@ -5,7 +5,7 @@
|
||||
/// 2. **Hard Updates**: Periodic full weight copy
|
||||
///
|
||||
/// Rainbow DQN uses Polyak averaging with τ=0.001 for smoother Q-value stability.
|
||||
use candle_core::{Result as CandleResult, Tensor, Var};
|
||||
use candle_core::{Result as CandleResult, Tensor};
|
||||
use candle_nn::VarMap;
|
||||
|
||||
/// Polyak averaging (soft target update)
|
||||
@@ -61,8 +61,9 @@ pub fn polyak_update(online_vars: &VarMap, target_vars: &VarMap, tau: f64) -> Ca
|
||||
// θ_target = (1-τ)*θ_target + τ*θ_online
|
||||
let online_t: &Tensor = online_tensor.as_ref();
|
||||
let target_t: &Tensor = target_tensor.as_ref();
|
||||
// In-place update reuses existing GPU buffer (no cudaMalloc per param per step)
|
||||
let new_target = ((target_t * (1.0 - tau))? + (online_t * tau)?)?;
|
||||
*target_tensor = Var::from_tensor(&new_target)?;
|
||||
target_tensor.set(&new_target)?;
|
||||
}
|
||||
}
|
||||
|
||||
@@ -204,7 +205,7 @@ pub fn compute_network_divergence(online_vars: &VarMap, target_vars: &VarMap) ->
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use candle_core::{DType, Device};
|
||||
use candle_core::{DType, Device, Var};
|
||||
|
||||
fn create_test_varmap(value: f32) -> VarMap {
|
||||
let varmap = VarMap::new();
|
||||
|
||||
@@ -2904,42 +2904,53 @@ impl DQNTrainer {
|
||||
let mut grad_acc = Tensor::new(0.0_f32, &self.device)
|
||||
.map_err(|e| anyhow::anyhow!("grad acc init: {e}"))?;
|
||||
|
||||
for _ in 0..num_training_steps {
|
||||
// Pre-sample batch (READ lock)
|
||||
let explicit_batch = {
|
||||
// Batch pre-sampling: sample K batches under one READ lock, train all
|
||||
// under one WRITE lock. Reduces async lock acquisitions from 2×N to
|
||||
// 2×ceil(N/K). Priority staleness across K steps is negligible (K=8).
|
||||
const PREFETCH_K: usize = 8;
|
||||
for chunk_start in (0..num_training_steps).step_by(PREFETCH_K) {
|
||||
let chunk_end = (chunk_start + PREFETCH_K).min(num_training_steps);
|
||||
|
||||
// Pre-sample chunk of batches (single READ lock)
|
||||
let batches = {
|
||||
let agent = self.agent.read().await;
|
||||
let buffer = agent.memory();
|
||||
buffer.can_sample(self.current_batch_size).then(|| {
|
||||
buffer
|
||||
.sample(self.current_batch_size)
|
||||
.map_err(|e| anyhow::anyhow!("Pre-sample: {e}"))
|
||||
}).transpose()?
|
||||
};
|
||||
|
||||
// GPU train step (WRITE lock) — zero CPU sync
|
||||
let gpu_result = {
|
||||
let mut agent = self.agent.write().await;
|
||||
match agent.train_step(explicit_batch) {
|
||||
Ok(r) => r,
|
||||
Err(e) => {
|
||||
let msg = e.to_string();
|
||||
if msg.contains("Early stopping") || msg.contains("Gradient collapse") {
|
||||
return Err(anyhow::anyhow!("{}", msg));
|
||||
}
|
||||
warn!("GPU train step failed: {e}, continuing...");
|
||||
continue;
|
||||
}
|
||||
let mut b = Vec::with_capacity(chunk_end - chunk_start);
|
||||
for _ in chunk_start..chunk_end {
|
||||
b.push(buffer.can_sample(self.current_batch_size).then(|| {
|
||||
buffer
|
||||
.sample(self.current_batch_size)
|
||||
.map_err(|e| anyhow::anyhow!("Pre-sample: {e}"))
|
||||
}).transpose()?);
|
||||
}
|
||||
b
|
||||
};
|
||||
|
||||
// GPU-side accumulation — detach() breaks the autograd chain so
|
||||
// candle doesn't keep every forward/backward graph alive across steps.
|
||||
// Without detach: ~32 MB/step leak (entire computation graph retained).
|
||||
loss_acc = (&loss_acc + &gpu_result.loss_gpu.detach())
|
||||
.map_err(|e| anyhow::anyhow!("loss acc: {e}"))?.detach();
|
||||
grad_acc = (&grad_acc + &gpu_result.grad_norm_gpu.detach())
|
||||
.map_err(|e| anyhow::anyhow!("grad acc: {e}"))?.detach();
|
||||
train_step_count += 1;
|
||||
// GPU train steps (single WRITE lock) — zero CPU sync
|
||||
{
|
||||
let mut agent = self.agent.write().await;
|
||||
for explicit_batch in batches {
|
||||
let gpu_result = match agent.train_step(explicit_batch) {
|
||||
Ok(r) => r,
|
||||
Err(e) => {
|
||||
let msg = e.to_string();
|
||||
if msg.contains("Early stopping") || msg.contains("Gradient collapse") {
|
||||
return Err(anyhow::anyhow!("{}", msg));
|
||||
}
|
||||
warn!("GPU train step failed: {e}, continuing...");
|
||||
continue;
|
||||
}
|
||||
};
|
||||
|
||||
// GPU-side accumulation — detach() breaks the autograd chain so
|
||||
// candle doesn't keep every forward/backward graph alive across steps.
|
||||
loss_acc = (&loss_acc + &gpu_result.loss_gpu.detach())
|
||||
.map_err(|e| anyhow::anyhow!("loss acc: {e}"))?.detach();
|
||||
grad_acc = (&grad_acc + &gpu_result.grad_norm_gpu.detach())
|
||||
.map_err(|e| anyhow::anyhow!("grad acc: {e}"))?.detach();
|
||||
train_step_count += 1;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// ═══ Single epoch-boundary readback ═══
|
||||
@@ -3028,40 +3039,47 @@ impl DQNTrainer {
|
||||
let mut grad_acc = Tensor::new(0.0_f32, &self.device)
|
||||
.map_err(|e| anyhow::anyhow!("grad acc init: {e}"))?;
|
||||
|
||||
for _ in 0..num_training_steps {
|
||||
// Pre-sample batch (READ lock)
|
||||
let explicit_batch = {
|
||||
// Batch pre-sampling: same K=8 chunking as GPU-PER path (see above)
|
||||
const PREFETCH_K: usize = 8;
|
||||
for chunk_start in (0..num_training_steps).step_by(PREFETCH_K) {
|
||||
let chunk_end = (chunk_start + PREFETCH_K).min(num_training_steps);
|
||||
|
||||
let batches = {
|
||||
let agent = self.agent.read().await;
|
||||
let buffer = agent.memory();
|
||||
buffer.can_sample(self.current_batch_size).then(|| {
|
||||
buffer
|
||||
.sample(self.current_batch_size)
|
||||
.map_err(|e| anyhow::anyhow!("Pre-sample: {e}"))
|
||||
}).transpose()?
|
||||
};
|
||||
|
||||
// GPU train step (WRITE lock) — zero CPU sync
|
||||
let gpu_result = {
|
||||
let mut agent = self.agent.write().await;
|
||||
match agent.train_step(explicit_batch) {
|
||||
Ok(r) => r,
|
||||
Err(e) => {
|
||||
let msg = e.to_string();
|
||||
if msg.contains("Early stopping") || msg.contains("Gradient collapse") {
|
||||
return Err(anyhow::anyhow!("{}", msg));
|
||||
}
|
||||
warn!("GPU train step failed: {e}, continuing...");
|
||||
continue;
|
||||
}
|
||||
let mut b = Vec::with_capacity(chunk_end - chunk_start);
|
||||
for _ in chunk_start..chunk_end {
|
||||
b.push(buffer.can_sample(self.current_batch_size).then(|| {
|
||||
buffer
|
||||
.sample(self.current_batch_size)
|
||||
.map_err(|e| anyhow::anyhow!("Pre-sample: {e}"))
|
||||
}).transpose()?);
|
||||
}
|
||||
b
|
||||
};
|
||||
|
||||
// GPU-side accumulation — detach to prevent autograd chain growth
|
||||
loss_acc = (&loss_acc + &gpu_result.loss_gpu.detach())
|
||||
.map_err(|e| anyhow::anyhow!("loss acc: {e}"))?.detach();
|
||||
grad_acc = (&grad_acc + &gpu_result.grad_norm_gpu.detach())
|
||||
.map_err(|e| anyhow::anyhow!("grad acc: {e}"))?.detach();
|
||||
train_step_count += 1;
|
||||
{
|
||||
let mut agent = self.agent.write().await;
|
||||
for explicit_batch in batches {
|
||||
let gpu_result = match agent.train_step(explicit_batch) {
|
||||
Ok(r) => r,
|
||||
Err(e) => {
|
||||
let msg = e.to_string();
|
||||
if msg.contains("Early stopping") || msg.contains("Gradient collapse") {
|
||||
return Err(anyhow::anyhow!("{}", msg));
|
||||
}
|
||||
warn!("GPU train step failed: {e}, continuing...");
|
||||
continue;
|
||||
}
|
||||
};
|
||||
|
||||
loss_acc = (&loss_acc + &gpu_result.loss_gpu.detach())
|
||||
.map_err(|e| anyhow::anyhow!("loss acc: {e}"))?.detach();
|
||||
grad_acc = (&grad_acc + &gpu_result.grad_norm_gpu.detach())
|
||||
.map_err(|e| anyhow::anyhow!("grad acc: {e}"))?.detach();
|
||||
train_step_count += 1;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// ═══ Single epoch-boundary readback ═══
|
||||
|
||||
Reference in New Issue
Block a user