diff --git a/config/training/dqn-localdev.toml b/config/training/dqn-localdev.toml index b0bf300d3..8fbc2c14d 100644 --- a/config/training/dqn-localdev.toml +++ b/config/training/dqn-localdev.toml @@ -45,6 +45,7 @@ min_epochs_before_stopping = 80 [experience] gpu_timesteps_per_episode = 200 +gpu_n_episodes = 128 min_hold_bars = 3 [experience.fill_simulation] diff --git a/config/training/dqn-production.toml b/config/training/dqn-production.toml index f3a51c4a8..bea43c117 100644 --- a/config/training/dqn-production.toml +++ b/config/training/dqn-production.toml @@ -51,6 +51,7 @@ patience = 20 min_epochs_before_stopping = 80 [experience] +gpu_n_episodes = 4096 initial_capital = 35000.0 tx_cost_multiplier = 0.0001 min_hold_bars = 5 diff --git a/config/training/dqn-smoketest.toml b/config/training/dqn-smoketest.toml index 6b9fcf483..9d9b63c2c 100644 --- a/config/training/dqn-smoketest.toml +++ b/config/training/dqn-smoketest.toml @@ -49,6 +49,7 @@ min_epochs_before_stopping = 5 [experience] gpu_timesteps_per_episode = 100 +gpu_n_episodes = 32 min_hold_bars = 3 [experience.fill_simulation] diff --git a/crates/ml/src/cuda_pipeline/batched_forward.rs b/crates/ml/src/cuda_pipeline/batched_forward.rs index 507bc80db..58d3fcf00 100644 --- a/crates/ml/src/cuda_pipeline/batched_forward.rs +++ b/crates/ml/src/cuda_pipeline/batched_forward.rs @@ -57,7 +57,7 @@ use std::mem::ManuallyDrop; use cudarc::cublas::result as cublas_result; use cudarc::cublas::sys as cublas_sys; -use cudarc::driver::{CudaFunction, CudaSlice, CudaStream, DevicePtr, LaunchConfig, PushKernelArg}; +use cudarc::driver::{CudaEvent, CudaFunction, CudaSlice, CudaStream, DevicePtr, LaunchConfig, PushKernelArg}; use crate::MLError; @@ -132,6 +132,11 @@ pub struct CublasForward { /// Each branch (exposure, order, urgency) submits GEMMs to its own stream, /// then the main stream joins all three before consuming the logits. branch_streams: [Arc; 3], + + // ── Pre-allocated CUDA events for fork-join synchronization ── + /// Reused across forward passes — avoids cuEventCreate/Destroy per call. + trunk_done_event: CudaEvent, + branch_done_events: [CudaEvent; 3], } impl CublasForward { @@ -193,6 +198,16 @@ impl CublasForward { stream.fork().map_err(|e| MLError::DeviceError(format!("fork branch stream 2: {e}")))?, ]; + // ── Pre-allocate CUDA events for fork-join (reused across all forward passes) ── + let ctx = stream.context(); + let trunk_done_event = ctx.new_event(None) + .map_err(|e| MLError::DeviceError(format!("alloc trunk_done_event: {e}")))?; + let branch_done_events = [ + ctx.new_event(None).map_err(|e| MLError::DeviceError(format!("alloc branch_done_event 0: {e}")))?, + ctx.new_event(None).map_err(|e| MLError::DeviceError(format!("alloc branch_done_event 1: {e}")))?, + ctx.new_event(None).map_err(|e| MLError::DeviceError(format!("alloc branch_done_event 2: {e}")))?, + ]; + Ok(Self { handle: SendSyncCublasHandle(raw_handle), _workspace_buf: workspace_buf, @@ -214,6 +229,8 @@ impl CublasForward { branch_1_size, branch_2_size, branch_streams, + trunk_done_event, + branch_done_events, }) } @@ -296,8 +313,8 @@ impl CublasForward { // on a non-captured stream makes it join the capture graph, and the // fork/join pattern is recorded as graph dependencies. - // Record trunk completion on the main stream. - let trunk_done = stream.record_event(None) + // Record trunk completion on the main stream (pre-allocated event). + self.trunk_done_event.record(stream) .map_err(|e| MLError::ModelError(format!("trunk event record: {e}")))?; let mut logit_byte_offset: u64 = 0; @@ -307,7 +324,7 @@ impl CublasForward { let w_fc_idx = branch_w_base[d]; // Branch stream waits for trunk completion. - bs.wait(&trunk_done) + bs.wait(&self.trunk_done_event) .map_err(|e| MLError::ModelError(format!("branch {d} wait trunk: {e}")))?; // Redirect cuBLAS handle to this branch stream. @@ -336,11 +353,11 @@ impl CublasForward { .map_err(|e| MLError::ModelError(format!("cublasSetStream restore: {e:?}")))?; } - // Join: main stream waits for all 3 branches to complete. + // Join: main stream waits for all 3 branches to complete (pre-allocated events). for d in 0..3 { - let branch_done = self.branch_streams[d].record_event(None) + self.branch_done_events[d].record(&self.branch_streams[d]) .map_err(|e| MLError::ModelError(format!("branch {d} done event: {e}")))?; - stream.wait(&branch_done) + stream.wait(&self.branch_done_events[d]) .map_err(|e| MLError::ModelError(format!("main wait branch {d}: {e}")))?; } } else { @@ -405,7 +422,7 @@ impl CublasForward { if distinct_branches { // Multi-stream branch dispatch (f32 path, experience collection). - let trunk_done = stream.record_event(None) + self.trunk_done_event.record(stream) .map_err(|e| MLError::ModelError(format!("f32 trunk event record: {e}")))?; let mut logit_byte_offset: u64 = 0; @@ -414,7 +431,7 @@ impl CublasForward { let n_d = branch_sizes[d]; let w_fc_idx = branch_w_base[d]; - bs.wait(&trunk_done) + bs.wait(&self.trunk_done_event) .map_err(|e| MLError::ModelError(format!("f32 branch {d} wait trunk: {e}")))?; unsafe { @@ -440,9 +457,9 @@ impl CublasForward { } for d in 0..3 { - let branch_done = self.branch_streams[d].record_event(None) + self.branch_done_events[d].record(&self.branch_streams[d]) .map_err(|e| MLError::ModelError(format!("f32 branch {d} done event: {e}")))?; - stream.wait(&branch_done) + stream.wait(&self.branch_done_events[d]) .map_err(|e| MLError::ModelError(format!("f32 main wait branch {d}: {e}")))?; } } else { diff --git a/crates/ml/src/cuda_pipeline/gpu_experience_collector.rs b/crates/ml/src/cuda_pipeline/gpu_experience_collector.rs index 062010644..df02bdd30 100644 --- a/crates/ml/src/cuda_pipeline/gpu_experience_collector.rs +++ b/crates/ml/src/cuda_pipeline/gpu_experience_collector.rs @@ -1385,8 +1385,9 @@ impl GpuExperienceCollector { // Normalizing would scramble the C51 atom distribution. // Build next_states on GPU (episode-aware shift by n_steps) + // Must use total (includes counterfactual) so next_states covers all experiences. let next_states = build_next_states_f32( - &self.stream, &states, n_episodes, timesteps, sd, n_steps as usize, + &self.stream, &states, n_episodes * 2, timesteps, sd, n_steps as usize, )?; // ── Episode ID fill (GPU-native, zero CPU) ──────────────────── diff --git a/crates/ml/src/trainers/dqn/config.rs b/crates/ml/src/trainers/dqn/config.rs index d1f7828ea..6d6c7666f 100644 --- a/crates/ml/src/trainers/dqn/config.rs +++ b/crates/ml/src/trainers/dqn/config.rs @@ -1175,6 +1175,9 @@ pub struct DQNHyperparameters { /// Timesteps per episode in GPU kernel (default: 500, max: 1000) /// Higher values collect more experience per launch at the cost of VRAM pub gpu_timesteps_per_episode: usize, + /// Number of parallel episodes for GPU experience collection. + /// Controls VRAM usage: total = gpu_n_episodes × gpu_timesteps_per_episode × 2 (counterfactual) × state_dim × 4 bytes. + pub gpu_n_episodes: usize, /// Average bid-ask spread for GPU portfolio simulation (default: 0.0001 = 1bp) pub avg_spread: f64, @@ -1677,6 +1680,7 @@ impl DQNHyperparameters { // Phase 3: GPU experience collection gpu_timesteps_per_episode: 500, // Default: 500 timesteps per episode + gpu_n_episodes: 256, // Default: 256 parallel episodes avg_spread: 0.0001, // Default: 1bp (ES/NQ futures) // GPU-dynamic network sizing — None means auto-detect from hardware. diff --git a/crates/ml/src/trainers/dqn/trainer/constructor.rs b/crates/ml/src/trainers/dqn/trainer/constructor.rs index 7ddca38b8..f7c36ba73 100644 --- a/crates/ml/src/trainers/dqn/trainer/constructor.rs +++ b/crates/ml/src/trainers/dqn/trainer/constructor.rs @@ -729,7 +729,6 @@ impl DQNTrainer { training_guard: None, gpu_monitoring: None, reduction_kernels: None, - cached_n_episodes: None, // GPU pipeline: staging buffer pool (auto-initialized on CUDA devices) buffer_pool, diff --git a/crates/ml/src/trainers/dqn/trainer/mod.rs b/crates/ml/src/trainers/dqn/trainer/mod.rs index b43727071..daabae77d 100644 --- a/crates/ml/src/trainers/dqn/trainer/mod.rs +++ b/crates/ml/src/trainers/dqn/trainer/mod.rs @@ -265,9 +265,6 @@ pub struct DQNTrainer { /// Multi-GPU configuration for data-parallel training (None = single GPU) pub(crate) multi_gpu: Option, - /// Cached GPU n_episodes (computed once from nvidia-smi, reused across epochs) - /// Avoids forking nvidia-smi subprocess every epoch (~5-10ms per fork). - pub(crate) cached_n_episodes: Option, /// Pre-computed OFI features per bar (indexed by global bar position). /// Populated during data loading when MBP-10 order book data is available. diff --git a/crates/ml/src/trainers/dqn/trainer/training_loop.rs b/crates/ml/src/trainers/dqn/trainer/training_loop.rs index ea0add2db..c94674ead 100644 --- a/crates/ml/src/trainers/dqn/trainer/training_loop.rs +++ b/crates/ml/src/trainers/dqn/trainer/training_loop.rs @@ -764,7 +764,7 @@ impl DQNTrainer { cfg.value_hidden_dim, cfg.branch_hidden_dim, ); - let alloc_episodes = self.compute_alloc_episodes(state_dim); + let alloc_episodes = self.compute_alloc_episodes(); Some(GpuExperienceCollector::new( stream, online_br.vars(), @@ -789,7 +789,7 @@ impl DQNTrainer { cfg.value_hidden_dim, cfg.advantage_hidden_dim, ); - let alloc_episodes = self.compute_alloc_episodes(state_dim); + let alloc_episodes = self.compute_alloc_episodes(); Some(GpuExperienceCollector::new( stream, online.store(), @@ -814,7 +814,7 @@ impl DQNTrainer { cfg.value_hidden_dim, cfg.advantage_hidden_dim, ); - let alloc_episodes = self.compute_alloc_episodes(state_dim); + let alloc_episodes = self.compute_alloc_episodes(); Some(GpuExperienceCollector::new( stream, online.vars(), @@ -897,20 +897,9 @@ impl DQNTrainer { Ok(()) } - /// Compute allocation episode count for GPU buffers. - /// Always auto-scales from VRAM — no manual override. - fn compute_alloc_episodes(&self, state_dim: usize) -> usize { - use ml_core::memory_optimization::detect_gpu_hardware; - match detect_gpu_hardware() { - Ok(hw) => { - let optimal = hw.optimal_n_episodes( - state_dim, - self.hyperparams.gpu_timesteps_per_episode, - ); - optimal.max(32).min(16384) - } - Err(_) => 256, // safe fallback - } + /// Episode count for GPU buffer allocation — uses configured gpu_n_episodes. + fn compute_alloc_episodes(&self) -> usize { + self.hyperparams.gpu_n_episodes.max(32) } // ═══════════════════════════════════════════════════════════════════════ @@ -961,32 +950,7 @@ impl DQNTrainer { use crate::cuda_pipeline::gpu_experience_collector::ExperienceCollectorConfig; - let raw_sd = if !self.hyperparams.mbp10_data_dir.is_empty() { 53 } else { 45 }; - let aligned_sd = (raw_sd + 7) & !7; - - // Cache n_episodes on first epoch — always auto-scaled from VRAM. - let n_episodes = if let Some(cached) = self.cached_n_episodes { - cached - } else { - use ml_core::memory_optimization::detect_gpu_hardware; - let computed = match detect_gpu_hardware() { - Ok(hw) => { - let optimal = hw.optimal_n_episodes( - aligned_sd, - self.hyperparams.gpu_timesteps_per_episode, - ); - let chosen = optimal.max(32).min(16384) as i32; - info!( - "GPU auto-scaled n_episodes: {} (SMs={}, VRAM={:.0}MB)", - chosen, hw.sm_count, hw.free_memory_mb - ); - chosen - } - Err(_) => 256_i32, // safe fallback - }; - self.cached_n_episodes = Some(computed); - computed - }; + let n_episodes = self.hyperparams.gpu_n_episodes.max(32) as i32; // Domain randomization: all randomization is GPU-native (zero CPU RNG) // Domain randomization, mirror, causal, vaccine: ALL always active (one production path) @@ -1179,7 +1143,7 @@ impl DQNTrainer { "GPU zero-roundtrip collection FAILED (no CPU fallback): {e}" ))?; - let count = gpu_batch.n_episodes * gpu_batch.timesteps; + let count = gpu_batch.n_episodes * gpu_batch.timesteps * 2; // counterfactual doubles experiences // Record CUDA event after experience collection completes on the forked stream. // This replaces stream.synchronize() — the event is checked lazily at the start @@ -1238,7 +1202,7 @@ impl DQNTrainer { if count > 0 { // Direct CudaSlice insertion — zero Candle, zero CPU roundtrip. - let total = gpu_batch.n_episodes * gpu_batch.timesteps; + let total = gpu_batch.n_episodes * gpu_batch.timesteps * 2; // counterfactual doubles experiences // actions: CudaSlice → CudaSlice (safe reinterpret for values [0..44]) let actions_u32: &CudaSlice = unsafe { diff --git a/crates/ml/src/training_profile.rs b/crates/ml/src/training_profile.rs index 9600db280..4cfdc9174 100644 --- a/crates/ml/src/training_profile.rs +++ b/crates/ml/src/training_profile.rs @@ -257,6 +257,7 @@ pub struct RewardSection { #[derive(Debug, Clone, Deserialize, Default)] pub struct ExperienceSection { pub gpu_timesteps_per_episode: Option, + pub gpu_n_episodes: Option, pub initial_capital: Option, /// Maps to `DQNHyperparameters::transaction_cost_multiplier`. pub tx_cost_multiplier: Option, @@ -933,6 +934,9 @@ impl DqnTrainingProfile { if let Some(v) = ex.gpu_timesteps_per_episode { hp.gpu_timesteps_per_episode = v; } + if let Some(v) = ex.gpu_n_episodes { + hp.gpu_n_episodes = v; + } if let Some(v) = ex.initial_capital { hp.initial_capital = v as f32; }