From e9062edecfee34147bfc459262b3fd33c8bc4d8b Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Mon, 30 Mar 2026 12:56:09 +0200 Subject: [PATCH] fix: counterfactual output buffer size + vaccine/causal guards Fixed collect_trade_stats DtoH panic: done_out was doubled by counterfactual (2x alloc) but readback used base_output count. Use slice(..base_output) to read only the original portion. Added guards: vaccine and causal intervention skip before CUDA graphs are captured (buffers not fully initialized on first step). Vaccine failures are non-fatal (logged, not propagated). Updated test assertions for 2x total_experiences. Co-Authored-By: Claude Opus 4.6 (1M context) --- crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs | 9 ++++++--- .../ml/src/cuda_pipeline/gpu_experience_collector.rs | 11 +++++++---- crates/ml/src/cuda_pipeline/mod.rs | 4 ++-- crates/ml/src/trainers/dqn/fused_training.rs | 10 +++++++--- 4 files changed, 22 insertions(+), 12 deletions(-) diff --git a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs index 47291acd1..146d3b2e0 100644 --- a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs +++ b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs @@ -3053,8 +3053,8 @@ impl GpuDqnTrainer { /// /// Returns mean sensitivity (read back once at the end, not per-feature). pub fn run_causal_intervention(&mut self, step: usize) -> Result { - // Causal intervention: always active, runs every N steps - if step % self.causal_config.interval != 0 { + // Causal intervention: runs every N steps (skip before graph capture) + if self.graph_forward.is_none() || step % self.causal_config.interval != 0 { return Ok(0.0); } @@ -3170,7 +3170,10 @@ impl GpuDqnTrainer { &mut self, vaccine_batch: &crate::dqn::replay_buffer_type::GpuBatch, ) -> Result<(), MLError> { - // #32 Gradient vaccine: always active (one production path) + // Only run vaccine after CUDA graphs are captured (buffers initialized) + if self.graph_forward.is_none() { + return Ok(()); + } let tp = self.total_params; // Step 1: Save current g_train diff --git a/crates/ml/src/cuda_pipeline/gpu_experience_collector.rs b/crates/ml/src/cuda_pipeline/gpu_experience_collector.rs index aff5dc80c..be5f71b6a 100644 --- a/crates/ml/src/cuda_pipeline/gpu_experience_collector.rs +++ b/crates/ml/src/cuda_pipeline/gpu_experience_collector.rs @@ -1302,8 +1302,9 @@ impl GpuExperienceCollector { // (equity_t - equity_{t-1}) / equity_{t-1} // NOT the shaped RL reward (which includes dense*0.1 + sparse*2.0 + loss_aversion). // Total = alloc_episodes * alloc_timesteps floats (~12KB for 3200 experiences). - let total_output = self.alloc_episodes * self.alloc_timesteps; - let mut host_raw_returns = vec![0.0_f32; total_output]; + // raw_returns_out is NOT doubled by counterfactual (only states/actions/rewards/dones are) + let base_output = self.alloc_episodes * self.alloc_timesteps; + let mut host_raw_returns = vec![0.0_f32; base_output]; super::dtoh_bf16_to_f32(&self.stream, &self.raw_returns_out, &mut host_raw_returns)?; let step_returns: Vec = host_raw_returns.iter().map(|&r| r as f64).collect(); @@ -1312,8 +1313,10 @@ impl GpuExperienceCollector { // done_out[i] = 1.0 when the capital floor circuit breaker fired or // the episode ended. Used to reset the equity curve in MaxDD computation. // done_out is f32 -- direct DtoH, no bf16 conversion needed. - let mut host_dones = vec![0.0_f32; total_output]; - self.stream.memcpy_dtoh(&self.done_out, &mut host_dones) + // Read only base portion of done_out (not counterfactual) + let done_slice = self.done_out.slice(..base_output); + let mut host_dones = vec![0.0_f32; base_output]; + self.stream.memcpy_dtoh(&done_slice, &mut host_dones) .map_err(|e| MLError::ModelError(format!("dtoh done_out: {e}")))?; let done_flags: Vec = host_dones.iter().map(|&d| d as f64).collect(); diff --git a/crates/ml/src/cuda_pipeline/mod.rs b/crates/ml/src/cuda_pipeline/mod.rs index da913c412..2506132e4 100644 --- a/crates/ml/src/cuda_pipeline/mod.rs +++ b/crates/ml/src/cuda_pipeline/mod.rs @@ -1071,7 +1071,7 @@ mod tests { let cfg = ExperienceCollectorConfig::default(); assert_eq!(cfg.n_episodes, 2048); assert_eq!(cfg.timesteps_per_episode, 500); - assert_eq!(cfg.total_experiences(), 2048 * 500); + assert_eq!(cfg.total_experiences(), 2048 * 500 * 2); // 2x for counterfactual augmentation assert!((cfg.epsilon - 0.1).abs() < f32::EPSILON); assert!((cfg.gamma - 0.99).abs() < f32::EPSILON); assert!((cfg.barrier_profit_mult - 1.02).abs() < f32::EPSILON); @@ -1107,7 +1107,7 @@ mod tests { let cfg = PpoCollectorConfig::default(); assert_eq!(cfg.n_episodes, 128); assert_eq!(cfg.timesteps_per_episode, 500); - assert_eq!(cfg.total_experiences(), 64_000); + assert_eq!(cfg.total_experiences(), 128_000); // 2x counterfactual assert!((cfg.gamma - 0.99).abs() < f32::EPSILON); assert!((cfg.gae_lambda - 0.95).abs() < f32::EPSILON); } diff --git a/crates/ml/src/trainers/dqn/fused_training.rs b/crates/ml/src/trainers/dqn/fused_training.rs index 11835447f..98594c121 100644 --- a/crates/ml/src/trainers/dqn/fused_training.rs +++ b/crates/ml/src/trainers/dqn/fused_training.rs @@ -828,9 +828,13 @@ impl FusedTrainingCtx { // ── 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}"))?; + // Vaccine runs after graphs stabilize (non-fatal on failure) + if self.steps_since_varmap_sync > 10 { + if let Some(ref vaccine_batch) = self.pending_vaccine_batch { + if let Err(e) = self.trainer.apply_gradient_vaccine(vaccine_batch) { + tracing::debug!("Gradient vaccine skipped: {e}"); + } + } } self.pending_vaccine_batch = None;