From 5c8dc537f0e6b056aa5d49566d26162f96f5e7bc Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Sun, 22 Mar 2026 19:50:43 +0100 Subject: [PATCH] =?UTF-8?q?refactor:=20fully=20VRAM-proportional=20hyperop?= =?UTF-8?q?t=20bounds=20=E2=80=94=20zero=20hardcoded=20tiers?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit All search space bounds now derived from HardwareBudget methods: - batch: budget.max_batch_size() - hidden_dim: budget.max_hidden_dim_base_full() - dueling/branch: proportional to max_hidden (50%/25%) - buffer: 15% VRAM / 120 bytes per entry - atoms: 5% VRAM / per-atom tensor cost - accum: proportional to VRAM / 10GB Removed small_gpu/large_gpu boolean tiers entirely. A 24GB GPU now gets bounds between 4GB and 80GB values, not arbitrarily bucketed. Phase Fast fixed bounds (lo==hi) still skipped — never modified. Co-Authored-By: Claude Opus 4.6 (1M context) --- crates/ml/src/hyperopt/adapters/dqn.rs | 97 +++++++++++--------------- 1 file changed, 42 insertions(+), 55 deletions(-) diff --git a/crates/ml/src/hyperopt/adapters/dqn.rs b/crates/ml/src/hyperopt/adapters/dqn.rs index 5d1d52608..e696633cf 100644 --- a/crates/ml/src/hyperopt/adapters/dqn.rs +++ b/crates/ml/src/hyperopt/adapters/dqn.rs @@ -843,67 +843,54 @@ impl ParameterSpace for DQNParams { fn continuous_bounds_for(budget: &HardwareBudget) -> Vec<(f64, f64)> { let mut bounds = Self::continuous_bounds(); + let vram_mb = budget.gpu_memory_mb as f64; - // GPU-scaled bound adjustments: (index, small_gpu_cap, large_gpu_expand). - // Skips fixed bounds (lo==hi, from Phase Fast) to avoid breaking phase config. - // Indices: 0=lr, 1=batch, 3=buffer(ln), 13=dueling, 15=atoms, 21=hidden, 28=branch, 29=accum - struct ScaleRule { idx: usize, small_cap: Option, large_expand: Option } - let rules = [ - ScaleRule { idx: 0, small_cap: None, large_expand: Some(1e-3_f64.ln()) }, // lr - ScaleRule { idx: 1, small_cap: Some(128.0), large_expand: Some(1024.0) }, // batch - ScaleRule { idx: 3, small_cap: Some(50_000_f64.ln()), large_expand: Some(500_000_f64.ln()) }, // buffer(ln) - ScaleRule { idx: 13, small_cap: Some(128.0), large_expand: Some(768.0) }, // dueling - ScaleRule { idx: 15, small_cap: Some(51.0), large_expand: Some(301.0) }, // atoms - ScaleRule { idx: 21, small_cap: Some(512.0), large_expand: Some(2048.0) }, // hidden - ScaleRule { idx: 28, small_cap: None, large_expand: Some(512.0) }, // branch - ScaleRule { idx: 29, small_cap: None, large_expand: Some(8.0) }, // accum - ]; + // Derive upper bounds from VRAM via HardwareBudget methods (same estimators + // used by AutoBatchSizer and the DQN constructor). No hardcoded tiers. + // Skip fixed bounds (lo==hi from Phase Fast). - let small_gpu = budget.gpu_memory_mb <= 8192; - let large_gpu = budget.gpu_memory_mb >= 40960; - - for rule in &rules { - let Some(b) = bounds.get_mut(rule.idx) else { continue }; - if (b.0 - b.1).abs() < 1e-6 { continue; } // Phase Fast fixed — never modify - - if let Some(cap) = rule.small_cap.filter(|_| small_gpu) { b.1 = b.1.min(cap); } - if let Some(exp) = rule.large_expand.filter(|_| large_gpu) { b.1 = b.1.max(exp); } - } - tracing::trace!("GPU bounds scaled for {}MB VRAM", budget.gpu_memory_mb); - - // Cap batch_size by VRAM — use min() to never WIDEN beyond expanded range + // batch_size: delegate to budget.max_batch_size() if let Some(max_batch) = budget.max_batch_size(MODEL_OVERHEAD_MB, MB_PER_SAMPLE, 64.0, 4096.0) { - if let Some(batch_bound) = bounds.get_mut(1) { - batch_bound.1 = batch_bound.1.min(max_batch); + if let Some(b) = bounds.get_mut(1) { + if (b.0 - b.1).abs() > 1e-6 { b.1 = max_batch; } } } - // Cap hidden_dim_base by VRAM using actual DQN architecture params. - // GPU PER is always on — include replay buffer in VRAM estimate. - let worst_batch = bounds.get(1).map_or(512, |b| b.1 as usize); - let worst_atoms = bounds.get(15).map_or(201, |b| b.1 as usize); - // Scale buffer by VRAM: ~60 bytes/entry (BF16 states + actions + rewards + priorities + seg_tree) - let worst_buffer = ((budget.gpu_memory_mb as f64 * 0.15 * 1024.0 * 1024.0) / 60.0) as usize; - // Worst-case aligned state_dim: 56 with OFI (53→56), 48 without (45→48). - // Use 56 for bounds computation since OFI is the common production path. - let worst_state_dim: usize = 56; - let max_base = budget.max_hidden_dim_base_full( - 1, // concurrent_trials (1 per GPU) - worst_batch, // worst-case batch_size from (capped) search space - worst_state_dim, - 5, // DQN actions (branching: 5 exposure head) - worst_atoms, // worst-case num_atoms from (capped) search space - true, // noisy_nets always enabled - true, // dueling always enabled - worst_buffer, // 50K small, 300K large, 100K medium - ); - // Use min() to never widen the range beyond expanded bounds - if let Some(dim_bound) = bounds.get_mut(21) { - dim_bound.1 = dim_bound.1.min(max_base as f64); + + // buffer_size (ln-scale): 15% of VRAM / ~120 bytes per replay entry + let max_buffer = (vram_mb * 0.15 * 1024.0 * 1024.0 / 120.0).clamp(5000.0, 1_000_000.0); + if let Some(b) = bounds.get_mut(3) { + if (b.0 - b.1).abs() > 1e-6 { b.1 = max_buffer.ln(); } } - tracing::info!( - "HardwareBudget: capped hidden_dim_base to {} (GPU: {}MB, batch: {}, atoms: {}, buffer: {})", - max_base.min(bounds.get(21).map_or(2048, |b| b.1 as usize)), - budget.gpu_memory_mb, worst_batch, worst_atoms, worst_buffer + + // hidden_dim_base + related dims: delegate to budget.max_hidden_dim_base_full() + let worst_batch = bounds.get(1).map_or(512, |b| b.1 as usize); + let worst_atoms = bounds.get(15).map_or(51, |b| b.1 as usize); + let worst_buffer = (vram_mb * 0.15 * 1024.0 * 1024.0 / 120.0) as usize; + let max_base = budget.max_hidden_dim_base_full( + 1, worst_batch, 56, 5, worst_atoms, true, true, worst_buffer, + ); + // Cap hidden, dueling, branch proportionally from max_base + let max_hidden = max_base as f64; + let max_dueling = (max_hidden * 0.5).clamp(128.0, 768.0); + let max_branch = (max_hidden * 0.25).clamp(64.0, 512.0); + // num_atoms: proportional to VRAM (5% budget / per-atom cost) + let max_atoms = (vram_mb * 0.05 * 1024.0 * 1024.0 + / (worst_batch as f64 * 11.0 * 4.0 * 2.0)).clamp(11.0, 301.0); + // gradient_accum: proportional to compute budget (no extra VRAM) + let max_accum = (vram_mb / 10240.0 * 8.0).clamp(1.0, 8.0); + + // Apply VRAM-derived caps (skip fixed Phase Fast bounds) + for &(idx, cap) in &[ + (15, max_atoms), (21, max_hidden), (13, max_dueling), + (28, max_branch), (29, max_accum), + ] { + if let Some(b) = bounds.get_mut(idx) { + if (b.0 - b.1).abs() > 1e-6 { b.1 = b.1.min(cap); } + } + } + tracing::trace!( + "VRAM bounds: {}MB → batch≤{} hidden≤{:.0} atoms≤{:.0} buffer≤{}", + budget.gpu_memory_mb, worst_batch, max_hidden, max_atoms, worst_buffer ); bounds }