foxhunt/crates/ml/tests/mamba2_gradient_extraction_test.rs

#![allow(
    clippy::assertions_on_constants,
    clippy::assertions_on_result_states,
    clippy::clone_on_copy,
    clippy::decimal_literal_representation,
    clippy::doc_markdown,
    clippy::empty_line_after_doc_comments,
    clippy::field_reassign_with_default,
    clippy::get_unwrap,
    clippy::identity_op,
    clippy::inconsistent_digit_grouping,
    clippy::indexing_slicing,
    clippy::integer_division,
    clippy::len_zero,
    clippy::let_underscore_must_use,
    clippy::manual_div_ceil,
    clippy::manual_let_else,
    clippy::manual_range_contains,
    clippy::modulo_arithmetic,
    clippy::needless_range_loop,
    clippy::non_ascii_literal,
    clippy::redundant_clone,
    clippy::shadow_reuse,
    clippy::shadow_same,
    clippy::shadow_unrelated,
    clippy::single_match_else,
    clippy::str_to_string,
    clippy::string_slice,
    clippy::tests_outside_test_module,
    clippy::too_many_lines,
    clippy::unnecessary_wraps,
    clippy::unseparated_literal_suffix,
    clippy::use_debug,
    clippy::useless_vec,
    clippy::wildcard_enum_match_arm,
    clippy::else_if_without_else,
    clippy::expect_used,
    clippy::missing_const_for_fn,
    clippy::similar_names,
    clippy::type_complexity,
    clippy::collapsible_else_if,
    clippy::doc_lazy_continuation,
    clippy::items_after_test_module,
    clippy::map_clone,
    clippy::multiple_unsafe_ops_per_block,
    clippy::unwrap_or_default,
    clippy::assign_op_pattern,
    clippy::needless_borrow,
    clippy::println_empty_string,
    clippy::unnecessary_cast,
    clippy::used_underscore_binding,
    clippy::create_dir,
    clippy::implicit_saturating_sub,
    clippy::exit,
    clippy::expect_fun_call,
    clippy::too_many_arguments,
    clippy::unnecessary_map_or,
    clippy::unwrap_used,
    dead_code,
    unused_imports,
    unused_variables,
    clippy::cloned_ref_to_slice_refs,
    clippy::neg_multiply,
    clippy::while_let_loop,
    clippy::bool_assert_comparison,
    clippy::excessive_precision,
    clippy::trivially_copy_pass_by_ref,
    clippy::op_ref,
    clippy::redundant_closure,
    clippy::unnecessary_lazy_evaluations,
    clippy::if_then_some_else_none,
    clippy::unnecessary_to_owned,
    clippy::single_component_path_imports,
)]
//! # MAMBA-2 Gradient Extraction Test (TDD)
//!
//! **Test-Driven Development**: This test verifies that gradients are properly computed
//! after backward_pass() using finite-difference pseudo-gradients (StreamTensor has no autograd).
//!
//! **Expected Behavior**:
//! 1. Call forward() to compute output
//! 2. Compute loss via host-side MSE
//! 3. Call backward_pass() to compute pseudo-gradients
//! 4. Verify gradients are non-zero and valid (not NaN/Inf)

#![allow(unused_crate_dependencies)]

// candle eliminated — test uses native cudarc StreamTensor APIs
use std::sync::OnceLock;
use ml::mamba::Mamba2SSM;
use ml::MLError;
use ml_core::cuda_autograd::stream_ops::{StreamTensor, gpu_sub, gpu_sqr, gpu_mean_all};
use ml_core::device::MlDevice;
use tracing::info;

static SHARED_CUDA: OnceLock<MlDevice> = OnceLock::new();

fn cuda_device() -> MlDevice {
    SHARED_CUDA.get_or_init(|| MlDevice::cuda(0).expect("CUDA required")).clone()
}

