foxhunt/docs/plans/2026-02-28-gpu-utilization-optimization-implementation.md

GPU Utilization Optimization Implementation Plan

For Claude: REQUIRED SUB-SKILL: Use superpowers:executing-plans to implement this plan task-by-task.

Goal: Maximize GPU utilization for DQN/PPO hyperopt by fixing batch_size bounds, adding VRAM-aware concurrency, and protecting against OOM.

Architecture: Three changes: (1) Fix DQN's hardcoded batch_size clamp so PSO's choices flow through, (2) Add HardwareBudget::plan_hyperopt() that computes optimal concurrent trials based on model size vs VRAM, (3) Wire concurrency into optimize_parallel() with 5-layer OOM protection.

Tech Stack: Rust, Candle (cudarc), argmin PSO, nvidia-smi

---

Task 1: Fix DQN from_continuous() Batch Size Clamp

The DQN adapter's from_continuous() hardcodes .min(160.0) at line 532, which throws away any wider bounds from continuous_bounds_for(). This is the most impactful single fix — without it, nothing else matters.

Files:

• Modify: crates/ml/src/hyperopt/adapters/dqn.rs:532
• Test: crates/ml/src/hyperopt/adapters/dqn.rs (existing tests in same file)

Step 1: Write the failing test

Add to the existing #[cfg(test)] mod tests block at the bottom of dqn.rs:

#[test]
fn test_batch_size_respects_wide_bounds() {
    // Simulate PSO choosing batch_size=2048 (within VRAM-aware bounds)
    let mut params = vec![0.0_f64; 42];
    params[1] = 2048.0; // batch_size (index 1)
    // Fill other required params with valid defaults
    params[0] = (1e-4_f64).ln(); // learning_rate
    params[2] = 0.99; // gamma
    params[3] = (100_000.0_f64).ln(); // buffer_size
    params[4] = 1.5; // hold_penalty_weight
    params[5] = 6.0; // max_position_absolute
    params[6] = (25.0_f64).ln(); // huber_delta
    params[7] = 0.01; // entropy_coefficient
    params[8] = 0.5; // transaction_cost_multiplier
    params[9] = 0.6; // per_alpha
    params[10] = 0.4; // per_beta_start
    // Rainbow extensions (indices 11-41) — use midpoint of bounds
    let bounds = DqnHyperparams::continuous_bounds();
    for i in 11..42 {
        params[i] = (bounds[i].0 + bounds[i].1) / 2.0;
    }
    let result = DqnHyperparams::from_continuous(&params).unwrap();
    assert!(result.batch_size >= 1024, "batch_size {} should be >= 1024 when PSO picks 2048", result.batch_size);
}

Step 2: Run test to verify it fails

Run: SQLX_OFFLINE=true cargo test -p ml --lib -- hyperopt::adapters::dqn::tests::test_batch_size_respects_wide_bounds -v Expected: FAIL — batch_size 160 should be >= 1024 because .min(160.0) clamps it.

Step 3: Fix the clamp

In crates/ml/src/hyperopt/adapters/dqn.rs, line 532, change:

// OLD:
let mut batch_size = x[1].round().max(64.0).min(160.0) as usize;

// NEW: Use bounds from continuous_bounds() — PSO controls batch_size
let batch_bounds = Self::continuous_bounds();
let mut batch_size = x[1].round().max(batch_bounds[1].0).min(batch_bounds[1].1) as usize;

Wait — from_continuous doesn't have access to the HardwareBudget-adjusted bounds. The bounds could be wider than the static ones. The correct fix is to just clamp to the static lower bound and let the upper bound be open (PSO already constrains within bounds):

// NEW: Only enforce floor. PSO + HardwareBudget control the upper bound.
let mut batch_size = x[1].round().max(64.0) as usize;

Also update continuous_bounds() static upper bound from 160 to 4096 (line 451):

// OLD:
(64.0, 160.0),   // 1: batch_size (linear, GPU constrained)
// NEW:
(64.0, 4096.0),  // 1: batch_size (linear, upper bound adjusted by HardwareBudget)

Step 4: Run test to verify it passes

Run: SQLX_OFFLINE=true cargo test -p ml --lib -- hyperopt::adapters::dqn::tests::test_batch_size_respects_wide_bounds -v Expected: PASS

Step 5: Run full DQN adapter test suite

Run: SQLX_OFFLINE=true cargo test -p ml --lib -- hyperopt::adapters::dqn -v Expected: All existing tests pass. Check that any test asserting batch_size <= 160 is updated.

Step 6: Commit

