From 483a07b5dd821f00a84105071c10eeedb45061e2 Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Thu, 2 Apr 2026 22:34:02 +0200 Subject: [PATCH 1/6] =?UTF-8?q?feat:=20DqnGpuData::upload=5Fslices=20?= =?UTF-8?q?=E2=80=94=20contiguous=20array=20upload=20from=20fxcache?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Co-Authored-By: Claude Sonnet 4.6 --- crates/ml/src/cuda_pipeline/mod.rs | 67 ++++++++++++++++++++++++++++++ 1 file changed, 67 insertions(+) diff --git a/crates/ml/src/cuda_pipeline/mod.rs b/crates/ml/src/cuda_pipeline/mod.rs index a97c2d87e..3f38271ca 100644 --- a/crates/ml/src/cuda_pipeline/mod.rs +++ b/crates/ml/src/cuda_pipeline/mod.rs @@ -298,6 +298,73 @@ impl DqnGpuData { }) } + /// Upload DQN training data to GPU from separate contiguous arrays. + /// + /// Accepts features, targets, and OFI as separate slices matching the fxcache + /// flat binary layout — no tuple allocation or per-bar heap Vec required. + /// OFI is uploaded only when at least one value is non-zero. + pub fn upload_slices( + features: &[[f64; 42]], + targets: &[[f64; 4]], + ofi: &[[f64; 8]], + stream: &Arc, + ) -> Result { + let num_bars = features.len(); + if num_bars == 0 { + return Err(MLError::ModelError("Empty training data".to_owned())); + } + + let feature_dim = 42; + let target_dim = 4; + + let estimated_bytes = estimate_vram_bytes(num_bars * (feature_dim + target_dim)); + if estimated_bytes > MAX_UPLOAD_BYTES { + return Err(MLError::ModelError(format!( + "Training data too large for GPU upload: {:.1} GB > 2.0 GB limit ({} bars)", + estimated_bytes as f64 / 1_073_741_824.0, + num_bars, + ))); + } + + let flat_features: Vec = features + .iter() + .flat_map(|f| f.iter().map(|&v| v as f32)) + .collect(); + + let flat_targets: Vec = targets + .iter() + .flat_map(|t| t.iter().map(|&v| v as f32)) + .collect(); + + let features_gpu = clone_htod_f32_to_bf16(stream, &flat_features)?; + let targets_gpu = clone_htod_f32_to_bf16(stream, &flat_targets)?; + + let ofi_features = if !ofi.is_empty() && ofi.iter().any(|row| row.iter().any(|&v| v != 0.0)) { + let mut flat_ofi = Vec::with_capacity(num_bars * 8); + for i in 0..num_bars { + if i < ofi.len() { + for &v in ofi[i].iter() { + flat_ofi.push(v as f32); + } + } else { + flat_ofi.extend_from_slice(&[0.0_f32; 8]); + } + } + Some(clone_htod_f32_to_bf16(stream, &flat_ofi)?) + } else { + None + }; + + Ok(Self { + features: features_gpu, + targets: targets_gpu, + ofi_features, + num_bars, + feature_dim, + aligned_state_dim: None, + }) + } + /// Upload OFI features (8 dims per bar) from MBP-10 order book data. /// /// Must be called after `upload()` with a slice matching `num_bars` length. From 74c84b3c48e05014f8b5c014c325f9cc6cd0f732 Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Thu, 2 Apr 2026 22:36:53 +0200 Subject: [PATCH 2/6] =?UTF-8?q?feat:=20init=5Ffrom=5Ffxcache=20+=20set=5Fv?= =?UTF-8?q?al=5Fdata=5Ffrom=5Fslices=20=E2=80=94=20one-time=20GPU=20upload?= =?UTF-8?q?=20API?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Co-Authored-By: Claude Sonnet 4.6 --- crates/ml/src/trainers/dqn/trainer/mod.rs | 52 +++++++++++++++++++ .../src/trainers/dqn/trainer/training_loop.rs | 36 +++++++++++++ 2 files changed, 88 insertions(+) diff --git a/crates/ml/src/trainers/dqn/trainer/mod.rs b/crates/ml/src/trainers/dqn/trainer/mod.rs index b43727071..485113e27 100644 --- a/crates/ml/src/trainers/dqn/trainer/mod.rs +++ b/crates/ml/src/trainers/dqn/trainer/mod.rs @@ -1030,6 +1030,58 @@ impl DQNTrainer { Ok(data) } + /// Upload full fxcache dataset to GPU ONCE. Call before the fold loop. + /// Replaces per-fold init_gpu_data + init_gpu_raw_buffers. + pub async fn init_from_fxcache( + &mut self, + features: &[[f64; 42]], + targets: &[[f64; 4]], + ofi: &[[f64; 8]], + ) -> Result<()> { + let stream = self.cuda_stream.as_ref() + .ok_or_else(|| anyhow::anyhow!("CUDA stream required for init_from_fxcache"))?; + + // Upload via DqnGpuData::upload_slices (contiguous, zero tuple unpacking) + let mut gpu_data = crate::cuda_pipeline::DqnGpuData::upload_slices( + features, targets, ofi, stream, + ).map_err(|e| anyhow::anyhow!("DqnGpuData::upload_slices failed: {e}"))?; + + let ofi_enabled = gpu_data.ofi_features.is_some(); + let raw_dim = if ofi_enabled { 53 } else { 45 }; + let aligned_dim = (raw_dim + 7) & !7; + gpu_data.set_aligned_state_dim(aligned_dim); + + tracing::info!( + "init_from_fxcache: {} bars uploaded to GPU ({:.1} MB, OFI={})", + features.len(), + (features.len() * (42 + 4 + 8) * 4) as f64 / 1_048_576.0, + ofi_enabled, + ); + self.gpu_data = Some(gpu_data); + + // Also upload raw buffers for experience collector + self.init_gpu_raw_buffers_from_slices(features, targets).await?; + + Ok(()) + } + + /// Set validation data from contiguous fxcache slices. + /// Converts to legacy Vec format internally — val set is small (~50K bars). + pub fn set_val_data_from_slices( + &mut self, + features: &[[f64; 42]], + targets: &[[f64; 4]], + ofi_val_offset: usize, + ) { + self.val_data = features.iter().zip(targets.iter()) + .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; + } + /// Inject pre-uploaded GPU data (e.g. from a `DoubleBufferedLoader`). /// /// The trainer's `train_epoch` lazily uploads data on first call. diff --git a/crates/ml/src/trainers/dqn/trainer/training_loop.rs b/crates/ml/src/trainers/dqn/trainer/training_loop.rs index ea0add2db..dc1d57f29 100644 --- a/crates/ml/src/trainers/dqn/trainer/training_loop.rs +++ b/crates/ml/src/trainers/dqn/trainer/training_loop.rs @@ -718,6 +718,42 @@ impl DQNTrainer { Ok(()) } + /// Upload raw features + targets from contiguous fxcache slices. + /// No tuple unpacking — flat_map directly from &[[f64; N]] arrays. + pub(crate) async fn init_gpu_raw_buffers_from_slices( + &mut self, + features: &[[f64; 42]], + targets: &[[f64; 4]], + ) -> Result<()> { + if self.targets_raw_cuda.is_some() { + return Ok(()); + } + let stream = match self.cuda_stream { + Some(ref s) => std::sync::Arc::clone(s), + None => return Ok(()), + }; + + let num_bars = features.len(); + let flat_targets: Vec = targets.iter() + .flat_map(|t| t.iter().map(|&v| v as f32)) + .collect(); + let flat_features: Vec = features.iter() + .flat_map(|f| f.iter().map(|&v| v as f32)) + .collect(); + + self.targets_raw_cuda = Some( + crate::cuda_pipeline::clone_htod_f32_to_bf16(&stream, &flat_targets) + .map_err(|e| anyhow::anyhow!("targets_raw upload: {e}"))? + ); + self.features_raw_cuda = Some( + crate::cuda_pipeline::clone_htod_f32_to_bf16(&stream, &flat_features) + .map_err(|e| anyhow::anyhow!("features_raw upload: {e}"))? + ); + + tracing::info!("CUDA raw buffers uploaded from slices: {} bars x 42+4", num_bars); + Ok(()) + } + // ═══════════════════════════════════════════════════════════════════════ // Helper: Phase 1c — GPU experience collector init (once) // ═══════════════════════════════════════════════════════════════════════ From 5c985df1477ead2a7606ae83913f44bab6592299 Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Thu, 2 Apr 2026 22:46:30 +0200 Subject: [PATCH 3/6] =?UTF-8?q?perf:=20zero-copy=20fold=20loop=20=E2=80=94?= =?UTF-8?q?=20fxcache=20to=20GPU=20once,=20index=20slices=20per=20fold?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Rewrites the train_baseline_rl fold loop to eliminate per-fold waste: - Data loading: fxcache -> fold index ranges from timestamps (no bar reconstruction). DBN fallback builds FxCacheData in-place. - DQN trainer created ONCE before fold loop, fxcache uploaded to GPU ONCE via init_from_fxcache. Each fold uses set_training_range + set_val_data_from_slices + reset_for_fold instead of recreating. - Tokio runtime created ONCE (not per fold). - Hyperparams construction extracted to build_dqn_hyperparams(). - Deleted: prepare_fold_data, FoldData type, prefetch thread, DoubleBufferedLoader GPU staging, features_to_trainer_format (old). - Added: generate_walk_forward_indices_from_timestamps (i64 ns timestamps, O(log n) partition_point, no OHLCVBar dependency). - Added: features_to_trainer_format_fast (fxcache targets directly). - PPO compatibility preserved: constructs minimal OHLCVBars from fxcache targets for train_ppo_fold (Task 6 will refactor). - Ensemble mode preserved with per-member trainers for k>0. Co-Authored-By: Claude Opus 4.6 (1M context) --- crates/ml/examples/train_baseline_rl.rs | 589 ++++++++++-------------- crates/ml/src/walk_forward.rs | 88 ++++ 2 files changed, 326 insertions(+), 351 deletions(-) diff --git a/crates/ml/examples/train_baseline_rl.rs b/crates/ml/examples/train_baseline_rl.rs index abd6358b3..2b914818d 100644 --- a/crates/ml/examples/train_baseline_rl.rs +++ b/crates/ml/examples/train_baseline_rl.rs @@ -102,8 +102,6 @@ use clap::Parser; use serde_json::Value; use tracing::{error, info, warn}; -use ml::cuda_pipeline::DqnGpuData; -use std::sync::Arc; use ml::trainers::dqn::{DQNHyperparameters, DQNTrainer}; use ml::trainers::ppo::{PpoHyperparameters, PpoTrainer}; @@ -115,7 +113,10 @@ use common::metrics::{server as metrics_server, training_metrics as metrics}; use ml::features::extraction::{extract_ml_features, FeatureVector}; use ml_core::gpu::profile::GpuProfile; use ml::types::OHLCVBar; -use ml::walk_forward::{generate_walk_forward_windows, NormStats, WalkForwardConfig}; +use ml::walk_forward::{ + generate_walk_forward_indices_from_timestamps, + NormStats, WalkForwardConfig, +}; // --------------------------------------------------------------------------- // CLI Arguments @@ -344,153 +345,38 @@ fn hp_bool(params: &Option, key: &str) -> Option { } // --------------------------------------------------------------------------- -// Fold data preparation (extracted for prefetching) +// Fold data helpers // --------------------------------------------------------------------------- -/// Prepared fold data: (`train_features`, `val_features`, `train_bars`, `val_bars`) -type FoldData = (Vec, Vec, Vec, Vec); - -#[allow(clippy::cognitive_complexity)] -/// Prepare a fold's data for training: extract features, normalize, align bars. -/// -/// Returns `None` if feature extraction fails or produces empty features. -fn prepare_fold_data( - window: &ml::walk_forward::WalkForwardWindow, - output_dir: &Path, -) -> Option { - let train_feat = match extract_ml_features(&window.train) { - Ok(f) => f, - Err(e) => { - warn!(" Fold {} -- train feature extraction failed: {}", window.fold, e); - return None; - } - }; - let val_feat = match extract_ml_features(&window.val) { - Ok(f) => f, - Err(e) => { - warn!(" Fold {} -- val feature extraction failed: {}", window.fold, e); - return None; - } - }; - if train_feat.is_empty() || val_feat.is_empty() { - warn!(" Fold {} -- empty features, skipping", window.fold); - return None; - } - - let norm_stats = NormStats::from_features(&train_feat); - let train_norm = norm_stats.normalize_batch(&train_feat); - let val_norm = norm_stats.normalize_batch(&val_feat); - - // Save NormStats for this fold - let norm_path = output_dir.join(format!("norm_stats_fold{}.json", window.fold)); - match serde_json::to_string_pretty(&norm_stats) { - Ok(json) => { - let norm_tmp = norm_path.with_extension("json.tmp"); - if let Err(e) = std::fs::write(&norm_tmp, &json) { - error!(" Failed to write NormStats tmp {}: {}", norm_tmp.display(), e); - return None; - } - if let Err(e) = std::fs::rename(&norm_tmp, &norm_path) { - error!(" Failed to rename NormStats {} -> {}: {}", norm_tmp.display(), norm_path.display(), e); - drop(std::fs::remove_file(&norm_tmp)); - return None; - } - info!(" Saved NormStats to {}", norm_path.display()); - } - Err(e) => { - error!(" Failed to serialize NormStats: {}", e); - return None; - } - } - - let train_warmup = window.train.len().saturating_sub(train_norm.len()); - let train_bars_aligned = window.train.get(train_warmup..).unwrap_or(&window.train).to_vec(); - let val_warmup = window.val.len().saturating_sub(val_norm.len()); - let val_bars_aligned = window.val.get(val_warmup..).unwrap_or(&window.val).to_vec(); - - Some((train_norm, val_norm, train_bars_aligned, val_bars_aligned)) -} - // --------------------------------------------------------------------------- // DQN Training // --------------------------------------------------------------------------- -/// Convert pre-aligned features and bars into the format expected by +/// Convert fxcache feature/target slices into the format expected by /// `DQNTrainer::train_with_preloaded_data`: `Vec<(FeatureVector, Vec)>`. /// -/// The `Vec` target contains OHLCV prices for the trainer's internal -/// reward computation and portfolio simulation. -fn features_to_trainer_format( - features: &[FeatureVector], - bars: &[OHLCVBar], +/// Zero bar reconstruction — targets already contain [close, next_close, raw_close, raw_next]. +fn features_to_trainer_format_fast( + features: &[[f64; 42]], + targets: &[[f64; 4]], ) -> Vec<(FeatureVector, Vec)> { features .iter() - .zip(bars.iter()) - .map(|(feat, bar)| { - (*feat, vec![bar.open, bar.high, bar.low, bar.close]) - }) + .zip(targets.iter()) + .map(|(f, t)| (*f, t.to_vec())) .collect() } -/// Train a DQN model on a single walk-forward fold using `DQNTrainer`. +/// Build DQN hyperparameters from args, hyperopt JSON, and GPU profile. /// -/// Delegates all GPU optimizations (mixed precision, dynamic batching, gradient -/// accumulation, Rainbow DQN components) to the production trainer. The walk-forward -/// fold structure stays in this binary; only per-fold training delegates to `DQNTrainer`. -/// -/// Returns the best validation loss achieved. -#[allow(clippy::cognitive_complexity, clippy::too_many_arguments)] -fn train_dqn_fold( - fold: usize, - train_features: &[FeatureVector], - val_features: &[FeatureVector], - train_bars: &[OHLCVBar], - val_bars: &[OHLCVBar], +/// Extracted from the old `train_dqn_fold` so it can be called ONCE before the fold +/// loop. The returned hyperparams are ready for `DQNTrainer::new`. +#[allow(clippy::cognitive_complexity)] +fn build_dqn_hyperparams( args: &Args, - output_dir: &Path, hp: &Option, - pre_uploaded_gpu_data: Option, - checkpoint_prefix: &str, -) -> Result { - // Caller must pre-align bars to features (skip warmup period before calling). - debug_assert_eq!( - train_bars.len(), - train_features.len(), - "train_bars ({}) must be pre-aligned to train_features ({})", - train_bars.len(), - train_features.len(), - ); - debug_assert_eq!( - val_bars.len(), - val_features.len(), - "val_bars ({}) must be pre-aligned to val_features ({})", - val_bars.len(), - val_features.len(), - ); - info!(" [DQN] Fold {} -- {} train, {} val features", fold, train_features.len(), val_features.len()); - - // Compute average spread slippage in bps from training bar prices. - // Same pattern as evaluate_baseline: total_cost = commission + spread. - let avg_price = if train_bars.is_empty() { - 0.0 - } else { - train_bars.iter().map(|b| b.close).sum::() / train_bars.len() as f64 - }; - let avg_spread_bps = spread_cost_bps(avg_price, args.tick_size, args.spread_ticks); - let total_cost_bps = args.tx_cost_bps + avg_spread_bps; - info!(" [DQN] Fold {} total tx cost: {:.2} bps (commission {:.1} + spread {:.2})", - fold, total_cost_bps, args.tx_cost_bps, avg_spread_bps); - - // Build DQNHyperparameters -- hyperopt JSON overrides defaults, CLI args override both. - // `..DQNHyperparameters::default()` enables all Rainbow components (PER, dueling, - // distributional C51, noisy nets, gradient accumulation, mixed precision auto-detect). - // transaction_cost_multiplier scales internal reward cost by total trading friction. - // - // hidden_dim_base: hyperopt JSON takes priority. When absent, use VRAM-aware - // scaling so that large GPUs (L40S 48GB, H100 80GB) get proportionally wider - // networks instead of the tiny CPU defaults. - // Load GPU profile for VRAM-aware defaults (replaces scattered if/else chains) + total_cost_bps: f64, +) -> DQNHyperparameters { let gpu_profile = GpuProfile::load(); let hp_hidden_base = hp_usize(hp, "hidden_dim_base"); @@ -499,9 +385,6 @@ fn train_dqn_fold( Some(gpu_profile.training.hidden_dim_base) }); - // Noisy nets are always on (mandatory feature). Epsilon-greedy must be - // low — otherwise epsilon=1.0 forces pure random actions for the entire - // run, preventing the model from ever learning. let epsilon_start = hp_f64(hp, "epsilon_start").unwrap_or(0.05); let mut hyperparams = DQNHyperparameters { @@ -519,7 +402,6 @@ fn train_dqn_fold( warmup_steps: 0, early_stopping_enabled: true, transaction_cost_multiplier: total_cost_bps, - // Pass through hyperopt architectural choices per_alpha: hp_f64(hp, "per_alpha").unwrap_or(0.6), per_beta_start: hp_f64(hp, "per_beta_start").unwrap_or(0.4), dueling_hidden_dim: hp_usize(hp, "dueling_hidden_dim").unwrap_or(128), @@ -538,7 +420,6 @@ fn train_dqn_fold( noisy_sigma_init: hp_f64(hp, "noisy_sigma_init").unwrap_or(0.5), num_quantiles: hp_usize(hp, "num_quantiles").unwrap_or(64), noisy_epsilon_floor: hp_f64(hp, "noisy_epsilon_floor").unwrap_or(0.05).into(), - // Reward-shaping & environment params — critical for trade generation hold_penalty_weight: hp_f64(hp, "hold_penalty_weight").unwrap_or(0.01), max_position_absolute: hp_f64(hp, "max_position_absolute").unwrap_or(2.0), huber_delta: hp_f64(hp, "huber_delta").unwrap_or(10.0), @@ -551,20 +432,15 @@ fn train_dqn_fold( trades_data_dir: args.trades_data_dir.as_ref().map(|p| p.to_string_lossy().into_owned()).unwrap_or_else(|| "test_data/futures-baseline-trades".to_string()), offline_mode: args.offline, dataset_path: args.dataset_path.as_ref().map(|p| p.to_string_lossy().into_owned()), - // GPU PER is mandatory on CUDA. VRAM fraction controls AutoReplaySizer. - // Small GPUs (RTX 3050 profile: buffer_size=5000 < 100K threshold) bypass - // AutoReplaySizer entirely, so VRAM fraction is irrelevant for them. replay_buffer_vram_fraction: gpu_profile.training.replay_buffer_vram_fraction, - // GPU experience collection: n_episodes auto-scales from VRAM gpu_timesteps_per_episode: gpu_profile.experience.gpu_timesteps_per_episode, ..DQNHyperparameters::default() }; // Load training profile (TOML) and apply to hyperparams. - // Profile values override the struct defaults above; CLI args re-applied below win. let profile = ml::training_profile::DqnTrainingProfile::load(&args.training_profile); profile.apply_to(&mut hyperparams); - // CLI args override profile: re-apply any arg that the user can set explicitly. + // CLI args override profile hyperparams.epochs = args.epochs; hyperparams.learning_rate = hp_f64(hp, "learning_rate").unwrap_or(args.learning_rate); hyperparams.initial_capital = args.initial_capital as f32; @@ -572,20 +448,39 @@ fn train_dqn_fold( hyperparams.min_hold_bars = mhb; } - // Create DQNTrainer -- auto-detects GPU, mixed precision, dynamic batch sizing - let mut trainer = DQNTrainer::new(hyperparams) - .context("Failed to create DQNTrainer")?; + hyperparams +} - // Inject pre-uploaded GPU data from DoubleBufferedLoader (overlapped with previous fold) - if let Some(gpu_data) = pre_uploaded_gpu_data { - info!(" [DQN] Fold {} -- injecting pre-uploaded GPU data ({} bars), skipping lazy upload", - fold, gpu_data.num_bars); - trainer.set_gpu_data(gpu_data); - } +/// Train a single DQN fold on a pre-initialized trainer. +/// +/// The trainer already has GPU data uploaded (via `init_from_fxcache`). +/// This function sets per-fold ranges, resets state, and runs training. +/// +/// Returns the best validation loss achieved. +#[allow(clippy::cognitive_complexity, clippy::too_many_arguments)] +fn train_dqn_fold( + rt: &tokio::runtime::Runtime, + trainer: &mut DQNTrainer, + fold: usize, + train_features: &[[f64; 42]], + val_features: &[[f64; 42]], + train_targets: &[[f64; 4]], + val_targets: &[[f64; 4]], + range: &ml::walk_forward::FoldRange, + output_dir: &Path, + checkpoint_prefix: &str, +) -> Result { + info!(" [DQN] Fold {} -- {} train, {} val features", fold, train_features.len(), val_features.len()); - // Convert features + bars to trainer format - let training_data = features_to_trainer_format(train_features, train_bars); - let val_data = features_to_trainer_format(val_features, val_bars); + // Set fold range + val data + reset + trainer.set_training_range(range.train_start, range.train_end, range.val_start, range.val_end); + trainer.set_val_data_from_slices(val_features, val_targets, range.val_start); + rt.block_on(trainer.reset_for_fold()) + .context("reset_for_fold failed")?; + + // Convert to Vec format for the training API (Task 5 will eliminate this) + let training_data = features_to_trainer_format_fast(train_features, train_targets); + let val_data = features_to_trainer_format_fast(val_features, val_targets); // Checkpoint callback: save best model to output directory let output_dir_owned = output_dir.to_path_buf(); @@ -593,7 +488,6 @@ fn train_dqn_fold( let checkpoint_callback = move |epoch: usize, data: Vec, is_best: bool| -> Result { let suffix = if is_best { "best" } else { &format!("epoch{}", epoch) }; let ckpt_path = output_dir_owned.join(format!("{}_fold{}_{}.safetensors", prefix_owned, fold, suffix)); - // Atomic write: write to .tmp then rename (POSIX rename is atomic) let tmp_path = ckpt_path.with_extension("safetensors.tmp"); std::fs::write(&tmp_path, &data) .with_context(|| format!("Failed to write checkpoint tmp: {}", tmp_path.display()))?; @@ -605,14 +499,6 @@ fn train_dqn_fold( Ok(ckpt_path.to_string_lossy().into_owned()) }; - // Run the async training loop via a tokio runtime. - // The trainer handles epochs, early stopping, epsilon decay, validation, and - // all Rainbow DQN components internally. - let rt = tokio::runtime::Builder::new_current_thread() - .enable_all() - .build() - .context("Failed to create tokio runtime for DQN training")?; - let metrics = rt.block_on( trainer.train_with_preloaded_data(training_data, val_data, checkpoint_callback) ).map_err(|e| { @@ -842,24 +728,11 @@ fn run_training(args: &Args) -> Result> { } }); - let (bars, all_features) = if let Some(cached) = fxcache_data { - // Reconstruct bars from cached targets (raw_close at index 2) + real timestamps - let n = cached.bar_count; - let mut bars = Vec::with_capacity(n); - for i in 0..n { - let close = cached.targets[i][2]; // raw_close - bars.push(ml::features::extraction::OHLCVBar { - timestamp: chrono::DateTime::from_timestamp_nanos(cached.timestamps[i]), - open: close, - high: close, - low: close, - close, - volume: 0.0, - }); - } + // Load data into fxcache-compatible arrays: features, targets, timestamps, ofi + let fxcache = if let Some(cached) = fxcache_data { info!(" Loaded {} bars + features from fxcache in {:.1}s", - n, data_load_start.elapsed().as_secs_f64()); - (bars, cached.features) + cached.bar_count, data_load_start.elapsed().as_secs_f64()); + cached } else { // Fall back to DBN loading — this is SLOW (148GB MBP-10 parsing) info!(" Loading OHLCV bars from DBN files..."); @@ -873,41 +746,50 @@ fn run_training(args: &Args) -> Result> { bars.last().map(|b| b.timestamp.to_string()).unwrap_or_default(), ); - // 2. Extract features info!(" Extracting {}-dimensional features...", args.feature_dim); let all_features = extract_ml_features(&bars) .context("Feature extraction failed")?; - info!(" Extracted {} feature vectors (warmup period consumed {} bars)", - all_features.len(), - bars.len().saturating_sub(all_features.len()), - ); - // Trim bars to align with features (skip warmup) let warmup_offset = bars.len().saturating_sub(all_features.len()); - let bars = bars[warmup_offset..].to_vec(); - (bars, all_features) + info!(" Extracted {} feature vectors (warmup period consumed {} bars)", + all_features.len(), warmup_offset); + + // Build FxCacheData from DBN results (features are already warmup-trimmed) + let aligned_bars = &bars[warmup_offset..]; + let n = all_features.len(); + let timestamps: Vec = aligned_bars.iter() + .map(|b| b.timestamp.timestamp_nanos_opt().unwrap_or(0)) + .collect(); + let targets: Vec<[f64; 4]> = aligned_bars.iter() + .map(|b| [b.close, b.close, b.close, b.close]) + .collect(); + let ofi = vec![[0.0_f64; 8]; n]; + + ml::fxcache::FxCacheData { + timestamps, + features: all_features, + targets, + ofi, + cache_key: [0u8; 32], + bar_count: n, + } }; - // Since features skip the warmup period, we need bars aligned to features. - // Features start at bar index warmup_offset (typically 50). - let warmup_offset = bars.len().saturating_sub(all_features.len()); - let aligned_bars = bars.get(warmup_offset..).unwrap_or(&bars); - - // 3. Generate walk-forward windows - info!("Step 3/5: Generating walk-forward windows..."); + // 2. Generate walk-forward fold ranges from timestamps (zero-copy) + info!("Step 2/5: Generating walk-forward fold ranges..."); let wf_config = WalkForwardConfig { initial_train_months: args.train_months, val_months: args.val_months, test_months: args.test_months, step_months: args.step_months, }; - let windows = generate_walk_forward_windows(aligned_bars, &wf_config); - if windows.is_empty() { + let fold_ranges = generate_walk_forward_indices_from_timestamps(&fxcache.timestamps, &wf_config); + if fold_ranges.is_empty() { anyhow::bail!( - "No walk-forward windows generated. Need at least {} months of data.", + "No walk-forward folds generated. Need at least {} months of data.", wf_config.initial_train_months + wf_config.val_months + wf_config.test_months ); } - info!(" Generated {} walk-forward folds", windows.len()); + info!(" Generated {} walk-forward folds (zero-copy index ranges)", fold_ranges.len()); // Record data loading + feature extraction time if train_dqn { @@ -921,77 +803,102 @@ fn run_training(args: &Args) -> Result> { std::fs::create_dir_all(&args.output_dir) .with_context(|| format!("Failed to create output dir: {}", args.output_dir.display()))?; + // 3. Create tokio runtime ONCE, create DQN trainer ONCE, upload fxcache to GPU ONCE + let rt = tokio::runtime::Builder::new_current_thread() + .enable_all() + .build() + .context("Failed to create tokio runtime")?; + + // Compute average spread slippage in bps from the full dataset + let avg_price = { + let sum: f64 = fxcache.targets.iter().map(|t| t[2]).sum(); // raw_close + if fxcache.bar_count > 0 { sum / fxcache.bar_count as f64 } else { 0.0 } + }; + let avg_spread_bps = spread_cost_bps(avg_price, args.tick_size, args.spread_ticks); + let total_cost_bps = args.tx_cost_bps + avg_spread_bps; + info!(" Total tx cost: {:.2} bps (commission {:.1} + spread {:.2})", + total_cost_bps, args.tx_cost_bps, avg_spread_bps); + + // Build DQN trainer ONCE (shared across folds) + let mut dqn_trainer = if train_dqn { + let hp = load_hyperopt_params(&args.hyperopt_params, "dqn"); + let hyperparams = build_dqn_hyperparams(args, &hp, total_cost_bps); + let mut trainer = DQNTrainer::new(hyperparams) + .context("Failed to create DQNTrainer")?; + + // Upload full fxcache to GPU ONCE — all folds index into this data + info!(" Uploading {} bars to GPU via init_from_fxcache...", fxcache.bar_count); + rt.block_on(trainer.init_from_fxcache( + &fxcache.features, &fxcache.targets, &fxcache.ofi, + )).context("init_from_fxcache failed")?; + info!(" GPU data uploaded — ready for fold loop"); + Some(trainer) + } else { + None + }; + // 4. Train each fold info!("Step 4/5: Training models on each fold..."); let mut dqn_results: Vec<(usize, f64)> = Vec::new(); let mut ppo_results: Vec<(usize, f64)> = Vec::new(); - let mut prefetched_data: Option = None; - // GPU double-buffer: pre-uploaded GPU data for the next DQN fold. - // While fold N trains on GPU, fold N+1's data is uploaded in background. - let mut dqn_gpu_staged: Option = None; - - for (fold_idx, window) in windows.iter().enumerate() { - info!("--- Fold {} ---", window.fold); + for (fold_idx, range) in fold_ranges.iter().enumerate() { + info!("--- Fold {} ---", range.fold); info!( - " Train: {} bars (up to {}), Val: {} bars (up to {}), Test: {} bars (up to {})", - window.train.len(), - window.train_end, - window.val.len(), - window.val_end, - window.test.len(), - window.test_end, + " Train: bars [{}..{}] ({} bars), Val: bars [{}..{}] ({} bars)", + range.train_start, range.train_end, + range.train_end - range.train_start, + range.val_start, range.val_end, + range.val_end - range.val_start, ); - // Get fold data: from prefetcher (fold 1+) or prepare fresh (fold 0) - let fold_data = if let Some(data) = prefetched_data.take() { - info!(" Using prefetched data for fold {}", window.fold); - Some(data) - } else { - prepare_fold_data(window, &args.output_dir) - }; + // Slice features/targets for this fold (zero-copy from fxcache arrays) + let train_feat = &fxcache.features[range.train_start..range.train_end]; + let val_feat = &fxcache.features[range.val_start..range.val_end]; + let train_tgt = &fxcache.targets[range.train_start..range.train_end]; + let val_tgt = &fxcache.targets[range.val_start..range.val_end]; - // Kick off prefetch for NEXT fold on background thread - let prefetch_rx = if fold_idx + 1 < windows.len() { - let next_window = windows.get(fold_idx + 1).cloned(); - if let Some(nw) = next_window { - let out_dir = args.output_dir.clone(); - let (tx, rx) = std::sync::mpsc::channel::>(); - let _handle = std::thread::Builder::new() - .name("fold-prefetch".into()) - .spawn(move || { - info!(" [Prefetch] Loading fold {} data on background thread", nw.fold); - let result = prepare_fold_data(&nw, &out_dir); - drop(tx.send(result)); - }); - Some(rx) - } else { - None - } - } else { - None - }; - - let Some((train_norm, val_norm, train_bars_aligned, val_bars_aligned)) = fold_data else { - // Collect prefetch result even if this fold failed - if let Some(rx) = prefetch_rx { - if let Ok(data) = rx.recv() { - prefetched_data = data; - } - } + if train_feat.is_empty() || val_feat.is_empty() { + warn!(" Fold {} -- empty features, skipping", range.fold); continue; - }; + } + + // Normalize features using training fold statistics + let norm_stats = NormStats::from_features(train_feat); + let train_norm = norm_stats.normalize_batch(train_feat); + let val_norm = norm_stats.normalize_batch(val_feat); + + // Save NormStats for this fold + let norm_path = args.output_dir.join(format!("norm_stats_fold{}.json", range.fold)); + match serde_json::to_string_pretty(&norm_stats) { + Ok(json) => { + let norm_tmp = norm_path.with_extension("json.tmp"); + if let Err(e) = std::fs::write(&norm_tmp, &json) { + error!(" Failed to write NormStats tmp {}: {}", norm_tmp.display(), e); + continue; + } + if let Err(e) = std::fs::rename(&norm_tmp, &norm_path) { + error!(" Failed to rename NormStats: {}", e); + drop(std::fs::remove_file(&norm_tmp)); + continue; + } + info!(" Saved NormStats to {}", norm_path.display()); + } + Err(e) => { + error!(" Failed to serialize NormStats: {}", e); + continue; + } + } // Train DQN - if train_dqn { + if let Some(ref mut trainer) = dqn_trainer { let fold_str = fold_idx.to_string(); let fold_start = std::time::Instant::now(); - // Take pre-uploaded GPU data from previous fold's background upload - let mut gpu_data_for_fold = dqn_gpu_staged.take(); - if args.ensemble_top_k > 1 && args.hyperopt_params.is_some() { - // Ensemble mode: train one model per top-K hyperopt param set + // Ensemble mode: train one model per top-K hyperopt param set. + // NOTE: ensemble requires separate trainers with different hyperparams, + // so we create per-ensemble trainers (only the primary shares GPU data). let param_sets = load_top_k_params( &args.hyperopt_params, "dqn", @@ -1000,62 +907,70 @@ fn run_training(args: &Args) -> Result> { for (k, hp) in param_sets.iter().enumerate() { info!( " [DQN] Training ensemble member {}/{} on fold {}", - k + 1, - param_sets.len(), - window.fold + k + 1, param_sets.len(), range.fold ); let prefix = format!("dqn_ensemble_{}", k); - // Only the first ensemble member uses the pre-uploaded GPU data - let gpu_data = if k == 0 { gpu_data_for_fold.take() } else { None }; - match train_dqn_fold( - window.fold, - &train_norm, - &val_norm, - &train_bars_aligned, - &val_bars_aligned, - args, - &args.output_dir, - hp, - gpu_data, - &prefix, - ) { - Ok(best_loss) => { - info!( - " [DQN] Ensemble member {} fold {} best_loss={:.6}", - k, window.fold, best_loss - ); - // Record the best ensemble member's loss as the fold result - if k == 0 { + if k == 0 { + // Primary ensemble member uses the shared trainer + match train_dqn_fold( + &rt, trainer, range.fold, + &train_norm, &val_norm, train_tgt, val_tgt, + range, &args.output_dir, &prefix, + ) { + Ok(best_loss) => { + info!(" [DQN] Ensemble member {} fold {} best_loss={:.6}", + k, range.fold, best_loss); let elapsed = fold_start.elapsed().as_secs_f64(); metrics::set_epoch("dqn", &fold_str, fold_idx as f64); metrics::set_epoch_loss("dqn", &fold_str, best_loss); metrics::set_validation_loss("dqn", &fold_str, best_loss); metrics::set_iteration_seconds("dqn", &fold_str, elapsed); - dqn_results.push((window.fold, best_loss)); + dqn_results.push((range.fold, best_loss)); + } + Err(e) => { + error!(" [DQN] Ensemble member {} fold {} failed: {}", + k, range.fold, e); } } - Err(e) => { - error!( - " [DQN] Ensemble member {} fold {} failed: {}", - k, window.fold, e - ); + } else { + // Secondary ensemble members create separate trainers + let ens_hyperparams = build_dqn_hyperparams(args, hp, total_cost_bps); + match DQNTrainer::new(ens_hyperparams) { + Ok(mut ens_trainer) => { + // Upload fxcache to this ensemble trainer too + if let Err(e) = rt.block_on(ens_trainer.init_from_fxcache( + &fxcache.features, &fxcache.targets, &fxcache.ofi, + )) { + error!(" [DQN] Ensemble member {} init_from_fxcache failed: {}", k, e); + continue; + } + match train_dqn_fold( + &rt, &mut ens_trainer, range.fold, + &train_norm, &val_norm, train_tgt, val_tgt, + range, &args.output_dir, &prefix, + ) { + Ok(best_loss) => { + info!(" [DQN] Ensemble member {} fold {} best_loss={:.6}", + k, range.fold, best_loss); + } + Err(e) => { + error!(" [DQN] Ensemble member {} fold {} failed: {}", + k, range.fold, e); + } + } + } + Err(e) => { + error!(" [DQN] Failed to create ensemble trainer {}: {}", k, e); + } } } } } else { // Single-model mode (default) - let hp = load_hyperopt_params(&args.hyperopt_params, "dqn"); match train_dqn_fold( - window.fold, - &train_norm, - &val_norm, - &train_bars_aligned, - &val_bars_aligned, - args, - &args.output_dir, - &hp, - gpu_data_for_fold, - "dqn", + &rt, trainer, range.fold, + &train_norm, &val_norm, train_tgt, val_tgt, + range, &args.output_dir, "dqn", ) { Ok(best_loss) => { let elapsed = fold_start.elapsed().as_secs_f64(); @@ -1063,26 +978,41 @@ fn run_training(args: &Args) -> Result> { metrics::set_epoch_loss("dqn", &fold_str, best_loss); metrics::set_validation_loss("dqn", &fold_str, best_loss); metrics::set_iteration_seconds("dqn", &fold_str, elapsed); - dqn_results.push((window.fold, best_loss)); + dqn_results.push((range.fold, best_loss)); } Err(e) => { - error!(" [DQN] Fold {} failed: {:#}", window.fold, e); + error!(" [DQN] Fold {} failed: {:#}", range.fold, e); } } } } - // Train PPO + // Train PPO (uses old path — constructs minimal bars from fxcache targets) if train_ppo { let hp = load_hyperopt_params(&args.hyperopt_params, "ppo"); let fold_str = fold_idx.to_string(); let fold_start = std::time::Instant::now(); + + // PPO compatibility: construct minimal OHLCVBars from fxcache targets + let train_bars: Vec = fxcache.timestamps[range.train_start..range.train_end] + .iter().zip(fxcache.targets[range.train_start..range.train_end].iter()) + .map(|(&ts, t)| OHLCVBar { + timestamp: chrono::DateTime::from_timestamp_nanos(ts), + open: t[0], high: t[0], low: t[0], close: t[2], volume: 0.0, + }).collect(); + let val_bars: Vec = fxcache.timestamps[range.val_start..range.val_end] + .iter().zip(fxcache.targets[range.val_start..range.val_end].iter()) + .map(|(&ts, t)| OHLCVBar { + timestamp: chrono::DateTime::from_timestamp_nanos(ts), + open: t[0], high: t[0], low: t[0], close: t[2], volume: 0.0, + }).collect(); + match train_ppo_fold( - window.fold, + range.fold, &train_norm, &val_norm, - &train_bars_aligned, - &val_bars_aligned, + &train_bars, + &val_bars, args, &args.output_dir, &hp, @@ -1093,54 +1023,10 @@ fn run_training(args: &Args) -> Result> { metrics::set_epoch_loss("ppo", &fold_str, best_loss); metrics::set_validation_loss("ppo", &fold_str, best_loss); metrics::set_iteration_seconds("ppo", &fold_str, elapsed); - ppo_results.push((window.fold, best_loss)); + ppo_results.push((range.fold, best_loss)); } Err(e) => { - error!(" [PPO] Fold {} failed: {:#}", window.fold, e); - } - } - } - - // Collect prefetch result for next iteration - if let Some(rx) = prefetch_rx { - match rx.recv() { - Ok(Some(data)) => { - info!(" Prefetched fold {} data ready ({} train features)", - windows.get(fold_idx + 1).map_or(0, |w| w.fold), data.0.len()); - - // GPU double-buffer: pre-upload next fold's DQN data to GPU now - // while the result is fresh and GPU is idle between folds. - if train_dqn { - let next_train_data = features_to_trainer_format(&data.0, &data.2); - let ctx = cudarc::driver::CudaContext::new(0); - if let Ok(ctx) = ctx { - let stream = ctx.new_stream(); - if let Ok(stream) = stream { - let stream = Arc::new(stream); - match DqnGpuData::upload(&next_train_data, &stream) { - Ok(gpu_data) => { - info!(" [DoubleBuffer] Pre-uploaded {} bars to GPU for next fold ({:.1} MB)", - gpu_data.num_bars, - gpu_data.vram_bytes() as f64 / 1_048_576.0); - dqn_gpu_staged = Some(gpu_data); - } - Err(e) => { - warn!(" [DoubleBuffer] GPU pre-upload failed, will upload lazily: {}", e); - } - } - } - } - } - - prefetched_data = Some(data); - } - Ok(None) => { - warn!(" Prefetch for next fold returned None (will load synchronously)"); - prefetched_data = None; - } - Err(_) => { - warn!(" Prefetch thread disconnected (will load synchronously)"); - prefetched_data = None; + error!(" [PPO] Fold {} failed: {:#}", range.fold, e); } } } @@ -1151,6 +1037,7 @@ fn run_training(args: &Args) -> Result> { info!(" ==================================="); let mut all_results = Vec::new(); + let num_folds = fold_ranges.len(); if train_dqn { info!(" DQN Results ({} folds):", dqn_results.len()); @@ -1165,7 +1052,7 @@ fn run_training(args: &Args) -> Result> { all_results.push(RlTrainingResult { model_name: "dqn".to_owned(), fold_results: dqn_results, - total_epochs: windows.len() * args.epochs, + total_epochs: num_folds * args.epochs, }); } if train_ppo { @@ -1181,7 +1068,7 @@ fn run_training(args: &Args) -> Result> { all_results.push(RlTrainingResult { model_name: "ppo".to_owned(), fold_results: ppo_results, - total_epochs: windows.len() * args.epochs, + total_epochs: num_folds * args.epochs, }); } info!(" Checkpoints saved to: {}", args.output_dir.display()); diff --git a/crates/ml/src/walk_forward.rs b/crates/ml/src/walk_forward.rs index 0d9933b65..e69c81431 100644 --- a/crates/ml/src/walk_forward.rs +++ b/crates/ml/src/walk_forward.rs @@ -525,6 +525,94 @@ pub fn generate_walk_forward_indices( ranges } +/// Generate walk-forward folds as index ranges from nanosecond timestamps. +/// +/// Same date-boundary logic as [`generate_walk_forward_indices`] but operates +/// on raw `i64` timestamps (nanoseconds since Unix epoch) from fxcache data +/// instead of `OHLCVBar`. Difficulty score is always 0.0 because we don't +/// have OHLC data to compute ADX. +pub fn generate_walk_forward_indices_from_timestamps( + timestamps: &[i64], + config: &WalkForwardConfig, +) -> Vec { + if timestamps.is_empty() { + return Vec::new(); + } + + let first_ts = timestamps[0]; + let last_ts = timestamps[timestamps.len() - 1]; + + let data_start = chrono::DateTime::from_timestamp_nanos(first_ts).date_naive(); + let data_end = chrono::DateTime::from_timestamp_nanos(last_ts).date_naive(); + + let mut ranges = Vec::new(); + let mut fold = 0_usize; + + loop { + let total_train_months = config + .initial_train_months + .saturating_add(fold as u32 * config.step_months); + let train_end = match data_start.checked_add_months(Months::new(total_train_months)) { + Some(d) => d, + None => break, + }; + let val_end = match train_end.checked_add_months(Months::new(config.val_months)) { + Some(d) => d, + None => break, + }; + let test_end = match val_end.checked_add_months(Months::new(config.test_months)) { + Some(d) => d, + None => break, + }; + + if test_end > data_end + chrono::Duration::days(1) { + break; + } + + // O(log n) index lookups via partition_point (timestamps are sorted) + let train_end_idx = timestamps.partition_point(|&ts| { + chrono::DateTime::from_timestamp_nanos(ts).date_naive() < train_end + }); + let val_end_idx = timestamps.partition_point(|&ts| { + chrono::DateTime::from_timestamp_nanos(ts).date_naive() < val_end + }); + let test_end_idx = timestamps.partition_point(|&ts| { + chrono::DateTime::from_timestamp_nanos(ts).date_naive() < test_end + }); + + let train_start = 0; // expanding window: always starts at beginning + let val_start = train_end_idx; + let val_end_bound = val_end_idx; + + // Skip folds where any split is too small + let train_len = train_end_idx - train_start; + let val_len = val_end_bound - val_start; + let test_len = test_end_idx - val_end_bound; + + if train_len < MIN_BARS_PER_SPLIT || val_len < MIN_BARS_PER_SPLIT || test_len < MIN_BARS_PER_SPLIT { + if val_len < MIN_BARS_PER_SPLIT || test_len < MIN_BARS_PER_SPLIT { + if fold > 0 { break; } + } + fold = fold.saturating_add(1); + continue; + } + + // No ADX computation — fxcache doesn't have OHLC bars for difficulty + ranges.push(FoldRange { + fold, + train_start, + train_end: train_end_idx, + val_start, + val_end: val_end_bound, + difficulty_score: 0.0, + }); + + fold = fold.saturating_add(1); + } + + ranges +} + /// Per-feature normalization statistics (mean and standard deviation). /// /// Computed from training data and applied to val/test data to prevent From 8d5af977f7cdaa4e4ba2f6e9e98f4e345fabe0c1 Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Thu, 2 Apr 2026 22:50:11 +0200 Subject: [PATCH 4/6] =?UTF-8?q?feat:=20train=5Ffold=5Ffrom=5Fslices=20?= =?UTF-8?q?=E2=80=94=20trainer=20accepts=20contiguous=20feature/target=20s?= =?UTF-8?q?lices?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Adds DQNTrainer::train_fold_from_slices(&[[f64;42]], &[[f64;4]]) so the fold loop in train_baseline_rl passes raw fxcache slices directly, without the caller constructing any Vec<(FeatureVector, Vec)>. Removes features_to_trainer_format_fast helper (no longer needed) and updates train_dqn_fold to call the new method. Co-Authored-By: Claude Sonnet 4.6 --- crates/ml/examples/train_baseline_rl.rs | 21 +------------ crates/ml/src/trainers/dqn/trainer/mod.rs | 38 +++++++++++++++++++++++ 2 files changed, 39 insertions(+), 20 deletions(-) diff --git a/crates/ml/examples/train_baseline_rl.rs b/crates/ml/examples/train_baseline_rl.rs index 2b914818d..779c74681 100644 --- a/crates/ml/examples/train_baseline_rl.rs +++ b/crates/ml/examples/train_baseline_rl.rs @@ -352,21 +352,6 @@ fn hp_bool(params: &Option, key: &str) -> Option { // DQN Training // --------------------------------------------------------------------------- -/// Convert fxcache feature/target slices into the format expected by -/// `DQNTrainer::train_with_preloaded_data`: `Vec<(FeatureVector, Vec)>`. -/// -/// Zero bar reconstruction — targets already contain [close, next_close, raw_close, raw_next]. -fn features_to_trainer_format_fast( - features: &[[f64; 42]], - targets: &[[f64; 4]], -) -> Vec<(FeatureVector, Vec)> { - features - .iter() - .zip(targets.iter()) - .map(|(f, t)| (*f, t.to_vec())) - .collect() -} - /// Build DQN hyperparameters from args, hyperopt JSON, and GPU profile. /// /// Extracted from the old `train_dqn_fold` so it can be called ONCE before the fold @@ -478,10 +463,6 @@ fn train_dqn_fold( rt.block_on(trainer.reset_for_fold()) .context("reset_for_fold failed")?; - // Convert to Vec format for the training API (Task 5 will eliminate this) - let training_data = features_to_trainer_format_fast(train_features, train_targets); - let val_data = features_to_trainer_format_fast(val_features, val_targets); - // Checkpoint callback: save best model to output directory let output_dir_owned = output_dir.to_path_buf(); let prefix_owned = checkpoint_prefix.to_owned(); @@ -500,7 +481,7 @@ fn train_dqn_fold( }; let metrics = rt.block_on( - trainer.train_with_preloaded_data(training_data, val_data, checkpoint_callback) + trainer.train_fold_from_slices(train_features, train_targets, checkpoint_callback) ).map_err(|e| { error!(" [DQN] Fold {} training error chain: {:#}", fold, e); e diff --git a/crates/ml/src/trainers/dqn/trainer/mod.rs b/crates/ml/src/trainers/dqn/trainer/mod.rs index 485113e27..4ae740dfe 100644 --- a/crates/ml/src/trainers/dqn/trainer/mod.rs +++ b/crates/ml/src/trainers/dqn/trainer/mod.rs @@ -527,6 +527,44 @@ impl DQNTrainer { .await } + /// Train one fold from contiguous feature/target slices. + /// Zero Vec allocation — targets read as fixed-size &[[f64; 4]]. + /// + /// Caller must call `set_val_data_from_slices` and `set_training_range` before + /// this method. Val data and OFI offset are already set on `self`. + pub async fn train_fold_from_slices( + &mut self, + features: &[[f64; 42]], + targets: &[[f64; 4]], + checkpoint_callback: F, + ) -> Result + where + F: FnMut(usize, Vec, bool) -> Result + Send, + { + // Clear stale CUDA errors + if let MlDevice::Cuda { ref context, .. } = self.device { + let _ = context.check_err(); + } + + info!("train_fold_from_slices: {} bars", features.len()); + + // Build lightweight training_data with fixed-size [f64; 4] converted to Vec + // This is the last remaining allocation — each [f64; 4].to_vec() is 32 bytes. + // TODO: eliminate by changing train_with_data_full_loop signature + let training_data: Vec<(FeatureVector, Vec)> = features + .iter() + .zip(targets.iter()) + .map(|(f, t)| (*f, t.to_vec())) + .collect(); + + self.val_features_gpu = None; + self.val_closes_gpu = None; + self.val_ofi_gpu = None; + + self.train_with_data_full_loop(&training_data, checkpoint_callback) + .await + } + /// Train with shared preloaded data (zero-copy for hyperopt). /// /// Same as [`train_with_preloaded_data`] but accepts `Arc`-wrapped data, From 3278cec29b2c5feec5ea411bf54b1f29ba1fa141 Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Thu, 2 Apr 2026 22:54:00 +0200 Subject: [PATCH 5/6] =?UTF-8?q?perf:=20PPO=20zero-copy=20fold=20loop=20?= =?UTF-8?q?=E2=80=94=20train=5Ffrom=5Fslices=20+=20delete=20train=5Fppo=5F?= =?UTF-8?q?fold?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Add PpoTrainer::train_from_slices(&[[f64; 42]]) that converts to Vec> internally (PPO train() API requires ownership). Replace the PPO fold loop in train_baseline_rl.rs: eliminates the OHLCVBar construction shim and calls train_from_slices directly on fxcache feature slices. Delete the standalone train_ppo_fold function entirely. Co-Authored-By: Claude Sonnet 4.6 --- crates/ml/examples/train_baseline_rl.rs | 231 +++++++----------------- crates/ml/src/trainers/ppo.rs | 16 ++ 2 files changed, 83 insertions(+), 164 deletions(-) diff --git a/crates/ml/examples/train_baseline_rl.rs b/crates/ml/examples/train_baseline_rl.rs index 779c74681..862cb288f 100644 --- a/crates/ml/examples/train_baseline_rl.rs +++ b/crates/ml/examples/train_baseline_rl.rs @@ -103,7 +103,7 @@ use serde_json::Value; use tracing::{error, info, warn}; use ml::trainers::dqn::{DQNHyperparameters, DQNTrainer}; -use ml::trainers::ppo::{PpoHyperparameters, PpoTrainer}; +use ml::trainers::ppo::{PpoHyperparameters, PpoTrainer, PpoTrainingMetrics}; #[allow(unreachable_pub)] mod baseline_common; @@ -495,134 +495,6 @@ fn train_dqn_fold( Ok(metrics.loss) } -// --------------------------------------------------------------------------- -// PPO Training -// --------------------------------------------------------------------------- - -/// Train a PPO model on a single walk-forward fold using `PpoTrainer`. -/// -/// Delegates all GPU optimizations (dynamic batch sizing, mixed precision when -/// available, GAE computation, trajectory collection, early stopping) to the -/// production trainer. The walk-forward fold structure stays in this binary; -/// only per-fold training delegates to `PpoTrainer`. -/// -/// Returns the best validation loss achieved (approximated by final `value_loss`). -#[allow(clippy::cognitive_complexity)] -fn train_ppo_fold( - fold: usize, - train_features: &[FeatureVector], - val_features: &[FeatureVector], - train_bars: &[OHLCVBar], - _val_bars: &[OHLCVBar], - args: &Args, - output_dir: &Path, - hp: &Option, -) -> Result { - info!(" [PPO] Fold {} -- {} train, {} val features", fold, train_features.len(), val_features.len()); - - let n_train = train_features.len(); - if n_train < 2 { - warn!(" [PPO] Fold {} -- insufficient training features ({})", fold, n_train); - return Ok(f64::MAX); - } - - // Compute average spread slippage in bps from training bar prices. - // Same pattern as evaluate_baseline: total_cost = commission + spread. - let avg_price = if train_bars.is_empty() { - 0.0 - } else { - train_bars.iter().map(|b| b.close).sum::() / train_bars.len() as f64 - }; - let avg_spread_bps = spread_cost_bps(avg_price, args.tick_size, args.spread_ticks); - let total_cost_bps = args.tx_cost_bps + avg_spread_bps; - info!(" [PPO] Fold {} total tx cost: {:.2} bps (commission {:.1} + spread {:.2})", - fold, total_cost_bps, args.tx_cost_bps, avg_spread_bps); - - // Build PpoHyperparameters -- hyperopt JSON overrides defaults, CLI args override both. - // Start from conservative() baseline so all fields have sane values. - // - // hidden_dim_base: hyperopt JSON takes priority. When absent, use VRAM-aware - // scaling so that large GPUs get proportionally wider networks. PPO uses the - // "ppo_policy" model type for base dim since the trainer internally scales - // the value network from the same base (value: [4*base, 3*base, 2*base, base, base/2]). - let hp_ppo_hidden_base = hp_usize(hp, "hidden_dim_base"); - let ppo_hidden_base = hp_ppo_hidden_base.or_else(|| { - let profile = ml_core::gpu::profile::GpuProfile::load(); - info!(" [PPO] hidden_dim_base: {} (from GPU profile)", profile.training.hidden_dim_base); - Some(profile.training.hidden_dim_base) - }); - - let hyperparams = PpoHyperparameters { - learning_rate: hp_f64(hp, "learning_rate").unwrap_or(args.learning_rate), - actor_learning_rate: Some(hp_f64(hp, "policy_learning_rate").unwrap_or(args.learning_rate)), - critic_learning_rate: Some(hp_f64(hp, "value_learning_rate").unwrap_or(args.learning_rate * 3.0)), - batch_size: ml_core::gpu::profile::GpuProfile::load().training.batch_size, - gamma: hp_f64(hp, "gamma").unwrap_or(0.99), - clip_epsilon: hp_f64(hp, "clip_epsilon").unwrap_or(0.2) as f32, - vf_coef: hp_f64(hp, "value_loss_coeff").unwrap_or(0.5) as f32, - ent_coef: hp_f64(hp, "entropy_coeff").unwrap_or(0.01) as f32, - gae_lambda: hp_f64(hp, "gae_lambda").unwrap_or(0.95) as f32, - rollout_steps: hp_usize(hp, "rollout_steps").unwrap_or(2048), - minibatch_size: hp_usize(hp, "minibatch_size").unwrap_or(64), - epochs: args.epochs, - early_stopping_enabled: true, - transaction_cost_bps: total_cost_bps / 100.0, // total (commission + spread) in bps → pct for trainer - hidden_dim_base: ppo_hidden_base, - ..PpoHyperparameters::conservative() - }; - - // Checkpoint directory for this fold - let fold_ckpt_dir = output_dir.join(format!("ppo_fold{}", fold)); - - // Create PpoTrainer -- auto-detects GPU, mixed precision, dynamic batch sizing - let trainer = PpoTrainer::new( - hyperparams, - args.feature_dim, - &fold_ckpt_dir, - true, // use_gpu = true, trainer auto-detects availability - None, // num_envs: standard single-env collection - ).context("Failed to create PpoTrainer")?; - - // Convert features from FeatureVector to Vec> for the trainer - let market_data: Vec> = train_features - .iter() - .map(|feat| feat.iter().map(|&v| v as f32).collect()) - .collect(); - - // Progress callback: log per-epoch metrics - let progress_callback = move |metrics: ml::trainers::ppo::PpoTrainingMetrics| { - info!( - " [PPO] Fold {} Epoch {}/{} -- policy_loss={:.6} value_loss={:.6} expl_var={:.4} mean_reward={:.4}", - fold, - metrics.epoch, - args.epochs, - metrics.policy_loss, - metrics.value_loss, - metrics.explained_variance, - metrics.mean_reward, - ); - }; - - // Run the async training loop via a tokio runtime. - // PpoTrainer handles epochs, early stopping, trajectory collection, GAE, - // checkpointing, and all GPU optimizations internally. - let rt = tokio::runtime::Builder::new_current_thread() - .enable_all() - .build() - .context("Failed to create tokio runtime for PPO training")?; - - let metrics = rt.block_on( - trainer.train(market_data, progress_callback) - ).context("PpoTrainer training failed")?; - - info!( - " [PPO] Fold {} complete -- value_loss={:.6} policy_loss={:.6} expl_var={:.4} epochs_trained={}", - fold, metrics.value_loss, metrics.policy_loss, metrics.explained_variance, metrics.epoch - ); - - Ok(metrics.value_loss as f64) -} - // --------------------------------------------------------------------------- // Training orchestration // --------------------------------------------------------------------------- @@ -968,47 +840,78 @@ fn run_training(args: &Args) -> Result> { } } - // Train PPO (uses old path — constructs minimal bars from fxcache targets) + // Train PPO (zero-copy: pass fxcache feature slices directly) if train_ppo { - let hp = load_hyperopt_params(&args.hyperopt_params, "ppo"); + let hp_ppo = load_hyperopt_params(&args.hyperopt_params, "ppo"); let fold_str = fold_idx.to_string(); let fold_start = std::time::Instant::now(); - // PPO compatibility: construct minimal OHLCVBars from fxcache targets - let train_bars: Vec = fxcache.timestamps[range.train_start..range.train_end] - .iter().zip(fxcache.targets[range.train_start..range.train_end].iter()) - .map(|(&ts, t)| OHLCVBar { - timestamp: chrono::DateTime::from_timestamp_nanos(ts), - open: t[0], high: t[0], low: t[0], close: t[2], volume: 0.0, - }).collect(); - let val_bars: Vec = fxcache.timestamps[range.val_start..range.val_end] - .iter().zip(fxcache.targets[range.val_start..range.val_end].iter()) - .map(|(&ts, t)| OHLCVBar { - timestamp: chrono::DateTime::from_timestamp_nanos(ts), - open: t[0], high: t[0], low: t[0], close: t[2], volume: 0.0, - }).collect(); + // Compute per-fold tx cost from fxcache targets (raw close at index 2) + let train_closes: Vec = fxcache.targets[range.train_start..range.train_end] + .iter().map(|t| t[2]).collect(); + let avg_price = if train_closes.is_empty() { + 0.0 + } else { + train_closes.iter().sum::() / train_closes.len() as f64 + }; + let avg_spread_bps_ppo = spread_cost_bps(avg_price, args.tick_size, args.spread_ticks); + let total_cost_bps_ppo = args.tx_cost_bps + avg_spread_bps_ppo; - match train_ppo_fold( - range.fold, - &train_norm, - &val_norm, - &train_bars, - &val_bars, - args, - &args.output_dir, - &hp, - ) { - Ok(best_loss) => { - let elapsed = fold_start.elapsed().as_secs_f64(); - metrics::set_epoch("ppo", &fold_str, fold_idx as f64); - metrics::set_epoch_loss("ppo", &fold_str, best_loss); - metrics::set_validation_loss("ppo", &fold_str, best_loss); - metrics::set_iteration_seconds("ppo", &fold_str, elapsed); - ppo_results.push((range.fold, best_loss)); - } - Err(e) => { - error!(" [PPO] Fold {} failed: {:#}", range.fold, e); + // Build PpoHyperparameters inline (same as old train_ppo_fold) + let hp_ppo_hidden_base = hp_usize(&hp_ppo, "hidden_dim_base"); + let ppo_hidden_base = hp_ppo_hidden_base.or_else(|| { + let profile = ml_core::gpu::profile::GpuProfile::load(); + info!(" [PPO] hidden_dim_base: {} (from GPU profile)", profile.training.hidden_dim_base); + Some(profile.training.hidden_dim_base) + }); + let ppo_hp = PpoHyperparameters { + learning_rate: hp_f64(&hp_ppo, "learning_rate").unwrap_or(args.learning_rate), + actor_learning_rate: Some(hp_f64(&hp_ppo, "policy_learning_rate").unwrap_or(args.learning_rate)), + critic_learning_rate: Some(hp_f64(&hp_ppo, "value_learning_rate").unwrap_or(args.learning_rate * 3.0)), + batch_size: ml_core::gpu::profile::GpuProfile::load().training.batch_size, + gamma: hp_f64(&hp_ppo, "gamma").unwrap_or(0.99), + clip_epsilon: hp_f64(&hp_ppo, "clip_epsilon").unwrap_or(0.2) as f32, + vf_coef: hp_f64(&hp_ppo, "value_loss_coeff").unwrap_or(0.5) as f32, + ent_coef: hp_f64(&hp_ppo, "entropy_coeff").unwrap_or(0.01) as f32, + gae_lambda: hp_f64(&hp_ppo, "gae_lambda").unwrap_or(0.95) as f32, + rollout_steps: hp_usize(&hp_ppo, "rollout_steps").unwrap_or(2048), + minibatch_size: hp_usize(&hp_ppo, "minibatch_size").unwrap_or(64), + epochs: args.epochs, + early_stopping_enabled: true, + transaction_cost_bps: total_cost_bps_ppo / 100.0, + hidden_dim_base: ppo_hidden_base, + ..PpoHyperparameters::conservative() + }; + + let fold_ckpt_dir = args.output_dir.join(format!("ppo_fold{}", range.fold)); + let ppo_trainer = PpoTrainer::new( + ppo_hp, + args.feature_dim, + &fold_ckpt_dir, + true, + None, + ); + match ppo_trainer { + Ok(trainer) => { + let fold_idx_cap = range.fold; + let epochs_cap = args.epochs; + let progress_cb = move |m: PpoTrainingMetrics| { + info!("[PPO] Fold {} Epoch {}/{} -- value_loss={:.6}", fold_idx_cap, m.epoch, epochs_cap, m.value_loss); + }; + match rt.block_on(trainer.train_from_slices(&train_norm, progress_cb)) { + Ok(metrics) => { + let best_loss = metrics.value_loss as f64; + let elapsed = fold_start.elapsed().as_secs_f64(); + metrics::set_epoch("ppo", &fold_str, fold_idx as f64); + metrics::set_epoch_loss("ppo", &fold_str, best_loss); + metrics::set_validation_loss("ppo", &fold_str, best_loss); + metrics::set_iteration_seconds("ppo", &fold_str, elapsed); + ppo_results.push((range.fold, best_loss)); + } + Err(e) => error!("PPO fold {} failed: {:#}", range.fold, e), + } } + Err(e) => error!(" [PPO] Fold {} trainer init failed: {:#}", range.fold, e), } } } diff --git a/crates/ml/src/trainers/ppo.rs b/crates/ml/src/trainers/ppo.rs index 694652389..6932a2433 100644 --- a/crates/ml/src/trainers/ppo.rs +++ b/crates/ml/src/trainers/ppo.rs @@ -438,6 +438,22 @@ impl PpoTrainer { self.train_gpu(market_data, progress_callback).await } + /// Train from contiguous feature slices (fxcache format). + /// Converts to Vec> internally — PPO's train() API requires it. + pub async fn train_from_slices( + &self, + features: &[[f64; 42]], + progress_callback: F, + ) -> Result + where + F: FnMut(PpoTrainingMetrics) + Send, + { + let market_data: Vec> = features.iter() + .map(|f| f.iter().map(|&v| v as f32).collect()) + .collect(); + self.train(market_data, progress_callback).await + } + async fn train_gpu( &self, market_data: Vec>, From 95fc94ae9c4faaf1dd484a13f18d8856f47b2216 Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Thu, 2 Apr 2026 23:21:45 +0200 Subject: [PATCH 6/6] =?UTF-8?q?fix:=20precommit=20review=20=E2=80=94=20ens?= =?UTF-8?q?emble=20reuse,=20keep=20gpu=5Fdata=20across=20folds,=20eliminat?= =?UTF-8?q?e=20Vec?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Fix 1 (HIGH): Ensemble trainers were created + data uploaded inside the fold loop, causing redundant GPU uploads per fold per ensemble member. Moved creation + init_from_fxcache before the fold loop; inside the loop we now reuse trainers via set_training_range + reset_for_fold. Fix 2 (MEDIUM): reset_for_fold cleared gpu_data = None, forcing re-upload on every fold even though the full dataset was already GPU-resident via init_from_fxcache. Removed the clearing — data stays on GPU across folds. Fix 3 (LOW): train_fold_from_slices allocated a Vec per bar via t.to_vec(). Added train_with_data_full_loop_slices that accepts &[([f64; 42], [f64; 4])] — both are Copy stack types, zero heap alloc per element. Added matching collect_gpu_experiences_slices and run_training_steps_slices helpers. Co-Authored-By: Claude Opus 4.6 (1M context) --- crates/ml/examples/train_baseline_rl.rs | 140 +-- crates/ml/src/trainers/dqn/trainer/mod.rs | 15 +- .../src/trainers/dqn/trainer/training_loop.rs | 916 ++++++++++++++++++ 3 files changed, 1003 insertions(+), 68 deletions(-) diff --git a/crates/ml/examples/train_baseline_rl.rs b/crates/ml/examples/train_baseline_rl.rs index 862cb288f..504081d5e 100644 --- a/crates/ml/examples/train_baseline_rl.rs +++ b/crates/ml/examples/train_baseline_rl.rs @@ -690,6 +690,37 @@ fn run_training(args: &Args) -> Result> { None }; + // Build ensemble trainers ONCE (shared across folds) — upload data once each. + // Previously these were created inside the fold loop, re-uploading per fold. + let mut ensemble_trainers: Vec = Vec::new(); + if train_dqn && args.ensemble_top_k > 1 && args.hyperopt_params.is_some() { + let param_sets = load_top_k_params( + &args.hyperopt_params, + "dqn", + args.ensemble_top_k, + ); + // k=0 is the primary trainer, so build trainers for k=1.. + for (k, hp) in param_sets.iter().enumerate().skip(1) { + let ens_hyperparams = build_dqn_hyperparams(args, hp, total_cost_bps); + match DQNTrainer::new(ens_hyperparams) { + Ok(mut ens_trainer) => { + info!(" Uploading fxcache to ensemble trainer {}...", k); + if let Err(e) = rt.block_on(ens_trainer.init_from_fxcache( + &fxcache.features, &fxcache.targets, &fxcache.ofi, + )) { + error!(" [DQN] Ensemble trainer {} init_from_fxcache failed: {}", k, e); + continue; + } + ensemble_trainers.push(ens_trainer); + } + Err(e) => { + error!(" [DQN] Failed to create ensemble trainer {}: {}", k, e); + } + } + } + info!(" Created {} ensemble trainers (data uploaded once each)", ensemble_trainers.len()); + } + // 4. Train each fold info!("Step 4/5: Training models on each fold..."); let mut dqn_results: Vec<(usize, f64)> = Vec::new(); @@ -750,71 +781,60 @@ fn run_training(args: &Args) -> Result> { if args.ensemble_top_k > 1 && args.hyperopt_params.is_some() { // Ensemble mode: train one model per top-K hyperopt param set. - // NOTE: ensemble requires separate trainers with different hyperparams, - // so we create per-ensemble trainers (only the primary shares GPU data). - let param_sets = load_top_k_params( - &args.hyperopt_params, - "dqn", - args.ensemble_top_k, - ); - for (k, hp) in param_sets.iter().enumerate() { + // Trainers were created + data uploaded BEFORE the fold loop. + // Here we just set_training_range + reset_for_fold on each. + let total_members = 1 + ensemble_trainers.len(); // k=0 (primary) + secondaries + + // k=0: Primary ensemble member uses the shared trainer + { + let k = 0; info!( " [DQN] Training ensemble member {}/{} on fold {}", - k + 1, param_sets.len(), range.fold + k + 1, total_members, range.fold ); let prefix = format!("dqn_ensemble_{}", k); - if k == 0 { - // Primary ensemble member uses the shared trainer - match train_dqn_fold( - &rt, trainer, range.fold, - &train_norm, &val_norm, train_tgt, val_tgt, - range, &args.output_dir, &prefix, - ) { - Ok(best_loss) => { - info!(" [DQN] Ensemble member {} fold {} best_loss={:.6}", - k, range.fold, best_loss); - let elapsed = fold_start.elapsed().as_secs_f64(); - metrics::set_epoch("dqn", &fold_str, fold_idx as f64); - metrics::set_epoch_loss("dqn", &fold_str, best_loss); - metrics::set_validation_loss("dqn", &fold_str, best_loss); - metrics::set_iteration_seconds("dqn", &fold_str, elapsed); - dqn_results.push((range.fold, best_loss)); - } - Err(e) => { - error!(" [DQN] Ensemble member {} fold {} failed: {}", - k, range.fold, e); - } + match train_dqn_fold( + &rt, trainer, range.fold, + &train_norm, &val_norm, train_tgt, val_tgt, + range, &args.output_dir, &prefix, + ) { + Ok(best_loss) => { + info!(" [DQN] Ensemble member {} fold {} best_loss={:.6}", + k, range.fold, best_loss); + let elapsed = fold_start.elapsed().as_secs_f64(); + metrics::set_epoch("dqn", &fold_str, fold_idx as f64); + metrics::set_epoch_loss("dqn", &fold_str, best_loss); + metrics::set_validation_loss("dqn", &fold_str, best_loss); + metrics::set_iteration_seconds("dqn", &fold_str, elapsed); + dqn_results.push((range.fold, best_loss)); } - } else { - // Secondary ensemble members create separate trainers - let ens_hyperparams = build_dqn_hyperparams(args, hp, total_cost_bps); - match DQNTrainer::new(ens_hyperparams) { - Ok(mut ens_trainer) => { - // Upload fxcache to this ensemble trainer too - if let Err(e) = rt.block_on(ens_trainer.init_from_fxcache( - &fxcache.features, &fxcache.targets, &fxcache.ofi, - )) { - error!(" [DQN] Ensemble member {} init_from_fxcache failed: {}", k, e); - continue; - } - match train_dqn_fold( - &rt, &mut ens_trainer, range.fold, - &train_norm, &val_norm, train_tgt, val_tgt, - range, &args.output_dir, &prefix, - ) { - Ok(best_loss) => { - info!(" [DQN] Ensemble member {} fold {} best_loss={:.6}", - k, range.fold, best_loss); - } - Err(e) => { - error!(" [DQN] Ensemble member {} fold {} failed: {}", - k, range.fold, e); - } - } - } - Err(e) => { - error!(" [DQN] Failed to create ensemble trainer {}: {}", k, e); - } + Err(e) => { + error!(" [DQN] Ensemble member {} fold {} failed: {}", + k, range.fold, e); + } + } + } + + // k=1..: Secondary ensemble members reuse pre-created trainers + for (ens_idx, ens_trainer) in ensemble_trainers.iter_mut().enumerate() { + let k = ens_idx + 1; + info!( + " [DQN] Training ensemble member {}/{} on fold {}", + k + 1, total_members, range.fold + ); + let prefix = format!("dqn_ensemble_{}", k); + match train_dqn_fold( + &rt, ens_trainer, range.fold, + &train_norm, &val_norm, train_tgt, val_tgt, + range, &args.output_dir, &prefix, + ) { + Ok(best_loss) => { + info!(" [DQN] Ensemble member {} fold {} best_loss={:.6}", + k, range.fold, best_loss); + } + Err(e) => { + error!(" [DQN] Ensemble member {} fold {} failed: {}", + k, range.fold, e); } } } diff --git a/crates/ml/src/trainers/dqn/trainer/mod.rs b/crates/ml/src/trainers/dqn/trainer/mod.rs index 4ae740dfe..13e8463b4 100644 --- a/crates/ml/src/trainers/dqn/trainer/mod.rs +++ b/crates/ml/src/trainers/dqn/trainer/mod.rs @@ -548,20 +548,19 @@ impl DQNTrainer { info!("train_fold_from_slices: {} bars", features.len()); - // Build lightweight training_data with fixed-size [f64; 4] converted to Vec - // This is the last remaining allocation — each [f64; 4].to_vec() is 32 bytes. - // TODO: eliminate by changing train_with_data_full_loop signature - let training_data: Vec<(FeatureVector, Vec)> = features + // Build lightweight training_data with fixed-size arrays — zero heap alloc per element. + // Both [f64; 42] and [f64; 4] are Copy types that live entirely on the stack. + let training_data: Vec<([f64; 42], [f64; 4])> = features .iter() .zip(targets.iter()) - .map(|(f, t)| (*f, t.to_vec())) + .map(|(f, t)| (*f, *t)) .collect(); self.val_features_gpu = None; self.val_closes_gpu = None; self.val_ofi_gpu = None; - self.train_with_data_full_loop(&training_data, checkpoint_callback) + self.train_with_data_full_loop_slices(&training_data, checkpoint_callback) .await } @@ -1198,8 +1197,8 @@ impl DQNTrainer { .map_err(|e| anyhow::anyhow!("guard reset: {e}"))?; } - // Clear gpu_data to force re-init with current fold range - self.gpu_data = None; + // Keep gpu_data — the full dataset is already on GPU via init_from_fxcache. + // Clearing it would force a redundant re-upload on the next fold. info!("DQNTrainer: reset for new fold (GPU infrastructure preserved)"); Ok(()) diff --git a/crates/ml/src/trainers/dqn/trainer/training_loop.rs b/crates/ml/src/trainers/dqn/trainer/training_loop.rs index dc1d57f29..5ed818c0f 100644 --- a/crates/ml/src/trainers/dqn/trainer/training_loop.rs +++ b/crates/ml/src/trainers/dqn/trainer/training_loop.rs @@ -517,6 +517,419 @@ impl DQNTrainer { Ok(metrics) } + // ═══════════════════════════════════════════════════════════════════════ + // Zero-alloc training loop — accepts fixed-size arrays, no Vec + // ═══════════════════════════════════════════════════════════════════════ + + /// Like `train_with_data_full_loop` but accepts `&[([f64; 42], [f64; 4])]`. + /// + /// Requires `gpu_data` and `targets_raw_cuda` to already be populated + /// (via `init_from_fxcache`). Skips GPU upload entirely — the data is + /// already resident on GPU. This eliminates the per-bar `Vec` heap + /// allocation that the `(FeatureVector, Vec)` path required. + pub(crate) async fn train_with_data_full_loop_slices( + &mut self, + training_data: &[([f64; 42], [f64; 4])], + mut checkpoint_callback: F, + ) -> Result + where + F: FnMut(usize, Vec, bool) -> Result + Send, + { + let start_time = std::time::Instant::now(); + let mut total_loss = 0.0; + let mut total_q_value = 0.0; + let mut total_gradient_norm = 0.0; + let mut total_reward = 0.0; + let mut total_action_counts = [0_usize; 9]; + let mut total_factored_action_counts = [0_usize; 81]; + + self.log_training_config().await; + + // Decision Transformer pre-training on offline trajectories + if self.hyperparams.dt_pretrain_epochs > 0 { + info!( + epochs = self.hyperparams.dt_pretrain_epochs, + context_len = self.hyperparams.dt_context_len, + embed_dim = self.hyperparams.dt_embed_dim, + "Starting Decision Transformer pre-training" + ); + + // GPU raw buffers already uploaded via init_from_fxcache + if self.targets_raw_cuda.is_none() || self.features_raw_cuda.is_none() { + warn!("DT pre-training: GPU features/targets not available, skipping"); + } else if let Some(ref stream) = self.cuda_stream { + let stream = Arc::clone(stream); + + let dt_state_dim: usize = 42; + + let dt_config = crate::cuda_pipeline::decision_transformer::DecisionTransformerConfig { + state_dim: dt_state_dim, + num_actions: 9, + embed_dim: self.hyperparams.dt_embed_dim, + num_layers: self.hyperparams.dt_num_layers, + num_heads: 4, + context_len: self.hyperparams.dt_context_len, + dropout: 0.1, + batch_size: self.hyperparams.batch_size.min(256), + }; + + let mut dt = crate::cuda_pipeline::decision_transformer::DecisionTransformer::new( + stream, dt_config, + ).map_err(|e| anyhow::anyhow!("DT init: {e}"))?; + + let num_bars = training_data.len(); + if let (Some(ref features_gpu), Some(ref targets_gpu)) = + (&self.features_raw_cuda, &self.targets_raw_cuda) + { + let (trajectories, target_actions, num_batches) = dt + .build_dt_trajectories( + features_gpu, + targets_gpu, + num_bars, + self.hyperparams.gamma, + ) + .map_err(|e| anyhow::anyhow!("DT trajectory build: {e}"))?; + + if num_batches == 0 { + warn!( + num_bars, + context_len = self.hyperparams.dt_context_len, + batch_size = dt.config().batch_size, + "DT pre-training: not enough data for a full batch, skipping" + ); + } else { + let batch_size = dt.config().batch_size; + let context_len = dt.config().context_len; + let input_dim = dt.config().state_dim + 2; + + for epoch in 0..self.hyperparams.dt_pretrain_epochs { + let mut epoch_loss = 0.0_f32; + let mut batch_count = 0_usize; + + for batch_idx in 0..num_batches { + let ep_offset = batch_idx * batch_size; + let traj_elem_offset = ep_offset * context_len * input_dim; + let act_elem_offset = ep_offset * context_len; + + let loss = dt.pretrain_step( + &trajectories, + &target_actions, + batch_size, + traj_elem_offset, + act_elem_offset, + ).map_err(|e| anyhow::anyhow!("DT pretrain_step: {e}"))?; + + epoch_loss += loss; + batch_count += 1; + } + + let avg_loss = if batch_count > 0 { + epoch_loss / batch_count as f32 + } else { + 0.0 + }; + + info!( + epoch = epoch + 1, + total_epochs = self.hyperparams.dt_pretrain_epochs, + avg_loss = format!("{avg_loss:.4}"), + batches = batch_count, + "DT pre-training" + ); + + training_metrics::set_epoch( + "dt_pretrain", + "loss", + avg_loss as f64, + ); + } + + info!( + epochs = self.hyperparams.dt_pretrain_epochs, + "DT pre-training complete" + ); + } + } else { + warn!("DT pre-training: GPU features/targets not available, skipping"); + } + } else { + warn!("DT pre-training: no CUDA stream available, skipping"); + } + } + + // Training loop + for epoch in 0..self.hyperparams.epochs { + self.current_epoch = epoch; + self.reset_epoch_state(epoch); + + // C51 warmup ramp + { + let alpha = if self.hyperparams.c51_warmup_epochs > 0 { + (epoch as f32 / self.hyperparams.c51_warmup_epochs as f32).min(1.0) + } else { + 1.0 + }; + if let Some(ref mut fused) = self.fused_ctx { + fused.set_c51_alpha(alpha); + } + } + + // Q-gap warmup + if self.hyperparams.q_gap_threshold > 0.0 { + let warmup_epochs = 5.0_f32; + let ramp = (epoch as f32 / warmup_epochs).min(1.0); + let ramped_threshold = self.hyperparams.q_gap_threshold as f32 * ramp; + if let Some(ref mut selector) = self.gpu_action_selector { + selector.set_q_gap_threshold(ramped_threshold); + } + } + + log_epoch_start(epoch + 1, self.hyperparams.epochs, self.hyperparams.learning_rate); + training_metrics::set_epoch("dqn", "current", (epoch + 1) as f64); + + let mut monitor = TrainingMonitor::new(epoch + 1); + + self.portfolio_tracker.reset_drawdown_tracking(); + + let epoch_start = std::time::Instant::now(); + + // ── Phase 1: GPU data already uploaded via init_from_fxcache ── + let phase1_start = std::time::Instant::now(); + // Assert gpu_data + raw buffers are pre-populated (no re-upload) + if self.gpu_data.is_none() { + return Err(anyhow::anyhow!( + "train_with_data_full_loop_slices requires gpu_data (call init_from_fxcache first)" + )); + } + if self.targets_raw_cuda.is_none() { + return Err(anyhow::anyhow!( + "train_with_data_full_loop_slices requires targets_raw_cuda (call init_from_fxcache first)" + )); + } + self.init_gpu_experience_collector().await?; + let phase1_ms = phase1_start.elapsed().as_secs_f64() * 1000.0; + + // Causal feature masking + { + let frac = self.hyperparams.feature_mask_fraction; + self.epoch_feature_mask = None; + let _ = frac; // suppress unused warning + } + + // Vol normalization + self.epoch_vol_normalizer = if training_data.len() > 20 { + let returns: Vec = training_data.iter() + .map(|(fv, _)| fv[3]) + .collect(); + let n = returns.len() as f64; + let mean = returns.iter().sum::() / n; + let var = returns.iter().map(|r| (r - mean).powi(2)).sum::() / (n - 1.0); + var.sqrt().max(1e-8) as f32 + } else { + 0.0 + }; + + // ── Phase 2: GPU experience collection ── + eprintln!("[DEBUG] Phase 2: starting GPU experience collection (epoch {})", epoch); + let phase2_start = std::time::Instant::now(); + let gpu_experiences_collected = self.collect_gpu_experiences_slices( + training_data, + ).await?; + let phase2_ms = phase2_start.elapsed().as_secs_f64() * 1000.0; + eprintln!("[DEBUG] Phase 2: done in {:.1}ms, collected={}", phase2_ms, gpu_experiences_collected); + + if !gpu_experiences_collected { + return Err(anyhow::anyhow!( + "GPU experience collector MUST be active for CUDA training. \ + No CPU fallback path exists in CUDA builds." + )); + } + + // Periodic shrink-and-perturb + let sp_interval = self.hyperparams.shrink_perturb_interval; + if sp_interval > 0 && epoch > 0 && epoch % sp_interval == 0 { + if let Some(ref mut fused) = self.fused_ctx { + let alpha = self.hyperparams.shrink_perturb_alpha as f32; + let sigma = self.hyperparams.shrink_perturb_sigma as f32; + match fused.shrink_and_perturb(alpha, sigma) { + Ok(()) => info!(epoch, alpha, sigma, "Periodic shrink-and-perturb applied"), + Err(e) => warn!(epoch, "Shrink-and-perturb failed (non-fatal): {e}"), + } + } + } + + // ── Phase 3: Batched training from replay buffer ── + eprintln!("[DEBUG] Phase 3: starting training steps (epoch {})", epoch); + let phase3_start = std::time::Instant::now(); + let train_step_count = self.run_training_steps_slices(training_data.len()).await?; + let phase3_ms = phase3_start.elapsed().as_secs_f64() * 1000.0; + eprintln!("[DEBUG] Phase 3: done in {:.1}ms, steps={}", phase3_ms, train_step_count); + + // ── Phase 4: Epoch boundary + validation ── + let phase4_start = std::time::Instant::now(); + let boundary = if train_step_count > 0 { + let b = self.process_epoch_boundary(epoch, train_step_count, &mut monitor).await?; + Some(b) + } else { + None + }; + + self.sync_gpu_weights().await?; + + // Three-Phase Training Pipeline + let total_epochs = self.hyperparams.epochs; + let phase3_start_epoch = total_epochs * 80 / 100; + let in_phase3 = epoch >= phase3_start_epoch && total_epochs > 4; + + if in_phase3 { + if let Some(ref mut c) = self.gpu_experience_collector { + c.set_expert_ratio(0.0); + } + if epoch == phase3_start_epoch { + if let Some(ref mut fused) = self.fused_ctx { + let alpha = 0.9_f32; + let sigma = 0.01_f32; + if let Err(e) = fused.shrink_and_perturb(alpha, sigma) { + tracing::warn!("Shrink-and-Perturb failed at Phase 3 start (non-fatal): {e}"); + } else { + tracing::info!( + epoch, + phase3_start = phase3_start_epoch, + alpha, + sigma, + "Phase 3 (Refinement): Shrink-and-Perturb applied at phase boundary" + ); + } + } + } + } + + if !in_phase3 && self.hyperparams.epochs > 4 { + let interval = (self.hyperparams.epochs / 4).max(1); + if epoch > 0 && epoch % interval == 0 { + if let Some(ref mut fused) = self.fused_ctx { + let alpha = 0.9_f32; + let sigma = 0.01_f32; + if let Err(e) = fused.shrink_and_perturb(alpha, sigma) { + tracing::warn!("Shrink-and-Perturb failed (non-fatal): {e}"); + } else { + tracing::info!( + epoch, + alpha, + sigma, + "Shrink-and-Perturb applied (plasticity maintenance)" + ); + } + } + } + } + + // CVaR scales from IQN head + if let Some(ref mut fused) = self.fused_ctx { + let bs = fused.batch_size(); + let actions_clone = { + let actions = fused.actions_buf(); + actions.clone() + }; + let cvar_ptr = fused.compute_cvar_device_ptr(&actions_clone, bs, 0.05) + .unwrap_or(0); + if let Some(ref mut collector) = self.gpu_experience_collector { + collector.set_cvar_scales(cvar_ptr); + } + } + + // Flush GPU-accumulated max priority + if train_step_count > 0 { + let agent = self.agent.read().await; + if let Err(e) = agent.flush_max_priority() { + debug!("GPU max_priority flush failed (non-fatal): {}", e); + } + } + let phase4_ms = phase4_start.elapsed().as_secs_f64() * 1000.0; + + let epoch_duration = epoch_start.elapsed(); + + info!( + "Epoch {}/{} phase breakdown: init={:.0}ms experience={:.0}ms training={:.0}ms validation={:.0}ms total={:.0}ms", + epoch + 1, self.hyperparams.epochs, + phase1_ms, phase2_ms, phase3_ms, phase4_ms, + epoch_duration.as_secs_f64() * 1000.0, + ); + + let log_output = self.log_epoch_metrics_and_financials( + epoch, + train_step_count, + &boundary, + &mut monitor, + epoch_duration, + &mut total_action_counts, + &mut total_factored_action_counts, + ).await?; + + total_loss += log_output.avg_loss; + total_q_value += log_output.avg_q_value; + total_gradient_norm += log_output.avg_grad_norm; + + let epoch_avg_reward = if !monitor.reward_history.is_empty() { + monitor.reward_history.iter().sum::() / monitor.reward_history.len() as f32 + } else { + 0.0 + }; + total_reward += epoch_avg_reward as f64; + + self.handle_epoch_checkpoints_and_early_stopping( + epoch, + train_step_count, + &log_output, + &mut checkpoint_callback, + ).await?; + + if (epoch + 1) % self.hyperparams.checkpoint_frequency == 0 { + info!( + "Saving periodic checkpoint at epoch {}/{}", + epoch + 1, self.hyperparams.epochs + ); + let checkpoint_data = self.serialize_model().await?; + let checkpoint_size = checkpoint_data.len(); + let checkpoint_path = checkpoint_callback(epoch + 1, checkpoint_data, false) + .context("Failed to save periodic checkpoint")?; + info!( + "Periodic checkpoint saved: {} ({} bytes)", + checkpoint_path, checkpoint_size + ); + } + } + + let training_duration = start_time.elapsed(); + + let metrics = self.create_final_metrics( + total_loss, total_q_value, total_gradient_norm, total_reward, + self.hyperparams.epochs, training_duration, false, + total_action_counts, total_factored_action_counts, + ).await?; + + { + let mut stored_metrics = self.metrics.write().await; + *stored_metrics = metrics.clone(); + } + + info!( + "Training completed in {:.2}s: final_loss={:.6}, avg_q_value={:.4}", + training_duration.as_secs_f64(), + metrics.loss, + metrics.additional_metrics.get("avg_q_value").unwrap_or(&0.0) + ); + + info!("Best model summary:"); + info!( + " Best Sharpe: {:.4} at epoch {} (val_loss={:.6})", + self.best_sharpe, self.best_epoch, self.best_val_loss + ); + info!(" Best model checkpoint: best_model.safetensors"); + + Ok(metrics) + } + // ═══════════════════════════════════════════════════════════════════════ // Helper: one-time training configuration logging // ═══════════════════════════════════════════════════════════════════════ @@ -1314,6 +1727,329 @@ impl DQNTrainer { Ok(true) } + // ═══════════════════════════════════════════════════════════════════════ + // Helper: Phase 3 — GPU experience collection (zero-alloc slice variant) + // ═══════════════════════════════════════════════════════════════════════ + + /// Like `collect_gpu_experiences` but accepts `&[([f64; 42], [f64; 4])]`. + /// Only `.len()` and `.0` (features) are accessed — the `[f64; 4]` targets + /// are never read here (GPU raw buffers are already uploaded). + pub(crate) async fn collect_gpu_experiences_slices( + &mut self, + training_data: &[([f64; 42], [f64; 4])], + ) -> Result { + // Feature normalization stats calculation epoch (kept for API compat) + let _stats_collection_epochs = { + let ratio_based = (self.hyperparams.epochs as f32 * self.hyperparams.feature_stats_collection_ratio) as usize; + let capped = match self.hyperparams.max_feature_stats_epochs { + Some(max_epochs) => ratio_based.min(max_epochs), + None => ratio_based, + }; + capped.max(1) + }; + + // Set expert demo ratio BEFORE borrowing collector (avoids borrow conflict) + if self.hyperparams.expert_demo_ratio > 0.0 { + use crate::trainers::dqn::expert_demos::ExpertDemoGenerator; + let effective = ExpertDemoGenerator::effective_ratio( + self.hyperparams.expert_demo_ratio, + self.hyperparams.expert_demo_decay_epochs, + self.current_epoch, + ) as f32; + if let Some(ref mut c) = self.gpu_experience_collector { + c.set_expert_ratio(effective); + } + } + + 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, + ) else { + return Ok(false); + }; + + 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, + }; + self.cached_n_episodes = Some(computed); + computed + }; + + 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); + + collector.generate_episode_starts_gpu( + n_episodes as usize, stride, usable_bars, + ).map_err(|e| anyhow::anyhow!("GPU episode starts: {e}"))?; + collector.generate_sim_params_gpu( + n_episodes as usize, + self.hyperparams.transaction_cost_multiplier as f32, + self.hyperparams.avg_spread as f32, + self.hyperparams.fill_ioc_fill_prob as f32, + self.hyperparams.fill_limit_fill_min as f32, + self.hyperparams.fill_limit_fill_max as f32, + ).map_err(|e| anyhow::anyhow!("GPU sim params: {e}"))?; + + let mut episode_starts: Vec = (0..n_episodes) + .map(|i| { + let base = (i * stride).rem_euclid(usable_bars); + base + }) + .collect(); + + // Curriculum learning: filter episode starts by difficulty + let curriculum = self.hyperparams.curriculum_phase; + if curriculum < 2 && !training_data.is_empty() { + let adx_idx = 40; + let adx_threshold = 30.0_f64; + let original = episode_starts.clone(); + episode_starts.clear(); + + for &start in &original { + let bar = (start as usize).min(training_data.len().saturating_sub(1)); + let features = &training_data[bar].0; + let adx = if features.len() > adx_idx { features[adx_idx] } else { 0.0 }; + + match curriculum { + 0 => { + if adx > adx_threshold { + episode_starts.push(start); + } + } + 1 => { + episode_starts.push(start); + if adx > adx_threshold { + episode_starts.push(start); + } + } + _ => episode_starts.push(start), + } + } + + if episode_starts.len() < 16 { + tracing::warn!( + curriculum, + filtered = episode_starts.len(), + original = original.len(), + "Curriculum phase filtered too aggressively, falling back to Full" + ); + episode_starts = original; + } else { + tracing::debug!( + curriculum, + filtered = episode_starts.len(), + original = original.len(), + "Curriculum learning: episode starts filtered by ADX difficulty" + ); + } + } + + // 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, + ) { + return Err(anyhow::anyhow!("GPU episode reset FAILED (no CPU fallback): {e}")); + } + + let agent = self.agent.read().await; + let epsilon = agent.get_epsilon(); + drop(agent); + + let adversarial = self.adversarial_active; + if adversarial { + info!("Adversarial regime ACTIVE this epoch: 3x spread, 2x tx_cost, 0.5x fill"); + } + + let config = ExperienceCollectorConfig { + n_episodes, + timesteps_per_episode: timesteps, + total_bars, + epsilon, + gamma: self.hyperparams.gamma as f32, + max_position: self.max_position as f32, + enable_action_masking: self.enable_action_masking, + curiosity_scale: 1.0, + loss_aversion: self.hyperparams.loss_aversion as f32, + tx_cost_multiplier: { + let base = self.hyperparams.transaction_cost_multiplier as f32; + if adversarial { base * 2.0 } else { base } + }, + count_bonus_coefficient: self.hyperparams.count_bonus_coefficient + .unwrap_or(0.0) as f32, + q_clip_min: self.hyperparams.q_clip_min as f32, + q_clip_max: self.hyperparams.q_clip_max as f32, + huber_kappa: if true { + self.hyperparams.huber_delta as f32 + } else { + 0.0 + }, + use_noisy_nets: true, + noisy_sigma_init: self.hyperparams.noisy_sigma_init as f32, + use_distributional: true, + num_atoms: self.hyperparams.num_atoms as i32, + v_min: self.hyperparams.v_min as f32, + v_max: self.hyperparams.v_max as f32, + fill_median_spread: { + let base = self.hyperparams.avg_spread as f32; + if adversarial { base * 3.0 } else { base } + }, + fill_median_vol: self.median_vol as f32, + fill_ioc_fill_prob: { + let base = self.hyperparams.fill_ioc_fill_prob as f32; + if adversarial { base * 0.5 } else { base } + }, + fill_limit_fill_min: self.hyperparams.fill_limit_fill_min as f32, + fill_limit_fill_max: self.hyperparams.fill_limit_fill_max as f32, + fill_spread_cost_frac: self.hyperparams.fill_spread_cost_frac as f32, + fill_spread_capture_frac: self.hyperparams.fill_spread_capture_frac as f32, + fill_simulation_enabled: self.median_vol > 0.0, + q_gap_threshold: { + let target = self.hyperparams.q_gap_threshold as f32; + let epoch = self.current_epoch.min(10) as f32; + target * (epoch / 10.0) + }, + dsr_eta: self.hyperparams.dsr_eta as f32, + n_steps: self.hyperparams.n_steps as i32, + min_hold_bars: self.hyperparams.min_hold_bars as i32, + spread_cost: { + let base = (self.hyperparams.tick_size * self.hyperparams.contract_multiplier + * self.hyperparams.fill_spread_cost_frac) as f32; + if adversarial { base * 3.0 } else { base } + }, + contract_multiplier: self.hyperparams.contract_multiplier as f32, + margin_pct: self.hyperparams.margin_pct as f32, + dd_threshold: self.hyperparams.dd_threshold as f32, + w_dd: self.hyperparams.w_dd as f32, + beta_penalty: self.hyperparams.beta_penalty_strength as f32, + time_reversal_mod: self.hyperparams.time_reversal_mod as i32, + regret_blend: self.hyperparams.regret_blend as f32, + mirror_active: self.current_epoch % 2 == 1, + position_entropy_weight: self.hyperparams.position_entropy_weight as f32, + trade_clustering_penalty: self.hyperparams.trade_clustering_penalty as f32, + feature_noise_scale: self.hyperparams.feature_noise_scale as f32, + vol_normalizer: self.epoch_vol_normalizer, + feature_mask: self.epoch_feature_mask.clone(), + ..Default::default() + }; + + let gpu_batch = collector.collect_experiences_gpu( + features_buf, targets_buf, &episode_starts, &config, + ).map_err(|e| anyhow::anyhow!( + "GPU zero-roundtrip collection FAILED (no CPU fallback): {e}" + ))?; + + let count = gpu_batch.n_episodes * gpu_batch.timesteps; + + if let Some(ref stream) = self.cuda_stream { + self.experience_done_event = Some(stream.record_event(None) + .map_err(|e| { eprintln!("!!! EVENT RECORD FAILED: {e}"); anyhow::anyhow!("experience event record: {e}") })?); + } + + if let Some(ref mut mon) = self.gpu_monitoring { + if let Err(e) = mon.reduce(collector.rewards_gpu(), collector.actions_gpu(), count) { + debug!("GPU monitoring reduce failed: {e}"); + } + if let Some(ref stream) = self.cuda_stream { + if stream.synchronize().is_err() { + warn!("GPU monitoring kernel crashed — disabling monitoring"); + self.gpu_monitoring = None; + } + } + } + + if count > 0 { + if let Err(e) = collector.train_curiosity_gpu( + gpu_batch.n_episodes, gpu_batch.timesteps, + ) { + debug!("GPU curiosity training failed (non-fatal): {e}"); + } + if let Some(ref stream) = self.cuda_stream { + match stream.record_event(None) { + Ok(event) => { + if !event.is_complete() { + if let Err(e) = event.synchronize() { + warn!("Curiosity kernel async error detected: {e} — disabling curiosity for this run"); + self.curiosity_module = None; + collector.disable_curiosity(); + } + } + } + Err(e) => { + warn!("Curiosity event record failed: {e} — disabling curiosity for this run"); + self.curiosity_module = None; + collector.disable_curiosity(); + } + } + } + } + + if count > 0 { + let total = gpu_batch.n_episodes * gpu_batch.timesteps; + + let actions_u32: &CudaSlice = unsafe { + &*(&gpu_batch.actions as *const CudaSlice as *const CudaSlice) + }; + + let agent = self.agent.read().await; + let mut gpu_buf = agent.memory().as_gpu_buffer() + .ok_or_else(|| anyhow::anyhow!("GPU PER buffer required"))?; + gpu_buf.gpu.insert_batch_bf16( + &gpu_batch.states, + &gpu_batch.next_states, + actions_u32, + &gpu_batch.rewards, + &gpu_batch.dones, + total, + ).map_err(|e| anyhow::anyhow!("GPU PER insert_batch: {e}"))?; + } + + match collector.read_epoch_dsr_state() { + Ok((dsr_a, dsr_b)) => { + self.reward_fn.sync_dsr_from_gpu(dsr_a as f64, dsr_b as f64); + debug!( + "DSR GPU->CPU sync: ema_return={:.6}, ema_return_sq={:.6}", + dsr_a, dsr_b + ); + } + Err(e) => { + debug!("DSR epoch_state readback failed (non-fatal): {e}"); + } + } + + Ok(true) + } + // ═══════════════════════════════════════════════════════════════════════ // Helper: Batched training steps with guard kernel // ═══════════════════════════════════════════════════════════════════════ @@ -1507,6 +2243,186 @@ impl DQNTrainer { Ok(train_step_count) } + /// Like `run_training_steps` but accepts `num_bars` directly instead of + /// extracting `.len()` from a `&[(FeatureVector, Vec)]` slice. + /// This avoids requiring a `Vec` allocation per bar. + pub(crate) async fn run_training_steps_slices( + &mut self, + num_bars: usize, + ) -> Result { + let batch_size = self.hyperparams.batch_size; + let num_training_steps = if self.can_train().await? { + (num_bars / batch_size).max(1) + } else { + 0 + }; + + let mut train_step_count = 0; + + // Lazy-init fused CUDA training context. + if num_training_steps > 0 { + let needs_init = match &self.fused_ctx { + None => self.device.is_cuda(), + Some(ctx) => ctx.batch_size() != self.current_batch_size, + }; + if needs_init { + if let Some(ref stream) = self.cuda_stream { + if self.fused_ctx.is_some() { + tracing::info!("Fused CUDA context: batch_size changed, recreating"); + self.fused_ctx = None; + } + let agent = self.agent.read().await; + match super::super::fused_training::FusedTrainingCtx::new( + &self.device, &*agent, &self.hyperparams, self.current_batch_size, std::sync::Arc::clone(stream), + ) { + Ok(ctx) => { self.fused_ctx = Some(ctx); } + Err(e) => { tracing::error!("Fused CUDA context init failed: {e}"); } + } + } + } + } + + // Lazy-init + reset guard accumulators for this epoch + { + if self.training_guard.is_none() && self.device.is_cuda() { + match crate::cuda_pipeline::gpu_training_guard::GpuTrainingGuard::from_stream(Arc::clone(self.cuda_stream.as_ref().ok_or_else(|| anyhow::anyhow!("CUDA stream for training guard"))?)) { + Ok(g) => { + info!("GPU training guard initialized (epoch loop)"); + self.training_guard = Some(g); + } + Err(e) => return Err(anyhow::anyhow!("GPU training guard init: {e}")), + } + } + if let Some(ref mut guard) = self.training_guard { + guard.reset_accumulators() + .map_err(|e| anyhow::anyhow!("guard reset: {e}"))?; + } + } + + let guard_collapse_thresh = + self.hyperparams.learning_rate as f32 + * self.hyperparams.gradient_collapse_multiplier as f32; + let guard_past_warmup = { + let ws = (self.collapse_warmup_buffer_size as f64 * 0.2) as u64; + self.gradient_logging_step as u64 > ws + }; + + let mut sample_total_us = 0_u64; + let mut fused_total_us = 0_u64; + let mut guard_total_us = 0_u64; + + // Wait for experience collection GPU kernels to complete before sampling. + if let Some(event) = self.experience_done_event.take() { + if !event.is_complete() { + event.synchronize() + .map_err(|e| anyhow::anyhow!("experience event wait: {e}"))?; + } + } + + for _step in 0..num_training_steps { + let sample_start = std::time::Instant::now(); + let (batch, vaccine_batch) = { + let agent = self.agent.read().await; + let buffer = agent.memory(); + if !buffer.can_sample(self.current_batch_size) { + break; + } + let b = buffer.sample(self.current_batch_size) + .map_err(|e| anyhow::anyhow!("PER sample: {e}"))?; + let vb = buffer.sample(self.current_batch_size).ok(); + (b, vb) + }; + sample_total_us += sample_start.elapsed().as_micros() as u64; + + { + let mut agent = self.agent.write().await; + + let fused_start = std::time::Instant::now(); + let _gpu_result = if let Some(ref mut fused) = self.fused_ctx { + if let Some(vb) = vaccine_batch { + fused.pending_vaccine_batch = vb.gpu_batch; + } + let result = fused.run_full_step(&batch, &mut *agent, &self.device) + .map_err(|e| { eprintln!("!!! FUSED STEP ERROR: {:#}", e); e }) + .context("Fused CUDA training step failed")?; + result + } else { + unreachable!("Fused CUDA training is the only production path") + }; + + fused_total_us += fused_start.elapsed().as_micros() as u64; + + let guard_start = std::time::Instant::now(); + if let Some(ref mut guard) = self.training_guard { + let (loss_raw, grad_raw) = if let Some(ref fused) = self.fused_ctx { + (fused.loss_gpu_buf().raw_ptr(), fused.grad_norm_gpu_buf().raw_ptr()) + } else { + let ls = _gpu_result.loss_cuda_slice() + .map_err(|e| anyhow::anyhow!("guard loss: {e}"))?; + let gs = _gpu_result.grad_norm_cuda_slice() + .map_err(|e| anyhow::anyhow!("guard grad: {e}"))?; + (ls.raw_ptr(), gs.raw_ptr()) + }; + let gr = guard.check_and_accumulate( + loss_raw, + grad_raw, + 1e6_f32, + guard_collapse_thresh, + !guard_past_warmup, + ).map_err(|e| anyhow::anyhow!("guard check: {e}"))?; + if gr.halt_nan { + return Err(anyhow::anyhow!( + "NaN/Inf at step {}: loss={}, grad={}", + train_step_count, gr.raw_loss, gr.raw_grad_norm + )); + } + if gr.halt_grad_collapse { + agent.check_gradient_collapse(gr.raw_grad_norm).map_err(|e| { + tracing::info!("Early stopping (gradient collapse): {}", e); + anyhow::anyhow!("Early stopping: {}", e) + })?; + } + } + guard_total_us += guard_start.elapsed().as_micros() as u64; + + if train_step_count % 50 == 0 { + if let Some(ref mut fused) = self.fused_ctx { + if let Ok(stats) = fused.reduce_current_q_stats() { + self.cached_avg_q = stats.q_mean as f64; + if stats.q_min < self.epoch_q_min { self.epoch_q_min = stats.q_min; } + if stats.q_max > self.epoch_q_max { self.epoch_q_max = stats.q_max; } + } + } + } + train_step_count += 1; + self.gradient_logging_step += 1; + } + } + + if let Some(ref mut fused) = self.fused_ctx { + let _ = fused.flush_readback(); + if let Ok(stats) = fused.flush_q_stats_readback() { + self.cached_avg_q = stats.q_mean as f64; + if stats.q_min < self.epoch_q_min { self.epoch_q_min = stats.q_min; } + if stats.q_max > self.epoch_q_max { self.epoch_q_max = stats.q_max; } + } + } + + if train_step_count > 0 { + info!( + "Training step breakdown ({} steps): sample={:.0}ms fused={:.0}ms guard={:.0}ms (per-step: sample={:.1}ms fused={:.1}ms guard={:.1}ms)", + train_step_count, + sample_total_us as f64 / 1000.0, + fused_total_us as f64 / 1000.0, + guard_total_us as f64 / 1000.0, + sample_total_us as f64 / 1000.0 / train_step_count as f64, + fused_total_us as f64 / 1000.0 / train_step_count as f64, + guard_total_us as f64 / 1000.0 / train_step_count as f64, + ); + } + Ok(train_step_count) + } + // ═══════════════════════════════════════════════════════════════════════ // Helper: Epoch boundary readback + safety checks // ═══════════════════════════════════════════════════════════════════════