680bb7df0b
The entire foxhunt workspace compiles with BF16: - 37/37 CUDA kernels: __nv_bfloat16 parameters + native arithmetic - All Rust types: CudaSlice<half::bf16> across all crates - ml-core: 0 errors (nvrtc removed, BF16 boundaries fixed) - ml-dqn: 0 errors (noisy layers, replay buffer, branching → BF16) - ml-ppo: 0 errors (stubbed, cold path) - ml-supervised: 0 errors - ml-ensemble, ml-explainability: 0 errors - ml: 0 errors (22 files fixed, BF16 conversion helpers added) BF16 conversion at system boundary only: - Host f32 data → half::bf16::from_f32() → GPU upload - GPU download → bf16.to_f32() → host f32 Remaining Phase 4: wire cublasGemmEx + run tests. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
1317 lines
49 KiB
Rust
1317 lines
49 KiB
Rust
//! MAMBA-2 Hyperparameter Optimization Adapter
|
|
//!
|
|
//! This module provides a production-ready adapter for optimizing MAMBA-2
|
|
//! hyperparameters using the generic optimization framework. It implements:
|
|
//!
|
|
//! - Parameter space with log-scale handling for learning rates
|
|
//! - Training wrapper that integrates with existing MAMBA-2 pipeline
|
|
//! - Metrics extraction for validation loss optimization
|
|
//!
|
|
//! ## Usage Example
|
|
//!
|
|
//! ```rust,no_run
|
|
//! use ml::hyperopt::ArgminOptimizer;
|
|
//! use ml::hyperopt::adapters::mamba2::{Mamba2Trainer, Mamba2Params};
|
|
//!
|
|
//! # async fn example() -> anyhow::Result<()> {
|
|
//! // Create trainer
|
|
//! let trainer = Mamba2Trainer::new(
|
|
//! "test_data/ES_FUT_180d.parquet",
|
|
//! 50, // epochs per trial
|
|
//! )?;
|
|
//!
|
|
//! // Run optimization
|
|
//! let optimizer = ArgminOptimizer::with_trials(30, 5);
|
|
//! let result = optimizer.optimize(trainer)?;
|
|
//!
|
|
//! println!("Best learning rate: {}", result.best_params.learning_rate);
|
|
//! println!("Best batch size: {}", result.best_params.batch_size);
|
|
//! println!("Best validation loss: {:.6}", result.best_objective);
|
|
//! # Ok(())
|
|
//! # }
|
|
//! ```
|
|
|
|
use anyhow::{Context, Result};
|
|
use ml_core::device::MlDevice;
|
|
use ml_core::cuda_autograd::GpuTensor;
|
|
use ml_core::cuda_autograd::stream_ops::StreamTensor;
|
|
use serde::{Deserialize, Serialize};
|
|
use std::fs::OpenOptions;
|
|
use std::io::Write as IoWrite;
|
|
use std::path::PathBuf;
|
|
use std::sync::Arc;
|
|
use tracing::{debug, info, warn};
|
|
|
|
use crate::features::extraction::OHLCVBar;
|
|
use crate::features::{extract_ml_features, FeatureConfig};
|
|
use crate::hyperopt::paths::TrainingPaths;
|
|
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)
|
|
/// Mamba2 with d_model=128: ~500K params, state-space model much lighter than attention
|
|
const MODEL_OVERHEAD_MB: f64 = 80.0;
|
|
/// Mamba2 per-sample memory in MB (compact state-space recurrence, no attention)
|
|
const MB_PER_SAMPLE: f64 = 0.005;
|
|
|
|
/// MAMBA-2 hyperparameter space
|
|
///
|
|
/// Defines the hyperparameters to optimize for MAMBA-2 training:
|
|
/// - Learning rate (log-scale: 1e-5 to 1e-2)
|
|
/// - Batch size (linear scale: 4 to 256, optimized for GPU utilization)
|
|
/// - Dropout rate (linear scale: 0.0 to 0.5)
|
|
/// - Weight decay (log-scale: 1e-6 to 1e-2)
|
|
///
|
|
/// ## Parameter Scaling
|
|
///
|
|
/// - **Log-scale**: Learning rate, weight decay (span multiple orders of magnitude)
|
|
/// - **Linear scale**: Batch size, dropout (span single order of magnitude)
|
|
///
|
|
/// This scaling ensures efficient exploration by egobox's Gaussian Process.
|
|
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
|
|
pub struct Mamba2Params {
|
|
/// Learning rate for Adam optimizer (log-scale)
|
|
pub learning_rate: f64,
|
|
/// Batch size for training (linear scale, integer)
|
|
pub batch_size: usize,
|
|
/// Dropout rate for regularization (linear scale)
|
|
pub dropout: f64,
|
|
/// Weight decay for L2 regularization (log-scale)
|
|
pub weight_decay: f64,
|
|
/// Gradient clipping threshold (log-scale)
|
|
pub grad_clip: f64,
|
|
/// Warmup steps for learning rate schedule (linear scale)
|
|
pub warmup_steps: usize,
|
|
/// P0: Adam beta1 momentum parameter (linear scale)
|
|
pub adam_beta1: f64,
|
|
/// P1: Adam beta2 parameter (linear scale)
|
|
pub adam_beta2: f64,
|
|
/// P1: Adam epsilon (log-scale)
|
|
pub adam_epsilon: f64,
|
|
/// P2: Lookback window (sequence length) (linear scale, integer)
|
|
pub lookback_window: usize,
|
|
/// P2: Sequence stride for overlapping windows (linear scale, integer)
|
|
pub sequence_stride: usize,
|
|
/// P2: Normalization epsilon for layer norm (log-scale)
|
|
pub norm_eps: f64,
|
|
/// Signal high threshold in basis points (GPU backtest action mapping)
|
|
pub signal_high_bps: f64,
|
|
/// Signal low threshold in basis points (GPU backtest action mapping)
|
|
pub signal_low_bps: f64,
|
|
}
|
|
|
|
impl Default for Mamba2Params {
|
|
fn default() -> Self {
|
|
Self {
|
|
learning_rate: 1e-4,
|
|
batch_size: 32,
|
|
dropout: 0.1,
|
|
weight_decay: 1e-4,
|
|
grad_clip: 1.0,
|
|
warmup_steps: 100,
|
|
adam_beta1: 0.9,
|
|
adam_beta2: 0.999,
|
|
adam_epsilon: 1e-8,
|
|
lookback_window: 60,
|
|
sequence_stride: 1,
|
|
norm_eps: 1e-5,
|
|
signal_high_bps: 10.0,
|
|
signal_low_bps: 5.0,
|
|
}
|
|
}
|
|
}
|
|
|
|
impl ParameterSpace for Mamba2Params {
|
|
fn continuous_bounds() -> Vec<(f64, f64)> {
|
|
vec![
|
|
(1e-5_f64.ln(), 1e-2_f64.ln()), // learning_rate (log scale)
|
|
(4.0, 1024.0), // batch_size (linear) - wide bounds, capped by VRAM in continuous_bounds_for
|
|
(0.0, 0.5), // dropout (linear)
|
|
(1e-6_f64.ln(), 1e-2_f64.ln()), // weight_decay (log scale)
|
|
(0.5_f64.ln(), 5.0_f64.ln()), // grad_clip (log scale)
|
|
(100.0, 2000.0), // warmup_steps (linear)
|
|
(0.85, 0.95), // adam_beta1 (linear)
|
|
(0.98, 0.999), // adam_beta2 (linear)
|
|
(1e-9_f64.ln(), 1e-7_f64.ln()), // adam_epsilon (log scale)
|
|
(30.0, 120.0), // lookback_window (linear)
|
|
(1.0, 5.0), // sequence_stride (linear)
|
|
(1e-6_f64.ln(), 1e-4_f64.ln()), // norm_eps (log scale)
|
|
(1.0, 30.0), // signal_high_bps
|
|
(0.5, 15.0), // signal_low_bps
|
|
]
|
|
}
|
|
|
|
fn from_continuous(x: &[f64]) -> Result<Self, MLError> {
|
|
if x.len() != 14 {
|
|
return Err(MLError::ConfigError(format!("Expected 14 parameters, got {}", x.len())));
|
|
}
|
|
|
|
Ok(Self {
|
|
learning_rate: x[0].exp(),
|
|
batch_size: x[1].round().max(1.0) as usize,
|
|
dropout: x[2].clamp(0.0, 0.5),
|
|
weight_decay: x[3].exp(),
|
|
grad_clip: x[4].exp(),
|
|
warmup_steps: x[5].round().max(1.0) as usize,
|
|
adam_beta1: x[6].clamp(0.85, 0.95),
|
|
adam_beta2: x[7].clamp(0.98, 0.999),
|
|
adam_epsilon: x[8].exp(),
|
|
lookback_window: x[9].round().max(30.0).min(120.0) as usize,
|
|
sequence_stride: x[10].round().max(1.0).min(5.0) as usize,
|
|
norm_eps: x[11].exp(),
|
|
signal_high_bps: x[12].clamp(1.0, 30.0),
|
|
signal_low_bps: x[13].clamp(0.5, 15.0),
|
|
})
|
|
}
|
|
|
|
fn to_continuous(&self) -> Vec<f64> {
|
|
vec![
|
|
self.learning_rate.ln(),
|
|
self.batch_size as f64,
|
|
self.dropout,
|
|
self.weight_decay.ln(),
|
|
self.grad_clip.ln(),
|
|
self.warmup_steps as f64,
|
|
self.adam_beta1,
|
|
self.adam_beta2,
|
|
self.adam_epsilon.ln(),
|
|
self.lookback_window as f64,
|
|
self.sequence_stride as f64,
|
|
self.norm_eps.ln(),
|
|
self.signal_high_bps,
|
|
self.signal_low_bps,
|
|
]
|
|
}
|
|
|
|
fn param_names() -> Vec<&'static str> {
|
|
vec![
|
|
"learning_rate",
|
|
"batch_size",
|
|
"dropout",
|
|
"weight_decay",
|
|
"grad_clip",
|
|
"warmup_steps",
|
|
"adam_beta1",
|
|
"adam_beta2",
|
|
"adam_epsilon",
|
|
"lookback_window",
|
|
"sequence_stride",
|
|
"norm_eps",
|
|
"signal_high_bps",
|
|
"signal_low_bps",
|
|
]
|
|
}
|
|
|
|
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, 4096.0) {
|
|
if let Some(batch_bound) = bounds.get_mut(1) {
|
|
batch_bound.1 = max_batch;
|
|
}
|
|
}
|
|
bounds
|
|
}
|
|
}
|
|
|
|
/// MAMBA-2 training metrics
|
|
///
|
|
/// Contains all relevant metrics from a MAMBA-2 training run.
|
|
/// The primary optimization target is validation loss.
|
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
|
pub struct Mamba2Metrics {
|
|
/// Final validation loss (optimization target)
|
|
pub val_loss: f64,
|
|
/// Final training loss
|
|
pub train_loss: f64,
|
|
/// Validation perplexity (exp(val_loss))
|
|
pub val_perplexity: f64,
|
|
/// Directional accuracy (% of correct direction predictions)
|
|
pub directional_accuracy: f64,
|
|
/// Mean Absolute Error
|
|
pub mae: f64,
|
|
/// Root Mean Squared Error
|
|
pub rmse: f64,
|
|
/// R² (Coefficient of Determination)
|
|
pub r_squared: f64,
|
|
/// Number of epochs completed
|
|
pub epochs_completed: usize,
|
|
/// GPU walk-forward backtest Sharpe ratio
|
|
pub backtest_sharpe: Option<f32>,
|
|
/// GPU walk-forward backtest total trades
|
|
pub backtest_trades: Option<u32>,
|
|
}
|
|
|
|
/// MAMBA-2 trainer for hyperparameter optimization
|
|
///
|
|
/// This struct wraps the MAMBA-2 training pipeline and implements
|
|
/// `HyperparameterOptimizable` for use with `ArgminOptimizer`.
|
|
///
|
|
/// ## Configuration
|
|
///
|
|
/// - **Data dir**: Directory with .dbn.zst files (Databento OHLCV data)
|
|
/// - **Epochs**: Number of training epochs per trial
|
|
/// - **Device**: CUDA GPU (falls back to CPU if unavailable)
|
|
/// - **Features**: 42 market features from extraction pipeline
|
|
///
|
|
/// ## Optimized Hyperparameters
|
|
///
|
|
/// The following are optimized by `Mamba2Params`:
|
|
/// - Learning rate
|
|
/// - Batch size
|
|
/// - Dropout
|
|
/// - Weight decay
|
|
#[derive(Debug)]
|
|
pub struct Mamba2Trainer {
|
|
data_dir: PathBuf,
|
|
epochs: usize,
|
|
device: MlDevice,
|
|
feature_config: FeatureConfig,
|
|
d_model: usize,
|
|
train_split: f64,
|
|
target_min: Option<f64>,
|
|
target_max: Option<f64>,
|
|
batch_size_min: f64,
|
|
batch_size_max: f64,
|
|
async_loading: bool,
|
|
prefetch_count: usize,
|
|
training_paths: TrainingPaths,
|
|
early_stopping_patience: usize,
|
|
early_stopping_min_epochs: usize,
|
|
trial_counter: usize,
|
|
preloaded_bars: Option<Arc<[OHLCVBar]>>,
|
|
}
|
|
|
|
impl Mamba2Trainer {
|
|
/// Create a new MAMBA-2 trainer
|
|
///
|
|
/// # Arguments
|
|
///
|
|
/// * `data_dir` - Directory containing .dbn.zst files with OHLCV data
|
|
/// * `epochs` - Number of training epochs per trial
|
|
pub fn new<P: Into<PathBuf>>(data_dir: P, epochs: usize) -> Result<Self> {
|
|
let data_dir = data_dir.into();
|
|
|
|
if !data_dir.exists() {
|
|
return Err(MLError::ConfigError(format!("Data directory not found: {}", data_dir.display()))
|
|
.into());
|
|
}
|
|
|
|
let device = MlDevice::cuda(0)
|
|
.map_err(|e| MLError::ConfigError(format!("CUDA GPU required for Mamba2 hyperopt: {}", e)))?;
|
|
|
|
let feature_config = FeatureConfig::wave_d();
|
|
let d_model = feature_config.feature_count();
|
|
|
|
info!("MAMBA-2 Trainer initialized:");
|
|
info!(" Device: {:?}", device);
|
|
info!(" Data: {}", data_dir.display());
|
|
info!(" Features: {} (42 market + 3 portfolio)", d_model);
|
|
info!(" Epochs per trial: {}", epochs);
|
|
|
|
let training_paths = TrainingPaths::new("/tmp/ml_training", "mamba2", "default");
|
|
|
|
Ok(Self {
|
|
data_dir,
|
|
epochs,
|
|
device,
|
|
feature_config,
|
|
d_model,
|
|
train_split: 0.8,
|
|
target_min: None,
|
|
target_max: None,
|
|
batch_size_min: 4.0,
|
|
batch_size_max: 96.0,
|
|
async_loading: true,
|
|
prefetch_count: 3,
|
|
training_paths,
|
|
early_stopping_patience: 5,
|
|
early_stopping_min_epochs: 5,
|
|
trial_counter: 0,
|
|
preloaded_bars: None,
|
|
})
|
|
}
|
|
|
|
/// Configure early stopping parameters (overrides defaults)
|
|
pub fn with_early_stopping(mut self, patience: usize, min_epochs: usize) -> Self {
|
|
self.early_stopping_patience = patience;
|
|
self.early_stopping_min_epochs = min_epochs;
|
|
self
|
|
}
|
|
|
|
/// Set train/validation split ratio
|
|
pub fn with_train_split(mut self, split: f64) -> Self {
|
|
assert!(split > 0.0 && split < 1.0, "Split must be in (0, 1)");
|
|
self.train_split = split;
|
|
self
|
|
}
|
|
|
|
/// Set batch size bounds for GPU memory constraints
|
|
///
|
|
/// # Arguments
|
|
///
|
|
/// * `min` - Minimum batch size (must be >= 1)
|
|
/// * `max` - Maximum batch size (must be > min)
|
|
///
|
|
/// # Example
|
|
///
|
|
/// ```no_run
|
|
/// # use ml::hyperopt::adapters::mamba2::Mamba2Trainer;
|
|
/// // Configure for RTX 4090 (24GB VRAM)
|
|
/// let trainer = Mamba2Trainer::new("data.parquet", 50)?
|
|
/// .with_batch_size_bounds(4.0, 1024.0);
|
|
///
|
|
/// // Configure for RTX 3050 Ti (4GB VRAM)
|
|
/// let trainer = Mamba2Trainer::new("data.parquet", 50)?
|
|
/// .with_batch_size_bounds(4.0, 32.0);
|
|
/// # Ok::<(), anyhow::Error>(())
|
|
/// ```
|
|
pub fn with_batch_size_bounds(mut self, min: f64, max: f64) -> Self {
|
|
assert!(min >= 1.0, "Minimum batch size must be >= 1");
|
|
assert!(max > min, "Maximum batch size must be > minimum");
|
|
info!("Configuring batch_size bounds: [{}, {}]", min, max);
|
|
self.batch_size_min = min;
|
|
self.batch_size_max = max;
|
|
self
|
|
}
|
|
|
|
/// Enable or disable async data loading (prefetching)
|
|
///
|
|
/// When enabled, batches are prepared on CPU and transferred to GPU in a
|
|
/// background thread while the GPU trains on the current batch. This improves
|
|
/// GPU utilization from ~78% to ~90-95% and reduces training time by 20-30%.
|
|
///
|
|
/// # Arguments
|
|
///
|
|
/// * `enabled` - Enable async loading
|
|
/// * `prefetch_count` - Number of batches to prefetch (2-3 recommended)
|
|
///
|
|
/// # Example
|
|
///
|
|
/// ```no_run
|
|
/// # use ml::hyperopt::adapters::mamba2::Mamba2Trainer;
|
|
/// // Enable async loading with 3 batch prefetch
|
|
/// let trainer = Mamba2Trainer::new("data.parquet", 50)?
|
|
/// .with_async_loading(true, 3);
|
|
///
|
|
/// // Disable async loading (sync mode)
|
|
/// let trainer = Mamba2Trainer::new("data.parquet", 50)?
|
|
/// .with_async_loading(false, 0);
|
|
/// # Ok::<(), anyhow::Error>(())
|
|
/// ```
|
|
pub fn with_async_loading(mut self, enabled: bool, prefetch_count: usize) -> Self {
|
|
if enabled {
|
|
assert!(
|
|
prefetch_count >= 2,
|
|
"Prefetch count must be >= 2 when async loading enabled"
|
|
);
|
|
assert!(
|
|
prefetch_count <= 10,
|
|
"Prefetch count must be <= 10 to avoid excessive memory"
|
|
);
|
|
}
|
|
info!(
|
|
"Configuring async data loading: enabled={}, prefetch={}",
|
|
enabled, prefetch_count
|
|
);
|
|
self.async_loading = enabled;
|
|
self.prefetch_count = prefetch_count;
|
|
self
|
|
}
|
|
|
|
/// Set checkpoint directory for saving best models
|
|
///
|
|
/// # Arguments
|
|
///
|
|
/// * `dir` - Path to checkpoint directory (will be created if it doesn't exist)
|
|
///
|
|
/// # Returns
|
|
///
|
|
/// Self for method chaining
|
|
///
|
|
/// # Deprecated
|
|
///
|
|
/// Use `with_training_paths()` instead for full path management
|
|
pub fn with_checkpoint_dir<P: Into<PathBuf>>(mut self, dir: P) -> Self {
|
|
// Update training_paths to use provided checkpoint directory
|
|
let checkpoint_dir = dir.into();
|
|
self.training_paths = TrainingPaths::new(
|
|
checkpoint_dir.parent().unwrap_or(&checkpoint_dir),
|
|
"mamba2",
|
|
"legacy",
|
|
);
|
|
info!(
|
|
"Checkpoint directory set to: {:?}",
|
|
self.training_paths.checkpoints_dir()
|
|
);
|
|
self
|
|
}
|
|
|
|
/// Set training paths configuration (recommended over with_checkpoint_dir)
|
|
///
|
|
/// # Arguments
|
|
///
|
|
/// * `paths` - Training paths configuration
|
|
///
|
|
/// # Returns
|
|
///
|
|
/// Self for method chaining
|
|
pub fn with_training_paths(mut self, paths: TrainingPaths) -> Self {
|
|
info!("Training paths configured:");
|
|
info!(" Run directory: {:?}", paths.run_dir());
|
|
info!(" Checkpoints: {:?}", paths.checkpoints_dir());
|
|
info!(" Logs: {:?}", paths.logs_dir());
|
|
info!(" Hyperopt: {:?}", paths.hyperopt_dir());
|
|
self.training_paths = paths;
|
|
self
|
|
}
|
|
|
|
/// Denormalize a prediction from [0,1] to original price scale
|
|
///
|
|
/// # Arguments
|
|
///
|
|
/// * `normalized` - Normalized prediction in [0,1] range
|
|
///
|
|
/// # Returns
|
|
///
|
|
/// `Some(price)` in original scale (e.g., $5000-6000 for ES futures),
|
|
/// or `None` if normalization params have not been set (i.e., training has not run yet).
|
|
pub fn denormalize_prediction(&self, normalized: f64) -> Option<f64> {
|
|
let min = self.target_min?;
|
|
let max = self.target_max?;
|
|
|
|
Some(normalized * (max - min) + min)
|
|
}
|
|
|
|
/// Preload DBN data once for reuse across all hyperopt trials.
|
|
pub fn preload_data(&mut self) -> Result<(), MLError> {
|
|
if self.preloaded_bars.is_some() {
|
|
return Ok(());
|
|
}
|
|
let bars = super::dbn_loader::load_bars_from_dbn_dir(&self.data_dir)
|
|
.map_err(|e| MLError::ModelError(format!("Failed to preload DBN data: {}", e)))?;
|
|
info!("Preloaded {} OHLCV bars for reuse across trials", bars.len());
|
|
self.preloaded_bars = Some(Arc::from(bars));
|
|
Ok(())
|
|
}
|
|
|
|
/// Check if data has been preloaded.
|
|
pub fn has_preloaded_data(&self) -> bool {
|
|
self.preloaded_bars.is_some()
|
|
}
|
|
|
|
/// Load and prepare training data from Parquet
|
|
///
|
|
/// Reads OHLCV bars, extracts features, creates sequences.
|
|
fn load_and_prepare_data(
|
|
&self,
|
|
seq_len: usize,
|
|
_stride: usize,
|
|
) -> Result<(Vec<(GpuTensor, GpuTensor)>, Vec<(GpuTensor, GpuTensor)>, f64, f64)> {
|
|
// Load bars from DBN files (use preloaded cache if available)
|
|
let all_ohlcv_bars = if let Some(ref preloaded) = self.preloaded_bars {
|
|
preloaded.to_vec() // cpu-side Arc<[OHLCVBar]> clone
|
|
} else {
|
|
super::dbn_loader::load_bars_from_dbn_dir(&self.data_dir)
|
|
.context("Failed to load DBN data")?
|
|
};
|
|
|
|
if all_ohlcv_bars.len() < seq_len + 1 {
|
|
return Err(MLError::ModelError(format!(
|
|
"Insufficient data: {} bars (need at least {} for seq_len={})",
|
|
all_ohlcv_bars.len(),
|
|
seq_len + 1,
|
|
seq_len
|
|
))
|
|
.into());
|
|
}
|
|
|
|
// Extract features
|
|
let features =
|
|
extract_ml_features(&all_ohlcv_bars).context("Failed to extract features")?;
|
|
|
|
if features.is_empty() {
|
|
return Err(
|
|
MLError::ModelError("No features extracted from data".to_owned()).into(),
|
|
);
|
|
}
|
|
|
|
// P0 FIX: Collect all target prices for normalization
|
|
let mut all_target_prices = Vec::new();
|
|
for window_idx in 0..features.len().saturating_sub(seq_len) {
|
|
let target_price = all_ohlcv_bars[window_idx + seq_len].close;
|
|
all_target_prices.push(target_price);
|
|
}
|
|
|
|
// Compute normalization parameters
|
|
let target_min = all_target_prices
|
|
.iter()
|
|
.copied()
|
|
.fold(f64::INFINITY, f64::min);
|
|
let target_max = all_target_prices
|
|
.iter()
|
|
.copied()
|
|
.fold(f64::NEG_INFINITY, f64::max);
|
|
|
|
if (target_max - target_min).abs() < 1e-6 {
|
|
return Err(MLError::ModelError(
|
|
"Target prices have zero variance - cannot normalize".to_owned(),
|
|
)
|
|
.into());
|
|
}
|
|
|
|
info!(
|
|
"Target normalization: min={:.2}, max={:.2}, range={:.2}",
|
|
target_min,
|
|
target_max,
|
|
target_max - target_min
|
|
);
|
|
|
|
// FIX: Apply percentile clipping BEFORE normalization to prevent outliers
|
|
// (e.g., OBV features with extreme values) from crushing other features
|
|
let all_feature_values: Vec<f64> =
|
|
features.iter().flat_map(|f| f.iter().copied()).collect();
|
|
|
|
// Compute 1st and 99th percentiles
|
|
let mut sorted_features = all_feature_values.clone();
|
|
sorted_features.sort_by(|a, b| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal));
|
|
|
|
let p1_idx = (sorted_features.len() as f64 * 0.01).round() as usize;
|
|
let p99_idx = (sorted_features.len() as f64 * 0.99).round() as usize;
|
|
let p1 = sorted_features[p1_idx.min(sorted_features.len() - 1)];
|
|
let p99 = sorted_features[p99_idx.min(sorted_features.len() - 1)];
|
|
|
|
info!("Feature percentile clipping: p1={:.2}, p99={:.2}", p1, p99);
|
|
|
|
// Clip outliers to [p1, p99] range
|
|
let clipped_feature_values: Vec<f64> = all_feature_values
|
|
.iter()
|
|
.map(|&x| x.clamp(p1, p99))
|
|
.collect();
|
|
|
|
// Now compute normalization parameters from clipped data
|
|
let feature_min = clipped_feature_values
|
|
.iter()
|
|
.copied()
|
|
.fold(f64::INFINITY, f64::min);
|
|
let feature_max = clipped_feature_values
|
|
.iter()
|
|
.copied()
|
|
.fold(f64::NEG_INFINITY, f64::max);
|
|
|
|
if (feature_max - feature_min).abs() < 1e-6 {
|
|
return Err(MLError::ModelError(
|
|
"Features have zero variance - cannot normalize".to_owned(),
|
|
)
|
|
.into());
|
|
}
|
|
|
|
info!(
|
|
"Feature normalization (after clipping): min={:.2}, max={:.2}, range={:.2}",
|
|
feature_min,
|
|
feature_max,
|
|
feature_max - feature_min
|
|
);
|
|
|
|
// Create sequences with normalized features and targets
|
|
let mut feature_sequences = Vec::new();
|
|
|
|
for (window_idx, &target_price) in all_target_prices.iter().enumerate() {
|
|
// FIX: Apply percentile clipping + normalization to [0, 1] range
|
|
let sequence: Vec<f64> = features[window_idx..window_idx + seq_len]
|
|
.iter()
|
|
.flat_map(|f| f.iter().copied())
|
|
.map(|val| {
|
|
// Clip to percentile range, then normalize
|
|
let clipped = val.clamp(p1, p99);
|
|
(clipped - feature_min) / (feature_max - feature_min)
|
|
})
|
|
.collect();
|
|
|
|
// Normalize target to [0,1]
|
|
let normalized_target = (target_price - target_min) / (target_max - target_min);
|
|
|
|
let sequence_f32: Vec<f32> = sequence.iter().map(|&v| v as f32).collect();
|
|
let stream = self.device.cuda_stream()?;
|
|
let input_tensor = GpuTensor::from_host(&sequence_f32, vec![sequence_f32.len()], stream)?.reshape(
|
|
vec![1, seq_len, self.d_model],
|
|
)?;
|
|
let target_tensor =
|
|
GpuTensor::from_host(&[normalized_target as f32], vec![1], stream)?.reshape(vec![1, 1, 1])?;
|
|
|
|
feature_sequences.push((input_tensor, target_tensor));
|
|
}
|
|
|
|
// Split train/validation
|
|
let split_idx = (feature_sequences.len() as f64 * self.train_split) as usize;
|
|
let train_data = feature_sequences[..split_idx].to_vec(); // cpu-side split
|
|
let val_data = feature_sequences[split_idx..].to_vec(); // cpu-side split
|
|
|
|
// P1 FIX: Validate validation set size
|
|
if val_data.len() < 10 {
|
|
return Err(MLError::ModelError(format!(
|
|
"Validation set too small: {} samples (need at least 10)",
|
|
val_data.len()
|
|
))
|
|
.into());
|
|
}
|
|
|
|
Ok((train_data, val_data, target_min, target_max))
|
|
}
|
|
|
|
/// Train model with async data loading (optimized for GPU utilization)
|
|
///
|
|
/// This method uses AsyncDataLoader to prefetch batches while GPU trains,
|
|
/// improving GPU utilization from ~78% to ~90-95% and reducing training
|
|
/// time by 20-30%.
|
|
///
|
|
/// # Architecture
|
|
///
|
|
/// ```text
|
|
/// CPU Thread (Background): GPU Thread (Main):
|
|
/// ┌─────────────────┐ ┌─────────────────┐
|
|
/// │ Load batch N+1 │ ────────> │ Train batch N │
|
|
/// │ Concat tensors │ │ Forward pass │
|
|
/// │ Transfer to GPU │ │ Backward pass │
|
|
/// └─────────────────┘ │ Optimizer step │
|
|
/// │ └─────────────────┘
|
|
/// ▼ │
|
|
/// ┌─────────────────┐ │
|
|
/// │ Load batch N+2 │ │
|
|
/// │ ... │ <───────────────────┘
|
|
/// └─────────────────┘
|
|
/// ```
|
|
///
|
|
/// # Arguments
|
|
///
|
|
/// * `model` - MAMBA-2 model to train
|
|
/// * `train_data` - Training data (already prepared tensors)
|
|
/// * `val_data` - Validation data
|
|
/// * `epochs` - Number of training epochs
|
|
/// * `batch_size` - Batch size for training
|
|
///
|
|
/// # Returns
|
|
///
|
|
/// Training history with metrics per epoch
|
|
async fn train_with_async_loading(
|
|
&self,
|
|
model: &mut Mamba2SSM,
|
|
train_data: &[(GpuTensor, GpuTensor)],
|
|
val_data: &[(GpuTensor, GpuTensor)],
|
|
epochs: usize,
|
|
batch_size: usize,
|
|
checkpoint_dir: Option<&std::path::Path>,
|
|
) -> Result<Vec<crate::mamba::TrainingEpoch>, MLError> {
|
|
info!(
|
|
"Async data loading enabled (prefetch={}, batch_size={})",
|
|
self.prefetch_count, batch_size
|
|
);
|
|
|
|
let stream = self.device.cuda_stream()
|
|
.map_err(|e| MLError::DeviceError(format!("CUDA stream: {e}")))?;
|
|
|
|
// Convert GpuTensor pairs to StreamTensor pairs for ml-supervised API
|
|
let train_stream: Vec<(StreamTensor, StreamTensor)> = train_data.iter()
|
|
.filter_map(|(x, y)| {
|
|
let sx = gpu_to_stream(x, stream).ok()?;
|
|
let sy = gpu_to_stream(y, stream).ok()?;
|
|
Some((sx, sy))
|
|
})
|
|
.collect();
|
|
let val_stream: Vec<(StreamTensor, StreamTensor)> = val_data.iter()
|
|
.filter_map(|(x, y)| {
|
|
let sx = gpu_to_stream(x, stream).ok()?;
|
|
let sy = gpu_to_stream(y, stream).ok()?;
|
|
Some((sx, sy))
|
|
})
|
|
.collect();
|
|
|
|
model
|
|
.train_async(
|
|
&train_stream,
|
|
&val_stream,
|
|
epochs,
|
|
batch_size,
|
|
self.prefetch_count,
|
|
checkpoint_dir,
|
|
)
|
|
.await
|
|
}
|
|
}
|
|
|
|
/// Convert a `GpuTensor` to a `StreamTensor` by cloning the data on the same device.
|
|
fn gpu_to_stream(
|
|
t: &GpuTensor,
|
|
stream: &Arc<cudarc::driver::CudaStream>,
|
|
) -> Result<StreamTensor, MLError> {
|
|
let host = t.to_host(stream)?;
|
|
let data = crate::cuda_pipeline::clone_htod_f32_to_bf16(stream, &host)?;
|
|
Ok(StreamTensor {
|
|
data,
|
|
shape: t.shape().to_vec(), // cpu-side shape clone
|
|
stream: stream.clone(),
|
|
})
|
|
}
|
|
|
|
/// Write a log entry to the training log file
|
|
fn write_training_log_mamba2(
|
|
logs_dir: &std::path::Path,
|
|
message: &str,
|
|
) -> Result<(), std::io::Error> {
|
|
let log_file = logs_dir.join("training.log");
|
|
let mut file = OpenOptions::new()
|
|
.create(true)
|
|
.append(true)
|
|
.open(log_file)?;
|
|
|
|
let timestamp = chrono::Utc::now().format("%Y-%m-%d %H:%M:%S");
|
|
writeln!(file, "[{}] {}", timestamp, message)?;
|
|
Ok(())
|
|
}
|
|
|
|
/// Write trial results to JSON file
|
|
fn write_trial_result_mamba2(
|
|
hyperopt_dir: &std::path::Path,
|
|
trial_result: &crate::hyperopt::traits::TrialResult<Mamba2Params>,
|
|
) -> Result<(), std::io::Error> {
|
|
let trials_file = hyperopt_dir.join("trials.json");
|
|
|
|
// Read existing trials (if any)
|
|
let mut all_trials = if trials_file.exists() {
|
|
let content = std::fs::read_to_string(&trials_file)?;
|
|
serde_json::from_str::<Vec<serde_json::Value>>(&content).unwrap_or_default()
|
|
} else {
|
|
Vec::new()
|
|
};
|
|
|
|
// Append new trial
|
|
let trial_json = serde_json::to_value(trial_result)
|
|
.map_err(|e| std::io::Error::new(std::io::ErrorKind::Other, e))?;
|
|
all_trials.push(trial_json);
|
|
|
|
// Write back to file (pretty printed)
|
|
let content = serde_json::to_string_pretty(&all_trials)
|
|
.map_err(|e| std::io::Error::new(std::io::ErrorKind::Other, e))?;
|
|
std::fs::write(&trials_file, content)?;
|
|
|
|
Ok(())
|
|
}
|
|
|
|
impl HyperparameterOptimizable for Mamba2Trainer {
|
|
type Params = Mamba2Params;
|
|
type Metrics = Mamba2Metrics;
|
|
|
|
#[allow(clippy::unwrap_in_result)]
|
|
fn train_with_params(&mut self, mut params: Self::Params) -> Result<Self::Metrics, MLError> {
|
|
// START: Add trial timing
|
|
let trial_start = std::time::Instant::now();
|
|
|
|
// Get current trial number and increment for next trial
|
|
let current_trial = self.trial_counter;
|
|
self.trial_counter += 1;
|
|
|
|
// Clamp batch_size to configured bounds (for GPU memory constraints)
|
|
let original_batch_size = params.batch_size;
|
|
let clamped_batch_size = (params.batch_size as f64)
|
|
.clamp(self.batch_size_min, self.batch_size_max)
|
|
.round() as usize;
|
|
|
|
if clamped_batch_size != original_batch_size {
|
|
warn!(
|
|
"Batch size clamped: {} → {} (bounds: [{}, {}])",
|
|
original_batch_size, clamped_batch_size, self.batch_size_min, self.batch_size_max
|
|
);
|
|
params.batch_size = clamped_batch_size;
|
|
}
|
|
|
|
info!("Training MAMBA-2 with 12 hyperparameters:");
|
|
info!(" Learning rate: {:.6}", params.learning_rate);
|
|
info!(
|
|
" Batch size: {} (bounds: [{}, {}])",
|
|
params.batch_size, self.batch_size_min, self.batch_size_max
|
|
);
|
|
info!(" Dropout: {:.3}", params.dropout);
|
|
info!(" Weight decay: {:.6}", params.weight_decay);
|
|
info!(" P0 - Grad clip: {:.3}", params.grad_clip);
|
|
info!(" P0 - Warmup steps: {}", params.warmup_steps);
|
|
info!(" P0 - Adam beta1: {:.4}", params.adam_beta1);
|
|
info!(" P1 - Adam beta2: {:.4}", params.adam_beta2);
|
|
info!(" P1 - Adam epsilon: {:.2e}", params.adam_epsilon);
|
|
info!(" P2 - Lookback window: {}", params.lookback_window);
|
|
info!(" P2 - Sequence stride: {}", params.sequence_stride);
|
|
info!(" P2 - Norm epsilon: {:.2e}", params.norm_eps);
|
|
|
|
// Log trial start (ensure directory exists first)
|
|
std::fs::create_dir_all(self.training_paths.logs_dir()).ok();
|
|
write_training_log_mamba2(
|
|
&self.training_paths.logs_dir(),
|
|
&format!("=== Starting MAMBA-2 Trial ===\nParams: {:#?}", params),
|
|
)
|
|
.ok();
|
|
|
|
// Load and prepare data FIRST (required to calculate total_decay_steps)
|
|
let (train_data, val_data, target_min, target_max) = self
|
|
.load_and_prepare_data(params.lookback_window, params.sequence_stride)
|
|
.map_err(|e| MLError::ModelError(format!("Data loading failed: {}", e)))?;
|
|
|
|
// Calculate total_decay_steps dynamically based on actual data size
|
|
let steps_per_epoch = (train_data.len() + params.batch_size - 1) / params.batch_size;
|
|
let total_decay_steps = self.epochs * steps_per_epoch;
|
|
|
|
info!(
|
|
"Calculated total_decay_steps: {} (epochs={}, steps_per_epoch={})",
|
|
total_decay_steps, self.epochs, steps_per_epoch
|
|
);
|
|
|
|
// Create MAMBA-2 config with trial hyperparameters (ALL 12 params + calculated total_decay_steps)
|
|
let mamba_config = Mamba2Config {
|
|
d_model: self.d_model,
|
|
d_state: 16,
|
|
d_head: self.d_model / 8,
|
|
num_heads: 8,
|
|
expand: 2,
|
|
num_layers: 6,
|
|
dropout: params.dropout,
|
|
use_ssd: true,
|
|
use_selective_state: true,
|
|
hardware_aware: true,
|
|
target_latency_us: 5,
|
|
max_seq_len: 120,
|
|
learning_rate: params.learning_rate,
|
|
weight_decay: params.weight_decay,
|
|
grad_clip: params.grad_clip, // P0
|
|
warmup_steps: params.warmup_steps, // P0
|
|
adam_beta1: params.adam_beta1, // P0
|
|
adam_beta2: params.adam_beta2, // P1
|
|
adam_epsilon: params.adam_epsilon, // P1
|
|
total_decay_steps, // P1 - CALCULATED DYNAMICALLY
|
|
batch_size: params.batch_size,
|
|
seq_len: params.lookback_window, // P2
|
|
shuffle_batches: false,
|
|
optimizer_type: OptimizerType::Adam,
|
|
sgd_momentum: 0.9,
|
|
sequence_stride: params.sequence_stride, // P2
|
|
norm_eps: params.norm_eps, // P2
|
|
early_stopping_enabled: true, // Enable early stopping for hyperopt
|
|
early_stopping_patience: self.early_stopping_patience,
|
|
early_stopping_min_delta: 1e-4, // Minimum improvement threshold
|
|
early_stopping_min_epochs: self.early_stopping_min_epochs,
|
|
spsa_epsilon: 0.01, // SPSA perturbation magnitude (Spall 1992)
|
|
};
|
|
|
|
// Store normalization params for inference
|
|
self.target_min = Some(target_min);
|
|
self.target_max = Some(target_max);
|
|
|
|
if train_data.is_empty() || val_data.is_empty() {
|
|
debug!("Empty training or validation data, returning penalty metrics");
|
|
return Ok(Mamba2Metrics {
|
|
val_loss: 1000.0, // Penalty
|
|
train_loss: 1000.0,
|
|
val_perplexity: f64::INFINITY,
|
|
directional_accuracy: 0.0,
|
|
mae: 1000.0,
|
|
rmse: 1000.0,
|
|
r_squared: 0.0,
|
|
epochs_completed: 0,
|
|
backtest_sharpe: None,
|
|
backtest_trades: None,
|
|
});
|
|
}
|
|
|
|
// Create all training directories
|
|
self.training_paths.create_all().map_err(|e| {
|
|
MLError::ModelError(format!("Failed to create training directories: {}", e))
|
|
})?;
|
|
|
|
info!("Training directories created:");
|
|
info!(" Checkpoints: {:?}", self.training_paths.checkpoints_dir());
|
|
|
|
// Create and train model
|
|
let stream = self.device.cuda_stream()
|
|
.map_err(|e| MLError::DeviceError(format!("CUDA stream: {e}")))?;
|
|
let mut model = Mamba2SSM::new(mamba_config.clone(), stream)
|
|
.map_err(|e| MLError::ModelError(format!("Failed to create model: {}", e)))?;
|
|
|
|
// P1 FIX: Wrap training in catch_unwind to handle CUDA OOM panics gracefully
|
|
// Run training (async or sync based on configuration)
|
|
let training_result = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
|
|
if self.async_loading {
|
|
info!(
|
|
"Using async data loading (prefetch={})",
|
|
self.prefetch_count
|
|
);
|
|
tokio::runtime::Runtime::new()
|
|
.map_err(|e| MLError::ModelError(format!("Failed to create async runtime: {}", e)))?
|
|
.block_on(self.train_with_async_loading(
|
|
&mut model,
|
|
&train_data,
|
|
&val_data,
|
|
self.epochs,
|
|
params.batch_size,
|
|
Some(&self.training_paths.checkpoints_dir()),
|
|
))
|
|
} else {
|
|
info!("Using synchronous data loading");
|
|
let stream = self.device.cuda_stream()
|
|
.map_err(|e| MLError::DeviceError(format!("CUDA stream: {e}")))?;
|
|
let train_s: Vec<(StreamTensor, StreamTensor)> = train_data.iter()
|
|
.filter_map(|(x, y)| {
|
|
let sx = gpu_to_stream(x, stream).ok()?;
|
|
let sy = gpu_to_stream(y, stream).ok()?;
|
|
Some((sx, sy))
|
|
})
|
|
.collect();
|
|
let val_s: Vec<(StreamTensor, StreamTensor)> = val_data.iter()
|
|
.filter_map(|(x, y)| {
|
|
let sx = gpu_to_stream(x, stream).ok()?;
|
|
let sy = gpu_to_stream(y, stream).ok()?;
|
|
Some((sx, sy))
|
|
})
|
|
.collect();
|
|
tokio::runtime::Runtime::new()
|
|
.map_err(|e| MLError::ModelError(format!("Failed to create async runtime: {}", e)))?
|
|
.block_on(model.train(
|
|
&train_s,
|
|
&val_s,
|
|
self.epochs,
|
|
Some(&self.training_paths.checkpoints_dir()),
|
|
))
|
|
}
|
|
}));
|
|
|
|
let training_history = match training_result {
|
|
Ok(Ok(history)) => history,
|
|
Ok(Err(e)) => {
|
|
warn!("Training failed (possibly OOM): {}", e);
|
|
// Return penalty metrics instead of propagating error
|
|
return Ok(Mamba2Metrics {
|
|
val_loss: 1000.0,
|
|
train_loss: 1000.0,
|
|
val_perplexity: f64::INFINITY,
|
|
directional_accuracy: 0.0,
|
|
mae: 1000.0,
|
|
rmse: 1000.0,
|
|
r_squared: 0.0,
|
|
epochs_completed: 0,
|
|
backtest_sharpe: None,
|
|
backtest_trades: None,
|
|
});
|
|
},
|
|
Err(_) => {
|
|
warn!("Training panicked (likely CUDA OOM or GPU error)");
|
|
// Return penalty metrics for optimizer
|
|
return Ok(Mamba2Metrics {
|
|
val_loss: 1000.0,
|
|
train_loss: 1000.0,
|
|
val_perplexity: f64::INFINITY,
|
|
directional_accuracy: 0.0,
|
|
mae: 1000.0,
|
|
rmse: 1000.0,
|
|
r_squared: 0.0,
|
|
epochs_completed: 0,
|
|
backtest_sharpe: None,
|
|
backtest_trades: None,
|
|
});
|
|
},
|
|
};
|
|
|
|
// Extract final metrics
|
|
let final_epoch = training_history
|
|
.last()
|
|
.ok_or_else(|| MLError::TrainingError("No training history".to_owned()))?;
|
|
|
|
// Map fields from TrainingEpoch (which only has loss/accuracy)
|
|
// to Mamba2Metrics (which expects detailed metrics)
|
|
let metrics = Mamba2Metrics {
|
|
val_loss: final_epoch.loss, // Use loss as val_loss
|
|
train_loss: final_epoch.loss, // Same value for train_loss (best available)
|
|
val_perplexity: final_epoch.loss.exp(),
|
|
directional_accuracy: final_epoch.accuracy, // Use accuracy as directional_accuracy
|
|
mae: final_epoch.loss, // Approximation
|
|
rmse: final_epoch.loss.sqrt(), // Approximation
|
|
r_squared: 1.0 - final_epoch.loss.min(1.0), // Approximation
|
|
epochs_completed: training_history.len(),
|
|
backtest_sharpe: None,
|
|
backtest_trades: None,
|
|
};
|
|
|
|
info!("Training completed:");
|
|
info!(" Training loss: {:.6}", metrics.train_loss);
|
|
info!(" Validation loss: {:.6}", metrics.val_loss);
|
|
info!(" Perplexity: {:.4}", metrics.val_perplexity);
|
|
info!(
|
|
" Directional accuracy: {:.2}%",
|
|
metrics.directional_accuracy * 100.0
|
|
);
|
|
info!(" MAE: {:.4}", metrics.mae);
|
|
info!(" RMSE: {:.4}", metrics.rmse);
|
|
info!(" R²: {:.4}", metrics.r_squared);
|
|
|
|
// END: Add trial completion logging
|
|
let duration_secs = trial_start.elapsed().as_secs_f64();
|
|
write_training_log_mamba2(
|
|
&self.training_paths.logs_dir(),
|
|
&format!(
|
|
"Training completed in {:.2}s: val_loss={:.6}, train_loss={:.6}, accuracy={:.2}%",
|
|
duration_secs,
|
|
metrics.val_loss,
|
|
metrics.train_loss,
|
|
metrics.directional_accuracy * 100.0
|
|
),
|
|
)
|
|
.ok();
|
|
|
|
// CRITICAL FIX: Explicit memory cleanup to prevent OOM between trials
|
|
// Drop model, data loaders, and sync CUDA to free GPU/RAM
|
|
info!("Cleaning up resources...");
|
|
drop(model);
|
|
drop(train_data);
|
|
drop(val_data);
|
|
|
|
// Sync CUDA to ensure GPU memory is freed
|
|
if self.device.is_cuda() {
|
|
use ml_core::device::MlDevice;
|
|
if let MlDevice::Cuda { .. } = &self.device {
|
|
// Force CUDA synchronization to release GPU memory
|
|
std::thread::sleep(std::time::Duration::from_millis(100));
|
|
}
|
|
}
|
|
info!("Resource cleanup complete");
|
|
|
|
// Write trial result to JSON (ensure directory exists first)
|
|
let trial_result = crate::hyperopt::traits::TrialResult {
|
|
trial_num: current_trial,
|
|
params,
|
|
objective: Self::extract_objective(&metrics),
|
|
duration_secs,
|
|
metrics: None,
|
|
};
|
|
|
|
std::fs::create_dir_all(self.training_paths.hyperopt_dir()).ok();
|
|
write_trial_result_mamba2(&self.training_paths.hyperopt_dir(), &trial_result).ok();
|
|
|
|
Ok(metrics)
|
|
}
|
|
|
|
fn extract_objective(metrics: &Self::Metrics) -> f64 {
|
|
if let (Some(sharpe), Some(trades)) = (metrics.backtest_sharpe, metrics.backtest_trades) {
|
|
return crate::cuda_pipeline::signal_adapter::backtest_fitness(sharpe, trades, 30, 0.0);
|
|
}
|
|
metrics.val_loss
|
|
}
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
|
|
#[test]
|
|
fn test_mamba2_params_roundtrip() {
|
|
let params = Mamba2Params {
|
|
learning_rate: 0.001,
|
|
batch_size: 64,
|
|
dropout: 0.2,
|
|
weight_decay: 0.0001,
|
|
grad_clip: 2.5,
|
|
warmup_steps: 500,
|
|
adam_beta1: 0.9,
|
|
adam_beta2: 0.999,
|
|
adam_epsilon: 1e-8,
|
|
lookback_window: 90,
|
|
sequence_stride: 3,
|
|
norm_eps: 5e-5,
|
|
signal_high_bps: 10.0,
|
|
signal_low_bps: 5.0,
|
|
};
|
|
|
|
let continuous = params.to_continuous();
|
|
assert_eq!(continuous.len(), 14, "Expected 14 continuous parameters");
|
|
|
|
let recovered = Mamba2Params::from_continuous(&continuous).unwrap();
|
|
|
|
assert!((recovered.learning_rate - params.learning_rate).abs() < 1e-10);
|
|
assert_eq!(recovered.batch_size, params.batch_size);
|
|
assert!((recovered.dropout - params.dropout).abs() < 1e-10);
|
|
assert!((recovered.weight_decay - params.weight_decay).abs() < 1e-10);
|
|
assert!((recovered.grad_clip - params.grad_clip).abs() < 1e-6);
|
|
assert_eq!(recovered.warmup_steps, params.warmup_steps);
|
|
assert!((recovered.adam_beta1 - params.adam_beta1).abs() < 1e-10);
|
|
assert!((recovered.adam_beta2 - params.adam_beta2).abs() < 1e-10);
|
|
assert!((recovered.adam_epsilon - params.adam_epsilon).abs() < 1e-12);
|
|
assert_eq!(recovered.lookback_window, params.lookback_window);
|
|
assert_eq!(recovered.sequence_stride, params.sequence_stride);
|
|
assert!((recovered.norm_eps - params.norm_eps).abs() < 1e-12);
|
|
assert!((recovered.signal_high_bps - params.signal_high_bps).abs() < 1e-10);
|
|
assert!((recovered.signal_low_bps - params.signal_low_bps).abs() < 1e-10);
|
|
}
|
|
|
|
#[test]
|
|
fn test_mamba2_params_bounds() {
|
|
let bounds = Mamba2Params::continuous_bounds();
|
|
assert_eq!(bounds.len(), 14, "Expected 14 parameter bounds");
|
|
|
|
// Check log-scale bounds are reasonable
|
|
assert!(bounds[0].0 < bounds[0].1); // learning_rate
|
|
assert!(bounds[3].0 < bounds[3].1); // weight_decay
|
|
assert!(bounds[4].0 < bounds[4].1); // grad_clip
|
|
assert!(bounds[8].0 < bounds[8].1); // adam_epsilon
|
|
assert!(bounds[11].0 < bounds[11].1); // norm_eps
|
|
|
|
// Check linear bounds
|
|
assert_eq!(bounds[1], (4.0, 1024.0)); // batch_size (wide bounds, capped by VRAM in continuous_bounds_for)
|
|
assert_eq!(bounds[2], (0.0, 0.5)); // dropout
|
|
assert_eq!(bounds[5], (100.0, 2000.0)); // warmup_steps
|
|
assert_eq!(bounds[6], (0.85, 0.95)); // adam_beta1
|
|
assert_eq!(bounds[7], (0.98, 0.999)); // adam_beta2
|
|
assert_eq!(bounds[9], (30.0, 120.0)); // lookback_window
|
|
assert_eq!(bounds[10], (1.0, 5.0)); // sequence_stride
|
|
}
|
|
|
|
#[test]
|
|
fn test_param_names() {
|
|
let names = Mamba2Params::param_names();
|
|
assert_eq!(names.len(), 14, "Expected 14 parameter names");
|
|
assert_eq!(names[0], "learning_rate");
|
|
assert_eq!(names[1], "batch_size");
|
|
assert_eq!(names[2], "dropout");
|
|
assert_eq!(names[3], "weight_decay");
|
|
assert_eq!(names[4], "grad_clip");
|
|
assert_eq!(names[5], "warmup_steps");
|
|
assert_eq!(names[6], "adam_beta1");
|
|
assert_eq!(names[7], "adam_beta2");
|
|
assert_eq!(names[8], "adam_epsilon");
|
|
assert_eq!(names[9], "lookback_window");
|
|
assert_eq!(names[10], "sequence_stride");
|
|
assert_eq!(names[11], "norm_eps");
|
|
assert_eq!(names[12], "signal_high_bps");
|
|
assert_eq!(names[13], "signal_low_bps");
|
|
}
|
|
|
|
#[test]
|
|
fn test_target_normalization() {
|
|
// Test normalization/denormalization with realistic ES price ranges
|
|
let target_min = 5000.0;
|
|
let target_max = 6000.0;
|
|
|
|
// Test min value
|
|
let normalized_min: f64 = (target_min - target_min) / (target_max - target_min);
|
|
assert!(
|
|
(normalized_min - 0.0).abs() < 1e-10,
|
|
"Min should normalize to 0"
|
|
);
|
|
|
|
// Test max value
|
|
let normalized_max: f64 = (target_max - target_min) / (target_max - target_min);
|
|
assert!(
|
|
(normalized_max - 1.0).abs() < 1e-10,
|
|
"Max should normalize to 1"
|
|
);
|
|
|
|
// Test mid value
|
|
let mid_price = 5500.0;
|
|
let normalized_mid: f64 = (mid_price - target_min) / (target_max - target_min);
|
|
assert!(
|
|
(normalized_mid - 0.5).abs() < 1e-10,
|
|
"Midpoint should normalize to 0.5"
|
|
);
|
|
|
|
// Test denormalization recovers original
|
|
let denormalized: f64 = normalized_mid * (target_max - target_min) + target_min;
|
|
assert!(
|
|
(denormalized - mid_price).abs() < 1e-6,
|
|
"Denormalization should recover original price"
|
|
);
|
|
}
|
|
|
|
#[test]
|
|
fn test_denormalize_prediction() {
|
|
// Create a trainer with normalization params set
|
|
let trainer = Mamba2Trainer {
|
|
data_dir: PathBuf::from("dummy_data"),
|
|
epochs: 1,
|
|
device: MlDevice::cuda(0).unwrap_or(MlDevice::Cpu),
|
|
feature_config: FeatureConfig::wave_d(),
|
|
d_model: 225,
|
|
train_split: 0.8,
|
|
target_min: Some(5000.0),
|
|
target_max: Some(6000.0),
|
|
batch_size_min: 4.0,
|
|
batch_size_max: 96.0,
|
|
async_loading: false,
|
|
prefetch_count: 0,
|
|
training_paths: TrainingPaths::new("/tmp/test", "mamba2", "test"),
|
|
early_stopping_patience: 10,
|
|
early_stopping_min_epochs: 5,
|
|
trial_counter: 0,
|
|
preloaded_bars: None,
|
|
};
|
|
|
|
// Test denormalization
|
|
assert!((trainer.denormalize_prediction(0.0).unwrap_or(0.0) - 5000.0).abs() < 1e-6);
|
|
assert!((trainer.denormalize_prediction(1.0).unwrap_or(0.0) - 6000.0).abs() < 1e-6);
|
|
assert!((trainer.denormalize_prediction(0.5).unwrap_or(0.0) - 5500.0).abs() < 1e-6);
|
|
assert!((trainer.denormalize_prediction(0.25).unwrap_or(0.0) - 5250.0).abs() < 1e-6);
|
|
}
|
|
|
|
#[test]
|
|
fn test_denormalize_before_training_returns_none() {
|
|
// Create trainer without normalization params
|
|
let trainer = Mamba2Trainer {
|
|
data_dir: PathBuf::from("dummy_data"),
|
|
epochs: 1,
|
|
device: MlDevice::cuda(0).unwrap_or(MlDevice::Cpu),
|
|
feature_config: FeatureConfig::wave_d(),
|
|
batch_size_min: 4.0,
|
|
batch_size_max: 96.0,
|
|
async_loading: false,
|
|
prefetch_count: 0,
|
|
d_model: 225,
|
|
train_split: 0.8,
|
|
target_min: None,
|
|
target_max: None,
|
|
training_paths: TrainingPaths::new("/tmp/test", "mamba2", "test"),
|
|
early_stopping_patience: 10,
|
|
early_stopping_min_epochs: 5,
|
|
trial_counter: 0,
|
|
preloaded_bars: None,
|
|
};
|
|
|
|
// Should return None when normalization params not set
|
|
assert!(trainer.denormalize_prediction(0.5).is_none());
|
|
}
|
|
|
|
#[test]
|
|
fn test_normalized_targets_in_range() {
|
|
// Verify that normalized targets are always in [0,1]
|
|
let target_min = 5000.0;
|
|
let target_max = 6000.0;
|
|
|
|
let test_prices = vec![5000.0, 5100.0, 5500.0, 5900.0, 6000.0];
|
|
|
|
for price in test_prices {
|
|
let normalized: f64 = (price - target_min) / (target_max - target_min);
|
|
assert!(
|
|
normalized >= 0.0,
|
|
"Normalized target {} should be >= 0",
|
|
normalized
|
|
);
|
|
assert!(
|
|
normalized <= 1.0,
|
|
"Normalized target {} should be <= 1",
|
|
normalized
|
|
);
|
|
|
|
// Verify round-trip
|
|
let denormalized: f64 = normalized * (target_max - target_min) + target_min;
|
|
assert!(
|
|
(denormalized - price).abs() < 1e-6,
|
|
"Round-trip failed: {} -> {} -> {}",
|
|
price,
|
|
normalized,
|
|
denormalized
|
|
);
|
|
}
|
|
}
|
|
}
|