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

refactor: fully VRAM-proportional hyperopt bounds — zero hardcoded tiers

Browse Source 
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) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-03-22 19:50:43 +01:00
parent f3b98f456a
commit 5c8dc537f0
1 changed files with 42 additions and 55 deletions
Download Patch File Download Diff File Expand all files Collapse all files

97
crates/ml/src/hyperopt/adapters/dqn.rs
View File

@@ -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<f64>, large_expand: Option<f64> }
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
}