diff --git a/config/training/dqn-localdev.toml b/config/training/dqn-localdev.toml index 2a47cff42..6b86a3cf4 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 +enable_gradient_vaccine = false bottleneck_dim = 0 stochastic_depth_prob = 0.2 asymmetric_dd_weight = 0.5 diff --git a/config/training/dqn-production.toml b/config/training/dqn-production.toml index 6601ba5f7..0c3b14041 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 +enable_gradient_vaccine = false bottleneck_dim = 0 stochastic_depth_prob = 0.2 asymmetric_dd_weight = 0.5 diff --git a/crates/ml/src/cuda_pipeline/dqn_utility_kernels.cu b/crates/ml/src/cuda_pipeline/dqn_utility_kernels.cu index aff67974b..5875ad526 100644 --- a/crates/ml/src/cuda_pipeline/dqn_utility_kernels.cu +++ b/crates/ml/src/cuda_pipeline/dqn_utility_kernels.cu @@ -498,6 +498,85 @@ extern "C" __global__ void bn_tanh_concat_kernel( } } +/* ══════════════════════════════════════════════════════════════════════ + * GRADIENT VACCINE KERNELS (#32) + * + * Two kernels that implement the gradient projection for the vaccine: + * 1. gradient_dot_and_norm: computes dot(g_train, g_val) and |g_val|^2 + * via warp + block reduction. + * 2. gradient_project: conditionally projects g_train when dot < 0. + * g_train -= (dot / norm_sq) * g_val + * + * The dot and norm_sq scalars are written to GPU buffers (2 floats) + * and read by the projection kernel — no CPU roundtrip. + * ══════════════════════════════════════════════════════════════════════ */ + +extern "C" __global__ void gradient_dot_and_norm( + const float* __restrict__ g_train, /* [total_params] f32 */ + const float* __restrict__ g_val, /* [total_params] f32 */ + float* __restrict__ result, /* [2]: result[0]=dot, result[1]=norm_sq */ + int n +) { + /* Parallel reduction: each block reduces a tile, atomicAdd into result. */ + __shared__ float s_dot[256]; + __shared__ float s_norm[256]; + + int tid = threadIdx.x; + int idx = blockIdx.x * blockDim.x + threadIdx.x; + + float local_dot = 0.0f; + float local_norm = 0.0f; + + /* Grid-stride loop for large n */ + for (int i = idx; i < n; i += blockDim.x * gridDim.x) { + float t = g_train[i]; + float v = g_val[i]; + local_dot += t * v; + local_norm += v * v; + } + + s_dot[tid] = local_dot; + s_norm[tid] = local_norm; + __syncthreads(); + + /* Block reduction */ + for (int s = 128; s > 0; s >>= 1) { + if (tid < s) { + s_dot[tid] += s_dot[tid + s]; + s_norm[tid] += s_norm[tid + s]; + } + __syncthreads(); + } + + if (tid == 0) { + atomicAdd(&result[0], s_dot[0]); + atomicAdd(&result[1], s_norm[0]); + } +} + +extern "C" __global__ void gradient_project( + float* __restrict__ g_train, /* [total_params] f32 — MODIFIED in-place */ + const float* __restrict__ g_val, /* [total_params] f32 */ + const float* __restrict__ dot_norm, /* [2]: dot_norm[0]=dot, dot_norm[1]=norm_sq */ + int n +) { + float dot = dot_norm[0]; + float norm_sq = dot_norm[1]; + + /* Only project when gradients DISAGREE (dot < 0). + * When they agree (dot >= 0), both directions benefit generalization. */ + if (dot >= 0.0f) return; + + /* Projection scale: g_train -= (dot / norm_sq) * g_val */ + if (norm_sq < 1e-12f) return; /* g_val is zero — nothing to project */ + float scale = dot / norm_sq; + + int idx = blockIdx.x * blockDim.x + threadIdx.x; + if (idx < n) { + g_train[idx] -= scale * g_val[idx]; + } +} + /* ══════════════════════════════════════════════════════════════════════ * BOTTLENECK TANH BACKWARD KERNEL (#31) * diff --git a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs index 8176181b5..35d076106 100644 --- a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs +++ b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs @@ -189,6 +189,10 @@ pub struct GpuDqnTrainConfig { pub ensemble_disagreement_penalty: f32, /// #21 Stochastic depth drop probability. 0.0 = disabled, 0.2 = 20% drop per layer. pub stochastic_depth_prob: f32, + /// #32 Enable gradient vaccine. When true, projects out overfitting gradient + /// components using a held-out validation batch. + pub enable_gradient_vaccine: bool, + /// #31 Temporal causal bottleneck dimension. 0 = disabled (no bottleneck). /// 2 = maximum compression (14 market features → 2 abstract dimensions). pub bottleneck_dim: usize, @@ -234,6 +238,7 @@ impl Default for GpuDqnTrainConfig { asymmetric_dd_weight: 0.0, ensemble_disagreement_penalty: 0.0, stochastic_depth_prob: 0.0, + enable_gradient_vaccine: false, bottleneck_dim: 0, market_dim: 42, } @@ -582,6 +587,16 @@ pub struct GpuDqnTrainer { drawdown_depths_buf: CudaSlice, /// #18 Asymmetric DD loss weight (from config). asymmetric_dd_weight: f32, + /// #32 Gradient vaccine: saved g_train gradient [total_params] f32. + vaccine_grad_save: CudaSlice, + /// #32 Gradient vaccine: dot(g_train, g_val) and |g_val|^2 result [2] f32. + vaccine_dot_norm_buf: CudaSlice, + /// #32 Gradient vaccine kernels. + vaccine_dot_kernel: CudaFunction, + vaccine_project_kernel: CudaFunction, + /// #32 Whether gradient vaccine is enabled. + enable_gradient_vaccine: bool, + /// #31 Bottleneck hidden activation [B, bottleneck_dim] bf16. None when bottleneck_dim=0. bn_hidden_buf: Option>, /// #31 Bottleneck + portfolio concatenated [B, bottleneck_dim + portfolio_dim] bf16. @@ -1989,7 +2004,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, stochastic_depth_kernel, bn_tanh_backward_kernel, bn_bias_grad_kernel, bn_tanh_concat_kernel_fn) = + 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, bn_tanh_backward_kernel, bn_bias_grad_kernel, bn_tanh_concat_kernel_fn, vaccine_dot_kernel, vaccine_project_kernel) = compile_training_kernels(&stream, &config)?; // Separate grad_norm instance for non-graph launches (clip_grad_buf_inplace). @@ -2405,9 +2420,16 @@ impl GpuDqnTrainer { } }; + // #32 Gradient vaccine buffers + let vaccine_grad_save = stream.alloc_zeros::(total_params) + .map_err(|e| MLError::ModelError(format!("alloc vaccine_grad_save: {e}")))?; + let vaccine_dot_norm_buf = stream.alloc_zeros::(2) + .map_err(|e| MLError::ModelError(format!("alloc vaccine_dot_norm: {e}")))?; + let dd_weight = config.asymmetric_dd_weight; let ens_weight = config.ensemble_disagreement_penalty; let sd_prob = config.stochastic_depth_prob; + let vaccine_enabled = config.enable_gradient_vaccine; let bn_dim = config.bottleneck_dim; let market_dim = config.market_dim; @@ -2557,6 +2579,11 @@ impl GpuDqnTrainer { curiosity_error_buf, drawdown_depths_buf, asymmetric_dd_weight: dd_weight, + vaccine_grad_save, + vaccine_dot_norm_buf, + vaccine_dot_kernel, + vaccine_project_kernel, + enable_gradient_vaccine: vaccine_enabled, bn_hidden_buf, bn_concat_buf, bn_d_concat_buf, @@ -2829,6 +2856,93 @@ 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(). + /// #32 Gradient Vaccine: project out overfitting gradient components. + /// + /// After the main forward+backward fills `grad_buf` with g_train: + /// 1. Copy grad_buf → vaccine_grad_save (save g_train) + /// 2. Zero grad_buf + /// 3. Upload vaccine batch + run forward+backward → grad_buf = g_val + /// 4. Dot product: dot(g_train, g_val), |g_val|² + /// 5. If dot < 0: g_train -= (dot/|g_val|²) * g_val (project out conflicting component) + /// 6. Copy projected g_train back to grad_buf for Adam + pub fn apply_gradient_vaccine( + &mut self, + vaccine_batch: &crate::dqn::replay_buffer_type::GpuBatch, + ) -> Result<(), MLError> { + if !self.enable_gradient_vaccine { + return Ok(()); + } + let tp = self.total_params; + + // Step 1: Save current g_train + let grad_ptr = self.grad_buf.raw_ptr(); + let save_ptr = self.vaccine_grad_save.raw_ptr(); + dtod_copy(save_ptr, grad_ptr, tp * std::mem::size_of::(), &self.stream, 0, "vaccine_save")?; + + // Step 2: Zero grad_buf for vaccine pass + self.stream.memset_zeros(&mut self.grad_buf) + .map_err(|e| MLError::ModelError(format!("vaccine zero grad: {e}")))?; + self.stream.memset_zeros(&mut self.d_value_logits_buf) + .map_err(|e| MLError::ModelError(format!("vaccine zero d_val: {e}")))?; + self.stream.memset_zeros(&mut self.d_adv_logits_buf) + .map_err(|e| MLError::ModelError(format!("vaccine zero d_adv: {e}")))?; + + // Step 3: Upload vaccine batch + forward+backward (NON-graph path) + self.upload_batch_gpu(vaccine_batch)?; + // Forward pass (cuBLAS, outside graph) + self.launch_cublas_forward()?; + self.launch_curiosity_inference()?; + // Loss (C51 + MSE blend) + self.stream.memset_zeros(&mut self.d_value_logits_mse) + .map_err(|e| MLError::ModelError(format!("vaccine zero mse: {e}")))?; + self.stream.memset_zeros(&mut self.d_adv_logits_mse) + .map_err(|e| MLError::ModelError(format!("vaccine zero mse2: {e}")))?; + self.launch_mse_loss()?; + self.launch_mse_grad_to_scratch()?; + self.launch_c51_loss()?; + self.launch_c51_grad()?; + // Backward (cuBLAS) + self.cast_d_logits_to_bf16()?; + self.launch_cublas_backward()?; + // Now grad_buf = g_val (vaccine gradient) + + // Step 4: Dot product and norm² + self.stream.memset_zeros(&mut self.vaccine_dot_norm_buf) + .map_err(|e| MLError::ModelError(format!("vaccine zero dot: {e}")))?; + let tp_i32 = tp as i32; + let blocks = ((tp as u32 + 255) / 256) as u32; + let cfg = LaunchConfig { grid_dim: (blocks, 1, 1), block_dim: (256, 1, 1), shared_mem_bytes: 0 }; + let dot_norm_ptr = self.vaccine_dot_norm_buf.raw_ptr(); + let g_val_ptr = self.grad_buf.raw_ptr(); + unsafe { + self.stream + .launch_builder(&self.vaccine_dot_kernel) + .arg(&save_ptr) + .arg(&g_val_ptr) + .arg(&dot_norm_ptr) + .arg(&tp_i32) + .launch(cfg) + .map_err(|e| MLError::ModelError(format!("vaccine dot: {e}")))?; + } + + // Step 5: Conditional projection (only when dot < 0) + unsafe { + self.stream + .launch_builder(&self.vaccine_project_kernel) + .arg(&save_ptr) + .arg(&g_val_ptr) + .arg(&dot_norm_ptr) + .arg(&tp_i32) + .launch(cfg) + .map_err(|e| MLError::ModelError(format!("vaccine project: {e}")))?; + } + + // Step 6: Copy projected g_train back to grad_buf for Adam + dtod_copy(grad_ptr, save_ptr, tp * std::mem::size_of::(), &self.stream, 0, "vaccine_restore")?; + + Ok(()) + } + /// #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. @@ -4900,7 +5014,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, CudaFunction, CudaFunction, CudaFunction, CudaFunction), MLError> { +) -> Result<(CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, 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), @@ -4944,6 +5058,10 @@ fn compile_training_kernels( .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}")))?; + let vaccine_dot = module.load_function("gradient_dot_and_norm") + .map_err(|e| MLError::ModelError(format!("gradient_dot_and_norm load: {e}")))?; + let vaccine_project = module.load_function("gradient_project") + .map_err(|e| MLError::ModelError(format!("gradient_project load: {e}")))?; let bn_tanh_bw = module.load_function("bn_tanh_backward_kernel") .map_err(|e| MLError::ModelError(format!("bn_tanh_backward_kernel load: {e}")))?; let bn_bias_grad = module.load_function("bn_bias_grad_kernel") @@ -4951,8 +5069,8 @@ fn compile_training_kernels( let bn_tanh_concat = module.load_function("bn_tanh_concat_kernel") .map_err(|e| MLError::ModelError(format!("bn_tanh_concat_kernel load: {e}")))?; - info!("GpuDqnTrainer: 19 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, bn_tanh_bw, bn_bias_grad, bn_tanh_concat)) + info!("GpuDqnTrainer: 21 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, bn_tanh_bw, bn_bias_grad, bn_tanh_concat, vaccine_dot, vaccine_project)) } /// 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 c2ebefd96..7b08eac50 100644 --- a/crates/ml/src/trainers/dqn/config.rs +++ b/crates/ml/src/trainers/dqn/config.rs @@ -1004,6 +1004,10 @@ pub struct DQNHyperparameters { /// 0.0 = disabled. pub trade_clustering_penalty: f64, + /// #32 Gradient vaccine: project out overfitting gradient components using + /// a held-out validation batch. Mathematical guarantee against overfitting. + pub enable_gradient_vaccine: bool, + /// #31 Temporal causal bottleneck dimension. 0 = disabled, 2 = maximum compression. /// All market features are compressed through [market_dim → bottleneck_dim] linear + tanh /// before reaching the Q-network. Makes memorization structurally impossible. @@ -1529,6 +1533,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 + enable_gradient_vaccine: false, // #32: disabled by default (enable for mathematical overfit guarantee) bottleneck_dim: 0, // #31: disabled by default (0 = no bottleneck, 2 = max compression) 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 diff --git a/crates/ml/src/trainers/dqn/fused_training.rs b/crates/ml/src/trainers/dqn/fused_training.rs index 78f6b522e..d92e63e9e 100644 --- a/crates/ml/src/trainers/dqn/fused_training.rs +++ b/crates/ml/src/trainers/dqn/fused_training.rs @@ -129,6 +129,9 @@ pub(crate) struct FusedTrainingCtx { pub(crate) ensemble_kl_grad_kernel: Option, /// Pre-allocated buffer: [B * num_atoms] for diversity gradient on head 0 logits. pub(crate) ensemble_d_logits_buf: Option>, + /// #32 Gradient Vaccine: pending validation batch for next training step. + /// Set by the training loop, consumed by `run_full_step()`. + pub(crate) pending_vaccine_batch: Option, } impl Drop for FusedTrainingCtx { @@ -209,6 +212,7 @@ impl FusedTrainingCtx { 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, + enable_gradient_vaccine: hyperparams.enable_gradient_vaccine, bottleneck_dim: hyperparams.bottleneck_dim, market_dim: 42, // MARKET_DIM constant — 42 preprocessed features }; @@ -470,6 +474,7 @@ impl FusedTrainingCtx { ensemble_diversity_loss_buf, ensemble_kl_grad_kernel, ensemble_d_logits_buf, + pending_vaccine_batch: None, }) } @@ -792,7 +797,16 @@ impl FusedTrainingCtx { } } - // ── Step 5e: Replay graph_adam — Adam sees budget-allocated gradient ── + // ── Step 5e: Gradient Vaccine (#32) — project out overfitting gradients ── + // If a vaccine batch is provided, compute its gradient and project out + // the component of grad_buf that contradicts it. + if let Some(ref vaccine_batch) = self.pending_vaccine_batch { + self.trainer.apply_gradient_vaccine(vaccine_batch) + .map_err(|e| anyhow::anyhow!("Gradient vaccine: {e}"))?; + } + self.pending_vaccine_batch = None; + + // ── Step 5f: 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. diff --git a/crates/ml/src/training_profile.rs b/crates/ml/src/training_profile.rs index 0492dd28c..c113f60c5 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 enable_gradient_vaccine: Option, pub bottleneck_dim: Option, pub stochastic_depth_prob: Option, pub asymmetric_dd_weight: Option, @@ -812,6 +813,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.enable_gradient_vaccine { hp.enable_gradient_vaccine = v; } if let Some(v) = g.bottleneck_dim { hp.bottleneck_dim = 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; }