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

fix(ml): keep DQN training pipeline in BF16 on Ampere+ GPUs

Browse Source 
Remove premature F32 output casts from network forward passes
(Sequential, NetworkLayers, DistributionalDueling) that were negating
tensor-core acceleration. Instead, cast F32 input tensors (rewards,
dones, gamma, weights, atoms) to training_dtype() at the boundary
and only escape to F32 at scalar extraction points (loss value,
TD errors, logging vectors).

465 DQN tests pass, 0 clippy warnings.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-03-04 13:28:06 +01:00
parent 7c8fa991cc
commit 4c88498b8b
5 changed files with 94 additions and 57 deletions
Download Patch File Download Diff File Expand all files Collapse all files

21
crates/ml/src/dqn/agent.rs
View File

@@ -306,7 +306,11 @@ impl DQNAgent {
}
// Compute loss with proper gradient tracking
let loss = self.compute_loss(&states, &actions, &rewards, &next_states, &dones)?;
let loss_raw = self.compute_loss(&states, &actions, &rewards, &next_states, &dones)?;
// Cast loss to F32 at boundary for scalar extraction and backward pass
let loss = loss_raw.to_dtype(candle_core::DType::F32)
.map_err(|e| MLError::TrainingError(format!("Failed to cast loss to F32: {}", e)))?;
// Extract loss value before backward pass
let loss_value = loss
@@ -387,10 +391,13 @@ impl DQNAgent {
// Compute target Q-values using Bellman equation (no gradients)
let max_next_q = next_q_values.max(1)?; // Get maximum values
// Create reward and done tensors
// Create reward and done tensors, cast to training dtype at the boundary
let dtype = training_dtype(device);
let reward_tensor =
Tensor::from_vec(rewards.to_vec(), batch_size, device).map_err(|e| {
MLError::TrainingError(format!("Failed to create reward tensor: {}", e))
})?.to_dtype(dtype).map_err(|e| {
MLError::TrainingError(format!("Failed to cast reward tensor: {}", e))
})?;
let done_tensor = Tensor::from_vec(
@@ -401,7 +408,10 @@ impl DQNAgent {
batch_size,
device,
)
.map_err(|e| MLError::TrainingError(format!("Failed to create done tensor: {}", e)))?;
.map_err(|e| MLError::TrainingError(format!("Failed to create done tensor: {}", e)))?
.to_dtype(dtype).map_err(|e| {
MLError::TrainingError(format!("Failed to cast done tensor: {}", e))
})?;
// Target = reward + gamma * max(next_q) * (1 - done)
let gamma_tensor = Tensor::from_vec(
@@ -409,7 +419,10 @@ impl DQNAgent {
batch_size,
device,
)
.map_err(|e| MLError::TrainingError(format!("Failed to create gamma tensor: {}", e)))?;
.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 discounted_future = max_next_q
.squeeze(1)?

5
crates/ml/src/dqn/distributional_dueling.rs
View File

@@ -327,11 +327,6 @@ impl DistributionalDuelingQNetwork {
MLError::ModelError(format!("Softmax over atoms failed: {}", e))
})?;
// Cast output back to F32 for API compatibility
let z_probs = z_probs.to_dtype(candle_core::DType::F32).map_err(|e| {
MLError::ModelError(format!("Output dtype cast failed: {}", e))
})?;
Ok(z_probs)
}

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

@@ -955,10 +955,6 @@ impl Sequential {
}
}
// Cast output back to F32 for API compatibility
let x = x.to_dtype(DType::F32)
.map_err(|e| MLError::ModelError(format!("Output dtype cast failed: {}", e)))?;
Ok(x)
}
@@ -1334,9 +1330,10 @@ impl DQN {
let batch_size = state.dim(0)?;
let num_actions = self.config.num_actions;
// Broadcast atoms to [batch, num_actions, num_atoms]
// Note: atoms stay F32 to match z_probs output (already cast to F32 by sub-networks)
// Broadcast atoms to [batch, num_actions, num_atoms], matching training dtype
let dtype = training_dtype(&self.device);
let atoms_tensor = Tensor::from_vec(atoms, num_atoms, &self.device)?
.to_dtype(dtype)?
.unsqueeze(0)?
.unsqueeze(0)?
.broadcast_as((batch_size, num_actions, num_atoms))?;
@@ -1379,18 +1376,18 @@ impl DQN {
// Log clipped values for monitoring (every 1000 steps)
if self.training_steps % 1000 == 0 {
let q_vec: Vec<f32> = q_values.to_vec2::<f32>()
let q_f32 = q_values.to_dtype(DType::F32).unwrap_or_else(|_| q_values.clone());
let q_vec: Vec<f32> = q_f32.to_vec2::<f32>()
.unwrap_or_else(|_| {
// Fallback for 1D tensors
vec![q_values.to_vec1::<f32>().unwrap_or_default()]
vec![q_f32.to_vec1::<f32>().unwrap_or_default()]
})
.into_iter()
.flatten()
.collect();
let clamped_vec: Vec<f32> = clamped.to_vec2::<f32>()
let clamped_f32 = clamped.to_dtype(DType::F32).unwrap_or_else(|_| clamped.clone());
let clamped_vec: Vec<f32> = clamped_f32.to_vec2::<f32>()
.unwrap_or_else(|_| {
// Fallback for 1D tensors
vec![clamped.to_vec1::<f32>().unwrap_or_default()]
vec![clamped_f32.to_vec1::<f32>().unwrap_or_default()]
})
.into_iter()
.flatten()
@@ -1702,6 +1699,8 @@ impl DQN {
// Get the max probability (the selected action's probability)
let max_prob = probs.max(0)
.map_err(|e| MLError::ModelError(format!("Max prob failed: {}", e)))?
.to_dtype(DType::F32)
.map_err(|e| MLError::ModelError(format!("F32 cast failed: {}", e)))?
.to_scalar::<f32>()
.map_err(|e| MLError::ModelError(format!("Scalar conversion failed: {}", e)))?;
@@ -2012,12 +2011,18 @@ impl DQN {
MLError::TrainingError(format!("Failed to create actions tensor: {}", e))
})?;
let dtype = training_dtype(device);
let rewards_tensor = Tensor::from_vec(rewards, batch_size, device).map_err(|e| {
MLError::TrainingError(format!("Failed to create rewards tensor: {}", e))
})?.to_dtype(dtype).map_err(|e| {
MLError::TrainingError(format!("Failed to cast 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)))?;
.map_err(|e| MLError::TrainingError(format!("Failed to create dones tensor: {}", e)))?
.to_dtype(dtype).map_err(|e| {
MLError::TrainingError(format!("Failed to cast dones tensor: {}", e))
})?;
// Wave 11.6: Forward pass through main network (hybrid > dueling > standard)
let current_q_values = if self.dist_dueling_q_network.is_some() {
@@ -2062,6 +2067,7 @@ impl DQN {
// Broadcast atoms to [batch, num_actions, num_atoms]
let atoms_tensor = Tensor::from_vec(atoms, num_atoms, device)?
.to_dtype(dtype)?
.unsqueeze(0)?
.unsqueeze(0)?
.broadcast_as((batch_size, num_actions, num_atoms))?;
@@ -2097,15 +2103,14 @@ impl DQN {
.gather(&next_actions_unsqueezed, 1)?
.squeeze(1)?
.detach() // Prevent gradient flow to frozen target network
.to_dtype(DType::F32)?
.to_dtype(dtype)?
} else {
// Standard DQN: use max Q-value from target network
// Note: max(1) already returns a 1D tensor, no need to squeeze
// Ensure F32 dtype to match other tensors
// BUG #41 FIX: Detach target network values BEFORE dtype conversion
next_q_values.max(1)?
.detach() // Prevent gradient flow to frozen target network
.to_dtype(DType::F32)?
.to_dtype(dtype)?
};
// Compute target values using Bellman equation
@@ -2117,20 +2122,15 @@ impl DQN {
let gamma_tensor =
Tensor::from_vec(vec![gamma_n; batch_size], 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::F32, device)? - &dones_tensor)?;
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 * &not_done)?;
let target_q_values = (&rewards_tensor + &discounted)?.detach(); // Stop gradient computation
// Ensure F32 for loss computation (always needed regardless of NaN checks)
let target_q_values_f32 = target_q_values.to_dtype(DType::F32)?;
// Compute TD errors for priority updates (before applying weights)
// Ensure both tensors have the same dtype (F32)
let target_q_values = target_q_values_f32;
let diff = state_action_values.sub(&target_q_values)?;
// BUG #14 FIX: NaN detection — periodic to avoid GPU sync stall every step.
@@ -2138,7 +2138,7 @@ impl DQN {
// 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_vec1()?;
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!(
@@ -2149,7 +2149,7 @@ impl DQN {
);
}
let target_vec: Vec<f32> = target_q_values.to_vec1()?;
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!(
@@ -2160,7 +2160,7 @@ impl DQN {
);
}
let diff_vec: Vec<f32> = diff.to_vec1()?;
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!(
@@ -2241,12 +2241,14 @@ impl DQN {
let rewards_broadcast = rewards_tensor
.unsqueeze(1)?
.broadcast_as((batch_size, num_quantiles))?;
let not_done_broadcast = (Tensor::ones(&[batch_size], DType::F32, &device)? - &dones_tensor)?
let not_done_broadcast = (Tensor::ones(&[batch_size], dtype, &device)? - &dones_tensor)?
.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)))?;
.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)?)?
.detach();
@@ -2261,6 +2263,8 @@ impl DQN {
// Apply PER importance-sampling weights (detached — not learnable)
let weights_tensor = Tensor::from_vec(weights.clone(), batch_size, &device)?
.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
@@ -2311,6 +2315,7 @@ impl DQN {
.map(|i| v_min + i as f32 * delta_z)
.collect();
let atoms_tensor = Tensor::from_vec(atoms, num_atoms, device)?
.to_dtype(dtype)?
.unsqueeze(0)?
.unsqueeze(0)?
.broadcast_as((batch_size, self.config.num_actions, num_atoms))?;
@@ -2330,9 +2335,10 @@ impl DQN {
.map(|i| v_min + i as f32 * delta_z)
.collect();
let atoms_tensor = Tensor::from_vec(atoms, num_atoms, device)?
.to_dtype(dtype)?
.unsqueeze(0)?
.unsqueeze(0)?
.broadcast_as((batch_size, self.config.num_actions, num_atoms))?;
.broadcast_as((batch_size, self.config.num_atoms, num_atoms))?;
(all_next_dists.clone() * atoms_tensor)?.sum(2)?
};
@@ -2395,17 +2401,17 @@ impl DQN {
// 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_vec2()?
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)
// 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)?;
// Now compute loss: clean target (constant) × log_probs (variable)
// Gradients flow ONLY through log_probs ✅
@@ -2418,6 +2424,8 @@ impl DQN {
// 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();
let weighted_loss = (per_sample_loss * weights_tensor)?.mean_all()?;
@@ -2430,6 +2438,8 @@ impl DQN {
// 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();
let weighted_diff = (&diff * &weights_tensor)?;
@@ -2442,6 +2452,10 @@ impl DQN {
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))
})?
.to_dtype(dtype)
.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
@@ -2449,6 +2463,10 @@ impl DQN {
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))
})?
.to_dtype(dtype)
.map_err(|e| {
MLError::TrainingError(format!("Failed to cast delta tensor: {}", e))
})?;
let linear_loss_term1 = (&abs_diff * &delta_tensor)?;
@@ -2460,12 +2478,16 @@ impl DQN {
)
.map_err(|e| {
MLError::TrainingError(format!("Failed to create linear term tensor: {}", e))
})?
.to_dtype(dtype)
.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)
// Condition: use squared if |x| <= delta, else linear
let mask = abs_diff.le(delta)?.to_dtype(DType::F32)?; // 1.0 if |x| <= delta, 0.0 otherwise
let one_minus_mask = (Tensor::ones(mask.shape(), DType::F32, device)? - &mask)?;
let mask = abs_diff.le(delta)?.to_dtype(dtype)?; // 1.0 if |x| <= delta, 0.0 otherwise
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()?
} else {
@@ -2521,7 +2543,10 @@ impl DQN {
};
// 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
.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)))?;
@@ -2536,7 +2561,7 @@ impl DQN {
// 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_vec1()?
diff.detach().to_dtype(DType::F32)?.to_vec1()?
} else {
Vec::new()
};
@@ -2866,7 +2891,7 @@ impl DQN {
let q_values = self.q_network.forward(&first_state)?;
// Extract Q-values from network output
let q_vec: Vec<f32> = q_values.to_vec2::<f32>()?.into_iter().flatten().collect();
let q_vec: Vec<f32> = q_values.to_dtype(DType::F32)?.to_vec2::<f32>()?.into_iter().flatten().collect();
let n_actions = q_vec.len();
// Compute statistics across ALL actions
@@ -3044,7 +3069,7 @@ impl DQN {
// Check each layer's weights
for (_name, var) in vars_data.iter() {
let tensor = var.as_tensor();
let values = tensor.flatten_all()?.to_vec1::<f32>()?;
let values = tensor.flatten_all()?.to_dtype(DType::F32)?.to_vec1::<f32>()?;
for &val in values.iter() {
total_count += 1;

12
crates/ml/src/dqn/mixed_precision.rs
View File

@@ -185,6 +185,13 @@ pub fn detect_from_gpu_name(gpu_name: &str) -> Option<MixedPrecisionConfig> {
/// [`detect_from_gpu_name`].
///
/// Returns `None` on CPU or if the GPU is too old for tensor-core mixed precision.
///
/// **Note**: Currently returns detection results for informational logging only.
/// Actual mixed precision in the training loop is disabled because Candle's
/// `linear.forward()` requires matching dtypes between weights and input.
/// Weights are always F32 (via [`training_dtype`]), so enabling AMP forward
/// would cause dtype mismatches. Re-enable once Candle supports mixed-dtype
/// matmul or we add per-layer weight casting.
pub fn detect_from_gpu_name_auto() -> Option<MixedPrecisionConfig> {
let caps = crate::gpu::capabilities::cached_capabilities();
if !caps.is_cuda {
@@ -201,6 +208,11 @@ pub fn detect_from_gpu_name_auto() -> Option<MixedPrecisionConfig> {
/// This is the single source of truth for "what dtype should my weight
/// tensors and forward-pass intermediaries use?" Call it once when
/// constructing a model or trainer and thread the result through.
///
/// **Important**: Any F32 tensors entering the training loop (replay buffer
/// samples, reward/done tensors, etc.) must be cast to this dtype at the
/// boundary via [`ensure_training_dtype`] before they interact with model
/// weight tensors.
pub fn training_dtype(device: &candle_core::Device) -> candle_core::DType {
match device {
candle_core::Device::Cuda(_) => {

12
crates/ml/src/dqn/network.rs
View File

@@ -2,7 +2,7 @@
use std::sync::atomic::{AtomicBool, AtomicU32, AtomicU64, Ordering};
use candle_core::{DType, Device, Result as CandleResult, Tensor};
use candle_core::{Device, Result as CandleResult, Tensor};
use crate::dqn::mixed_precision::training_dtype;
use candle_nn::Module;
use candle_nn::{ops::leaky_relu, Dropout, Linear, VarBuilder, VarMap};
@@ -196,7 +196,7 @@ impl NetworkLayers {
mixed_precision: &Option<crate::dqn::mixed_precision::MixedPrecisionConfig>,
) -> CandleResult<Tensor> {
let xs = crate::dqn::mixed_precision::ensure_training_dtype(xs)?;
let (x_input, use_amp) = match mixed_precision {
let (x_input, _use_amp) = match mixed_precision {
Some(mp) if mp.enabled => {
let target_dtype = mp.dtype.to_dtype();
match xs.to_dtype(target_dtype) {
@@ -217,11 +217,6 @@ impl NetworkLayers {
}
}
// Cast back to FP32 for loss computation and gradient stability
if use_amp {
x = x.to_dtype(DType::F32)?;
}
Ok(x)
}
}
@@ -243,9 +238,6 @@ impl Module for NetworkLayers {
}
}
// Cast output back to F32 for API compatibility
x = x.to_dtype(DType::F32)?;
Ok(x)
}
}