git add crates/ml/src/hyperopt/adapters/dqn.rs
git commit -m "fix(hyperopt): remove hardcoded batch_size=160 clamp from DQN adapter

PSO now controls batch_size within [64, 4096] range. HardwareBudget
adjusts upper bound based on detected VRAM. The .min(160.0) clamp was
silently overriding PSO's choices on any GPU larger than 3050 Ti."

---

Task 2: Add HyperoptStrategy and HardwareBudget::plan_hyperopt()

Files:

• Modify: crates/ml/src/hyperopt/traits.rs (where HardwareBudget lives)
• Modify: crates/ml/src/hyperopt/mod.rs (export HyperoptStrategy)

Step 1: Write the failing test

Add to the existing #[cfg(test)] mod tests block in traits.rs:

#[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);
    // DQN (300MB overhead, 0.0002 MB/sample) should get many concurrent trials
    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);
}

#[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);
    // TFT (2GB overhead, 0.05 MB/sample) should get fewer concurrent trials
    assert!(strategy.max_concurrent_trials >= 1,
        "should always allow at least 1 trial");
    assert!(strategy.max_concurrent_trials <= 9,
        "L4 should not fit 10+ TFT trials, got {}", strategy.max_concurrent_trials);
}

#[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);
    assert_eq!(strategy.max_concurrent_trials, 1,
        "CPU-only should fall back to sequential");
}

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

Step 2: Run tests to verify they fail

Run: SQLX_OFFLINE=true cargo test -p ml --lib -- hyperopt::traits::tests::test_plan_hyperopt -v Expected: FAIL — plan_hyperopt method doesn't exist.

Step 3: Implement HyperoptStrategy and plan_hyperopt()

In crates/ml/src/hyperopt/traits.rs, add the struct and method:

/// Strategy computed by HardwareBudget for optimal GPU utilization during hyperopt.
#[derive(Debug, Clone)]
pub struct HyperoptStrategy {
    /// Max number of trials to evaluate concurrently on the GPU.
    pub max_concurrent_trials: usize,
    /// VRAM-aware (lower, upper) batch_size bounds for PSO search.
    pub batch_size_bounds: (f64, f64),
    /// Estimated VRAM per trial in MB (for logging/monitoring).
    pub per_trial_memory_mb: f64,
    /// Total VRAM reserved for all concurrent trials (for watchdog).
    pub total_reserved_mb: f64,
}

impl HardwareBudget {
    /// Compute optimal hyperopt strategy based on model size vs available VRAM.
    ///
    /// # Arguments
    /// * `model_overhead_mb` - Fixed per-trial VRAM: model weights + optimizer + gradients + activations
    /// * `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 {
        // CPU-only: sequential, use static bounds
        if self.gpu_memory_mb == 0 {
            return HyperoptStrategy {
                max_concurrent_trials: 1,
                batch_size_bounds: (min_batch, max_batch),
                per_trial_memory_mb: model_overhead_mb,
                total_reserved_mb: model_overhead_mb,
            };
        }

        let safety_margin = 0.20; // Keep 20% free for CUDA allocator + fragmentation
        let available_mb = self.gpu_memory_mb as f64 * (1.0 - safety_margin);

        // Per-trial memory: model overhead + default batch (use min_batch for conservative estimate)
        let default_batch_mb = min_batch * per_sample_mb;
        let per_trial_mb = model_overhead_mb + default_batch_mb;

        // How many trials fit?
        let max_concurrent = if per_trial_mb <= 0.0 {
            1
        } else {
            let raw = (available_mb / per_trial_mb).floor() as usize;
            raw.clamp(1, swarm_size)
        };

        // Per-trial VRAM budget (divide available equally among concurrent trials)
        let per_trial_budget = available_mb / max_concurrent as f64;

        // Batch size upper bound: what fits in per_trial_budget after model overhead
        let batch_budget = per_trial_budget - model_overhead_mb;
        let batch_upper = if batch_budget > 0.0 && per_sample_mb > 0.0 {
            (batch_budget / per_sample_mb).clamp(min_batch, max_batch)
        } else {
            min_batch
        };

        let total_reserved = per_trial_mb * max_concurrent as f64;

        HyperoptStrategy {
            max_concurrent_trials: max_concurrent,
            batch_size_bounds: (min_batch, batch_upper),
            per_trial_memory_mb: per_trial_mb,
            total_reserved_mb: total_reserved,
        }
    }
}

Step 4: Export from mod.rs

In crates/ml/src/hyperopt/mod.rs, add to the pub use traits:: line:

pub use traits::{
    HardwareBudget, HyperoptStrategy, HyperparameterOptimizable, OptimizationResult,
    ParameterSpace, TrialResult,
};

