gitadmin/foxhunt
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

test: add supervised_gpu_smoke_test for all 8 models (UnifiedTrainable on CUDA)

Browse Source 
Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-03-13 11:55:45 +01:00
parent 1520c56e41
commit a3058a68de
1 changed files with 449 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

449
crates/ml/tests/supervised_gpu_smoke_test.rs Normal file
View File

@@ -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::<f32>().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");
}