foxhunt/ml/examples/mamba2_simple_train.rs

//! Simple MAMBA-2 Training Script
//!
//! **Agent 40: Simplified Production Run**
//!
//! This is a simplified version that works around the TFT compilation issue.
//! It directly uses MAMBA-2's training interface without the full trainer wrapper.

use anyhow::{Context, Result};
use candle_core::{Device, Tensor};
use std::time::Instant;
use tracing::{info, warn};

use ml::mamba::{Mamba2Config, Mamba2SSM};

#[tokio::main]
async fn main() -> Result<()> {
    // Initialize tracing
    tracing_subscriber::fmt()
        .with_max_level(tracing::Level::INFO)
        .with_target(false)
        .init();

    info!("=== MAMBA-2 Simple Training Script ===");

    // Configuration
    let config = Mamba2Config {
        d_model: 256,
        d_state: 32,
        d_head: 32,
        num_heads: 8,
        expand: 2,
        num_layers: 6,
        dropout: 0.1,
        use_ssd: true,
        use_selective_state: true,
        hardware_aware: true,
        target_latency_us: 5,
        max_seq_len: 256,
        learning_rate: 0.0001,
        weight_decay: 1e-4,
        grad_clip: 1.0,
        warmup_steps: 1000,
        batch_size: 16,
        seq_len: 128,
    };

    info!("Creating MAMBA-2 model...");
    let mut model = Mamba2SSM::new(config.clone())?;

    // Generate synthetic training data
    let device = Device::cuda_if_available(0).unwrap_or(Device::Cpu);
    info!("Using device: {:?}", device);

    let num_train = 800;
    let num_val = 200;

    info!("Generating {} training sequences...", num_train);
    let train_data: Vec<(Tensor, Tensor)> = (0..num_train)
        .map(|_| {
            let input = Tensor::randn(
                0.0,
                1.0,
                &[config.batch_size, config.seq_len, config.d_model],
                &device,
            )
            .unwrap();
            let target = Tensor::randn(0.0, 1.0, &[config.batch_size, 1], &device).unwrap();
            (input, target)
        })
        .collect();

    info!("Generating {} validation sequences...", num_val);
    let val_data: Vec<(Tensor, Tensor)> = (0..num_val)
        .map(|_| {
            let input = Tensor::randn(
                0.0,
                1.0,
                &[config.batch_size, config.seq_len, config.d_model],
                &device,
            )
            .unwrap();
            let target = Tensor::randn(0.0, 1.0, &[config.batch_size, 1], &device).unwrap();
            (input, target)
        })
        .collect();

    // Train model
    let epochs = 100; // Reduced from 500 for quick validation
    info!("Starting training for {} epochs...", epochs);

    let start_time = Instant::now();
    let history = model
        .train(&train_data, &val_data, epochs)
        .await
        .context("Training failed")?;

    let elapsed = start_time.elapsed();
    info!("Training completed in {:.2}s", elapsed.as_secs_f64());

    // Analyze results
    info!("\n=== Training Results ===");
    if let Some(first) = history.first() {
        info!("Initial loss: {:.6}", first.loss);
        info!("Initial accuracy: {:.4}", first.accuracy);
    }

    if let Some(last) = history.last() {
        info!("Final loss: {:.6}", last.loss);
        info!("Final accuracy: {:.4}", last.accuracy);
        info!("Final perplexity: {:.4}", last.loss.exp());
    }

    // Compute best loss
    let best_loss = history.iter().map(|e| e.loss).fold(f64::INFINITY, f64::min);
    info!("Best loss: {:.6}", best_loss);

    // Perplexity analysis
    if history.len() >= 2 {
        let initial_perplexity = history[0].loss.exp();
        let final_perplexity = history.last().unwrap().loss.exp();
        let reduction = ((initial_perplexity - final_perplexity) / initial_perplexity) * 100.0;

        info!("\n=== Perplexity Analysis ===");
        info!("Initial: {:.4}", initial_perplexity);
        info!("Final: {:.4}", final_perplexity);
        info!("Reduction: {:.2}%", reduction);

        if reduction > 10.0 {
            info!("✓ Perplexity decreased significantly");
        } else {
            warn!("⚠ Perplexity reduction < 10%");
        }
    }

    // Validate model performance
    info!("\n=== Model Statistics ===");
    let model_metrics = model.get_performance_metrics();
    for (key, value) in model_metrics.iter() {
        info!("{}: {:.4}", key, value);
    }

    info!("\n=== Training Complete ===");
    Ok(())
}