From be8f3310d51d7df2b3f81c230ec4aaac3a2a3240 Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Fri, 27 Mar 2026 19:52:44 +0100 Subject: [PATCH] feat: gradient stability + 9-agent audit bug fixes MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Per-component gradient clipping: - 2 new CUDA kernels (dqn_clipped_saxpy, dqn_clip_grad) - CQL gradient isolated into separate scratch buffer - Budget allocation: C51=70%, CQL=15%, IQN=10%, Ens=5% - Dynamic budget — inactive components' share goes to C51 Spectral norm extended to all 10 weight matrices: - Was trunk-only (W_s1, W_s2), now covers all heads - 16 u/v power iteration buffers, batched GOFF sync-back - Uniform sigma_max, end-to-end Lipschitz bounded Bug fixes from 9-agent audit: - Backtest episode reset: all 8 fields (was 5), max_equity updated before floor check - gradient_clip_norm unified: 10.0 everywhere (was 10.0 vs 1.0) - entropy_coefficient: Option → f64, single default 0.001 - Close price: unwrap_or(1.0) → direct indexing (no silent wrong rewards) - step_count: only increments during training (was incrementing on inference) - Quantile loss: single CUDA context (was 3×), silent fallbacks removed - MaybeNoisyLinear: single-variant enum removed → direct NoisyLinear - Budget fractions: exported as pub(crate) constants, referenced not hardcoded Monitoring: - Per-component gradient Prometheus gauges (C51 raw, combined) - Diagnostic logging every 1000 steps Tests: 7 new gradient budget tests, 1 gradient bounds smoke test All 488 tests pass (359 ml-dqn + 129 ml) Co-Authored-By: Claude Opus 4.6 (1M context) --- crates/common/src/metrics/training_metrics.rs | 18 + crates/ml-dqn/src/agent.rs | 10 +- crates/ml-dqn/src/branching.rs | 147 ++-- crates/ml-dqn/src/ensemble_network.rs | 2 +- crates/ml-dqn/src/network.rs | 28 +- crates/ml-dqn/src/quantile_regression.rs | 39 +- crates/ml/examples/train_baseline_rl.rs | 2 +- .../src/cuda_pipeline/backtest_env_kernel.cu | 36 +- .../src/cuda_pipeline/dqn_utility_kernels.cu | 59 ++ .../src/cuda_pipeline/gpu_action_selector.rs | 6 + .../ml/src/cuda_pipeline/gpu_dqn_trainer.rs | 413 ++++++++-- crates/ml/src/hyperopt/adapters/dqn.rs | 2 +- crates/ml/src/trainers/dqn/config.rs | 6 +- crates/ml/src/trainers/dqn/data_loading.rs | 4 +- crates/ml/src/trainers/dqn/fused_training.rs | 86 +- crates/ml/src/trainers/dqn/mod.rs | 2 +- .../dqn/smoke_tests/gradient_budget.rs | 366 +++++++++ crates/ml/src/trainers/dqn/smoke_tests/mod.rs | 2 + .../dqn/smoke_tests/training_stability.rs | 38 + crates/ml/src/trainers/dqn/trainer/action.rs | 2 +- .../src/trainers/dqn/trainer/constructor.rs | 4 +- crates/ml/src/trainers/dqn/trainer/state.rs | 2 +- .../src/trainers/dqn/trainer/training_loop.rs | 5 + crates/ml/src/training_profile.rs | 12 +- .../superpowers/plans/2026-03-27-bug-fixes.md | 501 ++++++++++++ .../2026-03-27-gpu-hot-path-performance.md | 732 +++++++++++++++++ .../plans/2026-03-27-gradient-stability.md | 769 ++++++++++++++++++ .../2026-03-27-bf16-performance-design.md | 522 ++++++++++++ .../specs/2026-03-27-bug-fixes-design.md | 703 ++++++++++++++++ .../specs/2026-03-27-code-hygiene-design.md | 490 +++++++++++ 30 files changed, 4765 insertions(+), 243 deletions(-) create mode 100644 crates/ml/src/trainers/dqn/smoke_tests/gradient_budget.rs create mode 100644 docs/superpowers/plans/2026-03-27-bug-fixes.md create mode 100644 docs/superpowers/plans/2026-03-27-gpu-hot-path-performance.md create mode 100644 docs/superpowers/plans/2026-03-27-gradient-stability.md create mode 100644 docs/superpowers/specs/2026-03-27-bf16-performance-design.md create mode 100644 docs/superpowers/specs/2026-03-27-bug-fixes-design.md create mode 100644 docs/superpowers/specs/2026-03-27-code-hygiene-design.md diff --git a/crates/common/src/metrics/training_metrics.rs b/crates/common/src/metrics/training_metrics.rs index bd1a4d95e..6ac0b67a9 100644 --- a/crates/common/src/metrics/training_metrics.rs +++ b/crates/common/src/metrics/training_metrics.rs @@ -168,6 +168,16 @@ pub fn init() { "L2 gradient norm after backward", mf, ); + _ = register_gauge_vec( + "foxhunt_training_grad_norm_c51", + "Per-component gradient L2 norm: C51 primary loss (before budget clip)", + mf, + ); + _ = register_gauge_vec( + "foxhunt_training_grad_norm_combined", + "Combined gradient L2 norm after all injections (should ≤ max_grad_norm)", + mf, + ); _ = register_gauge_vec("foxhunt_training_learning_rate", "Effective learning rate", mf); _ = register_gauge_vec( "foxhunt_training_epoch_duration_seconds", @@ -505,6 +515,14 @@ pub fn set_gradient_norm(model: &str, fold: &str, norm: f64) { set_gauge_vec("foxhunt_training_gradient_norm", &[model, fold], norm); } +pub fn set_grad_norm_c51(model: &str, fold: &str, norm: f64) { + set_gauge_vec("foxhunt_training_grad_norm_c51", &[model, fold], norm); +} + +pub fn set_grad_norm_combined(model: &str, fold: &str, norm: f64) { + set_gauge_vec("foxhunt_training_grad_norm_combined", &[model, fold], norm); +} + pub fn set_learning_rate(model: &str, fold: &str, lr: f64) { set_gauge_vec("foxhunt_training_learning_rate", &[model, fold], lr); } diff --git a/crates/ml-dqn/src/agent.rs b/crates/ml-dqn/src/agent.rs index 8f8082220..a75469a0d 100644 --- a/crates/ml-dqn/src/agent.rs +++ b/crates/ml-dqn/src/agent.rs @@ -304,12 +304,12 @@ impl DQNAgent { let reward = rewards.get(i).copied().unwrap_or(0.0); let done = dones.get(i).copied().unwrap_or(false); - // Forward current state - let q_values = self.q_network.forward(state)?; + // Forward current state (training path: loss computation) + let q_values = self.q_network.forward(state, true)?; let q_taken = q_values.get(action).copied().unwrap_or(0.0); - // Forward next state through target network - let next_q = self.target_network.forward(next_state)?; + // Forward next state through target network (training path: TD target) + let next_q = self.target_network.forward(next_state, true)?; let max_next_q = next_q.iter().copied().fold(f32::NEG_INFINITY, f32::max); // TD target @@ -684,7 +684,7 @@ impl DQNAgent { ) -> Result, MLError> { // Get raw Q-values from network (returns Vec) let state_vec = state.to_vector(); - let q_values = self.q_network.forward(&state_vec)?; + let q_values = self.q_network.forward(&state_vec, false)?; // Extract expected price from state features (technical_indicators[1]) let expected_price = if state.technical_indicators.len() > 1 { diff --git a/crates/ml-dqn/src/branching.rs b/crates/ml-dqn/src/branching.rs index b3e0629db..e4ac88534 100644 --- a/crates/ml-dqn/src/branching.rs +++ b/crates/ml-dqn/src/branching.rs @@ -248,85 +248,34 @@ impl BranchingConfig { } } -/// A `NoisyLinear` wrapper for value/branch heads (noisy nets always enabled). -enum MaybeNoisyLinear { - Noisy(NoisyLinear), -} +/// Register `NoisyLinear` mu (weight/bias) params in `GpuVarStore` under `{name}.weight` / `{name}.bias`. +/// +/// Clones the mu data from the `NoisyLinear` into the `GpuVarStore` so that the +/// GPU experience collector and serialization can find them by name. +/// Uses async `DtoD` memcpy -- zero CPU involvement. +fn register_noisy_mu_in_varstore(n: &NoisyLinear, vars: &mut GpuVarStore, name: &str) -> Result<(), MLError> { + let [w_mu, b_mu] = n.mu_slices(); + let out_f = n.out_features(); + let in_f = n.in_features(); + let stream = vars.stream().clone(); -impl MaybeNoisyLinear { - /// Forward pass through the noisy layer. - fn forward(&self, x: &GpuTensor) -> Result { - let Self::Noisy(n) = self; - n.forward(x) - } + // Clone weight_mu -> "{name}.weight" via DtoD memcpy + let w_len = w_mu.len(); + let mut w_clone = stream.alloc_zeros::(w_len).map_err(|e| { + MLError::ModelError(format!("register_mu alloc {name}.weight: {e}")) + })?; + dtod_copy_slice(w_mu, &mut w_clone, &stream, &format!("register_mu {name}.weight"))?; + vars.register(format!("{name}.weight"), w_clone, vec![out_f, in_f])?; - /// Resample noise. - fn reset_noise(&mut self) -> Result<(), MLError> { - let Self::Noisy(n) = self; - n.reset_noise() - } + // Clone bias_mu -> "{name}.bias" via DtoD memcpy + let b_len = b_mu.len(); + let mut b_clone = stream.alloc_zeros::(b_len).map_err(|e| { + MLError::ModelError(format!("register_mu alloc {name}.bias: {e}")) + })?; + dtod_copy_slice(b_mu, &mut b_clone, &stream, &format!("register_mu {name}.bias"))?; + vars.register(format!("{name}.bias"), b_clone, vec![out_f])?; - /// Resample noise with custom sigma scale. - fn reset_noise_with_sigma(&mut self, sigma_scale: f64) -> Result<(), MLError> { - let Self::Noisy(n) = self; - n.reset_noise_with_sigma(sigma_scale) - } - - /// Disable noise for evaluation. - fn disable_noise(&mut self) -> Result<(), MLError> { - let Self::Noisy(n) = self; - n.disable_noise() - } - - /// Convert sigma vars and epsilon buffers to F32 to match mu vars. - #[allow(dead_code)] - const fn ensure_f32(&mut self) -> Result<(), MLError> { - let Self::Noisy(n) = self; - n.ensure_f32() - } - - /// Register mu (weight/bias) params in `GpuVarStore` under `{name}.weight` / `{name}.bias`. - /// - /// Clones the mu data from the `NoisyLinear` into the `GpuVarStore` so that the - /// GPU experience collector and serialization can find them by name. - /// Uses async `DtoD` memcpy -- zero CPU involvement. - fn register_mu_in_varstore(&self, vars: &mut GpuVarStore, name: &str) -> Result<(), MLError> { - let Self::Noisy(n) = self; - let [w_mu, b_mu] = n.mu_slices(); - let out_f = n.out_features(); - let in_f = n.in_features(); - let stream = vars.stream().clone(); - - // Clone weight_mu -> "{name}.weight" via DtoD memcpy - let w_len = w_mu.len(); - let mut w_clone = stream.alloc_zeros::(w_len).map_err(|e| { - MLError::ModelError(format!("register_mu alloc {name}.weight: {e}")) - })?; - dtod_copy_slice(w_mu, &mut w_clone, &stream, &format!("register_mu {name}.weight"))?; - vars.register(format!("{name}.weight"), w_clone, vec![out_f, in_f])?; - - // Clone bias_mu -> "{name}.bias" via DtoD memcpy - let b_len = b_mu.len(); - let mut b_clone = stream.alloc_zeros::(b_len).map_err(|e| { - MLError::ModelError(format!("register_mu alloc {name}.bias: {e}")) - })?; - dtod_copy_slice(b_mu, &mut b_clone, &stream, &format!("register_mu {name}.bias"))?; - vars.register(format!("{name}.bias"), b_clone, vec![out_f])?; - - Ok(()) - } - - /// Get only sigma (noise std dev) parameter slices -- `weight_sigma`, `bias_sigma`. - const fn noisy_sigma_slices(&self) -> [&CudaSlice; 2] { - let Self::Noisy(n) = self; - n.sigma_slices() - } -} - -impl std::fmt::Debug for MaybeNoisyLinear { - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - write!(f, "Noisy(NoisyLinear)") - } + Ok(()) } /// Branching Dueling Q-Network with independent advantage heads. @@ -342,12 +291,12 @@ pub struct BranchingDuelingQNetwork { shared_layers: Vec, /// Value stream: hidden -> scalar (or `num_atoms` when distributional) - value_fc: MaybeNoisyLinear, - value_out: MaybeNoisyLinear, + value_fc: NoisyLinear, + value_out: NoisyLinear, /// Per-branch advantage streams: D x (hidden -> `n_d` or `n_d`*`num_atoms`) - branch_fcs: Vec, - branch_outs: Vec, + branch_fcs: Vec, + branch_outs: Vec, /// Configuration config: BranchingConfig, @@ -418,13 +367,13 @@ impl BranchingDuelingQNetwork { } // Register NoisyLinear mu params in GpuVarStore for GPU weight extraction - value_fc.register_mu_in_varstore(&mut vars, "value_fc")?; - value_out.register_mu_in_varstore(&mut vars, "value_out")?; + register_noisy_mu_in_varstore(&value_fc, &mut vars, "value_fc")?; + register_noisy_mu_in_varstore(&value_out, &mut vars, "value_out")?; for (d, fc) in branch_fcs.iter().enumerate() { - fc.register_mu_in_varstore(&mut vars, &format!("branch_{d}_fc"))?; + register_noisy_mu_in_varstore(fc, &mut vars, &format!("branch_{d}_fc"))?; } for (d, out) in branch_outs.iter().enumerate() { - out.register_mu_in_varstore(&mut vars, &format!("branch_{d}_out"))?; + register_noisy_mu_in_varstore(out, &mut vars, &format!("branch_{d}_out"))?; } // Compute C51 support atoms @@ -455,9 +404,8 @@ impl BranchingDuelingQNetwork { stream: &Arc, _name: &str, sigma_init: f64, - ) -> Result { - let noisy = NoisyLinear::new(fan_in, fan_out, stream.clone(), sigma_init)?; - Ok(MaybeNoisyLinear::Noisy(noisy)) + ) -> Result { + NoisyLinear::new(fan_in, fan_out, stream.clone(), sigma_init) } /// Compute C51 support atoms: `linspace(v_min, v_max, num_atoms)` on device. @@ -1009,19 +957,19 @@ impl BranchingDuelingQNetwork { vars.push(param.data.clone()); } // Only sigma vars -- mu already in GpuVarStore - for s in self.value_fc.noisy_sigma_slices() { + for s in self.value_fc.sigma_slices() { vars.push(s.clone()); } - for s in self.value_out.noisy_sigma_slices() { + for s in self.value_out.sigma_slices() { vars.push(s.clone()); } for fc in &self.branch_fcs { - for s in fc.noisy_sigma_slices() { + for s in fc.sigma_slices() { vars.push(s.clone()); } } for out in &self.branch_outs { - for s in out.noisy_sigma_slices() { + for s in out.sigma_slices() { vars.push(s.clone()); } } @@ -1038,19 +986,19 @@ impl BranchingDuelingQNetwork { /// produce the same order, so vars can be zipped for Polyak update. pub fn noisy_vars_ordered(&self) -> Vec> { let mut vars = Vec::new(); - for s in self.value_fc.noisy_sigma_slices() { + for s in self.value_fc.sigma_slices() { vars.push(s.clone()); } - for s in self.value_out.noisy_sigma_slices() { + for s in self.value_out.sigma_slices() { vars.push(s.clone()); } for fc in &self.branch_fcs { - for s in fc.noisy_sigma_slices() { + for s in fc.sigma_slices() { vars.push(s.clone()); } } for out in &self.branch_outs { - for s in out.noisy_sigma_slices() { + for s in out.sigma_slices() { vars.push(s.clone()); } } @@ -1113,13 +1061,10 @@ impl BranchingDuelingQNetwork { /// Copy `NoisyLinear` weights between matching layers via `DtoD` memcpy. fn copy_noisy_layer( - dst: &mut MaybeNoisyLinear, - src: &MaybeNoisyLinear, + dst: &mut NoisyLinear, + src: &NoisyLinear, ) -> Result<(), MLError> { - let MaybeNoisyLinear::Noisy(src_noisy) = src; - #[allow(irrefutable_let_patterns)] - let MaybeNoisyLinear::Noisy(dst_noisy) = dst; - dst_noisy.copy_params_from(src_noisy) + dst.copy_params_from(src) } } diff --git a/crates/ml-dqn/src/ensemble_network.rs b/crates/ml-dqn/src/ensemble_network.rs index 4229fda0b..79d66b045 100644 --- a/crates/ml-dqn/src/ensemble_network.rs +++ b/crates/ml-dqn/src/ensemble_network.rs @@ -112,7 +112,7 @@ impl EnsembleQNetwork { let mut q_values = Vec::with_capacity(self.num_networks); for network in &self.networks { - let q_vals = network.forward(state)?; + let q_vals = network.forward(state, false)?; q_values.push(q_vals); } diff --git a/crates/ml-dqn/src/network.rs b/crates/ml-dqn/src/network.rs index 3113aa653..94d199df9 100644 --- a/crates/ml-dqn/src/network.rs +++ b/crates/ml-dqn/src/network.rs @@ -205,8 +205,14 @@ impl QNetwork { Ok(layers) } - /// Forward pass through the network - pub fn forward(&self, state: &[f32]) -> Result, MLError> { + /// Forward pass through the network. + /// + /// # Arguments + /// * `state` - Input state vector of length `state_dim` + /// * `training` - When `true`, increments `step_count` and advances the + /// dropout scheduler. Pass `false` for inference / action-selection / + /// backtest paths to avoid corrupting the training schedule. + pub fn forward(&self, state: &[f32], training: bool) -> Result, MLError> { if state.len() != self.config.state_dim { return Err(MLError::InvalidInput(format!( "State dimension mismatch: expected {}, got {}", @@ -246,13 +252,15 @@ impl QNetwork { // Download output to CPU let output_vec = x.to_host(&self.stream)?; - // Update step count - let _step = self.step_count.fetch_add(1, Ordering::Relaxed); + // Only update step count and dropout scheduler during training + if training { + let _step = self.step_count.fetch_add(1, Ordering::Relaxed); - // Update dropout scheduler (Wave 26 P1.6) - if let Ok(mut scheduler_opt) = self.dropout_scheduler.lock() { - if let Some(scheduler) = scheduler_opt.as_mut() { - scheduler.step(1); + // Update dropout scheduler (Wave 26 P1.6) + if let Ok(mut scheduler_opt) = self.dropout_scheduler.lock() { + if let Some(scheduler) = scheduler_opt.as_mut() { + scheduler.step(1); + } } } @@ -324,7 +332,7 @@ impl QNetwork { /// Select action using greedy policy (exploration handled by noisy networks) pub fn select_action(&self, state: &[f32]) -> Result { - let q_values = self.forward(state)?; + let q_values = self.forward(state, false)?; let best_action = q_values .iter() .enumerate() @@ -422,7 +430,7 @@ mod tests { let state = vec![1.0, 2.0, 3.0, 4.0]; let q_values = network - .forward(&state) + .forward(&state, false) .map_err(|e| anyhow::anyhow!("Forward pass failed: {:?}", e))?; assert_eq!(q_values.len(), 9); // 9 exposure levels (default) diff --git a/crates/ml-dqn/src/quantile_regression.rs b/crates/ml-dqn/src/quantile_regression.rs index 4a4daee6e..92985ef70 100644 --- a/crates/ml-dqn/src/quantile_regression.rs +++ b/crates/ml-dqn/src/quantile_regression.rs @@ -309,6 +309,14 @@ pub fn quantile_huber_loss( taus: &GpuTensor, kappa: f32, ) -> Result { + // Create CUDA context + stream ONCE, reuse for all operations. + let ctx = cudarc::driver::CudaContext::new(0).map_err(|e| { + MLError::DeviceError(format!("CUDA context: {e}")) + })?; + let stream = ctx.new_stream().map_err(|e| { + MLError::DeviceError(format!("CUDA stream: {e}")) + })?; + // Compute quantile Huber loss via host-side calculation (cold path). // predicted, target, taus: [batch, num_quantiles] // Returns: scalar mean loss @@ -317,33 +325,10 @@ pub fn quantile_huber_loss( // Mean over batch let batch = per_sample.numel(); if batch == 0 { - let ctx = cudarc::driver::CudaContext::new(0).map_err(|e| { - MLError::DeviceError(format!("CUDA context: {e}")) - })?; - let stream = ctx.new_stream().map_err(|e| { - MLError::DeviceError(format!("CUDA stream: {e}")) - })?; return GpuTensor::scalar(0.0, &stream); } - let host = per_sample.to_host( - // We need a stream -- extract from tensor data context - // The tensor's CudaSlice carries its context - &{ - let ctx = cudarc::driver::CudaContext::new(0).map_err(|e| { - MLError::DeviceError(format!("CUDA context: {e}")) - })?; - ctx.new_stream().map_err(|e| { - MLError::DeviceError(format!("CUDA stream: {e}")) - })? - }, - )?; + let host = per_sample.to_host(&stream)?; let mean_loss: f32 = host.iter().sum::() / batch as f32; - let ctx = cudarc::driver::CudaContext::new(0).map_err(|e| { - MLError::DeviceError(format!("CUDA context: {e}")) - })?; - let stream = ctx.new_stream().map_err(|e| { - MLError::DeviceError(format!("CUDA stream: {e}")) - })?; GpuTensor::scalar(mean_loss, &stream) } @@ -382,9 +367,9 @@ pub fn quantile_huber_loss_per_sample( let mut sample_loss = 0.0_f32; for q in 0..num_q { let idx = b * num_q + q; - let pred_val = pred_host.get(idx).copied().unwrap_or(0.0); - let tgt_val = tgt_host.get(idx).copied().unwrap_or(0.0); - let tau_val = tau_host.get(idx).copied().unwrap_or(0.5); + let pred_val = pred_host[idx]; + let tgt_val = tgt_host[idx]; + let tau_val = tau_host[idx]; let diff = tgt_val - pred_val; let abs_diff = diff.abs(); diff --git a/crates/ml/examples/train_baseline_rl.rs b/crates/ml/examples/train_baseline_rl.rs index a38f7bef3..db98f99c5 100644 --- a/crates/ml/examples/train_baseline_rl.rs +++ b/crates/ml/examples/train_baseline_rl.rs @@ -539,7 +539,7 @@ fn train_dqn_fold( hold_penalty_weight: hp_f64(hp, "hold_penalty_weight").unwrap_or(0.01), max_position_absolute: hp_f64(hp, "max_position_absolute").unwrap_or(2.0), huber_delta: hp_f64(hp, "huber_delta").unwrap_or(10.0), - entropy_coefficient: hp_f64(hp, "entropy_coefficient").map(Some).unwrap_or(Some(0.01)), + entropy_coefficient: hp_f64(hp, "entropy_coefficient").unwrap_or(0.01), curiosity_weight: hp_f64(hp, "curiosity_weight").unwrap_or(0.1), weight_decay: hp_f64(hp, "weight_decay").unwrap_or(1e-4), kelly_fractional: hp_f64(hp, "kelly_fractional").unwrap_or(0.5), diff --git a/crates/ml/src/cuda_pipeline/backtest_env_kernel.cu b/crates/ml/src/cuda_pipeline/backtest_env_kernel.cu index ed781a10d..36919dcb1 100644 --- a/crates/ml/src/cuda_pipeline/backtest_env_kernel.cu +++ b/crates/ml/src/cuda_pipeline/backtest_env_kernel.cu @@ -102,14 +102,14 @@ extern "C" __global__ void backtest_env_step( step_rewards[w] = liq_ret; step_returns[w * max_len + current_step] = liq_ret; actions_history[w * max_len + current_step] = b0_size / 2; // Flat - portfolio_state[ps + 0] = liq_value; - portfolio_state[ps + 1] = 0.0f; - portfolio_state[ps + 2] = liq_value; - portfolio_state[ps + 3] = 0.0f; - portfolio_state[ps + 4] = max_equity; - portfolio_state[ps + 5] = 0.0f; - portfolio_state[ps + 6] = cum_return + liq_ret; - portfolio_state[ps + 7] += 1.0f; + portfolio_state[ps + 0] = liq_value; // value = new initial capital + portfolio_state[ps + 1] = 0.0f; // position = flat + portfolio_state[ps + 2] = liq_value; // cash = new initial capital + portfolio_state[ps + 3] = 0.0f; // entry_price = none + portfolio_state[ps + 4] = liq_value; // max_equity = RESET (was stale) + portfolio_state[ps + 5] = 0.0f; // hold_time = 0 + portfolio_state[ps + 6] = 0.0f; // cum_return = RESET (was accumulating) + portfolio_state[ps + 7] = 0.0f; // step_count = RESET (was incrementing) done_flags[w] = 1; return; } @@ -175,6 +175,10 @@ extern "C" __global__ void backtest_env_step( // Mark-to-market current position (notional model: equity = cash + position * price) float new_value = cash + position * close; + // Update max_equity BEFORE floor check — prevents stale peak from + // missing breaches or triggering false ones after profitable trades. + max_equity = fmaxf(max_equity, new_value); + // ── Post-trade floor check: catch intra-step breaches ──────────── // The pre-trade check (top of kernel) catches breaches from the // PREVIOUS step. This catches breaches from THIS step's price move @@ -196,14 +200,14 @@ extern "C" __global__ void backtest_env_step( step_rewards[w] = step_ret; step_returns[w * max_len + current_step] = step_ret; actions_history[w * max_len + current_step] = b0_size / 2; - portfolio_state[ps + 0] = new_value; - portfolio_state[ps + 1] = 0.0f; - portfolio_state[ps + 2] = new_value; - portfolio_state[ps + 3] = 0.0f; - portfolio_state[ps + 4] = max_equity; - portfolio_state[ps + 5] = 0.0f; - portfolio_state[ps + 6] = cum_return + step_ret; - portfolio_state[ps + 7] += 1.0f; + portfolio_state[ps + 0] = new_value; // value = new initial capital + portfolio_state[ps + 1] = 0.0f; // position = flat + portfolio_state[ps + 2] = new_value; // cash = new initial capital + portfolio_state[ps + 3] = 0.0f; // entry_price = none + portfolio_state[ps + 4] = new_value; // max_equity = RESET + portfolio_state[ps + 5] = 0.0f; // hold_time = 0 + portfolio_state[ps + 6] = 0.0f; // cum_return = RESET + portfolio_state[ps + 7] = 0.0f; // step_count = RESET done_flags[w] = 1; return; } diff --git a/crates/ml/src/cuda_pipeline/dqn_utility_kernels.cu b/crates/ml/src/cuda_pipeline/dqn_utility_kernels.cu index 9835e1966..d6cd113a9 100644 --- a/crates/ml/src/cuda_pipeline/dqn_utility_kernels.cu +++ b/crates/ml/src/cuda_pipeline/dqn_utility_kernels.cu @@ -133,6 +133,65 @@ extern "C" __global__ void dqn_saxpy_kernel( if (i < n) y[i] += alpha * x[i]; } +/* ══════════════════════════════════════════════════════════════════════ + * CLIPPED SAXPY KERNEL + * + * y[i] += alpha * clip(x_norm, max_norm) * x[i] + * + * Reads the pre-computed L2 norm-squared of x from x_norm_sq[0], + * computes clip_scale = min(1, max_norm / ||x||), and applies: + * y[i] += alpha * clip_scale * x[i] + * + * This prevents any auxiliary gradient source (IQN, CQL, ensemble) + * from dominating the combined gradient buffer. Each component is + * independently clipped BEFORE accumulation. + * + * Launch config: grid=(ceil(n/256), 1, 1), block=(256, 1, 1). + * ══════════════════════════════════════════════════════════════════════ */ + +extern "C" __global__ void dqn_clipped_saxpy_kernel( + float* __restrict__ y, + const float* __restrict__ x, + float alpha, + float max_norm, + const float* __restrict__ x_norm_sq, + int n +) { + int i = blockIdx.x * blockDim.x + threadIdx.x; + if (i >= n) return; + + float norm = sqrtf(*x_norm_sq + 1e-12f); + float clip = (norm > max_norm) ? (max_norm / norm) : 1.0f; + y[i] += alpha * clip * x[i]; +} + +/* ══════════════════════════════════════════════════════════════════════ + * IN-PLACE GRADIENT CLIPPING KERNEL + * + * grads[i] *= min(1, max_norm / ||grads||) + * + * Reads the pre-computed L2 norm-squared from grad_norm_sq[0]. + * Used to clip the combined C51+CQL gradient BEFORE IQN/ensemble + * injection, preventing gradient interference cascade. + * + * Launch config: grid=(ceil(n/256), 1, 1), block=(256, 1, 1). + * ══════════════════════════════════════════════════════════════════════ */ + +extern "C" __global__ void dqn_clip_grad_kernel( + float* __restrict__ grads, + const float* __restrict__ grad_norm_sq, + float max_norm, + int total_params +) { + int idx = blockIdx.x * blockDim.x + threadIdx.x; + if (idx >= total_params) return; + + float norm = sqrtf(*grad_norm_sq + 1e-12f); + if (norm > max_norm) { + grads[idx] *= max_norm / norm; + } +} + /* ══════════════════════════════════════════════════════════════════════ * ZERO KERNEL * diff --git a/crates/ml/src/cuda_pipeline/gpu_action_selector.rs b/crates/ml/src/cuda_pipeline/gpu_action_selector.rs index 7da74b516..b358417ee 100644 --- a/crates/ml/src/cuda_pipeline/gpu_action_selector.rs +++ b/crates/ml/src/cuda_pipeline/gpu_action_selector.rs @@ -7,6 +7,12 @@ //! //! Pure cudarc implementation with no Candle dependency. //! Kernel loaded from build-time precompiled cubin (zero runtime nvcc). +//! +//! NoisyNet exploration: Q-values received by the action selector already +//! include factorized Gaussian noise from NoisyLinear layers in the network +//! forward pass. Epsilon-greedy acts as a secondary fallback, not the primary +//! exploration mechanism when NoisyNet is enabled. Explicit noise injection +//! here would double-count exploration. use cudarc::driver::{CudaFunction, CudaSlice, CudaStream, DevicePtr, LaunchConfig, PushKernelArg}; use cudarc::nvrtc::Ptx; diff --git a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs index 50f8cfd67..c4070c065 100644 --- a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs +++ b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs @@ -338,11 +338,22 @@ pub struct GpuDqnTrainer { shrink_perturb_kernel: CudaFunction, relu_mask_kernel: CudaFunction, spectral_norm_kernel: CudaFunction, + clipped_saxpy_kernel: CudaFunction, + clip_grad_kernel: CudaFunction, /// Spectral norm left singular vectors: [out_dim] per weight matrix (W_s1, W_s2) spec_u_s1: CudaSlice, spec_v_s1: CudaSlice, spec_u_s2: CudaSlice, spec_v_s2: CudaSlice, + // Head spectral norm vectors (power iteration) + spec_u_v1: CudaSlice, spec_v_v1: CudaSlice, + spec_u_v2: CudaSlice, spec_v_v2: CudaSlice, + spec_u_a1: CudaSlice, spec_v_a1: CudaSlice, + spec_u_a2: CudaSlice, spec_v_a2: CudaSlice, + spec_u_bo1: CudaSlice, spec_v_bo1: CudaSlice, + spec_u_bo2: CudaSlice, spec_v_bo2: CudaSlice, + spec_u_bu1: CudaSlice, spec_v_bu1: CudaSlice, + spec_u_bu2: CudaSlice, spec_v_bu2: CudaSlice, /// IQN trunk gradient Adam first moment [trunk_params] (separate from C51 Adam). iqn_trunk_m: CudaSlice, /// IQN trunk gradient Adam second moment [trunk_params]. @@ -404,6 +415,7 @@ pub struct GpuDqnTrainer { m_buf: CudaSlice, // [TOTAL_PARAMS] Adam first moment v_buf: CudaSlice, // [TOTAL_PARAMS] Adam second moment grad_norm_buf: CudaSlice, // [1] pre-clip gradient L2 norm + cql_grad_scratch: CudaSlice, // [TOTAL_PARAMS] CQL gradient isolation buffer // ── Adam step counter on device (CUDA Graph cannot bake scalars) ─ t_buf: CudaSlice, // [1] current Adam step @@ -649,6 +661,23 @@ impl GpuDqnTrainer { raw_device_ptr(&self.grad_buf, &self.stream) } + /// Mutable reference to grad_buf for test injection. + #[cfg(test)] + pub fn grad_buf_mut(&mut self) -> &mut CudaSlice { + &mut self.grad_buf + } + + /// Mutable reference to CQL gradient scratch for test injection. + #[cfg(test)] + pub fn cql_grad_scratch_mut(&mut self) -> &mut CudaSlice { + &mut self.cql_grad_scratch + } + + /// Total number of trainable parameters. + pub fn total_params(&self) -> usize { + self.total_params + } + /// Apply IQN auxiliary gradient to the shared trunk via SAXPY into `grad_buf`. /// /// After `graph_forward` replay (C51 forward → backward) and IQN backward @@ -796,28 +825,52 @@ impl GpuDqnTrainer { )?; } - // ── 7. SAXPY: grad_buf[trunk] += iqn_lambda * scratch[trunk] ───── - // Adds scaled IQN trunk gradient to C51's gradient already in grad_buf. - // graph_adam will then see the combined gradient and apply single Adam. + // ── 7. Clipped SAXPY: grad_buf[trunk] += iqn_lambda * clip(scratch) ── + // Compute IQN trunk gradient norm, then add with per-component clipping. + // Prevents IQN from overwhelming C51's gradient in grad_buf. { - let grad_ptr = raw_device_ptr(&self.grad_buf, &self.stream); + // Zero norm accumulator + self.stream.memset_zeros(&mut self.iqn_trunk_grad_norm) + .map_err(|e| MLError::ModelError(format!("zero iqn_trunk_grad_norm: {e}")))?; + + // Compute IQN trunk gradient norm (sum of squares) let scratch_ptr = raw_device_ptr(&self.iqn_trunk_m, &self.stream); - let scale = self.config.iqn_lambda; + let norm_ptr = raw_device_ptr(&self.iqn_trunk_grad_norm, &self.stream); let n_i32 = trunk_grad_total as i32; let blocks = ((trunk_grad_total + 255) / 256) as u32; unsafe { self.stream - .launch_builder(&self.saxpy_kernel) + .launch_builder(&self.grad_norm_kernel) + .arg(&scratch_ptr) + .arg(&norm_ptr) + .arg(&n_i32) + .launch(LaunchConfig { + grid_dim: (blocks, 1, 1), + block_dim: (256, 1, 1), + shared_mem_bytes: 256, // 8 warps * 2 stride * 4 bytes + }) + .map_err(|e| MLError::ModelError(format!("IQN trunk grad_norm: {e}")))?; + } + + // Clipped SAXPY: grad_buf += iqn_lambda * clip(scratch, iqn_budget) + let grad_ptr = raw_device_ptr(&self.grad_buf, &self.stream); + let max_component_norm = self.config.max_grad_norm * crate::trainers::dqn::fused_training::IQN_GRAD_BUDGET; + let scale = self.config.iqn_lambda; + unsafe { + self.stream + .launch_builder(&self.clipped_saxpy_kernel) .arg(&grad_ptr) .arg(&scratch_ptr) .arg(&scale) + .arg(&max_component_norm) + .arg(&norm_ptr) .arg(&n_i32) .launch(LaunchConfig { grid_dim: (blocks, 1, 1), block_dim: (256, 1, 1), shared_mem_bytes: 0, }) - .map_err(|e| MLError::ModelError(format!("IQN trunk SAXPY: {e}")))?; + .map_err(|e| MLError::ModelError(format!("IQN clipped SAXPY: {e}")))?; } } @@ -1014,25 +1067,50 @@ impl GpuDqnTrainer { )?; } - // ── 9. SAXPY: grad_buf[trunk] += scale * scratch[trunk] ──────────── + // ── 9. Clipped SAXPY: grad_buf[trunk] += scale * clip(scratch) ──── + // Per-component clipping prevents ensemble diversity from overwhelming + // the primary C51 gradient — same pattern as IQN trunk gradient. { - let grad_ptr = raw_device_ptr(&self.grad_buf, &self.stream); + // Compute ensemble trunk gradient norm + self.stream.memset_zeros(&mut self.iqn_trunk_grad_norm) + .map_err(|e| MLError::ModelError(format!("zero ens_trunk_grad_norm: {e}")))?; + let scratch_ptr = raw_device_ptr(&self.iqn_trunk_m, &self.stream); + let norm_ptr = raw_device_ptr(&self.iqn_trunk_grad_norm, &self.stream); let n_i32 = trunk_grad_total as i32; let blocks = ((trunk_grad_total + 255) / 256) as u32; unsafe { self.stream - .launch_builder(&self.saxpy_kernel) + .launch_builder(&self.grad_norm_kernel) + .arg(&scratch_ptr) + .arg(&norm_ptr) + .arg(&n_i32) + .launch(LaunchConfig { + grid_dim: (blocks, 1, 1), + block_dim: (256, 1, 1), + shared_mem_bytes: 256, + }) + .map_err(|e| MLError::ModelError(format!("ens trunk grad_norm: {e}")))?; + } + + // Clipped SAXPY + let grad_ptr = raw_device_ptr(&self.grad_buf, &self.stream); + let max_component_norm = self.config.max_grad_norm * crate::trainers::dqn::fused_training::ENS_GRAD_BUDGET; + unsafe { + self.stream + .launch_builder(&self.clipped_saxpy_kernel) .arg(&grad_ptr) .arg(&scratch_ptr) .arg(&scale) + .arg(&max_component_norm) + .arg(&norm_ptr) .arg(&n_i32) .launch(LaunchConfig { grid_dim: (blocks, 1, 1), block_dim: (256, 1, 1), shared_mem_bytes: 0, }) - .map_err(|e| MLError::ModelError(format!("ensemble trunk SAXPY: {e}")))?; + .map_err(|e| MLError::ModelError(format!("ensemble clipped SAXPY: {e}")))?; } } @@ -1192,9 +1270,12 @@ impl GpuDqnTrainer { let scratch_d_h_b1 = raw_device_ptr(&self.bw_d_h_b1, &self.stream); let scratch_d_h_b2 = raw_device_ptr(&self.bw_d_h_b2, &self.stream); - // Run full backward pass with CQL logit gradients. - // This ACCUMULATES into grad_buf (beta=1.0 in cuBLAS SGEMM), - // adding CQL parameter gradients on top of C51's. + // Zero CQL scratch buffer (backward_full uses beta=1.0 accumulation) + self.stream.memset_zeros(&mut self.cql_grad_scratch) + .map_err(|e| MLError::ModelError(format!("zero cql_grad_scratch: {e}")))?; + + // Run full backward pass with CQL logit gradients into ISOLATED scratch buffer. + // Produces CQL parameter gradients WITHOUT mixing with C51's grad_buf. self.cublas_backward.backward_full( &self.stream, d_v_ptr, @@ -1203,7 +1284,7 @@ impl GpuDqnTrainer { h_s1_ptr, h_s2_ptr, h_v_ptr, &[h_b0_ptr, h_b1_ptr, h_b2_ptr], &w_ptrs, - raw_device_ptr(&self.grad_buf, &self.stream), + raw_device_ptr(&self.cql_grad_scratch, &self.stream), scratch_d_h_s2, scratch_d_h_s1, scratch_d_h_v, &[scratch_d_h_b0, scratch_d_h_b1, scratch_d_h_b2], ).map_err(|e| MLError::ModelError(format!("CQL backward_full: {e}")))?; @@ -1216,16 +1297,70 @@ impl GpuDqnTrainer { self.cql_logit_grad_kernel.is_some() && self.config.cql_alpha > 0.0 } - /// Apply spectral normalization to trunk weight matrices W_s1 and W_s2. + /// Clip CQL gradient scratch and add to grad_buf via clipped SAXPY. + /// + /// Called after `apply_cql_gradient` populated `cql_grad_scratch`. + /// Computes norm of scratch, clips to `cql_budget`, then SAXPYs into grad_buf. + pub fn apply_cql_clipped_saxpy(&mut self, cql_budget: f32) -> Result<(), MLError> { + let _evt_guard = EventTrackingGuard::new(self.stream.context()); + + // Compute CQL gradient norm + self.stream.memset_zeros(&mut self.grad_norm_buf) + .map_err(|e| MLError::ModelError(format!("zero cql_grad_norm: {e}")))?; + + let total = self.total_params as i32; + let blocks = ((self.total_params + 255) / 256) as u32; + let scratch_ptr = raw_device_ptr(&self.cql_grad_scratch, &self.stream); + let norm_ptr = raw_device_ptr(&self.grad_norm_buf, &self.stream); + + unsafe { + self.stream + .launch_builder(&self.grad_norm_kernel) + .arg(&scratch_ptr) + .arg(&norm_ptr) + .arg(&total) + .launch(LaunchConfig { + grid_dim: (blocks, 1, 1), + block_dim: (256, 1, 1), + shared_mem_bytes: 256, + }) + .map_err(|e| MLError::ModelError(format!("CQL grad_norm: {e}")))?; + } + + // Clipped SAXPY: grad_buf += 1.0 × clip(cql_scratch, cql_budget) + let grad_ptr = raw_device_ptr(&self.grad_buf, &self.stream); + let alpha = 1.0_f32; + unsafe { + self.stream + .launch_builder(&self.clipped_saxpy_kernel) + .arg(&grad_ptr) + .arg(&scratch_ptr) + .arg(&alpha) + .arg(&cql_budget) + .arg(&norm_ptr) + .arg(&total) + .launch(LaunchConfig { + grid_dim: (blocks, 1, 1), + block_dim: (256, 1, 1), + shared_mem_bytes: 0, + }) + .map_err(|e| MLError::ModelError(format!("CQL clipped SAXPY: {e}")))?; + } + + Ok(()) + } + + /// Apply spectral normalization to all 10 weight matrices (trunk + 8 heads). /// /// One step of power iteration per call (standard practice — single step /// per training step converges in practice). Constrains ||W||_σ ≤ 1.0, /// bounding the network's Lipschitz constant and preventing Q-value explosion. /// - /// Runs outside the CUDA Graph on `params_buf` and `DuelingWeightSet`. + /// Runs outside the CUDA Graph on `params_buf`, `DuelingWeightSet`, and `BranchingWeightSet`. pub fn apply_spectral_norm( &mut self, online_dueling: &mut DuelingWeightSet, + online_branching: &mut BranchingWeightSet, ) -> Result<(), MLError> { let sh1 = self.config.shared_h1 as i32; let sh2 = self.config.shared_h2 as i32; @@ -1272,41 +1407,167 @@ impl GpuDqnTrainer { } } - // Sync scaled weights back to params_buf (trunk portion) so the CUDA - // Graph's Adam sees the spectrally normalized weights. Without this, - // the graph reads stale unscaled weights from params_buf. + // ── Head spectral norm kernel launches ─────────────────────────── + let vh = self.config.value_h as i32; + let ah = self.config.adv_h as i32; + let na = self.config.num_atoms as i32; + let b0 = self.config.branch_0_size as i32; + let b1 = self.config.branch_1_size as i32; + let b2 = self.config.branch_2_size as i32; + + macro_rules! spec_norm { + ($w_slice:expr, $u_slice:expr, $v_slice:expr, $out_dim:expr, $in_dim:expr, $label:literal) => {{ + let (w_ptr, _gw) = $w_slice.device_ptr(&self.stream); + let (u_ptr, _gu) = $u_slice.device_ptr(&self.stream); + let (v_ptr, _gv) = $v_slice.device_ptr(&self.stream); + unsafe { + self.stream + .launch_builder(&self.spectral_norm_kernel) + .arg(&w_ptr).arg(&u_ptr).arg(&v_ptr) + .arg(&$out_dim).arg(&$in_dim).arg(&sigma_max) + .launch(LaunchConfig { + grid_dim: (1, 1, 1), + block_dim: (256, 1, 1), + shared_mem_bytes: 512 * 4, + }) + .map_err(|e| MLError::ModelError(format!("spectral_norm {}: {e}", $label)))?; + } + }}; + } + + // W_v1 [value_h, shared_h2] + spec_norm!(online_dueling.w_v1, self.spec_u_v1, self.spec_v_v1, vh, sh2, "W_v1"); + // W_v2 [num_atoms, value_h] + spec_norm!(online_dueling.w_v2, self.spec_u_v2, self.spec_v_v2, na, vh, "W_v2"); + // W_a1 [adv_h, shared_h2] + spec_norm!(online_dueling.w_a1, self.spec_u_a1, self.spec_v_a1, ah, sh2, "W_a1"); + // W_a2 [b0*num_atoms, adv_h] + spec_norm!(online_dueling.w_a2, self.spec_u_a2, self.spec_v_a2, b0 * na, ah, "W_a2"); + // W_bo1 [adv_h, shared_h2] + spec_norm!(online_branching.w_bo1, self.spec_u_bo1, self.spec_v_bo1, ah, sh2, "W_bo1"); + // W_bo2 [b1*num_atoms, adv_h] + spec_norm!(online_branching.w_bo2, self.spec_u_bo2, self.spec_v_bo2, b1 * na, ah, "W_bo2"); + // W_bu1 [adv_h, shared_h2] + spec_norm!(online_branching.w_bu1, self.spec_u_bu1, self.spec_v_bu1, ah, sh2, "W_bu1"); + // W_bu2 [b2*num_atoms, adv_h] + spec_norm!(online_branching.w_bu2, self.spec_u_bu2, self.spec_v_bu2, b2 * na, ah, "W_bu2"); + + // Sync scaled weights back to params_buf so the CUDA Graph's Adam + // sees the spectrally normalized weights. Without this, the graph reads + // stale unscaled weights from params_buf. { let f32_sz = std::mem::size_of::(); + let param_sizes = compute_param_sizes(&self.config); + let byte_offset = |idx: usize| -> u64 { + param_sizes[..idx].iter().sum::() as u64 * f32_sz as u64 + }; + let params_base = raw_device_ptr(&self.params_buf, &self.stream); + + macro_rules! sync_w { + ($w_slice:expr, $goff_idx:expr, $elem_count:expr, $label:literal) => {{ + let src = raw_device_ptr(&$w_slice, &self.stream); + let n_bytes = ($elem_count as usize) * f32_sz; + unsafe { + cudarc::driver::result::memcpy_dtod_async( + params_base + byte_offset($goff_idx), + src, n_bytes, self.stream.cu_stream() + ).map_err(|e| MLError::ModelError(format!("spectral sync {}: {e}", $label)))?; + } + }}; + } + let sh1_u = self.config.shared_h1; let sh2_u = self.config.shared_h2; let sd_u = self.config.state_dim; - let params_base = raw_device_ptr(&self.params_buf, &self.stream); + let vh_u = self.config.value_h; + let ah_u = self.config.adv_h; + let na_u = self.config.num_atoms; + let b0_u = self.config.branch_0_size; + let b1_u = self.config.branch_1_size; + let b2_u = self.config.branch_2_size; - // w_s1 → params_buf[goff_w_s1] - let w_s1_src = raw_device_ptr(&online_dueling.w_s1, &self.stream); - let w_s1_n = sh1_u * sd_u * f32_sz; - unsafe { - cudarc::driver::result::memcpy_dtod_async( - params_base, w_s1_src, w_s1_n, self.stream.cu_stream() - ).map_err(|e| MLError::ModelError(format!("spectral sync w_s1: {e}")))?; - } - - // w_s2 → params_buf[goff_w_s2] - let w_s2_src = raw_device_ptr(&online_dueling.w_s2, &self.stream); - // GOFF layout: w_s1(sh1*sd), b_s1(sh1), w_s2(sh2*sh1), b_s2(sh2) - let goff_w_s2_bytes = (sh1_u * sd_u + sh1_u) * f32_sz; - let w_s2_n = sh2_u * sh1_u * f32_sz; - unsafe { - cudarc::driver::result::memcpy_dtod_async( - params_base + goff_w_s2_bytes as u64, - w_s2_src, w_s2_n, self.stream.cu_stream() - ).map_err(|e| MLError::ModelError(format!("spectral sync w_s2: {e}")))?; - } + // Trunk (indices 0, 2) + sync_w!(online_dueling.w_s1, 0, sh1_u * sd_u, "w_s1"); + sync_w!(online_dueling.w_s2, 2, sh2_u * sh1_u, "w_s2"); + // Value head (indices 4, 6) + sync_w!(online_dueling.w_v1, 4, vh_u * sh2_u, "w_v1"); + sync_w!(online_dueling.w_v2, 6, na_u * vh_u, "w_v2"); + // Exposure/adv head (indices 8, 10) + sync_w!(online_dueling.w_a1, 8, ah_u * sh2_u, "w_a1"); + sync_w!(online_dueling.w_a2, 10, b0_u * na_u * ah_u, "w_a2"); + // Order head (indices 12, 14) + sync_w!(online_branching.w_bo1, 12, ah_u * sh2_u, "w_bo1"); + sync_w!(online_branching.w_bo2, 14, b1_u * na_u * ah_u, "w_bo2"); + // Urgency head (indices 16, 18) + sync_w!(online_branching.w_bu1, 16, ah_u * sh2_u, "w_bu1"); + sync_w!(online_branching.w_bu2, 18, b2_u * na_u * ah_u, "w_bu2"); } Ok(()) } + /// Clip `grad_buf` in-place: zero grad_norm_buf, compute norm, then scale down if > max_norm. + /// + /// Used between gradient injection phases to prevent any component combination + /// from overwhelming the subsequent auxiliary gradient additions. + /// All operations are async on the same stream — zero CPU sync. + pub fn clip_grad_buf_inplace(&mut self, max_norm: f32) -> Result<(), MLError> { + let _evt_guard = EventTrackingGuard::new(self.stream.context()); + + // Zero the grad_norm accumulator + self.stream + .memset_zeros(&mut self.grad_norm_buf) + .map_err(|e| MLError::ModelError(format!("zero grad_norm for clip: {e}")))?; + + // Compute grad_norm (sum of squares) + self.launch_grad_norm()?; + + // Clip in-place + let total = self.total_params as i32; + let blocks = ((self.total_params + 255) / 256) as u32; + unsafe { + self.stream + .launch_builder(&self.clip_grad_kernel) + .arg(&raw_device_ptr(&self.grad_buf, &self.stream)) + .arg(&raw_device_ptr(&self.grad_norm_buf, &self.stream)) + .arg(&max_norm) + .arg(&total) + .launch(LaunchConfig { + grid_dim: (blocks, 1, 1), + block_dim: (256, 1, 1), + shared_mem_bytes: 0, + }) + .map_err(|e| MLError::ModelError(format!("clip_grad_kernel: {e}")))?; + } + + Ok(()) + } + + /// Compute and return the current grad_buf L2 norm (synchronous readback). + /// + /// Used for per-component gradient diagnostics. Costs one stream sync + + /// 4-byte DtoH transfer. Call sparingly (e.g. once per epoch or every N steps). + pub fn read_grad_norm_sync(&mut self) -> Result { + let _evt_guard = EventTrackingGuard::new(self.stream.context()); + + self.stream + .memset_zeros(&mut self.grad_norm_buf) + .map_err(|e| MLError::ModelError(format!("zero grad_norm for diag: {e}")))?; + + self.launch_grad_norm()?; + + unsafe { cudarc::driver::sys::cuStreamSynchronize(self.stream.cu_stream()); } + + let mut norm_sq = [0.0_f32; 1]; + unsafe { + cudarc::driver::sys::cuMemcpyDtoH_v2( + norm_sq.as_mut_ptr().cast(), + raw_device_ptr(&self.grad_norm_buf, &self.stream), 4, + ); + } + Ok(norm_sq[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]` @@ -1423,7 +1684,7 @@ impl GpuDqnTrainer { // per array. Stack is set once in DQNTrainer::new() (64KB for all kernels). // ── Compile 4 utility kernels (grad_norm, adam_update, BF16 converters) ─ - let (grad_norm_kernel, adam_update_kernel, f32_to_bf16_kernel, bf16_to_f32_kernel, saxpy_kernel, zero_kernel, regime_scale_kernel, shrink_perturb, _relu_mask_in_module, spectral_norm_kernel) = + let (grad_norm_kernel, adam_update_kernel, f32_to_bf16_kernel, bf16_to_f32_kernel, saxpy_kernel, zero_kernel, regime_scale_kernel, shrink_perturb, _relu_mask_in_module, spectral_norm_kernel, clipped_saxpy_kernel, clip_grad_kernel) = compile_training_kernels(&stream, &config)?; // ── Compile EMA kernel (standalone — not captured in CUDA Graph) ── @@ -1474,6 +1735,7 @@ impl GpuDqnTrainer { let m_buf = alloc_f32(&stream, total_params, "adam_m")?; let v_buf = alloc_f32(&stream, total_params, "adam_v")?; let grad_norm_buf = alloc_f32(&stream, 1, "grad_norm")?; + let cql_grad_scratch = alloc_f32(&stream, total_params, "cql_grad_scratch")?; let t_buf = alloc_i32(&stream, 1, "adam_t")?; // ── Allocate flat BF16 weight mirror buffers ────────────────── @@ -1632,6 +1894,45 @@ impl GpuDqnTrainer { stream.memcpy_htod(&init_v_s2, &mut spec_v_s2) .map_err(|e| MLError::ModelError(format!("spec_v_s2 init: {e}")))?; + // Head spectral norm vectors: value head, adv/branch heads + let vh = config.value_h; + let ah = config.adv_h; + let na = config.num_atoms; + let sh2 = config.shared_h2; + let b0 = config.branch_0_size; + let b1 = config.branch_1_size; + let b2 = config.branch_2_size; + + macro_rules! alloc_spec_pair { + ($u_name:ident, $v_name:ident, $u_n:expr, $v_n:expr, $lbl_u:literal, $lbl_v:literal) => { + let mut $u_name = alloc_f32(&stream, $u_n, $lbl_u)?; + let mut $v_name = alloc_f32(&stream, $v_n, $lbl_v)?; + let init_u = rand_unit($u_n, &mut rng_state); + let init_v = rand_unit($v_n, &mut rng_state); + stream.memcpy_htod(&init_u, &mut $u_name) + .map_err(|e| MLError::ModelError(format!("{} init: {e}", $lbl_u)))?; + stream.memcpy_htod(&init_v, &mut $v_name) + .map_err(|e| MLError::ModelError(format!("{} init: {e}", $lbl_v)))?; + }; + } + + // W_v1 [value_h, shared_h2] + alloc_spec_pair!(spec_u_v1, spec_v_v1, vh, sh2, "spec_u_v1", "spec_v_v1"); + // W_v2 [num_atoms, value_h] + alloc_spec_pair!(spec_u_v2, spec_v_v2, na, vh, "spec_u_v2", "spec_v_v2"); + // W_a1 [adv_h, shared_h2] + alloc_spec_pair!(spec_u_a1, spec_v_a1, ah, sh2, "spec_u_a1", "spec_v_a1"); + // W_a2 [b0*num_atoms, adv_h] + alloc_spec_pair!(spec_u_a2, spec_v_a2, b0 * na, ah, "spec_u_a2", "spec_v_a2"); + // W_bo1 [adv_h, shared_h2] + alloc_spec_pair!(spec_u_bo1, spec_v_bo1, ah, sh2, "spec_u_bo1", "spec_v_bo1"); + // W_bo2 [b1*num_atoms, adv_h] + alloc_spec_pair!(spec_u_bo2, spec_v_bo2, b1 * na, ah, "spec_u_bo2", "spec_v_bo2"); + // W_bu1 [adv_h, shared_h2] + alloc_spec_pair!(spec_u_bu1, spec_v_bu1, ah, sh2, "spec_u_bu1", "spec_v_bu1"); + // W_bu2 [b2*num_atoms, adv_h] + alloc_spec_pair!(spec_u_bu2, spec_v_bu2, b2 * na, ah, "spec_u_bu2", "spec_v_bu2"); + // ── Precompute F32 GOFF byte offsets (for cuBLAS weight pointers) ── let mut f32_goff_byte_offsets = [0_u64; 20]; { @@ -1707,10 +2008,20 @@ impl GpuDqnTrainer { shrink_perturb_kernel: shrink_perturb, relu_mask_kernel: relu_mask_standalone, spectral_norm_kernel, + clipped_saxpy_kernel, + clip_grad_kernel, spec_u_s1, spec_v_s1, spec_u_s2, spec_v_s2, + spec_u_v1, spec_v_v1, + spec_u_v2, spec_v_v2, + spec_u_a1, spec_v_a1, + spec_u_a2, spec_v_a2, + spec_u_bo1, spec_v_bo1, + spec_u_bo2, spec_v_bo2, + spec_u_bu1, spec_v_bu1, + spec_u_bu2, spec_v_bu2, iqn_trunk_m, iqn_trunk_v, iqn_trunk_adam_step: 0, @@ -1746,6 +2057,7 @@ impl GpuDqnTrainer { m_buf, v_buf, grad_norm_buf, + cql_grad_scratch, t_buf, adam_step: 0, total_params, @@ -3703,17 +4015,18 @@ impl GpuDqnTrainer { // ── Compilation ───────────────────────────────────────────────────────────── -/// Load the 10 utility kernels from precompiled cubin (ZERO runtime nvcc). +/// Load the 12 utility kernels from precompiled cubin (ZERO runtime nvcc). /// /// The cubin is produced by build.rs from `common_device_functions.cuh` + /// `dqn_utility_kernels.cu`. Contains: grad_norm, adam_update, f32_to_bf16, -/// bf16_to_f32, saxpy, zero, regime_scale, shrink_perturb, relu_mask, spectral_norm. +/// bf16_to_f32, saxpy, zero, regime_scale, shrink_perturb, relu_mask, spectral_norm, +/// clipped_saxpy, clip_grad. /// -/// Returns `(grad_norm, adam_update, f32_to_bf16, bf16_to_f32, saxpy, zero, regime_scale, shrink_perturb, relu_mask, spectral_norm)`. +/// Returns `(grad_norm, adam_update, f32_to_bf16, bf16_to_f32, saxpy, zero, regime_scale, shrink_perturb, relu_mask, spectral_norm, clipped_saxpy, clip_grad)`. fn compile_training_kernels( stream: &Arc, config: &GpuDqnTrainConfig, -) -> Result<(CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction), MLError> { +) -> Result<(CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction), MLError> { info!( state_dim = config.state_dim, total_params = compute_total_params(config), @@ -3746,9 +4059,13 @@ fn compile_training_kernels( .map_err(|e| MLError::ModelError(format!("dqn_relu_mask_kernel load: {e}")))?; let spectral_norm = module.load_function("spectral_norm_kernel") .map_err(|e| MLError::ModelError(format!("spectral_norm_kernel load: {e}")))?; + let clipped_saxpy = module.load_function("dqn_clipped_saxpy_kernel") + .map_err(|e| MLError::ModelError(format!("dqn_clipped_saxpy_kernel load: {e}")))?; + let clip_grad = module.load_function("dqn_clip_grad_kernel") + .map_err(|e| MLError::ModelError(format!("dqn_clip_grad_kernel load: {e}")))?; - info!("GpuDqnTrainer: 10 utility kernels loaded from precompiled cubin"); - Ok((grad_norm, adam_update, f32_to_bf16, bf16_to_f32, saxpy, zero, regime_scale, shrink_perturb, _relu_mask_from_module, spectral_norm)) + info!("GpuDqnTrainer: 12 utility kernels loaded from precompiled cubin"); + Ok((grad_norm, adam_update, f32_to_bf16, bf16_to_f32, saxpy, zero, regime_scale, shrink_perturb, _relu_mask_from_module, spectral_norm, clipped_saxpy, clip_grad)) } /// Load the standalone Polyak EMA kernel from precompiled cubin. diff --git a/crates/ml/src/hyperopt/adapters/dqn.rs b/crates/ml/src/hyperopt/adapters/dqn.rs index 5243c044b..385a2d47c 100644 --- a/crates/ml/src/hyperopt/adapters/dqn.rs +++ b/crates/ml/src/hyperopt/adapters/dqn.rs @@ -2550,7 +2550,7 @@ impl HyperparameterOptimizable for DQNTrainer { max_position_absolute: params.max_position_absolute, // BLOCKER #2: Use hyperopt-tunable position limit // WAVE 17: Hyperopt-tuned parameters - entropy_coefficient: Some(params.entropy_coefficient), // Use hyperopt value + entropy_coefficient: params.entropy_coefficient, // Use hyperopt value noisy_epsilon_floor: Some(params.noisy_epsilon_floor), // Exploration floor for noisy nets count_bonus_coefficient: Some(params.count_bonus_coefficient), // UCB exploration bonus transaction_cost_multiplier: params.transaction_cost_multiplier, // Use hyperopt value diff --git a/crates/ml/src/trainers/dqn/config.rs b/crates/ml/src/trainers/dqn/config.rs index 01ea98501..2dd76f63a 100644 --- a/crates/ml/src/trainers/dqn/config.rs +++ b/crates/ml/src/trainers/dqn/config.rs @@ -956,8 +956,8 @@ pub struct DQNHyperparameters { pub max_position_absolute: f64, // WAVE 17: Hyperopt-tuned parameters - /// SAC-style entropy coefficient for Q-value softmax regularization (default: 0.01) - pub entropy_coefficient: Option, + /// SAC-style entropy coefficient for Q-value softmax regularization (default: 0.001) + pub entropy_coefficient: f64, /// Minimum epsilon when noisy nets are enabled (default: 0.05) pub noisy_epsilon_floor: Option, /// Count-based exploration bonus coefficient (default: 0.1, range: [0.01, 0.5]) @@ -1428,7 +1428,7 @@ impl DQNHyperparameters { max_position_absolute: 2.0, // Default: ±2.0 position limit (matches production) // WAVE 17: Hyperopt-tuned parameters - entropy_coefficient: Some(0.001), // Must be ≤ reward magnitude (~0.001). Old 0.01 was 10x reward → entropy dominated learning. + entropy_coefficient: 0.001, // Must be ≤ reward magnitude (~0.001). Old 0.01 was 10x reward → entropy dominated learning. noisy_epsilon_floor: Some(0.02), // NoisyNet handles exploration; low floor avoids interference count_bonus_coefficient: Some(0.1), transaction_cost_multiplier: 1.0, diff --git a/crates/ml/src/trainers/dqn/data_loading.rs b/crates/ml/src/trainers/dqn/data_loading.rs index 0ecf2522c..6b3e3c78a 100644 --- a/crates/ml/src/trainers/dqn/data_loading.rs +++ b/crates/ml/src/trainers/dqn/data_loading.rs @@ -303,8 +303,8 @@ impl DQNTrainer { let n = close_prices_f32.len(); let mut simple_ret = vec![0.0_f32]; for i in 1..n { - let prev = close_prices_f32.get(i - 1).copied().unwrap_or(1.0); - let curr = close_prices_f32.get(i).copied().unwrap_or(prev); + let prev = close_prices_f32[i - 1]; + let curr = close_prices_f32[i]; simple_ret.push(if prev.abs() > 1e-10 { (curr - prev) / prev } else { 0.0 }); } simple_ret diff --git a/crates/ml/src/trainers/dqn/fused_training.rs b/crates/ml/src/trainers/dqn/fused_training.rs index 34b3c635f..ebff14e42 100644 --- a/crates/ml/src/trainers/dqn/fused_training.rs +++ b/crates/ml/src/trainers/dqn/fused_training.rs @@ -44,6 +44,14 @@ use crate::dqn::replay_buffer_type::{BatchSample, GpuBatch}; use super::config::DQNAgentType; use super::DQNHyperparameters; +/// Per-component gradient norm budget fractions for auxiliary objectives. +/// C51 gets whatever remains: `1.0 - sum(active_auxiliary_budgets)`. +/// When all auxiliaries are active: C51=70%, CQL=15%, IQN=10%, Ens=5%. +/// When auxiliaries are disabled, their budget automatically goes to C51. +pub(crate) const CQL_GRAD_BUDGET: f32 = 0.15; +pub(crate) const IQN_GRAD_BUDGET: f32 = 0.10; +pub(crate) const ENS_GRAD_BUDGET: f32 = 0.05; + /// Fused CUDA training context -- owns the `GpuDqnTrainer` and extracted weight sets. /// /// Weight sets are extracted from the Candle GpuVarStore at initialization via direct @@ -69,6 +77,10 @@ pub(crate) struct FusedTrainingCtx { batch_size: usize, /// Steps since last GpuVarStore sync (deferred to epoch boundary). steps_since_varmap_sync: usize, + /// Last observed C51 raw gradient norm (before budget clip). Updated every 1000 steps. + last_c51_raw_norm: f32, + /// Last observed combined gradient norm (from Adam readback). + last_combined_norm: f32, /// GPU HER relabeler -- initialized when `her_ratio > 0.0`. /// When active, `run_full_step()` splits the batch into normal + HER portions, /// runs the relabel kernel on the HER portion, and merges before training. @@ -157,6 +169,10 @@ impl FusedTrainingCtx { // collector, fused trainer, portfolio sim, monitoring) share this single // stream via Arc::clone — no more dual-stream event tracking conflicts. + // Resolve gradient clip norm ONCE — single source of truth for all sub-trainers. + // Was: main=10.0, IQN=1.0, IQL=1.0 → inconsistent 10× mismatch. + let resolved_grad_norm = hyperparams.gradient_clip_norm.unwrap_or(10.0); + // Build config from DQN network dimensions let (shared_h1, shared_h2, value_h, adv_h) = agent.network_dims(); let config = GpuDqnTrainConfig { @@ -178,12 +194,8 @@ impl FusedTrainingCtx { beta1: 0.9, beta2: 0.999, epsilon: 1e-8, - weight_decay: 1e-5, - // Scale gradient clip by (1 + iqn_lambda) to account for combined gradient: - // grad_buf = C51_grad + iqn_lambda * IQN_grad. Without scaling, the clip - // is too tight and Adam sees inconsistent clipping, causing loss oscillation. - max_grad_norm: hyperparams.gradient_clip_norm.unwrap_or(10.0) as f32 - * (1.0 + hyperparams.iqn_lambda as f32), + weight_decay: hyperparams.weight_decay as f32, + max_grad_norm: resolved_grad_norm as f32, spectral_norm_sigma_max: hyperparams.spectral_norm_sigma_max as f32, iqn_lambda: hyperparams.iqn_lambda as f32, iqn_num_quantiles: if hyperparams.iqn_lambda > 0.0 { hyperparams.num_quantiles } else { 0 }, @@ -252,7 +264,7 @@ impl FusedTrainingCtx { state_dim: dqn.config.state_dim, batch_size, lr: hyperparams.learning_rate as f32, - max_grad_norm: hyperparams.gradient_clip_norm.unwrap_or(1.0) as f32, + max_grad_norm: resolved_grad_norm as f32, ..GpuIqlConfig::default() }; match GpuIqlTrainer::new(stream.clone(), iql_config) { @@ -287,7 +299,7 @@ impl FusedTrainingCtx { branch_2_size: dqn.config.num_urgency_levels, batch_size, lr: hyperparams.learning_rate as f32, - max_grad_norm: hyperparams.gradient_clip_norm.unwrap_or(1.0) as f32, + max_grad_norm: resolved_grad_norm as f32, gamma: hyperparams.gamma as f32, ..GpuIqnConfig::default() }; @@ -435,6 +447,8 @@ impl FusedTrainingCtx { stream, batch_size, steps_since_varmap_sync: 0, + last_c51_raw_norm: 0.0, + last_combined_norm: 0.0, gpu_her, gpu_iql, gpu_iqn, @@ -463,6 +477,12 @@ impl FusedTrainingCtx { self.steps_since_varmap_sync } + /// Per-component gradient norm diagnostics for epoch-level Prometheus reporting. + /// Returns (c51_raw_norm, combined_norm). + pub(crate) fn grad_norm_diagnostics(&self) -> (f32, f32) { + (self.last_c51_raw_norm, self.last_combined_norm) + } + /// Run one full fused training step — GPU-only, no CPU paths. /// /// All batch data stays GPU-resident. DtoD copies between GpuBatch tensors @@ -561,7 +581,7 @@ impl FusedTrainingCtx { // so the forward pass reads normalized weights. This is correct placement: // normalize → forward → loss → backward → Adam → (weights grow) → normalize → ... // Previously ran AFTER Adam which caused a tug-of-war. - self.trainer.apply_spectral_norm(&mut self.online_dueling) + self.trainer.apply_spectral_norm(&mut self.online_dueling, &mut self.online_branching) .map_err(|e| anyhow::anyhow!("Spectral norm (pre-forward): {e}"))?; // ── Step 2: Upload batch + replay graph_forward ONLY ───────────── @@ -573,6 +593,33 @@ impl FusedTrainingCtx { &self.target_dueling, &self.target_branching, ).map_err(|e| anyhow::anyhow!("Fused train_step_gpu (forward only): {e}"))?; + // ── Step 2b: Clip raw C51 gradient to its dynamic budget ──────── + // C51 gets whatever budget the active auxiliaries don't use. + // When all aux are active: C51=70%. When none: C51=100%. + { + 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 }; + let ens_frac = if !self.ensemble_extra_heads.is_empty() { ENS_GRAD_BUDGET } else { 0.0 }; + let c51_frac = 1.0 - cql_frac - iqn_frac - ens_frac; + let c51_budget = self.trainer.config().max_grad_norm * c51_frac; + + // Diagnostic: log C51 pre-clip norm every 1000 steps + if self.steps_since_varmap_sync % 1000 == 0 { + let pre_clip_norm = self.trainer.read_grad_norm_sync() + .unwrap_or(f32::NAN); + self.last_c51_raw_norm = pre_clip_norm; + tracing::info!( + c51_raw_grad_norm = pre_clip_norm, + c51_budget, + step = self.steps_since_varmap_sync, + "Per-component gradient diagnostic (C51 before budget clip)" + ); + } + + self.trainer.clip_grad_buf_inplace(c51_budget) + .map_err(|e| anyhow::anyhow!("C51 gradient budget clip: {e}"))?; + } + // ── Step 3: GPU-native Polyak EMA target update ────────────────── { let dqn = agent.primary_dqn_mut(); @@ -739,14 +786,17 @@ impl FusedTrainingCtx { // Spectral norm moved to Step 1b (before graph_forward) — correct placement. - // ── Step 5c: CQL conservative penalty (if enabled) ─────────────── - // Computes CQL logit gradients from current Q-values and adds - // parameter gradients into grad_buf via a second cuBLAS backward. - // Same pattern as IQN trunk gradient injection. + // ── Step 5c: CQL conservative penalty (isolated gradient) ───────── + // CQL backward runs into a SEPARATE scratch buffer (cql_grad_scratch). + // Its gradient is independently clipped to CQL's budget fraction, + // then added to grad_buf via clipped SAXPY. No mixing with C51. if self.trainer.has_cql() { match self.trainer.apply_cql_gradient() { Ok(true) => { - tracing::trace!("CQL gradient injected into grad_buf"); + let cql_budget = self.trainer.config().max_grad_norm * CQL_GRAD_BUDGET; + self.trainer.apply_cql_clipped_saxpy(cql_budget) + .map_err(|e| anyhow::anyhow!("CQL clipped SAXPY: {e}"))?; + tracing::trace!("CQL gradient: isolated → clipped → SAXPY into grad_buf"); } Ok(false) => {} // CQL disabled or alpha=0 Err(e) => { @@ -755,9 +805,10 @@ impl FusedTrainingCtx { } } - // ── Step 5e: Replay graph_adam — Adam sees combined gradient ──── - // All auxiliary gradients (IQN, attention, ensemble, CQL) are now in grad_buf. - // The single Adam update sees: C51 + iqn_lambda*IQN + attention + diversity + CQL. + // ── Step 5e: Replay graph_adam — Adam sees budget-allocated gradient ── + // grad_buf contains: C51 (≤70%) + CQL (≤15%) + IQN (≤10%) + ensemble (≤5%). + // Budgets sum to ≤100% of max_grad_norm → Adam safety clip should never fire. + // Attention has its own optimizer and does NOT contribute to grad_buf. let fused_result = self.trainer.replay_adam_and_readback() .map_err(|e| { eprintln!("!!! ADAM REPLAY FAILED: {e}"); anyhow::anyhow!("graph_adam replay: {e}") })?; @@ -793,6 +844,7 @@ impl FusedTrainingCtx { .map_err(|e| anyhow::anyhow!("Fused GPU bookkeeping: {e}"))?; self.steps_since_varmap_sync += 1; + self.last_combined_norm = fused_result.grad_norm; // ── Step 7: Wrap raw scalars into GpuTrainResult ───────────────── GpuTrainResult::from_fused_scalars( diff --git a/crates/ml/src/trainers/dqn/mod.rs b/crates/ml/src/trainers/dqn/mod.rs index baff2ad24..83f52aaaf 100644 --- a/crates/ml/src/trainers/dqn/mod.rs +++ b/crates/ml/src/trainers/dqn/mod.rs @@ -24,7 +24,7 @@ mod early_stopping; pub mod expert_demos; pub(crate) mod financials; mod features; -mod fused_training; +pub(crate) mod fused_training; pub mod lr_scheduler; mod monitoring; mod risk; diff --git a/crates/ml/src/trainers/dqn/smoke_tests/gradient_budget.rs b/crates/ml/src/trainers/dqn/smoke_tests/gradient_budget.rs new file mode 100644 index 000000000..0688f7000 --- /dev/null +++ b/crates/ml/src/trainers/dqn/smoke_tests/gradient_budget.rs @@ -0,0 +1,366 @@ +//! Gradient budget and spectral norm validation tests. +//! +//! These tests validate the per-component gradient clipping infrastructure +//! and spectral normalization at the GPU kernel level, without requiring +//! training data. They construct a GpuDqnTrainer directly and exercise +//! the clipping/norm kernels with synthetic gradient data. + +use std::sync::Arc; + +use super::helpers::cuda_device; +use crate::cuda_pipeline::gpu_dqn_trainer::{GpuDqnTrainConfig, GpuDqnTrainer}; +use crate::cuda_pipeline::gpu_weights::{DuelingWeightSet, BranchingWeightSet}; + +/// Small config for fast kernel-level tests (no training, just buffer ops). +fn test_config() -> GpuDqnTrainConfig { + GpuDqnTrainConfig { + state_dim: 16, + shared_h1: 32, + shared_h2: 32, + value_h: 16, + adv_h: 16, + num_atoms: 11, + v_min: -10.0, + v_max: 10.0, + branch_0_size: 9, + branch_1_size: 3, + branch_2_size: 3, + batch_size: 8, + max_grad_norm: 10.0, + spectral_norm_sigma_max: 3.0, + ..GpuDqnTrainConfig::default() + } +} + +/// Allocate a DuelingWeightSet with random values (not zero — spectral norm +/// needs non-degenerate matrices for power iteration to converge). +fn alloc_dueling(stream: &Arc, cfg: &GpuDqnTrainConfig) -> DuelingWeightSet { + let na = cfg.num_atoms; + let alloc = |n: usize| -> cudarc::driver::CudaSlice { + // Fill with 1.0 (non-zero, non-degenerate for spectral norm) + stream.clone_htod(&vec![0.1_f32; n]).expect("alloc dueling weight") + }; + DuelingWeightSet { + w_s1: alloc(cfg.shared_h1 * cfg.state_dim), + b_s1: alloc(cfg.shared_h1), + w_s2: alloc(cfg.shared_h2 * cfg.shared_h1), + b_s2: alloc(cfg.shared_h2), + w_v1: alloc(cfg.value_h * cfg.shared_h2), + b_v1: alloc(cfg.value_h), + w_v2: alloc(na * cfg.value_h), + b_v2: alloc(na), + w_a1: alloc(cfg.adv_h * cfg.shared_h2), + b_a1: alloc(cfg.adv_h), + w_a2: alloc(cfg.branch_0_size * na * cfg.adv_h), + b_a2: alloc(cfg.branch_0_size * na), + } +} + +fn alloc_branching(stream: &Arc, cfg: &GpuDqnTrainConfig) -> BranchingWeightSet { + let na = cfg.num_atoms; + let alloc = |n: usize| -> cudarc::driver::CudaSlice { + stream.clone_htod(&vec![0.1_f32; n]).expect("alloc branching weight") + }; + BranchingWeightSet { + w_bo1: alloc(cfg.adv_h * cfg.shared_h2), + b_bo1: alloc(cfg.adv_h), + w_bo2: alloc(cfg.branch_1_size * na * cfg.adv_h), + b_bo2: alloc(cfg.branch_1_size * na), + w_bu1: alloc(cfg.adv_h * cfg.shared_h2), + b_bu1: alloc(cfg.adv_h), + w_bu2: alloc(cfg.branch_2_size * na * cfg.adv_h), + b_bu2: alloc(cfg.branch_2_size * na), + } +} + +/// clip_grad_buf_inplace must reduce gradient norm to at most max_norm. +#[test] +fn test_clip_grad_buf_reduces_norm() -> anyhow::Result<()> { + let dev = cuda_device(); + let stream = Arc::clone(dev.cuda_stream().expect("cuda stream")); + let cfg = test_config(); + let mut trainer = GpuDqnTrainer::new(stream.clone(), cfg)?; + + // Inject a large gradient: fill grad_buf with 1.0 (norm = sqrt(total_params)) + let total_params = trainer.total_params(); + let big_grad = vec![1.0_f32; total_params]; + stream.memcpy_htod(&big_grad, trainer.grad_buf_mut()) + .map_err(|e| anyhow::anyhow!("{e}"))?; + + let norm_before = trainer.read_grad_norm_sync()?; + assert!(norm_before > 10.0, "pre-clip norm should be large, got {norm_before}"); + + // Clip to max_grad_norm = 10.0 + trainer.clip_grad_buf_inplace(10.0)?; + + let norm_after = trainer.read_grad_norm_sync()?; + assert!( + norm_after <= 10.0 + 0.01, // small tolerance for float precision + "post-clip norm should be ≤ 10.0, got {norm_after}" + ); + assert!( + norm_after > 9.9, + "post-clip norm should be close to 10.0 (not zero), got {norm_after}" + ); + + Ok(()) +} + +/// clip_grad_buf_inplace must NOT change gradients that are already within budget. +#[test] +fn test_clip_grad_buf_no_op_when_small() -> anyhow::Result<()> { + let dev = cuda_device(); + let stream = Arc::clone(dev.cuda_stream().expect("cuda stream")); + let cfg = test_config(); + let mut trainer = GpuDqnTrainer::new(stream.clone(), cfg)?; + + // Inject a small gradient: fill with tiny values + let total_params = trainer.total_params(); + let small_grad = vec![0.001_f32; total_params]; + stream.memcpy_htod(&small_grad, trainer.grad_buf_mut()) + .map_err(|e| anyhow::anyhow!("{e}"))?; + + let norm_before = trainer.read_grad_norm_sync()?; + assert!(norm_before < 10.0, "pre-clip norm should be small, got {norm_before}"); + + trainer.clip_grad_buf_inplace(10.0)?; + + let norm_after = trainer.read_grad_norm_sync()?; + let ratio = norm_after / norm_before; + assert!( + (ratio - 1.0).abs() < 0.01, + "clip should be no-op for small gradients, ratio={ratio}" + ); + + Ok(()) +} + +/// Spectral norm on all 10 weight matrices must not panic and must +/// produce finite weights after normalization. +#[test] +fn test_spectral_norm_all_heads_no_panic() -> anyhow::Result<()> { + let dev = cuda_device(); + let stream = Arc::clone(dev.cuda_stream().expect("cuda stream")); + let cfg = test_config(); + let mut trainer = GpuDqnTrainer::new(stream.clone(), cfg.clone())?; + let mut dueling = alloc_dueling(&stream, &cfg); + let mut branching = alloc_branching(&stream, &cfg); + + // Run spectral norm — should not panic with non-zero weights + trainer.apply_spectral_norm(&mut dueling, &mut branching)?; + + // Verify a few weights are still finite after normalization + let mut w_s1_host = vec![0.0_f32; cfg.shared_h1 * cfg.state_dim]; + stream.memcpy_dtoh(&dueling.w_s1, &mut w_s1_host) + .map_err(|e| anyhow::anyhow!("{e}"))?; + assert!( + w_s1_host.iter().all(|v| v.is_finite()), + "W_s1 has non-finite values after spectral norm" + ); + + let mut w_v1_host = vec![0.0_f32; cfg.value_h * cfg.shared_h2]; + stream.memcpy_dtoh(&dueling.w_v1, &mut w_v1_host) + .map_err(|e| anyhow::anyhow!("{e}"))?; + assert!( + w_v1_host.iter().all(|v| v.is_finite()), + "W_v1 has non-finite values after spectral norm" + ); + + let mut w_bo1_host = vec![0.0_f32; cfg.adv_h * cfg.shared_h2]; + stream.memcpy_dtoh(&branching.w_bo1, &mut w_bo1_host) + .map_err(|e| anyhow::anyhow!("{e}"))?; + assert!( + w_bo1_host.iter().all(|v| v.is_finite()), + "W_bo1 has non-finite values after spectral norm" + ); + + Ok(()) +} + +/// Spectral norm must actually constrain the spectral norm of weight matrices. +/// After applying spectral norm with σ_max=3.0, the L2 operator norm +/// of each weight matrix should be ≤ σ_max. +#[test] +fn test_spectral_norm_constrains_operator_norm() -> anyhow::Result<()> { + let dev = cuda_device(); + let stream = Arc::clone(dev.cuda_stream().expect("cuda stream")); + let cfg = test_config(); + let mut trainer = GpuDqnTrainer::new(stream.clone(), cfg.clone())?; + + // Create weights with large values — spectral norm should constrain them + let alloc_large = |n: usize| -> cudarc::driver::CudaSlice { + stream.clone_htod(&vec![5.0_f32; n]).expect("alloc large weight") + }; + let na = cfg.num_atoms; + let mut dueling = DuelingWeightSet { + w_s1: alloc_large(cfg.shared_h1 * cfg.state_dim), + b_s1: alloc_large(cfg.shared_h1), + w_s2: alloc_large(cfg.shared_h2 * cfg.shared_h1), + b_s2: alloc_large(cfg.shared_h2), + w_v1: alloc_large(cfg.value_h * cfg.shared_h2), + b_v1: alloc_large(cfg.value_h), + w_v2: alloc_large(na * cfg.value_h), + b_v2: alloc_large(na), + w_a1: alloc_large(cfg.adv_h * cfg.shared_h2), + b_a1: alloc_large(cfg.adv_h), + w_a2: alloc_large(cfg.branch_0_size * na * cfg.adv_h), + b_a2: alloc_large(cfg.branch_0_size * na), + }; + let mut branching = BranchingWeightSet { + w_bo1: alloc_large(cfg.adv_h * cfg.shared_h2), + b_bo1: alloc_large(cfg.adv_h), + w_bo2: alloc_large(cfg.branch_1_size * na * cfg.adv_h), + b_bo2: alloc_large(cfg.branch_1_size * na), + w_bu1: alloc_large(cfg.adv_h * cfg.shared_h2), + b_bu1: alloc_large(cfg.adv_h), + w_bu2: alloc_large(cfg.branch_2_size * na * cfg.adv_h), + b_bu2: alloc_large(cfg.branch_2_size * na), + }; + + // Run spectral norm multiple times (power iteration converges over iterations) + for _ in 0..5 { + trainer.apply_spectral_norm(&mut dueling, &mut branching)?; + } + + // Check W_s1 [32, 16]: Frobenius norm of a rank-1 matrix with σ_max=3.0 + // would be 3.0. A uniform matrix of 5.0 has σ = 5*sqrt(rows*cols), + // after spectral norm it should have σ ≤ 3.0. + // Proxy check: Frobenius norm should be significantly reduced from initial. + let rows = cfg.shared_h1; + let cols = cfg.state_dim; + let mut w_host = vec![0.0_f32; rows * cols]; + stream.memcpy_dtoh(&dueling.w_s1, &mut w_host) + .map_err(|e| anyhow::anyhow!("{e}"))?; + + let frobenius: f32 = w_host.iter().map(|x| x * x).sum::().sqrt(); + let initial_frobenius = 5.0 * (rows as f32 * cols as f32).sqrt(); + + assert!( + frobenius < initial_frobenius * 0.5, + "Frobenius norm should be significantly reduced: {frobenius:.2} vs initial {initial_frobenius:.2}" + ); + + // The spectral norm (largest singular value) should be ≤ σ_max. + // For a uniform matrix after spectral norm, the spectral norm equals + // σ_max if the original σ > σ_max. We can estimate σ ≈ Frobenius / sqrt(min(m,n)). + let approx_sigma = frobenius / (rows.min(cols) as f32).sqrt(); + assert!( + approx_sigma < cfg.spectral_norm_sigma_max * 1.5, // tolerance for approximation + "Approximate spectral norm {approx_sigma:.2} should be near σ_max={:.1}", + cfg.spectral_norm_sigma_max + ); + + Ok(()) +} + +/// CQL clipped SAXPY must add CQL gradient contribution bounded by cql_budget. +#[test] +fn test_cql_clipped_saxpy_respects_budget() -> anyhow::Result<()> { + let dev = cuda_device(); + let stream = Arc::clone(dev.cuda_stream().expect("cuda stream")); + let cfg = test_config(); + let mut trainer = GpuDqnTrainer::new(stream.clone(), cfg)?; + + let total_params = trainer.total_params(); + + // Start with zero grad_buf + stream.memcpy_htod(&vec![0.0_f32; total_params], trainer.grad_buf_mut()) + .map_err(|e| anyhow::anyhow!("{e}"))?; + + // Fill CQL scratch with large gradients (norm >> budget) + let large_cql = vec![1.0_f32; total_params]; + stream.memcpy_htod(&large_cql, trainer.cql_grad_scratch_mut()) + .map_err(|e| anyhow::anyhow!("{e}"))?; + + let cql_budget = 1.5_f32; // Small budget + trainer.apply_cql_clipped_saxpy(cql_budget)?; + + // grad_buf should now have the CQL contribution, clipped to budget + let norm = trainer.read_grad_norm_sync()?; + assert!( + norm <= cql_budget + 0.1, + "CQL contribution should be ≤ budget {cql_budget}, got {norm}" + ); + assert!( + norm > 0.1, + "CQL contribution should be non-zero, got {norm}" + ); + + Ok(()) +} + +/// Dynamic budget: auxiliary fractions must leave room for C51. +/// C51 gets 1.0 - sum(active auxiliary budgets). When all are active: C51=70%. +#[test] +fn test_budget_fractions_leave_room_for_c51() { + // Auxiliary budgets from fused_training.rs + let cql: f32 = 0.15; + let iqn: f32 = 0.10; + let ens: f32 = 0.05; + let aux_total = cql + iqn + ens; + + // C51 gets the remainder + let c51 = 1.0 - aux_total; + assert!(c51 > 0.5, "C51 should get at least 50%, got {c51}"); + assert!( + (c51 + aux_total - 1.0).abs() < 1e-6, + "Total budget must equal 1.0, got {}", c51 + aux_total + ); + + // When no auxiliaries: C51 gets 100% + let c51_solo: f32 = 1.0 - 0.0 - 0.0 - 0.0; + assert!((c51_solo - 1.0).abs() < 1e-6, "Solo C51 should be 1.0"); +} + +/// Combined gradient from all budgeted components must not exceed max_grad_norm. +/// Simulates the full pipeline: C51 clip → CQL SAXPY → check combined ≤ mgn. +#[test] +fn test_combined_gradient_within_max_norm() -> anyhow::Result<()> { + let dev = cuda_device(); + let stream = Arc::clone(dev.cuda_stream().expect("cuda stream")); + let cfg = test_config(); + let max_grad_norm = cfg.max_grad_norm; // 10.0 + let mut trainer = GpuDqnTrainer::new(stream.clone(), cfg)?; + + let total_params = trainer.total_params(); + + // Simulate C51 gradient: large values, will be clipped to 70% budget + let c51_grad = vec![1.0_f32; total_params]; + stream.memcpy_htod(&c51_grad, trainer.grad_buf_mut()) + .map_err(|e| anyhow::anyhow!("{e}"))?; + + let c51_budget = max_grad_norm * 0.70; + trainer.clip_grad_buf_inplace(c51_budget)?; + + let norm_after_c51 = trainer.read_grad_norm_sync()?; + assert!( + norm_after_c51 <= c51_budget + 0.01, + "C51 clip: {norm_after_c51} > budget {c51_budget}" + ); + + // Simulate CQL: large gradient, clipped to 15% budget + let cql_grad = vec![1.0_f32; total_params]; + stream.memcpy_htod(&cql_grad, trainer.cql_grad_scratch_mut()) + .map_err(|e| anyhow::anyhow!("{e}"))?; + + let cql_budget = max_grad_norm * 0.15; + trainer.apply_cql_clipped_saxpy(cql_budget)?; + + let norm_combined = trainer.read_grad_norm_sync()?; + + // Combined should be ≤ c51_budget + cql_budget = 0.85 × max_grad_norm + // (by triangle inequality, worst case is when both point same direction) + assert!( + norm_combined <= (c51_budget + cql_budget) + 0.1, + "Combined C51+CQL norm {norm_combined} exceeds sum of budgets {:.1}", + c51_budget + cql_budget + ); + + // Should be significantly less than max_grad_norm + assert!( + norm_combined < max_grad_norm, + "Combined norm {norm_combined} should be < max_grad_norm {max_grad_norm}" + ); + + Ok(()) +} diff --git a/crates/ml/src/trainers/dqn/smoke_tests/mod.rs b/crates/ml/src/trainers/dqn/smoke_tests/mod.rs index db7986f54..8da5042bf 100644 --- a/crates/ml/src/trainers/dqn/smoke_tests/mod.rs +++ b/crates/ml/src/trainers/dqn/smoke_tests/mod.rs @@ -8,3 +8,5 @@ mod training_stability; mod feature_coverage; #[cfg(test)] mod performance; +#[cfg(test)] +mod gradient_budget; diff --git a/crates/ml/src/trainers/dqn/smoke_tests/training_stability.rs b/crates/ml/src/trainers/dqn/smoke_tests/training_stability.rs index 817249612..4ed248d32 100644 --- a/crates/ml/src/trainers/dqn/smoke_tests/training_stability.rs +++ b/crates/ml/src/trainers/dqn/smoke_tests/training_stability.rs @@ -45,6 +45,44 @@ fn test_production_training_stability() -> anyhow::Result<()> { Ok(()) } +/// Gradient norms must stay bounded across all epochs. +/// +/// Before the per-component clipping fix, grad_norm grew from 89K → 3.6B +/// (40,000x over 8 epochs). This test verifies the fix: grad_norm must +/// stay below 100K throughout training (no exponential growth). +#[test] +#[ignore] // Loads real training data — run via nightly CI or manual trigger +fn test_gradient_norm_bounded_across_epochs() -> anyhow::Result<()> { + let data_dir = test_data_dir() + .expect("FOXHUNT_TEST_DATA or test_data/ must exist"); + + let mut p = smoke_params(); + p.epochs = 3; + + let mut trainer = smoke_trainer_with(p)?; + let rt = tokio::runtime::Builder::new_current_thread() + .enable_all() + .build()?; + let metrics = rt.block_on(trainer.train(&data_dir, |_epoch, _bytes, _best| { + Ok("skip".to_owned()) + }))?; + + let grad_norm = metrics.additional_metrics.get("avg_gradient_norm") + .copied().unwrap_or(f64::NAN); + + assert!(grad_norm.is_finite(), "avg_gradient_norm not finite: {grad_norm}"); + assert!(grad_norm > 0.0, "avg_gradient_norm should be positive: {grad_norm}"); + assert!( + grad_norm < 100_000.0, + "GRADIENT EXPLOSION: avg_gradient_norm {grad_norm:.1} exceeds 100K" + ); + assert_eq!(metrics.epochs_trained, 3, "should complete all 3 epochs"); + + drop(trainer); + drop(rt); + Ok(()) +} + /// PER importance-sampling weights must be finite and positive. #[test] fn test_per_weights_valid() -> anyhow::Result<()> { diff --git a/crates/ml/src/trainers/dqn/trainer/action.rs b/crates/ml/src/trainers/dqn/trainer/action.rs index 23c73c268..800335cba 100644 --- a/crates/ml/src/trainers/dqn/trainer/action.rs +++ b/crates/ml/src/trainers/dqn/trainer/action.rs @@ -158,7 +158,7 @@ impl DQNTrainer { let q_values = agent.forward(state)?; let stream = self.cuda_stream.as_ref().ok_or_else(|| anyhow::anyhow!("CUDA stream required"))?; let indices = q_values.argmax(1, stream).map_err(|e| anyhow::anyhow!("argmax: {e}"))?; - let best_action = indices.first().copied().unwrap_or(2) as usize; + let best_action = *indices.first().expect("argmax on non-empty Q-values") as usize; Ok(best_action) } } diff --git a/crates/ml/src/trainers/dqn/trainer/constructor.rs b/crates/ml/src/trainers/dqn/trainer/constructor.rs index bc4a37537..493029d69 100644 --- a/crates/ml/src/trainers/dqn/trainer/constructor.rs +++ b/crates/ml/src/trainers/dqn/trainer/constructor.rs @@ -311,7 +311,7 @@ impl DQNTrainer { minimum_profit_factor: hyperparams.minimum_profit_factor as f32, weight_decay: hyperparams.weight_decay, dropout_rate: if hyperparams.enable_dropout_scheduler { hyperparams.dropout_initial } else { 0.0 }, - entropy_coefficient: hyperparams.entropy_coefficient.unwrap_or(0.01), + entropy_coefficient: hyperparams.entropy_coefficient, noisy_epsilon_floor: hyperparams.noisy_epsilon_floor.unwrap_or(0.0) as f32, // C2: NoisyNet handles exploration use_count_bonus: hyperparams.count_bonus_coefficient.unwrap_or(0.0) > 0.0, // C3 FIX: enable when coefficient > 0 count_bonus_coefficient: hyperparams.count_bonus_coefficient.unwrap_or(0.0), @@ -396,7 +396,7 @@ impl DQNTrainer { // Entropy regularization: SAC-style computed directly on Q-values in DQN::compute_loss_internal if hyperparams.enable_entropy_regularization { - let coeff = hyperparams.entropy_coefficient.unwrap_or(0.01); + let coeff = hyperparams.entropy_coefficient; info!("Entropy regularization enabled (coefficient={coeff:.4}, applied to Q-value softmax in loss)"); } diff --git a/crates/ml/src/trainers/dqn/trainer/state.rs b/crates/ml/src/trainers/dqn/trainer/state.rs index 1831fcc1c..1b1892a16 100644 --- a/crates/ml/src/trainers/dqn/trainer/state.rs +++ b/crates/ml/src/trainers/dqn/trainer/state.rs @@ -70,7 +70,7 @@ impl DQNTrainer { // BUG #36 FIX: Use NORMALIZED portfolio features to prevent Q-value explosion let portfolio_features = if let Some(price) = close_price { - let price_f32 = price.to_f32().unwrap_or(0.0); + let price_f32 = price.to_f32().expect("trading price fits f32"); self.portfolio_tracker .get_portfolio_features(price_f32) .to_vec() // cpu-side portfolio features diff --git a/crates/ml/src/trainers/dqn/trainer/training_loop.rs b/crates/ml/src/trainers/dqn/trainer/training_loop.rs index 5f73d2eea..a4525962c 100644 --- a/crates/ml/src/trainers/dqn/trainer/training_loop.rs +++ b/crates/ml/src/trainers/dqn/trainer/training_loop.rs @@ -1990,6 +1990,11 @@ impl DQNTrainer { } training_metrics::set_q_value_stats("dqn", "current", q_mean, q_max); training_metrics::set_gradient_norm("dqn", "current", avg_grad_norm); + if let Some(ref fused) = self.fused_ctx { + let (c51, combined) = fused.grad_norm_diagnostics(); + training_metrics::set_grad_norm_c51("dqn", "current", c51 as f64); + training_metrics::set_grad_norm_combined("dqn", "current", combined as f64); + } training_metrics::set_epoch_duration("dqn", "current", epoch_duration.as_secs_f64()); { let agent = self.agent.read().await; diff --git a/crates/ml/src/training_profile.rs b/crates/ml/src/training_profile.rs index 635e682f0..809160cee 100644 --- a/crates/ml/src/training_profile.rs +++ b/crates/ml/src/training_profile.rs @@ -657,7 +657,7 @@ impl DqnTrainingProfile { hp.noisy_sigma_init = v; } if let Some(v) = e.entropy_coefficient { - hp.entropy_coefficient = Some(v); + hp.entropy_coefficient = v; } if let Some(v) = e.count_bonus_coefficient { hp.count_bonus_coefficient = Some(v); @@ -1160,10 +1160,10 @@ mod tests { "noisy_sigma_init should be 0.5 from smoketest [exploration], got {}", hp.noisy_sigma_init ); - assert_eq!( - hp.entropy_coefficient, - Some(0.001), - "entropy_coefficient should be Some(0.001) from smoketest [exploration]" + assert!( + (hp.entropy_coefficient - 0.001).abs() < 1e-9, + "entropy_coefficient should be 0.001 from smoketest [exploration], got {}", + hp.entropy_coefficient ); assert_eq!( hp.count_bonus_coefficient, @@ -1209,7 +1209,7 @@ mod tests { // [exploration] section — new fields assert!((hp.noisy_sigma_init - 0.5).abs() < 0.01); - assert_eq!(hp.entropy_coefficient, Some(0.001)); + assert!((hp.entropy_coefficient - 0.001).abs() < 1e-9); assert_eq!(hp.count_bonus_coefficient, Some(0.1)); // [risk] section — new loss_aversion field diff --git a/docs/superpowers/plans/2026-03-27-bug-fixes.md b/docs/superpowers/plans/2026-03-27-bug-fixes.md new file mode 100644 index 000000000..05f746517 --- /dev/null +++ b/docs/superpowers/plans/2026-03-27-bug-fixes.md @@ -0,0 +1,501 @@ +# Bug Fixes Implementation Plan (Spec B) + +> **For agentic workers:** REQUIRED SUB-SKILL: Use superpowers:subagent-driven-development (recommended) or superpowers:executing-plans to implement this plan task-by-task. Steps use checkbox (`- [ ]`) syntax for tracking. + +**Goal:** Fix all 9 CRITICAL bugs and top-priority IMPORTANT issues found by the 9-agent audit, eliminating wrong training results and silent error suppression. + +**Architecture:** Config-only fixes first (lowest risk), then Rust refactors, then CUDA kernel changes last. Each task is independently testable and committable. Order follows the spec's Phase 1 → Phase 2 risk-based sequencing. + +**Tech Stack:** Rust, CUDA (.cu kernels compiled by build.rs), cudarc 0.19.3, serde + +**Spec:** `docs/superpowers/specs/2026-03-27-bug-fixes-design.md` + +--- + +## File Map + +| File | Bugs | Action | +|---|---|---| +| `crates/ml/src/trainers/dqn/config.rs` | #4, #5 | Make gradient_clip_norm and entropy_coefficient non-optional | +| `crates/ml/src/trainers/dqn/fused_training.rs` | #4 | Remove unwrap_or on gradient_clip_norm | +| `crates/ml/src/trainers/dqn/trainer/constructor.rs` | #4, #5 | Remove unwrap_or fallbacks | +| `crates/ml/src/trainers/dqn/data_loading.rs` | #3 | Replace unwrap_or(1.0) with direct indexing | +| `crates/ml-dqn/src/network.rs` | #8 | Conditional step_count increment | +| `crates/ml-dqn/src/quantile_regression.rs` | #6, #7 | Single CUDA context, remove silent fallbacks | +| `crates/ml/src/cuda_pipeline/backtest_env_kernel.cu` | #1, #2 | Complete episode reset, update max_equity | +| `crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs` | #10 | Export budget fraction constants | +| `crates/ml-dqn/src/branching.rs` | #16 | Remove MaybeNoisyLinear enum | +| `crates/ml/src/trainers/dqn/trainer/action.rs` | #15 | Replace unwrap_or(2) with expect | +| `crates/ml/src/trainers/dqn/trainer/state.rs` | #15 | Replace unwrap_or(0.0) with expect | +| `crates/ml/src/cuda_pipeline/gpu_action_selector.rs` | #9 | Add NoisyNet architecture comment | + +--- + +### Task 1: Unify gradient_clip_norm — eliminate inconsistent defaults + +Bug #4: Main trainer clips at 10.0, IQN/IQL clip at 1.0. The fix: resolve `gradient_clip_norm` once at the top of `FusedTrainingCtx::new()` and pass the resolved value to ALL sub-configs. + +**Files:** +- Modify: `crates/ml/src/trainers/dqn/fused_training.rs` + +- [ ] **Step 1: Resolve gradient_clip_norm once at top of constructor** + +At the top of `FusedTrainingCtx::new()` (before any config construction), add: + +```rust +let resolved_grad_norm = hyperparams.gradient_clip_norm.unwrap_or(10.0); +``` + +Then replace all 3 sites: +- Line 194: `max_grad_norm: hyperparams.gradient_clip_norm.unwrap_or(10.0) as f32,` → `max_grad_norm: resolved_grad_norm as f32,` +- Line 263: `max_grad_norm: hyperparams.gradient_clip_norm.unwrap_or(1.0) as f32,` → `max_grad_norm: resolved_grad_norm as f32,` +- Line 298: `max_grad_norm: hyperparams.gradient_clip_norm.unwrap_or(1.0) as f32,` → `max_grad_norm: resolved_grad_norm as f32,` + +- [ ] **Step 2: Verify compilation + tests** + +Run: `SQLX_OFFLINE=true cargo check -p ml && SQLX_OFFLINE=true cargo test -p ml --lib -- dqn 2>&1 | tail -5` +Expected: Compiles, all tests pass. + +- [ ] **Step 3: Commit** + +```bash +git add crates/ml/src/trainers/dqn/fused_training.rs +git commit -m "fix: unify gradient_clip_norm — resolve once, pass to all sub-trainers (was 10.0 vs 1.0)" +``` + +--- + +### Task 2: Make entropy_coefficient non-optional + +Bug #5: `Option` with 3 different `unwrap_or` defaults (0.001, 0.01, 0.02). Fix: make it non-optional with a single canonical default. + +**Files:** +- Modify: `crates/ml/src/trainers/dqn/config.rs` +- Modify: `crates/ml/src/trainers/dqn/trainer/constructor.rs` + +- [ ] **Step 1: Change type in DQNHyperparameters** + +In `config.rs` (~line 960), change: +```rust +pub entropy_coefficient: Option, +``` +To: +```rust +#[serde(default = "default_entropy_coefficient")] +pub entropy_coefficient: f64, +``` + +Add the default function nearby: +```rust +fn default_entropy_coefficient() -> f64 { 0.001 } +``` + +In the `conservative()` and `default()` impl, change: +```rust +entropy_coefficient: Some(0.001), +``` +To: +```rust +entropy_coefficient: 0.001, +``` + +- [ ] **Step 2: Remove unwrap_or in constructor.rs** + +Line 314: `entropy_coefficient: hyperparams.entropy_coefficient.unwrap_or(0.01),` → `entropy_coefficient: hyperparams.entropy_coefficient,` + +Line 399: `let coeff = hyperparams.entropy_coefficient.unwrap_or(0.01);` → `let coeff = hyperparams.entropy_coefficient;` + +- [ ] **Step 3: Fix all compilation errors from type change** + +Run `SQLX_OFFLINE=true cargo check --workspace 2>&1 | head -40` — fix any sites that still use `Some(...)` or `.unwrap()` on entropy_coefficient. Common patterns: +- `Some(0.001)` → `0.001` +- `.unwrap_or(x)` → remove +- Pattern match on `Some(v)` → use directly + +- [ ] **Step 4: Verify compilation + tests** + +Run: `SQLX_OFFLINE=true cargo check --workspace && SQLX_OFFLINE=true cargo test -p ml --lib -- dqn && SQLX_OFFLINE=true cargo test -p ml-dqn --lib 2>&1 | tail -5` +Expected: All pass. + +- [ ] **Step 5: Commit** + +```bash +git add crates/ml/src/trainers/dqn/config.rs crates/ml/src/trainers/dqn/trainer/constructor.rs +git commit -m "fix: make entropy_coefficient non-optional — single canonical default 0.001" +``` + +--- + +### Task 3: Replace close price unwrap_or(1.0) with direct indexing + +Bug #3: Silent wrong rewards when price parse fails. The indices are always in-bounds (loop from 1..n), so direct indexing is safe and surfaces any future bugs. + +**Files:** +- Modify: `crates/ml/src/trainers/dqn/data_loading.rs` + +- [ ] **Step 1: Find and replace the unwrap_or(1.0)** + +At line 306, change: +```rust +let prev = close_prices_f32.get(i - 1).copied().unwrap_or(1.0); +let curr = close_prices_f32.get(i).copied().unwrap_or(prev); +``` +To: +```rust +let prev = close_prices_f32[i - 1]; +let curr = close_prices_f32[i]; +``` + +- [ ] **Step 2: Verify compilation + tests** + +Run: `SQLX_OFFLINE=true cargo check -p ml && SQLX_OFFLINE=true cargo test -p ml --lib -- dqn 2>&1 | tail -5` + +- [ ] **Step 3: Commit** + +```bash +git add crates/ml/src/trainers/dqn/data_loading.rs +git commit -m "fix: replace close price unwrap_or(1.0) with direct indexing — no silent wrong rewards" +``` + +--- + +### Task 4: Fix step_count incrementing during inference + +Bug #8: `forward()` unconditionally increments step_count, advancing dropout schedule during eval. + +**Files:** +- Modify: `crates/ml-dqn/src/network.rs` + +- [ ] **Step 1: Read current forward() to understand full context** + +Read `network.rs` around line 245-260 to see how `step_count` and `dropout_scheduler` are used. + +- [ ] **Step 2: Add `training: bool` parameter to forward()** + +Change the forward method to accept a training flag. Only increment step_count and advance dropout scheduler when training=true. Update the method signature and all call sites. + +If there's a trait constraint, add a `forward_eval` convenience method that calls `forward(state, false)`. + +- [ ] **Step 3: Update all call sites** + +Search for `.forward(` across the codebase. Inference sites (backtest, action selection) pass `false`. Training sites pass `true`. + +- [ ] **Step 4: Verify compilation + tests** + +Run: `SQLX_OFFLINE=true cargo check --workspace && SQLX_OFFLINE=true cargo test -p ml-dqn --lib 2>&1 | tail -5` + +- [ ] **Step 5: Commit** + +```bash +git add crates/ml-dqn/src/network.rs +git commit -m "fix: step_count only increments during training — no inference side effects" +``` + +--- + +### Task 5: Single CUDA context in quantile loss + +Bug #6: Creates CudaContext 3× per call. Refactor to create once and pass through. + +**Files:** +- Modify: `crates/ml-dqn/src/quantile_regression.rs` + +- [ ] **Step 1: Read the full quantile_huber_loss function** + +Understand the 3 context creation sites and the inner `_per_sample` function. + +- [ ] **Step 2: Refactor to single context** + +Create context + stream once at the top. Pass the stream to all operations that need it. If `to_host()` creates its own context internally, check if there's a `to_host_with_stream()` variant or use manual `memcpy_dtoh`. + +- [ ] **Step 3: Verify compilation + tests** + +Run: `SQLX_OFFLINE=true cargo check -p ml-dqn && SQLX_OFFLINE=true cargo test -p ml-dqn --lib 2>&1 | tail -5` + +- [ ] **Step 4: Commit** + +```bash +git add crates/ml-dqn/src/quantile_regression.rs +git commit -m "fix: single CUDA context per quantile loss call — was 3× context creation" +``` + +--- + +### Task 6: Remove silent fallbacks in quantile loss + +Bug #7: `unwrap_or(0.0)` / `unwrap_or(0.5)` on tensor indexing hides bugs. + +**Files:** +- Modify: `crates/ml-dqn/src/quantile_regression.rs` + +- [ ] **Step 1: Replace unwrap_or with direct indexing** + +At line ~385-387, change: +```rust +let pred_val = pred_host.get(idx).copied().unwrap_or(0.0); +let tgt_val = tgt_host.get(idx).copied().unwrap_or(0.0); +let tau_val = tau_host.get(idx).copied().unwrap_or(0.5); +``` +To: +```rust +let pred_val = pred_host[idx]; +let tgt_val = tgt_host[idx]; +let tau_val = tau_host[idx]; +``` + +- [ ] **Step 2: Verify compilation + tests** + +Run: `SQLX_OFFLINE=true cargo test -p ml-dqn --lib 2>&1 | tail -5` + +- [ ] **Step 3: Commit** + +```bash +git add crates/ml-dqn/src/quantile_regression.rs +git commit -m "fix: remove silent 0.0/0.5 fallbacks in quantile loss — panic on OOB" +``` + +--- + +### Task 7: Complete backtest episode reset + fix stale max_equity + +Bugs #1 + #2: The most impactful CUDA kernel changes. Reset all 8 portfolio fields, update max_equity before floor check. + +**Files:** +- Modify: `crates/ml/src/cuda_pipeline/backtest_env_kernel.cu` + +- [ ] **Step 1: Fix pre-trade floor reset (lines 105-112)** + +Replace: +```cuda +portfolio_state[ps + 0] = liq_value; +portfolio_state[ps + 1] = 0.0f; +portfolio_state[ps + 2] = liq_value; +portfolio_state[ps + 3] = 0.0f; +portfolio_state[ps + 4] = max_equity; +portfolio_state[ps + 5] = 0.0f; +portfolio_state[ps + 6] = cum_return + liq_ret; +portfolio_state[ps + 7] += 1.0f; +``` +With: +```cuda +portfolio_state[ps + 0] = liq_value; // value = new initial capital +portfolio_state[ps + 1] = 0.0f; // position = flat +portfolio_state[ps + 2] = liq_value; // cash = new initial capital +portfolio_state[ps + 3] = 0.0f; // entry_price = none +portfolio_state[ps + 4] = liq_value; // max_equity = RESET (was stale) +portfolio_state[ps + 5] = 0.0f; // hold_time = 0 +portfolio_state[ps + 6] = 0.0f; // cum_return = RESET (was accumulating) +portfolio_state[ps + 7] = 0.0f; // step_count = RESET (was incrementing) +``` + +- [ ] **Step 2: Fix max_equity update before post-trade floor check (line ~176-184)** + +After `float new_value = cash + position * close;` (line 176), insert BEFORE the floor check: +```cuda +// Update max_equity BEFORE floor check — prevents stale peak +max_equity = fmaxf(max_equity, new_value); +``` + +- [ ] **Step 3: Fix post-trade floor reset (lines 199-206)** + +Replace with complete reset (same pattern as pre-trade): +```cuda +portfolio_state[ps + 0] = new_value; +portfolio_state[ps + 1] = 0.0f; +portfolio_state[ps + 2] = new_value; +portfolio_state[ps + 3] = 0.0f; +portfolio_state[ps + 4] = new_value; // max_equity = RESET +portfolio_state[ps + 5] = 0.0f; +portfolio_state[ps + 6] = 0.0f; // cum_return = RESET +portfolio_state[ps + 7] = 0.0f; // step_count = RESET +``` + +- [ ] **Step 4: Verify build (cubins recompile via build.rs)** + +Run: `SQLX_OFFLINE=true cargo check -p ml 2>&1 | tail -5` +Expected: build.rs recompiles the .cu file, check passes. + +- [ ] **Step 5: Run tests** + +Run: `SQLX_OFFLINE=true cargo test -p ml --lib -- dqn 2>&1 | tail -5` + +- [ ] **Step 6: Commit** + +```bash +git add crates/ml/src/cuda_pipeline/backtest_env_kernel.cu +git commit -m "fix: complete backtest episode reset (all 8 fields) + update max_equity before floor check" +``` + +--- + +### Task 8: Export budget fraction constants + fix hardcoded values + +Bug #10: `0.10` and `0.05` hardcoded in gpu_dqn_trainer.rs should reference the constants from fused_training.rs. + +**Files:** +- Modify: `crates/ml/src/trainers/dqn/fused_training.rs` +- Modify: `crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs` + +- [ ] **Step 1: Make budget constants public** + +In `fused_training.rs`, change: +```rust +const CQL_GRAD_BUDGET: f32 = 0.15; +const IQN_GRAD_BUDGET: f32 = 0.10; +const ENS_GRAD_BUDGET: f32 = 0.05; +``` +To: +```rust +pub(crate) const CQL_GRAD_BUDGET: f32 = 0.15; +pub(crate) const IQN_GRAD_BUDGET: f32 = 0.10; +pub(crate) const ENS_GRAD_BUDGET: f32 = 0.05; +``` + +- [ ] **Step 2: Reference constants in gpu_dqn_trainer.rs** + +At line ~831 in `apply_iqn_trunk_gradient`, change: +```rust +let max_component_norm = self.config.max_grad_norm * 0.10; // IQN_GRAD_BUDGET +``` +To: +```rust +let max_component_norm = self.config.max_grad_norm * crate::trainers::dqn::fused_training::IQN_GRAD_BUDGET; +``` + +At line ~1072 in `apply_ensemble_trunk_gradient`, change: +```rust +let max_component_norm = self.config.max_grad_norm * 0.05; // ENS_GRAD_BUDGET +``` +To: +```rust +let max_component_norm = self.config.max_grad_norm * crate::trainers::dqn::fused_training::ENS_GRAD_BUDGET; +``` + +- [ ] **Step 3: Verify compilation + tests** + +Run: `SQLX_OFFLINE=true cargo check -p ml && SQLX_OFFLINE=true cargo test -p ml --lib -- dqn 2>&1 | tail -5` + +- [ ] **Step 4: Commit** + +```bash +git add crates/ml/src/trainers/dqn/fused_training.rs crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs +git commit -m "fix: reference budget fraction constants — eliminate hardcoded 0.10/0.05 in gpu_dqn_trainer" +``` + +--- + +### Task 9: Remove MaybeNoisyLinear single-variant enum + +Bug #16: Enum with only `Noisy` variant — unnecessary indirection. + +**Files:** +- Modify: `crates/ml-dqn/src/branching.rs` + +- [ ] **Step 1: Read current MaybeNoisyLinear usage** + +Search for all `MaybeNoisyLinear` references in branching.rs to understand the scope of changes. + +- [ ] **Step 2: Replace enum with direct NoisyLinear** + +- Remove the `MaybeNoisyLinear` enum definition +- Change all struct fields from `MaybeNoisyLinear` to `NoisyLinear` +- Remove all `let MaybeNoisyLinear::Noisy(n) = self;` destructuring — use the field directly +- Inline the thin wrapper methods (forward, reset_noise, etc.) or call directly on NoisyLinear + +- [ ] **Step 3: Verify compilation + tests** + +Run: `SQLX_OFFLINE=true cargo check --workspace && SQLX_OFFLINE=true cargo test -p ml-dqn --lib 2>&1 | tail -5` + +- [ ] **Step 4: Commit** + +```bash +git add crates/ml-dqn/src/branching.rs +git commit -m "fix: remove MaybeNoisyLinear single-variant enum — use NoisyLinear directly" +``` + +--- + +### Task 10: Fix Tier 1 silent unwrap_or fallbacks + +Bug #15: Critical silent fallbacks in action selection and state processing. + +**Files:** +- Modify: `crates/ml/src/trainers/dqn/trainer/action.rs` +- Modify: `crates/ml/src/trainers/dqn/trainer/state.rs` + +- [ ] **Step 1: Fix action.rs:161 — argmax fallback** + +Change `unwrap_or(2)` (silent Flat action) to `.expect("argmax on non-empty Q-values")`. + +- [ ] **Step 2: Fix state.rs:73 — price to_f32 fallback** + +Change `unwrap_or(0.0)` to `.expect("trading price fits f32")`. + +- [ ] **Step 3: Verify compilation + tests** + +Run: `SQLX_OFFLINE=true cargo check -p ml && SQLX_OFFLINE=true cargo test -p ml --lib -- dqn 2>&1 | tail -5` + +- [ ] **Step 4: Commit** + +```bash +git add crates/ml/src/trainers/dqn/trainer/action.rs crates/ml/src/trainers/dqn/trainer/state.rs +git commit -m "fix: replace silent unwrap_or fallbacks in action selection and state processing" +``` + +--- + +### Task 11: Add NoisyNet architecture comment to action selector + +Bug #9 (NOT A BUG): Document the architectural choice so future audits don't re-flag this. + +**Files:** +- Modify: `crates/ml/src/cuda_pipeline/gpu_action_selector.rs` + +- [ ] **Step 1: Add comment at top of file** + +After the module doc comment, add: +```rust +// NoisyNet exploration note: Q-values received by the action selector already +// include factorized Gaussian noise from NoisyLinear layers in the network +// forward pass. Epsilon-greedy acts as a secondary fallback, not the primary +// exploration mechanism when NoisyNet is enabled. Explicit noise injection +// here would double-count exploration. +``` + +- [ ] **Step 2: Commit** + +```bash +git add crates/ml/src/cuda_pipeline/gpu_action_selector.rs +git commit -m "docs: clarify NoisyNet exploration is in network forward, not action selector" +``` + +--- + +## Execution Dependencies + +``` +Task 1 (gradient_clip_norm) — independent +Task 2 (entropy_coefficient) — independent +Task 3 (close price) — independent +Task 4 (step_count) — independent +Task 5 (CUDA context) — independent +Task 6 (quantile fallbacks) — after Task 5 (same file) +Task 7 (backtest reset) — independent +Task 8 (budget constants) — independent +Task 9 (MaybeNoisyLinear) — independent +Task 10 (unwrap_or fixes) — independent +Task 11 (NoisyNet comment) — independent +``` + +Tasks 5-6 must be sequential (same file). All others are independent and can run in parallel. + +## Verification + +After ALL tasks: +```bash +SQLX_OFFLINE=true cargo check --workspace +SQLX_OFFLINE=true cargo test -p ml --lib -- dqn +SQLX_OFFLINE=true cargo test -p ml-dqn --lib +``` + +All must pass with zero failures. diff --git a/docs/superpowers/plans/2026-03-27-gpu-hot-path-performance.md b/docs/superpowers/plans/2026-03-27-gpu-hot-path-performance.md new file mode 100644 index 000000000..2cc1324bf --- /dev/null +++ b/docs/superpowers/plans/2026-03-27-gpu-hot-path-performance.md @@ -0,0 +1,732 @@ +# GPU Hot Path Performance — Zero-Sync Training Step + +> **For agentic workers:** REQUIRED SUB-SKILL: Use superpowers:subagent-driven-development (recommended) or superpowers:executing-plans to implement this plan task-by-task. Steps use checkbox (`- [ ]`) syntax for tracking. + +**Goal:** Eliminate all CPU sync points and F32 compute waste from the per-step training hot path, restoring 3s/epoch performance on H100. + +**Architecture:** Cache all CUDA device pointers at construction/graph-capture time into a `CachedPtrs` struct. Replace per-step `cuStreamSynchronize` + `EventTrackingGuard` + `raw_device_ptr` calls with cached u64 pointer lookups. Convert cuBLAS `cublasSgemm` (F32) to `cublasGemmEx` (BF16 compute with F32 accumulate) for 2-3× tensor core throughput. Batch 8 head spectral norm launches into a single multi-matrix kernel. + +**Tech Stack:** Rust, CUDA, cudarc 0.19.3, cuBLAS (cublasGemmEx), BF16 tensor cores + +**Verification requirement:** Every task MUST compile (`SQLX_OFFLINE=true cargo check -p ml`) and pass tests (`SQLX_OFFLINE=true cargo test -p ml --lib -- dqn`). No exceptions. No skipping. + +--- + +## Root Cause Analysis + +**Measured:** Epoch time regressed from 3s → 30s. Per-step hot path has: + +| Issue | Per-step cost | Source | Lines | +|---|---|---|---| +| `cuStreamSynchronize` in `apply_iqn_trunk_gradient` | ~50μs (stalls entire GPU pipeline) | Pre-existing | `gpu_dqn_trainer.rs:708` | +| `cuStreamSynchronize` in `apply_ensemble_trunk_gradient` | ~50μs | Pre-existing | `gpu_dqn_trainer.rs:911` | +| `cuStreamSynchronize` in `run_ensemble_step` | ~50μs | Pre-existing | `fused_training.rs:875` | +| `cuStreamSynchronize` in `replay_adam_and_readback` | ~50μs | Pre-existing | `gpu_dqn_trainer.rs:2369` | +| 99× `raw_device_ptr()` calls | ~1μs each but forces `device_ptr()` event machinery | Per-step | Throughout | +| 11× `EventTrackingGuard` create/drop | ~0.1μs each (atomic ops) | Per-step | Throughout | +| F32 `cublasSgemm` instead of BF16 `cublasGemmEx` | 2-3× slower SGEMM on H100 | Pre-existing | `batched_forward.rs:513` | +| 8× single-block spectral norm launches | ~5μs each | New | `gpu_dqn_trainer.rs:1439-1453` | +| 10× DtoD spectral norm sync-back | ~2μs each | New | `gpu_dqn_trainer.rs:1490-1503` | + +**Total per-step overhead: ~200μs sync stalls + 2-3× slower SGEMM.** Over 10K steps/epoch: 2s sync + 2-3× slower forward/backward = easily 30s epochs. + +--- + +## File Map + +| File | Action | +|---|---| +| `crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs` | Major: add `CachedPtrs` struct, remove syncs from IQN/ensemble/CQL paths, cache all pointers at capture time | +| `crates/ml/src/cuda_pipeline/batched_forward.rs` | Major: replace `cublasSgemm` with `cublasGemmEx` (BF16 input, F32 accumulate) | +| `crates/ml/src/cuda_pipeline/batched_backward.rs` | Major: same `cublasSgemm` → `cublasGemmEx` conversion | +| `crates/ml/src/trainers/dqn/fused_training.rs` | Moderate: remove `cuStreamSynchronize` from `run_ensemble_step`, dynamic budget fix | +| `crates/ml/src/cuda_pipeline/dqn_utility_kernels.cu` | Minor: add batched spectral norm kernel | +| `crates/ml/src/trainers/dqn/smoke_tests/gradient_budget.rs` | Minor: add performance regression test | + +--- + +### Task 1: Cache all device pointers at graph-capture time + +The root cause of most overhead: `raw_device_ptr()` calls `slice.device_ptr(stream)` which goes through cudarc's event tracking machinery. Even with events disabled, it creates a `SyncOnDrop` guard that's immediately leaked via `ManuallyDrop`. This happens 99 times per step. + +**Fix:** Compute all u64 pointers ONCE at construction (or re-capture), store them in a `CachedPtrs` struct, and use the cached values in all per-step methods. + +**Files:** +- Modify: `crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs` + +- [ ] **Step 1: Define `CachedPtrs` struct** + +After the `GpuDqnTrainConfig` struct (~line 219), add: + +```rust +/// Pre-computed raw CUDA device pointers for all buffers used in the per-step hot path. +/// Eliminates 99 `raw_device_ptr()` → `device_ptr()` calls per training step. +/// Pointers are invalidated when buffers are reallocated (which never happens after init). +struct CachedPtrs { + // Flat parameter buffers + params: u64, + target_params: u64, + grad: u64, + grad_norm: u64, + m: u64, + v: u64, + t: u64, + total_loss: u64, + cql_grad_scratch: u64, + + // Batch input buffers + states: u64, + next_states: u64, + actions: u64, + rewards: u64, + dones: u64, + is_weights: u64, + + // Activation save buffers + save_h_s1: u64, + save_h_s2: u64, + save_h_v: u64, + save_h_b0: u64, + save_h_b1: u64, + save_h_b2: u64, + + // Backward scratch + bw_d_h_s1: u64, + bw_d_h_s2: u64, + bw_d_h_v: u64, + bw_d_h_b0: u64, + bw_d_h_b1: u64, + bw_d_h_b2: u64, + + // IQN trunk scratch + iqn_trunk_m: u64, + iqn_trunk_grad_norm: u64, + + // Loss outputs + td_errors: u64, + + // Per-weight pointers (20 online + 20 target, computed from params base) + on_weight_ptrs: [u64; 20], + tg_weight_ptrs: [u64; 20], +} +``` + +- [ ] **Step 2: Build `CachedPtrs` in constructor** + +At the end of `new()`, before `Ok(Self { ... })`, compute all pointers. Since event tracking is enabled at this point and all buffers were just allocated on the same stream, `device_ptr()` is safe: + +```rust + let cached_ptrs = { + let _evt_guard = EventTrackingGuard::new(stream.context()); + let param_sizes_arr = compute_param_sizes(&config); + CachedPtrs { + params: raw_device_ptr(¶ms_buf, &stream), + target_params: raw_device_ptr(&target_params_buf, &stream), + grad: raw_device_ptr(&grad_buf, &stream), + grad_norm: raw_device_ptr(&grad_norm_buf, &stream), + m: raw_device_ptr(&m_buf, &stream), + v: raw_device_ptr(&v_buf, &stream), + t: raw_device_ptr_i32(&t_buf, &stream), + total_loss: raw_device_ptr(&total_loss_buf, &stream), + cql_grad_scratch: raw_device_ptr(&cql_grad_scratch, &stream), + states: raw_device_ptr(&states_buf, &stream), + next_states: raw_device_ptr(&next_states_buf, &stream), + actions: raw_device_ptr_i32(&actions_buf, &stream), + rewards: raw_device_ptr(&rewards_buf, &stream), + dones: raw_device_ptr(&dones_buf, &stream), + is_weights: raw_device_ptr(&is_weights_buf, &stream), + save_h_s1: raw_device_ptr(&save_h_s1, &stream), + save_h_s2: raw_device_ptr(&save_h_s2, &stream), + save_h_v: raw_device_ptr(&save_h_v, &stream), + save_h_b0: raw_device_ptr(&save_h_b0, &stream), + save_h_b1: raw_device_ptr(&save_h_b1, &stream), + save_h_b2: raw_device_ptr(&save_h_b2, &stream), + bw_d_h_s1: raw_device_ptr(&bw_d_h_s1, &stream), + bw_d_h_s2: raw_device_ptr(&bw_d_h_s2, &stream), + bw_d_h_v: raw_device_ptr(&bw_d_h_v, &stream), + bw_d_h_b0: raw_device_ptr(&bw_d_h_b0, &stream), + bw_d_h_b1: raw_device_ptr(&bw_d_h_b1, &stream), + bw_d_h_b2: raw_device_ptr(&bw_d_h_b2, &stream), + iqn_trunk_m: raw_device_ptr(&iqn_trunk_m, &stream), + iqn_trunk_grad_norm: raw_device_ptr(&iqn_trunk_grad_norm, &stream), + td_errors: raw_device_ptr(&td_errors_buf, &stream), + on_weight_ptrs: f32_weight_ptrs(¶ms_buf, ¶m_sizes_arr, &stream), + tg_weight_ptrs: f32_weight_ptrs(&target_params_buf, ¶m_sizes_arr, &stream), + } + }; +``` + +Add `cached_ptrs: CachedPtrs` to the struct and `Ok(Self { ... })`. + +- [ ] **Step 3: Replace `raw_device_ptr` calls in `apply_iqn_trunk_gradient`** + +Replace all `raw_device_ptr(&self.xxx, &self.stream)` with `self.cached_ptrs.xxx`. Remove the `cuStreamSynchronize` (line 708) and `EventTrackingGuard` (line 712) — they were only needed because `device_ptr()` could trigger event validation. With cached pointers, no events are involved. + +**CRITICAL:** The `cuStreamSynchronize` at line 708 has the comment "ensure graph_forward replay completed before we touch buffers." With cached pointers, we don't "touch" buffers via device_ptr — we just use pre-computed addresses. All kernel launches go to the same stream, so ordering is guaranteed by CUDA stream semantics. The sync is unnecessary. + +- [ ] **Step 4: Replace `raw_device_ptr` calls in `apply_ensemble_trunk_gradient`** + +Same pattern. Remove the `cuStreamSynchronize` (line 911) and `EventTrackingGuard` (line 914). + +- [ ] **Step 5: Replace `raw_device_ptr` calls in `apply_cql_gradient` and `apply_cql_clipped_saxpy`** + +Same pattern for both methods. + +- [ ] **Step 6: Replace `raw_device_ptr` calls in `clip_grad_buf_inplace` and `read_grad_norm_sync`** + +For `clip_grad_buf_inplace`: use `self.cached_ptrs.grad` and `self.cached_ptrs.grad_norm`. Remove `EventTrackingGuard`. + +For `read_grad_norm_sync`: keep the `cuStreamSynchronize` (this is a diagnostic method that explicitly syncs for readback) but use cached pointers. This method only runs every 1000 steps. + +- [ ] **Step 7: Replace `raw_device_ptr` calls in `apply_spectral_norm`** + +Use `self.cached_ptrs.params` for `params_base` in the sync-back section. The spectral norm kernel launches on individual weight set buffers (DuelingWeightSet/BranchingWeightSet) — these are passed as references, so their `device_ptr()` calls remain. BUT: spectral norm runs ONCE per step before graph_forward. The sync at the method start is the graph_forward guard — remove it and rely on stream ordering. + +- [ ] **Step 8: Replace `raw_device_ptr` in kernel launch methods** + +Update `launch_grad_norm`, `launch_adam_update`, and the SAXPY kernel launches (IQN clipped, ensemble clipped) to use cached pointers. Remove the `EventTrackingGuard` in each. + +- [ ] **Step 9: Verify compilation + tests** + +Run: `SQLX_OFFLINE=true cargo check -p ml && SQLX_OFFLINE=true cargo test -p ml --lib -- dqn` +Expected: Compiles clean, all 129+ tests pass. + +- [ ] **Step 10: Commit** + +```bash +git add crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs +git commit -m "perf: cache all CUDA device pointers — eliminate 99 raw_device_ptr + 4 cuStreamSync per step" +``` + +--- + +### Task 2: Remove `cuStreamSynchronize` from `run_ensemble_step` + +The ensemble step in `fused_training.rs:875` syncs before touching `save_h_s2`. With cached pointers and single-stream execution, this sync is unnecessary — CUDA guarantees ordering on the same stream. + +**Files:** +- Modify: `crates/ml/src/trainers/dqn/fused_training.rs` + +- [ ] **Step 1: Remove the sync + EventTrackingGuard in `run_ensemble_step`** + +At ~line 874-876, remove: +```rust +unsafe { cudarc::driver::sys::cuStreamSynchronize(self.stream.cu_stream()); } +let _ = self.stream.context().check_err(); +``` + +And the local `EvtGuard` struct + its creation (~lines 881-893). The ensemble heads use `device_ptr()` on their own weight buffers — since these are created after event tracking is disabled by the top-level guard in `apply_spectral_norm`, they have no events. Use raw u64 pointers from the ensemble head structs instead. + +Note: the `device_ptr()` calls in the ensemble loop (on `ensemble_extra_heads[k]` weights) cannot be cached in `CachedPtrs` because ensemble heads are optional and variable-count. For these, use a SINGLE `EventTrackingGuard` scope at the top of `run_ensemble_step` and call `device_ptr()` within it — but do NOT sync the stream. + +- [ ] **Step 2: Verify compilation + tests** + +Run: `SQLX_OFFLINE=true cargo check -p ml && SQLX_OFFLINE=true cargo test -p ml --lib -- dqn` +Expected: All pass. + +- [ ] **Step 3: Commit** + +```bash +git add crates/ml/src/trainers/dqn/fused_training.rs +git commit -m "perf: remove cuStreamSynchronize from run_ensemble_step" +``` + +--- + +### Task 3: Convert cuBLAS forward from `cublasSgemm` (F32) to `cublasGemmEx` (BF16 tensors, F32 accumulate) + +The forward pass uses `cublasSgemm` (F32×F32→F32). On H100, `cublasGemmEx` with BF16 inputs and F32 accumulate uses tensor cores at 2-3× the throughput. The BF16 weight buffers (`bf16_params_buf`, `bf16_target_params_buf`) already exist but are never used. + +**Key constraints:** +- Weights: BF16 (pre-converted from F32 via `f32_to_bf16_kernel`) +- Activations/states: MUST stay F32 (written by upstream kernels in F32) +- Accumulator: F32 (no precision loss in output) +- cuBLAS GemmEx can do BF16×F32→F32 using `CUBLAS_COMPUTE_32F` + +Actually — `cublasGemmEx` requires both A and B to have the same type when using tensor cores. We need BF16×BF16→F32. This means **states and activations must also be BF16** for the forward SGEMM. But: +- States come from the experience collector in F32 +- Activations (h_s1, h_s2) are written by SGEMM and consumed by the next layer + +**Practical approach:** Use **mixed-precision GemmEx** with: +- A = weights (BF16, from bf16_params_buf) +- B = input (F32, states/activations) +- C = output (F32, activations/logits) +- Compute: `CUBLAS_COMPUTE_32F` (uses tensor cores on H100 for BF16×F32 with TF32 accumulate) + +On H100, `CUBLAS_COMPUTE_32F` with mixed types uses TF32 tensor cores, which is ~1.5× faster than pure F32. For full BF16 tensor core speed (3×), both inputs must be BF16. That requires converting activations to BF16 too — a larger change. + +**This task implements the TF32 mixed-precision path (1.5× speedup, minimal changes).** + +**Files:** +- Modify: `crates/ml/src/cuda_pipeline/batched_forward.rs` + +- [ ] **Step 1: Replace `cublasSgemm` with `cublasGemmEx` using TF32 compute** + +In `sgemm_layer` (~line 495), change from: + +```rust +cublas_result::sgemm( + self.handle.0, + cublas_sys::cublasOperation_t::CUBLAS_OP_T, + cublas_sys::cublasOperation_t::CUBLAS_OP_N, + n as i32, b as i32, k as i32, + &alpha, + w_ptr as *const f32, k as i32, + a_ptr as *const f32, k as i32, + &beta, + c_ptr as *mut f32, n as i32, +) +``` + +To: + +```rust +cublas_sys::cublasGemmEx( + self.handle.0, + cublas_sys::cublasOperation_t::CUBLAS_OP_T, + cublas_sys::cublasOperation_t::CUBLAS_OP_N, + n as i32, b as i32, k as i32, + &alpha as *const f32 as *const std::ffi::c_void, + w_ptr as *const std::ffi::c_void, + cublas_sys::cudaDataType_t::CUDA_R_32F, // A type + k as i32, + a_ptr as *const std::ffi::c_void, + cublas_sys::cudaDataType_t::CUDA_R_32F, // B type + k as i32, + &beta as *const f32 as *const std::ffi::c_void, + c_ptr as *mut std::ffi::c_void, + cublas_sys::cudaDataType_t::CUDA_R_32F, // C type + n as i32, + cublas_sys::cublasComputeType_t::CUBLAS_COMPUTE_32F_FAST_TF32, // TF32 tensor cores + cublas_sys::cublasGemmAlgo_t::CUBLAS_GEMM_DEFAULT_TENSOR_OP, +) +``` + +`CUBLAS_COMPUTE_32F_FAST_TF32` enables TF32 tensor cores on H100/A100. With F32 inputs, cuBLAS internally truncates to TF19 mantissa for the multiply, accumulates in F32. ~1.5× speedup over pure F32 SGEMM. + +Do the same for `sgemm_layer_raw`. + +- [ ] **Step 2: Set cuBLAS math mode for tensor ops** + +In the `CublasForward::new()` constructor, set the math mode: + +```rust +unsafe { + cublas_sys::cublasSetMathMode( + handle.0, + cublas_sys::cublasMath_t::CUBLAS_TF32_TENSOR_OP_MATH, + ); +} +``` + +This allows `cublasSgemm` to use TF32 even without switching to `cublasGemmEx`. + +- [ ] **Step 3: Do the same for `batched_backward.rs`** + +The backward pass also uses `cublasSgemm`. Apply the same TF32 math mode or `cublasGemmEx` conversion. + +- [ ] **Step 4: Verify compilation + tests** + +Run: `SQLX_OFFLINE=true cargo check -p ml && SQLX_OFFLINE=true cargo test -p ml --lib -- dqn` +Expected: All pass. TF32 has ≤0.1% numerical difference from F32 — tests should still pass. + +- [ ] **Step 5: Commit** + +```bash +git add crates/ml/src/cuda_pipeline/batched_forward.rs crates/ml/src/cuda_pipeline/batched_backward.rs +git commit -m "perf: enable TF32 tensor cores for cuBLAS SGEMM — 1.5x forward+backward throughput" +``` + +--- + +### Task 4: Batch spectral norm into fewer kernel launches + +Currently 8 separate kernel launches for head spectral norm, each with `grid=(1,1,1)` — massively underutilized GPU. Batch them into a single kernel that processes all 8 matrices sequentially (same single block, but loop over matrices using a matrix descriptor array). + +**Files:** +- Modify: `crates/ml/src/cuda_pipeline/dqn_utility_kernels.cu` +- Modify: `crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs` + +- [ ] **Step 1: Add batched spectral norm kernel** + +In `dqn_utility_kernels.cu`, add: + +```cuda +/* ══════════════════════════════════════════════════════════════════════ + * BATCHED SPECTRAL NORM KERNEL + * + * Processes N weight matrices in a single kernel launch. + * Grid = (N, 1, 1), Block = (256, 1, 1). + * Each block handles one matrix via the same power iteration as + * spectral_norm_kernel but indexed by blockIdx.x into descriptor arrays. + * ══════════════════════════════════════════════════════════════════════ */ + +extern "C" __global__ void batched_spectral_norm_kernel( + /* Arrays of N pointers, one per matrix */ + float** __restrict__ W_arr, /* [N] pointers to weight matrices */ + float** __restrict__ u_arr, /* [N] pointers to left singular vectors */ + float** __restrict__ v_arr, /* [N] pointers to right singular vectors */ + const int* __restrict__ out_dims, /* [N] out_dim per matrix */ + const int* __restrict__ in_dims, /* [N] in_dim per matrix */ + float sigma_max, + int n_matrices +) { + int mat_idx = blockIdx.x; + if (mat_idx >= n_matrices) return; + + float* W = W_arr[mat_idx]; + float* u = u_arr[mat_idx]; + float* v = v_arr[mat_idx]; + int out_dim = out_dims[mat_idx]; + int in_dim = in_dims[mat_idx]; + int n_total = out_dim * in_dim; + + /* Same power iteration logic as spectral_norm_kernel... */ + /* (copy the body of spectral_norm_kernel here, replacing + the fixed arguments with the per-matrix values) */ + __shared__ float shmem[256]; + int tid = threadIdx.x; + int bd = blockDim.x; + + // Step 1: v_new = W^T u + for (int col = tid; col < in_dim; col += bd) { + float val = 0.0f; + for (int row = 0; row < out_dim; row++) + val += W[row * in_dim + col] * u[row]; + v[col] = val; + } + __syncthreads(); + + // Normalize v + float local_v2 = 0.0f; + for (int col = tid; col < in_dim; col += bd) + local_v2 += v[col] * v[col]; + shmem[tid] = local_v2; + __syncthreads(); + for (int s = bd / 2; s > 0; s >>= 1) { + if (tid < s) shmem[tid] += shmem[tid + s]; + __syncthreads(); + } + float v_norm = sqrtf(shmem[0] + 1e-12f); + for (int col = tid; col < in_dim; col += bd) + v[col] /= v_norm; + __syncthreads(); + + // Step 2: u_new = W v_new + for (int row = tid; row < out_dim; row += bd) { + float val = 0.0f; + for (int col = 0; col < in_dim; col++) + val += W[row * in_dim + col] * v[col]; + u[row] = val; + } + __syncthreads(); + + // sigma = ||u_new|| + float local_u2 = 0.0f; + for (int row = tid; row < out_dim; row += bd) + local_u2 += u[row] * u[row]; + shmem[tid] = local_u2; + __syncthreads(); + for (int s = bd / 2; s > 0; s >>= 1) { + if (tid < s) shmem[tid] += shmem[tid + s]; + __syncthreads(); + } + float sigma = sqrtf(shmem[0] + 1e-12f); + + // Normalize u + for (int row = tid; row < out_dim; row += bd) + u[row] /= (sigma + 1e-12f); + __syncthreads(); + + // Scale W if sigma > sigma_max + if (sigma < 1e-6f) return; + float scale = (sigma > sigma_max) ? (sigma_max / sigma) : 1.0f; + if (scale < 1.0f) { + for (int i = tid; i < n_total; i += bd) + W[i] *= scale; + } +} +``` + +- [ ] **Step 2: Load the kernel in Rust** + +Add `batched_spectral_norm_kernel` loading in `compile_training_kernels`. Add the function field to the struct. + +- [ ] **Step 3: Allocate descriptor arrays on GPU** + +Allocate `CudaSlice` arrays for the 10 W/u/v pointer arrays and `CudaSlice` for out_dims/in_dims. Upload once at construction. + +- [ ] **Step 4: Replace 10 individual spec_norm! calls with single batched launch** + +In `apply_spectral_norm`, replace the 2 trunk + 8 head launches with: + +```rust +// Single batched launch: 10 blocks × 256 threads +unsafe { + self.stream + .launch_builder(&self.batched_spectral_norm_kernel) + .arg(&self.spec_w_ptrs_ptr) + .arg(&self.spec_u_ptrs_ptr) + .arg(&self.spec_v_ptrs_ptr) + .arg(&self.spec_out_dims_ptr) + .arg(&self.spec_in_dims_ptr) + .arg(&sigma_max) + .arg(&10_i32) + .launch(LaunchConfig { + grid_dim: (10, 1, 1), + block_dim: (256, 1, 1), + shared_mem_bytes: 256 * 4, + })?; +} +``` + +This replaces 10 kernel launches with 1. + +- [ ] **Step 5: Verify compilation + tests** + +Run: `SQLX_OFFLINE=true cargo check -p ml && SQLX_OFFLINE=true cargo test -p ml --lib -- dqn` +Expected: All pass. + +- [ ] **Step 6: Commit** + +```bash +git add crates/ml/src/cuda_pipeline/dqn_utility_kernels.cu crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs +git commit -m "perf: batch 10 spectral norm launches into single kernel" +``` + +--- + +### Task 5: Dynamic gradient budget allocation + +When IQN/ensemble are disabled, their budget is wasted — C51 only gets 70% even as the sole gradient source. + +**Files:** +- Modify: `crates/ml/src/trainers/dqn/fused_training.rs` + +**NOTE:** This is already partially implemented (dynamic `c51_frac` computation at Step 2b). Verify it's correct and add a test. + +- [ ] **Step 1: Verify the dynamic budget code** + +Read the current Step 2b in `fused_training.rs` and confirm: +```rust +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 }; +let ens_frac = if !self.ensemble_extra_heads.is_empty() { ENS_GRAD_BUDGET } else { 0.0 }; +let c51_frac = 1.0 - cql_frac - iqn_frac - ens_frac; +``` + +This is correct. C51 gets 100% when nothing else is active. + +- [ ] **Step 2: Add test for dynamic budget** + +In `gradient_budget.rs`, add a test that constructs a trainer with CQL disabled (cql_alpha=0) and verifies C51 can use the full max_grad_norm: + +```rust +#[test] +fn test_c51_gets_full_budget_when_no_auxiliaries() -> anyhow::Result<()> { + let dev = cuda_device(); + let stream = Arc::clone(dev.cuda_stream().expect("cuda stream")); + let mut cfg = test_config(); + cfg.cql_alpha = 0.0; // Disable CQL + cfg.iqn_lambda = 0.0; // Disable IQN + let mut trainer = GpuDqnTrainer::new(stream.clone(), cfg)?; + + // Fill grad_buf with large gradient + let total_params = trainer.total_params(); + stream.memcpy_htod(&vec![1.0_f32; total_params], trainer.grad_buf_mut()) + .map_err(|e| anyhow::anyhow!("{e}"))?; + + // Clip to full max_grad_norm (what C51 should get when solo) + trainer.clip_grad_buf_inplace(10.0)?; // max_grad_norm * 1.0 + + let norm = trainer.read_grad_norm_sync()?; + assert!( + norm > 9.9, + "C51 solo should use full max_grad_norm, got {norm}" + ); + + Ok(()) +} +``` + +- [ ] **Step 3: Verify compilation + tests** + +Run: `SQLX_OFFLINE=true cargo test -p ml --lib -- gradient_budget` +Expected: All pass including new test. + +- [ ] **Step 4: Commit** + +```bash +git add crates/ml/src/trainers/dqn/fused_training.rs crates/ml/src/trainers/dqn/smoke_tests/gradient_budget.rs +git commit -m "fix: dynamic gradient budget — C51 gets full max_grad_norm when auxiliaries disabled" +``` + +--- + +### Task 6: Remove `cuStreamSynchronize` from `replay_adam_and_readback` + +The Adam readback (`replay_adam_and_readback`) syncs the stream and reads back loss + grad_norm (8 bytes DtoH). This is the ONLY legitimate sync in the per-step path — we need the loss for logging. But we can defer it to the NEXT step's start, overlapping GPU work with CPU logging. + +**Files:** +- Modify: `crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs` +- Modify: `crates/ml/src/trainers/dqn/fused_training.rs` + +- [ ] **Step 1: Convert to async readback with deferred sync** + +In `replay_adam_and_readback`, replace: +```rust +cuStreamSynchronize(self.stream.cu_stream()); +cuMemcpyDtoH_v2(loss_host..., loss_ptr, 4); +cuMemcpyDtoH_v2(norm_host..., norm_ptr, 4); +``` + +With async DtoH into a pinned host buffer, returning the PREVIOUS step's values: + +```rust +pub fn replay_adam_and_readback(&mut self) -> Result { + // Return PREVIOUS step's readback (already synced by this step's graph_forward) + let prev_loss = self.scalars_readback_host[0]; + let prev_norm = self.scalars_readback_host[1]; + + // Replay Adam (async) + self.replay_adam()?; + + // Queue async DtoH for THIS step's results (will be ready by next step) + unsafe { + cudarc::driver::sys::cuMemcpyDtoHAsync_v2( + self.scalars_readback_host.as_mut_ptr().cast(), + self.cached_ptrs.total_loss, 4, + self.stream.cu_stream(), + ); + cudarc::driver::sys::cuMemcpyDtoHAsync_v2( + self.scalars_readback_host.as_mut_ptr().add(1).cast(), + self.cached_ptrs.grad_norm, 4, + self.stream.cu_stream(), + ); + } + + Ok(FusedTrainScalars { + total_loss: prev_loss, + grad_norm: prev_norm.sqrt(), + }) +} +``` + +This eliminates the per-step sync. The readback has 1-step lag (reporting previous step's metrics) which is standard and acceptable for monitoring. + +- [ ] **Step 2: Allocate pinned host buffer** + +Replace `scalars_readback_host: Vec` with pinned memory for async DtoH: + +```rust +// In constructor: +let scalars_readback_host = vec![0.0_f32; 2]; // [loss, grad_norm_sq] +// cudarc doesn't expose cuMemAllocHost directly — keep Vec but ensure +// the async DtoH is followed by a stream sync at epoch boundary. +``` + +Actually — `cuMemcpyDtoHAsync_v2` requires the destination to be page-locked (pinned). With a regular Vec, this is technically UB. The safe approach: keep the sync but move it to the epoch boundary, NOT per-step. + +**Revised approach:** Batch the sync. Instead of syncing every step, sync once per N steps (e.g. every 100 steps or at epoch boundary). Use the existing `readback_host` Vec with synchronous DtoH but only when actually needed. + +- [ ] **Step 3: Move sync to epoch boundary** + +In `fused_training.rs`, remove the per-step sync from `replay_adam_and_readback`. Instead, sync only when `steps_since_varmap_sync % 100 == 0` or at epoch boundary. For intermediate steps, return the last known values. + +- [ ] **Step 4: Verify compilation + tests** + +Run: `SQLX_OFFLINE=true cargo check -p ml && SQLX_OFFLINE=true cargo test -p ml --lib -- dqn` +Expected: All pass. + +- [ ] **Step 5: Commit** + +```bash +git add crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs crates/ml/src/trainers/dqn/fused_training.rs +git commit -m "perf: defer Adam readback sync to every 100 steps — eliminate per-step cuStreamSync" +``` + +--- + +### Task 7: Performance regression smoke test + +Add a test that measures training step throughput and fails if it drops below a threshold. + +**Files:** +- Modify: `crates/ml/src/trainers/dqn/smoke_tests/gradient_budget.rs` + +- [ ] **Step 1: Add throughput test** + +```rust +/// Training step throughput must not regress. +/// At 3s/epoch with 10K steps: ~3333 steps/sec minimum on RTX 3050. +/// On H100: ~10K+ steps/sec. Test uses a conservative threshold. +#[test] +fn test_training_step_throughput() -> anyhow::Result<()> { + let dev = cuda_device(); + let stream = Arc::clone(dev.cuda_stream().expect("cuda stream")); + let cfg = test_config(); + let mut trainer = GpuDqnTrainer::new(stream.clone(), cfg.clone())?; + let mut dueling = alloc_dueling(&stream, &cfg); + let mut branching = alloc_branching(&stream, &cfg); + + // Warm up + trainer.apply_spectral_norm(&mut dueling, &mut branching)?; + trainer.clip_grad_buf_inplace(7.0)?; + + // Measure 1000 iterations of clip + spectral norm (the hot path additions) + let n_iters = 1000; + let start = std::time::Instant::now(); + for _ in 0..n_iters { + trainer.apply_spectral_norm(&mut dueling, &mut branching)?; + trainer.clip_grad_buf_inplace(7.0)?; + } + // Sync to ensure all GPU work completes + unsafe { cudarc::driver::sys::cuStreamSynchronize(stream.cu_stream()); } + let elapsed = start.elapsed(); + + let per_iter_us = elapsed.as_micros() as f64 / n_iters as f64; + assert!( + per_iter_us < 500.0, // 500μs budget for spectral norm + clip per step + "Spectral norm + clip took {per_iter_us:.0}μs/iter — budget is 500μs" + ); + + Ok(()) +} +``` + +- [ ] **Step 2: Verify compilation + run** + +Run: `SQLX_OFFLINE=true cargo test -p ml --lib -- test_training_step_throughput --nocapture` +Expected: PASS with per-iter time well under 500μs. + +- [ ] **Step 3: Commit** + +```bash +git add crates/ml/src/trainers/dqn/smoke_tests/gradient_budget.rs +git commit -m "test: training step throughput regression test" +``` + +--- + +## Execution Dependencies + +``` +Task 1 (cache pointers) → Task 2 (remove ensemble sync) → Task 6 (defer Adam readback) + → Task 3 (TF32 tensor cores) — independent + → Task 4 (batch spectral norm) — independent +Task 5 (dynamic budget) — independent +Task 7 (throughput test) — after Tasks 1-4 +``` + +Task 1 is the foundation — all other tasks benefit from cached pointers. Tasks 3 and 4 are independent GPU optimizations. Task 6 requires Task 1's cached pointers. Task 7 validates everything. + +## Expected Impact + +| Optimization | Estimated speedup | +|---|---| +| Eliminate 4× cuStreamSync (200μs/step) | 2s/epoch saved | +| TF32 tensor cores (1.5× SGEMM) | 30-40% forward+backward time | +| Batch spectral norm (10→1 launch) | ~50μs/step saved | +| Eliminate 99 raw_device_ptr calls | ~100μs/step saved | +| Defer Adam readback sync | ~50μs/step saved | +| **Combined** | **Target: ≤5s/epoch on H100** | diff --git a/docs/superpowers/plans/2026-03-27-gradient-stability.md b/docs/superpowers/plans/2026-03-27-gradient-stability.md new file mode 100644 index 000000000..84e13aa0b --- /dev/null +++ b/docs/superpowers/plans/2026-03-27-gradient-stability.md @@ -0,0 +1,769 @@ +# Gradient Stability — Proper Multi-Task Gradient Isolation + Full Spectral Norm + +> **For agentic workers:** REQUIRED SUB-SKILL: Use superpowers:subagent-driven-development (recommended) or superpowers:executing-plans to implement this plan task-by-task. Steps use checkbox (`- [ ]`) syntax for tracking. + +**Goal:** Eliminate gradient explosion by isolating every gradient source into its own buffer with budget-allocated clipping, extending spectral normalization to all weight matrices, and adding per-component monitoring. + +**Architecture:** Four gradient sources (C51, CQL, IQN, ensemble) each get an independent norm budget that sums to exactly `max_grad_norm`. CQL gets its own scratch buffer (currently it accumulates directly into `grad_buf` with beta=1.0, mixing with C51). Spectral norm power iteration constrains all 10 weight matrices (2 trunk + 8 heads) with uniform σ_max. Per-component Prometheus gauges provide real-time gradient interference detection. + +**Tech Stack:** Rust, CUDA (precompiled cubin), cudarc, cuBLAS, Prometheus, tokio + +**Prior work (already merged):** `dqn_clipped_saxpy_kernel`, `dqn_clip_grad_kernel`, `clip_grad_buf_inplace()`, clipped SAXPY for IQN/ensemble, `weight_decay` hardcode fix, intermediate gradient clipping between CQL and IQN. **This plan fixes design flaws in the prior work.** + +--- + +## Design — Why the current per-component clipping is wrong + +The already-implemented clipping has three flaws: + +**Flaw 1: CQL has no independent clipping.** CQL's `backward_full(beta=1.0)` accumulates directly into `grad_buf` on top of C51. The intermediate clip (Step 5d) then clips the MIXED C51+CQL. If CQL dominates, C51's gradient direction is destroyed by the clip. + +**Flaw 2: Unbudgeted clip thresholds.** Each component clips to the FULL `max_grad_norm`: +``` +After Step 5d: ||grad_buf|| ≤ max_grad_norm (C51+CQL) +After IQN: ||grad_buf|| ≤ max_grad_norm + iqn_lambda × max_grad_norm +After ensemble: ||grad_buf|| ≤ max_grad_norm + iqn_lambda×mgn + ens_scale×mgn +Adam final clip: ||grad_buf|| → min(||grad||, max_grad_norm) +``` +The combined norm is ~1.35× `max_grad_norm`. Adam clips 26% every step, wasting gradient signal and biasing the direction toward whichever component was largest. + +**Flaw 3: CQL and IQN share `bw_d_h_*` scratch buffers.** This works only because IQN runs before CQL. An ordering change would silently corrupt gradients. Undocumented fragile invariant. + +### Proper design + +Each source gets a **budget fraction** of `max_grad_norm`: + +| Component | Budget fraction | Rationale | +|---|---|---| +| C51/MSE (primary) | 0.70 | Primary RL loss — gets most of the gradient capacity | +| CQL (conservative penalty) | 0.15 | Regularizer — secondary importance | +| IQN (quantile risk) | 0.10 | Auxiliary head — small trunk contribution | +| Ensemble (diversity) | 0.05 | Weakest signal — diversity is a gentle push | +| **Total** | **1.00** | Adam clip NEVER fires → zero wasted gradient | + +Pipeline after fix: +``` +1. graph_forward → grad_buf = C51_raw +2. clip_grad_buf(0.70 × mgn) → ||grad_buf|| ≤ C51 budget +3. CQL backward → cql_scratch = CQL_raw (separate buffer) +4. clipped SAXPY → grad_buf += clip(cql_scratch, 0.15 × mgn) +5. IQN clipped SAXPY → grad_buf += λ × clip(iqn_scratch, 0.10 × mgn) +6. Ens clipped SAXPY → grad_buf += s × clip(ens_scratch, 0.05 × mgn) +7. graph_adam → Adam (safety clip at mgn, should never fire) +``` + +--- + +## File Map + +| File | Role | Action | +|---|---|---| +| `crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs` | GPU trainer — buffers, kernel launches, spectral norm | Modify: add `cql_grad_scratch` buffer, change `apply_cql_gradient` to use scratch, add 16 spec_u/v buffers, extend spectral norm to all heads | +| `crates/ml/src/trainers/dqn/fused_training.rs` | Fused training orchestration | Modify: add C51 clip after graph_forward, change CQL to use scratch+SAXPY, update all clip thresholds to budget fractions | +| `crates/common/src/metrics/training_metrics.rs` | Prometheus gauge definitions + setters | Modify: add 2 per-component gradient norm gauges | +| `crates/ml/src/trainers/dqn/trainer/training_loop.rs` | Training loop — epoch metrics | Modify: push per-component grad norms at epoch boundary | +| `crates/ml/src/trainers/dqn/smoke_tests/training_stability.rs` | Smoke tests | Modify: add gradient-bounded-across-epochs test | + +--- + +### Task 1: Isolate CQL gradient into separate scratch buffer + +CQL currently runs `backward_full(beta=1.0)` directly into `grad_buf`, mixing with C51 before any CQL-specific clipping. Fix: allocate a `total_params`-sized scratch buffer, run CQL backward into it, then independently clip and SAXPY into `grad_buf`. + +**Files:** +- Modify: `crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs` + +- [ ] **Step 1: Add `cql_grad_scratch` field to struct** + +After the existing `grad_norm_buf` field (~line 406), add: + +```rust + cql_grad_scratch: CudaSlice, // [TOTAL_PARAMS] CQL gradient isolation buffer +``` + +- [ ] **Step 2: Allocate in constructor** + +After the existing `grad_buf` allocation (~line 1530), add: + +```rust + let cql_grad_scratch = alloc_f32(&stream, total_params, "cql_grad_scratch")?; +``` + +- [ ] **Step 3: Add to struct construction** + +After `grad_buf,` in the `Ok(Self { ... })` block, add: + +```rust + cql_grad_scratch, +``` + +- [ ] **Step 4: Rewrite `apply_cql_gradient` to use scratch buffer** + +The core change: instead of passing `grad_buf` pointer to `backward_full`, pass `cql_grad_scratch`. Zero the scratch first (since backward uses beta=1.0 accumulation). + +In `apply_cql_gradient`, replace the backward_full call block (lines ~1195-1209). Change from: + +```rust + // Run full backward pass with CQL logit gradients. + // This ACCUMULATES into grad_buf (beta=1.0 in cuBLAS SGEMM), + // adding CQL parameter gradients on top of C51's. + self.cublas_backward.backward_full( + &self.stream, + d_v_ptr, + &d_adv_ptrs, + states_ptr_fw, + h_s1_ptr, h_s2_ptr, h_v_ptr, + &[h_b0_ptr, h_b1_ptr, h_b2_ptr], + &w_ptrs, + raw_device_ptr(&self.grad_buf, &self.stream), + scratch_d_h_s2, scratch_d_h_s1, scratch_d_h_v, + &[scratch_d_h_b0, scratch_d_h_b1, scratch_d_h_b2], + ).map_err(|e| MLError::ModelError(format!("CQL backward_full: {e}")))?; + + Ok(true) +``` + +To: + +```rust + // Zero CQL scratch buffer (backward_full uses beta=1.0 accumulation) + self.stream.memset_zeros(&mut self.cql_grad_scratch) + .map_err(|e| MLError::ModelError(format!("zero cql_grad_scratch: {e}")))?; + + // Run full backward pass with CQL logit gradients into ISOLATED scratch buffer. + // This produces CQL parameter gradients WITHOUT mixing with C51's grad_buf. + self.cublas_backward.backward_full( + &self.stream, + d_v_ptr, + &d_adv_ptrs, + states_ptr_fw, + h_s1_ptr, h_s2_ptr, h_v_ptr, + &[h_b0_ptr, h_b1_ptr, h_b2_ptr], + &w_ptrs, + raw_device_ptr(&self.cql_grad_scratch, &self.stream), + scratch_d_h_s2, scratch_d_h_s1, scratch_d_h_v, + &[scratch_d_h_b0, scratch_d_h_b1, scratch_d_h_b2], + ).map_err(|e| MLError::ModelError(format!("CQL backward_full: {e}")))?; + + Ok(true) +``` + +- [ ] **Step 5: Add `apply_cql_clipped_saxpy` method** + +Add a new public method after `apply_cql_gradient` that clips the CQL scratch and SAXPYs into `grad_buf`. This is called from `fused_training.rs` after `apply_cql_gradient` returns `Ok(true)`. + +```rust + /// Clip CQL gradient scratch buffer and add to grad_buf via SAXPY. + /// + /// Called after `apply_cql_gradient` which populated `cql_grad_scratch`. + /// Uses the same clipped SAXPY pattern as IQN/ensemble: compute norm of + /// scratch, then add with per-component clip. + /// + /// `cql_budget` is the CQL fraction of max_grad_norm (e.g. 0.15 × max_grad_norm). + pub fn apply_cql_clipped_saxpy(&mut self, cql_budget: f32) -> Result<(), MLError> { + let _evt_guard = EventTrackingGuard::new(self.stream.context()); + + // Compute CQL gradient norm + self.stream.memset_zeros(&mut self.grad_norm_buf) + .map_err(|e| MLError::ModelError(format!("zero cql_grad_norm: {e}")))?; + + let total = self.total_params as i32; + let blocks = ((self.total_params + 255) / 256) as u32; + let scratch_ptr = raw_device_ptr(&self.cql_grad_scratch, &self.stream); + let norm_ptr = raw_device_ptr(&self.grad_norm_buf, &self.stream); + + unsafe { + self.stream + .launch_builder(&self.grad_norm_kernel) + .arg(&scratch_ptr) + .arg(&norm_ptr) + .arg(&total) + .launch(LaunchConfig { + grid_dim: (blocks, 1, 1), + block_dim: (256, 1, 1), + shared_mem_bytes: 256, + }) + .map_err(|e| MLError::ModelError(format!("CQL grad_norm: {e}")))?; + } + + // Clipped SAXPY: grad_buf += 1.0 * clip(cql_scratch, cql_budget) + let grad_ptr = raw_device_ptr(&self.grad_buf, &self.stream); + let alpha = 1.0_f32; + unsafe { + self.stream + .launch_builder(&self.clipped_saxpy_kernel) + .arg(&grad_ptr) + .arg(&scratch_ptr) + .arg(&alpha) + .arg(&cql_budget) + .arg(&norm_ptr) + .arg(&total) + .launch(LaunchConfig { + grid_dim: (blocks, 1, 1), + block_dim: (256, 1, 1), + shared_mem_bytes: 0, + }) + .map_err(|e| MLError::ModelError(format!("CQL clipped SAXPY: {e}")))?; + } + + Ok(()) + } +``` + +- [ ] **Step 6: Verify compilation** + +Run: `SQLX_OFFLINE=true cargo check -p ml 2>&1 | tail -5` +Expected: `Finished` (new method is unused — wired in Task 2). + +- [ ] **Step 7: Commit** + +```bash +git add crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs +git commit -m "feat: isolate CQL gradient into separate scratch buffer with independent clipping" +``` + +--- + +### Task 2: Implement gradient budget allocation + +Change all clip thresholds to use fractional budgets that sum to `max_grad_norm`. Add C51 clipping after graph_forward. Wire up the new CQL isolated path. + +**Files:** +- Modify: `crates/ml/src/trainers/dqn/fused_training.rs` +- Modify: `crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs` + +**Budget constants** (defined in fused_training.rs): +```rust +/// Per-component gradient norm budget fractions. Must sum to 1.0. +/// The primary C51 loss gets 70% of max_grad_norm. Auxiliary objectives +/// (CQL, IQN, ensemble) share the remaining 30%. +const C51_GRAD_BUDGET: f32 = 0.70; +const CQL_GRAD_BUDGET: f32 = 0.15; +const IQN_GRAD_BUDGET: f32 = 0.10; +const ENS_GRAD_BUDGET: f32 = 0.05; +``` + +- [ ] **Step 1: Add budget constants** + +At the top of `fused_training.rs` (after the `use` imports, before struct definitions), add the 4 constants shown above. + +- [ ] **Step 2: Add C51 clip after graph_forward** + +After `train_step_gpu` returns (after `let _fused_placeholder = ...` ~line 570), add a clip of the raw C51 gradient: + +```rust + // ── Step 2b: Clip raw C51 gradient to its budget ────────────────── + // The graph_forward backward produces raw C51 gradients in grad_buf. + // Clip to C51's budget fraction BEFORE any auxiliary injection. + { + let c51_budget = self.trainer.config().max_grad_norm * C51_GRAD_BUDGET; + self.trainer.clip_grad_buf_inplace(c51_budget) + .map_err(|e| anyhow::anyhow!("C51 gradient budget clip: {e}"))?; + } +``` + +- [ ] **Step 3: Wire CQL isolated path** + +Replace the existing CQL block (~lines 742-756) with the new isolated CQL path: + +```rust + // ── Step 5c: CQL conservative penalty (isolated gradient) ───────── + // CQL backward runs into a SEPARATE scratch buffer (not grad_buf). + // Its gradient is independently clipped to CQL's budget fraction, + // then added to grad_buf via clipped SAXPY. + if self.trainer.has_cql() { + match self.trainer.apply_cql_gradient() { + Ok(true) => { + let cql_budget = self.trainer.config().max_grad_norm * CQL_GRAD_BUDGET; + self.trainer.apply_cql_clipped_saxpy(cql_budget) + .map_err(|e| anyhow::anyhow!("CQL clipped SAXPY: {e}"))?; + tracing::trace!("CQL gradient: isolated → clipped → SAXPY into grad_buf"); + } + Ok(false) => {} // CQL disabled or alpha=0 + Err(e) => { + tracing::warn!("CQL gradient failed (non-fatal): {e}"); + } + } + } +``` + +- [ ] **Step 4: Remove the old combined Step 5d clip** + +Delete the entire "Step 5d: Clip combined C51+CQL gradient before auxiliary injections" block (~lines 754-773). It's replaced by Steps 2b (C51 clip) and 5c (CQL isolation). Keep the diagnostic logging but move it to Step 2b: + +In the Step 2b block added above, add the diagnostic: + +```rust + // Diagnostic: log C51 pre-clip norm every 1000 steps + if self.steps_since_varmap_sync % 1000 == 0 { + let pre_clip_norm = self.trainer.read_grad_norm_sync() + .unwrap_or(f32::NAN); + tracing::info!( + c51_raw_grad_norm = pre_clip_norm, + c51_budget, + step = self.steps_since_varmap_sync, + "Per-component gradient diagnostic (C51 before budget clip)" + ); + } +``` + +- [ ] **Step 5: Update IQN clip threshold to budget fraction** + +In `gpu_dqn_trainer.rs`, in `apply_iqn_trunk_gradient` (~line 830), change: + +```rust + let max_component_norm = self.config.max_grad_norm; +``` + +To: + +```rust + let max_component_norm = self.config.max_grad_norm * 0.10; // IQN_GRAD_BUDGET +``` + +- [ ] **Step 6: Update ensemble clip threshold to budget fraction** + +In `gpu_dqn_trainer.rs`, in `apply_ensemble_trunk_gradient` (~line 1071), change: + +```rust + let max_component_norm = self.config.max_grad_norm; +``` + +To: + +```rust + let max_component_norm = self.config.max_grad_norm * 0.05; // ENS_GRAD_BUDGET +``` + +- [ ] **Step 7: Verify compilation** + +Run: `SQLX_OFFLINE=true cargo check -p ml 2>&1 | tail -5` +Expected: `Finished` with no errors. + +- [ ] **Step 8: Run existing tests** + +Run: `SQLX_OFFLINE=true cargo test -p ml --lib -- dqn 2>&1 | tail -5` +Expected: All tests pass (122+). + +- [ ] **Step 9: Commit** + +```bash +git add crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs crates/ml/src/trainers/dqn/fused_training.rs +git commit -m "feat: proper gradient budget allocation — C51 70%, CQL 15%, IQN 10%, ensemble 5%" +``` + +--- + +### Task 3: Allocate spectral norm u/v vectors for all 8 head weight matrices + +Each weight matrix W [out_dim, in_dim] needs u [out_dim] and v [in_dim] vectors for power iteration. The trunk already has `spec_u_s1/v_s1`, `spec_u_s2/v_s2`. We add 8 pairs for the heads. + +**Files:** +- Modify: `crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs` + +**Weight matrix dimensions** (from `DuelingWeightSet` at `gpu_weights.rs:109` and `BranchingWeightSet` at `gpu_weights.rs:141`): + +| Matrix | Shape [out, in] | u size | v size | GOFF index | +|---|---|---|---|---| +| W_v1 | [VALUE_H, SHARED_H2] | value_h | shared_h2 | 4 | +| W_v2 | [NUM_ATOMS, VALUE_H] | num_atoms | value_h | 6 | +| W_a1 | [ADV_H, SHARED_H2] | adv_h | shared_h2 | 8 | +| W_a2 | [B0×NUM_ATOMS, ADV_H] | b0×num_atoms | adv_h | 10 | +| W_bo1 | [ADV_H, SHARED_H2] | adv_h | shared_h2 | 12 | +| W_bo2 | [B1×NUM_ATOMS, ADV_H] | b1×num_atoms | adv_h | 14 | +| W_bu1 | [ADV_H, SHARED_H2] | adv_h | shared_h2 | 16 | +| W_bu2 | [B2×NUM_ATOMS, ADV_H] | b2×num_atoms | adv_h | 18 | + +GOFF layout verified at `gpu_dqn_trainer.rs:280-302` (`compute_param_sizes`). + +- [ ] **Step 1: Add 16 struct fields** + +After `spec_v_s2: CudaSlice,` (~line 347): + +```rust + // Head spectral norm vectors (power iteration) + // Value head + spec_u_v1: CudaSlice, // [VALUE_H] + spec_v_v1: CudaSlice, // [SHARED_H2] + spec_u_v2: CudaSlice, // [NUM_ATOMS] + spec_v_v2: CudaSlice, // [VALUE_H] + // Exposure head (branch 0) — in DuelingWeightSet + spec_u_a1: CudaSlice, // [ADV_H] + spec_v_a1: CudaSlice, // [SHARED_H2] + spec_u_a2: CudaSlice, // [B0*NUM_ATOMS] + spec_v_a2: CudaSlice, // [ADV_H] + // Order head (branch 1) — in BranchingWeightSet + spec_u_bo1: CudaSlice, // [ADV_H] + spec_v_bo1: CudaSlice, // [SHARED_H2] + spec_u_bo2: CudaSlice, // [B1*NUM_ATOMS] + spec_v_bo2: CudaSlice, // [ADV_H] + // Urgency head (branch 2) — in BranchingWeightSet + spec_u_bu1: CudaSlice, // [ADV_H] + spec_v_bu1: CudaSlice, // [SHARED_H2] + spec_u_bu2: CudaSlice, // [B2*NUM_ATOMS] + spec_v_bu2: CudaSlice, // [ADV_H] +``` + +- [ ] **Step 2: Allocate and initialize in constructor** + +After the existing `spec_v_s2` init block (~line 1746), using the existing `rand_unit` helper: + +```rust + // Head spectral norm u/v vectors (random unit initialization) + let vh = config.value_h; + let ah = config.adv_h; + let na = config.num_atoms; + let b0na = config.branch_0_size * na; + let b1na = config.branch_1_size * na; + let b2na = config.branch_2_size * na; + + let mut spec_u_v1 = alloc_f32(&stream, vh, "spec_u_v1")?; + let mut spec_v_v1 = alloc_f32(&stream, config.shared_h2, "spec_v_v1")?; + let mut spec_u_v2 = alloc_f32(&stream, na, "spec_u_v2")?; + let mut spec_v_v2 = alloc_f32(&stream, vh, "spec_v_v2")?; + let mut spec_u_a1 = alloc_f32(&stream, ah, "spec_u_a1")?; + let mut spec_v_a1 = alloc_f32(&stream, config.shared_h2, "spec_v_a1")?; + let mut spec_u_a2 = alloc_f32(&stream, b0na, "spec_u_a2")?; + let mut spec_v_a2 = alloc_f32(&stream, ah, "spec_v_a2")?; + let mut spec_u_bo1 = alloc_f32(&stream, ah, "spec_u_bo1")?; + let mut spec_v_bo1 = alloc_f32(&stream, config.shared_h2, "spec_v_bo1")?; + let mut spec_u_bo2 = alloc_f32(&stream, b1na, "spec_u_bo2")?; + let mut spec_v_bo2 = alloc_f32(&stream, ah, "spec_v_bo2")?; + let mut spec_u_bu1 = alloc_f32(&stream, ah, "spec_u_bu1")?; + let mut spec_v_bu1 = alloc_f32(&stream, config.shared_h2, "spec_v_bu1")?; + let mut spec_u_bu2 = alloc_f32(&stream, b2na, "spec_u_bu2")?; + let mut spec_v_bu2 = alloc_f32(&stream, ah, "spec_v_bu2")?; + + for (buf, len) in [ + (&mut spec_u_v1, vh), (&mut spec_v_v1, config.shared_h2), + (&mut spec_u_v2, na), (&mut spec_v_v2, vh), + (&mut spec_u_a1, ah), (&mut spec_v_a1, config.shared_h2), + (&mut spec_u_a2, b0na), (&mut spec_v_a2, ah), + (&mut spec_u_bo1, ah), (&mut spec_v_bo1, config.shared_h2), + (&mut spec_u_bo2, b1na), (&mut spec_v_bo2, ah), + (&mut spec_u_bu1, ah), (&mut spec_v_bu1, config.shared_h2), + (&mut spec_u_bu2, b2na), (&mut spec_v_bu2, ah), + ] { + let init = rand_unit(len, &mut rng_state); + stream.memcpy_htod(&init, buf) + .map_err(|e| MLError::ModelError(format!("spec head init: {e}")))?; + } +``` + +- [ ] **Step 3: Add to `Ok(Self { ... })`** + +After `spec_v_s2,`: + +```rust + spec_u_v1, spec_v_v1, spec_u_v2, spec_v_v2, + spec_u_a1, spec_v_a1, spec_u_a2, spec_v_a2, + spec_u_bo1, spec_v_bo1, spec_u_bo2, spec_v_bo2, + spec_u_bu1, spec_v_bu1, spec_u_bu2, spec_v_bu2, +``` + +- [ ] **Step 4: Verify compilation** + +Run: `SQLX_OFFLINE=true cargo check -p ml 2>&1 | tail -5` +Expected: `Finished` (unused fields warning OK — used in Task 4). + +- [ ] **Step 5: Commit** + +```bash +git add crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs +git commit -m "feat: allocate spectral norm u/v vectors for all 8 head weight matrices" +``` + +--- + +### Task 4: Extend `apply_spectral_norm` to all heads + sync params_buf + +Add 8 spectral norm kernel launches (one per head weight matrix) and sync the spectrally-normalized weights back to `params_buf` at their GOFF offsets. + +**Files:** +- Modify: `crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs` +- Modify: `crates/ml/src/trainers/dqn/fused_training.rs` + +**Design note:** Uniform σ_max for all layers. With σ_max=3.0 on 4 layers (trunk→head_fc→head_out), the end-to-end Lipschitz bound is 3⁴=81. The C51 logit spread under σ_max=3.0 is ~18 (for 101 atoms with typical activations), which is more than sufficient for sharp softmax distributions. Output layer spectral norm does NOT collapse C51 — verified analytically. + +- [ ] **Step 1: Change `apply_spectral_norm` signature** + +```rust +pub fn apply_spectral_norm( + &mut self, + online_dueling: &mut DuelingWeightSet, + online_branching: &mut BranchingWeightSet, +) -> Result<(), MLError> { +``` + +- [ ] **Step 2: Add 8 spectral norm kernel launches** + +After the existing W_s2 block (before the trunk sync-back), add: + +```rust + // Head spectral norm — uniform σ_max, same as trunk + let vh = self.config.value_h as i32; + let ah = self.config.adv_h as i32; + let na = self.config.num_atoms as i32; + let b0na = (self.config.branch_0_size * self.config.num_atoms) as i32; + let b1na = (self.config.branch_1_size * self.config.num_atoms) as i32; + let b2na = (self.config.branch_2_size * self.config.num_atoms) as i32; + + macro_rules! spec_norm { + ($w:expr, $u:expr, $v:expr, $od:expr, $id:expr, $name:literal) => {{ + let (w_ptr, _g) = $w.device_ptr(&self.stream); + let (u_ptr, _g2) = $u.device_ptr(&self.stream); + let (v_ptr, _g3) = $v.device_ptr(&self.stream); + unsafe { + self.stream + .launch_builder(&self.spectral_norm_kernel) + .arg(&w_ptr).arg(&u_ptr).arg(&v_ptr) + .arg(&$od).arg(&$id).arg(&sigma_max) + .launch(LaunchConfig { + grid_dim: (1, 1, 1), + block_dim: (256, 1, 1), + shared_mem_bytes: 512 * 4, + }) + .map_err(|e| MLError::ModelError( + format!(concat!("spectral_norm ", $name, ": {}"), e) + ))?; + } + }}; + } + + spec_norm!(online_dueling.w_v1, self.spec_u_v1, self.spec_v_v1, vh, sh2, "W_v1"); + spec_norm!(online_dueling.w_v2, self.spec_u_v2, self.spec_v_v2, na, vh, "W_v2"); + spec_norm!(online_dueling.w_a1, self.spec_u_a1, self.spec_v_a1, ah, sh2, "W_a1"); + spec_norm!(online_dueling.w_a2, self.spec_u_a2, self.spec_v_a2, b0na, ah, "W_a2"); + spec_norm!(online_branching.w_bo1, self.spec_u_bo1, self.spec_v_bo1, ah, sh2, "W_bo1"); + spec_norm!(online_branching.w_bo2, self.spec_u_bo2, self.spec_v_bo2, b1na, ah, "W_bo2"); + spec_norm!(online_branching.w_bu1, self.spec_u_bu1, self.spec_v_bu1, ah, sh2, "W_bu1"); + spec_norm!(online_branching.w_bu2, self.spec_u_bu2, self.spec_v_bu2, b2na, ah, "W_bu2"); +``` + +- [ ] **Step 3: Add head weight sync-back to params_buf** + +After the existing trunk sync-back section, add: + +```rust + // Sync head weights back to params_buf at their GOFF offsets + { + let param_sizes = compute_param_sizes(&self.config); + let f32_sz = std::mem::size_of::(); + let params_base = raw_device_ptr(&self.params_buf, &self.stream); + + let byte_offset = |idx: usize| -> u64 { + param_sizes[..idx].iter().sum::() as u64 * f32_sz as u64 + }; + + // GOFF indices: 4=w_v1, 6=w_v2, 8=w_a1, 10=w_a2, + // 12=w_bo1, 14=w_bo2, 16=w_bu1, 18=w_bu2 + let head_syncs: [(u64, u64, usize); 8] = [ + (raw_device_ptr(&online_dueling.w_v1, &self.stream), byte_offset(4), param_sizes[4]), + (raw_device_ptr(&online_dueling.w_v2, &self.stream), byte_offset(6), param_sizes[6]), + (raw_device_ptr(&online_dueling.w_a1, &self.stream), byte_offset(8), param_sizes[8]), + (raw_device_ptr(&online_dueling.w_a2, &self.stream), byte_offset(10), param_sizes[10]), + (raw_device_ptr(&online_branching.w_bo1, &self.stream), byte_offset(12), param_sizes[12]), + (raw_device_ptr(&online_branching.w_bo2, &self.stream), byte_offset(14), param_sizes[14]), + (raw_device_ptr(&online_branching.w_bu1, &self.stream), byte_offset(16), param_sizes[16]), + (raw_device_ptr(&online_branching.w_bu2, &self.stream), byte_offset(18), param_sizes[18]), + ]; + + for (src, dst_off, n_elems) in head_syncs { + unsafe { + cudarc::driver::result::memcpy_dtod_async( + params_base + dst_off, src, n_elems * f32_sz, + self.stream.cu_stream() + ).map_err(|e| MLError::ModelError(format!("spectral sync head: {e}")))?; + } + } + } +``` + +- [ ] **Step 4: Update call site in fused_training.rs** + +Change (~line 560): +```rust +self.trainer.apply_spectral_norm(&mut self.online_dueling) +``` +To: +```rust +self.trainer.apply_spectral_norm(&mut self.online_dueling, &mut self.online_branching) +``` + +- [ ] **Step 5: Verify compilation + run tests** + +Run: `SQLX_OFFLINE=true cargo check -p ml && SQLX_OFFLINE=true cargo test -p ml --lib -- dqn 2>&1 | tail -10` +Expected: Compiles, all 122+ tests pass. + +- [ ] **Step 6: Commit** + +```bash +git add crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs crates/ml/src/trainers/dqn/fused_training.rs +git commit -m "feat: spectral norm on all 10 weight matrices — uniform σ_max, Lipschitz bounded" +``` + +--- + +### Task 5: Per-component gradient Prometheus gauges + +**Files:** +- Modify: `crates/common/src/metrics/training_metrics.rs` +- Modify: `crates/ml/src/trainers/dqn/fused_training.rs` +- Modify: `crates/ml/src/trainers/dqn/trainer/training_loop.rs` + +- [ ] **Step 1: Register gauges and add setters in training_metrics.rs** + +In `register_all()` (after existing gradient gauge registrations): + +```rust + _ = register_gauge_vec( + "foxhunt_training_grad_norm_c51", + "Per-component gradient L2 norm: C51 primary loss (before budget clip)", + mf, + ); + _ = register_gauge_vec( + "foxhunt_training_grad_norm_combined", + "Combined gradient L2 norm after all injections (should ≤ max_grad_norm)", + mf, + ); +``` + +After `set_gradient_norm`: + +```rust +pub fn set_grad_norm_c51(model: &str, fold: &str, norm: f64) { + set_gauge_vec("foxhunt_training_grad_norm_c51", &[model, fold], norm); +} + +pub fn set_grad_norm_combined(model: &str, fold: &str, norm: f64) { + set_gauge_vec("foxhunt_training_grad_norm_combined", &[model, fold], norm); +} +``` + +- [ ] **Step 2: Track norms in FusedTrainingCtx** + +Add fields `last_c51_raw_norm: f32` and `last_combined_norm: f32` to `FusedTrainingCtx`. Initialize to 0.0. Update `last_c51_raw_norm` from the diagnostic readback in Step 2b. Update `last_combined_norm` from `fused_result.grad_norm` after graph_adam. + +Add accessor: +```rust + pub fn grad_norm_diagnostics(&self) -> (f32, f32) { + (self.last_c51_raw_norm, self.last_combined_norm) + } +``` + +- [ ] **Step 3: Push in training_loop.rs** + +After `training_metrics::set_gradient_norm(...)`: +```rust + if let Some(ref fused) = self.fused_ctx { + let (c51, combined) = fused.grad_norm_diagnostics(); + training_metrics::set_grad_norm_c51("dqn", "current", c51 as f64); + training_metrics::set_grad_norm_combined("dqn", "current", combined as f64); + } +``` + +- [ ] **Step 4: Verify compilation** + +Run: `SQLX_OFFLINE=true cargo check --workspace 2>&1 | tail -5` + +- [ ] **Step 5: Commit** + +```bash +git add crates/common/src/metrics/training_metrics.rs crates/ml/src/trainers/dqn/fused_training.rs crates/ml/src/trainers/dqn/trainer/training_loop.rs +git commit -m "feat: per-component gradient Prometheus gauges (C51 raw, combined)" +``` + +--- + +### Task 6: Smoke test — gradient norms bounded across epochs + +**Files:** +- Modify: `crates/ml/src/trainers/dqn/smoke_tests/training_stability.rs` + +- [ ] **Step 1: Write test** + +```rust +/// Gradient norms must stay bounded across all epochs. +/// +/// Before fix: grad_norm grew 89K → 3.6B (40,000x over 8 epochs). +/// After fix: per-component budget allocation ensures Adam clip never fires +/// and combined norm ≤ max_grad_norm at every step. +#[test] +#[ignore] // Loads real training data — run via nightly CI or manual trigger +fn test_gradient_norm_bounded_across_epochs() -> anyhow::Result<()> { + let data_dir = test_data_dir() + .expect("FOXHUNT_TEST_DATA or test_data/ must exist"); + + let mut p = smoke_params(); + p.epochs = 3; + + let mut trainer = smoke_trainer_with(p)?; + let rt = tokio::runtime::Builder::new_current_thread() + .enable_all() + .build()?; + let metrics = rt.block_on(trainer.train(&data_dir, |_epoch, _bytes, _best| { + Ok("skip".to_owned()) + }))?; + + let grad_norm = metrics.additional_metrics.get("avg_gradient_norm") + .copied().unwrap_or(f64::NAN); + + assert!(grad_norm.is_finite(), "avg_gradient_norm not finite: {grad_norm}"); + assert!(grad_norm > 0.0, "avg_gradient_norm should be positive: {grad_norm}"); + assert!( + grad_norm < 100_000.0, + "GRADIENT EXPLOSION: avg_gradient_norm {grad_norm:.1} exceeds 100K" + ); + assert_eq!(metrics.epochs_trained, 3, "should complete all 3 epochs"); + + drop(trainer); + drop(rt); + Ok(()) +} +``` + +- [ ] **Step 2: Compile + run** + +Run: `SQLX_OFFLINE=true cargo test -p ml --lib -- test_gradient_norm_bounded_across_epochs --no-run 2>&1 | tail -5` + +If GPU + test data available: +Run: `FOXHUNT_TEST_DATA=test_data/futures-baseline cargo test -p ml --lib -- test_gradient_norm_bounded_across_epochs --ignored --nocapture 2>&1 | tail -20` + +- [ ] **Step 3: Commit** + +```bash +git add crates/ml/src/trainers/dqn/smoke_tests/training_stability.rs +git commit -m "test: gradient norm bounded across epochs smoke test" +``` + +--- + +### Task 7: Research — PCGrad, LayerNorm, adaptive weighting analysis + +Pure research — no code. Write analysis of 3 additional architectural approaches. + +**Files:** +- Create: `docs/superpowers/plans/2026-03-27-gradient-stability-research.md` + +- [ ] **Step 1: Write analysis** covering PCGrad (Yu 2020), LayerNorm, and Uncertainty Weighting (Kendall 2018). Evaluate feasibility, impact, and interaction with our per-component budget + spectral norm approach. Recommend PCGrad as the best follow-up if gradient interference persists after validation. + +- [ ] **Step 2: Commit** + +```bash +git add docs/superpowers/plans/2026-03-27-gradient-stability-research.md +git commit -m "docs: gradient stability research — PCGrad, LayerNorm, adaptive weighting" +``` + +--- + +## Execution Dependencies + +``` +Task 1 (CQL isolation) → Task 2 (budget allocation) → Task 6 (smoke test) + ↗ +Task 3 (spec norm bufs) → Task 4 (spec norm launches) +Task 5 (Prometheus) ────────────────────────────────── +Task 7 (research) — independent +``` + +Tasks 1→2 fix the per-component clipping design. Tasks 3→4 add spectral norm to heads. Task 5 is independent monitoring. Task 6 validates everything. Task 7 is research. diff --git a/docs/superpowers/specs/2026-03-27-bf16-performance-design.md b/docs/superpowers/specs/2026-03-27-bf16-performance-design.md new file mode 100644 index 000000000..c122ea31b --- /dev/null +++ b/docs/superpowers/specs/2026-03-27-bf16-performance-design.md @@ -0,0 +1,522 @@ +# Spec A: Full BF16 Tensor Core Conversion + GPU Hot Path Performance + +**Date**: 2026-03-27 +**Spec**: A of 3 (DQN CUDA training pipeline overhaul) +**Target**: Epoch time <=5s on H100 (from ~30s), ~2.25x cuBLAS throughput gain +**Scope**: BF16 tensor cores, cached device pointers, zero-sync hot path, batched spectral norm, dead BF16 code removal + +## Problem Statement + +The DQN training pipeline uses `cublasSgemm` (F32) for all matrix multiplications in both forward and backward passes. On H100, BF16 tensor cores provide 989 TFLOPS vs 67 TFLOPS F32 scalar -- a 14.8x theoretical gap. With F32 accumulation (`CUBLAS_COMPUTE_32F`), practical BF16 throughput is ~3x over F32 SGEMM. + +BF16 weight mirror buffers (`bf16_params_buf`, `bf16_target_params_buf`) already exist in `GpuDqnTrainer` and are allocated at construction, but are **never read by the forward pass**. The forward reads F32 from `params_buf` via `cublasSgemm`. The conversion kernel `f32_to_bf16_kernel` exists in `common_device_functions.cuh` but is not invoked in the training loop. + +Additionally, the per-step hot path has severe overhead: + +| Issue | Count | Source | +|-------|------:|--------| +| `cuStreamSynchronize` calls | 4 | `apply_iqn_trunk_gradient:708`, `apply_ensemble_trunk_gradient:911`, `run_ensemble_step:875`, `replay_adam_and_readback:2369` | +| `raw_device_ptr()` calls | 110+ | cudarc event machinery per call | +| `EventTrackingGuard` create/drop | 11 | Disable/enable event tracking RAII cycles | +| Single-block spectral norm launches | 10 | `grid=(1,1,1)` each, serial kernel launch overhead | + +## Architecture + +### Current Data Flow (F32 Only) + +``` +Experience Collector (F32 states) + | + v +states_buf [B, SD] (F32) + | + v cublasSgemm (F32 x F32 -> F32) +CublasForward::forward_pass() + |-- h_s1 [B, SH1] = ReLU(states @ W_s1^T + b_s1) F32 + |-- h_s2 [B, SH2] = ReLU(h_s1 @ W_s2^T + b_s2) F32 + |-- h_v [B, VH] = ReLU(h_s2 @ W_v1^T + b_v1) F32 + |-- v_logits [B, NA] = h_v @ W_v2^T + b_v2 F32 + |-- h_a [B, AH] = ReLU(h_s2 @ W_a1^T + b_a1) F32 + |-- adv_logits [B, B0*NA] = h_a @ W_a2^T + b_a2 F32 + |-- (branches 1, 2 same pattern) + v +C51/MSE loss kernels (F32) + | + v cublasSgemm (F32 x F32 -> F32) +CublasBackward::backward_pass() + |-- dW = dY^T @ X (weight gradients, F32) + |-- dX = dY @ W^T (upstream gradients, F32) + v +grad_buf [TOTAL_PARAMS] (F32) + | + v +Adam update (F32) -> params_buf (F32) + | + v [NEVER HAPPENS] +f32_to_bf16_kernel -> bf16_params_buf (u16, UNUSED) +``` + +### Proposed Data Flow (BF16 Forward + Mixed Backward) + +``` +Experience Collector (F32 states) + | + v f32_to_bf16_kernel (in upload path) +states_bf16 [B, SD] (BF16) states_buf [B, SD] (F32, kept for loss) + | | + v cublasGemmEx (BF16 x BF16 -> F32) | +CublasForward::forward_pass() | + |-- h_s1_bf16 = BF16(ReLU(...)) | activations saved as BF16 + |-- h_s2_bf16 = BF16(ReLU(...)) | + |-- h_v_bf16 = BF16(ReLU(...)) | + |-- v_logits [B, NA] (F32) <---------' logit output stays F32 + |-- (branches same) for C51/MSE loss precision + v +C51/MSE loss kernels (F32 logits -> F32 d_logits) + | + v cublasGemmEx (BF16 weights x F32 grads -> F32, COMPUTE_32F = TF32) +CublasBackward::backward_pass() + |-- dW = d_output^T @ input_bf16 BF16 activations x F32 grads + |-- dX = d_output @ W_bf16^T BF16 weights x F32 grads + v +grad_buf [TOTAL_PARAMS] (F32) + | + v +Adam update (F32) -> params_buf (F32) + | + v f32_to_bf16_kernel (single launch, full buffer) +bf16_params_buf [TOTAL_PARAMS] (u16/BF16, used by NEXT forward) +``` + +### Precision Strategy + +| Path | Input A | Input B | Compute | Accumulate | Speedup | Rationale | +|------|---------|---------|---------|------------|---------|-----------| +| Forward SGEMM | BF16 weights | BF16 activations | `CUBLAS_COMPUTE_32F` | F32 | ~3x | Both inputs BF16 enables H100 tensor cores | +| Backward dW | F32 d_output | BF16 activations | `CUBLAS_COMPUTE_32F` | F32 | ~1.5x (TF32) | Mixed types -> TF32 path, preserves gradient precision | +| Backward dX | F32 d_output | BF16 weights | `CUBLAS_COMPUTE_32F` | F32 | ~1.5x (TF32) | Same as dW | +| Loss (C51/MSE) | F32 logits | -- | F32 | -- | 1x (no change) | Element-wise, not GEMM-bound | +| Adam | F32 grads | F32 moments | F32 | -- | 1x (no change) | Must be F32 for optimizer stability | +| F32->BF16 conv | F32 params | -- | -- | -- | N/A | Single kernel, ~0.1ms | + +**Combined GEMM speedup**: Forward 3x + Backward 1.5x = ~2.25x blended (forward is ~60% of GEMM time). + +## Detailed Changes + +### 1. BF16 Forward Pass (`batched_forward.rs`) + +Replace `cublasSgemm` with `cublasGemmEx` using BF16 inputs. + +**Current call** (`sgemm_layer`, line ~513): +```rust +cublas_result::sgemm( + self.handle.0, + CUBLAS_OP_T, CUBLAS_OP_N, + n as i32, b as i32, k as i32, + &1.0f32, + w_ptr as *const f32, k as i32, + a_ptr as *const f32, k as i32, + &0.0f32, + c_ptr as *mut f32, n as i32, +) +``` + +**Replacement** (`gemmex_bf16_layer`): +```rust +cublasGemmEx( + self.handle.0, + CUBLAS_OP_T, CUBLAS_OP_N, + n as i32, b as i32, k as i32, + &1.0f32 as *const f32 as *const c_void, // alpha (F32) + w_ptr as *const c_void, // A = BF16 weights + CUDA_R_16BF, // Atype + k as i32, // lda + a_ptr as *const c_void, // B = BF16 activations + CUDA_R_16BF, // Btype + k as i32, // ldb + &0.0f32 as *const f32 as *const c_void, // beta (F32) + c_ptr as *mut c_void, // C = F32 output + CUDA_R_32F, // Ctype + n as i32, // ldc + CUBLAS_COMPUTE_32F, // computeType + CUBLAS_GEMM_DEFAULT_TENSOR_OP, // algo +) +``` + +**New BF16 activation buffers** (allocated in `CublasForward::new()`): +- `states_bf16: CudaSlice` -- `[B, SD]` +- `h_s1_bf16: CudaSlice` -- `[B, SH1]` +- `h_s2_bf16: CudaSlice` -- `[B, SH2]` +- `h_v_bf16: CudaSlice` -- `[B, VH]` +- `h_b0_bf16, h_b1_bf16, h_b2_bf16: CudaSlice` -- `[B, AH]` each + +**Activation conversion**: After each bias+ReLU kernel (which outputs F32 into existing `h_s1`, etc.), add an inline F32->BF16 conversion step that writes into the `_bf16` variant. The F32 activation is kept for the backward pass (relu mask needs original F32 values). + +Alternatively, fuse the F32->BF16 cast into the `add_bias_relu_kernel`: +```cuda +// New: add_bias_relu_bf16_kernel +// Writes F32 to output AND BF16 to output_bf16 in one pass +extern "C" __global__ void add_bias_relu_bf16_kernel( + float* __restrict__ output, // F32 for backward relu mask + __nv_bfloat16* __restrict__ output_bf16, // BF16 for next layer's GEMM + const float* __restrict__ bias, + int out_dim, int total_elements) +``` + +**State upload conversion**: States arrive F32 from the experience collector. Add a `f32_to_bf16_kernel` launch in the upload path (between `cuMemcpyHtoDAsync` of states and the first SGEMM). This writes `states_bf16` from `states_buf`. Both buffers remain live -- `states_buf` (F32) is used by the backward pass, `states_bf16` by the forward GEMM. + +**Weight pointers**: Forward GEMM reads from `bf16_params_buf` (online) and `bf16_target_params_buf` (target) instead of `params_buf`. The `f32_weight_ptrs()` function is replaced with `bf16_weight_ptrs()` that returns u64 pointers into the BF16 flat buffer at precomputed byte offsets (already stored as `bf16_weight_offsets` in `GpuDqnTrainer`). + +**Logit output**: The final layer GEMMs (W_v2, W_a2, W_b1out, W_b2out) output F32 logits into existing buffers. `cublasGemmEx` with `Ctype=CUDA_R_32F` handles this natively -- the output type is independent of input types. + +### 2. Mixed-Precision Backward Pass (`batched_backward.rs`) + +**Weight gradient** `dW[out, in] = dY^T[out, B] @ X[B, in]`: +- `dY` is F32 (from C51/MSE loss gradient) +- `X` is BF16 (saved activation from forward) +- Call `cublasGemmEx` with `Atype=CUDA_R_32F`, `Btype=CUDA_R_16BF`, `Ctype=CUDA_R_32F`, `CUBLAS_COMPUTE_32F` +- cuBLAS internally uses TF32 for this mixed configuration (~1.5x over F32) + +**Upstream gradient** `dX[B, in] = dY[B, out] @ W[out, in]`: +- `dY` is F32 +- `W` is BF16 (from `bf16_params_buf`) +- Same `cublasGemmEx` mixed call +- Output `dX` stays F32 (needed for relu mask and next layer's backward) + +**No change**: bias gradient kernel, relu mask kernel, C51/MSE loss kernels, gradient accumulation -- all remain F32. + +### 3. Cached Device Pointers (`gpu_dqn_trainer.rs`) + +Create a `CachedPtrs` struct computed once at construction (and refreshed if buffers are reallocated, which never happens after init): + +```rust +/// Pre-resolved raw u64 device pointers for all GPU buffers. +/// Eliminates 110+ per-step `raw_device_ptr()` calls that go through +/// cudarc's event tracking machinery. +struct CachedPtrs { + // Flat parameter buffers + params_buf: u64, + target_params_buf: u64, + bf16_params_buf: u64, + bf16_target_params_buf: u64, + grad_buf: u64, + m_buf: u64, + v_buf: u64, + + // Activation buffers (F32) + states_buf: u64, + h_s1: u64, + h_s2: u64, + h_v: u64, + h_b0: u64, h_b1: u64, h_b2: u64, + on_v_logits: u64, + on_adv_logits: u64, + + // Activation buffers (BF16) -- new + states_bf16: u64, + h_s1_bf16: u64, + h_s2_bf16: u64, + h_v_bf16: u64, + h_b0_bf16: u64, h_b1_bf16: u64, h_b2_bf16: u64, + + // Target forward buffers + tgt_h_s1: u64, + tgt_h_s2: u64, + tgt_h_v: u64, + tgt_v_logits: u64, + tgt_adv_logits: u64, + + // Loss / optimizer + total_loss_buf: u64, + grad_norm_buf: u64, + td_errors_buf: u64, + t_buf: u64, + + // Spectral norm vectors + spec_u_s1: u64, spec_v_s1: u64, + spec_u_s2: u64, spec_v_s2: u64, + // ... (all 10 pairs) + + // IQN scratch + iqn_trunk_m: u64, + iqn_trunk_v: u64, + iqn_trunk_grad_norm: u64, + iqn_trunk_t_buf: u64, +} +``` + +**Construction**: After all `CudaSlice` allocations in `GpuDqnTrainer::new()`, resolve every pointer once with `raw_device_ptr()` and store in `CachedPtrs`. This requires disabling event tracking once during construction, not per-step. + +**Usage**: Replace every `raw_device_ptr(&self.some_buf, &self.stream)` with `self.ptrs.some_buf`. + +**EventTrackingGuard removal**: With cached pointers, no per-step `device_ptr()` calls go through cudarc's event machinery. Remove all 11 `EventTrackingGuard` instances from per-step methods. Keep the single disable/enable in the constructor. + +### 4. Zero-Sync Hot Path (`gpu_dqn_trainer.rs`, `fused_training.rs`) + +**Syncs to remove** (all on the same CUDA stream -- CUDA guarantees in-order execution): + +| Location | Line | Current Purpose | Why Safe to Remove | +|----------|------|-----------------|-------------------| +| `apply_iqn_trunk_gradient` | 708 | "Ensure graph_forward completed" | Same stream -- SAXPY kernel won't start until graph_forward's last kernel finishes | +| `apply_ensemble_trunk_gradient` | 911 | Same | Same stream ordering guarantee | +| `run_ensemble_step` | 875 | Same | Same stream ordering guarantee | + +**Sync to defer** (reads back scalars to CPU): + +| Location | Line | Current Purpose | Change | +|----------|------|-----------------|--------| +| `replay_adam_and_readback` | 2369 | Sync before DtoH scalar readback | Defer to every N steps or epoch boundary | + +**Deferred readback design**: Add `readback_interval: usize` config (default 100). Between readbacks, return the previous step's `(total_loss, grad_norm)` values. The loss/grad_norm are used for logging only -- stale values by 100 steps are acceptable for monitoring. At epoch boundaries, force a sync for accurate epoch-level metrics. + +```rust +pub fn replay_adam_and_readback(&mut self) -> Result { + self.replay_adam()?; + self.step_since_readback += 1; + + if self.step_since_readback >= self.readback_interval || self.force_readback { + unsafe { cudarc::driver::sys::cuStreamSynchronize(self.stream.cu_stream()); } + // DtoH readback of loss + grad_norm + unsafe { + cudarc::driver::sys::cuMemcpyDtoH_v2( + self.cached_loss.as_mut_ptr().cast(), + self.ptrs.total_loss_buf, 4, + ); + cudarc::driver::sys::cuMemcpyDtoH_v2( + self.cached_norm.as_mut_ptr().cast(), + self.ptrs.grad_norm_buf, 4, + ); + } + self.step_since_readback = 0; + self.force_readback = false; + } + + Ok(FusedTrainScalars { + total_loss: self.cached_loss[0], + grad_norm: self.cached_norm[0], + }) +} +``` + +### 5. Batched Spectral Norm (`dqn_utility_kernels.cu`, `gpu_dqn_trainer.rs`) + +**Current**: 10 individual kernel launches, each with `grid=(1,1,1), block=(256,1,1)`. + +**Proposed**: Single `batched_spectral_norm_kernel` with `grid=(10,1,1), block=(256,1,1)`. Each block handles one weight matrix, identified by `blockIdx.x` indexing into descriptor arrays. + +**New CUDA kernel**: +```cuda +struct SpectralNormDesc { + float* W; // weight matrix pointer + float* u; // left singular vector + float* v; // right singular vector + int out_dim; + int in_dim; + float sigma_max; +}; + +extern "C" __global__ void batched_spectral_norm_kernel( + SpectralNormDesc* __restrict__ descs, // [num_matrices] on device + int num_matrices +) { + int matrix_idx = blockIdx.x; + if (matrix_idx >= num_matrices) return; + + SpectralNormDesc d = descs[matrix_idx]; + // ... same power iteration logic as current spectral_norm_kernel, + // but reading W, u, v, out_dim, in_dim, sigma_max from d +} +``` + +**Descriptor upload**: Allocate `CudaSlice` at construction. Populate with cached device pointers from `CachedPtrs` and dimensions from config. Upload once (or recompute if pointers change, which they don't post-init). + +**Rust-side change** in `apply_spectral_norm()`: +```rust +// Before: 10 separate launch_builder() calls +// After: +unsafe { + self.stream + .launch_builder(&self.batched_spectral_norm_kernel) + .arg(&self.ptrs.spec_norm_descs) + .arg(&10i32) + .launch(LaunchConfig { + grid_dim: (10, 1, 1), + block_dim: (256, 1, 1), + shared_mem_bytes: 256 * 4, + })?; +} +``` + +This eliminates 9 kernel launch overhead cycles (~5-10us each on H100). + +### 6. Dead BF16 Infrastructure Removal (`gpu_weights.rs`, `gpu_dqn_trainer.rs`) + +**Delete from `gpu_weights.rs`** (~200 LOC): +- `struct DuelingWeightSetBf16` (line 473) -- 12 per-tensor `CudaSlice` fields +- `struct BranchingWeightSetBf16` (line 490) -- 8 per-tensor `CudaSlice` fields +- `struct CuriosityWeightSetBf16` (line 503) -- 4 per-tensor `CudaSlice` fields +- `fn alloc_bf16_mirror()` (line 511) -- allocates matching u16 buffer +- `pub fn convert_f32_to_bf16()` (line 527) -- per-tensor conversion loop +- `impl DuelingWeightSetBf16` (line 554) -- `from_f32()` + `sync_from_f32()` +- `impl BranchingWeightSetBf16` (line 596) -- `from_f32()` + `sync_from_f32()` +- `impl CuriosityWeightSetBf16` (line 630) -- `from_f32()` + `sync_from_f32()` +- `struct KernelWeightPackBf16` (line 670) -- 24-pointer pack for the old fused kernel + +**Delete from `gpu_dqn_trainer.rs`**: +- `fn ensure_bf16_mirrors()` (line 2144) -- allocates per-tensor BF16 mirrors +- `fn sync_online_bf16()` (line 2237) -- syncs per-tensor F32->BF16 + +**Keep** (used by the new BF16 forward path): +- `bf16_params_buf: CudaSlice` -- flat online BF16 buffer, read by `cublasGemmEx` +- `bf16_target_params_buf: CudaSlice` -- flat target BF16 buffer +- `bf16_weight_offsets: Vec` -- byte offsets into flat BF16 buffer +- `bf16_mirrors_initialized: bool` -- tracks if BF16 mirrors are populated +- `f32_to_bf16_kernel: CudaFunction` -- used for bulk conversion +- `bf16_to_f32_kernel: CudaFunction` -- used for checkpointing reverse sync + +**Rationale**: The per-tensor BF16 structs were designed for the old fused 1-warp-per-sample kernel that loaded weights from global memory. The cuBLAS path reads from the flat `bf16_params_buf` at precomputed offsets. The per-tensor mirrors are 100% redundant with the flat buffer. + +### 7. BF16 Conversion in Adam Graph (`gpu_dqn_trainer.rs`) + +After the Adam update writes to `params_buf` (F32), append an `f32_to_bf16_kernel` launch to `graph_adam` that converts the entire `params_buf` -> `bf16_params_buf`. This is captured in the CUDA Graph so it has zero launch overhead on replay. + +```rust +// In build_graph_adam() -- after adam_update_kernel, before graph capture ends: +unsafe { + self.stream + .launch_builder(&self.f32_to_bf16_kernel) + .arg(&self.ptrs.params_buf) + .arg(&self.ptrs.bf16_params_buf) + .arg(&(total_params as i32)) + .launch(LaunchConfig { + grid_dim: (ceildiv(total_params, 256) as u32, 1, 1), + block_dim: (256, 1, 1), + shared_mem_bytes: 0, + })?; +} +``` + +Target BF16 buffer (`bf16_target_params_buf`) is updated during EMA sync -- add the same `f32_to_bf16_kernel` launch after the existing `ema_kernel` that syncs `target_params_buf`. + +## Files Changed + +| File | Changes | +|------|---------| +| `crates/ml/src/cuda_pipeline/batched_forward.rs` | Replace `cublasSgemm` -> `cublasGemmEx` with BF16 inputs; add BF16 activation buffers; new `gemmex_bf16_layer()` method; new `bf16_weight_ptrs()` helper; fused `add_bias_relu_bf16_kernel` | +| `crates/ml/src/cuda_pipeline/batched_backward.rs` | Replace `cublasSgemm` -> `cublasGemmEx` with mixed BF16/F32 inputs; read BF16 activations and weights | +| `crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs` | Add `CachedPtrs` struct and initialization; remove 3 `cuStreamSynchronize` calls; add deferred readback with interval; replace 10 spectral norm launches with 1 batched launch; add `f32_to_bf16_kernel` to `graph_adam`; delete `ensure_bf16_mirrors()`, `sync_online_bf16()`; remove 11 `EventTrackingGuard` instances | +| `crates/ml/src/cuda_pipeline/gpu_weights.rs` | Delete `DuelingWeightSetBf16`, `BranchingWeightSetBf16`, `CuriosityWeightSetBf16`, `KernelWeightPackBf16`, `alloc_bf16_mirror()`, `convert_f32_to_bf16()` and all associated impls (~200 LOC) | +| `crates/ml/src/trainers/dqn/fused_training.rs` | Remove `cuStreamSynchronize` from `run_ensemble_step`; remove local `EvtGuard`; use `CachedPtrs` | +| `crates/ml/src/cuda_pipeline/dqn_utility_kernels.cu` | Add `SpectralNormDesc` struct and `batched_spectral_norm_kernel`; keep existing `spectral_norm_kernel` (backward compat for tests) | +| `crates/ml/src/cuda_pipeline/bias_kernels.cu` | Add `add_bias_relu_bf16_kernel` variant that outputs both F32 and BF16 | + +## Performance Budget + +Estimated per-step breakdown on H100 (300 steps/epoch, B=256, ~288K params): + +| Component | Current | After | Savings | +|-----------|--------:|------:|--------:| +| Forward (3 passes x 10 GEMMs) | ~40ms | ~13ms | 3x BF16 tensor cores | +| Backward (10 GEMMs) | ~30ms | ~20ms | 1.5x TF32 mixed | +| `cuStreamSynchronize` x4 | ~8ms | ~0ms | Removed 3, deferred 1 | +| `raw_device_ptr()` x110 | ~5ms | ~0ms | Cached u64 lookups | +| `EventTrackingGuard` x11 | ~2ms | ~0ms | Eliminated | +| Spectral norm (10 launches) | ~1ms | ~0.2ms | 1 batched launch | +| F32->BF16 conversion | 0ms | ~0.1ms | New, but in CUDA Graph | +| Loss + optimizer | ~15ms | ~15ms | No change | +| **Total per step** | **~101ms** | **~48ms** | **~2.1x** | +| **Per epoch (300 steps)** | **~30s** | **~14s** | -- | + +Note: The 14s estimate is conservative. Combined with experience collection overlap (already async) and the H100's 3 TB/s HBM3, actual epoch time should be <=5s. + +## Testing Plan + +### Unit Tests (must all pass, 129+ existing) + +```bash +SQLX_OFFLINE=true cargo test -p ml --lib +``` + +All existing tests use the same `GpuDqnTrainer` code path. BF16 conversion is invisible to test assertions since loss/Q-values are computed in F32. + +### Smoke Test (gradient stability) + +```bash +FOXHUNT_TEST_DATA=test_data/futures-baseline cargo test -p ml --lib -- smoke_tests --ignored --nocapture +``` + +Verify: +- Gradient norms remain bounded (no NaN/Inf from BF16 underflow) +- Q-values converge (loss decreasing trend over 10+ epochs) +- No CUDA errors in stream after 50 epochs + +### BF16 Precision Regression Test (new) + +Add a test that runs 50 training epochs with both F32-only and BF16 paths on the same seed, comparing: +- Final Sharpe ratio: must be within 5% relative difference +- Mean Q-value: must be within 10% relative difference +- Gradient norm statistics: mean and max must be within 20% + +This test validates that BF16 forward + TF32 backward does not degrade convergence quality. + +### Cached Pointer Correctness Test (new) + +Add a test that: +1. Constructs `GpuDqnTrainer` +2. Verifies all `CachedPtrs` fields match `raw_device_ptr()` calls +3. Runs 10 training steps +4. Verifies pointers haven't changed (buffers are fixed post-init) + +### Batched Spectral Norm Equivalence Test (new) + +Add a test that: +1. Initializes random weights + u/v vectors +2. Runs 10 individual `spectral_norm_kernel` launches (current path) +3. Runs 1 `batched_spectral_norm_kernel` launch (new path) +4. Compares all weight matrices element-wise (must be bitwise identical since F32 throughout) + +### Deferred Readback Test (new) + +Add a test that: +1. Runs 200 training steps with `readback_interval=100` +2. Verifies sync happens at steps 100 and 200 +3. Verifies returned scalars at step 100 match actual GPU values +4. Verifies steps 1-99 return cached values without blocking + +## Success Criteria + +1. **Epoch time <=5s** on H100 (80GB, SM90) with B=256, 300 steps/epoch +2. **All 129+ existing tests pass** without modification +3. **Gradient norms bounded**: no NaN/Inf over 50 epochs on smoke test data +4. **Convergence parity**: Sharpe ratio within 5% of F32 baseline over 50 epochs +5. **Zero regression**: no new CUDA errors, no memory leaks (same peak VRAM +/- 10%) + +## Non-Goals + +- Bug fixes in training logic (Spec B scope) +- Dead code cleanup beyond BF16 infrastructure (Spec C scope) +- Full BF16 backward pass (TF32 mixed precision is acceptable and preserves gradient quality) +- FP16 support (BF16 has superior dynamic range for training; FP16 risks gradient underflow) +- Multi-GPU changes (single-GPU H100 is the target) + +## Risks and Mitigations + +| Risk | Impact | Mitigation | +|------|--------|------------| +| `cublasGemmEx` not capturable in CUDA Graph | Breaks graph_forward capture | Verify with CUDA 12.x docs; fallback: call outside graph (lose ~2ms/step launch overhead) | +| BF16 activation underflow for small features | NaN propagation | Monitor activation statistics; BF16 range is +-3.4e38 (same exponent bits as F32), underflow only at <1e-38 which doesn't occur in normalized features | +| Cached pointer invalidation after cudarc GC | Wrong memory access | cudarc doesn't GC/move device allocations; verify with `assert_eq!(cached, fresh)` in debug builds | +| Batched spectral norm shared memory overflow | Kernel crash for large matrices | Current max matrix is 256x256; `shmem[256]` is sufficient; add `assert(out_dim <= 256 && in_dim <= 256)` in kernel | +| Deferred readback masking training divergence | Late detection of NaN loss | Force readback on NaN detection (check `isnan(cached_loss)` and force sync if true) | + +## Implementation Order + +1. **CachedPtrs** -- lowest risk, enables all subsequent changes +2. **Dead BF16 code removal** -- reduces confusion, no behavior change +3. **Batched spectral norm** -- standalone kernel change, easy to test +4. **Zero-sync hot path** -- remove 3 syncs, defer 1 +5. **BF16 forward pass** -- core change, requires activation buffers +6. **BF16 backward (mixed)** -- depends on forward BF16 activations +7. **F32->BF16 in graph_adam** -- closes the loop + +Steps 1-4 can be landed as a single commit (performance-only, no precision change). Steps 5-7 are a second commit (precision change, requires BF16 regression test). diff --git a/docs/superpowers/specs/2026-03-27-bug-fixes-design.md b/docs/superpowers/specs/2026-03-27-bug-fixes-design.md new file mode 100644 index 000000000..85520898e --- /dev/null +++ b/docs/superpowers/specs/2026-03-27-bug-fixes-design.md @@ -0,0 +1,703 @@ +# Design Spec B: DQN CUDA Training Pipeline -- Correctness Bug Fixes + +**Date**: 2026-03-27 +**Spec**: B of 3 (A = Performance, B = Correctness, C = Dead code) +**Scope**: 9 CRITICAL + 14 IMPORTANT bugs from 9-agent audit +**Non-goals**: Performance optimization (Spec A), dead code removal (Spec C) + +--- + +## 1. Problem Statement + +A 9-agent audit of the DQN CUDA training pipeline uncovered 23 correctness +issues. Nine are CRITICAL (produce wrong training results), and 14 are IMPORTANT +(hide problems via silent fallbacks). One item initially flagged as critical +(NoisyNet in GPU action selector) was reclassified as NOT A BUG after +investigation -- see Section 3.9 for the architectural rationale. + +The bugs fall into four categories: + +1. **Incomplete state resets** -- backtest kernel resets 5 of 8 portfolio fields + on episode termination, while the experience kernel resets all 20. Stale + `max_equity`, `hold_time`, and `entry_price` leak across episodes. +2. **Inconsistent hyperparameter defaults** -- gradient clip norm defaults to + 10.0 in the main trainer but 1.0 in IQN/IQL sub-trainers; entropy + coefficient is `Option` with three different fallback values across + constructors. +3. **Silent fallbacks masking data errors** -- `unwrap_or(1.0)` on close prices, + `unwrap_or(0.0)` on quantile predictions, `unwrap_or(2)` on action indices. + A single bad data row produces wrong rewards for an entire episode without + any log message. +4. **Side effects in read-only paths** -- `step_count` increments during + inference forward passes, and CUDA contexts are created 3x per quantile + loss call. + +--- + +## 2. Files Changed + +| File | Bugs | +|------|------| +| `crates/ml/src/cuda_pipeline/backtest_env_kernel.cu` | #1, #2, #11 | +| `crates/ml/src/trainers/dqn/data_loading.rs` | #3, #15 | +| `crates/ml/src/trainers/dqn/fused_training.rs` | #4 | +| `crates/ml/src/trainers/dqn/trainer/constructor.rs` | #5 | +| `crates/ml-dqn/src/quantile_regression.rs` | #6, #7 | +| `crates/ml-dqn/src/network.rs` | #8 | +| `crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs` | #10 | +| `crates/ml/src/cuda_pipeline/experience_kernels.cu` | #14 | +| `crates/ml-dqn/src/branching.rs` | #16 | +| `crates/ml/src/trainers/dqn/trainer/action.rs` | #15 | +| `crates/ml/src/trainers/dqn/trainer/state.rs` | #15 | +| `crates/ml/src/cuda_pipeline/gpu_backtest_evaluator.rs` | #12, #13 | +| `crates/ml/src/trainers/dqn/config.rs` | #5 | + +--- + +## 3. CRITICAL Bugs (9 items -- wrong results) + +### 3.1 Backtest incomplete episode reset + +**Location**: `backtest_env_kernel.cu:105-112` (pre-trade floor) and +`backtest_env_kernel.cu:199-206` (post-trade floor) + +**Problem**: The backtest kernel uses `PORTFOLIO_STATE_SIZE = 8` fields per +window. On capital floor breach, both reset blocks only partially reinitialize +the state: + +``` +// Current (WRONG) -- pre-trade floor reset (lines 105-112): +portfolio_state[ps + 0] = liq_value; // value +portfolio_state[ps + 1] = 0.0f; // position +portfolio_state[ps + 2] = liq_value; // cash +portfolio_state[ps + 3] = 0.0f; // entry_price +portfolio_state[ps + 4] = max_equity; // STALE -- not updated to new_value +portfolio_state[ps + 5] = 0.0f; // hold_time +portfolio_state[ps + 6] = cum_return + liq_ret; // ACCUMULATES across episodes +portfolio_state[ps + 7] += 1.0f; // step_count INCREMENTS, never resets +``` + +Compare with the experience kernel (`experience_kernels.cu:614-630`) which +resets all 20 fields to fresh-episode values: position=0, cash=initial_capital, +portfolio_value=initial_capital, peak_equity=initial_capital, and zeroes all +counters. + +**Impact**: After a capital floor breach, the backtest continues the next +episode with (a) stale `max_equity` from the blown episode, which prevents the +floor check from firing correctly in subsequent episodes, (b) accumulated +`cum_return` that conflates multiple episodes, and (c) a `step_count` that +never resets, making per-episode metrics meaningless. + +**Fix**: Reset all 8 backtest portfolio fields to initial-episode state at both +reset sites. Use `liq_value` (or `new_value` for the post-trade block) as the +new initial capital: + +``` +// Pre-trade floor reset (lines 105-112): +portfolio_state[ps + 0] = liq_value; // value = initial capital for next ep +portfolio_state[ps + 1] = 0.0f; // position = flat +portfolio_state[ps + 2] = liq_value; // cash = initial capital +portfolio_state[ps + 3] = 0.0f; // entry_price = none +portfolio_state[ps + 4] = liq_value; // max_equity = reset to new initial +portfolio_state[ps + 5] = 0.0f; // hold_time = 0 +portfolio_state[ps + 6] = 0.0f; // cum_return = reset for new episode +portfolio_state[ps + 7] = 0.0f; // step_count = reset for new episode + +// Post-trade floor reset (lines 199-206): identical pattern with new_value +portfolio_state[ps + 0] = new_value; +portfolio_state[ps + 1] = 0.0f; +portfolio_state[ps + 2] = new_value; +portfolio_state[ps + 3] = 0.0f; +portfolio_state[ps + 4] = new_value; // max_equity = reset +portfolio_state[ps + 5] = 0.0f; +portfolio_state[ps + 6] = 0.0f; // cum_return = reset +portfolio_state[ps + 7] = 0.0f; // step_count = reset +``` + +**Testing**: Add a backtest scenario where a window triggers the capital floor. +Verify that the next episode starts with fresh `max_equity` equal to the +liquidated equity, `cum_return = 0`, and `step_count = 0`. + +--- + +### 3.2 Backtest stale max_equity before post-trade floor check + +**Location**: `backtest_env_kernel.cu:184` + +**Problem**: The post-trade capital floor check at line 184 runs +`check_capital_floor(new_value, max_equity)` using the *pre-trade* `max_equity`. +If the trade was profitable enough to set a new high-water mark, the floor +check uses a stale denominator, potentially missing a breach or triggering a +false one. + +The normal (non-done) path updates `max_equity` at line 214: +`float new_max = fmaxf(max_equity, new_value);` -- but this happens *after* +the floor check. + +**Fix**: Insert the max_equity update before the floor check: + +``` +float new_value = cash + position * close; + +// Update max_equity BEFORE floor check +max_equity = fmaxf(max_equity, new_value); + +if (check_capital_floor(new_value, max_equity)) { + ... +``` + +**Testing**: Construct a scenario where a profitable trade raises equity above +the previous peak, then verify the floor check uses the updated peak. + +--- + +### 3.3 Close price defaults to 1.0 + +**Location**: `data_loading.rs:306` + +**Problem**: When computing simple returns, the code uses: +```rust +let prev = close_prices_f32.get(i - 1).copied().unwrap_or(1.0); +let curr = close_prices_f32.get(i).copied().unwrap_or(prev); +``` + +If `close_prices_f32` contains any parse failures that result in a missing +value, the fallback `1.0` is a valid-looking but completely wrong price +(ES trades around 5000). The resulting return `(5000 - 1) / 1 = 4999` or +`(1 - 5000) / 5000 = -0.9998` will dominate the entire episode's reward +signal. + +Since `close_prices_f32` is an indexed `Vec` accessed by in-bounds index +`i-1` and `i` where `i in 1..n`, the `.get()` calls can never actually return +`None` -- the indices are always valid. However, the `unwrap_or(1.0)` gives +a false sense of safety and suppresses any future index bugs. + +**Fix**: Replace `unwrap_or` with direct indexing (which panics on OOB, +correctly surfacing the bug): + +```rust +let prev = close_prices_f32[i - 1]; +let curr = close_prices_f32[i]; +``` + +Alternatively, use `expect("close price index in bounds")` if a message is +preferred over a bare panic. + +**Testing**: Existing tests cover the happy path. Add a `#[should_panic]` test +that calls the return computation with an empty `close_prices_f32` to confirm +the panic behavior. + +--- + +### 3.4 gradient_clip_norm inconsistent defaults + +**Location**: `fused_training.rs:194,263,298` + +**Problem**: Three config construction sites read `gradient_clip_norm` from the +same `Option` hyperparameter but use different defaults: + +| Site | Default | Component | +|------|---------|-----------| +| `fused_training.rs:194` | `unwrap_or(10.0)` | Main C51 trainer | +| `fused_training.rs:263` | `unwrap_or(1.0)` | IQL value trainer | +| `fused_training.rs:298` | `unwrap_or(1.0)` | IQN dual-head trainer | +| `constructor.rs:257` | `unwrap_or(10.0)` | CPU-side DQN config | + +When hyperopt omits `gradient_clip_norm`, the main trainer clips at 10.0 while +IQN/IQL clip at 1.0. The gradient budget fractions in `gpu_dqn_trainer.rs` +(line 857: `max_grad_norm * 0.10` for IQN, line 1098: `max_grad_norm * 0.05` +for ensemble) assume the main norm is the ceiling. If the sub-trainers +independently clip to 1.0, the total gradient budget accounting is wrong. + +**Fix**: Make `gradient_clip_norm` required (non-optional) in `DqnHyperparams`. +Set the canonical default in `DqnHyperparams::default()` to `10.0`. Remove all +`unwrap_or()` fallbacks -- the value is always present: + +```rust +// In DqnHyperparams: +pub gradient_clip_norm: f64, // was Option + +// All construction sites: +max_grad_norm: hyperparams.gradient_clip_norm as f32, +``` + +If making the field non-optional is too disruptive for serialization +compatibility, keep `Option` but resolve it once at the top of +`FusedTrainingCtx::new()` and pass the resolved value to all sub-configs: + +```rust +let grad_norm = hyperparams.gradient_clip_norm.unwrap_or(10.0); +// Use `grad_norm` everywhere -- single source of truth +``` + +**Testing**: Verify that a training run with `gradient_clip_norm = None` +produces the same effective clip norm (10.0) in all three sub-trainers. Add an +assertion in `FusedTrainingCtx::new()` that the IQN and IQL configs use the +same `max_grad_norm` as the main config. + +--- + +### 3.5 entropy_coefficient Option mismatch + +**Location**: `config.rs:960`, `constructor.rs:314,399`, `dqn.rs:287,740`, +`hyperopt/adapters/dqn.rs:418,2553` + +**Problem**: The `entropy_coefficient` field has a fragmented type/default chain: + +| Location | Type | Default | +|----------|------|---------| +| `DqnHyperparams` (`config.rs:960`) | `Option` | `None` | +| `DQNConfig` (`dqn.rs:287`) | `f64` (non-optional) | `0.001` | +| `constructor.rs:314` | unwraps Option | `unwrap_or(0.01)` | +| `constructor.rs:399` | unwraps Option | `unwrap_or(0.01)` | +| Hyperopt default (`adapters/dqn.rs:418`) | `f64` | `0.02` | +| Hyperopt search range (`adapters/dqn.rs:517`) | `f64` | `(0.05, 0.5)` | +| Hyperopt to hyperparams (`adapters/dqn.rs:2553`) | `Some(f64)` | from PSO | + +The constructor's `unwrap_or(0.01)` disagrees with `DQNConfig::default()` +(0.001) and with hyperopt's default (0.02). When hyperopt runs, it always +provides a value via `Some(params.entropy_coefficient)`, so the disagreement +only matters for manual (non-hyperopt) training runs. + +**Fix**: Make `entropy_coefficient` non-optional in `DqnHyperparams`. Set the +canonical default to match the hyperopt search space center. Remove the +`unwrap_or` in the constructor: + +```rust +// In DqnHyperparams (config.rs): +pub entropy_coefficient: f64, // was Option + +// In constructor.rs: +entropy_coefficient: hyperparams.entropy_coefficient, +``` + +Update `DqnHyperparams::default()` to use `0.01` as the default (midpoint +between DQNConfig's 0.001 and hyperopt's 0.02 on log scale, and matching the +old `unwrap_or`). + +**Testing**: Grep for all `entropy_coefficient.*unwrap` after the change and +confirm zero hits in production code. + +--- + +### 3.6 3x CUDA context creation in quantile loss + +**Location**: `quantile_regression.rs:320-346` + +**Problem**: `quantile_huber_loss()` creates a new `CudaContext::new(0)` three +separate times in a single call: + +1. Line 320-326: for the empty-batch early return `GpuTensor::scalar(0.0)` +2. Line 332-338: for `per_sample.to_host()` stream +3. Line 341-346: for the final `GpuTensor::scalar(mean_loss)` + +Each `CudaContext::new(0)` call invokes `cuCtxCreate` or +`cuDevicePrimaryCtxRetain`, which is a heavyweight driver call. On the hot +training path, this is called once per batch. + +Additionally, `quantile_huber_loss_per_sample()` creates its own context at +line 369-374, for a total of 4 context creations per quantile loss computation. + +**Fix**: Create the context once at the top of `quantile_huber_loss()` and pass +it (or the stream derived from it) to all operations: + +```rust +pub fn quantile_huber_loss(...) -> Result { + let ctx = cudarc::driver::CudaContext::new(0) + .map_err(|e| MLError::DeviceError(format!("CUDA context: {e}")))?; + let stream = ctx.new_stream() + .map_err(|e| MLError::DeviceError(format!("CUDA stream: {e}")))?; + + let per_sample = quantile_huber_loss_per_sample_with_stream( + predicted, target, taus, kappa, &stream + )?; + + let batch = per_sample.numel(); + if batch == 0 { + return GpuTensor::scalar(0.0, &stream); + } + + let host = per_sample.to_host(&stream)?; + let mean_loss: f32 = host.iter().sum::() / batch as f32; + GpuTensor::scalar(mean_loss, &stream) +} +``` + +Extract the inner loop of `quantile_huber_loss_per_sample` into a +`_with_stream` variant that accepts an existing `&CudaStream`. + +**Testing**: Existing quantile loss tests cover correctness. No new tests +needed -- this is a pure refactor. Verify tests pass after the change. + +--- + +### 3.7 Silent 0.0/0.5 fallback in quantile loss + +**Location**: `quantile_regression.rs:385-387` + +**Problem**: Inside the per-sample quantile Huber loss loop: + +```rust +let pred_val = pred_host.get(idx).copied().unwrap_or(0.0); +let tgt_val = tgt_host.get(idx).copied().unwrap_or(0.0); +let tau_val = tau_host.get(idx).copied().unwrap_or(0.5); +``` + +The index `idx = b * num_q + q` with `b in 0..batch` and `q in 0..num_q` +should always be in bounds since `pred_host` has exactly `batch * num_q` +elements. If it is ever out of bounds, the loss computation silently uses +`pred=0.0`, `target=0.0` (producing zero loss -- masking real errors) or +`tau=0.5` (symmetric quantile -- incorrect for risk-sensitive IQN). + +**Fix**: Replace `unwrap_or` with direct indexing or `expect`: + +```rust +let pred_val = pred_host[idx]; +let tgt_val = tgt_host[idx]; +let tau_val = tau_host[idx]; +``` + +If bounds checking is desired for safety, use: + +```rust +let pred_val = *pred_host.get(idx).ok_or_else(|| { + MLError::ModelError(format!( + "pred OOB: idx={idx}, len={}", pred_host.len() + )) +})?; +``` + +**Testing**: Existing tests. Add a test with mismatched tensor shapes to verify +the error is returned (not silently swallowed). + +--- + +### 3.8 step_count increments during inference + +**Location**: `network.rs:250` + +**Problem**: The `forward()` method on the DQN network unconditionally +increments `step_count`: + +```rust +let _step = self.step_count.fetch_add(1, Ordering::Relaxed); +``` + +This counter is used by the dropout scheduler (line 253-257) and potentially +by other scheduling logic. During inference (backtest, action selection during +epsilon-greedy exploration), the forward pass is read-only -- it should not +advance the training schedule. + +Every inference forward pass advances the dropout schedule by one step, causing +dropout to decay faster than intended. + +**Fix**: Add a `training: bool` parameter to `forward()`, or split into +`forward_train()` and `forward_eval()`: + +Option A -- parameter: +```rust +pub fn forward(&self, state: &[f32], training: bool) -> Result, MLError> { + // ... existing logic ... + if training { + let _step = self.step_count.fetch_add(1, Ordering::Relaxed); + if let Ok(mut scheduler_opt) = self.dropout_scheduler.lock() { + if let Some(scheduler) = scheduler_opt.as_mut() { + scheduler.step(1); + } + } + } + Ok(output_vec) +} +``` + +Option B -- separate methods (preferred for API clarity): +```rust +pub fn forward_impl(&self, state: &[f32], training: bool) -> Result, MLError> { + // ... shared logic with conditional step_count ... +} + +pub fn forward_train(&self, state: &[f32]) -> Result, MLError> { + self.forward_impl(state, true) +} + +pub fn forward_for_eval(&self, state: &[f32]) -> Result, MLError> { + self.forward_impl(state, false) +} +``` + +All inference call sites (backtest, action selection in greedy mode) must call +the non-training variant. All training call sites must call the train variant. + +**Testing**: Call the non-training forward 1000 times and verify `step_count` +does not change. Call the training forward once and verify it increments by 1. + +--- + +### 3.9 NoisyNet noise in GPU action selector -- NOT A BUG + +**Location**: `gpu_action_selector.rs` + +**Initial finding**: The GPU action selector uses pure epsilon-greedy; there is +no explicit Gaussian noise injection for NoisyNet exploration. + +**Analysis**: NoisyNet (Fortunato et al., 2018) adds parametric noise to +linear layer weights. In this codebase, `NoisyLinear` layers are part of the +`BranchingDuelingNetwork`. The noise is injected into the weight matrices +during the forward pass, which means the Q-values output by the network +*already incorporate NoisyNet exploration*. The action selector receives +noise-perturbed Q-values and simply takes the argmax (or applies +epsilon-greedy as a fallback). + +This is the standard NoisyNet architecture. The selector does not need +separate noise injection -- that would double-count the exploration. + +**Verdict**: Not a bug. Document this architectural choice with a comment in +`gpu_action_selector.rs`: + +```rust +// NoisyNet exploration note: Q-values received here already include +// factorized Gaussian noise from NoisyLinear layers in the forward pass. +// Epsilon-greedy acts as a secondary fallback, not the primary exploration +// mechanism when NoisyNet is enabled. +``` + +--- + +## 4. IMPORTANT Bugs (14 items -- hiding problems) + +### 4.10 Budget fractions hardcoded in gpu_dqn_trainer + +**Location**: `gpu_dqn_trainer.rs:857,1098` + +**Problem**: +```rust +let max_component_norm = self.config.max_grad_norm * 0.10; // IQN_GRAD_BUDGET +let max_component_norm = self.config.max_grad_norm * 0.05; // ENS_GRAD_BUDGET +``` + +These magic numbers define what fraction of the total gradient budget is +allocated to auxiliary components (IQN, ensemble). They should be configurable +or at least centralized as named constants. + +**Fix**: Define constants in `fused_training.rs`: + +```rust +/// Fraction of max_grad_norm allocated to IQN gradient updates. +pub const IQN_GRAD_BUDGET_FRACTION: f32 = 0.10; + +/// Fraction of max_grad_norm allocated to ensemble gradient updates. +pub const ENSEMBLE_GRAD_BUDGET_FRACTION: f32 = 0.05; +``` + +Reference these from `gpu_dqn_trainer.rs`. Alternatively, add them to the +`FusedConfig` struct if hyperopt should search over them. + +--- + +### 4.11 Trailing stop distance hardcoded + +**Location**: `backtest_env_kernel.cu:151` + +**Problem**: `check_trailing_stop(..., 0.005f, ...)` -- the 0.5% base distance +is ES-specific and not configurable via kernel arguments. + +**Fix**: Add `trailing_stop_base_dist` to the kernel parameter list and pass +it from `GpuBacktestConfig`. Default to `0.005` for backward compatibility. + +--- + +### 4.12-4.13 ES-specific defaults in GpuBacktestConfig + +**Location**: `gpu_backtest_evaluator.rs:205,211` + +**Problem**: `contract_multiplier: 50.0` and `margin_pct: 0.06` are ES-specific +defaults. These are already configurable fields in `GpuBacktestConfig`, so the +fix is documentation, not code. + +**Fix**: Add doc comments to both fields in the `Default` impl: + +```rust +/// ES mini S&P 500 default. Override for other instruments: +/// NQ=20.0, CL=1000.0, GC=100.0, etc. +contract_multiplier: 50.0, +/// ES initial margin percentage (CME). Override for other instruments. +margin_pct: 0.06, +``` + +--- + +### 4.14 Kelly min trades hardcoded + +**Location**: `experience_kernels.cu:670` + +**Problem**: `if (total_trades >= 20.0f)` -- the minimum number of trades +before Kelly criterion activates is hardcoded. For different bar frequencies +(daily vs 1-min), 20 trades may be too few or too many. + +**Fix**: Add `kelly_min_trades` as a kernel parameter. Default to `20` for +backward compatibility. Pass from `DqnBacktestConfig` or derive from +`bars_per_day`. + +--- + +### 4.15 Silent unwrap_or fallbacks in production paths + +**Priority tiers**: + +**Tier 1 -- data correctness (fix now)**: +- `data_loading.rs:306` -- `unwrap_or(1.0)` on close price (covered in #3) +- `action.rs:161` -- `unwrap_or(2)` on argmax result (flat action fallback). + The argmax should never fail on a non-empty tensor. Replace with + `.expect("argmax on non-empty Q-values")`. +- `state.rs:73` -- `unwrap_or(0.0)` on price `to_f32()`. A zero price + produces zero returns. Replace with `.expect("price fits f32")` since + trading prices always fit in f32. + +**Tier 2 -- metrics/monitoring (fix later, lower risk)**: +- `training_stability.rs` (15 instances) -- test code using `unwrap_or(f64::NAN)` + on metric lookups. These are acceptable since NaN propagates to test failures. +- `metrics.rs` (10 instances) -- monitoring aggregation using `unwrap_or(0.0)`. + Missing metrics produce zeros in dashboards, which is misleading but not + training-breaking. Replace with `unwrap_or(f64::NAN)` where appropriate. +- `config.rs:300,419,773,774,788,789` -- shape/dimension lookups with fallbacks + to 0 or 256. These are initialization-time and generally safe. + +**Tier 3 -- test code (leave as-is)**: +- `smoke_tests/helpers.rs`, `tests.rs` -- test utilities. Fallbacks are + acceptable in test code. + +--- + +### 4.16 MaybeNoisyLinear single-variant enum + +**Location**: `branching.rs:252-324` + +**Problem**: The `MaybeNoisyLinear` enum has exactly one variant: + +```rust +enum MaybeNoisyLinear { + Noisy(NoisyLinear), +} +``` + +A `Linear` variant was presumably removed when the architecture converged on +always-noisy layers. The enum wrapper adds indirection, match arms, and +cognitive overhead for no benefit. + +**Fix**: Replace all occurrences of `MaybeNoisyLinear` with `NoisyLinear` +directly. The methods on `MaybeNoisyLinear` (`forward`, `reset_noise`, +`reset_noise_with_sigma`, `disable_noise`, `register_mu_in_varstore`, +`noisy_sigma_slices`) are thin wrappers -- inline them. + +Fields in `BranchingDuelingNetwork` change from: +```rust +value_fc: MaybeNoisyLinear, +value_out: MaybeNoisyLinear, +branch_fcs: Vec, +branch_outs: Vec, +``` +to: +```rust +value_fc: NoisyLinear, +value_out: NoisyLinear, +branch_fcs: Vec, +branch_outs: Vec, +``` + +**Testing**: Compile-only change. All existing tests cover the behavior. + +--- + +### 4.17-4.20 Lower-priority silent fallbacks + +These items are acknowledged but deprioritized: + +- **4.17**: `action.rs:96-101` -- count bonus `unwrap_or(0.0)` on branch + exploration arrays. Safe when `count_bonus_coefficient = 0.0` (disabled). + Fix: bounds-check the slice access. +- **4.18**: `attention.rs` fallbacks -- attention mechanism defaults. + Low risk since attention is not on the critical training path. +- **4.19**: PER beta defaults -- `unwrap_or` on PER importance sampling beta. + Fix: make PER beta required in config (it always should be present when PER + is enabled). +- **4.20**: Hyperopt validation gaps -- PSO can produce out-of-range values + that get clamped silently. Fix: add `validate()` method to + `DqnPsoParams` that returns `Result` instead of clamping. + +--- + +## 5. Implementation Plan + +### Phase 1: CRITICAL fixes (bugs #1-#8) + +All changes in this phase affect training correctness and must be deployed +together. + +| Order | Bug | Risk | Reason for ordering | +|-------|-----|------|---------------------| +| 1 | #4 gradient_clip_norm | Low | Config-only, no logic change | +| 2 | #5 entropy_coefficient | Low | Config-only, no logic change | +| 3 | #3 close price fallback | Low | Replace unwrap_or with indexing | +| 4 | #8 step_count inference | Medium | API change to forward() | +| 5 | #6 CUDA context 3x | Low | Pure refactor | +| 6 | #7 quantile fallbacks | Low | Replace unwrap_or with error | +| 7 | #1 backtest episode reset | Medium | CUDA kernel change | +| 8 | #2 stale max_equity | Medium | CUDA kernel change, order matters | + +### Phase 2: IMPORTANT fixes (bugs #10-#16) + +These can be deployed incrementally. + +| Order | Bug | Risk | +|-------|-----|------| +| 1 | #16 MaybeNoisyLinear | Low (compile-only) | +| 2 | #10 budget fraction constants | Low | +| 3 | #15 Tier 1 unwrap_or fixes | Low | +| 4 | #11 trailing stop param | Medium (kernel API change) | +| 5 | #14 Kelly min trades param | Medium (kernel API change) | +| 6 | #12-#13 doc comments | None | + +### Phase 3: Lower priority (#17-#20) + +Deferred to a follow-up PR. + +--- + +## 6. Success Criteria + +1. **Zero `unwrap_or` on data-critical paths**: `grep -rn 'unwrap_or'` in + `data_loading.rs`, `quantile_regression.rs`, `action.rs:161`, and + `state.rs:73` returns zero hits on the specific lines identified. + +2. **Backtest episode reset matches experience kernel**: After capital floor + breach, all 8 backtest portfolio fields are reset to fresh-episode values. + `max_equity` is set to the new initial capital, not the stale peak. + +3. **Gradient clip norm consistency**: A single `gradient_clip_norm` value + propagates to all sub-trainers. Searching for `unwrap_or.*gradient_clip` + returns zero hits. + +4. **entropy_coefficient consistency**: The field is non-optional (or resolved + once). No `unwrap_or` on entropy_coefficient in production code. + +5. **No inference side effects**: The non-training forward pass does not + increment `step_count` or advance the dropout scheduler. + +6. **All 129+ existing tests pass**: `SQLX_OFFLINE=true cargo test -p ml --lib` + and `SQLX_OFFLINE=true cargo test -p ml-dqn --lib` both pass. + +7. **Smoke test passes**: `FOXHUNT_TEST_DATA=test_data/futures-baseline cargo + test -p ml --lib -- smoke_tests --ignored --nocapture` completes without + regression. + +--- + +## 7. Risks and Mitigations + +| Risk | Mitigation | +|------|------------| +| Making `gradient_clip_norm` non-optional breaks deserialization of saved hyperparams | Use `#[serde(default = "default_grad_clip")]` with `fn default_grad_clip() -> f64 { 10.0 }` | +| Making `entropy_coefficient` non-optional breaks deserialization | Same `#[serde(default)]` strategy | +| Changing `forward()` signature breaks call sites | Use `forward_train()` / non-training forward split; deprecate `forward()` | +| Backtest reset changes affect hyperopt metric consistency | Run a before/after comparison on the smoke test dataset to quantify metric drift | +| Kernel parameter additions (#11, #14) require host-side changes | Add the new parameters to `GpuBacktestConfig` with defaults matching current hardcoded values -- zero behavioral change by default | diff --git a/docs/superpowers/specs/2026-03-27-code-hygiene-design.md b/docs/superpowers/specs/2026-03-27-code-hygiene-design.md new file mode 100644 index 000000000..73c51c1c2 --- /dev/null +++ b/docs/superpowers/specs/2026-03-27-code-hygiene-design.md @@ -0,0 +1,490 @@ +# Spec C: Code Hygiene -- Dead Code Deletion, Hardcoded Value Elimination, Config Unification + +**Date:** 2026-03-27 +**Spec:** C of 3 (DQN CUDA training pipeline overhaul) +**Execution order:** LAST -- after Spec A (BF16 + perf) and Spec B (bugs) +**Estimated LOC removed:** ~400 +**Estimated hardcoded values eliminated:** ~30 + +--- + +## 1. Motivation + +The DQN CUDA training pipeline has accumulated dead code paths from iterative +development (WAVE/Phase labeling, abandoned BF16 per-struct mirrors, single-variant +enums). Hardcoded magic numbers -- CUDA block sizes, monitoring limits, expert demo +parameters, kernel `#define` constants -- are scattered across 12+ files, making the +codebase fragile and difficult to tune via hyperopt. This spec removes all of it in a +single coordinated pass. + +--- + +## 2. Dead Code Deletion + +### 2.1 BF16 Per-Struct Mirrors in `gpu_weights.rs` + +**File:** `crates/ml/src/cuda_pipeline/gpu_weights.rs` + +The flat-buffer BF16 path (lines 2140-2251 in `gpu_dqn_trainer.rs`) superseded +the per-struct BF16 mirror approach. The following are unused: + +| Item | Lines | LOC | +|------|-------|-----| +| `DuelingWeightSetBf16` struct + `alloc_from` + `sync_from_f32` | 473-593 | ~120 | +| `BranchingWeightSetBf16` struct + `alloc_from` + `sync_from_f32` | 490-627 | ~55 | +| `CuriosityWeightSetBf16` struct + `alloc_from` + `sync_from_f32` | 503-653 | ~45 | +| `KernelWeightPackBf16` struct + `build()` + `DeviceRepr` impl | 670-719+ | ~80 | +| `alloc_bf16_mirror()` helper | 511-519 | ~10 | +| `convert_f32_to_bf16()` helper | 527-552 | ~25 | +| `raw_device_ptr_bf16()` helper | 657-661 | ~5 | + +**Total: ~340 LOC deleted.** + +Spec A will introduce a replacement flat-buffer BF16 conversion path. These per-struct +types are dead regardless -- they are not called from any production code path today. +Coordinate with Spec A to ensure the new BF16 conversion lands before these are removed. + +**Action:** Delete all items listed above from `gpu_weights.rs`. Remove any `pub use` +re-exports from `crates/ml/src/cuda_pipeline/mod.rs` referencing these types. + +### 2.2 Uncalled BF16 Methods in `gpu_dqn_trainer.rs` + +**File:** `crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs` + +| Method | Line | LOC | Reason | +|--------|------|-----|--------| +| `ensure_bf16_mirrors()` | 2144 | ~25 | Defined but never called; replaced by flat-buffer path | +| `sync_online_bf16()` | 2237 | ~10 | Delegates to `launch_bf16_convert_online()` directly; Spec A replaces | + +Both methods accept `_online_d`, `_online_b`, `_target_d`, `_target_b` parameters +(note the leading underscores -- the compiler already flagged them unused). The actual +BF16 conversion uses `launch_segmented_bf16_convert()` which operates on flat buffers. + +**Action:** Delete both methods. Verify no call sites exist (grep confirms zero callers). + +### 2.3 `MaybeNoisyLinear` Single-Variant Enum in `branching.rs` + +**File:** `crates/ml-dqn/src/branching.rs`, lines 252-330 + +`MaybeNoisyLinear` is an enum with exactly one variant: + +```rust +enum MaybeNoisyLinear { + Noisy(NoisyLinear), +} +``` + +Every match arm is `Self::Noisy(n) => ...` -- pure ceremony. The `Deterministic` +variant was removed when NoisyNet became mandatory. + +**Action:** Replace `MaybeNoisyLinear` with `NoisyLinear` directly. Update all field +types in `BranchingDuelingQNetwork` (lines 345-350: `value_fc`, `value_out`, +`branch_fcs`, `branch_outs`) and the `make_noisy_linear()` factory (line 458) to +return `NoisyLinear` directly. Remove the `Debug` impl (line 326-330) since +`NoisyLinear` derives or implements `Debug`. Remove all wrapper delegation methods +(lines 256-324) -- callers invoke `NoisyLinear` methods directly. Update +`copy_noisy_params()` (line 1116-1121) to accept `&mut NoisyLinear` / `&NoisyLinear`. + +### 2.4 Dead Code Path in `config.rs` + +**File:** `crates/ml/src/trainers/dqn/config.rs`, lines 105-107 + +```rust +if !true { + return Ok(None); +} +``` + +This is unreachable. The condition `!true` is always `false`. + +**Action:** Delete these three lines. + +### 2.5 Empty `reset_episode()` in `agent.rs` + +**File:** `crates/ml-dqn/src/agent.rs`, line 553 + +```rust +pub const fn reset_episode(&mut self) { + // Reset any per-episode tracking if needed + // Network weights and replay buffer are preserved +} +``` + +Empty no-op function. No callers. + +**Action:** Delete the method and its doc comment (lines 552-556). + +--- + +## 3. Hardcoded Value Elimination + +### 3.1 CUDA Grid Dimension Helper + +**Problem:** 50+ occurrences of `((n + 255) / 256)` scattered across 13 files in +`crates/ml/src/cuda_pipeline/`. Each instance re-derives the ceiling division and +hardcodes the block dimension 256. + +**Files affected (sample of 50+ sites):** + +- `gpu_dqn_trainer.rs` -- 20 occurrences (lines 751, 793, 840, 956, 995, 1036, 1081, 1131, 1202, 1312, 1527, 1633, 2190, 2550, 2956, 2957, 3422, 3590, 3684, 3721, 3983, 4402) +- `gpu_backtest_evaluator.rs` -- 5 occurrences (lines 684, 1074, 1298, 1621, 1661) +- `gpu_iqn_head.rs` -- 6 occurrences (lines 349, 511, 534, 583, 666, 719) +- `gpu_experience_collector.rs` -- 3 occurrences (lines 1065, 1178, 1743) +- `gpu_her.rs` -- 1 occurrence (line 409) +- `gpu_curiosity_trainer.rs` -- 3 occurrences (lines 130, 344, 387) +- `gpu_attention.rs` -- 2 occurrences (lines 346, 364) +- `gpu_iql_trainer.rs` -- 2 occurrences (lines 308, 328) +- `batched_backward.rs` -- 2 occurrences (lines 316, 757) +- `batched_forward.rs` -- 3 occurrences (lines 572, 605, 651) +- `decision_transformer.rs` -- 5 occurrences (lines 652, 753, 853, 878, 999) +- `signal_adapter.rs` -- 3 occurrences (lines 75, 133, 202) + +**Solution:** Add to `crates/ml/src/cuda_pipeline/mod.rs`: + +```rust +/// Standard CUDA block dimension for 1D kernels. +pub(crate) const CUDA_BLOCK_DIM: u32 = 256; + +/// Compute the 1D grid dimension for `n` elements at `CUDA_BLOCK_DIM` threads per block. +#[inline(always)] +pub(crate) const fn cuda_grid_1d(n: usize) -> u32 { + ((n + (CUDA_BLOCK_DIM as usize) - 1) / (CUDA_BLOCK_DIM as usize)) as u32 +} +``` + +**Action:** Replace all 50+ manual grid computations with `cuda_grid_1d(n)`. +Replace all hardcoded `256` in `block_dim` with `CUDA_BLOCK_DIM`. + +### 3.2 Monitoring Limits + +**File:** `crates/ml/src/trainers/dqn/monitoring.rs` + +| Constant | Current value | Line(s) | Usage | +|----------|--------------|---------|-------| +| Reward history limit | `1000` | 64 | `reward_history.len() > 1000` | +| Reward drain amount | `500` | 65 | `reward_history.drain(0..500)` | +| Q-value history limit | `1000` | 130 | `q_value_history.len() > 1000` | +| Q-value drain amount | `500` | 131 | `q_value_history.drain(0..500)` | +| Constant reward threshold | `0.01` | 162 | `std < 0.01` | +| Max constant epochs | `5` | 171 | `consecutive_constant_epochs >= 5` | +| Q-value divergence threshold | `1000.0` | 235 | `(max_q - min_q).abs() > 1000.0` | + +**Solution:** Add a `MonitoringConfig` struct: + +```rust +pub(crate) struct MonitoringConfig { + pub history_capacity: usize, // default: 1000 + pub history_drain_count: usize, // default: 500 + pub constant_reward_std_threshold: f32, // default: 0.01 + pub max_constant_epochs: usize, // default: 5 + pub q_divergence_threshold: f64, // default: 1000.0 +} +``` + +Pass `MonitoringConfig` into `TrainingMonitor::new()`. Default values are preserved +but become tunable. Alternatively, embed these in `DQNHyperparameters` if hyperopt +should search over them. + +### 3.3 Factored Action Space Array Sizes + +**File:** `crates/ml/src/trainers/dqn/monitoring.rs` + +| Constant | Current value | Line(s) | +|----------|--------------|---------| +| Factored action count array | `[usize; 81]` | 22, 45 | +| Exposure action array | `[usize; 9]` | 15-17, 40 | +| Q-value arrays | `[f64; 9]`, `[usize; 9]` | 16-17, 41 | + +The array size 81 = 9 * 3 * 3 (exposure * order * urgency). The size 9 is +`branch_0_size` (exposure levels). + +**Solution:** Make `TrainingMonitor` accept branch sizes at construction time and +use `Vec` instead of fixed arrays, or derive the constant from `BranchingConfig`: + +```rust +impl TrainingMonitor { + pub(crate) fn new(epoch: usize, branch_sizes: &[usize]) -> Self { + let total_actions: usize = branch_sizes.iter().product(); + let exposure_size = branch_sizes.first().copied().unwrap_or(9); + Self { + factored_action_counts: vec![0; total_actions], + action_counts: vec![0; exposure_size], + q_value_sums: vec![0.0; exposure_size], + q_value_counts: vec![0; exposure_size], + .. + } + } +} +``` + +### 3.4 Expert Demo Parameters + +**File:** `crates/ml/src/trainers/dqn/expert_demos.rs` + +| Parameter | Current value | Line | +|-----------|--------------|------| +| Fast EMA period | `20` | 46 | +| Slow EMA period | `50` | 47 | +| ADX threshold | `25.0` | 48 | +| ADX period | `14` | 100, 200 | + +These are already fields on `ExpertDemoGenerator` (lines 34-41) with defaults in +the `Default` impl (lines 43-51). The ADX period is the only one hardcoded at line +100 (`let adx_period: usize = 14;`) and line 200 (`let period: usize = 14;`). + +**Action:** Add `adx_period: usize` field to `ExpertDemoGenerator`. Update +`Default` impl to set `adx_period: 14`. Replace the two hardcoded `14` values +with `self.adx_period`. Update `ExpertDemoGenerator::new()` signature to accept +`adx_period`. + +### 3.5 CUDA Kernel `#define` Constants + +**Files:** +- `crates/ml/src/cuda_pipeline/c51_loss_kernel.cu` +- `crates/ml/src/cuda_pipeline/mse_loss_kernel.cu` + +| Define | Value | File(s) | Issue | +|--------|-------|---------|-------| +| `MAX_PER_SAMPLE_CE` | `50.0f` | c51_loss_kernel.cu:25 | Should derive from `v_range` | +| `LABEL_SMOOTHING_EPS` | `0.01f` | c51_loss_kernel.cu:28 | Should be a kernel parameter | +| `MAX_ATOMS` | `128` | c51_loss_kernel.cu:32 | Should be validated against `num_atoms` at compile time | +| `MAX_BRANCH_SIZE` | `9` | c51_loss_kernel.cu:35, mse_loss_kernel.cu:19 | Should derive from branch config | +| `NUM_BRANCHES` | `3` | c51_loss_kernel.cu:36, mse_loss_kernel.cu:20 | Should be a kernel parameter | + +**Solution:** Pass these as NVRTC compile-time defines rather than hardcoding them in +the `.cu` source. The Rust NVRTC compilation step already accepts `--define-macro` +options. + +```rust +// In the NVRTC compilation step: +let defines = format!( + "-DMAX_PER_SAMPLE_CE={:.1}f -DLABEL_SMOOTHING_EPS={:.4}f \ + -DMAX_ATOMS={} -DMAX_BRANCH_SIZE={} -DNUM_BRANCHES={}", + v_range * 2.0, // derived from config + label_smoothing_eps, // from DQNHyperparameters + num_atoms.next_power_of_two(), // validated upper bound + branch_sizes.iter().max().unwrap(), // max branch size + branch_sizes.len(), // number of branches +); +``` + +The `.cu` files already guard with `#ifndef`, so this is backwards-compatible. + +Add a runtime assertion at trainer construction time: + +```rust +assert!( + config.num_atoms <= MAX_ATOMS, + "num_atoms ({}) exceeds kernel MAX_ATOMS ({})", + config.num_atoms, MAX_ATOMS +); +``` + +### 3.6 Training Loop Dimensions + +**File:** `crates/ml/src/trainers/dqn/trainer/training_loop.rs` + +| Constant | Value | Line | Issue | +|----------|-------|------|-------| +| `STATIC_MAX_BATCH_SIZE` | `8192` | (not found as named const, used implicitly) | Should be config | +| Auto-batch safety margin | `0.15` | (not found as named const, used implicitly) | Should be config | + +Note: Grep did not find `STATIC_MAX_BATCH_SIZE` or `AutoBatchSizer` as named items +in the current codebase. These may have been refactored or renamed. Verify before +implementation. If they exist as inline literals, extract to named constants in +`GpuDqnTrainerConfig` or a new `BatchSizerConfig`. + +--- + +## 4. Stale Comment Cleanup + +### 4.1 WAVE / Phase / BUG Labels + +Remove all development-era labels from comments across these files: + +| File | Count | Examples | +|------|-------|---------| +| `monitoring.rs` | 4 | `WAVE 9-11 production monitoring`, `WAVE P2: Episode length tracking` | +| `constructor.rs` | 15+ | `WAVE 26 P2.2`, `WAVE 16`, `BUG #37 FIX`, `WAVE 23 P0 Fix #1`, `WAVE 1.1`, `WAVE 16S`, `WAVE 24`, `WAVE 26 P1`, `WAVE 3 FIX #2` | +| `agent.rs` | 2 | `BUG #7 FIX` (lines 615, 669) | + +**Action:** Replace with descriptive comments that explain *what* and *why* without +referencing internal development phases. Examples: + +| Before | After | +|--------|-------| +| `// WAVE 9-11 production monitoring` | `// Q-value range tracking for divergence detection` | +| `// BUG #37 FIX: Q-value clipping` | `// Clip Q-values to prevent step-level explosions` | +| `// WAVE 26 P2.2: Validate gradient_accumulation_steps > 0` | `// Validate gradient accumulation config` | +| `// WAVE 1.1: Initialize triple barrier engine` | `// Initialize triple barrier engine (max 1000 active trackers)` | +| `// WAVE 3 FIX #2: Start with None` | `// Start with None; collect feature stats during initial epochs` | + +### 4.2 MEMORY LEAK FIX Labels + +Lines 61, 129 in `monitoring.rs` have `// MEMORY LEAK FIX:` comments. Once the +limits are extracted to `MonitoringConfig` (section 3.2), these comments become +self-documenting and should be shortened to explain the bounded-history invariant +rather than referencing the original bug. + +--- + +## 5. Config Unification + +### 5.1 `entropy_coefficient`: Remove `Option` Wrapper + +**File:** `crates/ml/src/trainers/dqn/config.rs`, line 960 + +Currently `pub entropy_coefficient: Option`. Every usage site unwraps with +`.unwrap_or(0.01)` or `.unwrap_or(0.001)` -- and these defaults **disagree**: + +| Call site | Default | +|-----------|---------| +| `constructor.rs:314` | `0.01` | +| `constructor.rs:399` | `0.01` | +| `gpu_dqn_trainer.rs:213` | `0.001` | +| `DQNHyperparameters::default()` config.rs:1431 | `Some(0.001)` | + +**Action:** Make the field `pub entropy_coefficient: f64` with default `0.001`. +Remove all `.unwrap_or()` calls. Update serialization/deserialization to handle +the schema migration (serde default attribute). + +### 5.2 `gradient_clip_norm`: Unify Default + +**File:** `crates/ml/src/trainers/dqn/config.rs`, line 867 + +Currently `pub gradient_clip_norm: Option` with inconsistent defaults: + +| Call site | Default | +|-----------|---------| +| `fused_training.rs:194` (main DQN) | `10.0` | +| `fused_training.rs:263` (IQN) | `1.0` | +| `fused_training.rs:298` (ensemble) | `1.0` | +| `constructor.rs:257` | `10.0` | +| `DQNHyperparameters::default()` config.rs:1382 | `Some(10.0)` | + +**Action:** Make the field `pub gradient_clip_norm: f64` with default `10.0`. +For IQN and ensemble paths that need a different norm, add dedicated fields +(e.g., `iqn_gradient_clip_norm`, `ensemble_gradient_clip_norm`) or use the +same unified value. Remove all `.unwrap_or()` calls. + +### 5.3 `noisy_epsilon_floor`: Wire Through Consistently + +Currently `pub noisy_epsilon_floor: Option` with disagreeing defaults: + +| Call site | Default | +|-----------|---------| +| `constructor.rs:315` | `0.0` | +| `training_loop.rs:520` | `0.05` | +| `hyperopt/adapters/dqn.rs:476` | `0.10` | +| `DQNHyperparameters::default()` config.rs:1432 | `Some(0.02)` | +| Test assertions (tests.rs:582) | assert `0.10` | + +**Action:** Make the field `pub noisy_epsilon_floor: f64` with default `0.10` +(matching the test assertions and hyperopt default). Remove all `.unwrap_or()` +calls and ensure a single source of truth. + +### 5.4 `count_bonus_coefficient`: Wire Through Consistently + +Currently `pub count_bonus_coefficient: Option` with defaults: + +| Call site | Default | +|-----------|---------| +| `constructor.rs:316-317` | `0.0` | +| `gpu_experience_collector.rs:278` | `0.0` | +| `DQNHyperparameters::default()` config.rs:1433 | `Some(0.1)` | +| Test assertions (tests.rs:600) | assert `0.05` | + +**Action:** Make the field `pub count_bonus_coefficient: f64` with default `0.05` +(matching test expectations). Remove all `.unwrap_or()` calls. + +--- + +## 6. Files Changed Summary + +| File | Deletions | Additions | Net | +|------|-----------|-----------|-----| +| `crates/ml/src/cuda_pipeline/gpu_weights.rs` | ~340 LOC | 0 | -340 | +| `crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs` | ~35 LOC methods | ~5 LOC (use helper) | -30 | +| `crates/ml/src/cuda_pipeline/mod.rs` | 0 | ~10 LOC (const + fn) | +10 | +| `crates/ml-dqn/src/branching.rs` | ~80 LOC (enum + impls) | ~5 LOC (type aliases) | -75 | +| `crates/ml-dqn/src/agent.rs` | ~5 LOC | 0 | -5 | +| `crates/ml/src/trainers/dqn/config.rs` | ~3 LOC dead path + Option wrappers | ~10 LOC (defaults) | +7 | +| `crates/ml/src/trainers/dqn/monitoring.rs` | ~5 LOC (stale comments) | ~15 LOC (MonitoringConfig) | +10 | +| `crates/ml/src/trainers/dqn/expert_demos.rs` | 0 | ~5 LOC (adx_period field) | +5 | +| `crates/ml/src/trainers/dqn/trainer/constructor.rs` | ~15 LOC (stale comments) | ~5 LOC (simplified unwraps) | -10 | +| `crates/ml/src/trainers/dqn/fused_training.rs` | 0 | ~5 LOC (remove unwrap_or) | +5 | +| `crates/ml/src/cuda_pipeline/c51_loss_kernel.cu` | 0 | ~5 LOC (guard comments) | +5 | +| `crates/ml/src/cuda_pipeline/mse_loss_kernel.cu` | 0 | ~3 LOC (guard comments) | +3 | +| 13 files for `cuda_grid_1d` migration | ~50 LOC (manual math) | ~50 LOC (helper calls) | 0 | +| **TOTAL** | **~530** | **~120** | **-415** | + +--- + +## 7. Success Criteria + +1. **Zero dead code warnings.** Compile with `#[deny(dead_code)]` on + `gpu_weights.rs`, `branching.rs`, `agent.rs`, and `config.rs` -- zero warnings. + +2. **Zero magic numbers in kernel launches.** Every `grid_dim` computation in + `crates/ml/src/cuda_pipeline/` uses `cuda_grid_1d()` or `CUDA_BLOCK_DIM`. + Grep for `+ 255) / 256` returns zero results. + +3. **All config values traceable.** Every training parameter flows from + `DQNHyperparameters` through `GpuDqnTrainerConfig` to the GPU kernel. No + `.unwrap_or()` calls with differing defaults for the same field. + +4. **No stale labels.** Grep for `WAVE \d`, `Phase \d`, `BUG #\d` across + `crates/ml/` and `crates/ml-dqn/` returns zero results. + +5. **All 129+ tests pass.** `SQLX_OFFLINE=true cargo test -p ml --lib` and + `SQLX_OFFLINE=true cargo test -p ml-dqn --lib` both green. + +6. **Clean compile.** `SQLX_OFFLINE=true cargo check --workspace` produces zero + warnings related to unused imports, dead code, or unreachable patterns from + these changes. + +--- + +## 8. Non-Goals + +- **Performance optimization** -- covered by Spec A (BF16, cuBLAS SGEMM, CUDA Graphs). +- **Bug fixes** -- covered by Spec B. +- **New features** -- no new training capabilities are introduced. +- **Kernel algorithm changes** -- `#define` values become parameters but the + kernel algorithms are unchanged. + +--- + +## 9. Risks and Mitigations + +| Risk | Likelihood | Impact | Mitigation | +|------|-----------|--------|------------| +| Spec A lands new BF16 types that reference deleted items | Medium | Build break | Execute Spec C after Spec A merges; coordinate on `gpu_weights.rs` | +| Removing `Option` wrappers breaks JSON deserialization of saved trials | Medium | Hyperopt regression | Add `#[serde(default)]` on all unwrapped fields; migration test | +| `MaybeNoisyLinear` removal misses a usage site | Low | Build break | Grep for `MaybeNoisyLinear` across entire workspace before deletion | +| NVRTC `--define-macro` changes break kernel compilation on older CUDA | Low | Runtime crash | Gate behind feature flag; keep `#ifndef` guards as fallback | + +--- + +## 10. Execution Checklist + +This is the recommended execution order within Spec C. Each step is independently +testable. + +- [ ] **C.1** Add `cuda_grid_1d()` and `CUDA_BLOCK_DIM` to `cuda_pipeline/mod.rs` +- [ ] **C.2** Migrate all 50+ grid computations to use the helper (13 files) +- [ ] **C.3** Delete BF16 per-struct types from `gpu_weights.rs` (~340 LOC) +- [ ] **C.4** Delete `ensure_bf16_mirrors()` and `sync_online_bf16()` from `gpu_dqn_trainer.rs` +- [ ] **C.5** Replace `MaybeNoisyLinear` with `NoisyLinear` in `branching.rs` +- [ ] **C.6** Delete `if !true` dead path from `config.rs` +- [ ] **C.7** Delete `reset_episode()` from `agent.rs` +- [ ] **C.8** Extract monitoring limits to `MonitoringConfig` +- [ ] **C.9** Add `adx_period` field to `ExpertDemoGenerator` +- [ ] **C.10** Pass CUDA kernel defines via NVRTC compile-time options +- [ ] **C.11** Make `entropy_coefficient` non-optional with default `0.001` +- [ ] **C.12** Make `gradient_clip_norm` non-optional with default `10.0` +- [ ] **C.13** Make `noisy_epsilon_floor` non-optional with default `0.10` +- [ ] **C.14** Make `count_bonus_coefficient` non-optional with default `0.05` +- [ ] **C.15** Clean all WAVE/Phase/BUG labels from comments +- [ ] **C.16** Make `TrainingMonitor` accept dynamic branch sizes +- [ ] **C.17** Run full test suite and verify zero warnings