diff --git a/config/training/dqn-production.toml b/config/training/dqn-production.toml index bea43c117..bd31dc404 100644 --- a/config/training/dqn-production.toml +++ b/config/training/dqn-production.toml @@ -8,7 +8,6 @@ [training] epochs = 200 -batch_size = 8192 # H100 80GB — matches config/gpu/h100.toml learning_rate = 0.0001 gamma = 0.99 weight_decay = 0.0001 diff --git a/crates/ml-dqn/src/replay_buffer_type.rs b/crates/ml-dqn/src/replay_buffer_type.rs index 7c3091244..66304bd9d 100644 --- a/crates/ml-dqn/src/replay_buffer_type.rs +++ b/crates/ml-dqn/src/replay_buffer_type.rs @@ -60,6 +60,8 @@ pub struct GpuBatch { pub indices_ptr: u64, /// [batch_size] i32 on GPU (episode IDs for HER) pub episode_ids_ptr: u64, + /// Actual number of samples in this batch (may differ from trainer's max batch_size) + pub batch_size: usize, } /// GPU-resident prioritized replay buffer wrapper. @@ -247,6 +249,7 @@ impl ReplayBufferType { weights_ptr: ptrs.weights_ptr, indices_ptr: ptrs.indices_ptr, episode_ids_ptr: ptrs.episode_ids_ptr, + batch_size, }), }) } diff --git a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs index 0530d34e2..a8a8fe9d2 100644 --- a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs +++ b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs @@ -3901,8 +3901,6 @@ impl GpuDqnTrainer { ) -> Result<(), MLError> { // Synchronize the stream before capture to ensure all pending work // (BF16 mirror sync, batch upload, adam_step memcpy) is complete. - self.stream.synchronize() - .map_err(|e| MLError::ModelError(format!("stream sync before capture: {e}")))?; // Disable event tracking during capture — cudarc's device_ptr records // CudaEvents which are DISALLOWED inside CUDA Graph capture. @@ -4229,7 +4227,7 @@ impl GpuDqnTrainer { &mut self, gpu_batch: &crate::dqn::replay_buffer_type::GpuBatch, ) -> Result<(), MLError> { - let b = self.config.batch_size; + let b = gpu_batch.batch_size; let f32_size = std::mem::size_of::(); diff --git a/crates/ml/src/cuda_pipeline/gpu_experience_collector.rs b/crates/ml/src/cuda_pipeline/gpu_experience_collector.rs index df02bdd30..055dd2656 100644 --- a/crates/ml/src/cuda_pipeline/gpu_experience_collector.rs +++ b/crates/ml/src/cuda_pipeline/gpu_experience_collector.rs @@ -2076,6 +2076,9 @@ impl GpuExperienceCollector { Ok((host[5], host[6])) } + /// Allocated episode buffer capacity. + pub fn alloc_episodes(&self) -> usize { self.alloc_episodes } + /// Get a reference to the CUDA stream used by this collector. pub fn stream(&self) -> &Arc { &self.stream diff --git a/crates/ml/src/trainers/dqn/smoke_tests/training_stability.rs b/crates/ml/src/trainers/dqn/smoke_tests/training_stability.rs index 6e34ecba2..90fbf7ced 100644 --- a/crates/ml/src/trainers/dqn/smoke_tests/training_stability.rs +++ b/crates/ml/src/trainers/dqn/smoke_tests/training_stability.rs @@ -418,3 +418,130 @@ fn test_gpu_collector_auto_initializes() -> anyhow::Result<()> { drop(rt); Ok(()) } + +/// Validate the zero-copy fxcache training path end-to-end. +/// +/// This is the EXACT code path that H100 production uses: +/// fxcache → init_from_fxcache (one GPU upload) → fold ranges from timestamps → +/// per fold: set_training_range + reset_for_fold + train_fold_from_slices. +/// +/// The old smoketest above uses train_with_preloaded_data → train_with_data_full_loop. +/// This test validates the NEW path: train_fold_from_slices → train_with_data_full_loop_slices. +#[test] +#[ignore] // Requires fxcache — run via: FOXHUNT_TEST_DATA=test_data/futures-baseline cargo test -p ml --lib -- test_fxcache_zero_copy_training --ignored --nocapture +fn test_fxcache_zero_copy_training() -> anyhow::Result<()> { + use crate::fxcache; + use crate::walk_forward::NormStats; + + // 1. Find fxcache directory: FOXHUNT_FEATURE_CACHE_DIR > sibling of FOXHUNT_TEST_DATA > workspace root + let cache_dir = std::env::var("FOXHUNT_FEATURE_CACHE_DIR") + .map(std::path::PathBuf::from) + .or_else(|_| { + std::env::var("FOXHUNT_TEST_DATA").map(|td| { + let p = std::path::PathBuf::from(&td); + p.parent().unwrap_or(p.as_path()).join("feature-cache") + }) + }) + .unwrap_or_else(|_| { + // Try workspace root (cargo test runs from crate dir) + let mut d = std::env::current_dir().unwrap(); + loop { + let candidate = d.join("test_data").join("feature-cache"); + if candidate.exists() { return candidate; } + if !d.pop() { break; } + } + std::path::PathBuf::from("test_data/feature-cache") + }); + assert!(cache_dir.exists(), "feature-cache dir not found at {:?} — run precompute_features first", cache_dir); + + let entries: Vec<_> = std::fs::read_dir(cache_dir)? + .filter_map(|e| e.ok()) + .filter(|e| e.path().extension().and_then(|s| s.to_str()) == Some("fxcache")) + .collect(); + assert!(!entries.is_empty(), "No .fxcache files in test_data/feature-cache/"); + + let cache_path = entries[0].path(); + let mut fxcache_data = fxcache::load_fxcache(&cache_path)?; + // 5K bars: validates code path, not training quality. 5K/64=78 steps, ~1s + let max_bars = 5_000.min(fxcache_data.bar_count); + fxcache_data.features.truncate(max_bars); + fxcache_data.targets.truncate(max_bars); + fxcache_data.ofi.truncate(max_bars); + fxcache_data.timestamps.truncate(max_bars); + fxcache_data.bar_count = max_bars; + eprintln!("[FXCACHE] Loaded {} bars (truncated to {}) from {:?}", max_bars, max_bars, cache_path); + + // 2. Single fold: train on first 80%, validate on last 20% + // (Walk-forward needs months of data; smoketest uses a small subset) + let n = fxcache_data.bar_count; + let train_end = (n * 80) / 100; + let fold_ranges = vec![crate::walk_forward::FoldRange { + fold: 0, + train_start: 0, + train_end, + val_start: train_end, + val_end: n, + difficulty_score: 0.0, + }]; + eprintln!("[FXCACHE] Single fold: train=0..{}, val={}..{}", train_end, train_end, n); + + // 3. Create DQN trainer ONCE (1 epoch for speed) + let mut params = smoke_params(); + params.epochs = 1; + let mut trainer = smoke_trainer_with(params)?; + + // 4. Upload fxcache to GPU ONCE + let rt = tokio::runtime::Builder::new_current_thread() + .enable_all() + .build()?; + rt.block_on(trainer.init_from_fxcache( + &fxcache_data.features, &fxcache_data.targets, &fxcache_data.ofi, + ))?; + eprintln!("[FXCACHE] GPU upload complete"); + + // 5. Train each fold using the production zero-copy path + let mut fold_losses = Vec::new(); + for range in &fold_ranges { + let train_feat = &fxcache_data.features[range.train_start..range.train_end]; + let val_feat = &fxcache_data.features[range.val_start..range.val_end]; + let train_targets = &fxcache_data.targets[range.train_start..range.train_end]; + let val_targets = &fxcache_data.targets[range.val_start..range.val_end]; + + 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); + + trainer.set_training_range( + range.train_start, range.train_end, + range.val_start, range.val_end, + ); + trainer.set_val_data_from_slices(&val_norm, val_targets, range.train_end - range.train_start); + rt.block_on(trainer.reset_for_fold())?; + + let fold = range.fold; + let result = rt.block_on(trainer.train_fold_from_slices( + &train_norm, train_targets, + |_epoch, _bytes, _is_best| Ok(String::new()), + )); + + match result { + Ok(metrics) => { + eprintln!("[FXCACHE] Fold {} complete: loss={:.4}, epochs={}", fold, metrics.loss, metrics.epochs_trained); + fold_losses.push(metrics.loss); + } + Err(e) => { + panic!("Fold {} FAILED on zero-copy path: {:#}", fold, e); + } + } + } + + // 6. Verify training produced valid results + assert!(!fold_losses.is_empty(), "No folds completed"); + for (i, &loss) in fold_losses.iter().enumerate() { + assert!(loss.is_finite(), "Fold {} loss is not finite: {}", i, loss); + assert!(loss > 0.0, "Fold {} loss is zero or negative: {}", i, loss); + } + + eprintln!("[FXCACHE] All {} folds passed. Losses: {:?}", fold_losses.len(), fold_losses); + Ok(()) +} diff --git a/crates/ml/src/trainers/dqn/trainer/training_loop.rs b/crates/ml/src/trainers/dqn/trainer/training_loop.rs index e9d429c25..832d13068 100644 --- a/crates/ml/src/trainers/dqn/trainer/training_loop.rs +++ b/crates/ml/src/trainers/dqn/trainer/training_loop.rs @@ -759,7 +759,7 @@ impl DQNTrainer { } // ── Phase 3: Batched training from replay buffer ── - eprintln!("[DEBUG] Phase 3: starting training steps (epoch {})", epoch); + eprintln!("[DEBUG] Phase 3: starting training steps (epoch {}, bars={}, batch={})", epoch, training_data.len(), self.hyperparams.batch_size); 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; @@ -1742,9 +1742,10 @@ impl DQNTrainer { 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. + // Cache n_episodes — auto-scaled from VRAM, capped at collector's alloc size. + let alloc_cap = collector.alloc_episodes() as i32; let n_episodes = if let Some(cached) = self.cached_n_episodes { - cached + cached.min(alloc_cap) } else { use ml_core::memory_optimization::detect_gpu_hardware; let computed = match detect_gpu_hardware() { @@ -1753,14 +1754,14 @@ impl DQNTrainer { aligned_sd, self.hyperparams.gpu_timesteps_per_episode, ); - let chosen = optimal.max(32).min(16384) as i32; + let chosen = optimal.max(32).min(16384).min(alloc_cap as usize) as i32; info!( - "GPU auto-scaled n_episodes: {} (SMs={}, VRAM={:.0}MB)", - chosen, hw.sm_count, hw.free_memory_mb + "GPU auto-scaled n_episodes: {} (SMs={}, VRAM={:.0}MB, alloc_cap={})", + chosen, hw.sm_count, hw.free_memory_mb, alloc_cap ); chosen } - Err(_) => 256_i32, + Err(_) => alloc_cap.min(256), }; self.cached_n_episodes = Some(computed); computed diff --git a/crates/ml/src/training_profile.rs b/crates/ml/src/training_profile.rs index 82d455986..4cfdc9174 100644 --- a/crates/ml/src/training_profile.rs +++ b/crates/ml/src/training_profile.rs @@ -670,8 +670,9 @@ impl DqnTrainingProfile { if let Some(v) = t.epochs { hp.epochs = v; } - // batch_size: NOT applied from training profile. - // GPU profile (config/gpu/*.toml) is the sole authority for batch_size. + if let Some(v) = t.batch_size { + hp.batch_size = v; + } if let Some(v) = t.learning_rate { hp.learning_rate = v; }