#[test]
fn test_mamba2_gradient_extraction_from_backward_pass() -> Result<(), MLError> {
    info!("=== MAMBA-2 Gradient Extraction Test ===");

    let device = cuda_device();
    let stream = device.cuda_stream()?;
    info!("Device: CUDA");

    // Create small MAMBA-2 model
    let mut model = Mamba2SSM::default_hft(stream)?;
    info!(num_parameters = model.metadata.num_parameters, "Model created");

    // Create dummy input and target
    let batch_size = model.config.batch_size;
    let seq_len = model.config.seq_len;
    let d_model = model.config.d_model;

    let input_data = vec![0.1_f32; batch_size * seq_len * d_model];
    let input = StreamTensor::from_vec(input_data, &[batch_size, seq_len, d_model], stream)?;

    let target_data = vec![0.5_f32; batch_size * seq_len];
    let target = StreamTensor::from_vec(target_data, &[batch_size, seq_len, 1], stream)?;

    info!(input_shape = ?input.shape, "Input shape");
    info!(target_shape = ?target.shape, "Target shape");

    // Forward pass
    let output = model.forward(&input)?;
    info!(output_shape = ?output.shape, "Output shape");

    // Compute loss (MSE) via StreamTensor free functions
    let diff = gpu_sub(&output, &target)?;
    let sq = gpu_sqr(&diff)?;
    let loss_value = gpu_mean_all(&sq)?;
    info!(loss_value, "Loss");

    // Create a scalar loss StreamTensor for backward_pass
    let loss_tensor = StreamTensor::from_vec(vec![loss_value as f32], &[1], stream)?;

    // backward_pass computes finite-difference pseudo-gradients for SSM parameters
    model.backward_pass(&loss_tensor, &input, &target)?;

    // Extract gradients from model.gradients HashMap
    info!("=== Extracting Gradients from model.gradients ===");

    let total_entries = model.gradients.len();
    info!(total_entries, "Total gradient entries");

    let mut vars_with_gradients = 0;
    let mut total_grad_norm = 0.0_f64;

    for (key, grad) in model.gradients.iter() {
        // Download gradient to host
        let grad_vec = grad.to_vec()?;
        let grad_norm: f64 = grad_vec.iter().map(|&g| (g as f64).powi(2)).sum::<f64>().sqrt();

        info!(key, grad_norm, "Gradient entry norm");

        // Verify gradient is valid
        assert!(!grad_norm.is_nan(), "Gradient {} is NaN", key);
        assert!(!grad_norm.is_infinite(), "Gradient {} is Inf", key);

        if grad_norm > 1e-9 {
            vars_with_gradients += 1;
            total_grad_norm += grad_norm;
        }
    }

    info!(
        vars_with_gradients,
        total_entries,
        total_grad_norm,
        "Gradient summary"
    );

    // CRITICAL ASSERTION: At least some parameters should have non-zero gradients
    assert!(
        vars_with_gradients > 0,
        "FAIL: No variables have gradients! backward_pass() did not compute gradients."
    );

    assert!(
        total_grad_norm > 1e-6,
        "FAIL: Total gradient norm is too small ({:.6}). Gradients may be zeros.",
        total_grad_norm
    );

    info!("TEST PASSED: Gradients extracted from model.gradients");
    Ok(())
}

#[test]
fn test_mamba2_backward_pass_extracts_real_gradients() -> Result<(), MLError> {
    info!("=== MAMBA-2 backward_pass() Real Gradient Test ===");

    let device = cuda_device();
    let stream = device.cuda_stream()?;
    let mut model = Mamba2SSM::default_hft(stream)?;

    // Create input/target (StreamTensor is always f32)
    let batch_size = model.config.batch_size;
    let seq_len = model.config.seq_len;
    let d_model = model.config.d_model;

    let input_data = vec![1.0_f32; batch_size * seq_len * d_model];
    let input = StreamTensor::from_vec(input_data, &[batch_size, seq_len, d_model], stream)?;

    let target_data = vec![1.0_f32; batch_size * seq_len];
    let target = StreamTensor::from_vec(target_data, &[batch_size, seq_len, 1], stream)?;

    // Forward + loss
    let output = model.forward(&input)?;
    let diff = gpu_sub(&output, &target)?;
    let sq = gpu_sqr(&diff)?;
    let loss_value = gpu_mean_all(&sq)?;

    info!(loss = loss_value, "Loss");

    // Create a scalar loss StreamTensor for backward_pass
    let loss_tensor = StreamTensor::from_vec(vec![loss_value as f32], &[1], stream)?;

    // Call backward_pass (computes finite-difference pseudo-gradients)
    model.backward_pass(&loss_tensor, &input, &target)?;

    // Check model.gradients HashMap
    info!(total_entries = model.gradients.len(), "=== Model Gradients HashMap ===");

    for (key, grad) in model.gradients.iter() {
        let grad_vec = grad.to_vec()?;
        let grad_norm: f64 = grad_vec.iter().map(|&g| (g as f64).powi(2)).sum::<f64>().sqrt();
        let is_nonzero = grad_norm > 1e-9;
        info!(key, grad_norm, is_nonzero, "Gradient entry");
    }

    // After backward_pass, gradients should be non-zero
    let total_grad_norm: f64 = model
        .gradients
        .values()
        .map(|grad| {
            let grad_vec = grad.to_vec().unwrap();
            grad_vec.iter().map(|&g| (g as f64).powi(2)).sum::<f64>().sqrt()
        })
        .sum();

    info!(total_grad_norm, "Total gradient norm in model.gradients");

    assert!(
        total_grad_norm > 1e-6,
        "FAIL: backward_pass() produced zero gradients. Need non-zero pseudo-gradients."
    );

    info!("TEST PASSED: backward_pass() extracts real gradients");
    Ok(())
}