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

feat(hyperopt): add VRAM-aware batch_size scaling to 9 existing adapters

Browse Source 
Override continuous_bounds_for() in all adapters that already have batch_size:
- TFT (idx 1): 150MB overhead, 0.05 MB/sample, cap 2048
- Mamba2 (idx 1): 200MB overhead, 0.03 MB/sample, cap 2048
- TGGN (idx 7): 100MB overhead, 0.08 MB/sample, cap 1024
- TLOB (idx 6): 120MB overhead, 0.10 MB/sample, cap 1024
- KAN (idx 7): 80MB overhead, 0.04 MB/sample, cap 1024
- xLSTM (idx 6): 180MB overhead, 0.06 MB/sample, cap 1024
- Diffusion (idx 6): 250MB overhead, 0.12 MB/sample, cap 512
- DQN (idx 1): 300MB overhead, 0.02 MB/sample, cap 4096
- ContinuousPPO (idx 9): 250MB overhead, 0.03 MB/sample, cap 4096

On CPU-only, all adapters fall back to existing static bounds.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-02-28 19:19:25 +01:00
parent de2169fbe4
commit b4239e8c43
9 changed files with 144 additions and 9 deletions
Download Patch File Download Diff File Expand all files Collapse all files

17
crates/ml/src/hyperopt/adapters/continuous_ppo.rs
View File

@@ -38,7 +38,7 @@ use std::io::Write as IoWrite;
use tracing::{info, warn};
use crate::hyperopt::paths::TrainingPaths;
use crate::hyperopt::traits::{HyperparameterOptimizable, ParameterSpace};
use crate::hyperopt::traits::{HardwareBudget, HyperparameterOptimizable, ParameterSpace};
use crate::ppo::continuous_ppo::{
ContinuousPPO, ContinuousPPOConfig, ContinuousTrajectory, ContinuousTrajectoryBatch,
ContinuousTrajectoryStep,
@@ -48,6 +48,11 @@ use crate::ppo::flow_policy::FlowPolicyConfig;
use crate::ppo::gae::GAEConfig;
use crate::MLError;
/// Continuous PPO model overhead in MB (Gaussian policy + value net + flow config)
const MODEL_OVERHEAD_MB: f64 = 250.0;
/// Continuous PPO per-sample memory in MB
const MB_PER_SAMPLE: f64 = 0.03;
/// Continuous PPO hyperparameter space
///
/// Defines the hyperparameters to optimize for continuous PPO training:
@@ -184,6 +189,16 @@ impl ParameterSpace for ContinuousPPOParams {
"learnable_std",
]
}
fn continuous_bounds_for(budget: &HardwareBudget) -> Vec<(f64, f64)> {
let mut bounds = Self::continuous_bounds();
if let Some(max_batch) = budget.max_batch_size(MODEL_OVERHEAD_MB, MB_PER_SAMPLE, 32.0, 4096.0) {
if let Some(batch_bound) = bounds.get_mut(9) {
batch_bound.1 = max_batch;
}
}
bounds
}
}
/// Continuous PPO training metrics

17
crates/ml/src/hyperopt/adapters/diffusion.rs
View File

@@ -12,10 +12,15 @@ use crate::diffusion::config::{DiffusionConfig, NoiseSchedule};
use crate::diffusion::trainable::DiffusionTrainableAdapter;
use crate::features::extract_ml_features;
use crate::hyperopt::paths::TrainingPaths;
use crate::hyperopt::traits::{HyperparameterOptimizable, ParameterSpace};
use crate::hyperopt::traits::{HardwareBudget, HyperparameterOptimizable, ParameterSpace};
use crate::training::unified_trainer::UnifiedTrainable;
use crate::MLError;
/// Diffusion model overhead in MB (weights + optimizer state, heavy due to denoiser)
const MODEL_OVERHEAD_MB: f64 = 250.0;
/// Diffusion per-sample memory in MB (noise + denoise activations)
const MB_PER_SAMPLE: f64 = 0.12;
/// Hyperparameters for Diffusion model hyperopt tuning.
#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
pub struct DiffusionParams {
@@ -110,6 +115,16 @@ impl ParameterSpace for DiffusionParams {
"batch_size", "weight_decay", "grad_clip",
]
}
fn continuous_bounds_for(budget: &HardwareBudget) -> Vec<(f64, f64)> {
let mut bounds = Self::continuous_bounds();
if let Some(max_batch) = budget.max_batch_size(MODEL_OVERHEAD_MB, MB_PER_SAMPLE, 4.0, 512.0) {
if let Some(batch_bound) = bounds.get_mut(6) {
batch_bound.1 = max_batch;
}
}
bounds
}
}
/// Metrics returned from Diffusion hyperopt training.

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

