gitadmin/foxhunt
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat(hyperopt): auto-scale PSO particles to GPU capacity

Browse Source 
Remove artificial swarm_size cap from plan_hyperopt() that limited
concurrent trials to n_particles (default 20). Now concurrency is
purely VRAM-driven with a hardware cap of 128 threads.

optimize_parallel() auto-scales n_particles to match GPU budget:
- L4 24GB: ~65 concurrent DQN trials (was 20)
- H100 80GB: 128 concurrent DQN trials (was 20)
- CPU/small GPU: falls back to configured n_particles

max_trials scales proportionally to ensure 3+ PSO iterations
for convergence with larger swarms.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-02-28 22:00:30 +01:00
parent c0b47d6657
commit 5d3efbd2cc
3 changed files with 50 additions and 22 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
crates/ml/examples/hyperopt_baseline_rl.rs
View File

@@ -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);

46
crates/ml/src/hyperopt/optimizer.rs
View File

@@ -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<f64> = bounds.iter().map(|(min, _)| *min).collect();
let upper_bounds: Vec<f64> = 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));
},

24
crates/ml/src/hyperopt/traits.rs
View File

@@ -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");
}
}