Step 5: Run tests to verify they pass

Run: SQLX_OFFLINE=true cargo test -p ml --lib -- hyperopt::traits::tests::test_plan_hyperopt -v Expected: All 4 tests PASS.

Step 6: Commit

git add crates/ml/src/hyperopt/traits.rs crates/ml/src/hyperopt/mod.rs
git commit -m "feat(hyperopt): add HyperoptStrategy and HardwareBudget::plan_hyperopt()

Computes optimal (concurrency, batch_size_bounds) based on model memory
footprint vs detected GPU VRAM. Scales from CPU-only (sequential) through
H100 80GB (full swarm parallel). 20% safety margin for CUDA allocator."

---

Task 3: Wire Concurrency Into optimize_parallel() with VRAM Watchdog

The optimize_parallel() method already exists and uses std::thread::scope + model cloning. Currently it uses rayon thread count (CPU-based). We change it to use HyperoptStrategy.max_concurrent_trials (VRAM-based) and add the ramp-up + watchdog logic.

Files:

• Modify: crates/ml/src/hyperopt/optimizer.rs:563-774 (the optimize_parallel method)
• Modify: crates/ml/src/hyperopt/traits.rs (add ParameterSpace::model_overhead_mb() + per_sample_mb())

Step 1: Add model memory constants to ParameterSpace trait

In crates/ml/src/hyperopt/traits.rs, add default methods to ParameterSpace:

/// Estimated fixed VRAM overhead per model instance in MB.
/// Override for accurate VRAM-based concurrency planning.
fn model_overhead_mb() -> f64 { 500.0 } // Conservative default

/// Estimated VRAM per batch sample in MB.
fn per_sample_mb() -> f64 { 0.02 } // Conservative default

Step 2: Implement in DQN and PPO adapters

In dqn.rs ParameterSpace impl:

fn model_overhead_mb() -> f64 { 300.0 } // MODEL_OVERHEAD_MB constant
fn per_sample_mb() -> f64 { 0.0002 } // 54 floats × 4 bytes = 216 bytes

In ppo.rs ParameterSpace impl:

fn model_overhead_mb() -> f64 { 200.0 } // MODEL_OVERHEAD_MB constant
fn per_sample_mb() -> f64 { 0.015 }

Step 3: Wire plan_hyperopt() into optimize_parallel()

In optimizer.rs, at the start of optimize_parallel() (after line 577 where budget is detected), add:

// Compute VRAM-aware concurrency strategy
let strategy = budget.plan_hyperopt(
    M::Params::model_overhead_mb(),
    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);

// OOM Layer 2: Ramp-up — start at ceil(max/2), increase if no OOM
let initial_concurrency = ((strategy.max_concurrent_trials + 1) / 2).max(1);
let concurrency = std::sync::atomic::AtomicUsize::new(initial_concurrency);

Step 4: Replace rayon thread count with VRAM-based concurrency in LHS phase

Change the LHS parallel evaluation (lines 630-654) to use chunked concurrency:

// OOM Layer 2: Chunked LHS evaluation with ramp-up
let chunk_size = concurrency.load(std::sync::atomic::Ordering::Relaxed);
info!("LHS phase: evaluating {} samples in chunks of {} (ramp-up from {})",
    initial_vecs.len(), chunk_size, strategy.max_concurrent_trials);

for chunk in initial_vecs.chunks(chunk_size) {
    // OOM Layer 4: VRAM watchdog — check free memory before spawning
    if let Ok((_, free_mb, _)) = crate::memory_optimization::auto_batch_size::detect_gpu_memory() {
        let total_mb = budget.gpu_memory_mb as f64;
        if total_mb > 0.0 && free_mb / total_mb < 0.10 {
            warn!("VRAM watchdog: only {:.0}MB free ({:.0}%), reducing concurrency",
                free_mb, (free_mb / total_mb) * 100.0);
            let current = concurrency.load(std::sync::atomic::Ordering::Relaxed);
            concurrency.store((current / 2).max(1), std::sync::atomic::Ordering::Relaxed);
        }
    }

    let lhs_errors: Vec<Option<anyhow::Error>> = std::thread::scope(|scope| {
        let handles: Vec<_> = chunk.iter().map(|continuous_vec| {
            let mut model_clone = model.clone();
            let results = Arc::clone(&trial_results);
            let counter = Arc::clone(&trial_counter);
            let names = param_names.clone();
            let vec = continuous_vec.clone();
            scope.spawn(move || {
                Self::evaluate_point(&vec, &mut model_clone, &results, &counter, &names).err()
            })
        }).collect();
        handles.into_iter().map(|h| h.join().unwrap_or(None)).collect()
    });

    for err in lhs_errors.into_iter().flatten() {
        warn!("LHS evaluation error (continuing): {}", err);
    }

    // OOM Layer 2: If first chunk succeeded, ramp up to full concurrency
    let current = concurrency.load(std::sync::atomic::Ordering::Relaxed);
    if current < strategy.max_concurrent_trials {
        info!("Ramp-up: {} → {} concurrent trials (first chunk succeeded)",
            current, strategy.max_concurrent_trials);
        concurrency.store(strategy.max_concurrent_trials, std::sync::atomic::Ordering::Relaxed);
    }
}