@@ -55,7 +55,7 @@ use chrono::Utc;
use crate::evaluation::engine::{Action, EvaluationEngine};
use crate::evaluation::metrics::{OHLCVBarF32, PerformanceMetrics};
use crate::hyperopt::paths::TrainingPaths;
use crate::hyperopt::traits::{HyperparameterOptimizable, ParameterSpace};
use crate::hyperopt::traits::{HardwareBudget, HyperparameterOptimizable, ParameterSpace};
use crate::trainers::dqn::{DQNHyperparameters, DQNTrainer as InternalDQNTrainer};
use crate::MLError;
@@ -135,6 +135,11 @@ pub struct BestTrialExport {
pub gradient_clip_norm: f64,
}
/// DQN Rainbow model overhead in MB (main + target + distributional atoms + optimizer)
const MODEL_OVERHEAD_MB: f64 = 300.0;
/// DQN per-sample memory in MB (54 features x 4 bytes x ~2 activation factor)
const MB_PER_SAMPLE: f64 = 0.02;
/// DQN hyperparameter space (17D continuous - WAVE 11 with Rainbow booleans hardcoded to TRUE)
///
/// Defines the hyperparameters to optimize for DQN training:
@@ -796,6 +801,16 @@ impl ParameterSpace for DQNParams {
// WAVE 11: Rainbow DQN boolean parameters REMOVED (always TRUE, not tunable)
]
}
fn continuous_bounds_for(budget: &HardwareBudget) -> Vec<(f64, f64)> {
let mut bounds = Self::continuous_bounds();
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 = max_batch;
}
}
bounds
}
}
impl DQNParams {

17
crates/ml/src/hyperopt/adapters/kan.rs
View File

@@ -11,12 +11,17 @@ use tracing::{info, warn};
use crate::features::extract_ml_features;
use crate::hyperopt::paths::TrainingPaths;
use crate::hyperopt::traits::{HyperparameterOptimizable, ParameterSpace};
use crate::hyperopt::traits::{HardwareBudget, HyperparameterOptimizable, ParameterSpace};
use crate::kan::config::KANConfig;
use crate::kan::trainable::KANTrainableAdapter;
use crate::training::unified_trainer::UnifiedTrainable;
use crate::MLError;
/// KAN model overhead in MB (weights + optimizer state)
const MODEL_OVERHEAD_MB: f64 = 80.0;
/// KAN per-sample memory in MB (spline-based, compact)
const MB_PER_SAMPLE: f64 = 0.04;
/// Hyperparameters for KAN hyperopt tuning.
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct KANParams {
@@ -113,6 +118,16 @@ impl ParameterSpace for KANParams {
"batch_size",
]
}
fn continuous_bounds_for(budget: &HardwareBudget) -> Vec<(f64, f64)> {
let mut bounds = Self::continuous_bounds();
if let Some(max_batch) = budget.max_batch_size(MODEL_OVERHEAD_MB, MB_PER_SAMPLE, 8.0, 1024.0) {
if let Some(batch_bound) = bounds.get_mut(7) {
batch_bound.1 = max_batch;
}
}
bounds
}
}
/// Metrics returned from KAN hyperopt training.

17
crates/ml/src/hyperopt/adapters/mamba2.rs
View File

