docs: GPU utilization optimization implementation plan

5 tasks: fix DQN batch clamp, add HyperoptStrategy, wire VRAM-aware concurrency, update binary, verify. TDD throughout. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
2026-02-28 21:24:51 +01:00
parent ada33cf181
commit 9bbb03dcf1
1 changed files with 482 additions and 0 deletions
--- a/docs/plans/2026-02-28-gpu-utilization-optimization-implementation.md
+++ b/docs/plans/2026-02-28-gpu-utilization-optimization-implementation.md
@@ -0,0 +1,482 @@
+# 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`:
+
+```rust
+#[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:
+```rust
+// 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):
+
+```rust
+// 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):
+```rust
+// 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**
+
+```bash
+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`:
+
+```rust
+#[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:
+
+```rust
+/// 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:
+```rust
+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**
+
+```bash
+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`:
+
+```rust
+/// 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:
+```rust
+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:
+```rust
+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:
+
+```rust
+// 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:
+
+```rust
+// 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):
+```rust
+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:
+
+```rust
+// 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**
+
+```bash
+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:
+
+```rust
+// 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**
+
+```bash
+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**
+
+```bash
+git push -u origin feature/gpu-utilization-optimization
+```