09c515e3e9
Final 31 ml crate fixes: unsafe_code allows, unused vars prefixed, boolean simplification, dead code removal, integer suffix, drop cleanup. cargo fix auto-removed ~30 unused imports from ml crate. Total clippy cleanup: 278 errors → 0 across all ML crates. Full workspace: `cargo clippy --workspace --lib -- -D warnings` = 0 errors. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
525 lines
19 KiB
Rust
525 lines
19 KiB
Rust
//! Hyperopt adapter for the Diffusion model.
|
|
//!
|
|
//! Defines `DiffusionParams` (ParameterSpace), `DiffusionMetrics`,
|
|
//! and `DiffusionTrainer` (HyperparameterOptimizable) for hyperparameter
|
|
//! optimization of the Diffusion model via the unified framework.
|
|
|
|
use std::sync::Arc;
|
|
use ml_core::device::MlDevice;
|
|
use std::path::PathBuf;
|
|
use tracing::{info, warn};
|
|
|
|
use crate::diffusion::config::{DiffusionConfig, NoiseSchedule};
|
|
use crate::diffusion::trainable::DiffusionTrainableAdapter;
|
|
use crate::features::extract_ml_features;
|
|
use crate::features::extraction::OHLCVBar;
|
|
use crate::hyperopt::paths::TrainingPaths;
|
|
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 = 100.0;
|
|
/// Diffusion per-sample memory in MB (noise + denoise activations)
|
|
const MB_PER_SAMPLE: f64 = 0.03;
|
|
|
|
/// Hyperparameters for Diffusion model hyperopt tuning.
|
|
#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
|
|
pub struct DiffusionParams {
|
|
/// Learning rate (log scale).
|
|
pub learning_rate: f64,
|
|
/// Number of diffusion timesteps.
|
|
pub num_timesteps: usize,
|
|
/// Number of DDIM sampling steps.
|
|
pub sampling_steps: usize,
|
|
/// Hidden dimension of the denoiser.
|
|
pub hidden_dim: usize,
|
|
/// Number of denoiser layers.
|
|
pub num_layers: usize,
|
|
/// Time embedding dimension.
|
|
pub time_embed_dim: usize,
|
|
/// Batch size for training.
|
|
pub batch_size: usize,
|
|
/// Weight decay (log scale).
|
|
pub weight_decay: f64,
|
|
/// Gradient clipping max norm (log scale).
|
|
pub grad_clip: 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 DiffusionParams {
|
|
fn default() -> Self {
|
|
Self {
|
|
learning_rate: 1e-4,
|
|
num_timesteps: 1000,
|
|
sampling_steps: 10,
|
|
hidden_dim: 128,
|
|
num_layers: 3,
|
|
time_embed_dim: 32,
|
|
batch_size: 32,
|
|
weight_decay: 1e-4,
|
|
grad_clip: 1.0,
|
|
signal_high_bps: 10.0,
|
|
signal_low_bps: 5.0,
|
|
}
|
|
}
|
|
}
|
|
|
|
impl ParameterSpace for DiffusionParams {
|
|
fn continuous_bounds() -> Vec<(f64, f64)> {
|
|
vec![
|
|
(1e-5_f64.ln(), 1e-3_f64.ln()), // learning_rate (log)
|
|
(100.0, 2000.0), // num_timesteps
|
|
(5.0, 50.0), // sampling_steps
|
|
(32.0, 2048.0), // hidden_dim
|
|
(1.0, 6.0), // num_layers
|
|
(8.0, 64.0), // time_embed_dim
|
|
(4.0, 256.0), // batch_size
|
|
(1e-6_f64.ln(), 1e-2_f64.ln()), // weight_decay (log)
|
|
(0.5_f64.ln(), 5.0_f64.ln()), // grad_clip (log)
|
|
(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() != 11 {
|
|
return Err(MLError::ConfigError(format!("Expected 11 params, got {}", x.len())));
|
|
}
|
|
Ok(Self {
|
|
learning_rate: x.first().copied().unwrap_or(-9.21).exp(),
|
|
num_timesteps: x.get(1).copied().unwrap_or(1000.0).round().max(100.0) as usize,
|
|
sampling_steps: x.get(2).copied().unwrap_or(10.0).round().max(5.0) as usize,
|
|
hidden_dim: x.get(3).copied().unwrap_or(128.0).round().max(32.0) as usize,
|
|
num_layers: x.get(4).copied().unwrap_or(3.0).round().max(1.0) as usize,
|
|
time_embed_dim: x.get(5).copied().unwrap_or(32.0).round().max(8.0) as usize,
|
|
batch_size: x.get(6).copied().unwrap_or(32.0).round().max(4.0) as usize,
|
|
weight_decay: x.get(7).copied().unwrap_or(-9.21).exp(),
|
|
grad_clip: x.get(8).copied().unwrap_or(0.0).exp(),
|
|
signal_high_bps: x.get(9).copied().unwrap_or(10.0).clamp(1.0, 30.0),
|
|
signal_low_bps: x.get(10).copied().unwrap_or(5.0).clamp(0.5, 15.0),
|
|
})
|
|
}
|
|
|
|
fn to_continuous(&self) -> Vec<f64> {
|
|
vec![
|
|
self.learning_rate.ln(),
|
|
self.num_timesteps as f64,
|
|
self.sampling_steps as f64,
|
|
self.hidden_dim as f64,
|
|
self.num_layers as f64,
|
|
self.time_embed_dim as f64,
|
|
self.batch_size as f64,
|
|
self.weight_decay.ln(),
|
|
self.grad_clip.ln(),
|
|
self.signal_high_bps,
|
|
self.signal_low_bps,
|
|
]
|
|
}
|
|
|
|
fn param_names() -> Vec<&'static str> {
|
|
vec![
|
|
"learning_rate", "num_timesteps", "sampling_steps",
|
|
"hidden_dim", "num_layers", "time_embed_dim",
|
|
"batch_size", "weight_decay", "grad_clip",
|
|
"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, 2048.0) {
|
|
if let Some(batch_bound) = bounds.get_mut(6) {
|
|
batch_bound.1 = max_batch;
|
|
}
|
|
}
|
|
// Cap hidden_dim by VRAM (index 3)
|
|
if budget.gpu_memory_mb < 8000 {
|
|
bounds[3] = (32.0, 256.0);
|
|
} else if budget.gpu_memory_mb < 16000 {
|
|
bounds[3] = (32.0, 512.0);
|
|
} else if budget.gpu_memory_mb < 40000 {
|
|
bounds[3] = (32.0, 1024.0);
|
|
} else {
|
|
// ≥40GB (L40S 46GB, H100 80GB): allow full 2048
|
|
}
|
|
bounds
|
|
}
|
|
}
|
|
|
|
/// Metrics returned from Diffusion hyperopt training.
|
|
#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
|
|
pub struct DiffusionMetrics {
|
|
/// Validation noise prediction loss.
|
|
pub val_loss: f64,
|
|
/// Training noise prediction loss.
|
|
pub train_loss: 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>,
|
|
}
|
|
|
|
impl Default for DiffusionMetrics {
|
|
fn default() -> Self {
|
|
Self {
|
|
val_loss: 0.0,
|
|
train_loss: 0.0,
|
|
epochs_completed: 0,
|
|
backtest_sharpe: None,
|
|
backtest_trades: None,
|
|
}
|
|
}
|
|
}
|
|
|
|
/// Diffusion model trainer for hyperparameter optimization.
|
|
///
|
|
/// Wraps the `DiffusionTrainableAdapter` (which implements `UnifiedTrainable`)
|
|
/// and drives it through the `HyperparameterOptimizable` interface.
|
|
///
|
|
/// The Diffusion model generates price paths via denoising diffusion.
|
|
/// Training uses sequence input: `(1, seq_len, feature_dim)`.
|
|
#[derive(Debug)]
|
|
pub struct DiffusionTrainer {
|
|
data_dir: PathBuf,
|
|
epochs: usize,
|
|
device: MlDevice,
|
|
training_paths: TrainingPaths,
|
|
early_stopping_patience: usize,
|
|
trial_counter: usize,
|
|
preloaded_bars: Option<Arc<[OHLCVBar]>>,
|
|
}
|
|
|
|
impl DiffusionTrainer {
|
|
/// Create a new Diffusion trainer.
|
|
///
|
|
/// # Arguments
|
|
/// * `data_dir` - Directory containing .dbn/.dbn.zst files
|
|
/// * `epochs` - Number of training epochs per trial
|
|
pub fn new<P: Into<PathBuf>>(data_dir: P, epochs: usize) -> Result<Self, MLError> {
|
|
let data_dir = data_dir.into();
|
|
if !data_dir.exists() {
|
|
return Err(MLError::ConfigError(format!("Data directory not found: {}", data_dir.display())));
|
|
}
|
|
|
|
let device = MlDevice::cuda(0)
|
|
.map_err(|e| MLError::ConfigError(format!("CUDA GPU required for Diffusion hyperopt: {}", e)))?;
|
|
|
|
info!(
|
|
"Diffusion Trainer initialized: Device={:?}, Data={}, Epochs={}",
|
|
device, data_dir.display(), epochs
|
|
);
|
|
|
|
Ok(Self {
|
|
data_dir,
|
|
epochs,
|
|
device,
|
|
training_paths: TrainingPaths::new("/tmp/ml_training", "diffusion", "default"),
|
|
early_stopping_patience: 10,
|
|
trial_counter: 0,
|
|
preloaded_bars: None,
|
|
})
|
|
}
|
|
|
|
/// Set training paths configuration.
|
|
pub fn with_training_paths(mut self, paths: TrainingPaths) -> Self {
|
|
self.training_paths = paths;
|
|
self
|
|
}
|
|
|
|
/// Configure early stopping patience.
|
|
pub fn with_early_stopping(mut self, patience: usize) -> Self {
|
|
self.early_stopping_patience = patience;
|
|
self
|
|
}
|
|
|
|
/// 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()
|
|
}
|
|
}
|
|
|
|
impl HyperparameterOptimizable for DiffusionTrainer {
|
|
type Params = DiffusionParams;
|
|
type Metrics = DiffusionMetrics;
|
|
|
|
fn train_with_params(&mut self, params: Self::Params) -> Result<Self::Metrics, MLError> {
|
|
let trial_start = std::time::Instant::now();
|
|
let current_trial = self.trial_counter;
|
|
self.trial_counter += 1;
|
|
|
|
info!(
|
|
"Training Diffusion: lr={:.6}, timesteps={}, sampling={}, hidden={}, layers={}, batch={}",
|
|
params.learning_rate, params.num_timesteps, params.sampling_steps,
|
|
params.hidden_dim, params.num_layers, params.batch_size
|
|
);
|
|
|
|
self.training_paths.create_all().map_err(|e| {
|
|
MLError::ModelError(format!("Failed to create training directories: {}", e))
|
|
})?;
|
|
|
|
let 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)
|
|
.map_err(|e| MLError::ModelError(format!("Failed to load DBN data: {}", e)))?
|
|
};
|
|
|
|
let features = extract_ml_features(&bars)
|
|
.map_err(|e| MLError::ModelError(format!("Failed to extract features: {}", e)))?;
|
|
|
|
if features.len() < 2 {
|
|
return Err(MLError::ModelError("Insufficient features extracted".to_owned()));
|
|
}
|
|
|
|
let feature_dim = 51;
|
|
|
|
// Diffusion uses flat input: data_dim = seq_len * feature_dim
|
|
// For hyperopt, we use feature_dim=1 and seq_len=feature_dim (51)
|
|
// so data_dim = 51, matching our feature vectors.
|
|
let pairs = crate::hyperopt::shared_data::build_flat_pairs(
|
|
&features, &bars, feature_dim, &self.device,
|
|
)?;
|
|
|
|
let split = (pairs.len() as f64 * 0.8) as usize;
|
|
let train_data = &pairs[..split];
|
|
let val_data = &pairs[split..];
|
|
|
|
if val_data.len() < 2 {
|
|
return Err(MLError::ModelError("Validation set too small".to_owned()));
|
|
}
|
|
|
|
let config = DiffusionConfig {
|
|
num_timesteps: params.num_timesteps,
|
|
sampling_steps: params.sampling_steps,
|
|
seq_len: feature_dim, // Treat feature vector as a sequence
|
|
feature_dim: 1, // Univariate per position
|
|
hidden_dim: params.hidden_dim,
|
|
num_layers: params.num_layers,
|
|
time_embed_dim: params.time_embed_dim,
|
|
schedule: NoiseSchedule::Cosine,
|
|
learning_rate: params.learning_rate,
|
|
weight_decay: params.weight_decay,
|
|
grad_clip: params.grad_clip,
|
|
};
|
|
|
|
let training_result = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| -> Result<DiffusionMetrics, MLError> {
|
|
let stream = self.device.cuda_stream()
|
|
.map_err(|e| MLError::DeviceError(format!("CUDA stream: {e}")))?;
|
|
let mut model = DiffusionTrainableAdapter::new(config, stream)?;
|
|
|
|
let mut best_val_loss = f64::MAX;
|
|
let mut patience_counter = 0_usize;
|
|
let mut last_train_loss = 0.0_f64;
|
|
|
|
for epoch in 0..self.epochs {
|
|
let mut loss_accum = 0.0_f64;
|
|
let mut batch_count = 0_usize;
|
|
|
|
let stream_ref = self.device.cuda_stream()
|
|
.map_err(|e| MLError::DeviceError(format!("CUDA stream: {e}")))?;
|
|
for (input, target) in train_data {
|
|
let input_host = input.to_host(stream_ref)?;
|
|
let target_host = target.to_host(stream_ref)?;
|
|
let loss_val = model.forward_loss(&input_host, &target_host)?;
|
|
|
|
if batch_count == 0 && (loss_val.is_nan() || loss_val.is_infinite()) {
|
|
return Ok(DiffusionMetrics {
|
|
val_loss: 1000.0,
|
|
train_loss: 1000.0,
|
|
epochs_completed: epoch,
|
|
backtest_sharpe: None,
|
|
backtest_trades: None,
|
|
});
|
|
}
|
|
|
|
model.backward(loss_val)?;
|
|
model.optimizer_step()?;
|
|
loss_accum += loss_val;
|
|
batch_count += 1;
|
|
}
|
|
|
|
last_train_loss = if batch_count > 0 {
|
|
loss_accum / batch_count as f64
|
|
} else {
|
|
0.0
|
|
};
|
|
|
|
let mut vl_accum = 0.0_f64;
|
|
let mut vl_count = 0_usize;
|
|
for (vi, vt) in val_data {
|
|
let vi_h = vi.to_host(stream_ref)?;
|
|
let vt_h = vt.to_host(stream_ref)?;
|
|
if let Ok(vl) = model.forward_loss(&vi_h, &vt_h) {
|
|
vl_accum += vl;
|
|
vl_count += 1;
|
|
}
|
|
}
|
|
let val_loss = if vl_count > 0 { vl_accum / vl_count as f64 } else { 1000.0 };
|
|
|
|
if val_loss < best_val_loss {
|
|
best_val_loss = val_loss;
|
|
patience_counter = 0;
|
|
} else {
|
|
patience_counter += 1;
|
|
}
|
|
|
|
if epoch % 10 == 0 {
|
|
info!(
|
|
" Epoch {}/{}: train_loss={:.6}, val_loss={:.6}",
|
|
epoch + 1, self.epochs, last_train_loss, val_loss
|
|
);
|
|
}
|
|
|
|
if patience_counter >= self.early_stopping_patience {
|
|
info!("Early stopping at epoch {}", epoch + 1);
|
|
return Ok(DiffusionMetrics {
|
|
val_loss: best_val_loss,
|
|
train_loss: last_train_loss,
|
|
epochs_completed: epoch + 1,
|
|
backtest_sharpe: None,
|
|
backtest_trades: None,
|
|
});
|
|
}
|
|
}
|
|
|
|
Ok(DiffusionMetrics {
|
|
val_loss: best_val_loss,
|
|
train_loss: last_train_loss,
|
|
epochs_completed: self.epochs,
|
|
backtest_sharpe: None,
|
|
backtest_trades: None,
|
|
})
|
|
}));
|
|
|
|
let metrics = match training_result {
|
|
Ok(Ok(m)) => m,
|
|
Ok(Err(e)) => {
|
|
warn!("Diffusion training failed: {}", e);
|
|
DiffusionMetrics { val_loss: 1000.0, train_loss: 1000.0, epochs_completed: 0, backtest_sharpe: None, backtest_trades: None }
|
|
}
|
|
Err(_) => {
|
|
warn!("Diffusion training panicked (likely OOM)");
|
|
DiffusionMetrics { val_loss: 1000.0, train_loss: 1000.0, epochs_completed: 0, backtest_sharpe: None, backtest_trades: None }
|
|
}
|
|
};
|
|
|
|
info!(
|
|
"Diffusion training completed: val_loss={:.6}, train_loss={:.6}, epochs={}",
|
|
metrics.val_loss, metrics.train_loss, metrics.epochs_completed
|
|
);
|
|
|
|
let duration_secs = trial_start.elapsed().as_secs_f64();
|
|
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();
|
|
crate::hyperopt::shared_data::write_trial_result_json(
|
|
&self.training_paths.hyperopt_dir(),
|
|
&trial_result,
|
|
).ok();
|
|
|
|
if self.device.is_cuda() {
|
|
std::thread::sleep(std::time::Duration::from_millis(100));
|
|
}
|
|
|
|
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_bounds_count_matches_param_names() {
|
|
let bounds = DiffusionParams::continuous_bounds();
|
|
let names = DiffusionParams::param_names();
|
|
assert_eq!(bounds.len(), names.len());
|
|
}
|
|
|
|
#[test]
|
|
fn test_roundtrip_continuous() {
|
|
let params = DiffusionParams::default();
|
|
let continuous = params.to_continuous();
|
|
let restored = DiffusionParams::from_continuous(&continuous).unwrap();
|
|
assert!((params.learning_rate - restored.learning_rate).abs() < 1e-6);
|
|
assert_eq!(params.num_timesteps, restored.num_timesteps);
|
|
assert_eq!(params.hidden_dim, restored.hidden_dim);
|
|
assert_eq!(params.num_layers, restored.num_layers);
|
|
}
|
|
|
|
#[test]
|
|
fn test_from_continuous_wrong_length_errors() {
|
|
let result = DiffusionParams::from_continuous(&[0.1, 0.2]);
|
|
assert!(result.is_err());
|
|
}
|
|
|
|
#[test]
|
|
fn test_bounds_are_valid() {
|
|
for (min, max) in DiffusionParams::continuous_bounds() {
|
|
assert!(min < max, "Invalid bounds: {min} >= {max}");
|
|
}
|
|
}
|
|
|
|
#[test]
|
|
fn test_default_within_bounds() {
|
|
let params = DiffusionParams::default();
|
|
let continuous = params.to_continuous();
|
|
let bounds = DiffusionParams::continuous_bounds();
|
|
for (i, (val, (min, max))) in continuous.iter().zip(bounds.iter()).enumerate() {
|
|
assert!(
|
|
*val >= *min && *val <= *max,
|
|
"Param {} ({}) = {} outside [{}, {}]",
|
|
i,
|
|
DiffusionParams::param_names().get(i).unwrap_or(&"?"),
|
|
val, min, max,
|
|
);
|
|
}
|
|
}
|
|
|
|
#[test]
|
|
fn test_metrics_default() {
|
|
let metrics = DiffusionMetrics::default();
|
|
assert_eq!(metrics.val_loss, 0.0);
|
|
assert_eq!(metrics.train_loss, 0.0);
|
|
assert_eq!(metrics.epochs_completed, 0);
|
|
}
|
|
|
|
#[test]
|
|
fn test_batch_size_static_upper_bound() {
|
|
let bounds = DiffusionParams::continuous_bounds();
|
|
// batch_size is param index 6
|
|
let (_, max_batch) = bounds.get(6).copied().unwrap_or((4.0, 256.0));
|
|
assert!(max_batch <= 256.0, "Max static batch_size should be ≤256");
|
|
}
|
|
}
|