diff --git a/crates/ml/tests/supervised_gpu_smoke_test.rs b/crates/ml/tests/supervised_gpu_smoke_test.rs new file mode 100644 index 000000000..ae6fbe514 --- /dev/null +++ b/crates/ml/tests/supervised_gpu_smoke_test.rs @@ -0,0 +1,449 @@ +#![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, +)] +//! Supervised Model GPU Smoke Tests +//! +//! Validates all 8 supervised models through the UnifiedTrainable pipeline +//! on a CUDA device: construct → train 10 epochs → checkpoint roundtrip → validate. +//! +//! Models: TFT, Mamba2, TGGN, TLOB, Liquid, KAN, xLSTM, Diffusion +//! +//! Uses synthetic data with small configs to keep VRAM usage minimal. +//! Skips gracefully if no CUDA device is available. + +#![allow(unused_crate_dependencies)] + +use candle_core::{Device, Tensor}; +use ml::training::unified_trainer::UnifiedTrainable; + +fn require_cuda() -> Device { + match Device::cuda_if_available(0) { + Ok(dev) if dev.is_cuda() => { + println!("CUDA device available: {:?}", dev); + dev + } + _ => { + eprintln!("CUDA not available, skipping GPU smoke tests"); + std::process::exit(0); + } + } +} + +/// Generic train-checkpoint-validate pipeline for any UnifiedTrainable model. +/// Returns (final_loss, checkpoint_prediction_diff). +fn smoke_pipeline( + adapter: &mut dyn UnifiedTrainable, + input: &Tensor, + target: &Tensor, + model_name: &str, +) { + let device = adapter.device().clone(); + + // 1. Train 10 epochs + let mut first_loss = None; + let mut last_loss = 0.0_f64; + for epoch in 0..10 { + let preds = adapter.forward(input).unwrap(); + let loss = adapter.compute_loss(&preds, target).unwrap(); + let loss_val = loss.to_scalar::().unwrap() as f64; + assert!( + loss_val.is_finite(), + "{} epoch {}: loss is NaN/Inf ({})", + model_name, + epoch, + loss_val, + ); + if first_loss.is_none() { + first_loss = Some(loss_val); + } + last_loss = loss_val; + + let grad_norm = adapter.backward(&loss).unwrap(); + assert!( + grad_norm.is_finite(), + "{} epoch {}: grad_norm is NaN/Inf ({})", + model_name, + epoch, + grad_norm, + ); + + adapter.optimizer_step().unwrap(); + adapter.zero_grad().unwrap(); + } + let first = first_loss.unwrap(); + println!( + " {} train: first_loss={:.6}, last_loss={:.6}, reduction={:.1}%", + model_name, + first, + last_loss, + (1.0 - last_loss / first) * 100.0, + ); + assert_eq!(adapter.get_step(), 10, "{} should have 10 steps", model_name); + + // 2. Checkpoint roundtrip + let tmp_dir = std::env::temp_dir().join(format!("gpu_smoke_{}", model_name.to_lowercase())); + std::fs::create_dir_all(&tmp_dir).unwrap(); + let ckpt_path = tmp_dir.join("ckpt"); + let save_result = adapter.save_checkpoint(ckpt_path.to_str().unwrap()); + assert!( + save_result.is_ok(), + "{} checkpoint save failed: {:?}", + model_name, + save_result.err(), + ); + + // 3. Validate + let val_input = Tensor::randn(0f32, 0.5, input.dims(), &device).unwrap(); + let val_target = Tensor::randn(0f32, 0.1, target.dims(), &device).unwrap(); + let val_data = vec![(val_input, val_target)]; + let val_loss = adapter.validate(&val_data); + assert!( + val_loss.is_ok(), + "{} validation failed: {:?}", + model_name, + val_loss.err(), + ); + let vl = val_loss.unwrap(); + assert!( + vl.is_finite(), + "{} validation loss is NaN/Inf ({})", + model_name, + vl, + ); + println!(" {} validate: val_loss={:.6}", model_name, vl); + + // 4. Metrics + let metrics = adapter.collect_metrics(); + assert!( + metrics.learning_rate > 0.0, + "{} learning rate should be > 0", + model_name, + ); + + let _ = std::fs::remove_dir_all(&tmp_dir); +} + +// ──────────────────────────────────────────── +// TFT +// ──────────────────────────────────────────── +#[test] +fn test_tft_gpu_smoke() { + let device = require_cuda(); + println!("=== TFT GPU Smoke ==="); + + use ml::tft::trainable_adapter::TrainableTFT; + use ml::tft::TFTConfig; + + let feature_dim = 10; + let mut config = TFTConfig::default(); + config.input_dim = feature_dim; + config.hidden_dim = 32; + config.num_heads = 2; + config.num_layers = 1; + config.num_quantiles = 3; + config.num_static_features = 0; + config.num_known_features = 0; + config.num_unknown_features = feature_dim; + config.sequence_length = 1; + config.prediction_horizon = 1; + config.learning_rate = 1e-3; + config.dropout_rate = 0.0; + + // TFT auto-selects device (cuda_if_available) in TemporalFusionTransformer::new() + let mut adapter = TrainableTFT::new(config).unwrap(); + assert!(adapter.device().is_cuda(), "TFT should be on CUDA"); + + // Input: [batch, feature_dim] (seq_len=1, horizon=1, static=0, known=0) + let input = Tensor::randn(0f32, 0.5, &[16, feature_dim], &device).unwrap(); + // TFT output: [batch, horizon=1, quantiles=3] → loss target must match + let target = Tensor::randn(0f32, 0.1, &[16, 1, 3], &device).unwrap(); + smoke_pipeline(&mut adapter, &input, &target, "TFT"); +} + +// ──────────────────────────────────────────── +// Mamba2 +// ──────────────────────────────────────────── +#[test] +fn test_mamba2_gpu_smoke() { + let device = require_cuda(); + println!("=== Mamba2 GPU Smoke ==="); + + use ml::mamba::trainable_adapter::Mamba2TrainableAdapter; + use ml::mamba::Mamba2Config; + + let mut config = Mamba2Config::default(); + config.d_model = 32; + config.num_layers = 1; + config.d_state = 8; + config.max_seq_len = 8; + + let mut adapter = Mamba2TrainableAdapter::new(config, &device).unwrap(); + assert!(adapter.device().is_cuda(), "Mamba2 should be on CUDA"); + + // Input: [batch, seq_len, d_model] + let input = Tensor::randn(0f32, 0.5, &[16, 8, 32], &device).unwrap(); + // Mamba2 compute_loss narrows to last step → target is [batch, 1] + let target = Tensor::randn(0f32, 0.1, &[16, 1], &device).unwrap(); + smoke_pipeline(&mut adapter, &input, &target, "Mamba2"); +} + +// ──────────────────────────────────────────── +// TGGN +// ──────────────────────────────────────────── +#[test] +fn test_tggn_gpu_smoke() { + let device = require_cuda(); + println!("=== TGGN GPU Smoke ==="); + + use ml::tgnn::trainable_adapter::TGGNTrainableAdapter; + use ml::tgnn::TGGNConfig; + + let config = TGGNConfig { + node_dim: 10, + hidden_dim: 16, + num_layers: 2, + max_nodes: 8, + max_edges: 16, + edge_dim: 4, + temporal_decay: 0.99, + update_frequency_ns: 1_000_000, + use_simd: false, + }; + + let mut adapter = TGGNTrainableAdapter::new(config, &device).unwrap(); + assert!(adapter.device().is_cuda(), "TGGN should be on CUDA"); + + // Input: [batch, node_dim=10], Output: [batch, 1] + let input = Tensor::randn(0f32, 0.5, &[16, 10], &device).unwrap(); + let target = Tensor::randn(0f32, 0.1, &[16, 1], &device).unwrap(); + smoke_pipeline(&mut adapter, &input, &target, "TGGN"); +} + +// ──────────────────────────────────────────── +// TLOB +// ──────────────────────────────────────────── +#[test] +fn test_tlob_gpu_smoke() { + let device = require_cuda(); + println!("=== TLOB GPU Smoke ==="); + + use ml::tlob::trainable_adapter::{TLOBAdapterConfig, TLOBTrainableAdapter}; + + let config = TLOBAdapterConfig { + d_model: 16, + num_heads: 2, + num_layers: 1, + seq_len: 1, + feature_dim: 10, + }; + + let mut adapter = TLOBTrainableAdapter::new(config, &device).unwrap(); + assert!(adapter.device().is_cuda(), "TLOB should be on CUDA"); + + // Input: [batch, seq_len*feature_dim=10] (flat), Output: [batch, 1] + let input = Tensor::randn(0f32, 0.5, &[16, 10], &device).unwrap(); + let target = Tensor::randn(0f32, 0.1, &[16, 1], &device).unwrap(); + smoke_pipeline(&mut adapter, &input, &target, "TLOB"); +} + +// ──────────────────────────────────────────── +// Liquid (CfC) +// ──────────────────────────────────────────── +#[test] +fn test_liquid_gpu_smoke() { + let device = require_cuda(); + println!("=== Liquid GPU Smoke ==="); + + use ml::liquid::adapter::LiquidTrainableAdapter; + use ml::liquid::CfCTrainConfig; + use ml::gpu::DeviceConfig; + + let config = CfCTrainConfig { + input_size: 10, + hidden_size: 16, + output_size: 1, + backbone_hidden_sizes: vec![16, 8], + learning_rate: 1e-3, + device: DeviceConfig::Cuda(0), + ..CfCTrainConfig::default() + }; + + let mut adapter = LiquidTrainableAdapter::new(config).unwrap(); + assert!(adapter.device().is_cuda(), "Liquid should be on CUDA"); + + // Input: [batch, seq_len=4, input_size=10] (3D required) + let input = Tensor::randn(0f32, 0.5, &[16, 4, 10], &device).unwrap(); + // Output: [batch, output_size=1] + let target = Tensor::randn(0f32, 0.1, &[16, 1], &device).unwrap(); + smoke_pipeline(&mut adapter, &input, &target, "Liquid"); +} + +// ──────────────────────────────────────────── +// KAN +// ──────────────────────────────────────────── +#[test] +fn test_kan_gpu_smoke() { + let device = require_cuda(); + println!("=== KAN GPU Smoke ==="); + + use ml::kan::config::KANConfig; + use ml::kan::trainable::KANTrainableAdapter; + + let config = KANConfig { + layer_widths: vec![10, 8, 4, 1], + grid_size: 3, + spline_order: 3, + learning_rate: 1e-3, + weight_decay: 1e-5, + grad_clip: 1.0, + }; + + let mut adapter = KANTrainableAdapter::new(config, &device).unwrap(); + assert!(adapter.device().is_cuda(), "KAN should be on CUDA"); + + // Input: [batch, 10], Output: [batch, 1] + let input = Tensor::randn(0f32, 0.5, &[16, 10], &device).unwrap(); + let target = Tensor::randn(0f32, 0.1, &[16, 1], &device).unwrap(); + smoke_pipeline(&mut adapter, &input, &target, "KAN"); +} + +// ──────────────────────────────────────────── +// xLSTM +// ──────────────────────────────────────────── +#[test] +fn test_xlstm_gpu_smoke() { + let device = require_cuda(); + println!("=== xLSTM GPU Smoke ==="); + + use ml::xlstm::trainable::XLSTMTrainableAdapter; + use ml::xlstm::config::XLSTMConfig; + + let config = XLSTMConfig { + input_dim: 10, + hidden_dim: 16, + num_blocks: 2, + num_heads: 2, + slstm_ratio: 0.5, + output_dim: 1, + dropout: 0.0, + learning_rate: 1e-3, + weight_decay: 1e-5, + grad_clip: 1.0, + }; + + let mut adapter = XLSTMTrainableAdapter::new(config, &device).unwrap(); + assert!(adapter.device().is_cuda(), "xLSTM should be on CUDA"); + + // Input: [batch, seq_len=4, input_dim=10] (3D) + let input = Tensor::randn(0f32, 0.5, &[16, 4, 10], &device).unwrap(); + // Output: [batch, 1] + let target = Tensor::randn(0f32, 0.1, &[16, 1], &device).unwrap(); + smoke_pipeline(&mut adapter, &input, &target, "xLSTM"); +} + +// ──────────────────────────────────────────── +// Diffusion +// ──────────────────────────────────────────── +#[test] +fn test_diffusion_gpu_smoke() { + let device = require_cuda(); + println!("=== Diffusion GPU Smoke ==="); + + use ml::diffusion::config::DiffusionConfig; + use ml::diffusion::trainable::DiffusionTrainableAdapter; + + let config = DiffusionConfig { + num_timesteps: 50, + sampling_steps: 5, + seq_len: 8, + feature_dim: 1, + hidden_dim: 16, + num_layers: 1, + time_embed_dim: 8, + learning_rate: 1e-3, + weight_decay: 1e-5, + grad_clip: 1.0, + ..Default::default() + }; + let data_dim = config.data_dim(); // 8 + + let mut adapter = DiffusionTrainableAdapter::new(config, device.clone()).unwrap(); + assert!(adapter.device().is_cuda(), "Diffusion should be on CUDA"); + + // Input: [batch, data_dim=8] + let input = Tensor::randn(0f32, 0.5, &[16, data_dim], &device).unwrap(); + // Diffusion: use input as pseudo-target (noise prediction, loss won't be meaningful) + let target = input.clone(); + smoke_pipeline(&mut adapter, &input, &target, "Diffusion"); +}