From 71f2fb1ec3a423cdb8067365b1e6f10dddad3b83 Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Sun, 1 Mar 2026 19:05:11 +0100 Subject: [PATCH] feat(ml): dynamic GPU batch sizing + tensor core alignment utility MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - DQN: Scale UP batch_size on large GPUs (HardwareBudget::detect), raise static cap 4096→8192 - PPO: Scale UP batch_size from conservative 64 when GPU supports more - Add align_to_tensor_cores() utility (round up to multiple of 8) - Hidden dim_base rounding already aligned (nearest 256 = multiples of 8) - Tests: 2422 pass, 0 failures Co-Authored-By: Claude Opus 4.6 --- crates/ml/src/cuda_pipeline/mod.rs | 19 ++++++++++++ crates/ml/src/hyperopt/traits.rs | 2 +- crates/ml/src/trainers/dqn/config.rs | 2 +- crates/ml/src/trainers/dqn/trainer.rs | 42 +++++++++++++++++++++++++-- crates/ml/src/trainers/ppo.rs | 36 +++++++++++++++++++++-- 5 files changed, 93 insertions(+), 8 deletions(-) diff --git a/crates/ml/src/cuda_pipeline/mod.rs b/crates/ml/src/cuda_pipeline/mod.rs index ddc4c1d07..2546d8b7a 100644 --- a/crates/ml/src/cuda_pipeline/mod.rs +++ b/crates/ml/src/cuda_pipeline/mod.rs @@ -20,6 +20,15 @@ pub mod gpu_ppo_collector; /// Maximum bytes allowed for a single GPU upload (2 GB safety limit). const MAX_UPLOAD_BYTES: usize = 2 * 1024 * 1024 * 1024; +/// Round dimension up to nearest multiple of 8 for tensor core alignment. +/// +/// Tensor cores on Ampere (A100, L4, L40S) and Hopper (H100) require +/// matrix dimensions divisible by 8 (FP16/BF16) for optimal throughput. +/// This pads dimensions like state_dim=54 → 56 to enable tensor core paths. +pub fn align_to_tensor_cores(dim: usize) -> usize { + (dim + 7) & !7 +} + /// Estimate VRAM usage in bytes for pre-uploaded f32 data. pub fn estimate_vram_bytes(num_elements: usize) -> usize { num_elements * std::mem::size_of::() @@ -327,6 +336,16 @@ mod tests { assert_eq!(state.dims(), &[1, 54]); } + #[test] + fn test_tensor_core_alignment() { + assert_eq!(align_to_tensor_cores(54), 56); + assert_eq!(align_to_tensor_cores(64), 64); // already aligned + assert_eq!(align_to_tensor_cores(51), 56); // pad up + assert_eq!(align_to_tensor_cores(1), 8); + assert_eq!(align_to_tensor_cores(8), 8); + assert_eq!(align_to_tensor_cores(0), 0); + } + #[test] fn test_ppo_gpu_data_upload_cpu() { let market_data: Vec> = (0..200) diff --git a/crates/ml/src/hyperopt/traits.rs b/crates/ml/src/hyperopt/traits.rs index 9b6f7e0b1..cabda391c 100644 --- a/crates/ml/src/hyperopt/traits.rs +++ b/crates/ml/src/hyperopt/traits.rs @@ -240,7 +240,7 @@ impl HardwareBudget { hi = mid - 1; } } - // Round down to nearest 256 + // Round down to nearest 256 (all values are already multiples of 8 for tensor cores) (lo / 256) * 256 } diff --git a/crates/ml/src/trainers/dqn/config.rs b/crates/ml/src/trainers/dqn/config.rs index ef7fa93de..b7eace412 100644 --- a/crates/ml/src/trainers/dqn/config.rs +++ b/crates/ml/src/trainers/dqn/config.rs @@ -343,7 +343,7 @@ impl DQNAgentType { pub struct DQNHyperparameters { /// Learning rate (typically 1e-4 to 1e-3) pub learning_rate: f64, - /// Batch size (must be ≤230 for RTX 3050 Ti 4GB) + /// Batch size (dynamically capped by AutoBatchSizer based on GPU VRAM) pub batch_size: usize, /// Discount factor (typically 0.95-0.99) pub gamma: f64, diff --git a/crates/ml/src/trainers/dqn/trainer.rs b/crates/ml/src/trainers/dqn/trainer.rs index a83983589..2a23eb08e 100644 --- a/crates/ml/src/trainers/dqn/trainer.rs +++ b/crates/ml/src/trainers/dqn/trainer.rs @@ -260,9 +260,9 @@ impl DQNTrainer { )); } - // VRAM safety cap: only clamp batch_size if it would OOM. - // The optimizer (PSO/hyperopt) controls batch_size — we just enforce the ceiling. - const STATIC_MAX_BATCH_SIZE: usize = 4096; + // Dynamic batch sizing: scale UP for larger GPUs, cap DOWN for smaller ones. + // Uses HardwareBudget for consistent sizing across DQN/PPO. + const STATIC_MAX_BATCH_SIZE: usize = 8192; let max_safe_batch = match AutoBatchSizer::new() { Ok(sizer) => { let config = BatchSizeConfig { @@ -286,6 +286,20 @@ 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(50.0, 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 if hyperparams.batch_size > max_safe_batch { info!( "Batch size {} exceeds GPU VRAM ceiling ({}), clamping", @@ -4110,4 +4124,26 @@ mod tests { reward_unclamped ); } + + #[test] + fn test_dynamic_batch_size_l4() { + // L4 has 24GB VRAM — HardwareBudget should allow batch_size >> 230 + let budget = crate::hyperopt::HardwareBudget { + gpu_memory_mb: 24_000, + gpu_name: "NVIDIA L4".to_string(), + }; + let batch = budget.max_batch_size(50.0, 0.0005, 64.0, 8192.0); + assert!(batch.unwrap_or(0.0) > 230.0, "L4 should support DQN batch > 230, got {:?}", batch); + } + + #[test] + fn test_dynamic_batch_size_h100() { + // H100 has 80GB VRAM — should hit the 8192 ceiling + let budget = crate::hyperopt::HardwareBudget { + gpu_memory_mb: 81_920, + gpu_name: "NVIDIA H100".to_string(), + }; + let batch = budget.max_batch_size(50.0, 0.0005, 64.0, 8192.0); + assert!((batch.unwrap_or(0.0) - 8192.0).abs() < 1.0, "H100 should hit 8192 ceiling, got {:?}", batch); + } } diff --git a/crates/ml/src/trainers/ppo.rs b/crates/ml/src/trainers/ppo.rs index f20768d84..84ceef19c 100644 --- a/crates/ml/src/trainers/ppo.rs +++ b/crates/ml/src/trainers/ppo.rs @@ -258,7 +258,7 @@ impl PpoTrainer { }); } - // Dynamic GPU validation: detect hardware and auto-shrink batch size if needed + // Dynamic GPU validation: scale UP for large GPUs, cap DOWN for small ones let (device, effective_batch_size) = if use_gpu { let caps = cached_capabilities(); let max_batch = resolve_batch_size( @@ -267,12 +267,31 @@ impl PpoTrainer { hyperparams.batch_size, ); let device = DeviceConfig::Auto.resolve().unwrap_or(Device::Cpu); + + // Scale UP when running with conservative defaults on large GPUs + let effective = if device.is_cuda() && hyperparams.batch_size <= 64 { + let budget = crate::hyperopt::HardwareBudget::detect(); + let gpu_optimal = budget + .max_batch_size(80.0, 0.0004, 64.0, 4096.0) + .unwrap_or(max_batch as f64) as usize; + let scaled = gpu_optimal.min(max_batch); + if scaled > hyperparams.batch_size { + info!( + "PPO batch_size scaled UP: {} → {} (GPU: {})", + hyperparams.batch_size, scaled, caps.device_name + ); + } + scaled + } else { + max_batch + }; + if device.is_cuda() { - info!("PPO using GPU: {} (batch_size: {})", caps.device_name, max_batch); + info!("PPO using GPU: {} (batch_size: {})", caps.device_name, effective); } else { warn!("GPU requested but unavailable, falling back to CPU"); } - (device, max_batch) + (device, effective) } else { (Device::Cpu, hyperparams.batch_size) }; @@ -1519,4 +1538,15 @@ mod tests { error_msg ); } + + #[test] + fn test_ppo_dynamic_batch_size_l4() { + // L4 has 24GB VRAM — HardwareBudget should allow PPO batch > 64 + let budget = crate::hyperopt::HardwareBudget { + gpu_memory_mb: 24_000, + gpu_name: "NVIDIA L4".to_string(), + }; + let batch = budget.max_batch_size(80.0, 0.0004, 64.0, 4096.0); + assert!(batch.unwrap_or(0.0) > 64.0, "L4 should support PPO batch > 64, got {:?}", batch); + } }