foxhunt/ml/tests/mamba2_adamw_test.rs

//! Test Suite: Mamba-2 AdamW Optimizer Implementation
//!
//! Purpose: Verify that Mamba-2 uses AdamW with decoupled weight decay instead of Adam.
//!
//! Critical Difference:
//! - Adam: weight_decay applied to gradients (interferes with SSM spectral constraints)
//! - AdamW: weight_decay applied directly to parameters (preserves SSM dynamics)
//!
//! Expected Behavior:
//! 1. OptimizerType::AdamW should be available and default
//! 2. Weight decay should be applied as: param = param * (1 - wd), NOT to gradients
//! 3. SSM matrices should maintain spectral radius constraints
//! 4. Better generalization vs Adam (10-20% improvement expected)

use anyhow::Result;
use candle_core::{Device, Tensor};
use ml::mamba::{Mamba2, Mamba2Config, OptimizerType};

/// Test 1: Verify AdamW is available in OptimizerType enum
#[test]
fn test_adamw_optimizer_type_available() -> Result<()> {
    println!("\n=== Test 1: AdamW OptimizerType Available ===");

    // Should compile without error
    let optimizer_type = OptimizerType::AdamW;

    println!("✅ OptimizerType::AdamW is available");
    println!("   Optimizer type: {:?}", optimizer_type);

    Ok(())
}

/// Test 2: Verify AdamW is the default optimizer
#[test]
fn test_adamw_is_default() -> Result<()> {
    println!("\n=== Test 2: AdamW is Default Optimizer ===");

    let config = Mamba2Config::default();

    assert_eq!(
        config.optimizer_type,
        OptimizerType::AdamW,
        "Default optimizer should be AdamW, not Adam"
    );

    println!("✅ Default optimizer is AdamW");
    println!("   Config optimizer_type: {:?}", config.optimizer_type);

    Ok(())
}

/// Test 3: Verify weight decay is decoupled (applied to params, not gradients)
///
/// This test verifies the critical difference between Adam and AdamW:
/// - Adam: grad_with_decay = grad + weight_decay * param (affects gradient flow)
/// - AdamW: param_new = param * (1 - weight_decay) - lr * grad (decoupled)
#[test]
fn test_adamw_decoupled_weight_decay() -> Result<()> {
    println!("\n=== Test 3: AdamW Decoupled Weight Decay ===");

    let device = Device::cuda_if_available(0)?;

    // Create small Mamba-2 model with weight decay
    let config = Mamba2Config {
        d_model: 32,
        d_state: 8,
        d_head: 8,
        num_heads: 1,
        expand: 1,
        num_layers: 1,
        dropout: 0.0,
        use_ssd: false,
        use_selective_state: false,
        hardware_aware: false,
        target_latency_us: 1000,
        max_seq_len: 64,
        learning_rate: 0.001,
        weight_decay: 0.01, // Non-zero weight decay
        grad_clip: 1.0,
        warmup_steps: 0,
        adam_beta1: 0.9,
        adam_beta2: 0.999,
        adam_epsilon: 1e-8,
        total_decay_steps: 1000,
        optimizer_type: OptimizerType::AdamW, // Use AdamW
        sgd_momentum: 0.9,
        batch_size: 2,
        seq_len: 16,
        shuffle_batches: false,
        sequence_stride: 1,
        norm_eps: 1e-5,
    };

    let mut model = Mamba2::new(config.clone(), &device)?;

    // Get initial parameter value (B matrix, layer 0)
    let initial_b = model.state.ssm_states[0].B.clone();
    let initial_b_data = initial_b.to_vec1::<f32>()?;
    println!(
        "Initial B matrix norm: {:.6}",
        initial_b_data.iter().map(|x| x * x).sum::<f32>().sqrt()
    );

    // Create synthetic input and target
    let batch_size = 2;
    let seq_len = 16;
    let input = Tensor::randn(0f32, 1.0, (batch_size, seq_len, config.d_model), &device)?;
    let target = Tensor::randn(0f32, 1.0, (batch_size, seq_len, config.d_model), &device)?;

    // Forward pass
    let output = model.forward(&input)?;

    // Compute loss (MSE)
    let diff = output.sub(&target)?;
    let loss = diff.sqr()?.mean_all()?;
    println!("Initial loss: {:.6}", loss.to_vec0::<f32>()?);

    // Backward pass
    model.backward_pass(&input, &target)?;

    // Get gradient before optimizer step
    let b_grad = model
        .gradients
        .get("B_0")
        .cloned()
        .ok_or_else(|| anyhow::anyhow!("Missing B gradient"))?;
    let b_grad_norm = b_grad
        .to_vec1::<f32>()?
        .iter()
        .map(|x| x * x)
        .sum::<f32>()
        .sqrt();
    println!("B gradient norm: {:.6}", b_grad_norm);

    // Optimizer step
    model.optimizer_step()?;

    // Get updated parameter
    let updated_b = model.state.ssm_states[0].B.clone();
    let updated_b_data = updated_b.to_vec1::<f32>()?;
    let updated_b_norm = updated_b_data.iter().map(|x| x * x).sum::<f32>().sqrt();
    println!("Updated B matrix norm: {:.6}", updated_b_norm);

    // CRITICAL CHECK: Weight decay should directly reduce parameter norm
    // Expected: updated_norm ≈ initial_norm * (1 - weight_decay) - gradient_contribution
    // This verifies decoupled weight decay (AdamW) vs coupled (Adam)

    let decay_factor = 1.0 - config.weight_decay as f32;
    let expected_decay_effect =
        initial_b_data.iter().map(|x| x * x).sum::<f32>().sqrt() * decay_factor;

    // The updated norm should be less than initial due to weight decay
    // (even without gradient updates, weight decay alone reduces norm)
    assert!(
        updated_b_norm < initial_b_data.iter().map(|x| x * x).sum::<f32>().sqrt(),
        "Weight decay should reduce parameter norm"
    );

    println!("\n✅ Weight decay is decoupled:");
    println!(
        "   Initial norm: {:.6}",
        initial_b_data.iter().map(|x| x * x).sum::<f32>().sqrt()
    );
    println!("   Expected decay effect: {:.6}", expected_decay_effect);
    println!("   Updated norm: {:.6}", updated_b_norm);
    println!(
        "   Norm reduction: {:.2}%",
        (1.0 - updated_b_norm / initial_b_data.iter().map(|x| x * x).sum::<f32>().sqrt()) * 100.0
    );

    Ok(())
}