Step 5: Log concurrency in PSO phase header

Update the PSO phase info log (around line 693):

let final_concurrency = concurrency.load(std::sync::atomic::Ordering::Relaxed);
info!("Execution mode: Parallel trials (clone-per-particle, {} VRAM-based concurrent, {} rayon threads)",
    final_concurrency, rayon_threads);

Step 6: Configure rayon pool to match VRAM concurrency

In hyperopt_baseline_rl.rs, the rayon pool is set to CPU thread count. We need to also cap it by VRAM concurrency. This is already handled because ParallelObjectiveFunction::cost() clones the model per call — argmin PSO naturally evaluates N particles per iteration, and rayon parallelizes them up to the thread pool size. By setting the rayon pool to strategy.max_concurrent_trials, we cap concurrent GPU allocations.

In optimize_parallel(), after computing strategy, set rayon::ThreadPoolBuilder if it hasn't been configured externally:

// Cap rayon threads to VRAM concurrency (if not already smaller)
let effective_threads = rayon_threads.min(strategy.max_concurrent_trials);
if effective_threads < rayon_threads {
    info!("Capping rayon threads from {} to {} (VRAM-limited)", rayon_threads, effective_threads);
}

Note: We can't reconfigure the global rayon pool mid-process. The concurrency limit for PSO is naturally handled by argmin evaluating particles — the ParallelObjectiveFunction already clones per-particle. The rayon pool should be configured at startup in the binary. Add a note in the log if VRAM concurrency < rayon threads.

Step 7: Run full hyperopt test suite

Run: SQLX_OFFLINE=true cargo test -p ml --lib -- hyperopt -v Expected: All 176+ tests pass. No regressions.

Step 8: Compile check workspace

Run: SQLX_OFFLINE=true cargo check --workspace Expected: Clean compile.

Step 9: Commit

git add crates/ml/src/hyperopt/optimizer.rs crates/ml/src/hyperopt/traits.rs \
       crates/ml/src/hyperopt/adapters/dqn.rs crates/ml/src/hyperopt/adapters/ppo.rs
git commit -m "feat(hyperopt): VRAM-aware parallel trial concurrency with OOM protection

optimize_parallel() now uses HardwareBudget::plan_hyperopt() to compute
max concurrent trials based on model VRAM footprint. 5-layer OOM protection:
conservative estimation, concurrency ramp-up, per-trial OOM catch,
VRAM watchdog, graceful degradation to sequential."

---

Task 4: Update Binary and Integration Verification

Files:

• Modify: crates/ml/examples/hyperopt_baseline_rl.rs (cap rayon pool by VRAM)

Step 1: Update hyperopt_baseline_rl.rs parallel auto-detect

The binary currently auto-detects parallel from CPU count. Add VRAM-aware capping:

// After line 281 (resolving parallel count), add:
// 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 parallel = parallel.min(vram_cap);
info!("Parallel: {} threads (CPU: {}, VRAM cap: {})", parallel, cpus, vram_cap);

Step 2: Compile check

Run: SQLX_OFFLINE=true cargo check --workspace Expected: Clean.

Step 3: Commit

git add crates/ml/examples/hyperopt_baseline_rl.rs
git commit -m "feat(hyperopt): cap parallel threads by VRAM budget in hyperopt binary

Auto-detect now considers both CPU count and GPU VRAM when choosing
concurrent trial count. Prevents OOM on smaller GPUs."

---

Task 5: Final Verification and Push

Step 1: Run full hyperopt test suite

Run: SQLX_OFFLINE=true cargo test -p ml --lib -- hyperopt 2>&1 | tail -10 Expected: 176+ tests pass, 0 failures.

Step 2: Run clippy

Run: SQLX_OFFLINE=true cargo clippy --workspace -- -D warnings 2>&1 | tail -5 Expected: 0 warnings.

Step 3: Push branch

git push -u origin feature/gpu-utilization-optimization