a5ea8713b6
Branching DQN (Tavakoli et al., 2018):
- 3 independent advantage heads (exposure=5, order=3, urgency=3) in CUDA kernel
- `use_branching` as hyperopt-tunable parameter (index 11 in 29D space)
- Branch weight pointers packed into KernelWeightPack struct
KernelWeightPack refactor:
- 87→38 CUDA kernel params by packing 48 weight pointers into 384-byte #[repr(C)] struct
- UNPACK_WEIGHT_PTRS macro in CUDA header for clean kernel-side access
- Single .arg(&weight_pack) replaces 48 individual .arg() calls
Hyperopt metrics in JSON output:
- Added `metrics: Option<serde_json::Value>` to TrialResult<P>
- `extract_metrics()` trait method with default None (DQN overrides)
- JSON output now includes per-trial backtest metrics (Sharpe, Sortino,
Calmar, Omega, drawdown, win_rate, trades) + top-level best_metrics
Prometheus backtest gauges (Rust-side, no CUDA):
- 8 new gauges: foxhunt_hyperopt_best_{sharpe,sortino,calmar,omega,
max_drawdown_pct,win_rate,total_trades,total_return_pct}
Dynamic episode scaling:
- Removed hardcoded 256 episode cap on small GPUs
- VRAM budget calculation in optimal_n_episodes() handles scaling naturally
- Removed stale .min(256) in PPO trainer
Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
1252 lines
48 KiB
Rust
1252 lines
48 KiB
Rust
//! Generic Traits for Hyperparameter Optimization
|
||
//!
|
||
//! This module provides a trait-based architecture for hyperparameter optimization
|
||
//! that works across all Foxhunt ML models (MAMBA-2, DQN, PPO, TFT). The design
|
||
//! follows these principles:
|
||
//!
|
||
//! - **Model-agnostic**: Models implement `HyperparameterOptimizable` trait
|
||
//! - **Type-safe**: Parameter spaces are strongly typed with compile-time validation
|
||
//! - **Composable**: Traits can be combined for complex optimization workflows
|
||
//! - **Efficient**: Zero-cost abstractions with no runtime overhead
|
||
//!
|
||
//! ## Architecture
|
||
//!
|
||
//! ```text
|
||
//! ┌─────────────────────────────────────────────────────────────┐
|
||
//! │ HyperparameterOptimizable │
|
||
//! │ ┌─────────────┐ ┌──────────────┐ ┌─────────────┐ │
|
||
//! │ │ MAMBA-2 │ │ DQN │ │ PPO │ │
|
||
//! │ └─────────────┘ └──────────────┘ └─────────────┘ │
|
||
//! └─────────────────────────────────────────────────────────────┘
|
||
//! ↓
|
||
//! ┌──────────────────┐
|
||
//! │ ArgminOptimizer │
|
||
//! │ (Generic impl) │
|
||
//! └──────────────────┘
|
||
//! ↓
|
||
//! ┌──────────────────┐
|
||
//! │ OptimizationResult│
|
||
//! │ (Trial history) │
|
||
//! └──────────────────┘
|
||
//! ```
|
||
//!
|
||
//! ## Usage Example
|
||
//!
|
||
//! ```rust,no_run
|
||
//! use ml::hyperopt::{ArgminOptimizer, HyperparameterOptimizable};
|
||
//! use ml::hyperopt::adapters::mamba2::{Mamba2Trainer, Mamba2Params};
|
||
//!
|
||
//! # async fn example() -> anyhow::Result<()> {
|
||
//! // Create model trainer
|
||
//! let trainer = Mamba2Trainer::new(
|
||
//! "test_data/ES_FUT_180d.parquet",
|
||
//! 50, // epochs
|
||
//! )?;
|
||
//!
|
||
//! // Run optimization
|
||
//! let optimizer = ArgminOptimizer::new(30, 5);
|
||
//! let result = optimizer.optimize(trainer)?;
|
||
//!
|
||
//! println!("Best params: {:?}", result.best_params);
|
||
//! println!("Best loss: {:.6}", result.best_objective);
|
||
//! # Ok(())
|
||
//! # }
|
||
//! ```
|
||
|
||
use serde::{Deserialize, Serialize};
|
||
use std::fmt::Debug;
|
||
use tracing::info;
|
||
|
||
use crate::MLError;
|
||
|
||
/// Hyperopt objective mode for two-phase optimization.
|
||
///
|
||
/// In two-phase optimization, Phase A uses `EpisodeReward` for fast convergence,
|
||
/// then Phase B switches to `Sharpe` for financial quality refinement.
|
||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||
#[derive(Default)]
|
||
pub enum ObjectiveMode {
|
||
/// Phase A: Optimize episode reward (fast convergence signal)
|
||
EpisodeReward,
|
||
/// Phase B: Optimize Sharpe ratio (financial quality)
|
||
#[default]
|
||
Sharpe,
|
||
}
|
||
|
||
|
||
/// 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,
|
||
}
|
||
|
||
/// Hardware budget passed to dynamic bounds computation.
|
||
///
|
||
/// Contains detected GPU memory so that `ParameterSpace` implementations
|
||
/// can scale batch_size upper bounds to match available VRAM, instead of
|
||
/// using hardcoded limits tuned for a specific GPU (e.g., RTX 3050 Ti 4GB).
|
||
#[derive(Debug, Clone)]
|
||
pub struct HardwareBudget {
|
||
/// Available GPU memory in MB (0 = CPU-only, use conservative defaults)
|
||
pub gpu_memory_mb: usize,
|
||
/// GPU device name (e.g. "NVIDIA L40S"), empty for CPU-only
|
||
pub gpu_name: String,
|
||
}
|
||
|
||
impl HardwareBudget {
|
||
/// Detect hardware budget from the current system.
|
||
///
|
||
/// Uses `nvidia-smi` via the existing `detect_gpu_memory()` utility.
|
||
/// Uses **free** VRAM (not total) so that display server, compositor,
|
||
/// and other GPU consumers are already subtracted.
|
||
/// Falls back to `cpu_only()` if detection fails (no GPU, driver error, etc.).
|
||
pub fn detect() -> Self {
|
||
if let Ok((total_mb, free_mb, name)) = ml_core::memory_optimization::auto_batch_size::detect_gpu_memory() {
|
||
// Use free VRAM — on desktop GPUs the display server can consume
|
||
// 300-1500 MB of the total. Using total_mb would overestimate
|
||
// available budget and cause OOM on small GPUs (e.g. RTX 3050 Ti 4GB).
|
||
let gpu_memory_mb = free_mb as usize;
|
||
info!(
|
||
"HardwareBudget: detected {} — {}MB total, {}MB free (using free as budget)",
|
||
name, total_mb as usize, gpu_memory_mb
|
||
);
|
||
Self { gpu_memory_mb, gpu_name: name }
|
||
} else {
|
||
info!("HardwareBudget: no GPU detected, using CPU-only defaults");
|
||
Self::cpu_only()
|
||
}
|
||
}
|
||
|
||
/// Fallback for CPU-only or detection failure.
|
||
pub fn cpu_only() -> Self {
|
||
Self { gpu_memory_mb: 0, gpu_name: String::new() }
|
||
}
|
||
|
||
/// Create with a specific memory value (for testing).
|
||
pub fn with_memory_mb(gpu_memory_mb: usize) -> Self {
|
||
Self { gpu_memory_mb, gpu_name: String::new() }
|
||
}
|
||
|
||
/// Compute the max batch_size for a given model memory profile.
|
||
///
|
||
/// Formula: `(available_vram * safety_factor - model_overhead) / per_sample_mb`
|
||
/// Result is clamped to `[min_batch, max_batch]`.
|
||
///
|
||
/// Returns `None` if GPU is not available (budget is 0).
|
||
pub fn max_batch_size(
|
||
&self,
|
||
model_overhead_mb: f64,
|
||
per_sample_mb: f64,
|
||
min_batch: f64,
|
||
max_batch: f64,
|
||
) -> Option<f64> {
|
||
if self.gpu_memory_mb == 0 {
|
||
return None;
|
||
}
|
||
// Adaptive safety factor: larger GPUs can safely use more VRAM
|
||
let safety_factor = match self.gpu_memory_mb {
|
||
0..=8192 => 0.70,
|
||
8193..=24576 => 0.75,
|
||
24577..=49152 => 0.80,
|
||
_ => 0.85,
|
||
};
|
||
let available = self.gpu_memory_mb as f64 * safety_factor - model_overhead_mb;
|
||
if available <= 0.0 || per_sample_mb <= 0.0 {
|
||
return Some(min_batch);
|
||
}
|
||
Some((available / per_sample_mb).clamp(min_batch, max_batch))
|
||
}
|
||
|
||
/// 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
|
||
pub fn plan_hyperopt(
|
||
&self,
|
||
model_overhead_mb: f64,
|
||
per_sample_mb: f64,
|
||
min_batch: f64,
|
||
max_batch: f64,
|
||
) -> HyperoptStrategy {
|
||
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;
|
||
let available_mb = self.gpu_memory_mb as f64 * (1.0 - safety_margin);
|
||
|
||
let default_batch_mb = min_batch * per_sample_mb;
|
||
let per_trial_mb = model_overhead_mb + default_batch_mb;
|
||
|
||
let max_concurrent = if per_trial_mb <= 0.0 {
|
||
1
|
||
} else {
|
||
let raw = (available_mb / per_trial_mb).floor() as usize;
|
||
raw.clamp(1, 128) // Hardware cap: avoid thread explosion
|
||
};
|
||
|
||
let per_trial_budget = available_mb / max_concurrent as f64;
|
||
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,
|
||
}
|
||
}
|
||
|
||
/// Whether the detected GPU supports BF16 tensor cores (Ampere+ architecture).
|
||
pub fn supports_bf16(&self) -> bool {
|
||
let name = self.gpu_name.to_uppercase();
|
||
name.contains("A100") || name.contains("A10G") || name.contains("A30")
|
||
|| name.contains("A40") || name.contains("L40") || name.contains("L4 ")
|
||
|| name.contains("H100") || name.contains("H200")
|
||
|| name.contains("RTX 30") || name.contains("RTX 40") || name.contains("RTX 50")
|
||
|| name.contains("ADA") || name.contains("HOPPER") || name.contains("BLACKWELL")
|
||
}
|
||
|
||
/// Whether the detected GPU supports FP16 tensor cores (Volta+ architecture).
|
||
pub fn supports_fp16(&self) -> bool {
|
||
self.supports_bf16() || {
|
||
let name = self.gpu_name.to_uppercase();
|
||
name.contains("V100") || name.contains("T4") || name.contains("RTX 20")
|
||
|| name.contains("TURING")
|
||
}
|
||
}
|
||
|
||
/// Compute max hidden_dim_base that fits in VRAM for the actual DQN architecture.
|
||
///
|
||
/// Accounts for: C51/QR distributional atoms, dueling streams, NoisyNet (2x params),
|
||
/// main+target networks, Adam optimizer, gradients, activations, and GPU PER buffer.
|
||
pub fn max_hidden_dim_base(
|
||
&self,
|
||
concurrent_trials: usize,
|
||
batch_size: usize,
|
||
state_dim: usize,
|
||
num_actions: usize,
|
||
) -> usize {
|
||
self.max_hidden_dim_base_full(
|
||
concurrent_trials,
|
||
batch_size,
|
||
state_dim,
|
||
num_actions,
|
||
1, // num_atoms (vanilla DQN)
|
||
false, // noisy_nets
|
||
false, // dueling
|
||
0, // replay_buffer_capacity
|
||
)
|
||
}
|
||
|
||
/// Compute max hidden_dim_base with full architecture parameters.
|
||
///
|
||
/// This is the accurate estimator that accounts for the real DQN architecture:
|
||
/// C51/QR atoms, dueling value+advantage streams, NoisyNet mu+sigma doubling,
|
||
/// GPU PER replay buffer pre-allocation, BF16 dtype on CUDA, and regime-conditional
|
||
/// head multiplier (3 independent DQN heads for Trending/Ranging/Volatile).
|
||
#[allow(clippy::too_many_arguments)]
|
||
pub fn max_hidden_dim_base_full(
|
||
&self,
|
||
concurrent_trials: usize,
|
||
batch_size: usize,
|
||
state_dim: usize,
|
||
num_actions: usize,
|
||
num_atoms: usize,
|
||
noisy_nets: bool,
|
||
dueling: bool,
|
||
replay_buffer_capacity: usize,
|
||
) -> usize {
|
||
if self.gpu_memory_mb == 0 {
|
||
return 256; // CPU fallback
|
||
}
|
||
let safety_margin = 0.20;
|
||
let available_per_trial = (self.gpu_memory_mb as f64 * (1.0 - safety_margin))
|
||
/ concurrent_trials.max(1) as f64;
|
||
|
||
// Binary search for largest base that fits in budget
|
||
let mut lo: usize = 256;
|
||
let mut hi: usize = 8192;
|
||
while lo < hi {
|
||
let mid = (lo + hi).div_ceil(2);
|
||
let vram = Self::estimate_trial_vram_mb_full(
|
||
mid,
|
||
batch_size,
|
||
state_dim,
|
||
num_actions,
|
||
num_atoms,
|
||
noisy_nets,
|
||
dueling,
|
||
replay_buffer_capacity,
|
||
);
|
||
if vram <= available_per_trial {
|
||
lo = mid;
|
||
} else {
|
||
hi = mid - 1;
|
||
}
|
||
}
|
||
// Round down to nearest 256 (all values are already multiples of 8 for tensor cores)
|
||
(lo / 256) * 256
|
||
}
|
||
|
||
/// Check whether a trial with the given parameters fits in GPU VRAM.
|
||
///
|
||
/// Returns `Ok(estimated_mb)` if it fits, `Err(message)` if it exceeds budget.
|
||
///
|
||
/// Uses a two-gate approach:
|
||
/// 1. **Static estimate**: architecture-based VRAM estimate with 3× fragmentation
|
||
/// multiplier must be under 80% of initially detected free VRAM.
|
||
/// 2. **Live measurement**: re-queries nvidia-smi for CURRENT free VRAM (accounts
|
||
/// for CUDA context already initialized, previous trial leftovers, etc.)
|
||
/// and checks the estimate against 80% of that.
|
||
///
|
||
/// Both gates must pass. This prevents stale initial measurements from allowing
|
||
/// trials that won't fit after CUDA context initialization.
|
||
#[allow(clippy::too_many_arguments)]
|
||
pub fn trial_fits_vram(
|
||
&self,
|
||
hidden_dim_base: usize,
|
||
batch_size: usize,
|
||
state_dim: usize,
|
||
num_actions: usize,
|
||
num_atoms: usize,
|
||
noisy_nets: bool,
|
||
dueling: bool,
|
||
replay_buffer_capacity: usize,
|
||
) -> Result<f64, String> {
|
||
if self.gpu_memory_mb == 0 {
|
||
return Ok(0.0); // CPU mode, no VRAM constraint
|
||
}
|
||
let estimated = Self::estimate_trial_vram_mb_full(
|
||
hidden_dim_base,
|
||
batch_size,
|
||
state_dim,
|
||
num_actions,
|
||
num_atoms,
|
||
noisy_nets,
|
||
dueling,
|
||
replay_buffer_capacity,
|
||
);
|
||
|
||
// Gate 1: static budget from initial detection
|
||
let static_budget = self.gpu_memory_mb as f64 * 0.80;
|
||
if estimated > static_budget {
|
||
return Err(format!(
|
||
"Trial VRAM estimate {estimated:.0} MB exceeds static budget {static_budget:.0} MB \
|
||
(GPU: {} MB initial free). Params: hidden={hidden_dim_base} atoms={num_atoms} \
|
||
batch={batch_size} buffer={replay_buffer_capacity}",
|
||
self.gpu_memory_mb
|
||
));
|
||
}
|
||
|
||
// Gate 2: live nvidia-smi measurement for current free VRAM
|
||
// By the time this runs, CUDA context may already be initialized (eating ~300-500 MB),
|
||
// so live free is typically less than the initial detection.
|
||
if let Ok((_total, live_free, _name)) =
|
||
ml_core::memory_optimization::auto_batch_size::detect_gpu_memory()
|
||
{
|
||
let live_budget = live_free * 0.80;
|
||
if estimated > live_budget {
|
||
return Err(format!(
|
||
"Trial VRAM estimate {estimated:.0} MB exceeds live budget {live_budget:.0} MB \
|
||
(GPU: {live_free:.0} MB currently free). Params: hidden={hidden_dim_base} \
|
||
atoms={num_atoms} batch={batch_size} buffer={replay_buffer_capacity}",
|
||
));
|
||
}
|
||
}
|
||
|
||
Ok(estimated)
|
||
}
|
||
|
||
/// Estimate total VRAM for one trial (vanilla DQN, no distributional/dueling/noisy).
|
||
///
|
||
/// For production DQN with C51/Dueling/NoisyNet, use `estimate_trial_vram_mb_full`.
|
||
pub fn estimate_trial_vram_mb(
|
||
hidden_dim_base: usize,
|
||
batch_size: usize,
|
||
state_dim: usize,
|
||
num_actions: usize,
|
||
) -> f64 {
|
||
Self::estimate_trial_vram_mb_full(
|
||
hidden_dim_base,
|
||
batch_size,
|
||
state_dim,
|
||
num_actions,
|
||
1, // vanilla DQN
|
||
false, // no noisy nets
|
||
false, // no dueling
|
||
0, // no GPU PER buffer
|
||
)
|
||
}
|
||
|
||
/// Accurate VRAM estimate for the real DQN architecture.
|
||
///
|
||
/// Model weights and replay buffer states use BF16 (2 bytes) on CUDA.
|
||
/// Optimizer moments and gradients are stored in F32 (4 bytes) for
|
||
/// numerical stability (AdamW moment accumulation in BF16 causes
|
||
/// catastrophic underflow). Loss computation is fully F32.
|
||
///
|
||
/// The final multiplier (3×) accounts for severe memory fragmentation
|
||
/// caused by the cudarc + Candle dual CUDA allocator architecture.
|
||
/// nvidia-smi may report GB of "free" memory during OOM because the
|
||
/// driver sees total free, but neither allocator can use the other's
|
||
/// freed blocks. Empirically validated: a trial estimating 750 MB
|
||
/// OOM'd on a 4 GB GPU with 2.9 GB "free" per nvidia-smi.
|
||
///
|
||
/// Accounts for:
|
||
/// - Constant-width hidden layers: `[base, base]` (NOT pyramid `[b, b/2, b/4]`)
|
||
/// - C51/QR distributional: output = num_actions × num_atoms (not just num_actions)
|
||
/// - Dueling: separate value stream (1 × num_atoms) + advantage stream (num_actions × num_atoms)
|
||
/// - NoisyNet: doubles parameters per layer (mu + sigma) + noise buffers
|
||
/// - GPU PER replay buffer: pre-allocated [capacity, state_dim] × 2 (states + next_states)
|
||
/// - Regime-conditional: 3 independent heads (Trending/Ranging/Volatile), each with
|
||
/// its own network weights, optimizer, gradients, and replay buffer
|
||
/// - GPU experience collector: dynamic output buffers (n_episodes × timesteps)
|
||
/// - cudarc weight copies: F32 duplicates of all model weights for CUDA kernels
|
||
/// - C51 distributional loss intermediates: projected distributions, softmax, cross-entropy
|
||
/// - Training data: Candle BF16 + cudarc F32 copies of market features/targets
|
||
/// - PER sampling intermediates: cumsum, index_select on GPU buffer
|
||
/// - CUDA context: driver, cuBLAS workspace, NVRTC modules, memory allocator overhead
|
||
/// - Backward pass autograd graph: Candle retains ALL intermediate tensors from forward
|
||
/// through loss through backward — 20× activation memory for C51 distributional
|
||
pub fn estimate_trial_vram_mb_full(
|
||
hidden_dim_base: usize,
|
||
batch_size: usize,
|
||
state_dim: usize,
|
||
num_actions: usize,
|
||
num_atoms: usize,
|
||
noisy_nets: bool,
|
||
dueling: bool,
|
||
replay_buffer_capacity: usize,
|
||
) -> f64 {
|
||
// Actual network: constant-width [base, base] (see trainer.rs vec![b, b])
|
||
let b = hidden_dim_base;
|
||
let output_dim = num_actions * num_atoms.max(1);
|
||
|
||
// Backbone: input → base → base (2 constant-width layers)
|
||
// + bias terms: b + b
|
||
let mut param_count = state_dim * b + b // layer 1: weights + bias
|
||
+ b * b + b; // layer 2: weights + bias
|
||
|
||
if dueling {
|
||
// Dueling: value stream (b → value_hidden → num_atoms)
|
||
// + advantage stream (b → adv_hidden → output_dim)
|
||
// value_hidden and adv_hidden default to 128 (dueling_hidden_dim)
|
||
// but scale with hidden_dim_base in hyperopt
|
||
let vh = (b / 4).max(128);
|
||
param_count += b * vh + vh // value hidden layer
|
||
+ vh * num_atoms.max(1) + num_atoms.max(1) // value output
|
||
+ b * vh + vh // advantage hidden layer
|
||
+ vh * output_dim + output_dim; // advantage output
|
||
} else {
|
||
// Single head: b → output_dim
|
||
param_count += b * output_dim + output_dim;
|
||
}
|
||
|
||
// NoisyNet: mu + sigma (2x trainable params) + noise buffers (0.5x, non-trainable)
|
||
let noisy_multiplier: f64 = if noisy_nets { 2.5 } else { 1.0 };
|
||
let effective_params = param_count as f64 * noisy_multiplier;
|
||
// Trainable params only (for optimizer/gradient accounting)
|
||
let trainable_multiplier: f64 = if noisy_nets { 2.0 } else { 1.0 };
|
||
let trainable_params = param_count as f64 * trainable_multiplier;
|
||
|
||
// BF16 = 2 bytes, F32 = 4 bytes
|
||
let bf16_bytes: f64 = 2.0;
|
||
let f32_bytes: f64 = 4.0;
|
||
|
||
// === Per-head costs (× 3 for regime-conditional) ===
|
||
|
||
// Weights: main + target networks (BF16 in Candle via training_dtype)
|
||
let weights_mb = effective_params * bf16_bytes * 2.0 / 1e6;
|
||
// Adam optimizer: 2 moments per trainable param.
|
||
// Stored in F32 for numerical stability — BF16 moment accumulation
|
||
// causes catastrophic underflow in the running variance (second moment).
|
||
let optimizer_mb = trainable_params * f32_bytes * 2.0 / 1e6;
|
||
// Gradients: F32 in Candle's autograd (backward() always produces F32
|
||
// GradStore, even when model tensors are BF16 — casts happen in optimizer step)
|
||
let gradients_mb = trainable_params * f32_bytes / 1e6;
|
||
|
||
// GPU PER replay buffer: pre-allocated ring buffer tensors
|
||
let replay_mb = if replay_buffer_capacity > 0 {
|
||
let cap = replay_buffer_capacity as f64;
|
||
let sd = state_dim as f64;
|
||
cap * sd * bf16_bytes * 2.0 / 1e6 // states + next_states (BF16)
|
||
+ cap * f32_bytes * 4.0 / 1e6 // actions(U32)/rewards/dones/priorities (F32)
|
||
} else {
|
||
0.0
|
||
};
|
||
|
||
// PER sampling intermediates: cumsum, pow, index_select create temporary tensors
|
||
// on top of the replay buffer. Worst case ~3× single column of the buffer.
|
||
let per_sampling_mb = if replay_buffer_capacity > 0 {
|
||
replay_buffer_capacity as f64 * f32_bytes * 3.0 / 1e6
|
||
} else {
|
||
0.0
|
||
};
|
||
|
||
// Regime-conditional: 3 independent heads
|
||
let num_heads: f64 = 3.0;
|
||
let per_head_mb = weights_mb + optimizer_mb + gradients_mb + replay_mb + per_sampling_mb;
|
||
|
||
// === Shared costs (not multiplied by heads) ===
|
||
|
||
// Activation memory: 2 hidden layers × batch_size × BF16 × 2 (forward+backward)
|
||
let total_hidden = (b * 2) as f64;
|
||
let activation_mb = total_hidden * batch_size as f64 * bf16_bytes * 2.0 / 1e6;
|
||
|
||
// Backward pass autograd graph: Candle retains ALL intermediate tensors
|
||
// from forward pass through loss computation until backward() completes.
|
||
// For C51 distributional with 100 atoms, the Tz projection creates ~20
|
||
// intermediate tensors ([batch, atoms] or [batch, actions, atoms]),
|
||
// softmax creates ~5 more, cross-entropy ~3 more, plus dueling stream
|
||
// intermediates. Empirically measured at 20× forward activation memory.
|
||
let backward_graph_mb = activation_mb * 20.0;
|
||
|
||
// Batch tensors: states + next_states (BF16) + actions/rewards/dones/weights (F32 scalars)
|
||
let batch_mb = batch_size as f64 * state_dim as f64 * bf16_bytes * 2.0 / 1e6
|
||
+ batch_size as f64 * f32_bytes * 4.0 / 1e6;
|
||
|
||
// C51 distributional loss intermediates (per training step):
|
||
// - projected distribution: [batch, atoms] × F32
|
||
// - target distribution: [batch, atoms] × F32
|
||
// - online logits: [batch, actions, atoms] × F32
|
||
// - softmax + cross-entropy intermediates: ~3× above (Tz clip, floor, ceil,
|
||
// offset lower, offset upper, weighted scatter, log_softmax, kl_div)
|
||
let c51_loss_mb = if num_atoms > 1 {
|
||
let atoms = num_atoms as f64;
|
||
let bs = batch_size as f64;
|
||
let na = num_actions as f64;
|
||
(bs * atoms * 2.0 + bs * na * atoms * 5.0) * f32_bytes / 1e6
|
||
} else {
|
||
0.0
|
||
};
|
||
|
||
// GPU-resident features: experience collector + cudarc weight copies + data staging.
|
||
// These are ONLY active when GPU PER is enabled (replay_buffer_capacity > 0 as proxy).
|
||
// On small GPUs (≤8 GB), GPU PER and experience collector are disabled to avoid
|
||
// cudarc/Candle dual-allocator fragmentation.
|
||
let gpu_features_active = replay_buffer_capacity > 0;
|
||
|
||
let collector_mb = if gpu_features_active {
|
||
// GPU experience collector: 128 episodes × 500 timesteps = 64K output slots
|
||
let collector_output: f64 = 64_000.0;
|
||
let collector_per_episode: f64 = 128.0;
|
||
collector_output * state_dim as f64 * f32_bytes / 1e6
|
||
+ collector_output * f32_bytes * 5.0 / 1e6
|
||
+ collector_per_episode * 114.0 * f32_bytes / 1e6
|
||
} else {
|
||
0.0
|
||
};
|
||
|
||
// cudarc F32 weight copies: only when GPU experience collector is active
|
||
let cudarc_weights_mb = if gpu_features_active {
|
||
effective_params * f32_bytes * 2.0 / 1e6
|
||
} else {
|
||
0.0
|
||
};
|
||
|
||
// Training data on GPU (Candle BF16 always present for training)
|
||
let training_bars = if gpu_features_active {
|
||
(replay_buffer_capacity as f64).max(50_000.0)
|
||
} else {
|
||
50_000.0 // minimum training data estimate
|
||
};
|
||
let data_candle_mb = training_bars * (42.0 + 4.0) * bf16_bytes / 1e6;
|
||
// cudarc F32 data copies: only when GPU experience collector is active
|
||
let data_cudarc_mb = if gpu_features_active {
|
||
training_bars * (42.0 + 4.0) * f32_bytes / 1e6
|
||
} else {
|
||
0.0
|
||
};
|
||
|
||
// GPU buffer pool: only when GPU experience collector is active
|
||
let buffer_pool_mb = if gpu_features_active {
|
||
100_000.0 * 42.0 * f32_bytes / 1e6
|
||
} else {
|
||
0.0
|
||
};
|
||
|
||
// CUDA context + cuBLAS workspace. Reduced when no NVRTC/cudarc overhead.
|
||
let cuda_overhead: f64 = if gpu_features_active { 500.0 } else { 350.0 };
|
||
|
||
let subtotal = per_head_mb * num_heads
|
||
+ activation_mb + backward_graph_mb + batch_mb
|
||
+ c51_loss_mb
|
||
+ collector_mb + cudarc_weights_mb
|
||
+ data_candle_mb + data_cudarc_mb
|
||
+ buffer_pool_mb
|
||
+ cuda_overhead;
|
||
|
||
// Fragmentation multiplier:
|
||
// - 3× when dual allocators active (cudarc + Candle cannot reuse each other's freed blocks)
|
||
// - 1.5× when only Candle allocator (still some fragmentation from autograd temporaries)
|
||
let fragmentation = if gpu_features_active { 3.0 } else { 1.5 };
|
||
subtotal * fragmentation
|
||
}
|
||
}
|
||
|
||
/// Trait for models that support hyperparameter optimization
|
||
///
|
||
/// This trait defines the interface for models that can be optimized using
|
||
/// Bayesian optimization. Models must:
|
||
/// 1. Define a parameter space type implementing `ParameterSpace`
|
||
/// 2. Define a metrics type for training results
|
||
/// 3. Implement training with arbitrary parameters
|
||
/// 4. Extract a scalar objective from metrics (lower is better)
|
||
///
|
||
/// ## Type Parameters
|
||
///
|
||
/// - `Params`: Parameter space type (must implement `ParameterSpace`)
|
||
/// - `Metrics`: Training metrics type (must be Clone + Debug)
|
||
///
|
||
/// ## Example Implementation
|
||
///
|
||
/// ```rust,no_run
|
||
/// use ml::hyperopt::{HyperparameterOptimizable, ParameterSpace};
|
||
/// use ml::MLError;
|
||
///
|
||
/// #[derive(Debug, Clone)]
|
||
/// struct MyParams { learning_rate: f64 }
|
||
///
|
||
/// impl ParameterSpace for MyParams {
|
||
/// fn continuous_bounds() -> Vec<(f64, f64)> {
|
||
/// vec![(1e-5_f64.ln(), 1e-2_f64.ln())]
|
||
/// }
|
||
/// fn from_continuous(x: &[f64]) -> Result<Self, MLError> {
|
||
/// Ok(Self { learning_rate: x[0].exp() })
|
||
/// }
|
||
/// fn to_continuous(&self) -> Vec<f64> {
|
||
/// vec![self.learning_rate.ln()]
|
||
/// }
|
||
/// fn param_names() -> Vec<&'static str> {
|
||
/// vec!["learning_rate"]
|
||
/// }
|
||
/// }
|
||
///
|
||
/// #[derive(Debug, Clone)]
|
||
/// struct MyMetrics { val_loss: f64 }
|
||
///
|
||
/// struct MyModel;
|
||
///
|
||
/// impl HyperparameterOptimizable for MyModel {
|
||
/// type Params = MyParams;
|
||
/// type Metrics = MyMetrics;
|
||
///
|
||
/// fn train_with_params(&mut self, params: Self::Params) -> Result<Self::Metrics, MLError> {
|
||
/// // Train model with params...
|
||
/// Ok(MyMetrics { val_loss: 0.5 })
|
||
/// }
|
||
///
|
||
/// fn extract_objective(metrics: &Self::Metrics) -> f64 {
|
||
/// metrics.val_loss
|
||
/// }
|
||
/// }
|
||
/// ```
|
||
pub trait HyperparameterOptimizable {
|
||
/// Parameter space type (must implement `ParameterSpace`)
|
||
type Params: ParameterSpace + Clone + Debug;
|
||
|
||
/// Metrics type returned by training (must be Clone + Debug for result tracking)
|
||
type Metrics: Clone + Debug;
|
||
|
||
/// Train model with given parameters and return metrics
|
||
///
|
||
/// This method should:
|
||
/// 1. Configure the model with the provided parameters
|
||
/// 2. Run the full training loop
|
||
/// 3. Return comprehensive metrics including validation loss
|
||
///
|
||
/// # Arguments
|
||
///
|
||
/// * `params` - Hyperparameters to use for training
|
||
///
|
||
/// # Returns
|
||
///
|
||
/// Training metrics including validation loss for optimization
|
||
///
|
||
/// # Errors
|
||
///
|
||
/// Returns `MLError` if training fails due to:
|
||
/// - Invalid parameter combinations
|
||
/// - Out of memory conditions
|
||
/// - Numerical instability (NaN/Inf)
|
||
/// - CUDA errors
|
||
fn train_with_params(&mut self, params: Self::Params) -> Result<Self::Metrics, MLError>;
|
||
|
||
/// Extract scalar objective value from metrics (lower is better)
|
||
///
|
||
/// This method defines what metric to optimize. Common choices:
|
||
/// - Validation loss (minimize)
|
||
/// - Negative accuracy (minimize negative = maximize accuracy)
|
||
/// - Custom scoring function (e.g., Sharpe ratio)
|
||
///
|
||
/// # Arguments
|
||
///
|
||
/// * `metrics` - Training metrics from `train_with_params`
|
||
///
|
||
/// # Returns
|
||
///
|
||
/// Scalar objective value where **lower is better**. The optimizer
|
||
/// will try to minimize this value across trials.
|
||
///
|
||
/// # Example
|
||
///
|
||
/// ```rust
|
||
/// # use ml::hyperopt::HyperparameterOptimizable;
|
||
/// # struct MyModel;
|
||
/// # #[derive(Clone, Debug)] struct MyParams;
|
||
/// # #[derive(Clone, Debug)] struct MyMetrics { val_loss: f64, accuracy: f64 }
|
||
/// # impl HyperparameterOptimizable for MyModel {
|
||
/// # type Params = MyParams;
|
||
/// # type Metrics = MyMetrics;
|
||
/// # fn train_with_params(&mut self, _: Self::Params) -> Result<Self::Metrics, ml::MLError> { unimplemented!() }
|
||
/// fn extract_objective(metrics: &Self::Metrics) -> f64 {
|
||
/// // Minimize validation loss
|
||
/// metrics.val_loss
|
||
///
|
||
/// // Or maximize accuracy (negate to minimize)
|
||
/// // -metrics.accuracy
|
||
/// }
|
||
/// # }
|
||
/// ```
|
||
fn extract_objective(metrics: &Self::Metrics) -> f64;
|
||
|
||
/// Extract per-trial metrics as opaque JSON for reporting.
|
||
///
|
||
/// Override this to attach domain-specific metrics (Sharpe, drawdown, etc.)
|
||
/// to each `TrialResult`. The default returns `None` (no extra metrics).
|
||
fn extract_metrics(_metrics: &Self::Metrics) -> Option<serde_json::Value> {
|
||
None
|
||
}
|
||
}
|
||
|
||
/// Trait for parameter space definition
|
||
///
|
||
/// This trait defines how to convert between parameter representations:
|
||
/// - **Continuous space** `[0, 1]^n`: Used by the optimizer
|
||
/// - **Original space**: Human-readable parameter values (e.g., learning_rate=0.001)
|
||
///
|
||
/// The continuous space allows the optimizer to:
|
||
/// - Sample uniformly across the space using Latin Hypercube
|
||
/// - Build Gaussian Process surrogates
|
||
/// - Perform efficient gradient-based acquisition optimization
|
||
///
|
||
/// ## Log-Scale Parameters
|
||
///
|
||
/// For parameters spanning multiple orders of magnitude (learning rates, weight decay),
|
||
/// use log-scale transformations in `from_continuous`/`to_continuous`:
|
||
///
|
||
/// ```rust
|
||
/// # use ml::hyperopt::ParameterSpace;
|
||
/// # use ml::MLError;
|
||
/// #[derive(Clone, Debug)]
|
||
/// struct Params { learning_rate: f64 }
|
||
///
|
||
/// impl ParameterSpace for Params {
|
||
/// fn continuous_bounds() -> Vec<(f64, f64)> {
|
||
/// vec![(1e-5_f64.ln(), 1e-2_f64.ln())] // Log-scale bounds
|
||
/// }
|
||
///
|
||
/// fn from_continuous(x: &[f64]) -> Result<Self, MLError> {
|
||
/// Ok(Self { learning_rate: x[0].exp() }) // exp() to recover original scale
|
||
/// }
|
||
///
|
||
/// fn to_continuous(&self) -> Vec<f64> {
|
||
/// vec![self.learning_rate.ln()] // ln() for continuous space
|
||
/// }
|
||
///
|
||
/// fn param_names() -> Vec<&'static str> {
|
||
/// vec!["learning_rate"]
|
||
/// }
|
||
/// }
|
||
/// ```
|
||
pub trait ParameterSpace: Sized {
|
||
/// Return bounds for continuous parameters: `[(min, max), ...]`
|
||
///
|
||
/// Each tuple defines the valid range for one continuous parameter.
|
||
/// These bounds are used by the optimizer to:
|
||
/// 1. Generate initial samples via Latin Hypercube Sampling
|
||
/// 2. Constrain Gaussian Process predictions
|
||
/// 3. Limit acquisition function optimization
|
||
///
|
||
/// # Log-Scale Parameters
|
||
///
|
||
/// For parameters like learning rates (1e-5 to 1e-2), use log-scale bounds:
|
||
/// ```rust
|
||
/// vec![
|
||
/// (1e-5_f64.ln(), 1e-2_f64.ln()), // learning_rate (log scale)
|
||
/// (16.0, 256.0), // batch_size (linear scale)
|
||
/// ]
|
||
/// ```
|
||
///
|
||
/// # Returns
|
||
///
|
||
/// Vector of `(min, max)` tuples, one per parameter. Length must match
|
||
/// the dimensionality returned by `to_continuous()`.
|
||
///
|
||
/// # Panics
|
||
///
|
||
/// Implementations should never panic. Invalid bounds should return
|
||
/// default values or use type-level guarantees.
|
||
fn continuous_bounds() -> Vec<(f64, f64)>;
|
||
|
||
/// Convert from continuous vector `[0, 1]^n` to parameters
|
||
///
|
||
/// This method reconstructs original parameter values from the continuous
|
||
/// representation used by the optimizer. It should:
|
||
/// 1. Denormalize continuous values to actual ranges
|
||
/// 2. Apply inverse transformations (e.g., `exp()` for log-scale params)
|
||
/// 3. Convert to discrete types where needed (e.g., rounding for batch_size)
|
||
///
|
||
/// # Arguments
|
||
///
|
||
/// * `x` - Continuous parameter vector from the optimizer
|
||
///
|
||
/// # Returns
|
||
///
|
||
/// Reconstructed parameters, or `MLError` if conversion fails
|
||
///
|
||
/// # Errors
|
||
///
|
||
/// Returns `MLError::ConfigError` if:
|
||
/// - Vector length doesn't match expected dimensionality
|
||
/// - Values are outside valid bounds (should be impossible with proper bounds)
|
||
/// - Conversion produces invalid parameter combinations
|
||
///
|
||
/// # Example
|
||
///
|
||
/// ```rust
|
||
/// # use ml::hyperopt::ParameterSpace;
|
||
/// # use ml::MLError;
|
||
/// # #[derive(Clone, Debug)]
|
||
/// # struct Params { learning_rate: f64, batch_size: usize }
|
||
/// # impl ParameterSpace for Params {
|
||
/// # fn continuous_bounds() -> Vec<(f64, f64)> { vec![] }
|
||
/// fn from_continuous(x: &[f64]) -> Result<Self, MLError> {
|
||
/// if x.len() != 2 {
|
||
/// return Err(MLError::ConfigError(
|
||
/// format!("Expected 2 params, got {}", x.len())
|
||
/// ));
|
||
/// }
|
||
///
|
||
/// Ok(Self {
|
||
/// learning_rate: x[0].exp(), // Log-scale
|
||
/// batch_size: x[1].round() as usize, // Discrete
|
||
/// })
|
||
/// }
|
||
/// # fn to_continuous(&self) -> Vec<f64> { vec![] }
|
||
/// # fn param_names() -> Vec<&'static str> { vec![] }
|
||
/// # }
|
||
/// ```
|
||
fn from_continuous(x: &[f64]) -> Result<Self, MLError>;
|
||
|
||
/// Convert parameters to continuous vector for the optimizer
|
||
///
|
||
/// This method should be the exact inverse of `from_continuous()`.
|
||
/// It converts parameters to the continuous representation used by the optimizer.
|
||
///
|
||
/// # Returns
|
||
///
|
||
/// Continuous parameter vector. Length must match `continuous_bounds()`.
|
||
///
|
||
/// # Invariants
|
||
///
|
||
/// The following must hold for all valid parameters:
|
||
/// ```rust,ignore
|
||
/// let params = MyParams { ... };
|
||
/// let continuous = params.to_continuous();
|
||
/// let recovered = MyParams::from_continuous(&continuous)?;
|
||
/// assert_eq!(params, recovered); // Approximate equality for floats
|
||
/// ```
|
||
fn to_continuous(&self) -> Vec<f64>;
|
||
|
||
/// Get parameter names for display and logging
|
||
///
|
||
/// Returns human-readable names for each parameter dimension.
|
||
/// Used for:
|
||
/// - Trial result logging
|
||
/// - Convergence plot labels
|
||
/// - Final result reporting
|
||
///
|
||
/// # Returns
|
||
///
|
||
/// Vector of parameter names. Length must match `continuous_bounds()`.
|
||
///
|
||
/// # Example
|
||
///
|
||
/// ```rust
|
||
/// # use ml::hyperopt::ParameterSpace;
|
||
/// # use ml::MLError;
|
||
/// # #[derive(Clone, Debug)]
|
||
/// # struct Params { learning_rate: f64, batch_size: usize }
|
||
/// # impl ParameterSpace for Params {
|
||
/// # fn continuous_bounds() -> Vec<(f64, f64)> { vec![] }
|
||
/// # fn from_continuous(_: &[f64]) -> Result<Self, MLError> { unimplemented!() }
|
||
/// # fn to_continuous(&self) -> Vec<f64> { vec![] }
|
||
/// fn param_names() -> Vec<&'static str> {
|
||
/// vec!["learning_rate", "batch_size"]
|
||
/// }
|
||
/// # }
|
||
/// ```
|
||
fn param_names() -> Vec<&'static str>;
|
||
|
||
/// Return hardware-aware bounds for continuous parameters.
|
||
///
|
||
/// Override this to scale batch_size (or other memory-sensitive params)
|
||
/// based on detected GPU VRAM. The default implementation delegates to
|
||
/// `continuous_bounds()`, preserving existing static behavior.
|
||
///
|
||
/// Adapters that override this should call `Self::continuous_bounds()`
|
||
/// first, then replace the batch_size bound using `budget.max_batch_size()`.
|
||
fn continuous_bounds_for(_budget: &HardwareBudget) -> Vec<(f64, f64)> {
|
||
Self::continuous_bounds()
|
||
}
|
||
|
||
/// Estimated fixed VRAM overhead per model instance in MB (weights + optimizer + gradients + activations).
|
||
/// Override for accurate VRAM-based concurrency planning.
|
||
fn model_overhead_mb() -> f64 {
|
||
500.0
|
||
}
|
||
|
||
/// Estimated VRAM per batch sample in MB.
|
||
fn per_sample_mb() -> f64 {
|
||
0.02
|
||
}
|
||
}
|
||
|
||
/// Result of a single optimization trial
|
||
///
|
||
/// Contains parameter values, objective score, and timing information
|
||
/// for one training run during optimization.
|
||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||
pub struct TrialResult<P> {
|
||
/// Trial number (1-indexed)
|
||
pub trial_num: usize,
|
||
/// Parameters used for this trial
|
||
pub params: P,
|
||
/// Objective value achieved (lower is better)
|
||
pub objective: f64,
|
||
/// Wall-clock time for training (seconds)
|
||
pub duration_secs: f64,
|
||
/// Optional domain-specific metrics (Sharpe, drawdown, etc.)
|
||
#[serde(skip_serializing_if = "Option::is_none")]
|
||
pub metrics: Option<serde_json::Value>,
|
||
}
|
||
|
||
/// Complete optimization result with best parameters and trial history
|
||
///
|
||
/// Contains all information from a hyperparameter optimization run:
|
||
/// - Best parameters found
|
||
/// - Best objective value
|
||
/// - Full trial history for analysis
|
||
/// - Convergence plot data
|
||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||
pub struct OptimizationResult<P> {
|
||
/// Best parameters found during optimization
|
||
pub best_params: P,
|
||
/// Best objective value (lower is better)
|
||
pub best_objective: f64,
|
||
/// All trial results in execution order
|
||
pub all_trials: Vec<TrialResult<P>>,
|
||
/// Convergence plot data: `(trial_num, best_objective_so_far)`
|
||
///
|
||
/// This tracks the best objective seen up to each trial, showing
|
||
/// convergence behavior over time. Can be used to:
|
||
/// - Plot optimization progress
|
||
/// - Detect early convergence
|
||
/// - Identify exploration vs exploitation phases
|
||
pub convergence_plot_data: Vec<(usize, f64)>,
|
||
}
|
||
|
||
impl<P: Clone> OptimizationResult<P> {
|
||
/// Create new optimization result from trial history
|
||
///
|
||
/// Automatically computes best parameters and convergence data.
|
||
///
|
||
/// # Arguments
|
||
///
|
||
/// * `trials` - Complete trial history
|
||
///
|
||
/// # Panics
|
||
///
|
||
/// Panics if `trials` is empty (should never happen in practice)
|
||
#[allow(clippy::indexing_slicing)] // assert! guarantees non-empty; min_by always returns Some
|
||
pub fn from_trials(trials: Vec<TrialResult<P>>) -> Self {
|
||
assert!(!trials.is_empty(), "Cannot create result from empty trials");
|
||
|
||
// Find best trial (minimum objective)
|
||
// assert! guarantees non-empty, so min_by always returns Some
|
||
let best_trial = trials
|
||
.iter()
|
||
.min_by(|a, b| a.objective.partial_cmp(&b.objective).unwrap_or(std::cmp::Ordering::Equal))
|
||
.unwrap_or(&trials[0]);
|
||
|
||
// Build convergence plot data
|
||
let mut best_so_far = f64::INFINITY;
|
||
let convergence_plot_data: Vec<_> = trials
|
||
.iter()
|
||
.map(|trial| {
|
||
if trial.objective < best_so_far {
|
||
best_so_far = trial.objective;
|
||
}
|
||
(trial.trial_num, best_so_far)
|
||
})
|
||
.collect();
|
||
|
||
Self {
|
||
best_params: best_trial.params.clone(),
|
||
best_objective: best_trial.objective,
|
||
all_trials: trials,
|
||
convergence_plot_data,
|
||
}
|
||
}
|
||
|
||
/// Get the trial with the best objective.
|
||
///
|
||
/// Returns `None` if `all_trials` is empty (should not happen for well-constructed results).
|
||
pub fn best_trial(&self) -> Option<&TrialResult<P>> {
|
||
self.all_trials
|
||
.iter()
|
||
.min_by(|a, b| a.objective.partial_cmp(&b.objective).unwrap_or(std::cmp::Ordering::Equal))
|
||
}
|
||
|
||
/// Get improvement from first to best trial.
|
||
///
|
||
/// Returns `0.0` if there are no trials.
|
||
pub fn total_improvement(&self) -> f64 {
|
||
match self.all_trials.first() {
|
||
Some(first) => self.best_objective - first.objective,
|
||
None => 0.0,
|
||
}
|
||
}
|
||
|
||
/// Get improvement percentage.
|
||
///
|
||
/// Returns `0.0` if there are no trials or the first trial objective is near zero.
|
||
pub fn improvement_percentage(&self) -> f64 {
|
||
match self.all_trials.first() {
|
||
Some(first) if first.objective.abs() >= 1e-10 => {
|
||
((first.objective - self.best_objective) / first.objective.abs()) * 100.0
|
||
}
|
||
_ => 0.0,
|
||
}
|
||
}
|
||
}
|
||
|
||
#[cfg(test)]
|
||
mod tests {
|
||
use super::*;
|
||
|
||
#[derive(Debug, Clone, PartialEq)]
|
||
struct TestParams {
|
||
learning_rate: f64,
|
||
batch_size: usize,
|
||
}
|
||
|
||
impl ParameterSpace for TestParams {
|
||
fn continuous_bounds() -> Vec<(f64, f64)> {
|
||
vec![
|
||
(1e-5_f64.ln(), 1e-2_f64.ln()), // lr (log scale)
|
||
(16.0, 256.0), // batch_size
|
||
]
|
||
}
|
||
|
||
fn from_continuous(x: &[f64]) -> Result<Self, MLError> {
|
||
if x.len() != 2 {
|
||
return Err(MLError::ConfigError(format!("Expected 2 params, got {}", x.len())));
|
||
}
|
||
|
||
Ok(Self {
|
||
learning_rate: x[0].exp(),
|
||
batch_size: x[1].round() as usize,
|
||
})
|
||
}
|
||
|
||
fn to_continuous(&self) -> Vec<f64> {
|
||
vec![self.learning_rate.ln(), self.batch_size as f64]
|
||
}
|
||
|
||
fn param_names() -> Vec<&'static str> {
|
||
vec!["learning_rate", "batch_size"]
|
||
}
|
||
}
|
||
|
||
#[test]
|
||
fn test_parameter_space_roundtrip() {
|
||
let params = TestParams {
|
||
learning_rate: 0.001,
|
||
batch_size: 64,
|
||
};
|
||
|
||
let continuous = params.to_continuous();
|
||
let recovered = TestParams::from_continuous(&continuous).unwrap();
|
||
|
||
assert!((recovered.learning_rate - params.learning_rate).abs() < 1e-10);
|
||
assert_eq!(recovered.batch_size, params.batch_size);
|
||
}
|
||
|
||
#[test]
|
||
fn test_optimization_result_convergence() {
|
||
let trials = vec![
|
||
TrialResult {
|
||
trial_num: 1,
|
||
params: TestParams {
|
||
learning_rate: 0.01,
|
||
batch_size: 32,
|
||
},
|
||
objective: 1.5,
|
||
duration_secs: 10.0,
|
||
metrics: None,
|
||
},
|
||
TrialResult {
|
||
trial_num: 2,
|
||
params: TestParams {
|
||
learning_rate: 0.001,
|
||
batch_size: 64,
|
||
},
|
||
objective: 0.8,
|
||
duration_secs: 12.0,
|
||
metrics: None,
|
||
},
|
||
TrialResult {
|
||
trial_num: 3,
|
||
params: TestParams {
|
||
learning_rate: 0.005,
|
||
batch_size: 128,
|
||
},
|
||
objective: 0.5,
|
||
duration_secs: 15.0,
|
||
metrics: None,
|
||
},
|
||
];
|
||
|
||
let result = OptimizationResult::from_trials(trials);
|
||
|
||
assert_eq!(result.best_objective, 0.5);
|
||
assert_eq!(result.best_params.batch_size, 128);
|
||
assert_eq!(result.convergence_plot_data.len(), 3);
|
||
assert_eq!(result.convergence_plot_data[0].1, 1.5);
|
||
assert_eq!(result.convergence_plot_data[1].1, 0.8);
|
||
assert_eq!(result.convergence_plot_data[2].1, 0.5);
|
||
}
|
||
|
||
#[test]
|
||
fn test_hardware_budget_cpu_only() {
|
||
let budget = HardwareBudget::cpu_only();
|
||
assert_eq!(budget.gpu_memory_mb, 0);
|
||
assert!(budget.max_batch_size(100.0, 0.05, 16.0, 4096.0).is_none());
|
||
}
|
||
|
||
#[test]
|
||
fn test_hardware_budget_with_memory() {
|
||
let budget = HardwareBudget::with_memory_mb(4096); // 4GB like RTX 3050 Ti
|
||
assert_eq!(budget.gpu_memory_mb, 4096);
|
||
|
||
// 4096 * 0.70 = 2867.2; 2867.2 - 150 overhead = 2717.2; 2717.2 / 0.05 = 54344
|
||
// clamped to max_batch=4096
|
||
let max = budget.max_batch_size(150.0, 0.05, 16.0, 4096.0).unwrap();
|
||
assert!((max - 4096.0).abs() < 0.01);
|
||
}
|
||
|
||
#[test]
|
||
fn test_hardware_budget_h100_wide_range() {
|
||
let budget = HardwareBudget::with_memory_mb(81920); // 80GB H100
|
||
// 81920 * 0.85 = 69632; 69632 - 150 = 69482; 69482 / 0.05 = 1_389_640
|
||
// clamped to 4096
|
||
let max = budget.max_batch_size(150.0, 0.05, 16.0, 4096.0).unwrap();
|
||
assert!((max - 4096.0).abs() < 0.01);
|
||
}
|
||
|
||
#[test]
|
||
fn test_hardware_budget_small_gpu() {
|
||
let budget = HardwareBudget::with_memory_mb(2048); // 2GB GPU
|
||
// 2048 * 0.70 = 1433.6; 1433.6 - 150 = 1283.6; 1283.6 / 0.05 = 25672
|
||
// clamped to 4096
|
||
let max = budget.max_batch_size(150.0, 0.05, 16.0, 4096.0).unwrap();
|
||
assert!((max - 4096.0).abs() < 0.01);
|
||
|
||
// With much heavier model overhead
|
||
// 1433.6 - 1500 = -66.4 → available <= 0 → returns min_batch
|
||
let max2 = budget.max_batch_size(1500.0, 0.05, 16.0, 4096.0).unwrap();
|
||
assert!((max2 - 16.0).abs() < 0.01);
|
||
}
|
||
|
||
#[test]
|
||
fn test_hardware_budget_clamps_to_min() {
|
||
let budget = HardwareBudget::with_memory_mb(256); // Tiny GPU
|
||
// 256 * 0.70 = 179.2; 179.2 - 200 = -20.8 → available <= 0 → returns min_batch
|
||
let max = budget.max_batch_size(200.0, 0.05, 16.0, 4096.0).unwrap();
|
||
assert!((max - 16.0).abs() < 0.01);
|
||
}
|
||
|
||
#[test]
|
||
fn test_continuous_bounds_for_default_delegates() {
|
||
// Default impl should return same as continuous_bounds()
|
||
let budget = HardwareBudget::with_memory_mb(81920);
|
||
let static_bounds = TestParams::continuous_bounds();
|
||
let dynamic_bounds = TestParams::continuous_bounds_for(&budget);
|
||
assert_eq!(static_bounds, dynamic_bounds);
|
||
}
|
||
|
||
#[test]
|
||
fn test_plan_hyperopt_tiny_model_high_concurrency() {
|
||
let budget = HardwareBudget { gpu_memory_mb: 24576, gpu_name: String::new() }; // L4 24GB
|
||
// 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 <= 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, gpu_name: String::new() }; // L4 24GB
|
||
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,
|
||
"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);
|
||
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, gpu_name: String::new() }; // H100 80GB
|
||
// 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");
|
||
}
|
||
}
|