From 30f46c04c1a0f1b635709e3e68e7fd3150867412 Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Thu, 9 Apr 2026 14:18:36 +0200 Subject: [PATCH] fix(critical): mean-reduce gradients + fix ExposureLevel 4-branch mismatch Three root causes found and fixed: 1. SUM-reduced gradients without 1/N: all CUDA loss gradient kernels (C51, MSE, IQN backward, CQL) accumulated per-sample gradients as raw SUM. At batch=16384 (H100) the raw norm was 282x larger than batch=58, causing gradient clipping to destroy signal-to-noise ratio and collapse training at epoch 2-3. Now all kernels multiply by 1/batch_size, making gradient scale batch-invariant. 2. ExposureLevel::target_exposure() used a flat 9-level scale that did not match the 4-branch dir*mag encoding. The Rust backtest evaluator computed wrong position sizes (e.g. 4x oversize for Short+Small). Now uses dir x mag formula. Also fixed is_buy/is_sell/is_hold and from_trading_action for 4-branch semantics. 3. Rust epsilon-greedy only explored 5/9 exposure combos (0..5 instead of dir*3+mag), ignored the magnitude branch on greedy, and used wrong indices for order/urgency (get(1)/get(2) instead of get(2)/get(3)). LR recalibrated: old gradient_clip_norm was accidentally a batch-size- dependent LR reducer (~60x at batch=58, ~16000x at batch=16384). With mean-reduced gradients the clip rarely fires, so LR is now the sole training speed control. Smoketest 1e-4 -> 2e-6, production 1e-4 -> 1e-5, hyperopt range [1e-5,3e-4] -> [1e-7,1e-4]. Diagnostics: FOXHUNT_GRAD_DIAG=1 enables per-stage gradient norm logging. Co-Authored-By: Claude Opus 4.6 (1M context) --- config/training/dqn-hyperopt.toml | 6 +- config/training/dqn-production.toml | 11 +- config/training/dqn-smoketest.toml | 2 +- crates/ml-core/src/action_space.rs | 32 +++-- crates/ml-core/src/common/action.rs | 116 +++++++++++------- crates/ml-dqn/src/dqn.rs | 21 +++- .../ml/src/cuda_pipeline/backward_kernels.cu | 4 + .../ml/src/cuda_pipeline/c51_grad_kernel.cu | 10 +- .../ml/src/cuda_pipeline/cql_grad_kernel.cu | 7 +- .../ml/src/cuda_pipeline/gpu_dqn_trainer.rs | 28 +++++ .../src/cuda_pipeline/iqn_dual_head_kernel.cu | 31 +++-- .../ml/src/cuda_pipeline/mse_grad_kernel.cu | 9 +- crates/ml/src/trainers/dqn/config.rs | 2 +- crates/ml/src/trainers/dqn/fused_training.rs | 63 +++++++++- 14 files changed, 252 insertions(+), 90 deletions(-) diff --git a/config/training/dqn-hyperopt.toml b/config/training/dqn-hyperopt.toml index eb76c34ba..da66ce079 100644 --- a/config/training/dqn-hyperopt.toml +++ b/config/training/dqn-hyperopt.toml @@ -1,10 +1,14 @@ # DQN Hyperopt Profile — PSO search space definition # All bounds are [min, max] ranges. The adapter reads these at runtime. # To constrain the search space, narrow the ranges here — no code changes needed. +# +# Mean-reduced gradients (2026-04-09): LR range shifted down because gradient +# clipping no longer acts as a hidden LR reducer. Old range [1e-5, 3e-4] had +# effective LR of [~5e-10, ~1.5e-8] due to SUM-reduced clipping at batch=16384. [search_space] # Base parameters -learning_rate = [0.00001, 0.0003] # log scale in adapter +learning_rate = [1e-7, 1e-4] # log scale in adapter — mean-reduced gradients batch_size = [4096, 16384] gamma = [0.90, 0.99] # wider range — v_range computed dynamically from gamma buffer_size = [50000, 100000] # log scale in adapter diff --git a/config/training/dqn-production.toml b/config/training/dqn-production.toml index 61f11cdbc..d4c4ed666 100644 --- a/config/training/dqn-production.toml +++ b/config/training/dqn-production.toml @@ -4,12 +4,17 @@ # Mixed-precision (NVIDIA AMP pattern): # - f32 master weights, Adam moments, gradients, loss accumulators # - bf16 shadow copies for cuBLAS GemmEx tensor core GEMM -# - Cosine LR decay: 1e-4 → 1e-5 over 100 epochs +# - Cosine LR decay: 1e-5 → 1e-6 over 200 epochs +# +# Mean-reduced gradients (2026-04-09): all loss kernels divide by batch_size. +# gradient_clip_norm=1.0 is a safety net (norm ≈ 1.0, clip rarely fires). +# LR is the sole training speed control. Old lr=1e-4 was masked by SUM-reduced +# clipping (effective LR was ~6e-9). Start at 1e-5 and tune from there. [training] epochs = 200 batch_size = 16384 -learning_rate = 0.0001 +learning_rate = 1e-5 gamma = 0.99 weight_decay = 0.0001 adam_epsilon = 1e-8 @@ -18,7 +23,7 @@ hidden_dim_base = 256 reward_scale = 1.0 huber_delta = 1.0 lr_decay_type = 2 -lr_min = 0.00001 +lr_min = 1e-6 max_bars = 0 # Data source: "ohlcv" (1-min candles) or "mbp10" (imbalance bars from MBP-10 order book) data_source = "mbp10" diff --git a/config/training/dqn-smoketest.toml b/config/training/dqn-smoketest.toml index 0a07d7007..e9dbc8d1f 100644 --- a/config/training/dqn-smoketest.toml +++ b/config/training/dqn-smoketest.toml @@ -12,7 +12,7 @@ [training] epochs = 3 batch_size = 64 -learning_rate = 0.0001 +learning_rate = 2e-6 gamma = 0.95 weight_decay = 0.0001 adam_epsilon = 1e-8 diff --git a/crates/ml-core/src/action_space.rs b/crates/ml-core/src/action_space.rs index f28510065..239e46680 100644 --- a/crates/ml-core/src/action_space.rs +++ b/crates/ml-core/src/action_space.rs @@ -77,11 +77,18 @@ mod tests { #[test] fn test_target_exposure_values() { - assert_eq!(ExposureLevel::Short100.target_exposure(), -1.0); - assert_eq!(ExposureLevel::Short50.target_exposure(), -0.5); - assert_eq!(ExposureLevel::Flat.target_exposure(), 0.0); - assert_eq!(ExposureLevel::Long50.target_exposure(), 0.5); - assert_eq!(ExposureLevel::Long100.target_exposure(), 1.0); + // 4-branch: dir × mag. Index = dir*3 + mag. + // dir: 0=Short(-1), 1=Flat(0), 2=Long(+1) + // mag: 0=Small(0.25), 1=Half(0.50), 2=Full(1.00) + assert_eq!(ExposureLevel::Short100.target_exposure(), -0.25); // Short×Small + assert_eq!(ExposureLevel::Short75.target_exposure(), -0.50); // Short×Half + assert_eq!(ExposureLevel::Short50.target_exposure(), -1.0); // Short×Full + assert_eq!(ExposureLevel::Short25.target_exposure(), 0.0); // Flat×Small + assert_eq!(ExposureLevel::Flat.target_exposure(), 0.0); // Flat×Half + assert_eq!(ExposureLevel::Long25.target_exposure(), 0.0); // Flat×Full + assert_eq!(ExposureLevel::Long50.target_exposure(), 0.25); // Long×Small + assert_eq!(ExposureLevel::Long75.target_exposure(), 0.50); // Long×Half + assert_eq!(ExposureLevel::Long100.target_exposure(), 1.0); // Long×Full } #[test] @@ -169,26 +176,27 @@ mod tests { #[test] fn test_extreme_actions() { - // Test Short100 aggressive market (index 0 * 9 + 0 * 3 + 2 = 2) + // 4-branch: Short100(=0) is Short×Small = -0.25 (NOT -1.0) + // Full short is Short50(=2) = Short×Full = -1.0 let short_extreme = FactoredAction::new( - ExposureLevel::Short100, + ExposureLevel::Short50, // Short×Full = -1.0 OrderType::Market, Urgency::Aggressive, ); - assert_eq!(short_extreme.to_index(), 2); + assert_eq!(short_extreme.to_index(), 2 * 9 + 0 * 3 + 2); // 20 assert_eq!(short_extreme.target_exposure(), -1.0); - assert_eq!(short_extreme.transaction_cost(), 0.0015); // Wave 2.5 calibration + assert_eq!(short_extreme.transaction_cost(), 0.0015); assert_eq!(short_extreme.urgency_weight(), 1.5); - // Test Long100 aggressive market (index 8 * 9 + 0 * 3 + 2 = 74) + // Long100(=8) is Long×Full = +1.0 (still correct) let long_extreme = FactoredAction::new( ExposureLevel::Long100, OrderType::Market, Urgency::Aggressive, ); - assert_eq!(long_extreme.to_index(), 74); + assert_eq!(long_extreme.to_index(), 8 * 9 + 0 * 3 + 2); // 74 assert_eq!(long_extreme.target_exposure(), 1.0); - assert_eq!(long_extreme.transaction_cost(), 0.0015); // Wave 2.5 calibration + assert_eq!(long_extreme.transaction_cost(), 0.0015); assert_eq!(long_extreme.urgency_weight(), 1.5); } diff --git a/crates/ml-core/src/common/action.rs b/crates/ml-core/src/common/action.rs index 48b72c49c..4774588cc 100644 --- a/crates/ml-core/src/common/action.rs +++ b/crates/ml-core/src/common/action.rs @@ -66,34 +66,48 @@ impl fmt::Display for OrderType { /// Exposure level for position sizing (-100% to +100%) #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)] pub enum ExposureLevel { - Short100 = 0, // -100% of max position - Short75 = 1, // -75% - Short50 = 2, // -50% - Short25 = 3, // -25% - Flat = 4, // 0% (neutral) - Long25 = 5, // +25% - Long50 = 6, // +50% - Long75 = 7, // +75% - Long100 = 8, // +100% + /// 4-branch factored DQN: discriminant = dir * 3 + mag + /// dir: 0=Short(-1), 1=Flat(0), 2=Long(+1) + /// mag: 0=Small(0.25), 1=Half(0.50), 2=Full(1.00) + /// + /// LEGACY NAMES: variant names reflect the old flat 9-level model. + /// Actual position = direction × magnitude (see target_exposure()). + Short100 = 0, // dir=0(Short), mag=0(Small) → -0.25 + Short75 = 1, // dir=0(Short), mag=1(Half) → -0.50 + Short50 = 2, // dir=0(Short), mag=2(Full) → -1.00 + Short25 = 3, // dir=1(Flat), mag=0(Small) → 0 + Flat = 4, // dir=1(Flat), mag=1(Half) → 0 + Long25 = 5, // dir=1(Flat), mag=2(Full) → 0 + Long50 = 6, // dir=2(Long), mag=0(Small) → +0.25 + Long75 = 7, // dir=2(Long), mag=1(Half) → +0.50 + Long100 = 8, // dir=2(Long), mag=2(Full) → +1.00 } impl ExposureLevel { - /// Get target portfolio value percentage (-1.0 to +1.0 in 0.25 steps) + /// Target exposure for the 4-branch factored DQN. + /// + /// Decodes from composite index `dir * 3 + mag`: + /// direction: [-1.0, 0.0, +1.0] + /// magnitude: [0.25, 0.50, 1.00] + /// result: direction × magnitude pub fn target_exposure(&self) -> f64 { - match self { - ExposureLevel::Short100 => -1.0, - ExposureLevel::Short75 => -0.75, - ExposureLevel::Short50 => -0.5, - ExposureLevel::Short25 => -0.25, - ExposureLevel::Flat => 0.0, - ExposureLevel::Long25 => 0.25, - ExposureLevel::Long50 => 0.5, - ExposureLevel::Long75 => 0.75, - ExposureLevel::Long100 => 1.0, - } + let idx = *self as usize; + let dir = idx / 3; + let mag = idx % 3; + let direction: f64 = match dir { + 0 => -1.0, + 2 => 1.0, + _ => 0.0, + }; + let magnitude: f64 = match mag { + 0 => 0.25, + 1 => 0.50, + _ => 1.00, + }; + direction * magnitude } - /// Convert from index (0-8) + /// Convert from composite index (0-8): dir * 3 + mag pub fn from_index(idx: usize) -> Result { match idx { 0 => Ok(ExposureLevel::Short100), @@ -115,16 +129,17 @@ impl ExposureLevel { impl fmt::Display for ExposureLevel { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + // Display as 4-branch dir×mag combo for clarity match self { - ExposureLevel::Short100 => write!(f, "Short100"), - ExposureLevel::Short75 => write!(f, "Short75"), - ExposureLevel::Short50 => write!(f, "Short50"), - ExposureLevel::Short25 => write!(f, "Short25"), - ExposureLevel::Flat => write!(f, "Flat"), - ExposureLevel::Long25 => write!(f, "Long25"), - ExposureLevel::Long50 => write!(f, "Long50"), - ExposureLevel::Long75 => write!(f, "Long75"), - ExposureLevel::Long100 => write!(f, "Long100"), + ExposureLevel::Short100 => write!(f, "S25"), // Short×Small + ExposureLevel::Short75 => write!(f, "S50"), // Short×Half + ExposureLevel::Short50 => write!(f, "S100"), // Short×Full + ExposureLevel::Short25 => write!(f, "F25"), // Flat×Small + ExposureLevel::Flat => write!(f, "F50"), // Flat×Half + ExposureLevel::Long25 => write!(f, "F100"), // Flat×Full + ExposureLevel::Long50 => write!(f, "L25"), // Long×Small + ExposureLevel::Long75 => write!(f, "L50"), // Long×Half + ExposureLevel::Long100 => write!(f, "L100"), // Long×Full } } } @@ -245,18 +260,21 @@ impl FactoredAction { } /// Create a FactoredAction from a TradingAction. - /// Maps Buy -> Long100/Market/Normal, Sell -> Short100/Market/Normal, Hold -> Flat/Market/Normal. + /// Buy → Long+Full(+1.0), Sell → Short+Full(-1.0), Hold → Flat+Half(0). + /// + /// In the 4-branch encoding: Long+Full = Long100(=8), Short+Full = Short50(=2), + /// Flat+Half = Flat(=4). Legacy variant names don't match 4-branch semantics. pub fn from_trading_action(action: crate::trading_action::TradingAction) -> Self { use crate::trading_action::TradingAction; match action { TradingAction::Buy => { - Self::new(ExposureLevel::Long100, OrderType::Market, Urgency::Normal) + Self::new(ExposureLevel::Long100, OrderType::Market, Urgency::Normal) // Long+Full = +1.0 } TradingAction::Sell => { - Self::new(ExposureLevel::Short100, OrderType::Market, Urgency::Normal) + Self::new(ExposureLevel::Short50, OrderType::Market, Urgency::Normal) // Short+Full = -1.0 } TradingAction::Hold => { - Self::new(ExposureLevel::Flat, OrderType::Market, Urgency::Normal) + Self::new(ExposureLevel::Flat, OrderType::Market, Urgency::Normal) // Flat+Half = 0 } } } @@ -292,25 +310,28 @@ impl FactoredAction { trade_value * self.transaction_cost() } - /// Check if this action is a buy (long exposure) + /// Check if this action is a buy (long direction, dir=2, indices 6-8) pub fn is_buy(&self) -> bool { matches!( self.exposure, - ExposureLevel::Long100 | ExposureLevel::Long50 + ExposureLevel::Long50 | ExposureLevel::Long75 | ExposureLevel::Long100 ) } - /// Check if this action is a sell (short exposure) + /// Check if this action is a sell (short direction, dir=0, indices 0-2) pub fn is_sell(&self) -> bool { matches!( self.exposure, - ExposureLevel::Short100 | ExposureLevel::Short75 | ExposureLevel::Short50 | ExposureLevel::Short25 + ExposureLevel::Short100 | ExposureLevel::Short75 | ExposureLevel::Short50 ) } - /// Check if this action is neutral (flat exposure) + /// Check if this action is neutral (flat direction, dir=1, indices 3-5) pub fn is_hold(&self) -> bool { - matches!(self.exposure, ExposureLevel::Flat) + matches!( + self.exposure, + ExposureLevel::Short25 | ExposureLevel::Flat | ExposureLevel::Long25 + ) } /// Convert action to position delta @@ -411,11 +432,12 @@ mod tests { #[test] fn test_target_exposure() { - assert_eq!(ExposureLevel::Short100.target_exposure(), -1.0); - assert_eq!(ExposureLevel::Short50.target_exposure(), -0.5); - assert_eq!(ExposureLevel::Flat.target_exposure(), 0.0); - assert_eq!(ExposureLevel::Long50.target_exposure(), 0.5); - assert_eq!(ExposureLevel::Long100.target_exposure(), 1.0); + // 4-branch: dir × mag formula + assert_eq!(ExposureLevel::Short100.target_exposure(), -0.25); // Short×Small + assert_eq!(ExposureLevel::Short50.target_exposure(), -1.0); // Short×Full + assert_eq!(ExposureLevel::Flat.target_exposure(), 0.0); // Flat×Half + assert_eq!(ExposureLevel::Long50.target_exposure(), 0.25); // Long×Small + assert_eq!(ExposureLevel::Long100.target_exposure(), 1.0); // Long×Full } #[test] @@ -465,6 +487,6 @@ mod tests { let action = FactoredAction::new(ExposureLevel::Long50, OrderType::LimitMaker, Urgency::Patient); let s = format!("{}", action); - assert_eq!(s, "Long50+LimitMaker+Patient"); + assert_eq!(s, "L25+LimitMaker+Patient"); // Long50 displays as L25 (Long×Small) } } diff --git a/crates/ml-dqn/src/dqn.rs b/crates/ml-dqn/src/dqn.rs index 98be42156..d6bad125d 100644 --- a/crates/ml-dqn/src/dqn.rs +++ b/crates/ml-dqn/src/dqn.rs @@ -1520,8 +1520,11 @@ impl DQN { let urg_random = in_warmup || rng.gen::() < effective_epsilon; if exp_random && ord_random && urg_random { - // All branches random — skip forward pass - let exposure = ExposureLevel::from_index(rng.gen_range(0..5_usize))?; + // All branches random — skip forward pass. + // 4-branch: exposure = dir*3 + mag (9 combos), order(3), urgency(3). + let dir = rng.gen_range(0..3_usize); + let mag = rng.gen_range(0..3_usize); + let exposure = ExposureLevel::from_index(dir * 3 + mag)?; let order = OrderType::from_index(rng.gen_range(0..3_usize))?; let urgency = Urgency::from_index(rng.gen_range(0..3_usize))?; FactoredAction { exposure, order, urgency } @@ -1536,20 +1539,26 @@ impl DQN { let output = branching_net.forward_branches_eval(&state_tensor)?; let greedy = super::branching::BranchingDuelingQNetwork::greedy_branch_actions(&output, &self.stream)?; + // 4-branch: greedy[0]=dir(0-2), greedy[1]=mag(0-2), composite = dir*3+mag let exposure = if exp_random { - ExposureLevel::from_index(rng.gen_range(0..5_usize))? + let dir = rng.gen_range(0..3_usize); + let mag = rng.gen_range(0..3_usize); + ExposureLevel::from_index(dir * 3 + mag)? } else { - ExposureLevel::from_index(greedy.first().copied().unwrap_or(2) as usize)? + let dir = greedy.first().copied().unwrap_or(1) as usize; + let mag = greedy.get(1).copied().unwrap_or(1) as usize; + ExposureLevel::from_index(dir * 3 + mag)? }; + // 4-branch greedy: [dir, mag, order, urgency] → order at [2], urgency at [3] let order = if ord_random { OrderType::from_index(rng.gen_range(0..3_usize))? } else { - OrderType::from_index(greedy.get(1).copied().unwrap_or(0) as usize)? + OrderType::from_index(greedy.get(2).copied().unwrap_or(0) as usize)? }; let urgency = if urg_random { Urgency::from_index(rng.gen_range(0..3_usize))? } else { - Urgency::from_index(greedy.get(2).copied().unwrap_or(1) as usize)? + Urgency::from_index(greedy.get(3).copied().unwrap_or(1) as usize)? }; FactoredAction { exposure, order, urgency } } diff --git a/crates/ml/src/cuda_pipeline/backward_kernels.cu b/crates/ml/src/cuda_pipeline/backward_kernels.cu index 9681a91f0..08456577f 100644 --- a/crates/ml/src/cuda_pipeline/backward_kernels.cu +++ b/crates/ml/src/cuda_pipeline/backward_kernels.cu @@ -68,5 +68,9 @@ extern "C" __global__ void bias_grad_reduce_kernel( for (int b = 0; b < batch_size; b++) { sum += (float)dy[b * out_dim + j]; } + /* No 1/N here: the upstream dY is already mean-reduced (1/N applied in + * c51_grad_kernel / mse_grad_kernel). Chain rule propagates it through + * cuBLAS dX = dY @ W, so dY at every layer is already 1/N-scaled. + * Dividing again would be 1/N². */ atomicAdd(&db[j], sum); } diff --git a/crates/ml/src/cuda_pipeline/c51_grad_kernel.cu b/crates/ml/src/cuda_pipeline/c51_grad_kernel.cu index d8b20abe5..ec4fa71b9 100644 --- a/crates/ml/src/cuda_pipeline/c51_grad_kernel.cu +++ b/crates/ml/src/cuda_pipeline/c51_grad_kernel.cu @@ -38,8 +38,14 @@ extern "C" __global__ void c51_grad_kernel( float proj = (float)projected[tid]; /* Cross-entropy gradient: d/d_logits(-Sigma proj * lp) = exp(lp) - proj - * Float exp() handles full range — no bf16 overflow */ - float d_combined = isw * (expf(lp) - proj); + * Float exp() handles full range — no bf16 overflow. + * + * MEAN-reduce: divide by batch_size so gradient scale is invariant to + * batch size. Without this, batch=16384 (H100) produces a 282× larger + * raw gradient sum than batch=58 (smoke test), causing the budget clip + * to destroy gradient SNR on large batches → epoch 2-3 collapse. */ + float inv_batch = 1.0f / (float)batch_size; + float d_combined = inv_batch * isw * (expf(lp) - proj); /* Entropy regularization — magnitude branch (d==1) gets 5× boost to prevent * atom distribution collapse. Standardized advantages prevent Q-value scale diff --git a/crates/ml/src/cuda_pipeline/cql_grad_kernel.cu b/crates/ml/src/cuda_pipeline/cql_grad_kernel.cu index 6160830d3..23eb83a82 100644 --- a/crates/ml/src/cuda_pipeline/cql_grad_kernel.cu +++ b/crates/ml/src/cuda_pipeline/cql_grad_kernel.cu @@ -28,6 +28,11 @@ extern "C" __global__ void cql_logit_grad_kernel( int i = blockIdx.x * blockDim.x + threadIdx.x; if (i >= N) return; + /* MEAN-reduce: scale all CQL gradient outputs by 1/N so downstream cuBLAS + * backward produces mean-reduced weight gradients. Keeps gradient scale + * invariant to batch size — consistent with C51/MSE/IQN kernels. */ + float inv_batch = 1.0f / (float)N; + int total_actions = b0_size + b1_size + b2_size + b3_size; float dz = (num_atoms > 1) ? (v_max - v_min) / (float)(num_atoms - 1) : 0.0f; @@ -181,7 +186,7 @@ extern "C" __global__ void cql_logit_grad_kernel( a_mean_j /= (float)bd; float p = expf(val[j] + adv[j] - a_mean_j - log_sum); float z = v_min + (float)j * dz; - float d_combined = d_cql_dq[a] * p * (z - eq); + float d_combined = inv_batch * d_cql_dq[a] * p * (z - eq); d_adv[j] = d_combined; if (j < 256) d_val_accum[j] += d_combined; diff --git a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs index 4452861fa..49f1c0ae8 100644 --- a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs +++ b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs @@ -2108,6 +2108,34 @@ impl GpuDqnTrainer { Ok(norm_bf16[0].to_f32()) } + /// Compute and return the current grad_buf L2 norm using the STANDALONE + /// kernel handle (safe after mega-graph capture). + /// + /// `read_grad_norm_sync()` uses the graph-captured kernel handle which + /// corrupts after `graph_mega` capture. This variant uses + /// `grad_norm_standalone` — a separately loaded handle of the same kernel. + /// + /// Costs one stream sync + 4-byte DtoH. For diagnostic use only. + pub fn read_grad_norm_standalone_sync(&mut self) -> Result { + let _evt_guard = EventTrackingGuard::new(self.stream.context()); + + self.launch_grad_norm_standalone()?; + self.launch_grad_norm_finalize()?; + + unsafe { cudarc::driver::sys::cuStreamSynchronize(self.stream.cu_stream()); } + + let mut norm_f32 = [0.0_f32; 1]; + unsafe { + cudarc::driver::sys::cuMemcpyDtoH_v2( + norm_f32.as_mut_ptr().cast(), + self.grad_norm_f32_buf.raw_ptr(), + std::mem::size_of::(), + ); + } + // grad_norm_f32_buf holds sum-of-squares; return L2 norm + Ok(norm_f32[0].sqrt()) + } + /// Apply GPU multi-head feature attention to `save_h_s2` (post-graph). /// /// Runs 4-head self-attention over the trunk output `h_s2 [B, SHARED_H2]` diff --git a/crates/ml/src/cuda_pipeline/iqn_dual_head_kernel.cu b/crates/ml/src/cuda_pipeline/iqn_dual_head_kernel.cu index c340a5a67..5f84865f5 100644 --- a/crates/ml/src/cuda_pipeline/iqn_dual_head_kernel.cu +++ b/crates/ml/src/cuda_pipeline/iqn_dual_head_kernel.cu @@ -562,6 +562,11 @@ void iqn_backward_kernel( /* Process each quantile τ_i */ float kappa = IQN_KAPPA; float inv_n_sq = 1.0f / ((float)IQN_NUM_QUANTILES * (float)IQN_NUM_QUANTILES); + /* MEAN-reduce: all atomicAdd contributions are scaled by 1/batch_size so + * gradient magnitude is invariant to batch size. Without this, batch=16384 + * produces ~282× larger raw gradients than batch=58, causing the budget + * clip system to destroy gradient SNR on H100. */ + float inv_batch = 1.0f / (float)batch_size; for (int ti = 0; ti < IQN_NUM_QUANTILES; ti++) { float tau_i = my_taus[ti]; @@ -590,33 +595,33 @@ void iqn_backward_kernel( for (int h = tid; h < hidden_dim; h += IQN_BLOCK_SIZE) dL_dcomb[h / IQN_BLOCK_SIZE] = 0.0f; - /* Branch 0: accumulate gradient */ + /* Branch 0: accumulate gradient (mean-reduced by inv_batch) */ for (int h = tid; h < hidden_dim; h += IQN_BLOCK_SIZE) { dL_dcomb[h / IQN_BLOCK_SIZE] += w_b0[a0 * hidden_dim + h] * dL_dq; - /* dL/dW_b0[a0, h] += dL/dq × combined[h] */ + /* dL/dW_b0[a0, h] += inv_batch × dL/dq × combined[h] */ atomicAdd(&grad_buf[off[2] + a0 * hidden_dim + h], - dL_dq * comb_dist[h / IQN_BLOCK_SIZE]); + inv_batch * dL_dq * comb_dist[h / IQN_BLOCK_SIZE]); } if (tid == 0) - atomicAdd(&grad_buf[off[3] + a0], dL_dq); + atomicAdd(&grad_buf[off[3] + a0], inv_batch * dL_dq); /* Branch 1 */ for (int h = tid; h < hidden_dim; h += IQN_BLOCK_SIZE) { dL_dcomb[h / IQN_BLOCK_SIZE] += w_b1[a1 * hidden_dim + h] * dL_dq; atomicAdd(&grad_buf[off[4] + a1 * hidden_dim + h], - dL_dq * comb_dist[h / IQN_BLOCK_SIZE]); + inv_batch * dL_dq * comb_dist[h / IQN_BLOCK_SIZE]); } if (tid == 0) - atomicAdd(&grad_buf[off[5] + a1], dL_dq); + atomicAdd(&grad_buf[off[5] + a1], inv_batch * dL_dq); /* Branch 2 */ for (int h = tid; h < hidden_dim; h += IQN_BLOCK_SIZE) { dL_dcomb[h / IQN_BLOCK_SIZE] += w_b2[a2 * hidden_dim + h] * dL_dq; atomicAdd(&grad_buf[off[6] + a2 * hidden_dim + h], - dL_dq * comb_dist[h / IQN_BLOCK_SIZE]); + inv_batch * dL_dq * comb_dist[h / IQN_BLOCK_SIZE]); } if (tid == 0) - atomicAdd(&grad_buf[off[7] + a2], dL_dq); + atomicAdd(&grad_buf[off[7] + a2], inv_batch * dL_dq); /* ── Gradient through element-wise product ── */ /* combined = h_s2 ⊙ embed @@ -625,9 +630,9 @@ void iqn_backward_kernel( float dL_dembed[IQN_DIST_MAX]; for (int h = tid; h < hidden_dim; h += IQN_BLOCK_SIZE) { dL_dembed[h / IQN_BLOCK_SIZE] = dL_dcomb[h / IQN_BLOCK_SIZE] * h_dist[h / IQN_BLOCK_SIZE]; - /* Accumulate dL/d(h_s2) across all quantiles. + /* Accumulate dL/d(h_s2) across all quantiles (mean-reduced). * This flows IQN's bounded Huber gradient to the shared trunk. */ - float d_trunk = dL_dcomb[h / IQN_BLOCK_SIZE] * embed_dist[h / IQN_BLOCK_SIZE]; + float d_trunk = inv_batch * dL_dcomb[h / IQN_BLOCK_SIZE] * embed_dist[h / IQN_BLOCK_SIZE]; if (isfinite(d_trunk)) atomicAdd(&d_h_s2_out[sample * hidden_dim + h], d_trunk); } @@ -647,12 +652,12 @@ void iqn_backward_kernel( for (int h = tid; h < hidden_dim; h += IQN_BLOCK_SIZE) { float dL_dpre = dL_dembed[h / IQN_BLOCK_SIZE]; if (!isfinite(dL_dpre)) continue; - /* Bias gradient */ - atomicAdd(&grad_buf[off[1] + h], dL_dpre); + /* Bias gradient (mean-reduced) */ + atomicAdd(&grad_buf[off[1] + h], inv_batch * dL_dpre); /* Weight gradient: outer product with precomputed cosine features */ for (int d = 0; d < embed_dim; d++) { atomicAdd(&grad_buf[off[0] + h * embed_dim + d], - dL_dpre * my_cos[d]); + inv_batch * dL_dpre * my_cos[d]); } } } diff --git a/crates/ml/src/cuda_pipeline/mse_grad_kernel.cu b/crates/ml/src/cuda_pipeline/mse_grad_kernel.cu index d60d41405..585e84dd9 100644 --- a/crates/ml/src/cuda_pipeline/mse_grad_kernel.cu +++ b/crates/ml/src/cuda_pipeline/mse_grad_kernel.cu @@ -42,8 +42,13 @@ extern "C" __global__ void mse_grad_kernel( float delta_z = (num_atoms > 1) ? (v_max - v_min) / (float)(num_atoms - 1) : 0.0f; float z_j = v_min + (float)j * delta_z; - /* Gradient of MSE loss through softmax expectation (float arithmetic) */ - float d_combined = isw * td_error * p_j * (z_j - e_q); + /* Gradient of MSE loss through softmax expectation (float arithmetic). + * + * MEAN-reduce: divide by batch_size so gradient scale is invariant to + * batch size. Without this, the budget clip is 282× more aggressive on + * H100 (batch=16384) than on smoke test (batch=58). */ + float inv_batch = 1.0f / (float)batch_size; + float d_combined = inv_batch * isw * td_error * p_j * (z_j - e_q); /* Magnitude entropy boost: prevent atom distribution collapse for d==1. * MSE path has no entropy param — hardcoded 0.005 = 5× base (0.001). */ diff --git a/crates/ml/src/trainers/dqn/config.rs b/crates/ml/src/trainers/dqn/config.rs index 2c618e800..fbd9b2ff8 100644 --- a/crates/ml/src/trainers/dqn/config.rs +++ b/crates/ml/src/trainers/dqn/config.rs @@ -1589,7 +1589,7 @@ impl DQNHyperparameters { time_decay_rate: 0.0005, q_gap_threshold: 0.05, // Tier 2 default: mild conviction gating (hyperopt searches [0.0, 0.5]) huber_delta: 100.0, // BUG #12 FIX: Scale delta 100x for gradient explosion fix (was 1.0) - gradient_clip_norm: Some(10.0), // C51 101-atoms × 3 branches → gradient naturally 300x larger than DQN + gradient_clip_norm: Some(1.0), // Safety net — mean-reduced gradients have norm ≈ 1.0 hold_penalty_weight: 0.01, // Default: 1% penalty weight movement_threshold: 0.02, // Default: 2% price movement threshold preprocessing_window: 50, // Default: 50-bar rolling window diff --git a/crates/ml/src/trainers/dqn/fused_training.rs b/crates/ml/src/trainers/dqn/fused_training.rs index 026246780..b9cd644ce 100644 --- a/crates/ml/src/trainers/dqn/fused_training.rs +++ b/crates/ml/src/trainers/dqn/fused_training.rs @@ -217,6 +217,10 @@ pub(crate) struct FusedTrainingCtx { pub(crate) hindsight_fraction: f64, /// v8: Curriculum learning enabled (false by default). pub(crate) curriculum_enabled: bool, + /// When true, disables mega-graph/graph_aux capture and logs per-stage + /// gradient norms (raw → post-primary-clip → post-IQN → post-CQL). + /// Costs ~4 stream syncs per step. For debugging H100 gradient collapse. + pub(crate) gradient_stage_diagnostics: bool, } impl Drop for FusedTrainingCtx { @@ -610,6 +614,7 @@ impl FusedTrainingCtx { popart_enabled: hyperparams.popart_enabled, hindsight_fraction: hyperparams.hindsight_fraction, curriculum_enabled: hyperparams.curriculum_enabled, + gradient_stage_diagnostics: std::env::var("FOXHUNT_GRAD_DIAG").is_ok(), }) } @@ -852,7 +857,8 @@ impl FusedTrainingCtx { // Capture mega-graph at step 2 (all sub-trainers initialized). // Mega-graph fuses spectral + forward + backward + aux → 1 launch. - if self.graph_mega.is_none() && self.steps_since_varmap_sync == 2 { + // Skip when gradient_stage_diagnostics is on — need ungraphed path for per-stage reads. + if self.graph_mega.is_none() && self.steps_since_varmap_sync == 2 && !self.gradient_stage_diagnostics { if let Err(e) = self.capture_graph_mega(agent, gpu_batch) { tracing::warn!( "graph_mega capture failed (non-fatal, falling back to individual graphs): {e}" @@ -916,6 +922,19 @@ impl FusedTrainingCtx { // gradient even after C51 warmup. The alpha ramp prevents MSE suppression // during warmup, but MUST NOT drop below 30% post-warmup. { + // ── DIAG: raw C51+MSE gradient norm (before any clip) ───────── + if self.gradient_stage_diagnostics { + let raw_norm = self.trainer.read_grad_norm_standalone_sync() + .unwrap_or(f32::NAN); + self.last_c51_raw_norm = raw_norm; + tracing::warn!( + batch_size = self.batch_size, + raw_norm, + stage = "pre_primary_clip", + "GRAD_DIAG: raw C51+MSE gradient norm BEFORE budget clip" + ); + } + let alpha = self.trainer.c51_alpha(); let cql_frac = if self.trainer.has_cql() { CQL_GRAD_BUDGET } else { 0.0 }; let iqn_frac = if self.gpu_iqn.is_some() { IQN_GRAD_BUDGET } else { 0.0 }; @@ -925,6 +944,20 @@ impl FusedTrainingCtx { let primary_budget = mgn * (1.0 - alpha) + mgn * c51_frac * alpha; self.trainer.clip_grad_buf_inplace(primary_budget) .map_err(|e| anyhow::anyhow!("Gradient budget clip: {e}"))?; + + // ── DIAG: post-primary-clip gradient norm ───────────────────── + if self.gradient_stage_diagnostics { + let post_clip_norm = self.trainer.read_grad_norm_standalone_sync() + .unwrap_or(f32::NAN); + tracing::warn!( + batch_size = self.batch_size, + post_clip_norm, + primary_budget, + c51_alpha = alpha, + stage = "post_primary_clip", + "GRAD_DIAG: gradient norm AFTER primary budget clip" + ); + } } // EMA target update. @@ -991,6 +1024,18 @@ impl FusedTrainingCtx { ); iqn.target_ema_update(tau as f32) .map_err(|e| anyhow::anyhow!("IQN EMA update: {e}"))?; + + // ── DIAG: post-IQN trunk SAXPY gradient norm ────────── + if self.gradient_stage_diagnostics { + let post_iqn_norm = self.trainer.read_grad_norm_standalone_sync() + .unwrap_or(f32::NAN); + tracing::warn!( + batch_size = self.batch_size, + post_iqn_norm, + stage = "post_iqn_saxpy", + "GRAD_DIAG: gradient norm AFTER IQN trunk clipped SAXPY" + ); + } } Err(e) => { tracing::warn!("IQN step failed (non-fatal): {e}"); @@ -1021,6 +1066,22 @@ impl FusedTrainingCtx { } } + // ── DIAG: final gradient norm (after all component injections) ──── + if self.gradient_stage_diagnostics { + let final_norm = self.trainer.read_grad_norm_standalone_sync() + .unwrap_or(f32::NAN); + tracing::warn!( + batch_size = self.batch_size, + final_norm, + raw_norm = self.last_c51_raw_norm, + clip_ratio = if self.last_c51_raw_norm > 0.0 { + final_norm / self.last_c51_raw_norm + } else { 0.0 }, + stage = "final_pre_adam", + "GRAD_DIAG: final gradient norm before Adam (all components injected)" + ); + } + // Regime-adaptive PER scaling. self.trainer.regime_scale_td_errors() .map_err(|e| anyhow::anyhow!("Regime PER scaling: {e}"))?;