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

perf(ml): eliminate CPU from DQN training hot path — GPU-resident ops only

Browse Source 
- Fix NaN detection: replace non-existent isnan() with ne(&self) (NaN≠NaN)
- C51 gradient strip: copy().detach() replaces to_vec2→from_vec GPU→CPU→GPU roundtrip
- GPU batch fast path: skip CPU fold + 5× from_vec when GpuBatch available
- Action validation: GPU clamp() replaces CPU loop in GPU batch path
- PER TD errors: keep as GPU Tensor when GpuPrioritized replay active
- PER indices: use GpuBatch.indices directly instead of CPU Vec→Tensor
- IS weights: cache GPU tensor from GpuBatch, reuse in both loss paths
- Logging: Q-values 10→500 steps, diagnostics 100→1000 steps

2758 tests pass, 0 clippy warnings.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-03-08 01:08:00 +01:00
parent 88f6d90576
commit 8f95908e16
1 changed files with 189 additions and 155 deletions
Download Patch File Download Diff File Expand all files Collapse all files

344
crates/ml/src/dqn/dqn.rs
View File

@@ -2062,8 +2062,8 @@ impl DQN {
/// * `batch` - Optional externally-provided batch. If `None`, samples from replay buffer.
fn compute_loss_internal(&mut self, batch: Option<super::replay_buffer_type::BatchSample>) -> Result<ComputeLossResult, MLError> {
// Get batch of experiences with PER IS-weights preserved when pre-sampled
let (experiences, weights, indices) = if let Some(batch_sample) = batch {
(batch_sample.experiences, batch_sample.weights, batch_sample.indices)
let batch_sample = if let Some(bs) = batch {
bs
} else {
// Sample from replay buffer (uniform or prioritized)
if !self.memory.can_sample(self.config.min_replay_size) {
@@ -2071,10 +2071,15 @@ impl DQN {
"Not enough experiences for training".to_owned(),
));
}
let batch_sample = self.memory.sample(self.config.batch_size)?;
(batch_sample.experiences, batch_sample.weights, batch_sample.indices)
self.memory.sample(self.config.batch_size)?
};
#[cfg(feature = "cuda")]
let gpu_batch_opt = batch_sample.gpu_batch;
let experiences = batch_sample.experiences;
let weights = batch_sample.weights;
let indices = batch_sample.indices;
// Initialize optimizer if not done
if self.optimizer.is_none() {
// WAVE 16H: Use Rainbow DQN Adam epsilon (1.5e-4) for numerical stability
@@ -2115,81 +2120,112 @@ impl DQN {
);
}
// Convert experiences to tensors
let batch_size = experiences.len();
let device = self.q_network.device();
// All loss-path arithmetic uses F32 to match network forward-pass outputs.
// Network weights use BF16 on Ampere+ via VarBuilder, but loss-path tensors
// (rewards, dones, gamma, atoms, PER weights) must be F32 (BUG #41).
let dtype = DType::F32;
// OPTIMIZATION: Single-pass data extraction for 5-10% throughput improvement
// Instead of 5 separate iterator passes, do one fold operation
let state_dim = self.config.state_dim;
let (states, next_states, actions, rewards, dones) = experiences.iter().fold(
(
Vec::with_capacity(batch_size * state_dim),
Vec::with_capacity(batch_size * state_dim),
Vec::with_capacity(batch_size),
Vec::with_capacity(batch_size),
Vec::with_capacity(batch_size),
),
|(mut s, mut ns, mut a, mut r, mut d), exp| {
s.extend_from_slice(&exp.state);
ns.extend_from_slice(&exp.next_state);
a.push(exp.action as u32);
r.push(exp.reward_f32());
d.push(if exp.done { 1.0_f32 } else { 0.0_f32 });
(s, ns, a, r, d)
},
);
// GPU FAST PATH: When GpuBatch is available, use pre-built GPU tensors directly.
// Eliminates: CPU fold over experiences, 5× Tensor::from_vec CPU→GPU transfers,
// and CPU action validation loop. Action bounds enforced via GPU clamp.
let (batch_size, states_tensor, next_states_tensor, actions_tensor,
rewards_tensor, dones_tensor, weights_tensor_cached) = 'tensor_prep: {
#[cfg(feature = "cuda")]
{
if let Some(ref gpu) = gpu_batch_opt {
let bs = gpu.states.dim(0).map_err(|e| {
MLError::TrainingError(format!("GPU batch dim error: {}", e))
})?;
let max_action = (self.config.num_actions.saturating_sub(1)) as f64;
break 'tensor_prep (
bs,
gpu.states.to_dtype(training_dtype(device)).map_err(|e| {
MLError::TrainingError(format!("GPU states dtype cast: {}", e))
})?,
gpu.next_states.to_dtype(training_dtype(device)).map_err(|e| {
MLError::TrainingError(format!("GPU next_states dtype cast: {}", e))
})?,
gpu.actions.clamp(0.0_f64, max_action).map_err(|e| {
MLError::TrainingError(format!("GPU action clamp: {}", e))
})?,
gpu.rewards.to_dtype(dtype).map_err(|e| {
MLError::TrainingError(format!("GPU rewards dtype cast: {}", e))
})?,
gpu.dones.to_dtype(dtype).map_err(|e| {
MLError::TrainingError(format!("GPU dones dtype cast: {}", e))
})?,
Some(gpu.weights.to_dtype(dtype).map_err(|e| {
MLError::TrainingError(format!("GPU weights dtype cast: {}", e))
})?.detach()),
);
}
}
// Create tensors
let states_tensor = Tensor::from_vec(states, (batch_size, self.config.state_dim), device)
.map_err(|e| {
MLError::TrainingError(format!("Failed to create states tensor: {}", e))
})?
.to_dtype(training_dtype(device))
.map_err(|e| {
MLError::TrainingError(format!("Failed to cast states tensor to training dtype: {}", e))
})?;
// CPU PATH: Extract experiences into flat Vecs → create GPU tensors
let batch_size = experiences.len();
let state_dim = self.config.state_dim;
let (states, next_states, actions, rewards, dones) = experiences.iter().fold(
(
Vec::with_capacity(batch_size * state_dim),
Vec::with_capacity(batch_size * state_dim),
Vec::with_capacity(batch_size),
Vec::with_capacity(batch_size),
Vec::with_capacity(batch_size),
),
|(mut s, mut ns, mut a, mut r, mut d), exp| {
s.extend_from_slice(&exp.state);
ns.extend_from_slice(&exp.next_state);
a.push(exp.action as u32);
r.push(exp.reward_f32());
d.push(if exp.done { 1.0_f32 } else { 0.0_f32 });
(s, ns, a, r, d)
},
);
let next_states_tensor =
Tensor::from_vec(next_states, (batch_size, self.config.state_dim), device).map_err(
|e| MLError::TrainingError(format!("Failed to create next states tensor: {}", e)),
)?
let states_tensor = Tensor::from_vec(states, (batch_size, self.config.state_dim), device)
.map_err(|e| {
MLError::TrainingError(format!("Failed to create states tensor: {}", e))
})?
.to_dtype(training_dtype(device))
.map_err(|e| {
MLError::TrainingError(format!("Failed to cast next states tensor to training dtype: {}", e))
MLError::TrainingError(format!("Failed to cast states tensor to training dtype: {}", e))
})?;
// Wave 11.7: Validate action indices before creating tensor (CUDA gather bounds check)
let num_actions = self.config.num_actions;
for (idx, &action) in actions.iter().enumerate() {
if action as usize >= num_actions {
return Err(MLError::TrainingError(format!(
"Invalid action {} at index {} (num_actions={})",
action, idx, num_actions
)));
let next_states_tensor =
Tensor::from_vec(next_states, (batch_size, self.config.state_dim), device).map_err(
|e| MLError::TrainingError(format!("Failed to create next states tensor: {}", e)),
)?
.to_dtype(training_dtype(device))
.map_err(|e| {
MLError::TrainingError(format!("Failed to cast next states tensor to training dtype: {}", e))
})?;
// Validate action indices before GPU transfer (CUDA gather bounds check)
let num_actions = self.config.num_actions;
for (idx, &action) in actions.iter().enumerate() {
if action as usize >= num_actions {
return Err(MLError::TrainingError(format!(
"Invalid action {} at index {} (num_actions={})",
action, idx, num_actions
)));
}
}
}
let actions_tensor = Tensor::from_vec(actions, batch_size, device).map_err(|e| {
MLError::TrainingError(format!("Failed to create actions tensor: {}", e))
})?;
let actions_tensor = Tensor::from_vec(actions, batch_size, device).map_err(|e| {
MLError::TrainingError(format!("Failed to create actions tensor: {}", e))
})?;
// BF16 FIX: Keep all loss-path tensors in F32. Network weights use BF16
// internally (via VarBuilder), but loss arithmetic must stay F32 because
// forward-pass outputs are F32 (BUG #41) and Candle forbids mixed-dtype ops.
// Casting intermediates to BF16 caused "dtype mismatch in mul/sub" on Ampere+.
// BF16 FIX: All loss-path arithmetic uses F32 to match network forward-pass
// outputs (kept F32 for autograd, see BUG #41). Network weights internally
// use training_dtype (BF16 on Ampere+) via VarBuilder, but the loss-path
// tensors (rewards, dones, gamma, atoms, PER weights, Huber constants) must
// be F32 to avoid Candle "dtype mismatch" errors in mul/sub/add.
let dtype = DType::F32;
let rewards_tensor = Tensor::from_vec(rewards, batch_size, device).map_err(|e| {
MLError::TrainingError(format!("Failed to create rewards tensor: {}", e))
})?;
let rewards_tensor = Tensor::from_vec(rewards, batch_size, device).map_err(|e| {
MLError::TrainingError(format!("Failed to create rewards tensor: {}", e))
})?;
let dones_tensor = Tensor::from_vec(dones, batch_size, device)
.map_err(|e| MLError::TrainingError(format!("Failed to create dones tensor: {}", e)))?;
let dones_tensor = Tensor::from_vec(dones, batch_size, device)
.map_err(|e| MLError::TrainingError(format!("Failed to create dones tensor: {}", e)))?;
(batch_size, states_tensor, next_states_tensor, actions_tensor,
rewards_tensor, dones_tensor, None::<Tensor>)
};
// Wave 11.6: Forward pass through main network (hybrid > dueling > standard)
// BF16 FIX: Cast forward-pass output to F32 for all loss-path arithmetic.
@@ -2304,42 +2340,24 @@ impl DQN {
.to_dtype(state_action_values.dtype())?; // Match forward pass dtype (F32) for sub
let diff = state_action_values.sub(&target_q_values)?;
// BUG #14 FIX: NaN detection — periodic to avoid GPU sync stall every step.
// Each .to_vec1() forces a GPU→CPU DMA transfer and pipeline flush. Running
// these checks every 100 steps keeps detection while allowing the GPU to
// pipeline freely between checks (reduces 3 syncs/step to 3 syncs/100 steps).
if self.training_steps % 100 == 0 {
let next_state_values_vec: Vec<f32> = next_state_values.to_dtype(DType::F32)?.to_vec1()?;
let nan_count_next = next_state_values_vec.iter().filter(|v| !v.is_finite()).count();
if nan_count_next > 0 {
tracing::warn!(
"⚠️ BUG #14: {}/{} next_state_values are NaN/Inf at step {}",
nan_count_next,
batch_size,
self.training_steps
);
// BUG #14 FIX: NaN detection — GPU-resident check, single scalar sync.
// Uses ne(self) on GPU (NaN != NaN) instead of to_vec1() which transferred
// entire batch to CPU. Now: 1 scalar DMA per tensor vs batch_size floats.
if self.training_steps % 500 == 0 {
let nsv_f32 = next_state_values.to_dtype(DType::F32)?;
let nan_next = nsv_f32.ne(&nsv_f32)?.to_dtype(DType::F32)?.sum_all()?.to_scalar::<f32>().unwrap_or(0.0) as usize;
if nan_next > 0 {
tracing::warn!("BUG #14: {}/{} next_state_values are NaN at step {}", nan_next, batch_size, self.training_steps);
}
let target_vec: Vec<f32> = target_q_values.to_dtype(DType::F32)?.to_vec1()?;
let nan_count_target = target_vec.iter().filter(|v| !v.is_finite()).count();
if nan_count_target > 0 {
tracing::warn!(
"⚠️ BUG #14: {}/{} target_q_values are NaN/Inf at step {}",
nan_count_target,
batch_size,
self.training_steps
);
let tqv_f32 = target_q_values.to_dtype(DType::F32)?;
let nan_target = tqv_f32.ne(&tqv_f32)?.to_dtype(DType::F32)?.sum_all()?.to_scalar::<f32>().unwrap_or(0.0) as usize;
if nan_target > 0 {
tracing::warn!("BUG #14: {}/{} target_q_values are NaN at step {}", nan_target, batch_size, self.training_steps);
}
let diff_vec: Vec<f32> = diff.to_dtype(DType::F32)?.to_vec1()?;
let nan_count_diff = diff_vec.iter().filter(|v| !v.is_finite()).count();
if nan_count_diff > 0 {
tracing::warn!(
"⚠️ BUG #14: {}/{} TD errors are NaN/Inf at step {}",
nan_count_diff,
batch_size,
self.training_steps
);
let diff_f32 = diff.to_dtype(DType::F32)?;
let nan_diff = diff_f32.ne(&diff_f32)?.to_dtype(DType::F32)?.sum_all()?.to_scalar::<f32>().unwrap_or(0.0) as usize;
if nan_diff > 0 {
tracing::warn!("BUG #14: {}/{} TD errors are NaN at step {}", nan_diff, batch_size, self.training_steps);
}
}
@@ -2568,25 +2586,12 @@ impl DQN {
let eps = 1e-8;
let log_probs = (current_dists + eps)?.log()?;
// BUG #37 GRADIENT FIX: Convert detached target to Vec to strip gradient metadata
// Problem: Candle zeros ALL gradients when mixing detached/attached tensors
// Solution: Create fresh tensor from Vec (no gradient history, treated as constant)
// CRITICAL FIX (BUG #41 PART 2): Even after .detach(), tensor operations can re-attach
// gradient metadata. Converting to Vec and creating fresh tensor ensures clean break.
// Without this, Candle's autograd sees mixed gradient/no-gradient paths and zeros
// ALL gradients (both current_dists AND log_probs), causing total gradient collapse.
// This fix complements line 1207 (.detach() on TD errors for PER).
let target_vec: Vec<f32> = target_dists_detached.to_dtype(DType::F32)?.to_vec2()?
.into_iter()
.flatten()
.collect();
// Create fresh tensor from vec (clean, no gradient metadata), cast back to training dtype
let target_clean = Tensor::from_vec(
target_vec,
target_dists_detached.dims(),
target_dists_detached.device()
)?.to_dtype(dtype)?;
// BUG #37 GRADIENT FIX: Strip gradient metadata via GPU-resident deep copy.
// Problem: Candle zeros ALL gradients when mixing detached/attached tensors.
// Solution: copy() allocates fresh GPU storage with no autograd ancestry,
// then detach() marks it as a leaf — equivalent to the old Vec roundtrip
// (to_vec2 → from_vec) but without any GPU→CPU→GPU data transfer.
let target_clean = target_dists_detached.to_dtype(dtype)?.copy()?.detach();
// Now compute loss: clean target (constant) × log_probs (variable)
// Gradients flow ONLY through log_probs ✅
@@ -2595,27 +2600,32 @@ impl DQN {
.neg()?
.squeeze(1)?; // [batch]
// BUG #41 FIX: Detach IS weights before loss multiplication
// IS weights are metadata from replay buffer, NOT learnable parameters
let weights_tensor = Tensor::from_vec(weights.clone(), batch_size, device)
.map_err(|e| MLError::TrainingError(format!("Failed to create weights tensor: {}", e)))?
.to_dtype(dtype)
.map_err(|e| MLError::TrainingError(format!("Failed to cast weights tensor: {}", e)))?
.detach();
// IS weights: use cached GPU tensor or create from Vec
let weights_tensor = if let Some(ref wt) = weights_tensor_cached {
wt.clone()
} else {
Tensor::from_vec(weights.clone(), batch_size, device)
.map_err(|e| MLError::TrainingError(format!("Failed to create weights tensor: {}", e)))?
.to_dtype(dtype)
.map_err(|e| MLError::TrainingError(format!("Failed to cast weights tensor: {}", e)))?
.detach()
};
let weighted_loss = (per_sample_loss * weights_tensor)?.mean_all()?;
weighted_loss
} else {
// STANDARD DQN SCALAR LOSS PATH (MSE/Huber)
// Apply importance sampling weights for PER (element-wise multiplication)
// BUG #41 FIX: Detach IS weights before loss multiplication
// IS weights are metadata from replay buffer, NOT learnable parameters
let weights_tensor = Tensor::from_vec(weights.clone(), batch_size, device)
.map_err(|e| MLError::TrainingError(format!("Failed to create weights tensor: {}", e)))?
.to_dtype(dtype)
.map_err(|e| MLError::TrainingError(format!("Failed to cast weights tensor: {}", e)))?
.detach();
// IS weights: use cached GPU tensor or create from Vec
let weights_tensor = if let Some(ref wt) = weights_tensor_cached {
wt.clone()
} else {
Tensor::from_vec(weights.clone(), batch_size, device)
.map_err(|e| MLError::TrainingError(format!("Failed to create weights tensor: {}", e)))?
.to_dtype(dtype)
.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 {
@@ -2725,20 +2735,45 @@ impl DQN {
.to_scalar::<f32>()
.map_err(|e| MLError::TrainingError(format!("Failed to extract loss: {}", e)))?;
// BUG #41 FIX: Detach diff before converting to Vec to prevent gradient graph fork.
// TD errors only need values for PER priority updates, not gradients.
// Without .detach(), Candle's autograd gets confused and produces zero gradients.
// Skip GPU→CPU sync entirely when PER is disabled (indices will be empty).
//
// GPU SATURATION: This to_vec1() is placed AFTER to_scalar() above, which already
// flushed the GPU pipeline. At this point the diff tensor data is either already in
// the DMA queue or resident in host-visible memory, making the transfer near-free.
// Previously this lived before loss computation, forcing a premature pipeline stall
// that serialized TD-error DMA with loss kernel dispatch.
let td_errors_vec: Vec<f32> = if self.config.use_per {
diff.detach().to_dtype(DType::F32)?.to_vec1()?
// 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
// as a Tensor — the trainer uses update_priorities_gpu() directly, eliminating
// the to_vec1() GPU→CPU transfer of the entire batch every step.
// For CPU PER, to_vec1() is placed AFTER to_scalar() which already flushed
// the GPU pipeline, making the additional transfer near-free.
#[cfg(feature = "cuda")]
let is_gpu_per = self.memory.is_gpu_prioritized();
#[cfg(not(feature = "cuda"))]
let is_gpu_per = false;
let (td_errors_vec, td_gpu, idx_gpu) = if self.config.use_per {
if is_gpu_per {
// GPU PER: keep TD errors on GPU, use GpuBatch indices directly
let td_tensor = diff.detach().to_dtype(DType::F32)?;
#[cfg(feature = "cuda")]
let idx_tensor = gpu_batch_opt.as_ref()
.map(|gpu| gpu.indices.clone())
.unwrap_or_else(|| {
// Fallback: create indices tensor from CPU Vec
Tensor::from_vec(
indices.iter().map(|&i| i as u32).collect::<Vec<_>>(),
indices.len(),
device,
).unwrap_or_else(|_| Tensor::zeros(indices.len(), DType::U32, device).expect("zeros"))
});
#[cfg(not(feature = "cuda"))]
let idx_tensor = Tensor::from_vec(
indices.iter().map(|&i| i as u32).collect::<Vec<_>>(),
indices.len(),
device,
)?;
(Vec::new(), Some(td_tensor), Some(idx_tensor))
} else {
// CPU PER: transfer TD errors after pipeline flush
(diff.detach().to_dtype(DType::F32)?.to_vec1()?, None, None)
}
} else {
Vec::new()
(Vec::new(), None, None)
};
Ok(ComputeLossResult {
@@ -2748,9 +2783,9 @@ impl DQN {
indices,
states_tensor,
#[cfg(feature = "cuda")]
td_errors_gpu: None,
td_errors_gpu: td_gpu,
#[cfg(feature = "cuda")]
indices_gpu: None,
indices_gpu: idx_gpu,
})
}
@@ -2793,13 +2828,12 @@ impl DQN {
// Step beta annealing for PER
self.memory.step();
// WAVE 10-A4: Real-time diagnostic monitoring
// Q-value monitoring every 10 steps
if self.training_steps % 10 == 0 {
// Q-value monitoring every 500 steps (GPU→CPU sync each call)
if self.training_steps % 500 == 0 {
self.log_q_values(&result.states_tensor)?;
}
// Dead neuron detection every 100 steps
if self.training_steps % 100 == 0 {
// Dead neuron detection every 1000 steps
if self.training_steps % 1000 == 0 {
self.log_diagnostics(grad_norm)?;
}