diff --git a/crates/ml/src/dqn/curiosity.rs b/crates/ml/src/dqn/curiosity.rs index 5970144c0..27d3d78e7 100644 --- a/crates/ml/src/dqn/curiosity.rs +++ b/crates/ml/src/dqn/curiosity.rs @@ -104,6 +104,11 @@ impl ForwardDynamicsModel { Ok(pred) } + /// Get the VarMap for weight extraction (GPU sync). + fn vars(&self) -> &VarMap { + &self.vars + } + /// Train forward model on (state, action, next_state) transition /// /// # Arguments @@ -219,6 +224,11 @@ impl CuriosityModule { Ok(novelty_bonus) } + + /// Get the forward model's VarMap for GPU weight extraction. + pub fn forward_model_vars(&self) -> &candle_nn::VarMap { + self.forward_model.vars() + } } #[cfg(test)] diff --git a/crates/ml/src/trainers/dqn/trainer.rs b/crates/ml/src/trainers/dqn/trainer.rs index 0ecaf5534..a7ad28964 100644 --- a/crates/ml/src/trainers/dqn/trainer.rs +++ b/crates/ml/src/trainers/dqn/trainer.rs @@ -196,6 +196,18 @@ pub struct DQNTrainer { /// Pre-uploaded GPU training data (set once, reused across epochs) gpu_data: Option, + + /// GPU portfolio simulator for CUDA-accelerated experience collection + #[cfg(feature = "cuda")] + gpu_portfolio_sim: Option, + + /// Raw cudarc targets buffer for CUDA kernel (parallel to candle Tensor in gpu_data) + #[cfg(feature = "cuda")] + targets_raw_cuda: Option>, + + /// GPU experience collector for zero-roundtrip CUDA kernel (Phase 2b) + #[cfg(feature = "cuda")] + gpu_experience_collector: Option, } impl std::fmt::Debug for DQNTrainer { @@ -745,6 +757,12 @@ impl DQNTrainer { // GPU pipeline: pre-uploaded training data (initialized lazily at first epoch) gpu_data: None, + #[cfg(feature = "cuda")] + gpu_portfolio_sim: None, + #[cfg(feature = "cuda")] + targets_raw_cuda: None, + #[cfg(feature = "cuda")] + gpu_experience_collector: None, }) } @@ -1307,6 +1325,110 @@ impl DQNTrainer { } } + // Phase 1b: Upload raw targets 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 num_bars = training_data.len(); + + // Build flat targets (same as DqnGpuData::upload but for cudarc) + let mut flat_targets = Vec::with_capacity(num_bars * target_dim); + for (_, targets) in &training_data { + for i in 0..target_dim { + flat_targets.push(targets.get(i).copied().unwrap_or(0.0) as f32); + } + } + + match stream.memcpy_stod(&flat_targets) { + Ok(buf) => { + info!("CUDA targets_raw uploaded: {} bars × 4 ({:.1} KB)", + num_bars, (num_bars * target_dim * 4) as f64 / 1024.0); + self.targets_raw_cuda = Some(buf); + } + Err(e) => { + debug!("CUDA targets_raw upload failed (CPU fallback): {}", e); + } + } + + // Initialize GPU portfolio simulator + if self.gpu_portfolio_sim.is_none() { + if let Some(ref _targets_buf) = self.targets_raw_cuda { + use crate::cuda_pipeline::gpu_portfolio::GpuPortfolioSimulator; + match GpuPortfolioSimulator::new( + stream, + self.hyperparams.initial_capital as f32, + 0.0001, // avg_spread + self.hyperparams.cash_reserve_percent as f32, + self.max_position as f32, + EPISODE_LENGTH, + training_data.len(), + ) { + Ok(sim) => { + info!("GPU portfolio simulator initialized"); + self.gpu_portfolio_sim = Some(sim); + } + Err(e) => { + warn!("GPU portfolio sim init failed (CPU fallback): {}", e); + } + } + } + } + } + } + + // Phase 1c: Initialize zero-roundtrip GPU experience collector (once) + #[cfg(feature = "cuda")] + if self.gpu_experience_collector.is_none() { + if let candle_core::Device::Cuda(ref cuda_dev) = self.device { + use crate::cuda_pipeline::gpu_experience_collector::GpuExperienceCollector; + let stream = cuda_dev.cuda_stream(); + + // Read-lock agent to check for dueling networks + let agent = self.agent.read().await; + let init_result = if let DQNAgentType::Standard(ref dqn) = &*agent { + match ( + dqn.dueling_q_network.as_ref(), + dqn.dueling_target_network.as_ref(), + self.curiosity_module.as_ref(), + ) { + (Some(online), Some(target), Some(curiosity)) => { + Some(GpuExperienceCollector::new( + stream, + online.vars(), + target.vars(), + curiosity.forward_model_vars(), + self.hyperparams.initial_capital as f32, + 0.0001_f32, // avg_spread + self.hyperparams.cash_reserve_percent as f32, + )) + } + _ => { + debug!("GPU experience collector skipped: requires dueling networks + curiosity module"); + None + } + } + } else { + debug!("GPU experience collector skipped: regime-conditional DQN not supported"); + None + }; + drop(agent); + + if let Some(result) = init_result { + match result { + Ok(collector) => { + info!("GPU experience collector initialized (zero-roundtrip CUDA kernel)"); + self.gpu_experience_collector = Some(collector); + } + Err(e) => { + warn!("GPU experience collector init failed (CPU fallback): {}", e); + } + } + } + } + } + let mut epoch_loss = 0.0; let mut epoch_q_value = 0.0; let mut epoch_gradient_norm = 0.0; @@ -1345,25 +1467,83 @@ impl DQNTrainer { let batch_end = ((batch_idx + 1) * ACTION_BATCH_SIZE).min(total_samples); let batch_indices: Vec = (batch_start..batch_end).collect(); - // Convert batch to states for batched action selection - let states: Result> = batch_indices - .iter() - .map(|&i| { - let target = &training_data[i].1; - let current_close = if target.len() >= 2 { - target[0] - } else { - training_data[i].0[3] - }; - let close_price = rust_decimal::Decimal::try_from(current_close) - .unwrap_or(rust_decimal::Decimal::ZERO); - self.feature_vector_to_state(&training_data[i].0, Some(close_price)) - }) - .collect(); - let states = states?; + // Build batch states — GPU path skips ~770 Vec allocs per batch + let (batch_tensor, states) = if let Some(ref gpu_data) = self.gpu_data { + // GPU path: build [batch_size, 54] directly from pre-uploaded features + let current_price_f32 = { + let t = gpu_data.bar_target_values(batch_start)?; + if t[2] != 0.0 { t[2] } else { t[0] } + }; + let portfolio_features = self.portfolio_tracker.get_portfolio_features(current_price_f32); + let bt = gpu_data.build_batch_states( + batch_start, + batch_end - batch_start, + &portfolio_features, + &self.device, + )?; - // Batched action selection (single GPU kernel launch) ✅ - let actions = self.select_actions_batch(&states).await?; + // Still need TradingStates for reward calculation and experience storage + let cpu_states: Result> = batch_indices + .iter() + .map(|&i| { + let target = &training_data[i].1; + let current_close = if target.len() >= 2 { + target[0] + } else { + training_data[i].0[3] + }; + let close_price = rust_decimal::Decimal::try_from(current_close) + .unwrap_or(rust_decimal::Decimal::ZERO); + self.feature_vector_to_state(&training_data[i].0, Some(close_price)) + }) + .collect(); + (Some(bt), cpu_states?) + } else { + // CPU fallback: original path + let cpu_states: Result> = batch_indices + .iter() + .map(|&i| { + let target = &training_data[i].1; + let current_close = if target.len() >= 2 { + target[0] + } else { + training_data[i].0[3] + }; + let close_price = rust_decimal::Decimal::try_from(current_close) + .unwrap_or(rust_decimal::Decimal::ZERO); + self.feature_vector_to_state(&training_data[i].0, Some(close_price)) + }) + .collect(); + (None, cpu_states?) + }; + + // Batched action selection — GPU tensor path skips state→Vec→flatten→Tensor ✅ + let actions = if let Some(ref bt) = batch_tensor { + self.select_actions_batch_gpu(bt).await? + } else { + self.select_actions_batch(&states).await? + }; + + // GPU portfolio simulation for the batch (if CUDA available) + #[cfg(feature = "cuda")] + // GPU portfolio sim results: computed for metrics/validation, CPU inner loop remains source of truth. + // Phase 2b will wire these into the inner loop after CPU/GPU result equivalence is validated. + let _gpu_sim_results = { + if let (Some(ref mut sim), Some(ref targets_buf)) = + (&mut self.gpu_portfolio_sim, &self.targets_raw_cuda) + { + let action_indices: Vec = actions.iter().map(|a| a.to_index() as i32).collect(); + match sim.simulate_batch(targets_buf, &action_indices, batch_start) { + Ok(result) => Some(result), + Err(e) => { + warn!("GPU portfolio sim failed (CPU fallback): {}", e); + None + } + } + } else { + None + } + }; // Store experiences with batched actions for (idx_in_batch, &i) in batch_indices.iter().enumerate() { @@ -1974,6 +2154,31 @@ impl DQNTrainer { } } + // Phase 2b: Sync GPU weight copies after training updates + #[cfg(feature = "cuda")] + if let Some(ref mut collector) = self.gpu_experience_collector { + let agent = self.agent.read().await; + if let DQNAgentType::Standard(ref dqn) = &*agent { + if let Some(ref online) = dqn.dueling_q_network { + if let Err(e) = collector.sync_online_weights(online.vars()) { + warn!("GPU online weight sync failed: {}", e); + } + } + if let Some(ref target) = dqn.dueling_target_network { + if let Err(e) = collector.sync_target_weights(target.vars()) { + warn!("GPU target weight sync failed: {}", e); + } + } + } + drop(agent); + + if let Some(ref curiosity) = self.curiosity_module { + if let Err(e) = collector.sync_curiosity_weights_from(curiosity.forward_model_vars()) { + warn!("GPU curiosity weight sync failed: {}", e); + } + } + } + let epoch_duration = epoch_start.elapsed(); // Calculate epoch metrics (average over training steps, not samples) @@ -2609,6 +2814,53 @@ impl DQNTrainer { Ok(actions) } + /// GPU-optimized batch action selection using pre-built state tensor. + /// + /// Skips the state→Vec→flatten→Tensor pipeline (~130 allocs per batch). + async fn select_actions_batch_gpu(&self, batch_tensor: &Tensor) -> Result> { + let batch_size = batch_tensor.dims()[0]; + if batch_size == 0 { + return Ok(Vec::new()); + } + + let agent = self.agent.read().await; + + // WAVE 16S: Volatility-adjusted epsilon + let base_epsilon = agent.get_epsilon() as f64; + let adjusted_epsilon = self.calculate_volatility_adjusted_epsilon(base_epsilon); + let epsilon = adjusted_epsilon as f32; + debug!("Epsilon (GPU path): base={:.4}, volatility-adjusted={:.4}", base_epsilon, adjusted_epsilon); + + // Single forward pass — tensor already on GPU, no construction needed + let batch_q_values = agent + .forward(batch_tensor) + .map_err(|e| anyhow::anyhow!("GPU batched forward pass failed: {}", e))?; + + drop(agent); + + // GPU argmax + epsilon-greedy (same as CPU path) + let greedy_action_indices = batch_q_values + .argmax(1) + .map_err(|e| anyhow::anyhow!("Failed to compute argmax on GPU: {}", e))? + .to_vec1::() + .map_err(|e| anyhow::anyhow!("Failed to transfer argmax results to CPU: {}", e))?; + + let mut actions = Vec::with_capacity(batch_size); + let mut rng = rand::thread_rng(); + for i in 0..batch_size { + use rand::Rng; + let action_idx = if rng.gen::() < epsilon { + rng.gen_range(0..45) + } else { + greedy_action_indices[i] as usize + }; + let action = FactoredAction::from_index(action_idx) + .map_err(|e| anyhow::anyhow!("Invalid action index {}: {}", action_idx, e))?; + actions.push(action); + } + Ok(actions) + } + /// Epsilon-greedy action selection async fn epsilon_greedy_action(&self, state: &Tensor) -> Result { use rand::Rng;