foxhunt/ml/examples/train_mamba2.rs

//! MAMBA-2 Training Example
//!
//! Trains a MAMBA-2 state space model on real market data from DBN files.
//! Uses continuous price sequences with microstructure features for next-timestep prediction.
//!
//! # Usage
//!
//! ```bash
//! # Train with default parameters (100 epochs, real DBN data)
//! cargo run -p ml --example train_mamba2 --release --features cuda
//!
//! # Custom parameters with specific DBN directory
//! cargo run -p ml --example train_mamba2 --release --features cuda -- \
//!   --epochs 500 \
//!   --d-model 256 \
//!   --n-layers 6 \
//!   --seq-len 60 \
//!   --dbn-dir test_data/real/databento/ml_training_small
//!
//! # Quick test with fewer epochs
//! cargo run -p ml --example train_mamba2 --release --features cuda -- \
//!   --epochs 10 \
//!   --batch-size 4
//! ```
//!
//! # Data Requirements
//!
//! - DBN files with OHLCV data (1-minute bars recommended)
//! - At least 60-128 consecutive timesteps per symbol
//! - Multiple files per symbol for better training data

use anyhow::{Context, Result};
use clap::Parser;
use std::path::PathBuf;
use tracing::info;
use tracing_subscriber::FmtSubscriber;

use ml::data_loaders::DbnSequenceLoader;
use ml::trainers::mamba2::{Mamba2Hyperparameters, Mamba2Trainer};

#[derive(Debug, Parser)]
#[command(name = "train_mamba2", about = "Train MAMBA-2 model on market data")]
struct Opts {
    /// Number of training epochs
    #[arg(long, default_value = "100")]
    epochs: usize,

    /// Learning rate
    #[arg(long, default_value = "0.0001")]
    learning_rate: f64,

    /// Batch size (1-16 for 4GB VRAM)
    #[arg(long, default_value = "8")]
    batch_size: usize,

    /// Model dimension (256, 512, 1024)
    #[arg(long, default_value = "256")]
    d_model: usize,

    /// Number of layers (4-12)
    #[arg(long, default_value = "6")]
    n_layers: usize,

    /// Sequence length
    #[arg(long, default_value = "128")]
    seq_len: usize,

    /// Output directory for trained model
    #[arg(long, default_value = "ml/trained_models")]
    output_dir: String,

    /// DBN data directory (contains .dbn files)
    #[arg(long, default_value = "test_data/real/databento/ml_training_small")]
    dbn_dir: String,

    /// Train/validation split ratio
    #[arg(long, default_value = "0.9")]
    train_split: f64,

    /// Verbose logging
    #[arg(short, long)]
    verbose: bool,
}

#[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")?;

    info!("🚀 Starting MAMBA-2 Training");
    info!("Configuration:");
    info!("  • Epochs: {}", opts.epochs);
    info!("  • Learning rate: {}", opts.learning_rate);
    info!("  • Batch size: {}", opts.batch_size);
    info!("  • Model dimension: {}", opts.d_model);
    info!("  • Number of layers: {}", opts.n_layers);
    info!("  • Sequence length: {}", opts.seq_len);
    info!("  • Output directory: {}", opts.output_dir);

    // 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);
    }

    // Configure MAMBA-2 hyperparameters
    let hyperparams = Mamba2Hyperparameters {
        learning_rate: opts.learning_rate,
        batch_size: opts.batch_size,
        d_model: opts.d_model,
        n_layers: opts.n_layers,
        state_size: 32,
        dropout: 0.1,
        epochs: opts.epochs,
        seq_len: opts.seq_len,
        grad_clip: 1.0,
        weight_decay: 1e-4,
        warmup_steps: 1000,
    };

    // Validate hyperparameters for VRAM constraint
    hyperparams
        .validate()
        .context("Invalid hyperparameters for 4GB VRAM")?;

    info!(
        "✅ Hyperparameters validated (estimated VRAM: {}MB)",
        hyperparams.estimate_memory_usage()
    );

    // Create MAMBA-2 trainer
    let checkpoint_path = format!("{}/mamba2", opts.output_dir);
    let mut trainer = Mamba2Trainer::new(hyperparams.clone(), Some(checkpoint_path))
        .context("Failed to create MAMBA-2 trainer")?;

    info!(
        "✅ MAMBA-2 trainer initialized (job_id: {})",
        trainer.job_id
    );

    // Load real DBN market data sequences
    info!("\n📊 Loading DBN market data sequences...");
    info!("  • DBN directory: {}", opts.dbn_dir);
    info!("  • Sequence length: {}", opts.seq_len);
    info!("  • Feature dimension: {}", opts.d_model);
    info!(
        "  • Train/val split: {:.1}/{:.1}",
        opts.train_split * 100.0,
        (1.0 - opts.train_split) * 100.0
    );

    let mut loader = DbnSequenceLoader::new(opts.seq_len, opts.d_model)
        .await
        .context("Failed to create DBN sequence loader")?;

    let (train_data, val_data) = loader
        .load_sequences(&opts.dbn_dir, opts.train_split)
        .await
        .context("Failed to load DBN sequences")?;

    info!(
        "✅ Loaded {} training sequences, {} validation sequences",
        train_data.len(),
        val_data.len()
    );

    if train_data.is_empty() {
        return Err(anyhow::anyhow!(
            "No training sequences loaded! Check DBN directory: {}",
            opts.dbn_dir
        ));
    }

    // Log sequence shape information
    if let Some((input, target)) = train_data.first() {
        info!("  • Input shape: {:?}", input.dims());
        info!("  • Target shape: {:?}", target.dims());
    }

    // Set progress callback
    let progress_callback =
        std::sync::Arc::new(move |progress: ml::trainers::mamba2::TrainingProgress| {
            if progress.epoch % 10 == 0 {
                info!(
                    "📊 Epoch {}/{} ({:.1}%): loss={:.6}, perplexity={:.2}",
                    progress.epoch,
                    progress.total_epochs,
                    progress.progress_percentage,
                    progress.metrics.loss,
                    progress.metrics.perplexity
                );
            }
        });

    trainer.set_progress_callback(progress_callback);

    // Train the model
    info!("\n🏋️  Starting training...\n");
    let start_time = std::time::Instant::now();

    let training_history = trainer
        .train(&train_data, &val_data)
        .await
        .context("Training failed")?;

    let training_duration = start_time.elapsed();

    // Print final metrics
    info!("\n✅ Training completed successfully!");
    info!("\n📊 Final Metrics:");
    if let Some(final_epoch) = training_history.last() {
        info!("  • Final loss: {:.6}", final_epoch.loss);
        info!("  • Perplexity: {:.2}", final_epoch.loss.exp());
    }
    info!("  • Best validation loss: {:.6}", trainer.best_val_loss);
    info!("  • Epochs trained: {}", training_history.len());
    info!(
        "  • Training time: {:.1}s ({:.1} min)",
        training_duration.as_secs_f64(),
        training_duration.as_secs_f64() / 60.0
    );

    // Get training statistics
    let stats = trainer.get_training_statistics();
    info!("\n📈 Training Statistics:");
    if let Some(&memory_mb) = stats.get("estimated_memory_mb") {
        info!("  • Memory usage: {:.1}MB", memory_mb);
    }
    if let Some(&throughput) = stats.get("throughput_pps") {
        info!("  • Throughput: {:.0} predictions/sec", throughput);
    }

    info!(
        "\n💾 Model checkpoints saved to: {}",
        trainer.checkpoint_path
    );
    info!("\n🎉 MAMBA-2 training complete!");

    Ok(())
}