diff --git a/crates/ml/src/trainers/dqn/trainer/constructor.rs b/crates/ml/src/trainers/dqn/trainer/constructor.rs index e22681523..ed2fb985f 100644 --- a/crates/ml/src/trainers/dqn/trainer/constructor.rs +++ b/crates/ml/src/trainers/dqn/trainer/constructor.rs @@ -752,9 +752,7 @@ impl DQNTrainer { // OFI features: populated during data loading when MBP-10 data is available ofi_features: None, ofi_val_offset: 0, - val_features_gpu: None, - val_closes_gpu: None, - val_ofi_gpu: None, + gpu_evaluator: None, // Phase C: Fill simulation and smart order routing fill_simulator: FillSimulator::default(), diff --git a/crates/ml/src/trainers/dqn/trainer/metrics.rs b/crates/ml/src/trainers/dqn/trainer/metrics.rs index d3aaa83fc..a12091d6c 100644 --- a/crates/ml/src/trainers/dqn/trainer/metrics.rs +++ b/crates/ml/src/trainers/dqn/trainer/metrics.rs @@ -387,203 +387,152 @@ impl DQNTrainer { None } - /// Compute validation loss (negative Sharpe proxy) using cuBLAS forward — zero GpuTensor. + /// Compute validation loss (negative Sharpe) using the GPU backtest evaluator. /// - /// Builds the state tensor on CPU (features + portfolio broadcast + OFI + alignment - /// padding), uploads it once as a raw `CudaSlice`, runs cuBLAS forward via - /// `fused_ctx`, then downloads Q-values and performs all action selection and Sharpe - /// computation on CPU. This avoids GpuTensor elementwise kernels (add, broadcast_mul, - /// cat, etc.) that cache `CudaFunction` handles on the shared CUDA primary context, - /// which causes sequential test failures. + /// Runs the full backtest step loop on GPU: gather kernel -> cuBLAS forward -> + /// env kernel -> metrics reduction. Only the final metrics (1 window x 14 floats) + /// are downloaded to CPU. Replaces the old CPU path (GPU->CPU Q-value download, + /// CPU argmax, CPU Sharpe) which took ~2s/epoch vs ~50ms here. /// - /// Epoch-boundary cold path — simplicity over GPU utilisation. + /// The evaluator is lazy-initialized on the first call and cached across epochs. + /// Weight pointers change each epoch (training updates weights), so we call + /// `invalidate_dqn_graph()` to force CUDA Graph re-capture. pub(crate) async fn compute_validation_loss(&mut self) -> Result { + use crate::cuda_pipeline::gpu_backtest_evaluator::{ + DqnBacktestConfig, GpuBacktestConfig, GpuBacktestEvaluator, + }; + if self.val_data.is_empty() { return Ok(0.0); } - // Cap sample_size at fused_ctx batch_size to avoid exceeding cuBLAS buffer - let fused_batch = self.fused_ctx.as_ref().map(|f| f.batch_size()).unwrap_or(0); - if fused_batch == 0 { - return Ok(0.0); // fused_ctx not initialized — skip validation - } - - // Save current epsilon and force to 0 for deterministic evaluation - let original_epsilon = self.get_epsilon().await?; - self.set_epsilon(0.0).await?; // Pure greedy selection - - let sample_size = self.val_data.len().min(1000).min(fused_batch); - - let aligned_state_dim = { - let agent = self.agent.read().await; - agent.get_state_dim() + let stream = match self.cuda_stream.as_ref() { + Some(s) => s, + None => return Ok(0.0), // CPU device — skip GPU validation }; - let stream = self.cuda_stream.as_ref() - .ok_or_else(|| anyhow::anyhow!("cuda_stream required for validation"))?; + // ── Lazy-init the GPU evaluator (once per fold, reused across epochs) ── + if self.gpu_evaluator.is_none() { + let ofi_enabled = self.ofi_features.is_some(); + let market_dim: usize = 42; + let feature_dim: usize = if ofi_enabled { 50 } else { market_dim }; - // ── CPU-cached validation data (lazy init, built once per training run) ── - // val_features_gpu stores flat features as Vec [sample_size * 42] on CPU. - // val_closes_gpu stores (current_closes, next_closes) as Vec on CPU. - // val_ofi_gpu stores flat OFI features as Vec [sample_size * 8] on CPU. - if self.val_features_gpu.is_none() { - let mut flat_features: Vec = Vec::with_capacity(sample_size * 42); - let mut current_closes: Vec = Vec::with_capacity(sample_size); - let mut next_closes: Vec = Vec::with_capacity(sample_size); + // Build a single window from all val_data + let mut prices: Vec<[f32; 4]> = Vec::with_capacity(self.val_data.len()); + let mut features: Vec> = Vec::with_capacity(self.val_data.len()); - for (feature_vec, target) in self.val_data.iter().take(sample_size) { - for &v in feature_vec.iter() { - flat_features.push(v as f32); - } - let cur = if target.len() >= 2 { target[0] } else { feature_vec[3] }; - let nxt = if target.len() >= 2 { target[1] } else { cur }; - current_closes.push(cur as f32); - next_closes.push(nxt as f32); - } + for (i, (fv, target)) in self.val_data.iter().enumerate() { + // Prices: [close, close, close, close] — same convention as hyperopt adapter + let close = if target.len() >= 2 { target[0] as f32 } else { fv[3] as f32 }; + prices.push([close, close, close, close]); - self.val_features_gpu = Some(flat_features); - self.val_closes_gpu = Some((current_closes, next_closes)); - - if let Some(ref ofi) = self.ofi_features { - let mut flat_ofi: Vec = Vec::with_capacity(sample_size * 8); - for i in 0..sample_size { - let idx = self.ofi_val_offset + i; - if let Some(row) = ofi.get(idx) { - for &v in row.iter() { - flat_ofi.push(v as f32); + // Features: 42 market features, plus 8 OFI if enabled + let fv_slice = &fv[..market_dim.min(fv.len())]; + let mut fv_f32: Vec = fv_slice.iter().map(|&v| v as f32).collect(); + if ofi_enabled { + let ofi_idx = self.ofi_val_offset + i; + if let Some(ofi_row) = self.ofi_features.as_ref().and_then(|o| o.get(ofi_idx)) { + for &v in ofi_row.iter() { + fv_f32.push(v as f32); } } else { - flat_ofi.extend_from_slice(&[0.0_f32; 8]); + fv_f32.resize(feature_dim, 0.0); } } - self.val_ofi_gpu = Some(flat_ofi); - } - } - - let flat_features = self.val_features_gpu.as_ref() - .ok_or_else(|| anyhow::anyhow!("val_features_gpu must be initialized"))?; - let (current_closes, next_closes) = self.val_closes_gpu.as_ref() - .ok_or_else(|| anyhow::anyhow!("val_closes_gpu must be initialized"))?; - - // ── Build complete state tensor on CPU per validation call ── - // Portfolio features change every epoch (position, P&L), so we rebuild each call. - // Feature layout per row: [42 market features | 3 portfolio | 8 OFI (opt) | padding] - let val_price = self.val_data.first() - .map(|(fv, tgt)| if tgt.len() >= 2 { tgt[0] as f32 } else { fv[3] as f32 }) - .unwrap_or(0.0); - let portfolio_f = self.portfolio_tracker.get_portfolio_features(val_price); - - // State layout: [42 market | 8 portfolio | 16 MTF | 8 OFI | padding] - // We only have 42 features + 3 portfolio + 8 OFI = 53 values. - // Pad the rest with zeros to match state_dim_padded (pad128) per row. - // cuBLAS GemmEx reads with ldb=state_dim_padded (CUTLASS K-tile=128). - let state_dim_padded = (aligned_state_dim + 127) & !127; - let ofi_cols: usize = if self.val_ofi_gpu.is_some() { 8 } else { 0 }; - let content_cols: usize = 42 + 3 + ofi_cols; // what we actually write - let pad_cols = state_dim_padded.saturating_sub(content_cols); - - let mut state_flat: Vec = Vec::with_capacity(sample_size * state_dim_padded); - let feat_cols = flat_features.len() / sample_size; // 42 - for i in 0..sample_size { - // Market features - let feat_start = i * feat_cols; - let feat_end = feat_start + feat_cols; - let feat_slice = flat_features.get(feat_start..feat_end) - .ok_or_else(|| anyhow::anyhow!("val features OOB at row {i}"))?; - state_flat.extend_from_slice(feat_slice); - - // Portfolio features (same for all rows — broadcast) - state_flat.extend_from_slice(&portfolio_f[..3]); - - // OFI features (optional, 8 floats per row) - if let Some(ref flat_ofi) = self.val_ofi_gpu { - let ofi_start = i * 8; - let ofi_end = ofi_start + 8; - let ofi_slice = flat_ofi.get(ofi_start..ofi_end) - .ok_or_else(|| anyhow::anyhow!("val OFI OOB at row {i}"))?; - state_flat.extend_from_slice(ofi_slice); + features.push(fv_f32); } - // Alignment padding (zeros) - state_flat.extend(std::iter::repeat(0.0_f32).take(pad_cols)); - } + let window_prices = vec![prices]; + let window_features = vec![features]; - // Upload state to GPU as raw CudaSlice — no GpuTensor, no elementwise kernels. - let state_gpu = crate::cuda_pipeline::clone_htod_f32_to_bf16(stream, &state_flat) - .map_err(|e| anyhow::anyhow!("val state HtoD: {e}"))?; - - // cuBLAS forward via fused_ctx — zero GpuTensor involvement - let fused = self.fused_ctx.as_mut() - .ok_or_else(|| anyhow::anyhow!("fused_ctx required for validation forward"))?; - let total_actions = fused.total_actions(); - let q_out = fused.compute_q_values(&state_gpu, sample_size) - .map_err(|e| anyhow::anyhow!("cuBLAS validation forward: {e}"))?; - - // Download Q-values to CPU for action selection + Sharpe computation. - // q_out is [config.batch_size, total_actions]; allocate to match, then truncate. - let mut host_q = vec![0.0_f32; q_out.len()]; - crate::cuda_pipeline::dtoh_bf16_to_f32(stream, q_out, &mut host_q) - .map_err(|e| anyhow::anyhow!("Q-val DtoH: {e}"))?; - host_q.truncate(sample_size * total_actions); - - // ── CPU greedy action selection ── - // Direction LUT maps exposure head index → position size direction. - // 0=Short100(-1.0), 1=Short50(-0.5), 2=Flat(0.0), 3=Long50(0.5), 4=Long100(1.0) - const DIRECTION_LUT: [f32; 9] = [-1.0, -0.75, -0.5, -0.25, 0.0, 0.25, 0.5, 0.75, 1.0]; - let use_branching = true && total_actions >= 15; - - let rewards: Vec = (0..sample_size).map(|i| { - let row_start = i * total_actions; - let row = &host_q[row_start..row_start + total_actions]; - - // Greedy exposure action selection (argmax over exposure head, 5 values) - let exposure_idx = if use_branching { - // Branching DQN: exposure head is first 9 Q-values - row.get(..9).unwrap_or(row) - .iter() - .enumerate() - .max_by(|a, b| a.1.partial_cmp(b.1).unwrap_or(std::cmp::Ordering::Equal)) - .map(|(j, _)| j) - .expect("argmax must exist for non-empty Q-values") - } else { - // Non-branching: argmax over all Q-values maps to exposure - row.iter() - .enumerate() - .max_by(|a, b| a.1.partial_cmp(b.1).unwrap_or(std::cmp::Ordering::Equal)) - .map(|(j, _)| j.min(4)) - .expect("argmax must exist for non-empty Q-values") + let hp = &self.hyperparams; + #[allow(clippy::cast_possible_truncation)] + let config = GpuBacktestConfig { + max_position: hp.max_position_absolute as f32, + tx_cost_bps: hp.transaction_cost_multiplier as f32, + spread_cost: (hp.tick_size * hp.contract_multiplier * hp.fill_spread_cost_frac) as f32, + initial_capital: hp.initial_capital as f32, + contract_multiplier: hp.contract_multiplier as f32, + margin_pct: hp.margin_pct as f32, + max_leverage: 0.0, // Disabled: match training env (no leverage cap) + ofi_dim: if ofi_enabled { 8 } else { 0 }, + min_hold_bars: hp.min_hold_bars as i32, + bars_per_day: hp.bars_per_day as f32, + trading_days_per_year: hp.trading_days_per_year as f32, }; - let direction = DIRECTION_LUT[exposure_idx.min(8)]; - // PnL-based reward: direction × (next_close - cur_close) / cur_close - let cur = current_closes.get(i).copied().unwrap_or(0.0); - let nxt = next_closes.get(i).copied().unwrap_or(cur); - if cur.abs() > 1e-10 { - direction * (nxt - cur) / cur - } else { - 0.0 - } - }).collect(); + let evaluator = GpuBacktestEvaluator::new( + &window_prices, + &window_features, + feature_dim, + config, + stream, + ).map_err(|e| anyhow::anyhow!("GpuBacktestEvaluator init: {e}"))?; - // ── CPU Sharpe computation ── - let n = rewards.len() as f64; - let mean_scalar: f64 = rewards.iter().map(|&v| v as f64).sum::() / n; - let var_scalar: f64 = if n > 1.0 { - rewards.iter().map(|&v| { - let d = v as f64 - mean_scalar; - d * d - }).sum::() / n + self.gpu_evaluator = Some(evaluator); + } + + // ── Extract current weights from agent ── + let agent = self.agent.read().await; + let is_branching = agent.is_using_branching(); + let network_dims = agent.network_dims(); + let vars = agent.get_q_network_vars(); + + let evaluator = self.gpu_evaluator.as_mut() + .ok_or_else(|| anyhow::anyhow!("gpu_evaluator must be initialized"))?; + let eval_stream = evaluator.stream(); + + let online_weights = if is_branching { + crate::cuda_pipeline::gpu_weights::extract_dueling_weights_branching( + vars, eval_stream, + ).map_err(|e| anyhow::anyhow!("extract dueling weights (branching): {e}"))? } else { - 0.0 - }; - let std_val = var_scalar.sqrt(); - let val_sharpe = if std_val > 1e-10 { - (mean_scalar / std_val) * n.sqrt() - } else { - 0.0 + crate::cuda_pipeline::gpu_weights::extract_dueling_weights( + vars, eval_stream, + ).map_err(|e| anyhow::anyhow!("extract dueling weights: {e}"))? }; - self.set_epsilon(original_epsilon).await?; + let branching_weights = if is_branching { + Some(crate::cuda_pipeline::gpu_weights::extract_branching_weights( + vars, eval_stream, + ).map_err(|e| anyhow::anyhow!("extract branching weights: {e}"))?) + } else { + None + }; + + let hp = &self.hyperparams; + let (b0, b1, b2) = agent.branch_sizes(); + #[allow(clippy::cast_possible_truncation)] + let dqn_cfg = DqnBacktestConfig { + shared_h1: network_dims.0, + shared_h2: network_dims.1, + value_h: network_dims.2, + adv_h: network_dims.3, + num_atoms: hp.num_atoms, + branch_0_size: b0, + branch_1_size: if is_branching { b1 } else { b0 }, + branch_2_size: if is_branching { b2 } else { b0 }, + v_min: hp.v_min as f32, + v_max: hp.v_max as f32, + }; + + // Drop agent read guard before mutable evaluator call + drop(agent); + + // Weights changed since last epoch — invalidate the cached CUDA graph + // so evaluate_dqn_graphed re-captures with current weight pointers. + evaluator.invalidate_dqn_graph(); + + let metrics = evaluator.evaluate_dqn_graphed( + &online_weights, + branching_weights.as_ref(), + &dqn_cfg, + ).map_err(|e| anyhow::anyhow!("GPU backtest evaluate_dqn_graphed: {e}"))?; + + // Single window — take its Sharpe directly + let val_sharpe = metrics.first() + .map(|m| m.sharpe as f64) + .unwrap_or(0.0); Ok(-val_sharpe) } diff --git a/crates/ml/src/trainers/dqn/trainer/mod.rs b/crates/ml/src/trainers/dqn/trainer/mod.rs index 61d6eed8d..540250d23 100644 --- a/crates/ml/src/trainers/dqn/trainer/mod.rs +++ b/crates/ml/src/trainers/dqn/trainer/mod.rs @@ -275,12 +275,10 @@ pub struct DQNTrainer { /// val_data[i] corresponds to ofi_features[ofi_val_offset + i]. pub(crate) ofi_val_offset: usize, - /// CPU-cached validation features flat vec [sample_size * 42] — built once, reused per epoch - pub(crate) val_features_gpu: Option>, - /// CPU-cached validation close prices (current, next) — built once per training run - pub(crate) val_closes_gpu: Option<(Vec, Vec)>, - /// CPU-cached validation OFI features flat vec [sample_size * 8] — built once - pub(crate) val_ofi_gpu: Option>, + /// GPU backtest evaluator for validation loss (replaces CPU argmax+Sharpe path). + /// Lazy-initialized on first `compute_validation_loss()` call from val_data. + /// Invalidated (set to None) when val_data changes (fold transitions). + pub(crate) gpu_evaluator: Option, // Phase C: Fill simulation and smart order routing /// Fill simulator for order type-dependent execution modeling @@ -487,9 +485,7 @@ impl DQNTrainer { self.init_from_fxcache(&features, &targets, &[]).await?; self.set_val_data_from_slices(&val_features, &val_targets, features.len()); - self.val_features_gpu = None; - self.val_closes_gpu = None; - self.val_ofi_gpu = None; + self.gpu_evaluator = None; self.train_fold_from_slices(&features, &targets, checkpoint_callback).await } @@ -522,9 +518,7 @@ impl DQNTrainer { .map(|(f, t)| (*f, *t)) .collect(); - self.val_features_gpu = None; - self.val_closes_gpu = None; - self.val_ofi_gpu = None; + self.gpu_evaluator = None; self.train_with_data_full_loop_slices(&training_data, checkpoint_callback) .await @@ -1065,9 +1059,7 @@ impl DQNTrainer { .map(|(f, t)| (*f, t.to_vec())) .collect(); self.ofi_val_offset = ofi_val_offset; - self.val_features_gpu = None; - self.val_closes_gpu = None; - self.val_ofi_gpu = None; + self.gpu_evaluator = None; } /// Inject pre-uploaded GPU data (e.g. from a `DoubleBufferedLoader`).