From 80bc8a2ee1f4ad38e83c68c1aea8b16f5007c7ab Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Mon, 30 Mar 2026 09:17:25 +0200 Subject: [PATCH] =?UTF-8?q?feat(generalization):=20#21=20stochastic=20dept?= =?UTF-8?q?h=20=E2=80=94=20CUDA=20Graph=20compatible=20layer=20dropout?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Per-layer drop with 20% probability during training. Each hidden layer (h_s1, h_s2, h_v) is independently scaled by 0 (dropped) or 1/(1-p) (kept, expected-value correction). Only applied to online forward pass on current states — target and Double-DQN passes use full network. Implementation: - New stochastic_depth_scale kernel in dqn_utility_kernels.cu - Per-layer scale buffer [3] f32 — written by host before each graph replay (CUDA Graphs capture addresses, not contents) - Scale kernel inserted in launch_cublas_forward after online Pass 1 - update_stochastic_depth_mask() generates random 0/keep scales per step - At inference (experience collection), all layers active (no drop) Forces every layer to produce useful features independently — prevents deep compositional memorization where removal of any single layer would collapse the output. First RL trading application of stochastic depth (proven in Vision Transformers, novel in DQN). Co-Authored-By: Claude Opus 4.6 (1M context) --- config/training/dqn-localdev.toml | 1 + config/training/dqn-production.toml | 1 + .../src/cuda_pipeline/dqn_utility_kernels.cu | 28 ++++++ .../ml/src/cuda_pipeline/gpu_dqn_trainer.rs | 87 ++++++++++++++++++- crates/ml/src/trainers/dqn/config.rs | 6 ++ crates/ml/src/trainers/dqn/fused_training.rs | 1 + crates/ml/src/training_profile.rs | 2 + 7 files changed, 122 insertions(+), 4 deletions(-) diff --git a/config/training/dqn-localdev.toml b/config/training/dqn-localdev.toml index c49d3d720..c3076ffd9 100644 --- a/config/training/dqn-localdev.toml +++ b/config/training/dqn-localdev.toml @@ -82,6 +82,7 @@ anti_lr_sharpe_threshold = 0.3 feature_mask_fraction = 0.3 feature_noise_scale = 0.1 enable_vol_normalization = true +stochastic_depth_prob = 0.2 asymmetric_dd_weight = 0.5 time_reversal_mod = 5 regret_blend = 0.3 diff --git a/config/training/dqn-production.toml b/config/training/dqn-production.toml index 6df76c954..8561f09f6 100644 --- a/config/training/dqn-production.toml +++ b/config/training/dqn-production.toml @@ -87,6 +87,7 @@ anti_lr_sharpe_threshold = 0.3 feature_mask_fraction = 0.3 feature_noise_scale = 0.1 enable_vol_normalization = true +stochastic_depth_prob = 0.2 asymmetric_dd_weight = 0.5 time_reversal_mod = 5 regret_blend = 0.3 diff --git a/crates/ml/src/cuda_pipeline/dqn_utility_kernels.cu b/crates/ml/src/cuda_pipeline/dqn_utility_kernels.cu index 871967c5a..fde717e1d 100644 --- a/crates/ml/src/cuda_pipeline/dqn_utility_kernels.cu +++ b/crates/ml/src/cuda_pipeline/dqn_utility_kernels.cu @@ -454,3 +454,31 @@ extern "C" __global__ void pad_states_kernel( int j = idx % padded_sd; dst[idx] = (j < sd) ? src[b * sd + j] : bf16_zero(); } + +/* ══════════════════════════════════════════════════════════════════════ + * STOCHASTIC DEPTH KERNEL (#21) + * + * Scales a bf16 hidden activation buffer by a per-layer scalar. + * Used for stochastic depth: each hidden layer is either kept (scaled + * by 1/(1-p) for expected-value correction) or dropped (scaled by 0). + * + * The scale buffer [num_layers] is written by the host BEFORE each + * CUDA Graph replay. The graph reads the updated values — this is legal + * because CUDA Graphs capture buffer addresses, not contents. + * + * buf[i] *= scale_buf[layer_idx] for i = 0..n-1 + * + * Launch config: grid=(ceil(n/256), 1, 1), block=(256, 1, 1). + * ══════════════════════════════════════════════════════════════════════ */ +extern "C" __global__ void stochastic_depth_scale( + __nv_bfloat16* __restrict__ buf, + const float* __restrict__ scale_buf, + int layer_idx, + int n +) { + int i = blockIdx.x * blockDim.x + threadIdx.x; + if (i >= n) return; + float scale = scale_buf[layer_idx]; + float val = (float)buf[i] * scale; + buf[i] = __float2bfloat16(val); +} diff --git a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs index 668d0c64b..dcc9c08e3 100644 --- a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs +++ b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs @@ -187,6 +187,8 @@ pub struct GpuDqnTrainConfig { /// #27 Ensemble disagreement penalty weight. Q_target -= weight * ensemble_std. /// 0.0 = disabled. pub ensemble_disagreement_penalty: f32, + /// #21 Stochastic depth drop probability. 0.0 = disabled, 0.2 = 20% drop per layer. + pub stochastic_depth_prob: f32, } impl Default for GpuDqnTrainConfig { @@ -225,6 +227,7 @@ impl Default for GpuDqnTrainConfig { curiosity_q_penalty_lambda: 0.0, asymmetric_dd_weight: 0.0, ensemble_disagreement_penalty: 0.0, + stochastic_depth_prob: 0.0, } } } @@ -548,6 +551,14 @@ pub struct GpuDqnTrainer { drawdown_depths_buf: CudaSlice, /// #18 Asymmetric DD loss weight (from config). asymmetric_dd_weight: f32, + /// #21 Stochastic depth: per-layer scale buffer [3] (h_s1, h_s2, h_v). + /// Written by host before each graph replay. 0.0=drop, 1/(1-p)=keep. + stochastic_depth_scale_buf: CudaSlice, + /// #21 Stochastic depth kernel function. + stochastic_depth_kernel: CudaFunction, + /// #21 Stochastic depth drop probability (0.0=disabled, 0.2=20% drop). + stochastic_depth_prob: f32, + /// #27 Per-sample ensemble std for disagreement Q-penalty. /// Populated by `upload_ensemble_std()` before loss kernel launch. /// Zero-filled when ensemble is disabled. @@ -1930,7 +1941,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, grad_norm_finalize_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, pad_states_kernel, saxpy_f32_kernel) = + let (grad_norm_kernel, grad_norm_finalize_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, pad_states_kernel, saxpy_f32_kernel, stochastic_depth_kernel) = compile_training_kernels(&stream, &config)?; // Separate grad_norm instance for non-graph launches (clip_grad_buf_inplace). @@ -2276,6 +2287,11 @@ impl GpuDqnTrainer { .map_err(|e| MLError::ModelError(format!("alloc drawdown_depths_buf: {e}")))?; let ensemble_std_buf = stream.alloc_zeros::(b) .map_err(|e| MLError::ModelError(format!("alloc ensemble_std_buf: {e}")))?; + // #21 Stochastic depth: 3 layer scales (h_s1, h_s2, h_v). Init to 1.0 (keep all). + let mut stochastic_depth_scale_buf = stream.alloc_zeros::(3) + .map_err(|e| MLError::ModelError(format!("alloc stochastic_depth_scale: {e}")))?; + stream.memcpy_htod(&[1.0_f32, 1.0, 1.0], &mut stochastic_depth_scale_buf) + .map_err(|e| MLError::ModelError(format!("init stochastic_depth_scale: {e}")))?; let curiosity_inference_func = { static CURIOSITY_INFERENCE_CUBIN: &[u8] = include_bytes!( concat!(env!("OUT_DIR"), "/curiosity_inference_kernel.cubin") @@ -2343,6 +2359,7 @@ impl GpuDqnTrainer { let dd_weight = config.asymmetric_dd_weight; let ens_weight = config.ensemble_disagreement_penalty; + let sd_prob = config.stochastic_depth_prob; Ok(Self { config, stream, @@ -2473,6 +2490,9 @@ impl GpuDqnTrainer { curiosity_error_buf, drawdown_depths_buf, asymmetric_dd_weight: dd_weight, + stochastic_depth_scale_buf, + stochastic_depth_kernel, + stochastic_depth_prob: sd_prob, ensemble_std_buf, ensemble_disagreement_weight: ens_weight, curiosity_inference_func, @@ -2597,6 +2617,7 @@ impl GpuDqnTrainer { if self.graph_forward.is_none() { self.capture_training_graphs(online_dueling, online_branching)?; } + self.update_stochastic_depth_mask()?; self.replay_forward()?; Ok(FusedTrainScalars { total_loss: 0.0, grad_norm: 0.0 }) } @@ -2691,6 +2712,7 @@ impl GpuDqnTrainer { if self.graph_forward.is_none() { self.capture_training_graphs(online_dueling, online_branching)?; } + self.update_stochastic_depth_mask()?; self.replay_forward()?; // --- INJECTION POINT: caller should inject auxiliary gradients here --- self.replay_adam()?; @@ -2733,6 +2755,27 @@ impl GpuDqnTrainer { /// Replay graph_forward only (zero → forward → loss → grad → backward). /// After this call, grad_buf contains C51 gradients. The caller can inject /// auxiliary gradients (IQN, attention, ensemble) before calling replay_adam(). + /// #21 Stochastic depth: write random per-layer drop/keep scales before graph replay. + /// Must be called BEFORE replay_forward() each training step. + /// CUDA Graphs capture buffer addresses, not contents — this is safe. + pub fn update_stochastic_depth_mask(&mut self) -> Result<(), MLError> { + if self.stochastic_depth_prob <= 0.0 { + return Ok(()); // Disabled — mask stays [1.0, 1.0, 1.0] + } + use rand::Rng; + let mut rng = rand::thread_rng(); + let p = self.stochastic_depth_prob; + let keep_scale = 1.0 / (1.0 - p); // Expected-value correction + let mask: [f32; 3] = [ + if rng.gen::() > p { keep_scale } else { 0.0 }, + if rng.gen::() > p { keep_scale } else { 0.0 }, + if rng.gen::() > p { keep_scale } else { 0.0 }, + ]; + self.stream.memcpy_htod(&mask, &mut self.stochastic_depth_scale_buf) + .map_err(|e| MLError::ModelError(format!("stochastic_depth mask upload: {e}")))?; + Ok(()) + } + pub fn replay_forward(&self) -> Result<(), MLError> { if let Some(ref graph) = self.graph_forward { graph.0.launch().map_err(|e| { @@ -3619,6 +3662,40 @@ impl GpuDqnTrainer { self.ptrs.on_v_logits_buf, self.ptrs.on_b_logits_buf, )?; + // ── #21 Stochastic depth: scale hidden activations by per-layer mask ── + // Only applied to Pass 1 (online on current states). Target and Double-DQN + // passes use full network (no stochastic depth at inference). + if self.stochastic_depth_prob > 0.0 { + let b = self.config.batch_size; + let scale_ptr = self.stochastic_depth_scale_buf.raw_ptr(); + let layers: [(u64, usize); 3] = [ + (self.ptrs.save_h_s1, self.config.shared_h1), // layer 0: first shared hidden + (self.ptrs.save_h_s2, self.config.shared_h2), // layer 1: second shared hidden + (self.ptrs.save_h_v, self.config.value_h), // layer 2: value head hidden + ]; + for (layer_idx, &(buf_ptr, dim)) in layers.iter().enumerate() { + let n = (b * dim) as i32; + let blocks = ((n as u32 + 255) / 256) as u32; + let li = layer_idx as i32; + unsafe { + self.stream + .launch_builder(&self.stochastic_depth_kernel) + .arg(&buf_ptr) + .arg(&scale_ptr) + .arg(&li) + .arg(&n) + .launch(LaunchConfig { + grid_dim: (blocks, 1, 1), + block_dim: (256, 1, 1), + shared_mem_bytes: 0, + }) + .map_err(|e| MLError::ModelError(format!( + "stochastic_depth layer {layer_idx}: {e}" + )))?; + } + } + } + // ── Pass 2: Target forward on NEXT_STATES (graph-safe) cublas.forward_target_raw( &self.stream, self.ptrs.next_states_buf, &tg_w_ptrs, @@ -4600,7 +4677,7 @@ impl GpuDqnTrainer { fn compile_training_kernels( stream: &Arc, config: &GpuDqnTrainConfig, -) -> Result<(CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction), MLError> { +) -> Result<(CudaFunction, CudaFunction, CudaFunction, CudaFunction, 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), @@ -4642,9 +4719,11 @@ fn compile_training_kernels( .map_err(|e| MLError::ModelError(format!("dqn_clip_grad_kernel load: {e}")))?; let pad_states = module.load_function("pad_states_kernel") .map_err(|e| MLError::ModelError(format!("pad_states_kernel load: {e}")))?; + let stochastic_depth = module.load_function("stochastic_depth_scale") + .map_err(|e| MLError::ModelError(format!("stochastic_depth_scale load: {e}")))?; - info!("GpuDqnTrainer: 14 utility kernels loaded from precompiled cubin"); - Ok((grad_norm, grad_norm_finalize, adam_update, f32_to_bf16, bf16_to_f32, saxpy, zero, regime_scale, shrink_perturb, _relu_mask_from_module, spectral_norm, clipped_saxpy, clip_grad, pad_states, saxpy_f32)) + info!("GpuDqnTrainer: 15 utility kernels loaded from precompiled cubin"); + Ok((grad_norm, grad_norm_finalize, adam_update, f32_to_bf16, bf16_to_f32, saxpy, zero, regime_scale, shrink_perturb, _relu_mask_from_module, spectral_norm, clipped_saxpy, clip_grad, pad_states, saxpy_f32, stochastic_depth)) } /// Load the standalone Polyak EMA kernel from precompiled cubin. diff --git a/crates/ml/src/trainers/dqn/config.rs b/crates/ml/src/trainers/dqn/config.rs index db3990170..afeb8457a 100644 --- a/crates/ml/src/trainers/dqn/config.rs +++ b/crates/ml/src/trainers/dqn/config.rs @@ -1004,6 +1004,11 @@ pub struct DQNHyperparameters { /// 0.0 = disabled. pub trade_clustering_penalty: f64, + /// #21 Stochastic depth: drop probability per hidden layer during training. + /// Each layer is independently dropped with this probability. Surviving layers + /// are scaled by 1/(1-p) for expected-value correction. 0.0 = disabled. + pub stochastic_depth_prob: f64, + /// #18 Asymmetric drawdown loss: extra penalty on Q-overestimation when /// portfolio is in drawdown. Weight for L_dd = max(0, Q_pred-Q_target) * dd^2. /// 0.0 = disabled. @@ -1519,6 +1524,7 @@ impl DQNHyperparameters { enable_mirror_universe: true, // #10: alternate mirrored/normal epochs position_entropy_weight: 0.01, // #19: reward += 0.01 * H(position_histogram) trade_clustering_penalty: 0.05, // #25: penalize temporally clustered trades + stochastic_depth_prob: 0.2, // #21: 20% drop probability per hidden layer asymmetric_dd_weight: 0.5, // #18: extra loss on Q-overestimation in drawdown // Wave 16 Portfolio Features (default: ALL ENABLED) diff --git a/crates/ml/src/trainers/dqn/fused_training.rs b/crates/ml/src/trainers/dqn/fused_training.rs index 108ad8ec8..6d52b7336 100644 --- a/crates/ml/src/trainers/dqn/fused_training.rs +++ b/crates/ml/src/trainers/dqn/fused_training.rs @@ -208,6 +208,7 @@ impl FusedTrainingCtx { curiosity_q_penalty_lambda: hyperparams.curiosity_q_penalty_lambda as f32, asymmetric_dd_weight: hyperparams.asymmetric_dd_weight as f32, ensemble_disagreement_penalty: hyperparams.ensemble_disagreement_penalty as f32, + stochastic_depth_prob: hyperparams.stochastic_depth_prob as f32, }; // Extract weight sets from VarMaps (online + target) diff --git a/crates/ml/src/training_profile.rs b/crates/ml/src/training_profile.rs index 8f6b72160..8bc659a20 100644 --- a/crates/ml/src/training_profile.rs +++ b/crates/ml/src/training_profile.rs @@ -161,6 +161,7 @@ pub struct GeneralizationSection { pub feature_mask_fraction: Option, pub feature_noise_scale: Option, pub enable_vol_normalization: Option, + pub stochastic_depth_prob: Option, pub asymmetric_dd_weight: Option, pub time_reversal_mod: Option, pub regret_blend: Option, @@ -810,6 +811,7 @@ impl DqnTrainingProfile { if let Some(v) = g.feature_mask_fraction { hp.feature_mask_fraction = v; } if let Some(v) = g.feature_noise_scale { hp.feature_noise_scale = v; } if let Some(v) = g.enable_vol_normalization { hp.enable_vol_normalization = v; } + if let Some(v) = g.stochastic_depth_prob { hp.stochastic_depth_prob = v; } if let Some(v) = g.asymmetric_dd_weight { hp.asymmetric_dd_weight = v; } if let Some(v) = g.time_reversal_mod { hp.time_reversal_mod = v; } if let Some(v) = g.regret_blend { hp.regret_blend = v; }