From 9b2804f9ecd9eba7e2b71620bda8a378a801cfd2 Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Sat, 28 Feb 2026 14:14:48 +0100 Subject: [PATCH] feat(cuda): wire GPU experience collection into DQN & PPO training loops MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Phase 3 of the CUDA pipeline: both trainers now take a GPU-first path for experience collection (128×500 = 64K experiences per kernel launch, zero CPU-GPU roundtrips per timestep) with automatic CPU fallback. - DQN: upload features alongside targets, GPU collection branch before CPU loop, gpu_batch_to_experiences() conversion into replay buffer - PPO: set_raw_market_data() for CudaSlice upload, GPU collection branch bypasses collect_rollouts + prepare_training_batch entirely, gpu_batch_to_trajectory_batch() conversion with in-kernel GAE - Configurable GPU batch sizes (gpu_n_episodes, gpu_timesteps_per_episode) and trading params (initial_capital, avg_spread) via hyperparameters - SAFETY training diagnostics downgraded from warn! to debug! - 4 new batch index-math validation tests in cuda_pipeline Co-Authored-By: Claude Opus 4.6 --- crates/ml/src/cuda_pipeline/mod.rs | 222 +++++++++++++++++ crates/ml/src/hyperopt/adapters/dqn.rs | 5 + crates/ml/src/trainers/dqn/config.rs | 15 ++ crates/ml/src/trainers/dqn/trainer.rs | 186 ++++++++++++-- crates/ml/src/trainers/ppo.rs | 229 +++++++++++++++++- .../tests/ppo_recurrent_performance_tests.rs | 4 + 6 files changed, 631 insertions(+), 30 deletions(-) diff --git a/crates/ml/src/cuda_pipeline/mod.rs b/crates/ml/src/cuda_pipeline/mod.rs index ac1131c93..ddc4c1d07 100644 --- a/crates/ml/src/cuda_pipeline/mod.rs +++ b/crates/ml/src/cuda_pipeline/mod.rs @@ -515,4 +515,226 @@ mod tests { assert!((cfg.gamma - 0.99).abs() < f32::EPSILON); assert!((cfg.gae_lambda - 0.95).abs() < f32::EPSILON); } + + // ── Phase 3: Batch struct layout & index-math validation ────────── + + #[cfg(feature = "cuda")] + #[test] + fn test_experience_batch_index_math() { + use super::gpu_experience_collector::ExperienceBatch; + + let n_episodes = 4; + let timesteps = 10; + let state_dim = 54; + let total = n_episodes * timesteps; + + // Build synthetic batch matching kernel output layout + let states: Vec = (0..total * state_dim) + .map(|i| i as f32 * 0.001) + .collect(); + let actions: Vec = (0..total).map(|i| (i % 45) as i32).collect(); + let rewards: Vec = (0..total).map(|i| i as f32 * 0.01).collect(); + let done_flags: Vec = (0..total) + .map(|i| if i % timesteps == timesteps - 1 { 1 } else { 0 }) + .collect(); + let target_q_values = vec![0.5_f32; total]; + let td_errors = vec![0.1_f32; total]; + + let batch = ExperienceBatch { + states, + actions, + rewards, + done_flags, + target_q_values, + td_errors, + n_episodes, + timesteps, + }; + + assert_eq!(batch.states.len(), total * state_dim); + assert_eq!(batch.actions.len(), total); + assert_eq!(batch.rewards.len(), total); + assert_eq!(batch.done_flags.len(), total); + + // Verify episode indexing: batch[ep][t] = ep * timesteps + t + for ep in 0..n_episodes { + for t in 0..timesteps { + let idx = ep * timesteps + t; + assert!(idx < total, "idx {} >= total {}", idx, total); + + let state_start = idx * state_dim; + let state_end = state_start + state_dim; + assert!(state_end <= batch.states.len()); + + // Actions in 0..44 + let a = batch.actions[idx]; + assert!((0..45).contains(&a)); + } + + // Last step of each episode should be done + let last_idx = ep * timesteps + timesteps - 1; + assert_eq!(batch.done_flags[last_idx], 1); + } + } + + #[cfg(feature = "cuda")] + #[test] + fn test_ppo_experience_batch_index_math() { + use super::gpu_ppo_collector::PpoExperienceBatch; + + let n_episodes = 4; + let timesteps = 10; + let state_dim = 54; + let total = n_episodes * timesteps; + + let states: Vec = (0..total * state_dim) + .map(|i| i as f32 * 0.001) + .collect(); + let actions: Vec = (0..total).map(|i| (i % 45) as i32).collect(); + let log_probs: Vec = vec![-1.5_f32; total]; + let advantages: Vec = (0..total).map(|i| i as f32 * 0.1 - 2.0).collect(); + let returns: Vec = (0..total).map(|i| i as f32 * 0.5).collect(); + let done_flags: Vec = (0..total) + .map(|i| if i % timesteps == timesteps - 1 { 1 } else { 0 }) + .collect(); + + let batch = PpoExperienceBatch { + states, + actions, + log_probs, + advantages, + returns, + done_flags, + n_episodes, + timesteps, + }; + + assert_eq!(batch.states.len(), total * state_dim); + assert_eq!(batch.actions.len(), total); + assert_eq!(batch.log_probs.len(), total); + assert_eq!(batch.advantages.len(), total); + assert_eq!(batch.returns.len(), total); + assert_eq!(batch.done_flags.len(), total); + + // Verify V(s) = R(s) - A(s) reconstruction + for i in 0..total { + let value = batch.returns[i] - batch.advantages[i]; + assert!(value.is_finite(), "V(s) should be finite at idx {}", i); + } + + // Verify state chunking: each chunk should have exactly state_dim elements + let chunks: Vec<&[f32]> = batch.states.chunks(state_dim).collect(); + assert_eq!(chunks.len(), total); + for (i, chunk) in chunks.iter().enumerate() { + assert_eq!( + chunk.len(), + state_dim, + "Chunk {} has wrong length: {}", + i, + chunk.len() + ); + } + } + + #[cfg(feature = "cuda")] + #[test] + fn test_experience_batch_next_state_boundary() { + use super::gpu_experience_collector::ExperienceBatch; + + let n_episodes = 2; + let timesteps = 5; + let state_dim = 54; + let total = n_episodes * timesteps; + + // Give each step a distinct state for easy verification + let states: Vec = (0..total) + .flat_map(|idx| std::iter::repeat(idx as f32).take(state_dim)) + .collect(); + let done_flags: Vec = (0..total) + .map(|i| if i % timesteps == timesteps - 1 { 1 } else { 0 }) + .collect(); + + let batch = ExperienceBatch { + states, + actions: vec![0; total], + rewards: vec![0.0; total], + done_flags, + target_q_values: vec![0.0; total], + td_errors: vec![0.0; total], + n_episodes, + timesteps, + }; + + // For the DQN conversion: next_state[t] = state[t+1] unless done or last step + for ep in 0..n_episodes { + for t in 0..timesteps { + let idx = ep * timesteps + t; + let done = batch.done_flags[idx] != 0; + + if !done && t + 1 < timesteps { + let next_idx = ep * timesteps + t + 1; + let ns_start = next_idx * state_dim; + let ns_end = ns_start + state_dim; + assert!(ns_end <= batch.states.len()); + + // Next state should be distinct from current state + let s_val = batch.states[idx * state_dim]; + let ns_val = batch.states[ns_start]; + assert!( + (ns_val - s_val - 1.0).abs() < f32::EPSILON, + "next_state should be current + 1 step" + ); + } + } + } + } + + #[cfg(feature = "cuda")] + #[test] + fn test_ppo_batch_gae_consistency() { + use super::gpu_ppo_collector::PpoExperienceBatch; + + let n_episodes = 2; + let timesteps = 10; + let total = n_episodes * timesteps; + let state_dim = 54; + + // Simulate GAE output: advantages should sum roughly to 0 within each episode + // (not exact, but the kernel normalizes per-episode) + let mut advantages = Vec::with_capacity(total); + for ep in 0..n_episodes { + let _base = ep as f32 * 0.5; + for t in 0..timesteps { + // Triangular pattern: positive early, negative late + let a = (timesteps as f32 / 2.0 - t as f32) * 0.1; + advantages.push(a); + } + } + + let returns: Vec = advantages.iter().map(|&a| a + 5.0).collect(); + + let batch = PpoExperienceBatch { + states: vec![0.0_f32; total * state_dim], + actions: vec![1; total], + log_probs: vec![-1.0; total], + advantages, + returns, + done_flags: (0..total) + .map(|i| if i % timesteps == timesteps - 1 { 1 } else { 0 }) + .collect(), + n_episodes, + timesteps, + }; + + // All values (R - A) should be approximately 5.0 + for i in 0..total { + let v = batch.returns[i] - batch.advantages[i]; + assert!( + (v - 5.0).abs() < f32::EPSILON, + "Value at {} should be ~5.0, got {}", + i, + v + ); + } + } } diff --git a/crates/ml/src/hyperopt/adapters/dqn.rs b/crates/ml/src/hyperopt/adapters/dqn.rs index c522ff6a1..a4a27ef9b 100644 --- a/crates/ml/src/hyperopt/adapters/dqn.rs +++ b/crates/ml/src/hyperopt/adapters/dqn.rs @@ -2375,6 +2375,11 @@ impl HyperparameterOptimizable for DQNTrainer { use_qr_dqn: params.use_qr_dqn, num_quantiles: params.num_quantiles, qr_kappa: params.qr_kappa, + + // Phase 3: GPU experience collection (defaults, not in search space) + gpu_n_episodes: 128, + gpu_timesteps_per_episode: 500, + avg_spread: 0.0001, }; let data_path_str = self diff --git a/crates/ml/src/trainers/dqn/config.rs b/crates/ml/src/trainers/dqn/config.rs index e085e1960..aba7adb6f 100644 --- a/crates/ml/src/trainers/dqn/config.rs +++ b/crates/ml/src/trainers/dqn/config.rs @@ -590,6 +590,16 @@ pub struct DQNHyperparameters { pub num_quantiles: usize, /// Kappa for quantile Huber loss (0.5-2.0) pub qr_kappa: f64, + + // Phase 3: GPU experience collection kernel configuration + /// Number of parallel episodes per GPU kernel launch (default: 128, max: 256) + /// Higher values improve GPU utilization on larger GPUs (e.g. H100: 256) + pub gpu_n_episodes: usize, + /// 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, + /// Average bid-ask spread for GPU portfolio simulation (default: 0.0001 = 1bp) + pub avg_spread: f64, } impl Default for DQNHyperparameters { @@ -764,6 +774,11 @@ impl DQNHyperparameters { use_qr_dqn: true, // Default: enabled (IQN distributional RL) num_quantiles: 32, // Default: 32 quantiles qr_kappa: 1.0, // Default: 1.0 (standard quantile Huber loss) + + // Phase 3: GPU experience collection + gpu_n_episodes: 128, // Default: 128 (good for 4-8GB VRAM GPUs) + gpu_timesteps_per_episode: 500, // Default: 500 timesteps per episode + avg_spread: 0.0001, // Default: 1bp (ES/NQ futures) } } } diff --git a/crates/ml/src/trainers/dqn/trainer.rs b/crates/ml/src/trainers/dqn/trainer.rs index a7ad28964..d3c36a071 100644 --- a/crates/ml/src/trainers/dqn/trainer.rs +++ b/crates/ml/src/trainers/dqn/trainer.rs @@ -205,6 +205,10 @@ pub struct DQNTrainer { #[cfg(feature = "cuda")] targets_raw_cuda: Option>, + /// Raw cudarc features buffer for CUDA experience kernel [num_bars * 51] + #[cfg(feature = "cuda")] + features_raw_cuda: Option>, + /// GPU experience collector for zero-roundtrip CUDA kernel (Phase 2b) #[cfg(feature = "cuda")] gpu_experience_collector: Option, @@ -762,6 +766,8 @@ impl DQNTrainer { #[cfg(feature = "cuda")] targets_raw_cuda: None, #[cfg(feature = "cuda")] + features_raw_cuda: None, + #[cfg(feature = "cuda")] gpu_experience_collector: None, }) } @@ -1325,12 +1331,13 @@ impl DQNTrainer { } } - // Phase 1b: Upload raw targets via cudarc + init GPU portfolio sim (once) + // Phase 1b: Upload raw targets + features via cudarc + init GPU portfolio sim (once) #[cfg(feature = "cuda")] if self.targets_raw_cuda.is_none() { if let candle_core::Device::Cuda(ref cuda_dev) = self.device { let stream = cuda_dev.cuda_stream(); let target_dim = 4; + let feature_dim = 51; let num_bars = training_data.len(); // Build flat targets (same as DqnGpuData::upload but for cudarc) @@ -1352,6 +1359,26 @@ impl DQNTrainer { } } + // Build flat features [num_bars * 51] for GPU experience kernel + if self.features_raw_cuda.is_none() { + let mut flat_features = Vec::with_capacity(num_bars * feature_dim); + for (features, _) in &training_data { + for &v in features.iter() { + flat_features.push(v as f32); + } + } + match stream.memcpy_stod(&flat_features) { + Ok(buf) => { + info!("CUDA features_raw uploaded: {} bars × {} ({:.1} KB)", + num_bars, feature_dim, (num_bars * feature_dim * 4) as f64 / 1024.0); + self.features_raw_cuda = Some(buf); + } + Err(e) => { + debug!("CUDA features_raw upload failed: {}", e); + } + } + } + // Initialize GPU portfolio simulator if self.gpu_portfolio_sim.is_none() { if let Some(ref _targets_buf) = self.targets_raw_cuda { @@ -1359,7 +1386,7 @@ impl DQNTrainer { match GpuPortfolioSimulator::new( stream, self.hyperparams.initial_capital as f32, - 0.0001, // avg_spread + self.hyperparams.avg_spread as f32, self.hyperparams.cash_reserve_percent as f32, self.max_position as f32, EPISODE_LENGTH, @@ -1400,7 +1427,7 @@ impl DQNTrainer { target.vars(), curiosity.forward_model_vars(), self.hyperparams.initial_capital as f32, - 0.0001_f32, // avg_spread + self.hyperparams.avg_spread as f32, self.hyperparams.cash_reserve_percent as f32, )) } @@ -1456,6 +1483,79 @@ impl DQNTrainer { capped.max(1) // Always collect at least 1 epoch }; + // Phase 3: GPU experience collection via zero-roundtrip CUDA kernel + // Collects N episodes × L timesteps entirely on GPU, then bulk-inserts into replay buffer. + // Falls back to CPU path if GPU collection fails or is unavailable. + #[cfg(feature = "cuda")] + let gpu_experiences_collected = if let ( + Some(ref mut collector), + Some(ref features_buf), + Some(ref targets_buf), + ) = ( + &mut self.gpu_experience_collector, + &self.features_raw_cuda, + &self.targets_raw_cuda, + ) { + use crate::cuda_pipeline::gpu_experience_collector::ExperienceCollectorConfig; + + let n_episodes = self.hyperparams.gpu_n_episodes.min(256) as i32; + let timesteps = self.hyperparams.gpu_timesteps_per_episode.min(1000) as i32; + let total_bars = training_data.len() as i32; + let usable_bars = (total_bars - timesteps).max(1); + let stride = (usable_bars / n_episodes).max(1); + let episode_starts: Vec = (0..n_episodes) + .map(|i| (i * stride).rem_euclid(usable_bars)) + .collect(); + + // Reset per-episode state before each epoch + if let Err(e) = collector.reset_episodes( + self.hyperparams.initial_capital as f32, + self.hyperparams.avg_spread as f32, + self.hyperparams.cash_reserve_percent as f32, + ) { + warn!("GPU episode reset failed: {e}"); + false + } else { + let agent = self.agent.read().await; + let epsilon = agent.get_epsilon(); + drop(agent); + + let config = ExperienceCollectorConfig { + n_episodes, + timesteps_per_episode: timesteps, + total_bars, + epsilon, + gamma: self.hyperparams.gamma as f32, + ..Default::default() + }; + + match collector.collect_experiences(features_buf, targets_buf, &episode_starts, &config) { + Ok(batch) => { + let experiences = gpu_batch_to_experiences(&batch); + let count = experiences.len(); + info!("GPU collected {} experiences ({} episodes × {} timesteps)", + count, batch.n_episodes, batch.timesteps); + for exp in experiences { + self.store_experience(exp).await?; + } + true + } + Err(e) => { + warn!("GPU experience collection failed, falling back to CPU: {e}"); + false + } + } + } + } else { + false + }; + + #[cfg(not(feature = "cuda"))] + let gpu_experiences_collected = false; + + // Skip CPU experience collection if GPU already filled the replay buffer + if !gpu_experiences_collected { + // **PHASE 1: GPU-Optimized Experience Collection with Batched Action Selection** // Fill replay buffer with batched action selection (125× fewer GPU kernel launches) const ACTION_BATCH_SIZE: usize = 128; @@ -1789,10 +1889,7 @@ impl DQNTrainer { crate::safety::SafetyLevel::Strict => { return Err(anyhow::anyhow!("{} (stopping training)", msg)); }, - crate::safety::SafetyLevel::Normal => { - warn!("{}", msg); - }, - crate::safety::SafetyLevel::Permissive => { + crate::safety::SafetyLevel::Normal | crate::safety::SafetyLevel::Permissive => { debug!("{}", msg); }, } @@ -1816,7 +1913,7 @@ impl DQNTrainer { if usage_ratio > 0.9 { let usage_gb = current_usage as f64 / 1e9; let limit_gb = limit as f64 / 1e9; - warn!( + debug!( "SAFETY: GPU memory usage high at step {}: {:.1}% ({:.2} GB / {:.2} GB)", self.safety_step_counter, usage_ratio * 100.0, @@ -2003,6 +2100,8 @@ impl DQNTrainer { } } + } // end if !gpu_experiences_collected + // **PHASE 2: Batched Training from Replay Buffer** // Now that buffer is populated, perform batched training // This reduces train_step() calls from 1000×/epoch to ~8×/epoch (125× reduction) @@ -2030,13 +2129,10 @@ impl DQNTrainer { crate::safety::SafetyLevel::Strict => { return Err(anyhow::anyhow!("{} (stopping training)", msg)); }, - crate::safety::SafetyLevel::Normal => { - warn!("{} (continuing training)", msg); + crate::safety::SafetyLevel::Normal | crate::safety::SafetyLevel::Permissive => { + debug!("{} (continuing training)", msg); continue; // Skip this step }, - crate::safety::SafetyLevel::Permissive => { - debug!("{} (permissive mode)", msg); - }, } } @@ -2051,11 +2147,11 @@ impl DQNTrainer { return Err(anyhow::anyhow!("{} (stopping training)", msg)); }, crate::safety::SafetyLevel::Normal | crate::safety::SafetyLevel::Permissive => { - warn!("{}", msg); + debug!("{}", msg); }, } } else if grad_norm < 1e-6 && epoch > 10 { - warn!( + debug!( "SAFETY: Gradient vanishing detected at epoch {} - norm={:.2e} < 1e-6 threshold", epoch, grad_norm ); @@ -2065,7 +2161,7 @@ impl DQNTrainer { // WAVE 1.2 SAFETY #3: Q-Value Bounds Check (±1M limits) if q_value < -1_000_000.0 || q_value > 1_000_000.0 { - warn!( + debug!( "SAFETY: Q-value out of bounds at epoch {}: {:.2e} (bounds: ±1M)", epoch, q_value ); @@ -2088,7 +2184,7 @@ impl DQNTrainer { return Err(anyhow::anyhow!("{} (stopping training)", msg)); }, crate::safety::SafetyLevel::Normal | crate::safety::SafetyLevel::Permissive => { - warn!("{}", msg); + debug!("{}", msg); }, } } @@ -2265,7 +2361,7 @@ impl DQNTrainer { self.safety_loss_plateau_counter += 1; if self.safety_loss_plateau_counter >= 10 { - warn!( + debug!( "SAFETY: Training stuck for {} epochs (loss variance: {:.6}, mean: {:.6})", self.safety_loss_plateau_counter, std_dev, mean ); @@ -2480,7 +2576,7 @@ impl DQNTrainer { return Err(anyhow::anyhow!("{}", msg)); }, crate::safety::SafetyLevel::Normal => { - warn!("{} (continuing anyway)", msg); + debug!("{} (continuing anyway)", msg); }, _ => {}, } @@ -3415,6 +3511,58 @@ impl DQNTrainer { } } +// --------------------------------------------------------------------------- +// GPU batch → Experience conversion (Phase 3) +// --------------------------------------------------------------------------- + +/// Convert a GPU `ExperienceBatch` (flat arrays from CUDA kernel) into `Vec` +/// for insertion into the DQN replay buffer. +/// +/// State dimension is 54 (51 market + 3 portfolio features, assembled in-kernel). +/// Actions are mapped to `u8` for the `Experience` struct. Rewards are stored as +/// fixed-point (×10000) by `Experience::new`. +#[cfg(feature = "cuda")] +fn gpu_batch_to_experiences( + batch: &crate::cuda_pipeline::gpu_experience_collector::ExperienceBatch, +) -> Vec { + let state_dim = 54; + let total = batch.n_episodes * batch.timesteps; + let mut experiences = Vec::with_capacity(total); + + for ep in 0..batch.n_episodes { + for t in 0..batch.timesteps { + let idx = ep * batch.timesteps + t; + let state_start = idx * state_dim; + let state_end = state_start + state_dim; + let state: Vec = batch.states + .get(state_start..state_end) + .unwrap_or(&[]) + .to_vec(); + + // Clamp action to valid u8 range (0..45) + let action = batch.actions.get(idx).copied().unwrap_or(0).clamp(0, 44) as u8; + let reward = batch.rewards.get(idx).copied().unwrap_or(0.0); + let done = batch.done_flags.get(idx).copied().unwrap_or(0) != 0; + + // next_state = states[t+1] within same episode, or current state if done/last + let next_state = if !done && t + 1 < batch.timesteps { + let next_idx = ep * batch.timesteps + t + 1; + let ns_start = next_idx * state_dim; + let ns_end = ns_start + state_dim; + batch.states.get(ns_start..ns_end).unwrap_or(&[]).to_vec() + } else { + state.clone() + }; + + // Skip if state is empty (shouldn't happen, but defensive) + if state.len() == state_dim && next_state.len() == state_dim { + experiences.push(Experience::new(state, action, reward, next_state, done)); + } + } + } + experiences +} + #[cfg(test)] mod tests { use super::*; diff --git a/crates/ml/src/trainers/ppo.rs b/crates/ml/src/trainers/ppo.rs index 0ed3f5121..7d608caf6 100644 --- a/crates/ml/src/trainers/ppo.rs +++ b/crates/ml/src/trainers/ppo.rs @@ -63,6 +63,16 @@ pub struct PpoHyperparameters { pub sequence_length: usize, /// Maximum gradient norm for LSTM (default: 0.5) pub max_grad_norm_lstm: f64, + + // Phase 3: GPU experience collection kernel configuration + /// Number of parallel episodes per GPU kernel launch (default: 128, max: 256) + pub gpu_n_episodes: usize, + /// Timesteps per episode in GPU kernel (default: 500, max: 1000) + pub gpu_timesteps_per_episode: usize, + /// Initial capital for portfolio simulation (default: 1_000_000.0) + pub initial_capital: f64, + /// Average bid-ask spread for GPU portfolio simulation (default: 0.0001 = 1bp) + pub avg_spread: f64, } // REMOVED: Default implementation for PpoHyperparameters @@ -107,6 +117,12 @@ impl PpoHyperparameters { circuit_breaker_threshold: 5, sequence_length: 16, // 16 timesteps per BPTT window max_grad_norm_lstm: 0.5, // Tighter clipping for LSTM vs MLP (10.0) + + // Phase 3: GPU experience collection + gpu_n_episodes: 128, // Default: 128 (good for 4-8GB VRAM GPUs) + gpu_timesteps_per_episode: 500, // Default: 500 timesteps per episode + initial_capital: 1_000_000.0, // Default: $1M + avg_spread: 0.0001, // Default: 1bp (ES/NQ futures) } } } @@ -182,6 +198,15 @@ pub struct PpoTrainer { explained_variance_history: Arc>>, #[cfg(feature = "cuda")] gpu_ppo_collector: Option, + /// Raw cudarc features buffer for GPU experience kernel [num_bars * 51] + #[cfg(feature = "cuda")] + features_raw_cuda: Option>, + /// Raw cudarc targets buffer for GPU experience kernel [num_bars * 4] + #[cfg(feature = "cuda")] + targets_raw_cuda: Option>, + /// Number of bars in the raw data buffers (needed to configure kernel) + #[cfg(feature = "cuda")] + raw_data_num_bars: usize, } impl std::fmt::Debug for PpoTrainer { @@ -274,9 +299,65 @@ impl PpoTrainer { explained_variance_history: Arc::new(Mutex::new(VecDeque::new())), #[cfg(feature = "cuda")] gpu_ppo_collector: None, + #[cfg(feature = "cuda")] + features_raw_cuda: None, + #[cfg(feature = "cuda")] + targets_raw_cuda: None, + #[cfg(feature = "cuda")] + raw_data_num_bars: 0, }) } + /// Upload raw market data to GPU for the experience collection kernel. + /// + /// Call this before `train()` when raw `([f64; 51], Vec)` data is available. + /// The kernel needs flat `[num_bars * 51]` features and `[num_bars * 4]` targets + /// as `CudaSlice` buffers (not candle Tensors). + /// + /// No-op if not on a CUDA device. + #[cfg(feature = "cuda")] + pub fn set_raw_market_data( + &mut self, + data: &[([f64; 51], Vec)], + ) -> Result<(), MLError> { + if !self.device.is_cuda() { + return Ok(()); + } + let cuda_dev = match &self.device { + candle_core::Device::Cuda(d) => d, + candle_core::Device::Cpu | candle_core::Device::Metal(_) => return Ok(()), + }; + let stream = cuda_dev.cuda_stream(); + let num_bars = data.len(); + + // Upload features [num_bars * 51] + let mut flat_features = Vec::with_capacity(num_bars * 51); + for (features, _) in data { + for &v in features.iter() { + flat_features.push(v as f32); + } + } + let features_buf = stream.memcpy_stod(&flat_features) + .map_err(|e| MLError::ModelError(format!("CUDA features upload: {e}")))?; + self.features_raw_cuda = Some(features_buf); + + // Upload targets [num_bars * 4] + let mut flat_targets = Vec::with_capacity(num_bars * 4); + for (_, targets) in data { + for i in 0..4 { + flat_targets.push(targets.get(i).copied().unwrap_or(0.0) as f32); + } + } + let targets_buf = stream.memcpy_stod(&flat_targets) + .map_err(|e| MLError::ModelError(format!("CUDA targets upload: {e}")))?; + self.targets_raw_cuda = Some(targets_buf); + self.raw_data_num_bars = num_bars; + + info!("PPO raw market data uploaded: {} bars ({:.1} MB)", + num_bars, (num_bars * 55 * 4) as f64 / 1_048_576.0); + Ok(()) + } + /// Train PPO model /// /// # Arguments @@ -333,6 +414,9 @@ impl PpoTrainer { let model = self.model.lock().await; let actor_vars = model.actor.vars(); let critic_vars = model.critic.vars(); + let initial_capital = self.hyperparams.initial_capital as f32; + let avg_spread = self.hyperparams.avg_spread as f32; + let cash_reserve_pct = self.hyperparams.cash_reserve_pct as f32 / 100.0; match (|| -> Result<_, MLError> { let cuda_device = match &self.device { candle_core::Device::Cuda(d) => d, @@ -344,9 +428,9 @@ impl PpoTrainer { actor_vars, critic_vars, &candle_nn::VarMap::new(), // curiosity placeholder — PPO curiosity integration is future work - 1_000_000.0, // initial_capital - 0.0001, // avg_spread - 0.20, // cash_reserve_pct + initial_capital, + avg_spread, + cash_reserve_pct, ) })() { Ok(collector) => { @@ -365,19 +449,80 @@ impl PpoTrainer { for epoch in 0..self.hyperparams.epochs { debug!("Training epoch {}/{}", epoch + 1, self.hyperparams.epochs); - // Step 1: Collect rollouts (trajectories) - let trajectories = if let Some(n) = self.num_envs { - if n > 1 { - self.collect_vectorized_rollouts(&market_data).await? + // Phase 3: GPU experience collection — bypasses CPU collect_rollouts + prepare_training_batch + #[cfg(feature = "cuda")] + let gpu_batch_result: Option = if let ( + Some(ref mut collector), + Some(ref features_buf), + Some(ref targets_buf), + ) = ( + &mut gpu_ppo_collector, + &self.features_raw_cuda, + &self.targets_raw_cuda, + ) { + use crate::cuda_pipeline::gpu_ppo_collector::PpoCollectorConfig; + + let n_episodes = self.hyperparams.gpu_n_episodes.min(256) as i32; + let timesteps = self.hyperparams.gpu_timesteps_per_episode.min(1000) as i32; + let total_bars = self.raw_data_num_bars as i32; + let usable_bars = (total_bars - timesteps).max(1); + let stride = (usable_bars / n_episodes).max(1); + let episode_starts: Vec = (0..n_episodes) + .map(|i| (i * stride).rem_euclid(usable_bars)) + .collect(); + + let config = PpoCollectorConfig { + n_episodes, + timesteps_per_episode: timesteps, + total_bars, + gamma: self.hyperparams.gamma as f32, + gae_lambda: self.hyperparams.gae_lambda, + ..Default::default() + }; + + // Reset per-episode state before each epoch + let initial_capital = self.hyperparams.initial_capital as f32; + let avg_spread = self.hyperparams.avg_spread as f32; + let cash_reserve_pct = self.hyperparams.cash_reserve_pct as f32 / 100.0; + + if let Err(e) = collector.reset_episodes(initial_capital, avg_spread, cash_reserve_pct) { + warn!("GPU PPO episode reset failed: {e}"); + None } else { - self.collect_rollouts(&market_data).await? + match collector.collect_experiences(features_buf, targets_buf, &episode_starts, &config) { + Ok(batch) => { + info!("GPU PPO collected {} experiences ({} episodes × {} timesteps)", + batch.n_episodes * batch.timesteps, batch.n_episodes, batch.timesteps); + Some(gpu_batch_to_trajectory_batch(&batch)) + } + Err(e) => { + warn!("GPU PPO collection failed, CPU fallback: {e}"); + None + } + } } } else { - self.collect_rollouts(&market_data).await? + None }; - // Step 2: Prepare training batch with GAE - let mut training_batch = self.prepare_training_batch(trajectories)?; + #[cfg(not(feature = "cuda"))] + let gpu_batch_result: Option = None; + + // Step 1+2: Use GPU batch if available, otherwise CPU collect_rollouts + prepare_training_batch + let mut training_batch = if let Some(batch) = gpu_batch_result { + batch + } else { + let trajectories = if let Some(n) = self.num_envs { + if n > 1 { + self.collect_vectorized_rollouts(&market_data).await? + } else { + self.collect_rollouts(&market_data).await? + } + } else { + self.collect_rollouts(&market_data).await? + }; + self.prepare_training_batch(trajectories)? + }; // Step 2.5: Pre-train value network (first 10 epochs only) if epoch < 10 { @@ -1129,6 +1274,68 @@ impl PpoTrainer { } } +/// Convert GPU-collected PPO experience batch to TrajectoryBatch for model.update(). +/// +/// The GPU kernel computes advantages and returns in-kernel via GAE, so this +/// bypasses the CPU collect_rollouts() + prepare_training_batch() pipeline entirely. +/// The `trajectories` field is left empty since `update()` only uses flattened fields. +#[cfg(feature = "cuda")] +fn gpu_batch_to_trajectory_batch( + batch: &crate::cuda_pipeline::gpu_ppo_collector::PpoExperienceBatch, +) -> TrajectoryBatch { + let state_dim = 54; + let total = batch.n_episodes * batch.timesteps; + + let states: Vec> = batch + .states + .chunks(state_dim) + .take(total) + .map(|c| c.to_vec()) + .collect(); + + let actions: Vec = batch + .actions + .iter() + .take(total) + .map(|&a| { + TradingAction::from_int(a.clamp(0, 44) as u8).unwrap_or(TradingAction::Hold) + }) + .collect(); + + let log_probs: Vec = batch.log_probs.iter().take(total).copied().collect(); + + // V(s) = R(s) - A(s), since the kernel stores R = A + V + let values: Vec = batch + .returns + .iter() + .zip(&batch.advantages) + .take(total) + .map(|(r, a)| r - a) + .collect(); + + let dones: Vec = batch + .done_flags + .iter() + .take(total) + .map(|&d| d != 0) + .collect(); + + // Rewards not needed — advantages/returns already computed by kernel + let rewards = vec![0.0_f32; states.len()]; + + TrajectoryBatch { + trajectories: vec![], + states, + actions, + log_probs, + values, + rewards, + dones, + advantages: batch.advantages.iter().take(total).copied().collect(), + returns: batch.returns.iter().take(total).copied().collect(), + } +} + #[cfg(test)] mod tests { use super::*; diff --git a/crates/ml/tests/ppo_recurrent_performance_tests.rs b/crates/ml/tests/ppo_recurrent_performance_tests.rs index edcf2c071..223b3dbf8 100644 --- a/crates/ml/tests/ppo_recurrent_performance_tests.rs +++ b/crates/ml/tests/ppo_recurrent_performance_tests.rs @@ -96,6 +96,10 @@ fn create_test_hyperparams(_use_lstm: bool, sequence_length: usize) -> PpoHyperp circuit_breaker_threshold: 5, sequence_length, max_grad_norm_lstm: 0.5, + gpu_n_episodes: 128, + gpu_timesteps_per_episode: 500, + initial_capital: 1_000_000.0, + avg_spread: 0.0001, } }