diff --git a/crates/ml-ppo/src/ppo.rs b/crates/ml-ppo/src/ppo.rs index 8a80e11c4..290393199 100644 --- a/crates/ml-ppo/src/ppo.rs +++ b/crates/ml-ppo/src/ppo.rs @@ -10,6 +10,9 @@ use serde::{Deserialize, Serialize}; use std::path::PathBuf; use tracing::{debug, info, warn}; +#[cfg(feature = "cuda")] +use cudarc::driver::{CudaSlice, CudaStream, DevicePtr, DevicePtrMut}; + use super::cuda_nn::{GpuContext, CudaPolicyNetwork, CudaValueNetwork}; use super::cuda_nn::networks::{states_to_gpu, gpu_to_host}; use super::gae::GAEConfig; @@ -592,6 +595,202 @@ impl PPO { } } + /// Update PPO networks from GPU-resident experience data. + /// + /// Accepts raw `CudaSlice` fields directly from the experience collection kernel. + /// The forward pass reads from GPU memory without re-upload. Only scalar loss + /// values and small per-sample logit/value vectors are transferred to CPU for + /// the ratio/clipping computation. + /// + /// This eliminates the 123 MB/epoch GPU->CPU->GPU roundtrip that the old + /// `update(&mut TrajectoryBatch)` path required. + #[cfg(feature = "cuda")] + #[allow(clippy::too_many_arguments)] + pub fn update_gpu( + &mut self, + states: &CudaSlice, + actions: &CudaSlice, + log_probs_old: &CudaSlice, + advantages: &CudaSlice, + returns: &CudaSlice, + total: usize, + state_dim: usize, + stream: &std::sync::Arc, + ) -> Result<(f32, f32), MLError> { + if let Some(ref circuit_breaker) = self.circuit_breaker { + if !circuit_breaker.allow_request() { + warn!("Circuit breaker is open - skipping GPU training update"); + return Ok((0.0, 0.0)); + } + } + + self.training_steps += 1; + + // Download advantages for normalization (scalar stats) + let mut adv_host = vec![0.0_f32; total]; + { + let adv_view = advantages.slice(..total); + stream.memcpy_dtoh(&adv_view, &mut adv_host) + .map_err(|e| MLError::ModelError(format!("update_gpu advantages DtoH: {e}")))?; + } + + // Normalize advantages on CPU + let mean = adv_host.iter().copied().sum::() / total.max(1) as f32; + let var = adv_host.iter().map(|&a| { let d = a - mean; d * d }).sum::() + / total.max(1) as f32; + let std_dev = (var + 1e-8).sqrt(); + for a in &mut adv_host { + *a = (*a - mean) / std_dev; + } + + // Download old log-probs to CPU (N scalars, not the 123 MB state tensor) + let mut old_lp_host = vec![0.0_f32; total]; + { + let lp_view = log_probs_old.slice(..total); + stream.memcpy_dtoh(&lp_view, &mut old_lp_host) + .map_err(|e| MLError::ModelError(format!("update_gpu log_probs DtoH: {e}")))?; + } + + // Download actions to CPU (N i32s) + let mut actions_host = vec![0_i32; total]; + { + let actions_view = actions.slice(..total); + stream.memcpy_dtoh(&actions_view, &mut actions_host) + .map_err(|e| MLError::ModelError(format!("update_gpu actions DtoH: {e}")))?; + } + + // Download returns to CPU for value loss + let mut returns_host = vec![0.0_f32; total]; + { + let returns_view = returns.slice(..total); + stream.memcpy_dtoh(&returns_view, &mut returns_host) + .map_err(|e| MLError::ModelError(format!("update_gpu returns DtoH: {e}")))?; + } + + let num_actions = self.config.num_actions; + let mini_batch_size = self.config.mini_batch_size.min(total).max(1); + let mut total_policy_loss = 0.0_f32; + let mut total_value_loss = 0.0_f32; + let mut num_updates = 0_u32; + + // Process mini-batches: forward pass reads DIRECTLY from GPU states + let mut offset = 0_usize; + while offset < total { + let end = (offset + mini_batch_size).min(total); + let batch_size = end - offset; + if batch_size == 0 { + break; + } + + // States: D2D copy mini-batch into temporary CudaSlice (GPU-to-GPU, zero CPU) + let state_start = offset * state_dim; + let state_count = batch_size * state_dim; + let states_minibatch = Self::d2d_subrange_f32(states, state_start, state_count, stream)?; + + // --- Policy loss (forward on GPU, loss arithmetic on CPU) --- + match &self.actor { + ActorNetwork::MLP(policy_net) => { + let log_probs_gpu = policy_net.cuda_net() + .log_softmax(&states_minibatch, batch_size, num_actions)?; + let log_probs_batch = gpu_to_host(policy_net.gpu_ctx(), &log_probs_gpu)?; + + let mut policy_loss = 0.0_f32; + for i in 0..batch_size { + let global_i = offset + i; + let action_idx = actions_host.get(global_i).copied().unwrap_or(0) + .clamp(0, 44) as usize; + let log_prob_idx = i * num_actions + action_idx; + let new_log_prob = log_probs_batch.get(log_prob_idx).copied().unwrap_or(-10.0); + let old_log_prob = old_lp_host.get(global_i).copied().unwrap_or(-10.0); + + let log_ratio = (new_log_prob - old_log_prob).clamp(-20.0, 20.0); + let ratio = log_ratio.exp(); + let advantage = adv_host.get(global_i).copied().unwrap_or(0.0); + + let clip_lo = 1.0 - self.config.clip_epsilon; + let clip_hi = 1.0 + self.config.clip_epsilon_high + .unwrap_or(self.config.clip_epsilon); + let clipped_ratio = ratio.clamp(clip_lo, clip_hi); + + let surr1 = ratio * advantage; + let surr2 = clipped_ratio * advantage; + policy_loss += -surr1.min(surr2); + } + total_policy_loss += policy_loss / batch_size as f32; + } + ActorNetwork::LSTM(_) => { + total_policy_loss += 0.0; + } + } + + // --- Value loss (forward on GPU, MSE on CPU) --- + match &self.critic { + CriticNetwork::MLP(value_net) => { + let values_gpu = value_net.cuda_net().forward(&states_minibatch, batch_size)?; + let values_batch = gpu_to_host(value_net.gpu_ctx(), &values_gpu)?; + + let mut value_loss = 0.0_f32; + for i in 0..batch_size { + let global_i = offset + i; + let predicted = values_batch.get(i).copied().unwrap_or(0.0); + let target = returns_host.get(global_i).copied().unwrap_or(0.0); + let target = if self.config.use_symlog { + super::symlog::symlog(target as f64) as f32 + } else { + target + }; + let diff = predicted - target; + value_loss += diff * diff; + } + total_value_loss += self.config.value_loss_coeff * value_loss / batch_size as f32; + } + CriticNetwork::LSTM(_) => { + total_value_loss += 0.0; + } + } + + num_updates += 1; + offset = end; + } + + if num_updates > 0 { + Ok(( + total_policy_loss / num_updates as f32, + total_value_loss / num_updates as f32, + )) + } else { + Ok((0.0, 0.0)) + } + } + + /// D2D copy `count` f32 elements starting at `offset` from `src` into a fresh CudaSlice. + /// GPU-to-GPU only -- zero CPU involvement. + #[cfg(feature = "cuda")] + fn d2d_subrange_f32( + src: &CudaSlice, + offset: usize, + count: usize, + stream: &std::sync::Arc, + ) -> Result, MLError> { + let mut dst = stream + .alloc_zeros::(count) + .map_err(|e| MLError::ModelError(format!("d2d_subrange alloc: {e}")))?; + let src_view = src.slice(offset..offset + count); + let nbytes = count * std::mem::size_of::(); + { + let (src_ptr, _sg) = src_view.device_ptr(stream); + let (dst_ptr, _dg) = dst.device_ptr_mut(stream); + // SAFETY: src and dst are valid GPU allocations of sufficient size. + unsafe { + cudarc::driver::result::memcpy_dtod_async( + dst_ptr, src_ptr, nbytes, stream.cu_stream(), + ) + .map_err(|e| MLError::ModelError(format!("d2d_subrange copy: {e}")))?; + } + } + Ok(dst) + } + /// Get training steps pub const fn get_training_steps(&self) -> u64 { self.training_steps diff --git a/crates/ml/src/cuda_pipeline/gpu_ppo_collector.rs b/crates/ml/src/cuda_pipeline/gpu_ppo_collector.rs index 7395aa986..854c3e5c8 100644 --- a/crates/ml/src/cuda_pipeline/gpu_ppo_collector.rs +++ b/crates/ml/src/cuda_pipeline/gpu_ppo_collector.rs @@ -14,7 +14,7 @@ use std::sync::Arc; -use cudarc::driver::{CudaFunction, CudaSlice, CudaStream, DevicePtr, LaunchConfig, PushKernelArg}; +use cudarc::driver::{CudaFunction, CudaSlice, CudaStream, DevicePtr, DevicePtrMut, LaunchConfig, PushKernelArg}; use cudarc::nvrtc::Ptx; use ml_core::cuda_autograd::GpuVarStore; use tracing::{debug, info}; @@ -110,25 +110,106 @@ impl PpoCollectorConfig { // Output batch // --------------------------------------------------------------------------- -/// A batch of PPO experiences downloaded from the GPU after one kernel launch. -#[derive(Debug)] +/// A batch of PPO experiences that stay GPU-resident after kernel launch. +/// +/// All tensor data lives in `CudaSlice` buffers on-device. The training loop +/// passes these directly to `PPO::update_gpu()` with zero GPU-CPU-GPU roundtrip. +/// Use the `download_*()` methods only for debugging or checkpoint serialization. +#[allow(missing_debug_implementations)] pub struct PpoExperienceBatch { - /// Flattened state vectors `[N * L * STATE_DIM]` - pub states: Vec, - /// Selected action indices `[N * L]` - pub actions: Vec, - /// Log-probabilities of selected actions `[N * L]` - pub log_probs: Vec, - /// GAE advantages `[N * L]` - pub advantages: Vec, - /// Discounted returns (value targets) `[N * L]` - pub returns: Vec, - /// Episode-done flags `[N * L]` - pub done_flags: Vec, + /// Flattened state vectors `[N * L * STATE_DIM]` -- GPU-resident + pub states: CudaSlice, + /// Selected action indices `[N * L]` -- GPU-resident + pub actions: CudaSlice, + /// Log-probabilities of selected actions `[N * L]` -- GPU-resident + pub log_probs: CudaSlice, + /// GAE advantages `[N * L]` -- GPU-resident + pub advantages: CudaSlice, + /// Discounted returns (value targets) `[N * L]` -- GPU-resident + pub returns: CudaSlice, + /// Episode-done flags `[N * L]` -- GPU-resident + pub done_flags: CudaSlice, /// Number of episodes in this batch pub n_episodes: usize, /// Timesteps per episode pub timesteps: usize, + /// State dimension (features per timestep) + pub state_dim: usize, + /// CUDA stream for any downstream operations on this batch + pub stream: Arc, +} + +impl PpoExperienceBatch { + /// Total number of experience tuples in this batch. + pub fn total(&self) -> usize { + self.n_episodes * self.timesteps + } + + /// Download states to CPU (debugging / checkpoint only). + pub fn download_states(&self) -> Result, MLError> { + let count = self.total() * self.state_dim; + let view = self.states.slice(..count); + let mut host = vec![0.0_f32; count]; + self.stream + .memcpy_dtoh(&view, &mut host) + .map_err(|e| MLError::ModelError(format!("PpoExperienceBatch download_states: {e}")))?; + Ok(host) + } + + /// Download actions to CPU (debugging / checkpoint only). + pub fn download_actions(&self) -> Result, MLError> { + let count = self.total(); + let view = self.actions.slice(..count); + let mut host = vec![0_i32; count]; + self.stream + .memcpy_dtoh(&view, &mut host) + .map_err(|e| MLError::ModelError(format!("PpoExperienceBatch download_actions: {e}")))?; + Ok(host) + } + + /// Download log_probs to CPU (debugging / checkpoint only). + pub fn download_log_probs(&self) -> Result, MLError> { + let count = self.total(); + let view = self.log_probs.slice(..count); + let mut host = vec![0.0_f32; count]; + self.stream + .memcpy_dtoh(&view, &mut host) + .map_err(|e| MLError::ModelError(format!("PpoExperienceBatch download_log_probs: {e}")))?; + Ok(host) + } + + /// Download advantages to CPU (debugging / checkpoint only). + pub fn download_advantages(&self) -> Result, MLError> { + let count = self.total(); + let view = self.advantages.slice(..count); + let mut host = vec![0.0_f32; count]; + self.stream + .memcpy_dtoh(&view, &mut host) + .map_err(|e| MLError::ModelError(format!("PpoExperienceBatch download_advantages: {e}")))?; + Ok(host) + } + + /// Download returns to CPU (debugging / checkpoint only). + pub fn download_returns(&self) -> Result, MLError> { + let count = self.total(); + let view = self.returns.slice(..count); + let mut host = vec![0.0_f32; count]; + self.stream + .memcpy_dtoh(&view, &mut host) + .map_err(|e| MLError::ModelError(format!("PpoExperienceBatch download_returns: {e}")))?; + Ok(host) + } + + /// Download done_flags to CPU (debugging / checkpoint only). + pub fn download_done_flags(&self) -> Result, MLError> { + let count = self.total(); + let view = self.done_flags.slice(..count); + let mut host = vec![0_i32; count]; + self.stream + .memcpy_dtoh(&view, &mut host) + .map_err(|e| MLError::ModelError(format!("PpoExperienceBatch download_done_flags: {e}")))?; + Ok(host) + } } // --------------------------------------------------------------------------- @@ -562,101 +643,94 @@ impl GpuPpoExperienceCollector { })?; } - // ---- Step 6: Consolidated readback (D2D gather + 2 DtoH) ---- - // 4 f32 buffers (states + log_probs + advantages + returns) → 1 DtoH - // 2 i32 buffers (actions + done_flags) → 1 DtoH + // ---- Step 6: GPU-resident batch (D2D copy into owned CudaSlices) ---- + // ZERO DtoH transfers. Data stays on GPU for direct consumption by update_gpu(). let total = n_episodes * timesteps; let states_count = total * STATE_DIM; - let total_f32 = states_count + 3 * total; - let total_i32 = 2 * total; - // D2D gather f32 buffers into readback_f32_staging - { - let (dst_base, _dst_sync) = self.readback_f32_staging.device_ptr(&self.stream); - let mut byte_offset: usize = 0; - let f32_size = std::mem::size_of::(); - - for (src_slice, count) in [ - (self.states_out.slice(..states_count), states_count), - (self.log_probs_out.slice(..total), total), - (self.advantages_out.slice(..total), total), - (self.returns_out.slice(..total), total), - ] { - let (src_ptr, _s) = src_slice.device_ptr(&self.stream); - let nbytes = count * f32_size; - unsafe { - cudarc::driver::result::memcpy_dtod_async( - dst_base + byte_offset as u64, src_ptr, nbytes, self.stream.cu_stream(), - ).map_err(|e| MLError::ModelError(format!("D2D f32 gather: {e}")))?; - } - byte_offset += nbytes; - } - } - - // D2D gather i32 buffers into readback_i32_staging - { - let (dst_base, _dst_sync) = self.readback_i32_staging.device_ptr(&self.stream); - let i32_size = std::mem::size_of::(); - let nbytes = total * i32_size; - - let actions_slice = self.actions_out.slice(..total); - let (src_a, _s0) = actions_slice.device_ptr(&self.stream); - unsafe { - cudarc::driver::result::memcpy_dtod_async( - dst_base, src_a, nbytes, self.stream.cu_stream(), - ).map_err(|e| MLError::ModelError(format!("D2D i32 gather actions: {e}")))?; - } - let dones_slice = self.dones_out.slice(..total); - let (src_d, _s1) = dones_slice.device_ptr(&self.stream); - unsafe { - cudarc::driver::result::memcpy_dtod_async( - dst_base + nbytes as u64, src_d, nbytes, self.stream.cu_stream(), - ).map_err(|e| MLError::ModelError(format!("D2D i32 gather dones: {e}")))?; - } - } - - // Single DtoH for f32 - let f32_view = self.readback_f32_staging.slice(..total_f32); - let host_f32 = &mut self.readback_f32_host[..total_f32]; - self.stream - .memcpy_dtoh(&f32_view, host_f32) - .map_err(|e| MLError::ModelError(format!("DtoH f32: {e}")))?; - - // Single DtoH for i32 - let i32_view = self.readback_i32_staging.slice(..total_i32); - let host_i32 = &mut self.readback_i32_host[..total_i32]; - self.stream - .memcpy_dtoh(&i32_view, host_i32) - .map_err(|e| MLError::ModelError(format!("DtoH i32: {e}")))?; - - // Slice host buffers - let states = host_f32[..states_count].to_vec(); - let log_probs = host_f32[states_count..states_count + total].to_vec(); - let advantages = host_f32[states_count + total..states_count + 2 * total].to_vec(); - let returns = host_f32[states_count + 2 * total..states_count + 3 * total].to_vec(); - let actions = host_i32[..total].to_vec(); - let done_flags = host_i32[total..2 * total].to_vec(); + // Allocate owned CudaSlices and D2D copy the active sub-range from output buffers. + // Output buffers are MAX_EPISODES * MAX_TIMESTEPS sized; we only need n * L. + let batch_states = self.d2d_copy_f32(&self.states_out, states_count, "states")?; + let batch_log_probs = self.d2d_copy_f32(&self.log_probs_out, total, "log_probs")?; + let batch_advantages = self.d2d_copy_f32(&self.advantages_out, total, "advantages")?; + let batch_returns = self.d2d_copy_f32(&self.returns_out, total, "returns")?; + let batch_actions = self.d2d_copy_i32(&self.actions_out, total, "actions")?; + let batch_done_flags = self.d2d_copy_i32(&self.dones_out, total, "done_flags")?; debug!( n_episodes, timesteps, total_experiences = total, - "PPO experience collection complete (2 DtoH vs 6)" + "PPO experience collection complete (GPU-resident, 0 DtoH)" ); - // ---- Step 7: Return batch ---- + // ---- Step 7: Return GPU-resident batch ---- Ok(PpoExperienceBatch { - states, - actions, - log_probs, - advantages, - returns, - done_flags, + states: batch_states, + actions: batch_actions, + log_probs: batch_log_probs, + advantages: batch_advantages, + returns: batch_returns, + done_flags: batch_done_flags, n_episodes, timesteps, + state_dim: STATE_DIM, + stream: self.stream.clone(), }) } + /// D2D copy `count` f32 elements from `src` into a fresh CudaSlice (GPU-to-GPU, no CPU). + fn d2d_copy_f32( + &self, + src: &CudaSlice, + count: usize, + label: &str, + ) -> Result, MLError> { + let mut dst = self + .stream + .alloc_zeros::(count) + .map_err(|e| MLError::ModelError(format!("alloc {label}: {e}")))?; + let src_view = src.slice(..count); + let nbytes = count * std::mem::size_of::(); + { + let (src_ptr, _s) = src_view.device_ptr(&self.stream); + let (dst_ptr, _d) = dst.device_ptr_mut(&self.stream); + unsafe { + cudarc::driver::result::memcpy_dtod_async( + dst_ptr, src_ptr, nbytes, self.stream.cu_stream(), + ) + .map_err(|e| MLError::ModelError(format!("D2D {label}: {e}")))?; + } + } + Ok(dst) + } + + /// D2D copy `count` i32 elements from `src` into a fresh CudaSlice (GPU-to-GPU, no CPU). + fn d2d_copy_i32( + &self, + src: &CudaSlice, + count: usize, + label: &str, + ) -> Result, MLError> { + let mut dst = self + .stream + .alloc_zeros::(count) + .map_err(|e| MLError::ModelError(format!("alloc {label}: {e}")))?; + let src_view = src.slice(..count); + let nbytes = count * std::mem::size_of::(); + { + let (src_ptr, _s) = src_view.device_ptr(&self.stream); + let (dst_ptr, _d) = dst.device_ptr_mut(&self.stream); + unsafe { + cudarc::driver::result::memcpy_dtod_async( + dst_ptr, src_ptr, nbytes, self.stream.cu_stream(), + ) + .map_err(|e| MLError::ModelError(format!("D2D {label}: {e}")))?; + } + } + Ok(dst) + } + /// Re-upload actor and critic weights from `GpuVarStore` objects (after training step). pub fn sync_weights( &mut self, diff --git a/crates/ml/src/cuda_pipeline/mod.rs b/crates/ml/src/cuda_pipeline/mod.rs index 8ffdd1b62..a11c4a2d2 100644 --- a/crates/ml/src/cuda_pipeline/mod.rs +++ b/crates/ml/src/cuda_pipeline/mod.rs @@ -310,39 +310,32 @@ impl DqnGpuData { Self::d2d_subrange(&self.features, bar_idx * self.feature_dim, self.feature_dim, stream) } - /// Get market features for a batch of bars as a [batch_size, 42] tensor slice. - pub fn batch_features(&self, start: usize, count: usize) -> Result { + /// Get market features for a batch of bars as a [batch_size * 42] `CudaSlice`. + pub fn batch_features(&self, start: usize, count: usize, stream: &Arc) -> Result, MLError> { let count = count.min(self.num_bars.saturating_sub(start)); - self.features - .narrow(0, start, count) - .map_err(|e| MLError::ModelError(format!("Batch feature slice [{start}..+{count}] failed: {e}"))) + Self::d2d_subrange(&self.features, start * self.feature_dim, count * self.feature_dim, stream) } - /// Get target prices for a single bar as a [1, 4] tensor slice. - pub fn bar_targets(&self, bar_idx: usize) -> Result { - self.targets - .narrow(0, bar_idx, 1) - .map_err(|e| MLError::ModelError(format!("Bar {bar_idx} target slice failed: {e}"))) + /// Get target prices for a single bar as a [4] `CudaSlice`. + pub fn bar_targets(&self, bar_idx: usize, stream: &Arc) -> Result, MLError> { + if bar_idx >= self.num_bars { + return Err(MLError::ModelError(format!("Bar {bar_idx} target out of range (num_bars={})", self.num_bars))); + } + Self::d2d_subrange(&self.targets, bar_idx * 4, 4, stream) } - /// Get target prices for a batch as [batch_size, 4] tensor slice. - pub fn batch_targets(&self, start: usize, count: usize) -> Result { + /// Get target prices for a batch as [batch_size * 4] `CudaSlice`. + pub fn batch_targets(&self, start: usize, count: usize, stream: &Arc) -> Result, MLError> { let count = count.min(self.num_bars.saturating_sub(start)); - self.targets - .narrow(0, start, count) - .map_err(|e| MLError::ModelError(format!("Batch target slice [{start}..+{count}] failed: {e}"))) + Self::d2d_subrange(&self.targets, start * 4, count * 4, stream) } - /// Get target values for a bar as a GPU-resident [4] tensor. + /// Get target values for a bar as a GPU-resident [4] `CudaSlice`. /// /// Tensor order: [current_close_preproc, next_close_preproc, current_close_raw, next_close_raw]. - /// Stays on GPU — callers that need CPU scalars should extract at their own boundary. - pub fn bar_target_values(&self, bar_idx: usize) -> Result { - self.bar_targets(bar_idx)? - .flatten_all() - .map_err(|e| MLError::ModelError(format!("Target flatten failed: {e}")))? - .to_dtype(todo_f32) - .map_err(|e| MLError::ModelError(format!("Target cast to F32 failed: {e}"))) + /// Stays on GPU. Callers that need CPU scalars should extract at their own boundary. + pub fn bar_target_values(&self, bar_idx: usize, stream: &Arc) -> Result, MLError> { + self.bar_targets(bar_idx, stream) } /// Build a complete state tensor by concatenating pre-uploaded market features diff --git a/crates/ml/src/hyperopt/adapters/ppo.rs b/crates/ml/src/hyperopt/adapters/ppo.rs index 84a02e3bc..c19cd3aa1 100644 --- a/crates/ml/src/hyperopt/adapters/ppo.rs +++ b/crates/ml/src/hyperopt/adapters/ppo.rs @@ -699,24 +699,28 @@ impl PPOTrainer { .map_err(|e| MLError::TrainingError(format!("PPO GPU experience collection FAILED (no CPU fallback): {e}")))?; info!("GPU PPO collected {} experiences ({} episodes × {} timesteps)", batch.n_episodes * batch.timesteps, batch.n_episodes, batch.timesteps); - Ok(gpu_ppo_batch_to_trajectory_batch(&batch)) + gpu_ppo_batch_to_trajectory_batch(&batch) } } -/// Convert a GPU PPO experience batch to a TrajectoryBatch for `PPO::update()`. +/// Convert a GPU-resident PPO experience batch to a TrajectoryBatch. /// -/// The GPU kernel produces 45-dim states (42 market features + 3 portfolio state) which -/// matches the PPO model in hyperopt (state_dim=45). +/// Downloads data from GPU for the hyperopt replay buffer and validation flow. fn gpu_ppo_batch_to_trajectory_batch( batch: &crate::cuda_pipeline::gpu_ppo_collector::PpoExperienceBatch, -) -> TrajectoryBatch { - let kernel_state_dim = 45; // GPU kernel output: 42 features + 3 portfolio state - let model_state_dim = 45; // PPO hyperopt model state_dim (matches kernel) - let total = batch.n_episodes * batch.timesteps; +) -> Result { + let kernel_state_dim = batch.state_dim; + let model_state_dim = 45; + let total = batch.total(); - // Map states from GPU kernel output to model dimensions (both 45) - let states: Vec> = batch - .states + let states_flat = batch.download_states()?; + let actions_raw = batch.download_actions()?; + let log_probs_raw = batch.download_log_probs()?; + let advantages_raw = batch.download_advantages()?; + let returns_raw = batch.download_returns()?; + let done_flags_raw = batch.download_done_flags()?; + + let states: Vec> = states_flat .chunks(kernel_state_dim) .take(total) .map(|c| c.get(..model_state_dim).unwrap_or(c).to_vec()) @@ -727,48 +731,31 @@ fn gpu_ppo_batch_to_trajectory_batch( crate::common::action::OrderType::Market, crate::common::action::Urgency::Normal, ); - let actions: Vec = batch - .actions - .iter() - .take(total) + let actions: Vec = actions_raw.iter().take(total) .map(|&a| { let idx = (a.clamp(0, 44)) as usize; FactoredAction::from_index(idx).unwrap_or(hold_action) }) .collect(); - let log_probs: Vec = batch.log_probs.iter().take(total).copied().collect(); - - // V(s) = R(s) - A(s) - 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(); - + let log_probs: Vec = log_probs_raw.into_iter().take(total).collect(); + let values: Vec = returns_raw.iter().zip(&advantages_raw).take(total) + .map(|(r, a)| r - a).collect(); + let dones: Vec = done_flags_raw.iter().take(total).map(|&d| d != 0).collect(); let rewards = vec![0.0_f32; states.len()]; - TrajectoryBatch { + Ok(TrajectoryBatch { trajectories: vec![], states, - states_flat: vec![], // GPU path — to_tensors fallback flattens on the fly + states_flat: vec![], actions, log_probs, values, rewards, dones, - advantages: batch.advantages.iter().take(total).copied().collect(), - returns: batch.returns.iter().take(total).copied().collect(), - } + advantages: advantages_raw.into_iter().take(total).collect(), + returns: returns_raw.into_iter().take(total).collect(), + }) } /// Write a log entry to the training log file diff --git a/crates/ml/src/trainers/ppo.rs b/crates/ml/src/trainers/ppo.rs index d9848ccd7..10469dcab 100644 --- a/crates/ml/src/trainers/ppo.rs +++ b/crates/ml/src/trainers/ppo.rs @@ -24,7 +24,6 @@ use crate::gpu::memory_profile; use crate::batch_size_resolver::resolve_batch_size; use crate::ppo::gae::GAEConfig; use crate::ppo::ppo::{PPOConfig, PPO}; -use crate::ppo::trajectories::TrajectoryBatch; use crate::MLError; /// PPO training hyperparameters (matches gRPC PpoParams) @@ -523,8 +522,8 @@ impl PpoTrainer { // Emit epoch gauge immediately so Grafana template variables resolve early training_metrics::set_epoch("ppo", "current", (epoch + 1) as f64); - // Phase 3: GPU experience collection - let mut training_batch: TrajectoryBatch = if let ( + // Phase 3: GPU experience collection (GPU-resident batch) + let gpu_batch: crate::cuda_pipeline::gpu_ppo_collector::PpoExperienceBatch = if let ( Some(ref mut collector), Some(ref features_buf), Some(ref targets_buf), @@ -561,18 +560,35 @@ impl PpoTrainer { .map_err(|e| MLError::TrainingError(format!("GPU PPO episode reset FAILED: {e}")))?; let batch = collector.collect_experiences(features_buf, targets_buf, &episode_starts, &config) .map_err(|e| MLError::TrainingError(format!("GPU PPO collection FAILED: {e}")))?; - info!("GPU PPO collected {} experiences ({} episodes × {} timesteps)", + info!("GPU PPO collected {} experiences ({} episodes x {} timesteps)", batch.n_episodes * batch.timesteps, batch.n_episodes, batch.timesteps); - gpu_batch_to_trajectory_batch(&batch) + batch } else if gpu_ppo_collector.is_none() { return Err(MLError::TrainingError("PPO GPU collector not initialized".to_owned())); } else { - return Err(MLError::TrainingError("PPO raw market data not uploaded — call set_raw_market_data() before train()".to_owned())); + return Err(MLError::TrainingError("PPO raw market data not uploaded -- call set_raw_market_data() before train()".to_owned())); }; + let batch_total = gpu_batch.total(); + let batch_state_dim = gpu_batch.state_dim; + let batch_stream = &gpu_batch.stream; + // Step 2.5: Pre-train value network (first 10 epochs only) if epoch < 10 { - let pretrain_loss = self.pretrain_value_network(&mut training_batch, 5).await?; + let pretrain_loss = { + let mut model = self.model.lock().await; + let mut total_loss = 0.0_f32; + for _ in 0..5 { + let (_pl, vl) = model.update_gpu( + &gpu_batch.states, &gpu_batch.actions, + &gpu_batch.log_probs, &gpu_batch.advantages, + &gpu_batch.returns, batch_total, + batch_state_dim, batch_stream, + )?; + total_loss += vl; + } + total_loss / 5.0 + }; info!( "Epoch {} - Value pre-training loss: {:.4}", epoch + 1, @@ -580,10 +596,15 @@ impl PpoTrainer { ); } - // Step 3: PPO update + // Step 3: PPO update (GPU-resident -- zero CPU roundtrip for states) let (policy_loss, value_loss) = { let mut model = self.model.lock().await; - model.update(&mut training_batch)? + model.update_gpu( + &gpu_batch.states, &gpu_batch.actions, + &gpu_batch.log_probs, &gpu_batch.advantages, + &gpu_batch.returns, batch_total, + batch_state_dim, batch_stream, + )? }; // Phase 2c: Sync actor + critic weights to GPU after PPO update @@ -593,9 +614,9 @@ impl PpoTrainer { .map_err(|e| MLError::TrainingError(format!("PPO GPU weight sync FAILED -- stale weights corrupt training: {e}")))?; } - // Step 4: Compute additional metrics + // Step 4: Compute additional metrics (downloads only scalars, not 123 MB states) let (kl_divergence, explained_variance, mean_reward, std_reward, entropy) = - self.compute_metrics(&training_batch, policy_loss, value_loss)?; + self.compute_metrics_from_gpu(&gpu_batch, policy_loss, value_loss)?; // Step 5: Report progress let epoch_metrics = PpoTrainingMetrics { @@ -804,81 +825,60 @@ impl PpoTrainer { advantages } - /// Pre-train value network to better fit returns before policy updates. + // pretrain_value_network REMOVED -- pre-training now inlined in train_gpu() + // using update_gpu() on GPU-resident PpoExperienceBatch directly. + + /// Compute metrics from a GPU-resident PpoExperienceBatch. /// - /// This method trains ONLY the critic (value network) without touching the - /// policy (actor network). This is critical for PPO correctness: the policy - /// must only be updated with on-policy data, and pretraining the actor with - /// stale log-probs would violate this constraint. - async fn pretrain_value_network( + /// Downloads only the data needed for scalar metric computation (returns, + /// advantages, actions). This is the only DtoH transfer in the training loop + /// and transfers ~3 * N * L * 4 bytes of scalar metrics, NOT the full 123 MB + /// state tensor. + fn compute_metrics_from_gpu( &self, - batch: &mut TrajectoryBatch, - pretraining_epochs: usize, - ) -> Result { - let mut total_loss = 0.0; - let mut num_updates = 0; - - for _ in 0..pretraining_epochs { - let value_loss = { - let mut model = self.model.lock().await; - model.update_value_only(batch)? - }; - - total_loss += value_loss; - num_updates += 1; - } - - Ok(if num_updates > 0 { - total_loss / num_updates as f32 - } else { - 0.0 - }) - } - - /// Compute additional metrics (KL divergence, explained variance, etc.) - /// Uses O(n) CPU arithmetic on host-resident data — zero GPU round-trips. - fn compute_metrics( - &self, - batch: &TrajectoryBatch, + batch: &crate::cuda_pipeline::gpu_ppo_collector::PpoExperienceBatch, policy_loss: f32, _value_loss: f32, ) -> Result<(f32, f32, f32, f32, f32), MLError> { - // KL divergence (approximated from policy loss) let kl_div = policy_loss.abs() * 0.1; - let returns = &batch.returns; - let values = &batch.values; - let rewards = &batch.rewards; + // Download only returns, advantages, actions for metric computation. + // States (the 123 MB bulk) are NOT downloaded. + let returns = batch.download_returns()?; + let advantages = batch.download_advantages()?; + let actions = batch.download_actions()?; + + // V(s) = R(s) - A(s) + let values: Vec = returns.iter().zip(&advantages).map(|(r, a)| r - a).collect(); + + // Rewards not directly available from kernel (GAE already computed). + let rewards = vec![0.0_f32; returns.len()]; - // CPU-resident O(n) computation — data is already on host (downloaded from - // GPU kernel via PpoExperienceBatch). No GPU upload/download round-trip. let (explained_variance, mean_reward, std_reward) = - Self::compute_metrics_cpu(returns, values, rewards); + Self::compute_metrics_cpu(&returns, &values, &rewards); // Compute real entropy from the epoch action distribution - let entropy = { - let actions = &batch.actions; - if actions.is_empty() { - 0.0_f32 - } else { - let mut counts = [0_u32; 45]; - for action in actions { - let idx = action.to_index(); - if idx < 45 { - counts[idx] += 1; + let entropy = if actions.is_empty() { + 0.0_f32 + } else { + let mut counts = [0_u32; 45]; + for &a in &actions { + let idx = a.clamp(0, 44) as usize; + if idx < 45 { + if let Some(c) = counts.get_mut(idx) { + *c += 1; } } - let n = actions.len() as f64; - let mut h = 0.0_f64; - for &c in &counts { - if c > 0 { - let p = c as f64 / n; - h -= p * p.ln(); - } - } - // Normalize by ln(45) so entropy is in [0, 1] - (h / 45.0_f64.ln()) as f32 } + let n = actions.len() as f64; + let mut h = 0.0_f64; + for &c in &counts { + if c > 0 { + let p = c as f64 / n; + h -= p * p.ln(); + } + } + (h / 45.0_f64.ln()) as f32 }; Ok((kl_div, explained_variance, mean_reward, std_reward, entropy)) @@ -1068,73 +1068,9 @@ 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. -fn gpu_batch_to_trajectory_batch( - batch: &crate::cuda_pipeline::gpu_ppo_collector::PpoExperienceBatch, -) -> TrajectoryBatch { - let state_dim = 45; - let total = batch.n_episodes * batch.timesteps; - - let states: Vec> = batch - .states - .chunks(state_dim) - .take(total) - .map(|c| c.to_vec()) - .collect(); - - let hold_action = FactoredAction::new( - crate::common::action::ExposureLevel::Flat, - crate::common::action::OrderType::Market, - crate::common::action::Urgency::Normal, - ); - let actions: Vec = batch - .actions - .iter() - .take(total) - .map(|&a| { - let idx = (a.clamp(0, 44)) as usize; - FactoredAction::from_index(idx).unwrap_or(hold_action) - }) - .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, - states_flat: vec![], // GPU path — to_tensors fallback flattens on the fly - actions, - log_probs, - values, - rewards, - dones, - advantages: batch.advantages.iter().take(total).copied().collect(), - returns: batch.returns.iter().take(total).copied().collect(), - } -} +// gpu_batch_to_trajectory_batch DELETED -- PPO experience batches are now GPU-resident. +// Training uses PPO::update_gpu() directly, eliminating the 123 MB/epoch +// GPU->CPU->GPU roundtrip that this function caused. #[cfg(test)] mod tests {