foxhunt/crates/ml/tests/mamba2_weight_update_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 Weight Update Integration Test
//!
//! Verifies that the gradient fix enables actual weight updates during training.
//!
//! **Test Strategy**:
//! 1. Record initial SSM parameters (A, B, C, delta matrices)
//! 2. Run one training step (forward + backward + optimizer_step)
//! 3. Verify parameters changed
//! 4. Verify loss is finite (model is computing gradients)
//!
//! NOTE: ml-supervised uses StreamTensor (re-exported as GpuTensor inside the crate).
//! The readback method is `.to_vec()` (no stream arg — uses internal stream).

#![allow(unused_crate_dependencies)]

use std::sync::Arc;

use ml::mamba::Mamba2SSM;
use ml::MLError;
use ml_core::cuda_autograd::stream_ops::{StreamTensor, gpu_sub, gpu_sqr};
use tracing::info;

#[test]
fn test_mamba2_weights_update_after_training_step() -> Result<(), MLError> {
    info!("MAMBA-2 Weight Update Integration Test");

    let ctx = cudarc::driver::CudaContext::new(0)
        .map_err(|e| MLError::DeviceError(format!("CUDA context: {e}")))?;
    let stream = ctx
        .new_stream()
        .map_err(|e| MLError::DeviceError(format!("CUDA stream: {e}")))?;
    info!("Device: CUDA:0");

    let mut model = Mamba2SSM::default_hft(&stream)?;
    info!(num_parameters = model.metadata.num_parameters, "Model created");

    // Create dummy data — Mamba2SSM::forward takes StreamTensor [batch, d_model]
    let batch_size = 8;
    let d_model = model.config.d_model;

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

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

    info!("Recording Initial SSM Parameters");

    // Capture initial SSM state (A, B, C, delta matrices) on host via to_vec()
    let mut initial_params: Vec<Vec<f32>> = Vec::new();
    let num_layers = model.state.ssm_states.len();
    for layer_idx in 0..num_layers {
        let state = &model.state.ssm_states[layer_idx];
        let a_host = state.A.to_vec()?;
        let b_host = state.B.to_vec()?;
        let c_host = state.C.to_vec()?;
        let delta_host = state.delta.to_vec()?;
        initial_params.push(a_host);
        initial_params.push(b_host);
        initial_params.push(c_host);
        initial_params.push(delta_host);
        info!(
            layer = layer_idx,
            params = initial_params.len(),
            "Initial SSM state recorded"
        );
    }

    info!("Running Training Step 1");

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

    // Compute MSE loss on host (cold path for test)
    let output_host = output.to_vec()?;
    let target_host = target.to_vec()?;
    let loss_before: f32 = output_host
        .iter()
        .zip(target_host.iter())
        .map(|(o, t)| (o - t).powi(2))
        .sum::<f32>()
        / output_host.len() as f32;
    info!(loss_before, "Loss before training");

    // Create loss tensor for backward pass — gpu_sub and gpu_sqr are free functions on StreamTensor
    let loss_tensor = gpu_sub(&output, &target)?;
    let loss_sq = gpu_sqr(&loss_tensor)?;

    // Backward pass (pseudo-gradient computation for SSM parameters)
    model.backward_pass(&loss_sq, &input, &target)?;

    // Optimizer step (update SSM parameters)
    model.optimizer_step()?;

    info!("Running Training Step 2");

    // Forward pass again
    let output2 = model.forward(&input)?;
    let output2_host = output2.to_vec()?;
    let loss_after: f32 = output2_host
        .iter()
        .zip(target_host.iter())
        .map(|(o, t)| (o - t).powi(2))
        .sum::<f32>()
        / output2_host.len() as f32;
    info!(loss_after, "Loss after training");

    info!("Verifying Parameter Updates");

    // Capture updated SSM parameters and compare
    let mut params_changed = 0;
    let mut total_delta = 0.0_f64;

    let mut param_idx = 0;
    for layer_idx in 0..num_layers {
        let state = &model.state.ssm_states[layer_idx];
        let matrices = [
            state.A.to_vec()?,
            state.B.to_vec()?,
            state.C.to_vec()?,
            state.delta.to_vec()?,
        ];

        for matrix_host in &matrices {
            let initial = &initial_params[param_idx];
            let delta_norm: f64 = matrix_host
                .iter()
                .zip(initial.iter())
                .map(|(w_new, w_old)| (*w_new as f64 - *w_old as f64).powi(2))
                .sum::<f64>()
                .sqrt();

            if delta_norm > 1e-9 {
                params_changed += 1;
                total_delta += delta_norm;
                info!(layer = layer_idx, param_idx, delta_norm, "Param changed");
            } else {
                info!(layer = layer_idx, param_idx, delta_norm, "Param unchanged");
            }
            param_idx += 1;
        }
    }

    info!(
        params_changed,
        total = initial_params.len(),
        "Parameters changed"
    );
    info!(total_delta, "Total parameter delta norm");
    info!(
        loss_before,
        loss_after,
        loss_delta = loss_after - loss_before,
        "Loss change"
    );

    // ASSERTIONS
    assert!(
        params_changed > 0,
        "FAIL: No SSM parameters changed after training step! optimizer_step() may be broken."
    );

    assert!(
        total_delta > 1e-6,
        "FAIL: Total parameter delta too small ({:.6}). Parameters barely changed.",
        total_delta
    );

    // Note: Loss might increase in first step due to random initialization
    // but parameters MUST change if gradients are flowing
    info!("TEST PASSED: SSM parameters updated after training step");
    info!("Gradient fix enables learning");

    Ok(())
}