@@ -41,10 +41,15 @@ use tracing::{info, warn};
use crate::features::{extract_ml_features, FeatureConfig};
use crate::hyperopt::paths::TrainingPaths;
use crate::hyperopt::traits::{HyperparameterOptimizable, ParameterSpace};
use crate::hyperopt::traits::{HardwareBudget, HyperparameterOptimizable, ParameterSpace};
use crate::mamba::{Mamba2Config, Mamba2SSM, OptimizerType};
use crate::MLError;
/// Mamba2 SSM model overhead in MB (weights + optimizer state)
const MODEL_OVERHEAD_MB: f64 = 200.0;
/// Mamba2 per-sample memory in MB (compact state-space model)
const MB_PER_SAMPLE: f64 = 0.03;
/// MAMBA-2 hyperparameter space
///
/// Defines the hyperparameters to optimize for MAMBA-2 training:
@@ -180,6 +185,16 @@ impl ParameterSpace for Mamba2Params {
"norm_eps",
]
}
fn continuous_bounds_for(budget: &HardwareBudget) -> Vec<(f64, f64)> {
let mut bounds = Self::continuous_bounds();
if let Some(max_batch) = budget.max_batch_size(MODEL_OVERHEAD_MB, MB_PER_SAMPLE, 4.0, 2048.0) {
if let Some(batch_bound) = bounds.get_mut(1) {
batch_bound.1 = max_batch;
}
}
bounds
}
}
/// MAMBA-2 training metrics

17
crates/ml/src/hyperopt/adapters/tft.rs
View File

@@ -40,10 +40,15 @@ use std::path::PathBuf;
use tracing::{info, warn};
use crate::hyperopt::paths::TrainingPaths;
use crate::hyperopt::traits::{HyperparameterOptimizable, ParameterSpace};
use crate::hyperopt::traits::{HardwareBudget, HyperparameterOptimizable, ParameterSpace};
use crate::trainers::tft::{TFTTrainer as RealTFTTrainer, TFTTrainerConfig};
use crate::MLError;
/// TFT model overhead in MB (weights + optimizer state)
const MODEL_OVERHEAD_MB: f64 = 150.0;
/// TFT per-sample memory in MB (225 features x 32 seq x 4 bytes x ~2 activation factor)
const MB_PER_SAMPLE: f64 = 0.05;
/// TFT hyperparameter space
///
/// Defines the hyperparameters to optimize for TFT training:
@@ -154,6 +159,16 @@ impl ParameterSpace for TFTParams {
"dropout",
]
}
fn continuous_bounds_for(budget: &HardwareBudget) -> Vec<(f64, f64)> {
let mut bounds = Self::continuous_bounds();
if let Some(max_batch) = budget.max_batch_size(MODEL_OVERHEAD_MB, MB_PER_SAMPLE, 16.0, 2048.0) {
if let Some(batch_bound) = bounds.get_mut(1) {
batch_bound.1 = max_batch;
}
}
bounds
}
}
/// TFT training metrics

17
crates/ml/src/hyperopt/adapters/tggn.rs
View File

@@ -11,12 +11,17 @@ use tracing::{info, warn};
use crate::features::extract_ml_features;
use crate::hyperopt::paths::TrainingPaths;
use crate::hyperopt::traits::{HyperparameterOptimizable, ParameterSpace};
use crate::hyperopt::traits::{HardwareBudget, HyperparameterOptimizable, ParameterSpace};
use crate::tgnn::TGGNConfig;
use crate::tgnn::trainable_adapter::TGGNTrainableAdapter;
use crate::training::unified_trainer::UnifiedTrainable;
use crate::MLError;
/// TGGN graph network model overhead in MB (weights + optimizer state)
const MODEL_OVERHEAD_MB: f64 = 100.0;
/// TGGN per-sample memory in MB (graph edges + node features)
const MB_PER_SAMPLE: f64 = 0.08;
/// Hyperparameters for TGGN hyperopt tuning.
///
/// Controls the temporal graph neural network's architecture and training:
@@ -133,6 +138,16 @@ impl ParameterSpace for TGGNParams {
"grad_clip",
]
}
fn continuous_bounds_for(budget: &HardwareBudget) -> Vec<(f64, f64)> {
let mut bounds = Self::continuous_bounds();
if let Some(max_batch) = budget.max_batch_size(MODEL_OVERHEAD_MB, MB_PER_SAMPLE, 8.0, 1024.0) {
if let Some(batch_bound) = bounds.get_mut(7) {
batch_bound.1 = max_batch;
}
}
bounds
}
}
/// Metrics returned from TGGN hyperopt training.

