diff --git a/crates/ml/examples/hyperopt_baseline_rl.rs b/crates/ml/examples/hyperopt_baseline_rl.rs index 92527addd..1709d95d4 100644 --- a/crates/ml/examples/hyperopt_baseline_rl.rs +++ b/crates/ml/examples/hyperopt_baseline_rl.rs @@ -286,7 +286,7 @@ fn main() -> Result<()> { // Cap by VRAM budget — no point having more threads than concurrent GPU trials let budget = ml::hyperopt::HardwareBudget::detect(); - let vram_cap = budget.plan_hyperopt(300.0, 0.0002, 64.0, 4096.0, 20).max_concurrent_trials; + let vram_cap = budget.plan_hyperopt(300.0, 0.0002, 64.0, 4096.0).max_concurrent_trials; let parallel = parallel.min(vram_cap); info!("Parallel: {} threads (CPU: {}, VRAM cap: {})", parallel, cpus, vram_cap); diff --git a/crates/ml/src/hyperopt/optimizer.rs b/crates/ml/src/hyperopt/optimizer.rs index 206e781f2..2b0cd4858 100644 --- a/crates/ml/src/hyperopt/optimizer.rs +++ b/crates/ml/src/hyperopt/optimizer.rs @@ -582,13 +582,39 @@ impl ArgminOptimizer { M::Params::per_sample_mb(), 64.0, 4096.0, - self.n_particles, ); info!("HyperoptStrategy: {} concurrent trials, batch bounds [{}, {}], {:.0}MB/trial, {:.0}MB total", strategy.max_concurrent_trials, strategy.batch_size_bounds.0, strategy.batch_size_bounds.1, strategy.per_trial_memory_mb, strategy.total_reserved_mb); + // Auto-scale PSO particle count to match GPU concurrency. + // More particles = better exploration when we have the GPU budget. + // Fall back to configured n_particles on CPU or small GPUs. + let effective_particles = if strategy.max_concurrent_trials > self.n_particles { + let scaled = strategy.max_concurrent_trials; + info!("Auto-scaled PSO particles: {} -> {} (GPU can fit {} concurrent trials)", + self.n_particles, scaled, strategy.max_concurrent_trials); + scaled + } else { + self.n_particles + }; + + // Scale total trials proportionally -- more particles need more iterations + // to converge. Each PSO iteration evaluates all particles, so we want + // enough iterations (typically 3-5) for convergence. + let effective_max_trials = if effective_particles > self.n_particles { + let min_iterations = 3; + let scaled = (effective_particles * min_iterations).max(self.max_trials); + if scaled > self.max_trials { + info!("Auto-scaled max_trials: {} -> {} (3 iterations x {} particles)", + self.max_trials, scaled, effective_particles); + } + scaled + } else { + self.max_trials + }; + let bounds = M::Params::continuous_bounds_for(&budget); let n_params = bounds.len(); @@ -600,9 +626,9 @@ impl ArgminOptimizer { } info!("Configuration:"); - info!(" Max Trials: {}", self.max_trials); + info!(" Max Trials: {} (configured: {})", effective_max_trials, self.max_trials); info!(" Initial Samples: {}", self.n_initial); - info!(" Swarm Particles: {}", self.n_particles); + info!(" Swarm Particles: {} (configured: {})", effective_particles, self.n_particles); info!(" Parameters: {}", n_params); info!(" Max Iters/Restart: {}", self.max_iters_per_restart); info!(" Parallel Threads: {}", rayon_threads); @@ -683,7 +709,7 @@ impl ArgminOptimizer { } // Create trial budget observer BEFORE creating cost function - let observer = crate::hyperopt::TrialBudgetObserver::new(self.max_trials); + let observer = crate::hyperopt::TrialBudgetObserver::new(effective_max_trials); // Create parallel cost function — clones model per evaluation instead of holding mutex let cost_fn = ParallelObjectiveFunction { @@ -724,21 +750,21 @@ impl ArgminOptimizer { // Determine PSO iterations let trials_used = trials.len(); - let remaining_trials = self.max_trials.saturating_sub(trials_used); + let remaining_trials = effective_max_trials.saturating_sub(trials_used); let max_iters_by_budget = - ((remaining_trials as f64) / (self.n_particles as f64)).ceil() as usize; + ((remaining_trials as f64) / (effective_particles as f64)).ceil() as usize; let max_iters = max_iters_by_budget.min(self.max_iters_per_restart); info!( "PSO Budget: {} iterations ({} remaining trials / {} particles = {} max iters)", - max_iters, remaining_trials, self.n_particles, max_iters_by_budget + max_iters, remaining_trials, effective_particles, max_iters_by_budget ); if max_iters > 0 { let lower_bounds: Vec = bounds.iter().map(|(min, _)| *min).collect(); let upper_bounds: Vec = bounds.iter().map(|(_, max)| *max).collect(); - let solver = ParticleSwarm::new((lower_bounds, upper_bounds), self.n_particles); + let solver = ParticleSwarm::new((lower_bounds, upper_bounds), effective_particles); let res = Executor::new(cost_fn, solver) .configure(|state| state.max_iters(max_iters as u64)) @@ -757,8 +783,8 @@ impl ArgminOptimizer { }, Err(e) if e.to_string().contains("Trial budget exhausted") => { info!( - "PSO terminated: trial budget reached ({} trials)", - self.max_trials + "PSO terminated: trial budget reached ({} trials, configured {})", + effective_max_trials, self.max_trials ); info!(" Evaluations: {}", trial_counter.load(std::sync::atomic::Ordering::Relaxed)); }, diff --git a/crates/ml/src/hyperopt/traits.rs b/crates/ml/src/hyperopt/traits.rs index 9005e9eaa..6382da5c1 100644 --- a/crates/ml/src/hyperopt/traits.rs +++ b/crates/ml/src/hyperopt/traits.rs @@ -142,14 +142,12 @@ impl HardwareBudget { /// * `per_sample_mb` - VRAM per sample in a batch /// * `min_batch` - Minimum batch_size floor /// * `max_batch` - Maximum batch_size ceiling - /// * `swarm_size` - PSO swarm size (caps max concurrency) pub fn plan_hyperopt( &self, model_overhead_mb: f64, per_sample_mb: f64, min_batch: f64, max_batch: f64, - swarm_size: usize, ) -> HyperoptStrategy { if self.gpu_memory_mb == 0 { return HyperoptStrategy { @@ -170,7 +168,7 @@ impl HardwareBudget { 1 } else { let raw = (available_mb / per_trial_mb).floor() as usize; - raw.clamp(1, swarm_size) + raw.clamp(1, 128) // Hardware cap: avoid thread explosion }; let per_trial_budget = available_mb / max_concurrent as f64; @@ -786,17 +784,19 @@ mod tests { #[test] fn test_plan_hyperopt_tiny_model_high_concurrency() { let budget = HardwareBudget { gpu_memory_mb: 24576 }; // L4 24GB - let strategy = budget.plan_hyperopt(300.0, 0.0002, 64.0, 4096.0, 20); + // L4: 24576 * 0.80 = 19660MB available, per_trial = 300 + 64*0.0002 = ~300MB + // raw = 19660 / 300 = ~65, clamped to [1, 128] + let strategy = budget.plan_hyperopt(300.0, 0.0002, 64.0, 4096.0); assert!(strategy.max_concurrent_trials >= 10, "L4 should support 10+ concurrent DQN trials, got {}", strategy.max_concurrent_trials); - assert!(strategy.max_concurrent_trials <= 20, - "concurrency should be capped at swarm_size=20, got {}", strategy.max_concurrent_trials); + assert!(strategy.max_concurrent_trials <= 128, + "concurrency should be capped at hardware limit 128, got {}", strategy.max_concurrent_trials); } #[test] fn test_plan_hyperopt_large_model_low_concurrency() { let budget = HardwareBudget { gpu_memory_mb: 24576 }; // L4 24GB - let strategy = budget.plan_hyperopt(2000.0, 0.05, 64.0, 4096.0, 20); + let strategy = budget.plan_hyperopt(2000.0, 0.05, 64.0, 4096.0); assert!(strategy.max_concurrent_trials >= 1, "should always allow at least 1 trial"); assert!(strategy.max_concurrent_trials <= 9, @@ -806,7 +806,7 @@ mod tests { #[test] fn test_plan_hyperopt_cpu_only_sequential() { let budget = HardwareBudget::cpu_only(); - let strategy = budget.plan_hyperopt(300.0, 0.0002, 64.0, 4096.0, 20); + let strategy = budget.plan_hyperopt(300.0, 0.0002, 64.0, 4096.0); assert_eq!(strategy.max_concurrent_trials, 1, "CPU-only should fall back to sequential"); } @@ -814,8 +814,10 @@ mod tests { #[test] fn test_plan_hyperopt_h100_maxes_out() { let budget = HardwareBudget { gpu_memory_mb: 81920 }; // H100 80GB - let strategy = budget.plan_hyperopt(300.0, 0.0002, 64.0, 4096.0, 20); - assert_eq!(strategy.max_concurrent_trials, 20, - "H100 should max out at swarm_size=20 for DQN"); + // H100: 81920 * 0.80 = 65536MB available, per_trial ~300MB + // raw = 65536 / 300 = ~218, clamped to 128 + let strategy = budget.plan_hyperopt(300.0, 0.0002, 64.0, 4096.0); + assert_eq!(strategy.max_concurrent_trials, 128, + "H100 should max out at hardware cap 128 for DQN"); } }