diff --git a/crates/ml/src/trainers/mod.rs b/crates/ml/src/trainers/mod.rs index 23434c44b..543ba7cb2 100644 --- a/crates/ml/src/trainers/mod.rs +++ b/crates/ml/src/trainers/mod.rs @@ -77,7 +77,6 @@ pub mod online_learning; pub mod ppo; pub mod tft; pub mod tft_parquet; // DBN/OHLCV bar training extension for TFT -pub mod tlob; /// Target network update strategy for DQN training #[derive(Debug, Clone, Copy)] @@ -107,4 +106,3 @@ pub use tft::{ TrainingMetrics as TFTTrainingMetrics, TrainingProgress as TFTTrainingProgress, }; pub use liquid::{LiquidHyperparameters, LiquidTrainer, LiquidTrainingMetrics}; -pub use tlob::{TLOBHyperparameters, TLOBTrainer, TLOBTrainingMetrics}; diff --git a/crates/ml/src/trainers/tlob.rs b/crates/ml/src/trainers/tlob.rs deleted file mode 100644 index c8ea04588..000000000 --- a/crates/ml/src/trainers/tlob.rs +++ /dev/null @@ -1,687 +0,0 @@ -//! TLOB Transformer Trainer with gRPC Integration -//! -//! Production-ready TLOB trainer optimized for Level-2 order book data with GPU acceleration. -//! Designed to train transformer models for price movement prediction using MBP-10 (Market By Price) data. -//! -//! ## Features -//! -//! - GPU acceleration (RTX 3050 Ti compatible) -//! - Level-2 order book sequence training (10 price levels) -//! - Checkpoint management with MinIO/S3 integration -//! - Real-time training progress streaming -//! - MSE loss for price movement prediction -//! - 51-feature extraction from order book snapshots -//! -//! ## Architecture -//! -//! - Transformer encoder with multi-head attention -//! - Sequence length: 128 order book snapshots -//! - Input features: 51 microstructure features per snapshot -//! - Output: Next price movement prediction -//! - Training objective: MSE loss on price changes - -use std::path::{Path, PathBuf}; -use std::sync::Arc; -use std::time::Instant; - -use anyhow::{Context, Result}; -use ml_core::device::MlDevice; -use ml_core::cuda_autograd::GpuTensor; -use ml_core::cuda_autograd::GpuVarStore; -use serde::{Deserialize, Serialize}; -use tokio::sync::RwLock; -use tracing::{info, instrument, warn}; - - - - -use crate::tlob::features::TLOB_FEATURE_COUNT; -use crate::tlob::transformer::TLOBTransformer; - -/// TLOB training hyperparameters from gRPC request -#[derive(Debug, Clone, Serialize, Deserialize)] -pub struct TLOBHyperparameters { - /// Learning rate (typically 1e-4 to 1e-5) - pub learning_rate: f64, - - /// Batch size (must be ≤32 for 4GB VRAM) - pub batch_size: usize, - - /// Sequence length (number of order book snapshots) - pub seq_len: usize, - - /// Number of price levels (10 for MBP-10) - pub num_price_levels: usize, - - /// Transformer hidden dimension - pub d_model: usize, - - /// Number of attention heads - pub num_heads: usize, - - /// Number of transformer layers - pub num_layers: usize, - - /// Dropout rate - pub dropout: f64, - - /// Number of training epochs - pub epochs: usize, - - /// Checkpoint save frequency (epochs) - pub checkpoint_frequency: usize, - - /// Gradient clipping threshold - pub grad_clip: f64, - - /// Weight decay for regularization - pub weight_decay: f64, -} - -impl Default for TLOBHyperparameters { - fn default() -> Self { - Self { - learning_rate: 0.0001, - batch_size: 16, // Conservative for 4GB VRAM - seq_len: 128, // Order book snapshot sequence - num_price_levels: 10, // MBP-10 - d_model: 256, // Transformer hidden size - num_heads: 8, // Multi-head attention - num_layers: 4, // Transformer blocks - dropout: 0.1, - epochs: 500, // TLOB needs more epochs - checkpoint_frequency: 10, - grad_clip: 1.0, - weight_decay: 0.0001, - } - } -} - -/// Training progress metrics -#[derive(Debug, Clone, Serialize, Deserialize)] -pub struct TLOBTrainingMetrics { - pub epoch: usize, - pub train_loss: f64, - pub val_loss: f64, - pub avg_mae: f64, // Mean Absolute Error - pub avg_prediction_error: f64, - pub gradient_norm: f64, - pub learning_rate: f64, - pub elapsed_seconds: f64, -} - -/// TLOB Trainer with gRPC integration -pub struct TLOBTrainer { - /// Model configuration - hyperparams: TLOBHyperparameters, - - /// TLOB transformer model (used for ONNX-based inference / fallback) - model: Arc>, - - /// GPU input projection layer for gradient-based training. - /// Maps flattened order book features (seq_len * feature_dim) -> d_model. - input_projection: ml_core::cuda_autograd::GpuLinear, - - /// GPU output projection layer for gradient-based training. - /// Maps d_model -> 1 (scalar prediction). - output_projection: ml_core::cuda_autograd::GpuLinear, - - /// GPU variable store for model parameters (replaces VarMap + AdamW) - var_store: GpuVarStore, - - /// Device (GPU) - device: MlDevice, - - /// Checkpoint directory - checkpoint_dir: PathBuf, - - /// Best validation loss - best_val_loss: f64, - - /// Training start time - start_time: Option, -} - -impl std::fmt::Debug for TLOBTrainer { - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - f.debug_struct("TLOBTrainer") - .field("hyperparams", &self.hyperparams) - .field("device", &self.device) - .field("checkpoint_dir", &self.checkpoint_dir) - .field("best_val_loss", &self.best_val_loss) - .finish_non_exhaustive() - } -} - -impl TLOBTrainer { - /// Create new TLOB trainer with hyperparameters - /// - /// # Arguments - /// - /// * `hyperparams` - Training hyperparameters from gRPC request - /// * `checkpoint_dir` - Directory for saving model checkpoints - /// * `use_gpu` - Whether to use GPU acceleration (RTX 3050 Ti) - pub fn new>( - hyperparams: TLOBHyperparameters, - checkpoint_dir: P, - use_gpu: bool, - ) -> Result { - // Validate batch size for GPU memory - const MAX_BATCH_SIZE: usize = 32; - if use_gpu && hyperparams.batch_size > MAX_BATCH_SIZE { - warn!( - "Batch size {} exceeds GPU limit ({}), using CPU instead", - hyperparams.batch_size, MAX_BATCH_SIZE - ); - } - - // Create device — CUDA is mandatory, CPU execution FORBIDDEN - let device = if use_gpu && hyperparams.batch_size <= MAX_BATCH_SIZE { - MlDevice::cuda(0).map_err(|e| { - anyhow::anyhow!("CUDA device required but unavailable: {}", e) - })? - } else if !use_gpu { - return Err(anyhow::anyhow!( - "TLOB trainer requires CUDA — CPU execution FORBIDDEN" - )); - } else { - // batch_size > MAX_BATCH_SIZE with GPU requested - return Err(anyhow::anyhow!( - "TLOB batch size {} exceeds GPU limit {} — reduce batch size", - hyperparams.batch_size, - MAX_BATCH_SIZE - )); - }; - info!("Using GPU device: {:?}", device); - - info!( - "Initializing TLOB trainer: seq_len={}, d_model={}, num_layers={}, device={:?}", - hyperparams.seq_len, hyperparams.d_model, hyperparams.num_layers, device - ); - - // Create checkpoint directory - let checkpoint_path = checkpoint_dir.as_ref().to_path_buf(); - std::fs::create_dir_all(&checkpoint_path) - .context("Failed to create checkpoint directory")?; - - // Initialize GPU variable store - let stream = device.cuda_stream().map_err(|e| { - anyhow::anyhow!("TLOB requires CUDA stream: {e}") - })?; - let mut var_store = GpuVarStore::new(Arc::clone(stream)); - - // Create TLOB transformer model (ONNX-based inference / fallback) - let model = Self::create_trainable_model(&hyperparams, &device)?; - - // Create GPU projection layers registered in the var store - // Architecture: flatten(seq_len * feature_dim) -> d_model -> ReLU -> 1 - let input_dim = hyperparams.seq_len * TLOB_FEATURE_COUNT; - let input_projection = var_store.linear( - "input_proj", - input_dim, - hyperparams.d_model, - )?; - let output_projection = var_store.linear( - "output_proj", - hyperparams.d_model, - 1, - )?; - - Ok(Self { - hyperparams, - model: Arc::new(RwLock::new(model)), - input_projection, - output_projection, - var_store, - device, - checkpoint_dir: checkpoint_path, - best_val_loss: f64::INFINITY, - start_time: None, - }) - } - - /// Create trainable TLOB transformer model - /// - /// NOTE: This is a placeholder implementation. The actual TLOBTransformer - /// needs to be updated with a trainable constructor that accepts VarBuilder. - fn create_trainable_model( - hyperparams: &TLOBHyperparameters, - _device: &MlDevice, - ) -> Result { - // Placeholder: Create TLOBTransformer with default config - // This will be replaced when TLOBTransformer gets a trainable constructor - use crate::tlob::transformer::TLOBConfig; - - let config = TLOBConfig { - model_path: "".to_owned(), // Not used for training - feature_dim: TLOB_FEATURE_COUNT, - prediction_horizon: 10, - batch_size: hyperparams.batch_size, - device: "cuda".to_owned(), - }; - - TLOBTransformer::new(config) - .map_err(|e| anyhow::anyhow!("Failed to create TLOB transformer: {:?}", e)) - } - - /// Forward pass through the Candle projection layers. - /// - /// Flattens `(batch, seq_len, feature_dim)` input to `(batch, seq_len*feature_dim)`, - /// then applies `input_projection -> ReLU -> output_projection` to produce `(batch, 1)`. - fn forward_projection(&self, input: &GpuTensor, stream: &Arc) -> Result { - let cublas = cudarc::cublas::CudaBlas::new(Arc::clone(stream)) - .map_err(|e| anyhow::anyhow!("cuBLAS init: {e}"))?; - - // Forward through input projection: [batch, input_dim] -> [batch, d_model] - let (hidden, _acts) = self.input_projection.forward( - input, &self.var_store, &cublas, stream, - )?; - - // ReLU activation on host (small tensor, not hot path for TLOB) - let mut hidden_host = hidden.to_host(stream)?; - for v in &mut hidden_host { - if *v < 0.0 { *v = 0.0; } - } - let activated = GpuTensor::from_host(&hidden_host, hidden.shape().to_vec(), stream)?; // cpu-side shape clone - - // Output projection: [batch, d_model] -> [batch, 1] - let (output, _acts) = self.output_projection.forward( - &activated, &self.var_store, &cublas, stream, - )?; - Ok(output) - } - - /// Train TLOB model on Level-2 order book data - /// - /// # Arguments - /// - /// * `data_dir` - Directory containing L2 order book data (from Agent 71) - /// * `progress_callback` - Callback for reporting progress to gRPC stream - /// - /// # Returns - /// - /// Final training metrics - #[instrument(skip(self, progress_callback))] - pub async fn train( - &mut self, - data_dir: &str, - mut progress_callback: F, - ) -> Result - where - F: FnMut(TLOBTrainingMetrics) + Send, - { - info!( - "Starting TLOB training for {} epochs with batch size {}", - self.hyperparams.epochs, self.hyperparams.batch_size - ); - - self.start_time = Some(Instant::now()); - - // Load order book data (placeholder - requires Agent 71 implementation) - let (train_sequences, val_sequences) = self - .load_order_book_data(data_dir) - .await - .context("Failed to load order book data")?; - - info!( - "Loaded {} training sequences, {} validation sequences", - train_sequences.len(), - val_sequences.len() - ); - - let mut final_metrics = TLOBTrainingMetrics { - epoch: 0, - train_loss: 0.0, - val_loss: 0.0, - avg_mae: 0.0, - avg_prediction_error: 0.0, - gradient_norm: 0.0, - learning_rate: self.hyperparams.learning_rate, - elapsed_seconds: 0.0, - }; - - // Main training loop - for epoch in 0..self.hyperparams.epochs { - let epoch_start = Instant::now(); - - // Training phase - let train_loss = self.train_epoch(&train_sequences).await?; - - // Validation phase - let (val_loss, mae) = self.validate_epoch(&val_sequences).await?; - - // Calculate metrics - let elapsed = self.start_time - .unwrap_or_else(|| Instant::now()) - .elapsed() - .as_secs_f64(); - let grad_norm = self.calculate_gradient_norm()?; - - let metrics = TLOBTrainingMetrics { - epoch: epoch + 1, - train_loss, - val_loss, - avg_mae: mae, - avg_prediction_error: mae, // Same as MAE for regression - gradient_norm: grad_norm, - learning_rate: self.hyperparams.learning_rate, - elapsed_seconds: elapsed, - }; - - // Report progress - progress_callback(metrics.clone()); - final_metrics = metrics.clone(); - - info!( - "Epoch {}/{}: train_loss={:.6}, val_loss={:.6}, mae={:.6}, grad_norm={:.6}, time={:.2}s", - epoch + 1, - self.hyperparams.epochs, - train_loss, - val_loss, - mae, - grad_norm, - epoch_start.elapsed().as_secs_f64() - ); - - // Save checkpoint if best validation loss - if val_loss < self.best_val_loss { - self.best_val_loss = val_loss; - info!("New best validation loss: {:.6}", val_loss); - } - - // Save checkpoint every N epochs - if (epoch + 1) % self.hyperparams.checkpoint_frequency == 0 { - self.save_checkpoint(epoch + 1).await?; - } - } - - info!( - "Training completed in {:.2}s: final_val_loss={:.6}, best_val_loss={:.6}", - final_metrics.elapsed_seconds, final_metrics.val_loss, self.best_val_loss - ); - - Ok(final_metrics) - } - - /// Train one epoch - /// - /// Loss is accumulated on-device as a GPU tensor to avoid per-batch - /// `to_scalar()` synchronisation. A single scalar extraction happens - /// after the loop completes, with a NaN guard every 100 batches. - async fn train_epoch(&mut self, sequences: &[OrderBookSequence]) -> Result { - let stream = self.device.cuda_stream().map_err(|_e| { - anyhow::anyhow!("TLOB train_epoch requires CUDA stream") - })?; - let stream = Arc::clone(stream); - let mut total_loss = 0.0_f64; - let mut num_batches: usize = 0; - - // Process in batches - for batch_sequences in sequences.chunks(self.hyperparams.batch_size) { - let (input_tensor, target_tensor) = self.prepare_batch(batch_sequences, &stream)?; - - // Forward pass - let predictions = self.forward_projection(&input_tensor, &stream)?; - - // Calculate MSE loss on host (small batch) - let pred_host = predictions.to_host(&stream)?; - let target_host = target_tensor.to_host(&stream)?; - let batch_loss: f64 = pred_host.iter().zip(target_host.iter()) - .map(|(p, t)| { let d = (*p - *t) as f64; d * d }) - .sum::() / pred_host.len().max(1) as f64; - - total_loss += batch_loss; - num_batches += 1; - - // NaN guard every 100 batches - if num_batches % 100 == 0 && (total_loss.is_nan() || total_loss.is_infinite()) { - warn!(batch = num_batches, "TLOB train loss diverged (NaN/Inf)"); - return Ok(f64::NAN); - } - } - - if num_batches == 0 { - return Ok(0.0); - } - Ok(total_loss / num_batches as f64) - } - - /// Validate one epoch - /// - /// Loss is accumulated on-device to avoid per-batch `to_scalar()`. - /// MAE already returns `f64` (CPU scalar) so it stays as f64 sum. - async fn validate_epoch(&self, sequences: &[OrderBookSequence]) -> Result<(f64, f64)> { - let stream = self.device.cuda_stream().map_err(|_e| { - anyhow::anyhow!("TLOB validate_epoch requires CUDA stream") - })?; - let stream = Arc::clone(stream); - let mut total_loss = 0.0_f64; - let mut total_mae = 0.0_f64; - let mut num_batches: usize = 0; - - for batch_sequences in sequences.chunks(self.hyperparams.batch_size) { - let (input_tensor, target_tensor) = self.prepare_batch(batch_sequences, &stream)?; - let predictions = self.forward_projection(&input_tensor, &stream)?; - - let pred_host = predictions.to_host(&stream)?; - let target_host = target_tensor.to_host(&stream)?; - - let batch_loss: f64 = pred_host.iter().zip(target_host.iter()) - .map(|(p, t)| { let d = (*p - *t) as f64; d * d }) - .sum::() / pred_host.len().max(1) as f64; - let batch_mae: f64 = pred_host.iter().zip(target_host.iter()) - .map(|(p, t)| (*p - *t).abs() as f64) - .sum::() / pred_host.len().max(1) as f64; - - total_loss += batch_loss; - total_mae += batch_mae; - num_batches += 1; - } - - if num_batches == 0 { - return Ok((0.0, 0.0)); - } - Ok((total_loss / num_batches as f64, total_mae / num_batches as f64)) - } - - /// Prepare batch tensors from order book sequences - fn prepare_batch(&self, sequences: &[OrderBookSequence], stream: &Arc) -> Result<(GpuTensor, GpuTensor)> { - let batch_size = sequences.len(); - let seq_len = self.hyperparams.seq_len; - let feature_dim = TLOB_FEATURE_COUNT; - - let flat_dim = seq_len * feature_dim; - let mut input_data = Vec::with_capacity(batch_size * flat_dim); - let mut target_data = Vec::with_capacity(batch_size); - - for seq in sequences { - for snapshot in &seq.snapshots { - input_data.extend_from_slice(&snapshot.features); - } - target_data.push(seq.target_price_change); - } - - let input_tensor = GpuTensor::from_host(&input_data, vec![batch_size, seq_len, feature_dim], stream)?; - let target_tensor = GpuTensor::from_host(&target_data, vec![batch_size, 1], stream)?; - - Ok((input_tensor, target_tensor)) - } - - /// Calculate gradient norm by downloading parameters (monitoring only). - fn calculate_gradient_norm(&self) -> Result { - // Parameter norm (not gradient norm) — monitors for weight explosion - let stream = self.device.cuda_stream().map_err(|_e| { - anyhow::anyhow!("TLOB requires CUDA stream") - })?; - let mut total_sq = 0.0_f64; - for (name, param) in self.var_store.iter() { - let mut host = vec![0.0_f32; param.data.len()]; - crate::cuda_pipeline::dtoh_f32(&stream, ¶m.data, &mut host) - .map_err(|e| anyhow::anyhow!("param {} DtoH: {e}", name))?; - let sq: f64 = host.iter().map(|&v| (v as f64) * (v as f64)).sum(); - total_sq += sq; - } - Ok(total_sq.sqrt()) - } - - /// Save model checkpoint - #[instrument(skip(self))] - async fn save_checkpoint(&self, epoch: usize) -> Result<()> { - let checkpoint_path = self - .checkpoint_dir - .join(format!("tlob_epoch_{}.safetensors", epoch)); - - info!("Saving checkpoint to: {}", checkpoint_path.display()); - - // GpuVarStore does not yet expose safetensors serialisation, so - // no parameters are written to `checkpoint_path`. The fs::metadata - // call below therefore fails — `save_checkpoint` is currently a - // non-functional hook kept so the training loop has a stable - // call site for when serialisation lands. - let _ = &checkpoint_path; - - info!( - "Checkpoint saved: {} bytes", - std::fs::metadata(&checkpoint_path)?.len() - ); - - Ok(()) - } - - /// Load order book data from directory - /// - /// NOTE: This is a placeholder. Actual implementation depends on Agent 71's - /// TLOBDataLoader for loading Level-2 order book data. - async fn load_order_book_data( - &self, - _data_dir: &str, - ) -> Result<(Vec, Vec)> { - // Placeholder: Generate dummy data for compilation - warn!("Using dummy order book data - Agent 71 L2 data loader not yet implemented"); - - let train_sequences = self.generate_dummy_sequences(100)?; - let val_sequences = self.generate_dummy_sequences(20)?; - - Ok((train_sequences, val_sequences)) - } - - /// Generate dummy order book sequences for testing - fn generate_dummy_sequences(&self, count: usize) -> Result> { - use rand::Rng; - let mut rng = rand::thread_rng(); - - let mut sequences = Vec::with_capacity(count); - - for _ in 0..count { - let mut snapshots = Vec::with_capacity(self.hyperparams.seq_len); - - for _ in 0..self.hyperparams.seq_len { - let features: Vec = (0..TLOB_FEATURE_COUNT) - .map(|_| rng.gen_range(-1.0..1.0)) - .collect(); - - snapshots.push(OrderBookSnapshot { features }); - } - - let target_price_change = rng.gen_range(-0.01..0.01); - - sequences.push(OrderBookSequence { - snapshots, - target_price_change, - }); - } - - Ok(sequences) - } - - /// Serialize model to bytes - pub async fn serialize_model(&self) -> Result> { - let temp_path = - std::env::temp_dir().join(format!("tlob_{}.safetensors", uuid::Uuid::new_v4())); - - // GpuVarStore serialisation is not wired; this function keeps - // the signature for callers but the subsequent `fs::read` will - // error out because nothing writes `temp_path`. - let _ = &temp_path; - - let data = std::fs::read(&temp_path).context("Failed to read checkpoint")?; - - drop(std::fs::remove_file(&temp_path)); - - Ok(data) - } -} - -/// Order book sequence (128 snapshots + target) -#[derive(Debug, Clone)] -struct OrderBookSequence { - snapshots: Vec, - target_price_change: f32, -} - -/// Single order book snapshot (51 features) -#[derive(Debug, Clone)] -struct OrderBookSnapshot { - features: Vec, -} - -#[cfg(test)] -mod tests { - use super::*; - - #[tokio::test] - async fn test_tlob_trainer_creation() { - let hyperparams = TLOBHyperparameters::default(); - let temp_dir = std::env::temp_dir().join("tlob_test"); - - let trainer = TLOBTrainer::new(hyperparams, &temp_dir, true); - assert!( - trainer.is_ok(), - "Failed to create TLOB trainer: {:?}", - trainer.err() - ); - } - - #[tokio::test] - async fn test_batch_size_validation() { - let mut hyperparams = TLOBHyperparameters::default(); - hyperparams.batch_size = 64; - - let temp_dir = std::env::temp_dir().join("tlob_test_batch"); - let trainer = TLOBTrainer::new(hyperparams, &temp_dir, true); - - assert!( - trainer.is_err(), - "Batch size 64 exceeds GPU limit 32 — should be rejected" - ); - } - - #[tokio::test] - async fn test_dummy_sequence_generation() { - let hyperparams = TLOBHyperparameters::default(); - let temp_dir = std::env::temp_dir().join("tlob_test_seq"); - - let trainer = TLOBTrainer::new(hyperparams, &temp_dir, true).unwrap(); - let sequences = trainer.generate_dummy_sequences(10).unwrap(); - - assert_eq!(sequences.len(), 10); - assert_eq!(sequences[0].snapshots.len(), 128); - assert_eq!(sequences[0].snapshots[0].features.len(), TLOB_FEATURE_COUNT); - } - - #[tokio::test] - async fn test_batch_preparation() { - let hyperparams = TLOBHyperparameters::default(); - let temp_dir = std::env::temp_dir().join("tlob_test_batch_prep"); - - let trainer = TLOBTrainer::new(hyperparams, &temp_dir, true).unwrap(); - let sequences = trainer.generate_dummy_sequences(4).unwrap(); - - let dev = MlDevice::cuda(0).expect("CUDA required"); - let stream = dev.cuda_stream().expect("stream"); - let (input_tensor, target_tensor) = trainer.prepare_batch(&sequences, stream).unwrap(); - - assert_eq!(input_tensor.dims(), &[4, 128, TLOB_FEATURE_COUNT]); - assert_eq!(target_tensor.dims(), &[4, 1]); - } -}