From 4c88498b8b1ffebfd30e40304f4773a54ae37210 Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Wed, 4 Mar 2026 13:28:06 +0100 Subject: [PATCH] fix(ml): keep DQN training pipeline in BF16 on Ampere+ GPUs 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 --- crates/ml/src/dqn/agent.rs | 21 +++- crates/ml/src/dqn/distributional_dueling.rs | 5 - crates/ml/src/dqn/dqn.rs | 101 ++++++++++++-------- crates/ml/src/dqn/mixed_precision.rs | 12 +++ crates/ml/src/dqn/network.rs | 12 +-- 5 files changed, 94 insertions(+), 57 deletions(-) diff --git a/crates/ml/src/dqn/agent.rs b/crates/ml/src/dqn/agent.rs index eb819b8d1..ab02d8186 100644 --- a/crates/ml/src/dqn/agent.rs +++ b/crates/ml/src/dqn/agent.rs @@ -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)? diff --git a/crates/ml/src/dqn/distributional_dueling.rs b/crates/ml/src/dqn/distributional_dueling.rs index 0b1f7bbc4..c295e8ed6 100644 --- a/crates/ml/src/dqn/distributional_dueling.rs +++ b/crates/ml/src/dqn/distributional_dueling.rs @@ -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) } diff --git a/crates/ml/src/dqn/dqn.rs b/crates/ml/src/dqn/dqn.rs index 915bcd7b0..45b2e00e8 100644 --- a/crates/ml/src/dqn/dqn.rs +++ b/crates/ml/src/dqn/dqn.rs @@ -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 = q_values.to_vec2::() + let q_f32 = q_values.to_dtype(DType::F32).unwrap_or_else(|_| q_values.clone()); + let q_vec: Vec = q_f32.to_vec2::() .unwrap_or_else(|_| { - // Fallback for 1D tensors - vec![q_values.to_vec1::().unwrap_or_default()] + vec![q_f32.to_vec1::().unwrap_or_default()] }) .into_iter() .flatten() .collect(); - let clamped_vec: Vec = clamped.to_vec2::() + let clamped_f32 = clamped.to_dtype(DType::F32).unwrap_or_else(|_| clamped.clone()); + let clamped_vec: Vec = clamped_f32.to_vec2::() .unwrap_or_else(|_| { - // Fallback for 1D tensors - vec![clamped.to_vec1::().unwrap_or_default()] + vec![clamped_f32.to_vec1::().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::() .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 * ¬_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 = next_state_values.to_vec1()?; + let next_state_values_vec: Vec = 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 = target_q_values.to_vec1()?; + let target_vec: Vec = 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 = diff.to_vec1()?; + let diff_vec: Vec = 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 = target_dists_detached.to_vec2()? + let target_vec: Vec = 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::() .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 = 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 = q_values.to_vec2::()?.into_iter().flatten().collect(); + let q_vec: Vec = q_values.to_dtype(DType::F32)?.to_vec2::()?.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::()?; + let values = tensor.flatten_all()?.to_dtype(DType::F32)?.to_vec1::()?; for &val in values.iter() { total_count += 1; diff --git a/crates/ml/src/dqn/mixed_precision.rs b/crates/ml/src/dqn/mixed_precision.rs index 1a8ee31f5..676940178 100644 --- a/crates/ml/src/dqn/mixed_precision.rs +++ b/crates/ml/src/dqn/mixed_precision.rs @@ -185,6 +185,13 @@ pub fn detect_from_gpu_name(gpu_name: &str) -> Option { /// [`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 { let caps = crate::gpu::capabilities::cached_capabilities(); if !caps.is_cuda { @@ -201,6 +208,11 @@ pub fn detect_from_gpu_name_auto() -> Option { /// 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(_) => { diff --git a/crates/ml/src/dqn/network.rs b/crates/ml/src/dqn/network.rs index 932e0d740..8eded805d 100644 --- a/crates/ml/src/dqn/network.rs +++ b/crates/ml/src/dqn/network.rs @@ -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, ) -> CandleResult { 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) } }