From fb18e0f1dcbdb8add53d7517962a505f4ae0bfff Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Mon, 9 Mar 2026 21:05:42 +0100 Subject: [PATCH] fix(dqn): prevent BF16 softmax overflow, increase batch size for H100 MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit BF16's 7-bit mantissa overflows on exp(50+) in softmax → Inf → NaN. Cast distributional logits to F32 before softmax in both dueling and rainbow network heads, then cast back to original dtype. Increase default batch_size 128→1024 (standard), 128→256 (conservative), 512→2048 (aggressive) to saturate H100 tensor cores. AutoBatchSizer already caps to VRAM ceiling for smaller GPUs. Co-Authored-By: Claude Opus 4.6 --- crates/ml-dqn/src/distributional_dueling.rs | 22 +++++++++++++--- crates/ml-dqn/src/rainbow_network.rs | 29 ++++++++++++++++++--- crates/ml/src/trainers/dqn/config.rs | 14 +++++----- 3 files changed, 52 insertions(+), 13 deletions(-) diff --git a/crates/ml-dqn/src/distributional_dueling.rs b/crates/ml-dqn/src/distributional_dueling.rs index 70a20b8ce..2907f9fbb 100644 --- a/crates/ml-dqn/src/distributional_dueling.rs +++ b/crates/ml-dqn/src/distributional_dueling.rs @@ -362,11 +362,27 @@ impl DistributionalDuelingQNetwork { MLError::ModelError(format!("Distribution combination failed: {}", e)) })?; - // Apply softmax across atoms to get valid probability distributions - // Softmax over last dimension (num_atoms) - let z_probs = candle_nn::ops::softmax(&z_dist, z_dist.rank() - 1).map_err(|e| { + // Apply softmax across atoms to get valid probability distributions. + // Cast to F32 before softmax to prevent BF16 overflow (7-bit mantissa + // can't represent exp(50+) — produces Inf → NaN after normalization). + let orig_dtype = z_dist.dtype(); + let z_dist_f32 = if orig_dtype != candle_core::DType::F32 { + z_dist.to_dtype(candle_core::DType::F32).map_err(|e| { + MLError::ModelError(format!("Cast to F32 for softmax failed: {}", e)) + })? + } else { + z_dist + }; + let z_probs_f32 = candle_nn::ops::softmax(&z_dist_f32, z_dist_f32.rank() - 1).map_err(|e| { MLError::ModelError(format!("Softmax over atoms failed: {}", e)) })?; + let z_probs = if orig_dtype != candle_core::DType::F32 { + z_probs_f32.to_dtype(orig_dtype).map_err(|e| { + MLError::ModelError(format!("Cast back from F32 after softmax failed: {}", e)) + })? + } else { + z_probs_f32 + }; Ok(z_probs) } diff --git a/crates/ml-dqn/src/rainbow_network.rs b/crates/ml-dqn/src/rainbow_network.rs index dc7ad9e50..3015c556b 100644 --- a/crates/ml-dqn/src/rainbow_network.rs +++ b/crates/ml-dqn/src/rainbow_network.rs @@ -314,9 +314,21 @@ impl RainbowNetwork { .add(&advantage_reshaped)? .sub(&advantage_mean_broadcasted)?; - // Apply softmax to get valid distributions + // Apply softmax to get valid distributions. + // Cast to F32 before softmax to prevent BF16 overflow. let q_dist_flat = q_dist.reshape((batch_size * num_actions, num_atoms))?; - let q_dist_softmax = candle_nn::ops::softmax_last_dim(&q_dist_flat)?; + let orig_dtype = q_dist_flat.dtype(); + let q_flat_f32 = if orig_dtype != candle_core::DType::F32 { + q_dist_flat.to_dtype(candle_core::DType::F32)? + } else { + q_dist_flat + }; + let q_dist_softmax = candle_nn::ops::softmax_last_dim(&q_flat_f32)?; + let q_dist_softmax = if orig_dtype != candle_core::DType::F32 { + q_dist_softmax.to_dtype(orig_dtype)? + } else { + q_dist_softmax + }; q_dist_softmax.reshape((batch_size, num_actions, num_atoms)) } else { // Standard DQN with distributional output @@ -326,7 +338,18 @@ impl RainbowNetwork { let num_atoms = self.config.distributional.num_atoms; let q_dist_reshaped = q_dist.reshape((batch_size * num_actions, num_atoms))?; - let q_dist_softmax = candle_nn::ops::softmax_last_dim(&q_dist_reshaped)?; + let orig_dtype = q_dist_reshaped.dtype(); + let q_flat_f32 = if orig_dtype != candle_core::DType::F32 { + q_dist_reshaped.to_dtype(candle_core::DType::F32)? + } else { + q_dist_reshaped + }; + let q_dist_softmax = candle_nn::ops::softmax_last_dim(&q_flat_f32)?; + let q_dist_softmax = if orig_dtype != candle_core::DType::F32 { + q_dist_softmax.to_dtype(orig_dtype)? + } else { + q_dist_softmax + }; q_dist_softmax.reshape((batch_size, num_actions, num_atoms)) } } diff --git a/crates/ml/src/trainers/dqn/config.rs b/crates/ml/src/trainers/dqn/config.rs index 8a194d236..9315517d9 100644 --- a/crates/ml/src/trainers/dqn/config.rs +++ b/crates/ml/src/trainers/dqn/config.rs @@ -944,7 +944,7 @@ impl DQNHyperparameters { pub fn conservative() -> Self { Self { learning_rate: 0.0001, - batch_size: 128, + batch_size: 1024, // H100 80GB target; AutoBatchSizer caps for smaller GPUs gamma: 0.99, epsilon_start: 1.0, epsilon_end: 0.01, @@ -1202,7 +1202,7 @@ pub(crate) fn dqn_config_2025() -> DQNConfig { // Training Hyperparameters learning_rate: 1e-4, warmup_steps: 5000, - batch_size: 256, + batch_size: 1024, // H100 80GB easily fits 1024; AutoBatchSizer caps for smaller GPUs gamma: 0.99, gradient_clip_norm: 100.0, huber_delta: 10.0, @@ -1281,7 +1281,7 @@ pub(crate) fn dqn_config_2025_hft() -> DQNConfig { // HFT-specific overrides config.warmup_steps = 1000; - config.batch_size = 512; + config.batch_size = 1024; config.tau = 0.005; config @@ -1293,14 +1293,14 @@ pub(crate) fn dqn_config_2025_hft() -> DQNConfig { /// - Smaller network (256-128-64) /// - Lower learning rate (5e-5) /// - More aggressive exploration decay -/// - Smaller batch size (128) +/// - Smaller batch size (256) pub(crate) fn dqn_config_2025_conservative() -> DQNConfig { let mut config = dqn_config_2025(); // Conservative overrides config.hidden_dims = vec![256, 128, 64]; config.learning_rate = 5e-5; - config.batch_size = 128; + config.batch_size = 256; config.epsilon_decay = 0.999; // Faster decay to conservative policy config @@ -1312,14 +1312,14 @@ pub(crate) fn dqn_config_2025_conservative() -> DQNConfig { /// - Larger network (768-512-256) /// - Higher learning rate (3e-4) /// - Slower exploration decay -/// - Larger batch size (512) +/// - Larger batch size (2048) pub(crate) fn dqn_config_2025_aggressive() -> DQNConfig { let mut config = dqn_config_2025(); // Aggressive overrides config.hidden_dims = vec![768, 512, 256]; config.learning_rate = 3e-4; - config.batch_size = 512; + config.batch_size = 2048; config.epsilon_decay = 0.99995; // Slower decay, more exploration config