foxhunt/ml/examples/train_mamba2.rs

//! MAMBA-2 Training Example
//!
//! Trains a MAMBA-2 state space model on market sequences and saves checkpoints.
//!
//! # Usage
//!
//! ```bash
//! # Train with default parameters (100 epochs)
//! cargo run -p ml --example train_mamba2 --release --features cuda
//!
//! # Custom epochs and model size
//! cargo run -p ml --example train_mamba2 --release --features cuda -- \
//!   --epochs 500 \
//!   --d-model 256 \
//!   --n-layers 6
//! ```

use anyhow::{Context, Result};
use candle_core::Tensor;
use std::path::PathBuf;
use structopt::StructOpt;
use tracing::{info};
use tracing_subscriber::FmtSubscriber;

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

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

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

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

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

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

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

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

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

#[tokio::main]
async fn main() -> Result<()> {
    // Parse CLI options
    let opts = Opts::from_args();

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

    // Generate synthetic sequence data
    info!("\n📊 Generating training sequences...");
    let num_sequences = 1000;
    let mut train_data = Vec::with_capacity(num_sequences);
    let mut val_data = Vec::with_capacity(num_sequences / 10);

    let device = candle_core::Device::cuda_if_available(0)
        .unwrap_or(candle_core::Device::Cpu);

    for i in 0..num_sequences {
        // Generate synthetic sequence (input, target) pairs
        let seq_data: Vec<f32> = (0..opts.seq_len)
            .map(|j| (i as f32 * 0.01 + j as f32 * 0.1).sin())
            .collect();

        let input = Tensor::from_slice(&seq_data, (1, opts.seq_len), &device)
            .context("Failed to create input tensor")?;

        // Target is input shifted by 1
        let mut target_data = seq_data.clone();
        target_data.rotate_left(1);
        let target = Tensor::from_slice(&target_data, (1, opts.seq_len), &device)
            .context("Failed to create target tensor")?;

        // 90% train, 10% validation
        if i < (num_sequences * 9 / 10) {
            train_data.push((input, target));
        } else {
            val_data.push((input, target));
        }
    }

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

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