f946dcd952
WAVE 22: All examples, benchmarks, and data loaders updated Files Modified (41 files): - DQN examples: 7 files (train_dqn, evaluate_dqn, validate_dqn, etc.) - PPO examples: 6 files (train_ppo, continuous_ppo, benchmark_ppo, etc.) - TFT examples: 9 files (train_tft, validate_tft, benchmark_tft, etc.) - MAMBA-2 examples: 3 files (train_mamba2, verify_dimensions, etc.) - Benchmarks: 5 files (cuda_speedup, weight_caching, future_decoder, etc.) - Data loaders: 7 files (parquet_utils, dbn_sequence_loader, tlob_loader, etc.) - Integration: 4 files (load_parquet_data, streaming loaders, etc.) Key Changes: - state_dim: 225 → 54 (DQN, PPO) - input_dim: 225 → 54 (TFT) - d_model: 225 → 54 (MAMBA-2) - Memory: 1.8KB → 0.43KB per vector (76% reduction) - All tensor shapes updated: (batch, 225) → (batch, 54) Agents Deployed: 5 parallel agents Validation: cargo check PASSING Generated with Claude Code Co-Authored-By: Claude <noreply@anthropic.com>
487 lines
17 KiB
Rust
487 lines
17 KiB
Rust
//! TFT (Temporal Fusion Transformer) Training with Parquet Data
|
||
//!
|
||
//! Trains a TFT model using market data from Parquet files with lazy batch loading
|
||
//! to avoid OOM issues on large datasets. Uses the TFTParquetExt trait for efficient
|
||
//! memory management and 54-feature extraction (Wave E).
|
||
//!
|
||
//! # Usage
|
||
//!
|
||
//! ```bash
|
||
//! # Train with default parameters (20 epochs)
|
||
//! cargo run -p ml --example train_tft_parquet --release --features cuda
|
||
//!
|
||
//! # Custom configuration
|
||
//! cargo run -p ml --example train_tft_parquet --release --features cuda -- \
|
||
//! --parquet-file test_data/ES_FUT_180d.parquet \
|
||
//! --epochs 50 \
|
||
//! --batch-size 32 \
|
||
//! --lookback-window 60 \
|
||
//! --forecast-horizon 10
|
||
//! ```
|
||
//!
|
||
//! # Features
|
||
//!
|
||
//! - Lazy batch loading (10,000 rows at a time) to avoid OOM crashes
|
||
//! - 54-feature extraction (Wave E) from OHLCV bars
|
||
//! - Sliding window creation (configurable lookback/horizon)
|
||
//! - GPU-accelerated training (RTX 3050 Ti, 4GB VRAM)
|
||
//! - Automatic train/validation split (80/20)
|
||
//! - Model checkpointing and early stopping
|
||
//!
|
||
//! # Parquet Schema Requirements
|
||
//!
|
||
//! The Parquet file must follow Databento schema:
|
||
//! - Column 3: open (Float64)
|
||
//! - Column 4: high (Float64)
|
||
//! - Column 5: low (Float64)
|
||
//! - Column 6: close (Float64)
|
||
//! - Column 7: volume (UInt64)
|
||
//! - Column 9: ts_event (Timestamp[ns, UTC])
|
||
|
||
// Suppress warnings for unused dependencies in this example
|
||
// (examples have access to all crate dependencies but typically only use a subset)
|
||
#![allow(unused_crate_dependencies)]
|
||
|
||
// Use mimalloc allocator for 10-25% performance improvement
|
||
#[cfg(feature = "mimalloc-allocator")]
|
||
use mimalloc::MiMalloc;
|
||
#[cfg(feature = "mimalloc-allocator")]
|
||
#[global_allocator]
|
||
static GLOBAL: MiMalloc = MiMalloc;
|
||
|
||
use anyhow::{Context, Result};
|
||
use clap::Parser;
|
||
use std::path::PathBuf;
|
||
use tokio::sync::mpsc;
|
||
use tracing::{info, warn};
|
||
use tracing_subscriber::FmtSubscriber;
|
||
|
||
use ml::checkpoint::FileSystemStorage;
|
||
use ml::trainers::tft::{TFTTrainer, TFTTrainerConfig};
|
||
|
||
#[derive(Debug, Parser)]
|
||
#[command(
|
||
name = "train_tft_parquet",
|
||
about = "Train TFT model on Parquet market data with lazy loading"
|
||
)]
|
||
struct Opts {
|
||
/// Parquet file path containing OHLCV bars (Databento schema)
|
||
#[arg(long, default_value = "test_data/ES_FUT_small.parquet")]
|
||
parquet_file: String,
|
||
|
||
/// Number of training epochs
|
||
#[arg(long, default_value = "3")]
|
||
epochs: usize,
|
||
|
||
/// Learning rate
|
||
#[arg(long, default_value = "0.001")]
|
||
learning_rate: f64,
|
||
|
||
/// Batch size (max 32 for 4GB VRAM)
|
||
#[arg(long, default_value = "32")]
|
||
batch_size: usize,
|
||
|
||
/// Validation batch size (defaults to match training batch_size)
|
||
#[arg(long)]
|
||
validation_batch_size: Option<usize>,
|
||
|
||
/// Hidden dimension
|
||
#[arg(long, default_value = "256")]
|
||
hidden_dim: usize,
|
||
|
||
/// Number of attention heads
|
||
#[arg(long, default_value = "8")]
|
||
num_attention_heads: usize,
|
||
|
||
/// Lookback window (historical bars)
|
||
#[arg(long, default_value = "60")]
|
||
lookback_window: usize,
|
||
|
||
/// Forecast horizon (future bars)
|
||
#[arg(long, default_value = "10")]
|
||
forecast_horizon: usize,
|
||
|
||
/// Dropout rate for regularization
|
||
#[arg(long, default_value = "0.1")]
|
||
dropout_rate: f64,
|
||
|
||
/// Number of LSTM layers
|
||
#[arg(long, default_value = "2")]
|
||
lstm_layers: usize,
|
||
|
||
/// Quantiles for probabilistic forecasting (comma-separated)
|
||
#[arg(long, default_value = "0.1,0.5,0.9")]
|
||
quantiles: String,
|
||
|
||
/// Output directory for trained model checkpoints
|
||
#[arg(long, default_value = "ml/trained_models")]
|
||
output_dir: String,
|
||
|
||
/// Use GPU for training (CUDA required)
|
||
#[arg(long)]
|
||
use_gpu: bool,
|
||
|
||
/// Use INT8 quantization for memory efficiency (reduces VRAM usage by 3-8x)
|
||
#[arg(long)]
|
||
use_int8: bool,
|
||
|
||
/// Use Quantization-Aware Training (1-2% better accuracy than PTQ)
|
||
/// Trains with fake quantization, converts to INT8 at the end
|
||
#[arg(long)]
|
||
use_qat: bool,
|
||
|
||
/// Number of batches for QAT calibration (default: 100)
|
||
/// Higher values improve accuracy but increase training time
|
||
#[arg(long, default_value = "100")]
|
||
qat_calibration_batches: usize,
|
||
|
||
/// Minimum batch size for QAT calibration OOM recovery (default: 2)
|
||
/// If OOM occurs, batch size is halved automatically. Training aborts if below this threshold.
|
||
#[arg(long, default_value = "2")]
|
||
qat_min_batch_size: usize,
|
||
|
||
/// ⚠️ WARNING: Gradient checkpointing NOT IMPLEMENTED for QAT (see QAT_GRADIENT_CHECKPOINTING_WORKAROUND.md)
|
||
///
|
||
/// This flag is IGNORED when --use-qat is enabled. For QAT memory reduction:
|
||
/// 1. Use 2-phase workaround (calibrate → freeze observers → train with checkpointing)
|
||
/// 2. See ml/docs/QAT_GRADIENT_CHECKPOINTING_WORKAROUND.md for details
|
||
///
|
||
/// For non-QAT training: Reduces GPU memory usage by 30-40% but increases training time by ~20%
|
||
#[arg(long)]
|
||
use_gradient_checkpointing: bool,
|
||
|
||
/// Auto-detect optimal batch size based on available GPU memory
|
||
/// Overrides --batch-size if enabled. Prevents OOM errors and maximizes GPU utilization.
|
||
#[arg(long)]
|
||
auto_batch_size: bool,
|
||
|
||
/// Verbose logging (debug level)
|
||
#[arg(short, long)]
|
||
verbose: bool,
|
||
|
||
/// Maximum validation batches to run (default: unlimited, use 50 for 4GB GPUs)
|
||
/// Limits validation to N batches to reduce memory usage. Each batch uses ~10MB,
|
||
/// so 50 batches = ~500MB vs 1760MB for full validation (176 batches).
|
||
#[arg(long)]
|
||
max_validation_batches: Option<usize>,
|
||
}
|
||
|
||
#[tokio::main]
|
||
async fn main() -> Result<()> {
|
||
// Parse CLI options
|
||
let opts = Opts::parse();
|
||
|
||
// Setup logging
|
||
let level = if opts.verbose {
|
||
tracing::Level::DEBUG
|
||
} else {
|
||
tracing::Level::INFO
|
||
};
|
||
|
||
let subscriber = FmtSubscriber::builder().with_max_level(level).finish();
|
||
tracing::subscriber::set_global_default(subscriber)
|
||
.context("Failed to set tracing subscriber")?;
|
||
|
||
#[cfg(feature = "mimalloc-allocator")]
|
||
info!("🚀 Using mimalloc allocator for improved performance");
|
||
#[cfg(not(feature = "mimalloc-allocator"))]
|
||
info!("ℹ️ Using system allocator (consider --features mimalloc-allocator for 10-25% speedup)");
|
||
info!("🚀 Starting TFT Training with Parquet Data (Lazy Loading)");
|
||
info!("");
|
||
info!("Configuration:");
|
||
info!(" • Parquet file: {}", opts.parquet_file);
|
||
info!(" • Epochs: {}", opts.epochs);
|
||
info!(" • Learning rate: {}", opts.learning_rate);
|
||
info!(" • Batch size: {}", opts.batch_size);
|
||
info!(
|
||
" • Validation batch size: {}",
|
||
opts.validation_batch_size.unwrap_or(opts.batch_size)
|
||
);
|
||
if let Some(max_val_batches) = opts.max_validation_batches {
|
||
info!(
|
||
" • Max validation batches: {} (memory optimization)",
|
||
max_val_batches
|
||
);
|
||
} else {
|
||
info!(" • Max validation batches: unlimited");
|
||
}
|
||
info!(" • Hidden dimension: {}", opts.hidden_dim);
|
||
info!(" • Attention heads: {}", opts.num_attention_heads);
|
||
info!(" • Lookback window: {}", opts.lookback_window);
|
||
info!(" • Forecast horizon: {}", opts.forecast_horizon);
|
||
info!(" • Dropout rate: {}", opts.dropout_rate);
|
||
info!(" • LSTM layers: {}", opts.lstm_layers);
|
||
info!(" • Quantiles: {}", opts.quantiles);
|
||
info!(" • Feature count: 54 (Wave E)");
|
||
info!(" • GPU enabled: {}", opts.use_gpu);
|
||
info!(" • INT8 quantization: {}", opts.use_int8);
|
||
info!(" • Quantization-Aware Training: {}", opts.use_qat);
|
||
if opts.use_qat {
|
||
info!(
|
||
" • QAT calibration batches: {}",
|
||
opts.qat_calibration_batches
|
||
);
|
||
}
|
||
info!(
|
||
" • Gradient checkpointing: {}",
|
||
opts.use_gradient_checkpointing
|
||
);
|
||
if opts.use_gradient_checkpointing {
|
||
if opts.use_qat {
|
||
warn!("⚠️ WARNING: --use-gradient-checkpointing is IGNORED with --use-qat (not implemented)");
|
||
warn!(" → For QAT memory reduction, use 2-phase workaround:");
|
||
warn!(" → See ml/docs/QAT_GRADIENT_CHECKPOINTING_WORKAROUND.md");
|
||
} else {
|
||
info!(" → Expected: 30-40% memory reduction, ~20% slower training");
|
||
}
|
||
}
|
||
info!(" • Output directory: {}", opts.output_dir);
|
||
info!("");
|
||
|
||
// Verify Parquet file exists
|
||
let parquet_path = PathBuf::from(&opts.parquet_file);
|
||
if !parquet_path.exists() {
|
||
return Err(anyhow::anyhow!(
|
||
"Parquet file not found: {}",
|
||
opts.parquet_file
|
||
));
|
||
}
|
||
|
||
// Create output directory
|
||
let output_path = PathBuf::from(&opts.output_dir);
|
||
if !output_path.exists() {
|
||
std::fs::create_dir_all(&output_path).context("Failed to create output directory")?;
|
||
info!("✅ Created output directory: {}", opts.output_dir);
|
||
}
|
||
|
||
// Parse quantiles
|
||
let quantiles: Vec<f64> = opts
|
||
.quantiles
|
||
.split(',')
|
||
.map(|s| s.trim().parse::<f64>())
|
||
.collect::<Result<Vec<_>, _>>()
|
||
.context("Failed to parse quantiles")?;
|
||
|
||
if quantiles.is_empty() || quantiles.len() > 10 {
|
||
return Err(anyhow::anyhow!(
|
||
"Invalid number of quantiles: {} (must be 1-10)",
|
||
quantiles.len()
|
||
));
|
||
}
|
||
|
||
info!(
|
||
"📊 Quantiles for probabilistic forecasting: {:?}",
|
||
quantiles
|
||
);
|
||
|
||
// Configure TFT trainer
|
||
// Static features: 5 (symbol metadata)
|
||
// Historical features: 39 (Wave E 54 features - 5 static - 10 known)
|
||
// Future features: 10 (calendar features, time-based)
|
||
// Total input features: 54 (5 + 10 + 39)
|
||
let trainer_config = TFTTrainerConfig {
|
||
epochs: opts.epochs,
|
||
learning_rate: opts.learning_rate,
|
||
batch_size: opts.batch_size,
|
||
auto_batch_size: opts.auto_batch_size,
|
||
validation_batch_size: opts.validation_batch_size.unwrap_or(opts.batch_size),
|
||
hidden_dim: opts.hidden_dim,
|
||
num_attention_heads: opts.num_attention_heads,
|
||
dropout_rate: opts.dropout_rate,
|
||
lstm_layers: opts.lstm_layers,
|
||
quantiles,
|
||
lookback_window: opts.lookback_window,
|
||
forecast_horizon: opts.forecast_horizon,
|
||
use_gpu: opts.use_gpu,
|
||
use_int8_quantization: opts.use_int8,
|
||
use_qat: opts.use_qat,
|
||
qat_calibration_batches: opts.qat_calibration_batches,
|
||
qat_min_batch_size: opts.qat_min_batch_size,
|
||
qat_warmup_epochs: 10, // Default: 10 epochs LR warmup after calibration
|
||
qat_cooldown_factor: 0.1, // Default: 10x LR reduction in final 10% of training
|
||
use_gradient_checkpointing: opts.use_gradient_checkpointing,
|
||
max_validation_batches: opts.max_validation_batches,
|
||
validation_frequency: 1, // Validate every epoch
|
||
checkpoint_dir: opts.output_dir.clone(),
|
||
};
|
||
|
||
// Create checkpoint storage
|
||
let storage = std::sync::Arc::new(FileSystemStorage::new(output_path.clone()));
|
||
|
||
// Create TFT trainer
|
||
let mut trainer =
|
||
TFTTrainer::new(trainer_config.clone(), storage).context("Failed to create TFT trainer")?;
|
||
|
||
info!(
|
||
"✅ TFT trainer initialized with {} quantiles",
|
||
trainer_config.quantiles.len()
|
||
);
|
||
|
||
if opts.use_qat {
|
||
info!("🧠 Quantization-Aware Training (QAT) enabled");
|
||
info!(
|
||
" Phase 1: Calibration ({} batches) - collecting activation statistics",
|
||
opts.qat_calibration_batches
|
||
);
|
||
info!(" Phase 2: Training with fake quantization - simulating INT8 ops");
|
||
info!(" Phase 3: Conversion to true INT8 model");
|
||
info!(" Expected: 1-2% better accuracy than post-training quantization");
|
||
} else if opts.use_int8 {
|
||
info!("⚡ INT8 quantization enabled (PTQ mode) - expect 3-8x memory reduction");
|
||
info!(" Memory usage: ~125MB (vs ~1GB FP32)");
|
||
}
|
||
|
||
// Setup progress callback
|
||
let (progress_tx, mut progress_rx) = mpsc::unbounded_channel();
|
||
trainer.set_progress_callback(progress_tx);
|
||
|
||
// Spawn progress monitor task
|
||
let monitor_task = tokio::spawn(async move {
|
||
while let Some(progress) = progress_rx.recv().await {
|
||
info!("{}", progress.message);
|
||
if let Some(loss) = progress.metrics.get("train_loss") {
|
||
info!(" • Train loss: {:.6}", loss);
|
||
}
|
||
if let Some(val_loss) = progress.metrics.get("val_loss") {
|
||
info!(" • Val loss: {:.6}", val_loss);
|
||
}
|
||
if let Some(quantile_loss) = progress.metrics.get("quantile_loss") {
|
||
info!(" • Quantile loss: {:.6}", quantile_loss);
|
||
}
|
||
if let Some(rmse) = progress.metrics.get("rmse") {
|
||
info!(" • RMSE: {:.6}", rmse);
|
||
}
|
||
if let Some(attention_entropy) = progress.metrics.get("attention_entropy") {
|
||
info!(" • Attention entropy: {:.4}", attention_entropy);
|
||
}
|
||
}
|
||
});
|
||
|
||
// Train the model using lazy-loading Parquet pipeline
|
||
info!("");
|
||
info!("🏋️ Starting training with lazy-loading Parquet pipeline...");
|
||
info!(" (Loading 10,000 rows at a time to avoid OOM)");
|
||
info!("");
|
||
let start_time = std::time::Instant::now();
|
||
|
||
let final_metrics = trainer
|
||
.train_from_parquet(&opts.parquet_file)
|
||
.await
|
||
.context("Training failed")?;
|
||
|
||
let training_duration = start_time.elapsed();
|
||
|
||
// Wait for progress monitor to finish
|
||
drop(trainer); // Drop trainer to close progress channel
|
||
let _ = monitor_task.await;
|
||
|
||
// Print final metrics
|
||
info!("");
|
||
info!("✅ Training completed successfully!");
|
||
info!("");
|
||
info!("📊 Final Metrics:");
|
||
info!(" • Training loss: {:.6}", final_metrics.train_loss);
|
||
info!(" • Validation loss: {:.6}", final_metrics.val_loss);
|
||
info!(" • Quantile loss: {:.6}", final_metrics.quantile_loss);
|
||
info!(" • RMSE: {:.6}", final_metrics.rmse);
|
||
info!(
|
||
" • Attention entropy: {:.4}",
|
||
final_metrics.attention_entropy
|
||
);
|
||
info!(
|
||
" • Training duration: {:.1}s ({:.1} min)",
|
||
training_duration.as_secs_f64(),
|
||
training_duration.as_secs_f64() / 60.0
|
||
);
|
||
info!(
|
||
" • Reported training time: {:.1}s ({:.1} min)",
|
||
final_metrics.training_time_seconds,
|
||
final_metrics.training_time_seconds / 60.0
|
||
);
|
||
info!("");
|
||
info!("💾 Model checkpoints saved to: {}", opts.output_dir);
|
||
info!("");
|
||
info!("🎉 TFT training with Parquet data complete!");
|
||
|
||
Ok(())
|
||
}
|
||
|
||
#[cfg(test)]
|
||
mod tests {
|
||
use super::*;
|
||
|
||
#[test]
|
||
fn test_cli_parsing() {
|
||
// Test default parameters
|
||
let args = vec!["train_tft_parquet"];
|
||
let opts = Opts::try_parse_from(args).expect("Failed to parse default args");
|
||
|
||
assert_eq!(opts.parquet_file, "test_data/ES_FUT_180d.parquet");
|
||
assert_eq!(opts.epochs, 20);
|
||
assert_eq!(opts.batch_size, 32);
|
||
assert_eq!(opts.lookback_window, 60);
|
||
assert_eq!(opts.forecast_horizon, 10);
|
||
assert_eq!(opts.learning_rate, 0.001);
|
||
assert!(!opts.use_gpu); // Default is false (flag not set)
|
||
}
|
||
|
||
#[test]
|
||
fn test_cli_custom_parameters() {
|
||
// Test custom parameters with GPU enabled
|
||
let args = vec![
|
||
"train_tft_parquet",
|
||
"--parquet-file",
|
||
"test_data/NQ_FUT_180d.parquet",
|
||
"--epochs",
|
||
"50",
|
||
"--batch-size",
|
||
"16",
|
||
"--lookback-window",
|
||
"120",
|
||
"--forecast-horizon",
|
||
"20",
|
||
"--learning-rate",
|
||
"0.0005",
|
||
"--use-gpu", // Flag to enable GPU
|
||
"--verbose",
|
||
];
|
||
|
||
let opts = Opts::try_parse_from(args).expect("Failed to parse custom args");
|
||
|
||
assert_eq!(opts.parquet_file, "test_data/NQ_FUT_180d.parquet");
|
||
assert_eq!(opts.epochs, 50);
|
||
assert_eq!(opts.batch_size, 16);
|
||
assert_eq!(opts.lookback_window, 120);
|
||
assert_eq!(opts.forecast_horizon, 20);
|
||
assert_eq!(opts.learning_rate, 0.0005);
|
||
assert!(opts.use_gpu); // GPU should be enabled
|
||
assert!(opts.verbose);
|
||
}
|
||
|
||
#[test]
|
||
fn test_quantile_parsing() {
|
||
let quantiles_str = "0.1,0.5,0.9";
|
||
let quantiles: Vec<f64> = quantiles_str
|
||
.split(',')
|
||
.map(|s| s.trim().parse::<f64>())
|
||
.collect::<Result<Vec<_>, _>>()
|
||
.expect("Failed to parse quantiles");
|
||
|
||
assert_eq!(quantiles.len(), 3);
|
||
assert_eq!(quantiles[0], 0.1);
|
||
assert_eq!(quantiles[1], 0.5);
|
||
assert_eq!(quantiles[2], 0.9);
|
||
}
|
||
|
||
#[test]
|
||
fn test_invalid_quantiles() {
|
||
let quantiles_str = "invalid,0.5,0.9";
|
||
let result: Result<Vec<f64>, _> = quantiles_str
|
||
.split(',')
|
||
.map(|s| s.trim().parse::<f64>())
|
||
.collect();
|
||
|
||
assert!(result.is_err(), "Should fail on invalid quantile");
|
||
}
|
||
}
|