/// Test 4: Verify AdamW preserves SSM spectral radius constraints
///
/// SSM A matrices must maintain spectral radius < 1 for stability.
/// AdamW's decoupled weight decay should NOT interfere with this constraint.
#[test]
fn test_adamw_preserves_spectral_radius() -> Result<()> {
    println!("\n=== Test 4: AdamW Preserves SSM Spectral Radius ===");

    let device = Device::cuda_if_available(0)?;

    let config = Mamba2Config {
        d_model: 32,
        d_state: 8,
        d_head: 8,
        num_heads: 1,
        expand: 1,
        num_layers: 2, // Multiple layers
        dropout: 0.0,
        use_ssd: false,
        use_selective_state: false,
        hardware_aware: false,
        target_latency_us: 1000,
        max_seq_len: 64,
        learning_rate: 0.001,
        weight_decay: 0.01,
        grad_clip: 1.0,
        warmup_steps: 0,
        adam_beta1: 0.9,
        adam_beta2: 0.999,
        adam_epsilon: 1e-8,
        total_decay_steps: 1000,
        optimizer_type: OptimizerType::AdamW,
        sgd_momentum: 0.9,
        batch_size: 2,
        seq_len: 16,
        shuffle_batches: false,
        sequence_stride: 1,
        norm_eps: 1e-5,
    };

    let mut model = Mamba2::new(config.clone(), &device)?;

    // Run 10 training steps
    for step in 0..10 {
        let input = Tensor::randn(0f32, 1.0, (2, 16, config.d_model), &device)?;
        let target = Tensor::randn(0f32, 1.0, (2, 16, config.d_model), &device)?;

        let output = model.forward(&input)?;
        let diff = output.sub(&target)?;
        let loss = diff.sqr()?.mean_all()?;

        model.backward_pass(&input, &target)?;
        model.optimizer_step()?;

        // Check spectral radius of A matrices
        for (layer_idx, ssm_state) in model.state.ssm_states.iter().enumerate() {
            let a_data = ssm_state.A.to_vec1::<f32>()?;
            let max_eigenvalue = a_data.iter().map(|x| x.abs()).fold(0.0f32, f32::max);

            // Spectral radius should be < 1 for stability
            assert!(
                max_eigenvalue < 1.0,
                "Layer {} A matrix spectral radius ({:.6}) exceeds 1.0 after step {}",
                layer_idx,
                max_eigenvalue,
                step
            );

            if step == 9 {
                println!(
                    "   Layer {} A spectral radius: {:.6} ✓",
                    layer_idx, max_eigenvalue
                );
            }
        }

        if step == 0 || step == 9 {
            println!("Step {}: loss = {:.6}", step, loss.to_vec0::<f32>()?);
        }
    }

    println!("\n✅ AdamW preserves SSM spectral radius constraints");
    println!("   All layers maintain spectral radius < 1.0 after 10 training steps");

    Ok(())
}

