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

fix(dqn): prevent BF16 softmax overflow, increase batch size for H100

Browse Source 
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 <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-03-09 21:05:42 +01:00
parent e89fbc2b4d
commit fb18e0f1dc
3 changed files with 52 additions and 13 deletions
Download Patch File Download Diff File Expand all files Collapse all files

22
crates/ml-dqn/src/distributional_dueling.rs
View File

@@ -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)
}

29
crates/ml-dqn/src/rainbow_network.rs
View File

@@ -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))
}
}

14
crates/ml/src/trainers/dqn/config.rs
View File

@@ -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