From d63190b9b921f29ab60b99dbc7fd3946352cda54 Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Mon, 2 Mar 2026 11:42:38 +0100 Subject: [PATCH] =?UTF-8?q?feat(ml):=20GPU=20full=20saturation=20=E2=80=94?= =?UTF-8?q?=20remove=20artificial=20caps=20and=20enable=20VRAM-aware=20sca?= =?UTF-8?q?ling?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Phase 1: Remove artificial caps - TFT benchmark: VRAM-scaled batch sizes (4/16/32/64) replacing hardcoded max_batch=4 - Liquid CUDA: VRAM-aware config defaults (batch 256-2048, pool 10% VRAM) - DQN trainer: remove double-clamp between AutoBatchSizer and HardwareBudget Phase 2: Mixed precision in training - DQN agent: add BF16/FP16 dtype casting in forward_with_gradients and forward_without_gradients (training was bypassing forward_mixed) Phase 3: Reduce CPU round-trips - DQN trainer: flat buffer select_actions_batch (eliminate Vec>) - DQN trainer: early-skip experience extraction (avoid .to_vec() on invalid) - EpochPrefetcher: AtomicBool is_ready() so callers can detect completion Phase 4: Adaptive scaling - HardwareBudget: tiered safety factor (0.70-0.85 by GPU size) - AutoBatchSizer: VRAM-proportional batch_overhead_mb (1.5% instead of fixed 250MB) 2451 tests pass, 0 clippy warnings. Co-Authored-By: Claude Opus 4.6 --- crates/ml/src/benchmark/tft_benchmark.rs | 44 +++++++----- crates/ml/src/cuda_pipeline/prefetch.rs | 30 ++++++-- crates/ml/src/dqn/agent.rs | 45 +++++++++++- crates/ml/src/hyperopt/traits.rs | 23 +++--- crates/ml/src/liquid/cuda/mod.rs | 44 +++++++++++- .../memory_optimization/auto_batch_size.rs | 21 +++--- crates/ml/src/trainers/dqn/trainer.rs | 70 ++++++++----------- 7 files changed, 193 insertions(+), 84 deletions(-) diff --git a/crates/ml/src/benchmark/tft_benchmark.rs b/crates/ml/src/benchmark/tft_benchmark.rs index b4a3de853..20b6349f5 100644 --- a/crates/ml/src/benchmark/tft_benchmark.rs +++ b/crates/ml/src/benchmark/tft_benchmark.rs @@ -450,46 +450,56 @@ impl TftBenchmarkRunner { Ok(samples) } - /// Find optimal batch size (CONSTRAINED to max=4 for TFT!) + /// Find optimal batch size for TFT, scaled to GPU VRAM. /// - /// TFT is memory-intensive due to attention mechanism. - /// We enforce max_batch_size=4 and use gradient accumulation to simulate larger batches. + /// TFT is memory-intensive due to multi-head attention. + /// Batch size is scaled conservatively based on available VRAM, + /// with gradient accumulation to reach an effective batch of 64. fn find_optimal_batch_size(&self) -> Result { - info!("Finding optimal batch size (TFT constrained to max=4)..."); + // Scale TFT max batch based on GPU VRAM + let caps = crate::gpu::capabilities::cached_capabilities(); + let max_batch: usize = match caps.free_vram_mb as u64 { + 0..=4096 => 4, // 4GB: keep at 4 + 4097..=12288 => 16, // 12GB: allow 16 + 12289..=25600 => 32, // 24GB: allow 32 + _ => 64, // 48GB+: allow 64 + }; + info!("Finding optimal batch size (TFT max={} for {:.0}MB VRAM)...", max_batch, caps.free_vram_mb); let device = self.gpu_manager.device(); - // CRITICAL: Set max_batch_size=4 for TFT (not 256!) let finder = BatchSizeFinder::with_params( device.clone(), - 1, // min_batch: Start from 1 - 4, // max_batch: TFT constrained to 4 due to memory - 64, // target_effective_batch - 0.9, // safety_margin + 1, // min_batch + max_batch, // VRAM-scaled max + 64, // target_effective_batch + 0.9, // safety_margin ); // Test function for batch size finding let config = finder .find_optimal_batch_size(|batch_size| { - // Simple memory test - return true if batch size <= 4 - Ok(batch_size <= 4) + Ok(batch_size <= max_batch) }) .map_err(|e| anyhow::anyhow!("Failed to find optimal batch size: {}", e))?; - // Enforce gradient accumulation to simulate larger batches + // Adjust gradient accumulation to reach effective batch of 64 let mut config_with_accumulation = config; - if config_with_accumulation.gradient_accumulation_steps < 16 { - config_with_accumulation.gradient_accumulation_steps = 16; + let effective_target = 64_usize; + let needed_steps = (effective_target / config_with_accumulation.batch_size.max(1)).max(1); + if config_with_accumulation.gradient_accumulation_steps < needed_steps { + config_with_accumulation.gradient_accumulation_steps = needed_steps; info!( - "Increased gradient accumulation steps to 16 (effective batch_size={})", - config_with_accumulation.batch_size * 16 + "Gradient accumulation steps={} (effective batch_size={})", + needed_steps, + config_with_accumulation.batch_size * needed_steps ); } Ok(config_with_accumulation) } - /// Create TFT model config optimized for 4GB GPU + /// Create TFT model config fn create_tft_config(&self, batch_size: usize) -> Result { Ok(TFTConfig { // Model architecture (reduced for memory) diff --git a/crates/ml/src/cuda_pipeline/prefetch.rs b/crates/ml/src/cuda_pipeline/prefetch.rs index 0c657ccd2..f27ae83a6 100644 --- a/crates/ml/src/cuda_pipeline/prefetch.rs +++ b/crates/ml/src/cuda_pipeline/prefetch.rs @@ -4,7 +4,9 @@ //! walk-forward fold while the current fold trains on GPU. The result //! can be collected via `try_take()` (non-blocking) or `take()` (blocking). +use std::sync::atomic::{AtomicBool, Ordering}; use std::sync::mpsc; +use std::sync::Arc; use std::thread::{self, JoinHandle}; use tracing::{info, warn}; @@ -31,12 +33,14 @@ pub struct EpochPrefetcher { receiver: mpsc::Receiver>, #[allow(dead_code)] handle: Option>, + completed: Arc, } impl std::fmt::Debug for EpochPrefetcher { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { f.debug_struct("EpochPrefetcher") .field("handle", &self.handle.as_ref().map(|_| "...")) + .field("completed", &self.completed.load(Ordering::Acquire)) .finish() } } @@ -52,6 +56,8 @@ impl EpochPrefetcher { F: FnOnce() -> Result + Send + 'static, { let (tx, rx) = mpsc::channel(); + let completed = Arc::new(AtomicBool::new(false)); + let completed_clone = Arc::clone(&completed); let handle = thread::Builder::new() .name("epoch-prefetch".into()) .spawn(move || { @@ -64,6 +70,9 @@ impl EpochPrefetcher { ), Err(e) => warn!("EpochPrefetcher: load failed: {e}"), } + // Signal completion before sending so is_ready() becomes + // true even if the receiver has been dropped. + completed_clone.store(true, Ordering::Release); // Receiver may have been dropped if the trainer finished early drop(tx.send(result)); }) @@ -72,6 +81,7 @@ impl EpochPrefetcher { Self { receiver: rx, handle, + completed, } } @@ -92,12 +102,12 @@ impl EpochPrefetcher { } /// Whether the prefetch has completed (result is ready to take). + /// + /// Non-destructive: uses an atomic flag set by the background thread + /// after `load_fn()` returns, so this can be called repeatedly without + /// consuming the result. pub fn is_ready(&self) -> bool { - // Peek without consuming — unfortunately mpsc doesn't support peek, - // so we just check if try_recv succeeds. Since we can't put it back, - // this is a destructive check. Use try_take() instead for real use. - // This method is only for logging/diagnostics. - false // Conservative: caller should use try_take() + self.completed.load(Ordering::Acquire) } } @@ -144,6 +154,16 @@ mod tests { let _maybe = prefetcher.try_take(); } + #[test] + fn test_is_ready_returns_true_when_complete() { + let prefetcher = EpochPrefetcher::spawn(|| Ok(vec![])); + std::thread::sleep(std::time::Duration::from_millis(50)); + assert!( + prefetcher.is_ready(), + "is_ready should return true after task completes" + ); + } + #[test] fn test_prefetch_large_data() { let prefetcher = EpochPrefetcher::spawn(|| { diff --git a/crates/ml/src/dqn/agent.rs b/crates/ml/src/dqn/agent.rs index ec92989f1..10f76db02 100644 --- a/crates/ml/src/dqn/agent.rs +++ b/crates/ml/src/dqn/agent.rs @@ -423,6 +423,9 @@ impl DQNAgent { } /// Forward pass through network with gradient tracking + /// + /// Supports mixed precision: casts input to BF16/FP16 for compute, + /// casts output back to FP32 for loss calculation. fn forward_with_gradients( &self, input: &Tensor, @@ -430,6 +433,18 @@ impl DQNAgent { ) -> Result { use candle_nn::linear; + // Mixed precision: cast input to reduced precision for 2x compute speedup + let (x_input, use_amp) = match &self.config.mixed_precision { + Some(mp) if mp.enabled => { + let target_dtype = mp.dtype.to_dtype(); + match input.to_dtype(target_dtype) { + Ok(converted) => (converted, true), + Err(_) => (input.clone(), false), + } + } + _ => (input.clone(), false), + }; + let mut layers = Vec::new(); let mut input_dim = self.config.state_dim; @@ -451,7 +466,7 @@ impl DQNAgent { layers.push(output_layer); // Forward pass with ReLU activations - let mut x = input.clone(); + let mut x = x_input; let num_layers = layers.len(); for (i, layer) in layers.iter().enumerate() { x = layer.forward(&x)?; @@ -466,10 +481,18 @@ impl DQNAgent { } } + // Mixed precision: cast output back to FP32 for loss computation + if use_amp { + x = x.to_dtype(candle_core::DType::F32)?; + } + Ok(x) } /// Forward pass through network without gradient tracking (for target network) + /// + /// Supports mixed precision: casts input to BF16/FP16 for compute, + /// casts output back to FP32 for target Q-value calculation. fn forward_without_gradients( &self, input: &Tensor, @@ -477,6 +500,18 @@ impl DQNAgent { ) -> Result { use candle_nn::linear; + // Mixed precision: cast input to reduced precision + let (x_input, use_amp) = match &self.config.mixed_precision { + Some(mp) if mp.enabled => { + let target_dtype = mp.dtype.to_dtype(); + match input.to_dtype(target_dtype) { + Ok(converted) => (converted, true), + Err(_) => (input.clone(), false), + } + } + _ => (input.clone(), false), + }; + let mut layers = Vec::new(); let mut input_dim = self.config.state_dim; @@ -498,8 +533,7 @@ impl DQNAgent { layers.push(output_layer); // Forward pass with LeakyReLU activations (no dropout for target network) - // Bug #11 fix: LeakyReLU prevents dead neurons - let mut x = input.clone(); + let mut x = x_input; let num_layers = layers.len(); for (i, layer) in layers.iter().enumerate() { x = layer.forward(&x)?; @@ -510,6 +544,11 @@ impl DQNAgent { } } + // Mixed precision: cast output back to FP32 + if use_amp { + x = x.to_dtype(candle_core::DType::F32)?; + } + // Detach from gradient computation Ok(x.detach()) } diff --git a/crates/ml/src/hyperopt/traits.rs b/crates/ml/src/hyperopt/traits.rs index cabda391c..3cc249e10 100644 --- a/crates/ml/src/hyperopt/traits.rs +++ b/crates/ml/src/hyperopt/traits.rs @@ -116,7 +116,7 @@ impl HardwareBudget { /// Compute the max batch_size for a given model memory profile. /// - /// Formula: `(available_vram * 0.75 - model_overhead) / per_sample_mb` + /// Formula: `(available_vram * safety_factor - model_overhead) / per_sample_mb` /// Result is clamped to `[min_batch, max_batch]`. /// /// Returns `None` if GPU is not available (budget is 0). @@ -130,7 +130,14 @@ impl HardwareBudget { if self.gpu_memory_mb == 0 { return None; } - let available = self.gpu_memory_mb as f64 * 0.75 - model_overhead_mb; + // Adaptive safety factor: larger GPUs can safely use more VRAM + let safety_factor = match self.gpu_memory_mb { + 0..=8192 => 0.70, + 8193..=24576 => 0.75, + 24577..=49152 => 0.80, + _ => 0.85, + }; + let available = self.gpu_memory_mb as f64 * safety_factor - model_overhead_mb; if available <= 0.0 || per_sample_mb <= 0.0 { return Some(min_batch); } @@ -818,7 +825,7 @@ mod tests { let budget = HardwareBudget::with_memory_mb(4096); // 4GB like RTX 3050 Ti assert_eq!(budget.gpu_memory_mb, 4096); - // 4096 * 0.75 = 3072; 3072 - 150 overhead = 2922; 2922 / 0.05 = 58440 + // 4096 * 0.70 = 2867.2; 2867.2 - 150 overhead = 2717.2; 2717.2 / 0.05 = 54344 // clamped to max_batch=4096 let max = budget.max_batch_size(150.0, 0.05, 16.0, 4096.0).unwrap(); assert!((max - 4096.0).abs() < 0.01); @@ -827,7 +834,7 @@ mod tests { #[test] fn test_hardware_budget_h100_wide_range() { let budget = HardwareBudget::with_memory_mb(81920); // 80GB H100 - // 81920 * 0.75 = 61440; 61440 - 150 = 61290; 61290 / 0.05 = 1_225_800 + // 81920 * 0.85 = 69632; 69632 - 150 = 69482; 69482 / 0.05 = 1_389_640 // clamped to 4096 let max = budget.max_batch_size(150.0, 0.05, 16.0, 4096.0).unwrap(); assert!((max - 4096.0).abs() < 0.01); @@ -836,21 +843,21 @@ mod tests { #[test] fn test_hardware_budget_small_gpu() { let budget = HardwareBudget::with_memory_mb(2048); // 2GB GPU - // 2048 * 0.75 = 1536; 1536 - 150 = 1386; 1386 / 0.05 = 27720 + // 2048 * 0.70 = 1433.6; 1433.6 - 150 = 1283.6; 1283.6 / 0.05 = 25672 // clamped to 4096 let max = budget.max_batch_size(150.0, 0.05, 16.0, 4096.0).unwrap(); assert!((max - 4096.0).abs() < 0.01); // With much heavier model overhead - // 1536 - 1500 = 36; 36 / 0.05 = 720 + // 1433.6 - 1500 = -66.4 → available <= 0 → returns min_batch let max2 = budget.max_batch_size(1500.0, 0.05, 16.0, 4096.0).unwrap(); - assert!((max2 - 720.0).abs() < 0.01); + assert!((max2 - 16.0).abs() < 0.01); } #[test] fn test_hardware_budget_clamps_to_min() { let budget = HardwareBudget::with_memory_mb(256); // Tiny GPU - // 256 * 0.75 = 192; 192 - 200 = -8 → available <= 0 → returns min_batch + // 256 * 0.70 = 179.2; 179.2 - 200 = -20.8 → available <= 0 → returns min_batch let max = budget.max_batch_size(200.0, 0.05, 16.0, 4096.0).unwrap(); assert!((max - 16.0).abs() < 0.01); } diff --git a/crates/ml/src/liquid/cuda/mod.rs b/crates/ml/src/liquid/cuda/mod.rs index e1666a5f7..568ac163b 100644 --- a/crates/ml/src/liquid/cuda/mod.rs +++ b/crates/ml/src/liquid/cuda/mod.rs @@ -36,12 +36,23 @@ pub struct CudaLiquidConfig { impl Default for CudaLiquidConfig { fn default() -> Self { + let caps = crate::gpu::capabilities::cached_capabilities(); + // Scale batch size to GPU VRAM: + // ≤4GB: 256, 5-12GB: 512, 13-25GB: 1024, 26+GB: 2048 + let max_batch_size = match caps.free_vram_mb as u64 { + 0..=4096 => 256, + 4097..=12288 => 512, + 12289..=25600 => 1024, + _ => 2048, + }; + // Scale memory pool to GPU VRAM (10% of free VRAM, min 256MB) + let memory_pool_size_mb = ((caps.free_vram_mb * 0.10) as usize).max(256); Self { device_id: 0, - max_batch_size: 32, + max_batch_size, use_shared_memory: true, stream_count: 4, - memory_pool_size_mb: 256, + memory_pool_size_mb, enable_profiling: false, } } @@ -599,4 +610,31 @@ fn compile_liquid_kernels() -> Result { }) } -// TECHNICAL DEBT ELIMINATED - Use cudarc types directly \ No newline at end of file +// TECHNICAL DEBT ELIMINATED - Use cudarc types directly + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_cuda_liquid_config_default_scales_to_vram() { + let config = CudaLiquidConfig::default(); + // Should be at least 256 on any system (was hardcoded 32) + assert!( + config.max_batch_size >= 256, + "max_batch_size should be at least 256, got {}", + config.max_batch_size + ); + } + + #[test] + fn test_cuda_liquid_config_memory_pool_scales() { + let config = CudaLiquidConfig::default(); + // Memory pool should be at least 256MB (minimum floor) + assert!( + config.memory_pool_size_mb >= 256, + "memory_pool_size_mb should be at least 256, got {}", + config.memory_pool_size_mb + ); + } +} \ No newline at end of file diff --git a/crates/ml/src/memory_optimization/auto_batch_size.rs b/crates/ml/src/memory_optimization/auto_batch_size.rs index 2b46c4a3f..dcbe451f4 100644 --- a/crates/ml/src/memory_optimization/auto_batch_size.rs +++ b/crates/ml/src/memory_optimization/auto_batch_size.rs @@ -234,11 +234,14 @@ impl AutoBatchSizer { // Add batch-level overhead (calculated before error checking) // This accounts for intermediate buffers that don't scale linearly with batch size - let batch_overhead_mb = match config.model_precision { - ModelPrecision::FP32 => 250.0, // FP32: ~250MB per batch (attention, workspace) - ModelPrecision::INT8 => 75.0, // INT8: ~75MB per batch - ModelPrecision::QAT => 500.0, // QAT: ~500MB per batch (FP32 base + FakeQuantize intermediate tensors) + // Batch overhead scales with GPU VRAM (attention workspaces, cuDNN scratch) + // Base: ~1.5% of total VRAM for FP32, with per-precision multipliers + let base_overhead_pct = match config.model_precision { + ModelPrecision::FP32 => 0.015, + ModelPrecision::INT8 => 0.005, + ModelPrecision::QAT => 0.030, }; + let batch_overhead_mb = (self.total_memory_mb * base_overhead_pct).max(50.0); debug!( "Fixed overhead: Model={:.1}MB, Optimizer={:.1}MB, Gradients={:.1}MB, Activations={:.1}MB, Total={:.1}MB, Batch overhead={:.1}MB", @@ -344,11 +347,13 @@ impl AutoBatchSizer { }; let activation_mb = model_mb * activation_multiplier; let fixed_overhead_mb = model_mb + optimizer_mb + gradient_mb + activation_mb; - let batch_overhead_mb = match config.model_precision { - ModelPrecision::FP32 => 250.0, - ModelPrecision::INT8 => 75.0, - ModelPrecision::QAT => 500.0, + // Batch overhead scales with GPU VRAM (attention workspaces, cuDNN scratch) + let base_overhead_pct = match config.model_precision { + ModelPrecision::FP32 => 0.015, + ModelPrecision::INT8 => 0.005, + ModelPrecision::QAT => 0.030, }; + let batch_overhead_mb = (self.total_memory_mb * base_overhead_pct).max(50.0); let available = usable_memory_mb - fixed_overhead_mb - batch_overhead_mb; if available <= 0.0 { diff --git a/crates/ml/src/trainers/dqn/trainer.rs b/crates/ml/src/trainers/dqn/trainer.rs index d9faae8c8..1170dfdbe 100644 --- a/crates/ml/src/trainers/dqn/trainer.rs +++ b/crates/ml/src/trainers/dqn/trainer.rs @@ -333,23 +333,11 @@ impl DQNTrainer { } }; - // Scale UP when running with conservative defaults on large GPUs - let budget = crate::hyperopt::HardwareBudget::detect(); - let gpu_optimal = budget - .max_batch_size(model_overhead_mb, 0.0005, 64.0, 8192.0) - .unwrap_or(hyperparams.batch_size as f64) as usize; - if hyperparams.batch_size < gpu_optimal && hyperparams.batch_size <= 128 { - info!( - "DQN batch_size scaled UP: {} → {} (GPU: {})", - hyperparams.batch_size, gpu_optimal.min(max_safe_batch), budget.gpu_name - ); - hyperparams.batch_size = gpu_optimal.min(max_safe_batch); - } - - // Cap DOWN if still exceeding VRAM ceiling + // Cap to VRAM ceiling from AutoBatchSizer (no separate scale-UP — + // AutoBatchSizer already accounts for model size and available VRAM) if hyperparams.batch_size > max_safe_batch { info!( - "Batch size {} exceeds GPU VRAM ceiling ({}), clamping", + "DQN batch_size capped from {} → {} (VRAM ceiling)", hyperparams.batch_size, max_safe_batch ); hyperparams.batch_size = max_safe_batch; @@ -3038,18 +3026,21 @@ impl DQNTrainer { } let agent = self.agent.read().await; - - // Convert all states to vectors - let state_vecs: Vec> = states.iter().map(|s| s.to_vector()).collect(); - - // Validate all states have consistent dimensions (first state sets the dimension) let batch_size = states.len(); - if batch_size == 0 { - return Ok(Vec::new()); - } - let state_dim = state_vecs[0].len(); - for (i, vec) in state_vecs.iter().enumerate().skip(1) { + // Get state dimension from first state + let first_vec = states + .first() + .map(|s| s.to_vector()) + .ok_or_else(|| anyhow::anyhow!("Empty states slice"))?; + let state_dim = first_vec.len(); + + // Pre-allocate flat buffer and fill directly (no intermediate Vec>) + let mut flat_states = Vec::with_capacity(batch_size * state_dim); + flat_states.extend_from_slice(&first_vec); + + for (i, state) in states.iter().enumerate().skip(1) { + let vec = state.to_vector(); if vec.len() != state_dim { return Err(anyhow::anyhow!( "State {} dimension mismatch: expected {}, got {}", @@ -3058,13 +3049,11 @@ impl DQNTrainer { vec.len() )); } + flat_states.extend_from_slice(&vec); } - // Flatten all states into single tensor [batch_size, state_dim] - let batched_states: Vec = state_vecs.into_iter().flat_map(|v| v.into_iter()).collect(); - - // Create batched tensor - let batch_tensor = Tensor::from_vec(batched_states, (batch_size, state_dim), &self.device) + // Create batched tensor directly from flat buffer + let batch_tensor = Tensor::from_vec(flat_states, (batch_size, state_dim), &self.device) .map_err(|e| anyhow::anyhow!("Failed to create batched state tensor: {}", e))?; // WAVE 16S: Get volatility-adjusted epsilon for exploration @@ -3774,10 +3763,11 @@ fn gpu_batch_to_experiences( let idx = ep * batch.timesteps + t; let state_start = idx * state_dim; let state_end = state_start + state_dim; - let state: Vec = batch.states - .get(state_start..state_end) - .unwrap_or(&[]) - .to_vec(); + // Skip out-of-bounds entries early instead of creating empty vecs + let state_slice = match batch.states.get(state_start..state_end) { + Some(s) if s.len() == state_dim => s, + _ => continue, + }; // Clamp action to valid u8 range (0..45) let action = batch.actions.get(idx).copied().unwrap_or(0).clamp(0, 44) as u8; @@ -3789,15 +3779,15 @@ fn gpu_batch_to_experiences( let next_idx = ep * batch.timesteps + t + 1; let ns_start = next_idx * state_dim; let ns_end = ns_start + state_dim; - batch.states.get(ns_start..ns_end).unwrap_or(&[]).to_vec() + match batch.states.get(ns_start..ns_end) { + Some(s) if s.len() == state_dim => s.to_vec(), + _ => state_slice.to_vec(), // fallback to current state + } } else { - state.clone() + state_slice.to_vec() }; - // Skip if state is empty (shouldn't happen, but defensive) - if state.len() == state_dim && next_state.len() == state_dim { - experiences.push(Experience::new(state, action, reward, next_state, done)); - } + experiences.push(Experience::new(state_slice.to_vec(), action, reward, next_state, done)); } } experiences