/// Test 5: Compare Adam vs AdamW convergence (optional, informational)
///
/// This test demonstrates that AdamW should converge better than Adam for SSMs.
/// Expected: AdamW achieves lower final loss and better generalization.
#[test]
#[ignore] // Expensive test, run with --ignored
fn test_adamw_vs_adam_convergence() -> Result<()> {
    println!("\n=== Test 5: AdamW vs Adam Convergence Comparison ===");

    let device = Device::cuda_if_available(0)?;

    let base_config = Mamba2Config {
        d_model: 32,
        d_state: 8,
        d_head: 8,
        num_heads: 1,
        expand: 1,
        num_layers: 2,
        dropout: 0.0,
        use_ssd: false,
        use_selective_state: false,
        hardware_aware: false,
        target_latency_us: 1000,
        max_seq_len: 64,
        learning_rate: 0.001,
        weight_decay: 0.01,
        grad_clip: 1.0,
        warmup_steps: 0,
        adam_beta1: 0.9,
        adam_beta2: 0.999,
        adam_epsilon: 1e-8,
        total_decay_steps: 1000,
        optimizer_type: OptimizerType::AdamW, // Will override
        sgd_momentum: 0.9,
        batch_size: 4,
        seq_len: 32,
        shuffle_batches: false,
        sequence_stride: 1,
        norm_eps: 1e-5,
    };

    // Train with Adam
    let mut adam_config = base_config.clone();
    adam_config.optimizer_type = OptimizerType::Adam;
    let mut adam_model = Mamba2::new(adam_config.clone(), &device)?;

    // Train with AdamW
    let mut adamw_config = base_config.clone();
    adamw_config.optimizer_type = OptimizerType::AdamW;
    let mut adamw_model = Mamba2::new(adamw_config.clone(), &device)?;

    let num_steps = 100;
    let mut adam_losses = Vec::new();
    let mut adamw_losses = Vec::new();

    // Fixed dataset for fair comparison
    let mut train_data = Vec::new();
    for _ in 0..num_steps {
        let input = Tensor::randn(0f32, 1.0, (4, 32, base_config.d_model), &device)?;
        let target = Tensor::randn(0f32, 1.0, (4, 32, base_config.d_model), &device)?;
        train_data.push((input, target));
    }

    // Train Adam model
    println!("\nTraining with Adam optimizer...");
    for (step, (input, target)) in train_data.iter().enumerate() {
        let output = adam_model.forward(input)?;
        let diff = output.sub(target)?;
        let loss = diff.sqr()?.mean_all()?;
        adam_losses.push(loss.to_vec0::<f32>()?);

        adam_model.backward_pass(input, target)?;
        adam_model.optimizer_step()?;

        if step % 20 == 0 || step == num_steps - 1 {
            println!("  Step {}: loss = {:.6}", step, adam_losses[step]);
        }
    }

    // Train AdamW model
    println!("\nTraining with AdamW optimizer...");
    for (step, (input, target)) in train_data.iter().enumerate() {
        let output = adamw_model.forward(input)?;
        let diff = output.sub(target)?;
        let loss = diff.sqr()?.mean_all()?;
        adamw_losses.push(loss.to_vec0::<f32>()?);

        adamw_model.backward_pass(input, target)?;
        adamw_model.optimizer_step()?;

        if step % 20 == 0 || step == num_steps - 1 {
            println!("  Step {}: loss = {:.6}", step, adamw_losses[step]);
        }
    }

    // Compare final losses
    let adam_final = adam_losses[num_steps - 1];
    let adamw_final = adamw_losses[num_steps - 1];
    let improvement = ((adam_final - adamw_final) / adam_final) * 100.0;

    println!("\n=== Convergence Comparison ===");
    println!("Adam final loss:  {:.6}", adam_final);
    println!("AdamW final loss: {:.6}", adamw_final);
    println!("Improvement:      {:.2}%", improvement);

    // AdamW should achieve better or equal final loss
    assert!(
        adamw_final <= adam_final * 1.1,
        "AdamW should achieve comparable or better convergence than Adam"
    );

    println!("\n✅ AdamW convergence test passed");
    if improvement > 0.0 {
        println!("   AdamW converged {:.2}% better than Adam", improvement);
    } else {
        println!("   AdamW and Adam achieved similar convergence (expected for short training)");
    }

    Ok(())
}