17
crates/ml/src/hyperopt/adapters/tlob.rs
View File

@@ -11,11 +11,16 @@ use tracing::{info, warn};
use crate::features::extract_ml_features;
use crate::hyperopt::paths::TrainingPaths;
use crate::hyperopt::traits::{HyperparameterOptimizable, ParameterSpace};
use crate::hyperopt::traits::{HardwareBudget, HyperparameterOptimizable, ParameterSpace};
use crate::tlob::trainable_adapter::{TLOBAdapterConfig, TLOBTrainableAdapter};
use crate::training::unified_trainer::UnifiedTrainable;
use crate::MLError;
/// TLOB transformer model overhead in MB (weights + optimizer state)
const MODEL_OVERHEAD_MB: f64 = 120.0;
/// TLOB per-sample memory in MB (LOB depth x features x seq)
const MB_PER_SAMPLE: f64 = 0.10;
/// Hyperparameters for TLOB hyperopt tuning.
///
/// Controls the TLOB transformer's architecture and training:
@@ -124,6 +129,16 @@ impl ParameterSpace for TLOBParams {
"grad_clip",
]
}
fn continuous_bounds_for(budget: &HardwareBudget) -> Vec<(f64, f64)> {
let mut bounds = Self::continuous_bounds();
if let Some(max_batch) = budget.max_batch_size(MODEL_OVERHEAD_MB, MB_PER_SAMPLE, 8.0, 1024.0) {
if let Some(batch_bound) = bounds.get_mut(6) {
batch_bound.1 = max_batch;
}
}
bounds
}
}
/// Metrics returned from TLOB hyperopt training.

17
crates/ml/src/hyperopt/adapters/xlstm.rs
View File

@@ -10,12 +10,17 @@ use tracing::{info, warn};
use crate::features::extract_ml_features;
use crate::hyperopt::paths::TrainingPaths;
use crate::hyperopt::traits::{HyperparameterOptimizable, ParameterSpace};
use crate::hyperopt::traits::{HardwareBudget, HyperparameterOptimizable, ParameterSpace};
use crate::training::unified_trainer::UnifiedTrainable;
use crate::xlstm::config::XLSTMConfig;
use crate::xlstm::trainable::XLSTMTrainableAdapter;
use crate::MLError;
/// xLSTM model overhead in MB (weights + optimizer state)
const MODEL_OVERHEAD_MB: f64 = 180.0;
/// xLSTM per-sample memory in MB (extended LSTM states)
const MB_PER_SAMPLE: f64 = 0.06;
/// Hyperparameters for xLSTM hyperopt tuning.
#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
pub struct XLSTMParams {
@@ -104,6 +109,16 @@ impl ParameterSpace for XLSTMParams {
"slstm_ratio", "dropout", "batch_size", "weight_decay", "grad_clip",
]
}
fn continuous_bounds_for(budget: &HardwareBudget) -> Vec<(f64, f64)> {
let mut bounds = Self::continuous_bounds();
if let Some(max_batch) = budget.max_batch_size(MODEL_OVERHEAD_MB, MB_PER_SAMPLE, 8.0, 1024.0) {
if let Some(batch_bound) = bounds.get_mut(6) {
batch_bound.1 = max_batch;
}
}
bounds
}
}
/// Metrics returned from xLSTM hyperopt training.