From 3db7f4828bd871bbe47dfcff7aee544b34f49bbd Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Sun, 8 Mar 2026 09:11:16 +0100 Subject: [PATCH] refactor(ml): extract 8 supervised models into ml-supervised crate (task 8) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Move TFT, Mamba-2, Liquid, TGGN, TLOB, KAN, xLSTM, and Diffusion model implementations to ml-supervised. Bridge files (UnifiedTrainable adapters, Checkpointable impls) stay in ml. Delete AsyncDataLoader (replaced by StreamingDbnLoader + simple .chunks() batching). Remove empty ml-infra scaffold — the remaining ml modules are too tightly coupled for clean extraction, so ml stays as the orchestration facade. - ml-supervised: 234 tests, 0 failures - ml: 1687 tests, 0 failures - Workspace: 0 compilation errors Co-Authored-By: Claude Opus 4.6 --- Cargo.lock | 36 +- Cargo.toml | 2 - crates/ml-infra/Cargo.toml | 23 - crates/ml-infra/src/lib.rs | 1 - crates/ml-supervised/Cargo.toml | 50 + .../src/diffusion/config.rs | 0 .../src/diffusion/denoiser.rs | 8 +- crates/ml-supervised/src/diffusion/mod.rs | 17 + .../src/diffusion/noise.rs | 2 +- .../src/diffusion/sampler.rs | 4 +- .../{ml => ml-supervised}/src/kan/config.rs | 0 crates/{ml => ml-supervised}/src/kan/layer.rs | 4 +- crates/ml-supervised/src/kan/mod.rs | 13 + .../{ml => ml-supervised}/src/kan/network.rs | 6 +- .../{ml => ml-supervised}/src/kan/spline.rs | 2 +- crates/ml-supervised/src/lib.rs | 20 +- .../src/liquid/activation.rs | 0 .../src/liquid/candle_cfc.rs | 14 +- .../{ml => ml-supervised}/src/liquid/cells.rs | 0 crates/ml-supervised/src/liquid/mod.rs | 210 + .../src/liquid/network.rs | 2 +- .../src/liquid/ode_solvers.rs | 0 .../{ml => ml-supervised}/src/liquid/tests.rs | 0 .../src/liquid/training.rs | 4 +- .../src/mamba/cuda/selective_scan.cu | 0 .../src/mamba/hardware_aware.rs | 4 +- .../{ml => ml-supervised}/src/mamba/loss.rs | 2 +- crates/ml-supervised/src/mamba/mod.rs | 3782 +++++++++++++++++ .../src/mamba/scan_algorithms.rs | 2 +- .../src/mamba/selective_state.rs | 2 +- .../src/mamba/ssd_layer.rs | 4 +- .../src/tft/gated_residual.rs | 4 +- .../src/tft/hft_optimizations.rs | 2 +- .../src/tft/lstm_encoder.rs | 4 +- crates/ml-supervised/src/tft/mod.rs | 1410 ++++++ .../{ml => ml-supervised}/src/tft/qat_tft.rs | 2 +- .../src/tft/quantile_outputs.rs | 2 +- .../src/tft/quantized_attention.rs | 6 +- .../src/tft/quantized_grn.rs | 12 +- .../src/tft/quantized_lstm.rs | 12 +- .../src/tft/quantized_tft.rs | 6 +- .../src/tft/quantized_vsn.rs | 10 +- .../src/tft/temporal_attention.rs | 4 +- .../src/tft/variable_selection.rs | 2 +- .../src/tft/varmap_quantization.rs | 10 +- .../{ml => ml-supervised}/src/tgnn/gating.rs | 2 +- .../{ml => ml-supervised}/src/tgnn/graph.rs | 2 +- .../src/tgnn/message_passing.rs | 2 +- crates/ml-supervised/src/tgnn/mod.rs | 1310 ++++++ .../{ml => ml-supervised}/src/tgnn/traits.rs | 2 +- .../{ml => ml-supervised}/src/tgnn/types.rs | 0 .../src/tlob/analytics.rs | 0 .../src/tlob/features.rs | 2 +- .../src/tlob/mbp10_feature_extractor.rs | 2 +- crates/ml-supervised/src/tlob/mod.rs | 23 + .../src/tlob/performance.rs | 0 .../src/tlob/transformer.rs | 4 +- .../{ml => ml-supervised}/src/xlstm/block.rs | 4 +- .../{ml => ml-supervised}/src/xlstm/config.rs | 0 .../{ml => ml-supervised}/src/xlstm/mlstm.rs | 6 +- crates/ml-supervised/src/xlstm/mod.rs | 17 + .../src/xlstm/network.rs | 8 +- .../{ml => ml-supervised}/src/xlstm/slstm.rs | 6 +- crates/ml/Cargo.toml | 1 + .../src/checkpoint/model_implementations.rs | 711 ---- crates/ml/src/diffusion/mod.rs | 22 +- .../hyperopt/adapters/async_data_loader.rs | 646 --- crates/ml/src/hyperopt/adapters/mod.rs | 2 - crates/ml/src/kan/mod.rs | 18 +- crates/ml/src/liquid/cuda/liquid_kernels.cu | 513 --- crates/ml/src/liquid/cuda/memory.rs | 319 -- crates/ml/src/liquid/cuda/mod.rs | 640 --- crates/ml/src/liquid/mod.rs | 211 +- crates/ml/src/mamba/mod.rs | 3196 +------------- crates/ml/src/mamba/trainable_adapter.rs | 19 - crates/ml/src/tft/mod.rs | 1525 +------ crates/ml/src/tft/quantized_tft_forward.rs | 311 -- crates/ml/src/tgnn/mod.rs | 1204 +----- crates/ml/src/tlob/mod.rs | 26 +- crates/ml/src/xlstm/mod.rs | 22 +- .../ml/tests/async_data_loading_benchmark.rs | 265 -- 81 files changed, 7023 insertions(+), 9718 deletions(-) delete mode 100644 crates/ml-infra/Cargo.toml delete mode 100644 crates/ml-infra/src/lib.rs rename crates/{ml => ml-supervised}/src/diffusion/config.rs (100%) rename crates/{ml => ml-supervised}/src/diffusion/denoiser.rs (98%) create mode 100644 crates/ml-supervised/src/diffusion/mod.rs rename crates/{ml => ml-supervised}/src/diffusion/noise.rs (99%) rename crates/{ml => ml-supervised}/src/diffusion/sampler.rs (99%) rename crates/{ml => ml-supervised}/src/kan/config.rs (100%) rename crates/{ml => ml-supervised}/src/kan/layer.rs (98%) create mode 100644 crates/ml-supervised/src/kan/mod.rs rename crates/{ml => ml-supervised}/src/kan/network.rs (96%) rename crates/{ml => ml-supervised}/src/kan/spline.rs (99%) rename crates/{ml => ml-supervised}/src/liquid/activation.rs (100%) rename crates/{ml => ml-supervised}/src/liquid/candle_cfc.rs (98%) rename crates/{ml => ml-supervised}/src/liquid/cells.rs (100%) create mode 100644 crates/ml-supervised/src/liquid/mod.rs rename crates/{ml => ml-supervised}/src/liquid/network.rs (99%) rename crates/{ml => ml-supervised}/src/liquid/ode_solvers.rs (100%) rename crates/{ml => ml-supervised}/src/liquid/tests.rs (100%) rename crates/{ml => ml-supervised}/src/liquid/training.rs (99%) rename crates/{ml => ml-supervised}/src/mamba/cuda/selective_scan.cu (100%) rename crates/{ml => ml-supervised}/src/mamba/hardware_aware.rs (99%) rename crates/{ml => ml-supervised}/src/mamba/loss.rs (99%) create mode 100644 crates/ml-supervised/src/mamba/mod.rs rename crates/{ml => ml-supervised}/src/mamba/scan_algorithms.rs (99%) rename crates/{ml => ml-supervised}/src/mamba/selective_state.rs (99%) rename crates/{ml => ml-supervised}/src/mamba/ssd_layer.rs (99%) rename crates/{ml => ml-supervised}/src/tft/gated_residual.rs (99%) rename crates/{ml => ml-supervised}/src/tft/hft_optimizations.rs (99%) rename crates/{ml => ml-supervised}/src/tft/lstm_encoder.rs (99%) create mode 100644 crates/ml-supervised/src/tft/mod.rs rename crates/{ml => ml-supervised}/src/tft/qat_tft.rs (99%) rename crates/{ml => ml-supervised}/src/tft/quantile_outputs.rs (99%) rename crates/{ml => ml-supervised}/src/tft/quantized_attention.rs (99%) rename crates/{ml => ml-supervised}/src/tft/quantized_grn.rs (97%) rename crates/{ml => ml-supervised}/src/tft/quantized_lstm.rs (98%) rename crates/{ml => ml-supervised}/src/tft/quantized_tft.rs (99%) rename crates/{ml => ml-supervised}/src/tft/quantized_vsn.rs (97%) rename crates/{ml => ml-supervised}/src/tft/temporal_attention.rs (99%) rename crates/{ml => ml-supervised}/src/tft/variable_selection.rs (99%) rename crates/{ml => ml-supervised}/src/tft/varmap_quantization.rs (99%) rename crates/{ml => ml-supervised}/src/tgnn/gating.rs (99%) rename crates/{ml => ml-supervised}/src/tgnn/graph.rs (99%) rename crates/{ml => ml-supervised}/src/tgnn/message_passing.rs (99%) create mode 100644 crates/ml-supervised/src/tgnn/mod.rs rename crates/{ml => ml-supervised}/src/tgnn/traits.rs (97%) rename crates/{ml => ml-supervised}/src/tgnn/types.rs (100%) rename crates/{ml => ml-supervised}/src/tlob/analytics.rs (100%) rename crates/{ml => ml-supervised}/src/tlob/features.rs (99%) rename crates/{ml => ml-supervised}/src/tlob/mbp10_feature_extractor.rs (99%) create mode 100644 crates/ml-supervised/src/tlob/mod.rs rename crates/{ml => ml-supervised}/src/tlob/performance.rs (100%) rename crates/{ml => ml-supervised}/src/tlob/transformer.rs (99%) rename crates/{ml => ml-supervised}/src/xlstm/block.rs (98%) rename crates/{ml => ml-supervised}/src/xlstm/config.rs (100%) rename crates/{ml => ml-supervised}/src/xlstm/mlstm.rs (98%) create mode 100644 crates/ml-supervised/src/xlstm/mod.rs rename crates/{ml => ml-supervised}/src/xlstm/network.rs (97%) rename crates/{ml => ml-supervised}/src/xlstm/slstm.rs (98%) delete mode 100644 crates/ml/src/hyperopt/adapters/async_data_loader.rs delete mode 100644 crates/ml/src/liquid/cuda/liquid_kernels.cu delete mode 100644 crates/ml/src/liquid/cuda/memory.rs delete mode 100644 crates/ml/src/liquid/cuda/mod.rs delete mode 100644 crates/ml/src/tft/quantized_tft_forward.rs delete mode 100644 crates/ml/tests/async_data_loading_benchmark.rs diff --git a/Cargo.lock b/Cargo.lock index 565003ac0..98f703a24 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -6265,6 +6265,7 @@ dependencies = [ "ml-core", "ml-dqn", "ml-ppo", + "ml-supervised", "nalgebra 0.33.2", "ndarray", "num", @@ -6402,16 +6403,6 @@ dependencies = [ "tracing", ] -[[package]] -name = "ml-infra" -version = "1.0.0" -dependencies = [ - "ml-core", - "ml-dqn", - "ml-ppo", - "ml-supervised", -] - [[package]] name = "ml-ppo" version = "1.0.0" @@ -6440,7 +6431,32 @@ dependencies = [ name = "ml-supervised" version = "1.0.0" dependencies = [ + "anyhow", + "approx", + "async-trait", + "candle-core", + "candle-nn", + "common", + "config", + "dashmap 6.1.0", + "data", + "libc", + "lru", "ml-core", + "nalgebra 0.33.2", + "ndarray", + "num_cpus", + "parking_lot 0.12.5", + "petgraph 0.6.5", + "rand 0.8.5", + "rayon", + "serde", + "serde_json", + "tempfile", + "thiserror 1.0.69", + "tokio", + "tracing", + "uuid", ] [[package]] diff --git a/Cargo.toml b/Cargo.toml index 3dcd78e74..94b9aa5d5 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -113,7 +113,6 @@ members = [ "crates/ml-dqn", "crates/ml-ppo", "crates/ml-supervised", - "crates/ml-infra", "crates/ml-data", "crates/data", "crates/backtesting", @@ -393,7 +392,6 @@ ml-core = { path = "crates/ml-core" } ml-dqn = { path = "crates/ml-dqn" } ml-ppo = { path = "crates/ml-ppo" } ml-supervised = { path = "crates/ml-supervised" } -ml-infra = { path = "crates/ml-infra" } common = { path = "crates/common" } storage = { path = "crates/storage" } market-data = { path = "crates/market-data" } diff --git a/crates/ml-infra/Cargo.toml b/crates/ml-infra/Cargo.toml deleted file mode 100644 index 0cb9e194a..000000000 --- a/crates/ml-infra/Cargo.toml +++ /dev/null @@ -1,23 +0,0 @@ -[package] -name = "ml-infra" -version.workspace = true -edition.workspace = true -rust-version.workspace = true -authors.workspace = true -license.workspace = true -repository.workspace = true -homepage.workspace = true -documentation.workspace = true -publish.workspace = true -keywords.workspace = true -categories.workspace = true -description = "Training infrastructure (hyperopt, ensemble, trainers)" - -[dependencies] -ml-core.workspace = true -ml-dqn.workspace = true -ml-ppo.workspace = true -ml-supervised.workspace = true - -[lints] -workspace = true diff --git a/crates/ml-infra/src/lib.rs b/crates/ml-infra/src/lib.rs deleted file mode 100644 index 6ef21459c..000000000 --- a/crates/ml-infra/src/lib.rs +++ /dev/null @@ -1 +0,0 @@ -// Modules will be moved here from ml crate diff --git a/crates/ml-supervised/Cargo.toml b/crates/ml-supervised/Cargo.toml index 581a47f48..4f7a34b63 100644 --- a/crates/ml-supervised/Cargo.toml +++ b/crates/ml-supervised/Cargo.toml @@ -13,8 +13,58 @@ keywords.workspace = true categories.workspace = true description = "Supervised models (TFT, Mamba, Liquid, TGGN, TLOB, KAN, xLSTM, Diffusion)" +[features] +default = ["cuda"] +cuda = ["candle-core/cuda", "candle-core/cudnn", "candle-nn/cuda", "candle-nn/cudnn"] + [dependencies] ml-core.workspace = true +common.workspace = true +config.workspace = true +data.workspace = true + +# Async +tokio.workspace = true + +# ML frameworks +candle-core = { git = "https://github.com/huggingface/candle", rev = "671de1db" } +candle-nn = { git = "https://github.com/huggingface/candle", rev = "671de1db" } + +# Serialization +serde = { workspace = true, features = ["derive"] } +serde_json.workspace = true + +# Core utilities +thiserror.workspace = true +anyhow.workspace = true +tracing.workspace = true +rand.workspace = true +uuid.workspace = true +async-trait.workspace = true + +# Numerics +ndarray = { workspace = true, features = ["rayon"] } +nalgebra = { version = "0.33", features = ["serde-serialize"] } + +# Graph support (TGGN) +petgraph = { version = "0.6", features = ["serde"] } + +# Caching (TFT) +lru.workspace = true + +# System +libc = "0.2" +num_cpus = "1.16" + +# Concurrency +rayon.workspace = true +dashmap = { workspace = true } +parking_lot = { version = "0.12", features = ["hardware-lock-elision"] } + +[dev-dependencies] +tokio = { workspace = true, features = ["test-util", "macros"] } +approx.workspace = true +tempfile = "3" [lints] workspace = true diff --git a/crates/ml/src/diffusion/config.rs b/crates/ml-supervised/src/diffusion/config.rs similarity index 100% rename from crates/ml/src/diffusion/config.rs rename to crates/ml-supervised/src/diffusion/config.rs diff --git a/crates/ml/src/diffusion/denoiser.rs b/crates/ml-supervised/src/diffusion/denoiser.rs similarity index 98% rename from crates/ml/src/diffusion/denoiser.rs rename to crates/ml-supervised/src/diffusion/denoiser.rs index 33faeea6e..2608c5c31 100644 --- a/crates/ml/src/diffusion/denoiser.rs +++ b/crates/ml-supervised/src/diffusion/denoiser.rs @@ -4,7 +4,7 @@ //! This is more memory-efficient than Conv1D U-Net while still effective //! for price sequence denoising at small sequence lengths (64-128). -use crate::MLError; +use ml_core::MLError; use candle_core::{DType, Device, Tensor}; use candle_nn::{linear, Linear, Module, VarBuilder}; @@ -64,7 +64,7 @@ impl TimeEmbedding { .map_err(|e| MLError::ModelError(e.to_string()))?; // Cast to training dtype before projection through BF16 weights - let emb = crate::dqn::mixed_precision::ensure_training_dtype(&emb) + let emb = ml_core::mixed_precision::ensure_training_dtype(&emb) .map_err(|e| MLError::ModelError(e.to_string()))?; // Project to hidden_dim @@ -203,7 +203,7 @@ impl Denoiser { /// /// Input x: (batch, data_dim), t: (batch,) → Output: (batch, data_dim) pub fn forward(&self, x: &Tensor, t: &Tensor) -> Result { - let x = crate::dqn::mixed_precision::ensure_training_dtype(x) + let x = ml_core::mixed_precision::ensure_training_dtype(x) .map_err(|e| MLError::ModelError(e.to_string()))?; let map_err = |e: candle_core::Error| MLError::ModelError(e.to_string()); @@ -231,7 +231,7 @@ impl Denoiser { mod tests { use super::*; use candle_nn::VarMap; - use crate::dqn::mixed_precision::training_dtype; + use ml_core::mixed_precision::training_dtype; #[test] fn test_time_embedding_shape() { diff --git a/crates/ml-supervised/src/diffusion/mod.rs b/crates/ml-supervised/src/diffusion/mod.rs new file mode 100644 index 000000000..35aacfc9b --- /dev/null +++ b/crates/ml-supervised/src/diffusion/mod.rs @@ -0,0 +1,17 @@ +//! Diffusion model (DDPM/DDIM) for price path generation. +//! +//! Implements a denoising diffusion probabilistic model with: +//! - Cosine/linear noise schedules +//! - Fully-connected denoiser with sinusoidal time embedding +//! - DDIM sampling for fast inference +//! - UnifiedTrainable adapter for the training pipeline + +pub mod config; +pub mod denoiser; +pub mod noise; +pub mod sampler; + +pub use config::{DiffusionConfig, NoiseSchedule}; +pub use denoiser::Denoiser; +pub use noise::NoiseScheduler; +pub use sampler::DDIMSampler; diff --git a/crates/ml/src/diffusion/noise.rs b/crates/ml-supervised/src/diffusion/noise.rs similarity index 99% rename from crates/ml/src/diffusion/noise.rs rename to crates/ml-supervised/src/diffusion/noise.rs index bf6e4e5a8..34b13b917 100644 --- a/crates/ml/src/diffusion/noise.rs +++ b/crates/ml-supervised/src/diffusion/noise.rs @@ -3,7 +3,7 @@ //! Precomputes alpha_bar_t for all timesteps and provides //! forward process (add noise) operations. -use crate::MLError; +use ml_core::MLError; use candle_core::{Device, Tensor}; use super::config::NoiseSchedule; diff --git a/crates/ml/src/diffusion/sampler.rs b/crates/ml-supervised/src/diffusion/sampler.rs similarity index 99% rename from crates/ml/src/diffusion/sampler.rs rename to crates/ml-supervised/src/diffusion/sampler.rs index 5891fd612..1270b8060 100644 --- a/crates/ml/src/diffusion/sampler.rs +++ b/crates/ml-supervised/src/diffusion/sampler.rs @@ -3,7 +3,7 @@ //! Provides deterministic, fast sampling from a trained diffusion model //! using a small number of steps (e.g., 10) instead of the full T=1000. -use crate::MLError; +use ml_core::MLError; use candle_core::{Device, Tensor}; use super::denoiser::Denoiser; @@ -144,7 +144,7 @@ mod tests { use super::*; use super::super::config::{DiffusionConfig, NoiseSchedule}; use candle_nn::{VarBuilder, VarMap}; - use crate::dqn::mixed_precision::training_dtype; + use ml_core::mixed_precision::training_dtype; fn make_test_components() -> (Denoiser, NoiseScheduler, DDIMSampler) { let dev = Device::Cpu; diff --git a/crates/ml/src/kan/config.rs b/crates/ml-supervised/src/kan/config.rs similarity index 100% rename from crates/ml/src/kan/config.rs rename to crates/ml-supervised/src/kan/config.rs diff --git a/crates/ml/src/kan/layer.rs b/crates/ml-supervised/src/kan/layer.rs similarity index 98% rename from crates/ml/src/kan/layer.rs rename to crates/ml-supervised/src/kan/layer.rs index 0b3fa74b6..8157a2be0 100644 --- a/crates/ml/src/kan/layer.rs +++ b/crates/ml-supervised/src/kan/layer.rs @@ -6,7 +6,7 @@ use candle_core::Tensor; use candle_nn::VarBuilder; -use crate::MLError; +use ml_core::MLError; use super::spline::BSplineBasis; @@ -143,7 +143,7 @@ mod tests { use super::*; use candle_core::Device; use candle_nn::VarMap; - use crate::dqn::mixed_precision::training_dtype; + use ml_core::mixed_precision::training_dtype; #[test] fn test_kan_layer_output_shape() { diff --git a/crates/ml-supervised/src/kan/mod.rs b/crates/ml-supervised/src/kan/mod.rs new file mode 100644 index 000000000..2ce5cbe18 --- /dev/null +++ b/crates/ml-supervised/src/kan/mod.rs @@ -0,0 +1,13 @@ +//! KAN (Kolmogorov-Arnold Network) module. +//! +//! Implements KAN with learnable B-spline activation functions on each edge +//! of the network, replacing fixed activations (ReLU) with data-driven +//! non-linearities. + +pub mod config; +pub mod layer; +pub mod network; +pub mod spline; + +pub use config::KANConfig; +pub use network::KANNetwork; diff --git a/crates/ml/src/kan/network.rs b/crates/ml-supervised/src/kan/network.rs similarity index 96% rename from crates/ml/src/kan/network.rs rename to crates/ml-supervised/src/kan/network.rs index 27d1f60f2..9874f59fa 100644 --- a/crates/ml/src/kan/network.rs +++ b/crates/ml-supervised/src/kan/network.rs @@ -6,7 +6,7 @@ use candle_core::Tensor; use candle_nn::VarBuilder; -use crate::MLError; +use ml_core::MLError; use super::config::KANConfig; use super::layer::KANLayer; @@ -54,7 +54,7 @@ impl KANNetwork { /// Input shape: `(batch, layer_widths[0])` /// Output shape: `(batch, layer_widths[last])` pub fn forward(&self, input: &Tensor) -> Result { - let input = crate::dqn::mixed_precision::ensure_training_dtype(input) + let input = ml_core::mixed_precision::ensure_training_dtype(input) .map_err(|e| MLError::ModelError(e.to_string()))?; let mut x = input.clone(); for layer in &self.layers { @@ -70,7 +70,7 @@ impl KANNetwork { mod tests { use super::*; use candle_core::Device; - use crate::dqn::mixed_precision::training_dtype; + use ml_core::mixed_precision::training_dtype; use candle_nn::VarMap; fn make_config() -> KANConfig { diff --git a/crates/ml/src/kan/spline.rs b/crates/ml-supervised/src/kan/spline.rs similarity index 99% rename from crates/ml/src/kan/spline.rs rename to crates/ml-supervised/src/kan/spline.rs index b49f696c0..71bc2f0c0 100644 --- a/crates/ml/src/kan/spline.rs +++ b/crates/ml-supervised/src/kan/spline.rs @@ -9,7 +9,7 @@ //! single GPU gather + linear interpolation, keeping the GPU saturated. use candle_core::{DType, Device, Tensor}; -use crate::MLError; +use ml_core::MLError; /// Default lookup table resolution (number of uniformly spaced grid points). const DEFAULT_GRID_RESOLUTION: usize = 1024; diff --git a/crates/ml-supervised/src/lib.rs b/crates/ml-supervised/src/lib.rs index 6ef21459c..a287b4611 100644 --- a/crates/ml-supervised/src/lib.rs +++ b/crates/ml-supervised/src/lib.rs @@ -1 +1,19 @@ -// Modules will be moved here from ml crate +//! Supervised ML models for Foxhunt HFT trading +//! +//! Contains 8 model architectures: TFT, Mamba-2, Liquid, TGGN, TLOB, KAN, xLSTM, Diffusion. +//! Each model directory contains the core network, layers, and self-contained training logic. +//! Bridge modules (UnifiedTrainable adapters, Checkpointable impls) live in the `ml` crate. + +// Re-export shared types from ml-core +pub use ml_core::cuda_compat; +pub use ml_core::mixed_precision; +pub use ml_core::xavier_init; + +pub mod tft; +pub mod mamba; +pub mod liquid; +pub mod tgnn; +pub mod tlob; +pub mod kan; +pub mod xlstm; +pub mod diffusion; diff --git a/crates/ml/src/liquid/activation.rs b/crates/ml-supervised/src/liquid/activation.rs similarity index 100% rename from crates/ml/src/liquid/activation.rs rename to crates/ml-supervised/src/liquid/activation.rs diff --git a/crates/ml/src/liquid/candle_cfc.rs b/crates/ml-supervised/src/liquid/candle_cfc.rs similarity index 98% rename from crates/ml/src/liquid/candle_cfc.rs rename to crates/ml-supervised/src/liquid/candle_cfc.rs index 049d57a27..9879d7702 100644 --- a/crates/ml/src/liquid/candle_cfc.rs +++ b/crates/ml-supervised/src/liquid/candle_cfc.rs @@ -8,11 +8,11 @@ use candle_core::Tensor; use candle_nn::{Linear, Module, VarBuilder}; use serde::{Deserialize, Serialize}; -use crate::cuda_compat::manual_sigmoid; -use crate::MLError; +use ml_core::cuda_compat::manual_sigmoid; +use ml_core::MLError; /// Device configuration for CfC training -- re-exported from central gpu module. -pub use crate::gpu::DeviceConfig; +pub use ml_core::gpu::DeviceConfig; /// CfC v2 training configuration /// @@ -121,7 +121,7 @@ impl BackboneMLP { /// - `f_out` is tanh-activated (bounded in [-1, 1]) /// - `tau_out` is raw (sigmoid scaling applied in the CfC cell) pub fn forward(&self, input: &Tensor) -> Result<(Tensor, Tensor), MLError> { - let input = crate::dqn::mixed_precision::ensure_training_dtype(input) + let input = ml_core::mixed_precision::ensure_training_dtype(input) .map_err(|e| MLError::InferenceError(e.to_string()))?; let mut x = input.clone(); for layer in &self.layers { @@ -283,7 +283,7 @@ impl CandleCfCNetwork { let mut h = Tensor::zeros( (batch_size, self.config.hidden_size), - crate::dqn::mixed_precision::training_dtype(device), + ml_core::mixed_precision::training_dtype(device), device, ) .map_err(|e| MLError::InferenceError(format!("CfC init hidden: {}", e)))?; @@ -306,7 +306,7 @@ impl CandleCfCNetwork { /// Forward compatible with UnifiedTrainable (3D input, default dt=0.01) pub fn forward(&self, input: &Tensor) -> Result { - let input = crate::dqn::mixed_precision::ensure_training_dtype(input) + let input = ml_core::mixed_precision::ensure_training_dtype(input) .map_err(|e| MLError::InferenceError(e.to_string()))?; let output = self.forward_sequence(&input, 0.01)?; // Cast output back to F32 for API compatibility @@ -350,7 +350,7 @@ mod tests { use super::*; use candle_core::{DType, Device}; use candle_nn::VarMap; - use crate::dqn::mixed_precision::training_dtype; + use ml_core::mixed_precision::training_dtype; #[test] fn test_cfc_train_config_default() { diff --git a/crates/ml/src/liquid/cells.rs b/crates/ml-supervised/src/liquid/cells.rs similarity index 100% rename from crates/ml/src/liquid/cells.rs rename to crates/ml-supervised/src/liquid/cells.rs diff --git a/crates/ml-supervised/src/liquid/mod.rs b/crates/ml-supervised/src/liquid/mod.rs new file mode 100644 index 000000000..32f917216 --- /dev/null +++ b/crates/ml-supervised/src/liquid/mod.rs @@ -0,0 +1,210 @@ +//! Liquid Neural Networks for Ultra-Low Latency HFT +//! +//! Implementation of Liquid Time-constant (LTC) and Closed-form Continuous-time (CfC) +//! neural networks with fixed-point arithmetic for sub-100μs inference. + +use std::error::Error; +use std::fmt; + +use serde::{Deserialize, Serialize}; + +// Import MarketRegime from core types to avoid type conflicts +use ml_core::MLError; +use common::trading::MarketRegime; + +pub mod activation; +pub mod candle_cfc; +pub mod cells; +pub mod network; +pub mod ode_solvers; +pub mod training; + +#[cfg(test)] +mod tests; + +// Re-export main types for external usage +pub use activation::ActivationType; +pub use candle_cfc::{BackboneMLP, CandleCfCNetwork, CfCCell, CfCTrainConfig, DeviceConfig}; +pub use cells::{CfCConfig, LTCConfig}; +pub use network::{LayerConfig, LiquidNetwork, LiquidNetworkConfig, OutputLayerConfig}; +pub use ode_solvers::SolverType; +pub use training::{ + CandleCfCTrainer, CfCTrainerConfig, LiquidTrainer, LiquidTrainingConfig, TrainingBatch, + TrainingMetrics, TrainingSample, TrainingUtils, +}; + +/// Fixed-point arithmetic for ultra-low latency inference +pub const PRECISION: i64 = 100_000_000; // 8 decimal places + +#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize, Default)] +pub struct FixedPoint(pub i64); + +impl FixedPoint { + pub fn from_f64(value: f64) -> Self { + FixedPoint((value * PRECISION as f64) as i64) + } + + pub fn to_f64(self) -> f64 { + self.0 as f64 / PRECISION as f64 + } + + pub fn zero() -> Self { + FixedPoint(0) + } + + pub fn one() -> Self { + FixedPoint(PRECISION) + } + + pub fn is_finite(&self) -> bool { + self.0.abs() < i64::MAX / 2 + } +} + +impl std::ops::Add for FixedPoint { + type Output = Result; + + fn add(self, rhs: FixedPoint) -> Self::Output { + self.0 + .checked_add(rhs.0) + .map(FixedPoint) + .ok_or(LiquidError::Overflow("Addition overflow".to_owned())) + } +} + +impl std::ops::Sub for FixedPoint { + type Output = Result; + + fn sub(self, rhs: FixedPoint) -> Self::Output { + self.0 + .checked_sub(rhs.0) + .map(FixedPoint) + .ok_or(LiquidError::Overflow("Subtraction overflow".to_owned())) + } +} + +impl std::ops::Mul for FixedPoint { + type Output = Result; + + fn mul(self, rhs: FixedPoint) -> Self::Output { + let result = ((self.0 as i128) * (rhs.0 as i128)) / (PRECISION as i128); + if result > i64::MAX as i128 || result < i64::MIN as i128 { + Err(LiquidError::Overflow("Multiplication overflow".to_owned())) + } else { + Ok(FixedPoint(result as i64)) + } + } +} + +impl std::ops::Div for FixedPoint { + type Output = Result; + + fn div(self, rhs: FixedPoint) -> Self::Output { + if rhs.0 == 0 { + return Err(LiquidError::DivisionByZero); + } + let result = ((self.0 as i128) * (PRECISION as i128)) / (rhs.0 as i128); + if result > i64::MAX as i128 || result < i64::MIN as i128 { + Err(LiquidError::Overflow("Division overflow".to_owned())) + } else { + Ok(FixedPoint(result as i64)) + } + } +} + +/// Liquid Neural Network specific errors +#[derive(Debug, Clone, Serialize, Deserialize)] +pub enum LiquidError { + InvalidConfiguration(String), + InvalidInput(String), + Overflow(String), + DivisionByZero, + InferenceError(String), + TrainingError(String), + SolverError(String), +} + +impl fmt::Display for LiquidError { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + match self { + LiquidError::InvalidConfiguration(msg) => write!(f, "Invalid configuration: {}", msg), + LiquidError::InvalidInput(msg) => write!(f, "Invalid input: {}", msg), + LiquidError::Overflow(msg) => write!(f, "Overflow error: {}", msg), + LiquidError::DivisionByZero => write!(f, "Division by zero"), + LiquidError::InferenceError(msg) => write!(f, "Inference error: {}", msg), + LiquidError::TrainingError(msg) => write!(f, "Training error: {}", msg), + LiquidError::SolverError(msg) => write!(f, "ODE solver error: {}", msg), + } + } +} + +impl Error for LiquidError {} + +impl From for MLError { + fn from(err: LiquidError) -> Self { + match err { + LiquidError::InvalidConfiguration(msg) => MLError::ConfigError(msg), + LiquidError::InvalidInput(msg) => MLError::InvalidInput(msg), + LiquidError::InferenceError(msg) => MLError::InferenceError(msg), + LiquidError::TrainingError(msg) => MLError::TrainingError(msg), + LiquidError::Overflow(_) + | LiquidError::DivisionByZero + | LiquidError::SolverError(_) => MLError::ModelError(err.to_string()), + } + } +} + +impl From for LiquidError { + fn from(err: MLError) -> Self { + match err { + MLError::ConfigError(msg) => { + LiquidError::InvalidConfiguration(msg) + }, + MLError::InvalidInput(msg) => LiquidError::InvalidInput(msg), + MLError::InferenceError(msg) => LiquidError::InferenceError(msg), + MLError::TrainingError(msg) => LiquidError::TrainingError(msg), + MLError::DimensionMismatch { .. } + | MLError::GraphError { .. } + | MLError::ResourceLimit { .. } + | MLError::SerializationError { .. } + | MLError::ValidationError { .. } + | MLError::ConcurrencyError { .. } + | MLError::InitializationError { .. } + | MLError::ModelError(_) + | MLError::NotTrained(_) + | MLError::AnyhowError(_) + | MLError::TensorCreationError { .. } + | MLError::TensorOperationError(_) + | MLError::LockError(_) + | MLError::ModelNotFound(_) + | MLError::InsufficientData(_) + | MLError::CheckpointError(_) + | MLError::DeviceError(_) => LiquidError::InferenceError(err.to_string()), + } + } +} + +pub type Result = std::result::Result; + +/// Network type for liquid neural networks +#[derive(Debug, Clone, Serialize, Deserialize)] +pub enum NetworkType { + LTC, // Liquid Time-constant + CfC, // Closed-form Continuous-time + Mixed, // Combination of LTC and CfC layers +} + +// REMOVED: MarketRegime enum - now using common::MarketRegime instead +// This eliminates the type conflict and ensures consistency across the entire system + +/// Performance metrics for liquid networks +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct PerformanceMetrics { + pub total_inferences: u64, + pub average_inference_time_ns: u64, + pub average_inference_time_us: f64, + pub total_parameters: usize, + pub current_regime: MarketRegime, // Now uses core MarketRegime enum + pub regime_switches: u32, + pub last_adaptation_time: Option, // Store as timestamp millis instead of Instant +} diff --git a/crates/ml/src/liquid/network.rs b/crates/ml-supervised/src/liquid/network.rs similarity index 99% rename from crates/ml/src/liquid/network.rs rename to crates/ml-supervised/src/liquid/network.rs index 5d318ce35..626c83f1a 100644 --- a/crates/ml/src/liquid/network.rs +++ b/crates/ml-supervised/src/liquid/network.rs @@ -13,7 +13,7 @@ use super::cells::{CfCCell, CfCConfig, LTCCell, LTCConfig}; use super::{ FixedPoint, LiquidError, MarketRegime, NetworkType, PerformanceMetrics, Result, PRECISION, }; -use crate::{MLError, MLResult}; +use ml_core::{MLError, MLResult}; /// Layer configuration for liquid networks #[derive(Debug, Clone, Serialize, Deserialize)] diff --git a/crates/ml/src/liquid/ode_solvers.rs b/crates/ml-supervised/src/liquid/ode_solvers.rs similarity index 100% rename from crates/ml/src/liquid/ode_solvers.rs rename to crates/ml-supervised/src/liquid/ode_solvers.rs diff --git a/crates/ml/src/liquid/tests.rs b/crates/ml-supervised/src/liquid/tests.rs similarity index 100% rename from crates/ml/src/liquid/tests.rs rename to crates/ml-supervised/src/liquid/tests.rs diff --git a/crates/ml/src/liquid/training.rs b/crates/ml-supervised/src/liquid/training.rs similarity index 99% rename from crates/ml/src/liquid/training.rs rename to crates/ml-supervised/src/liquid/training.rs index e9eb0eb84..fea37ab3f 100644 --- a/crates/ml/src/liquid/training.rs +++ b/crates/ml-supervised/src/liquid/training.rs @@ -466,8 +466,8 @@ use candle_core::{Device, Tensor}; use candle_nn::{AdamW, Optimizer, ParamsAdamW, VarBuilder, VarMap}; use super::candle_cfc::{CandleCfCNetwork, CfCTrainConfig}; -use crate::dqn::mixed_precision::training_dtype; -use crate::MLError; +use ml_core::mixed_precision::training_dtype; +use ml_core::MLError; /// Configuration for the Candle-based CfC trainer #[derive(Debug, Clone, Serialize, Deserialize)] diff --git a/crates/ml/src/mamba/cuda/selective_scan.cu b/crates/ml-supervised/src/mamba/cuda/selective_scan.cu similarity index 100% rename from crates/ml/src/mamba/cuda/selective_scan.cu rename to crates/ml-supervised/src/mamba/cuda/selective_scan.cu diff --git a/crates/ml/src/mamba/hardware_aware.rs b/crates/ml-supervised/src/mamba/hardware_aware.rs similarity index 99% rename from crates/ml/src/mamba/hardware_aware.rs rename to crates/ml-supervised/src/mamba/hardware_aware.rs index b3a6b3232..30ab712c3 100644 --- a/crates/ml/src/mamba/hardware_aware.rs +++ b/crates/ml-supervised/src/mamba/hardware_aware.rs @@ -22,8 +22,8 @@ use nalgebra::DMatrix; use tracing::info; use super::Mamba2Config; -use crate::MLError; -use crate::PRECISION_FACTOR; +use ml_core::MLError; +use ml_core::PRECISION_FACTOR; // Platform-specific imports #[cfg(target_arch = "x86_64")] diff --git a/crates/ml/src/mamba/loss.rs b/crates/ml-supervised/src/mamba/loss.rs similarity index 99% rename from crates/ml/src/mamba/loss.rs rename to crates/ml-supervised/src/mamba/loss.rs index 8e3dc478f..6204df69a 100644 --- a/crates/ml/src/mamba/loss.rs +++ b/crates/ml-supervised/src/mamba/loss.rs @@ -5,7 +5,7 @@ //! than magnitude. use candle_core::Tensor; -use crate::{MLError, MLResult}; +use ml_core::{MLError, MLResult}; /// Standard MSE loss (baseline for comparison) pub fn mse_loss(predictions: &Tensor, targets: &Tensor) -> MLResult { diff --git a/crates/ml-supervised/src/mamba/mod.rs b/crates/ml-supervised/src/mamba/mod.rs new file mode 100644 index 000000000..94370357e --- /dev/null +++ b/crates/ml-supervised/src/mamba/mod.rs @@ -0,0 +1,3782 @@ +//! # Mamba-2 State-Space Model for HFT +//! +//! Next-generation Mamba-2 implementation with Structured State Duality (SSD), +//! hardware-aware algorithms, and 5x performance improvements over Mamba-1. +//! +//! **Note**: This module uses mathematical notation (A, B, C for state-space matrices). +//! Non-snake-case warnings are allowed for mathematical clarity. + +#![allow(non_snake_case)] +//! +//! ## Key Features +//! +//! - **SSD Layers**: Structured State Duality for linear attention mechanisms +//! - **Hardware-aware**: Optimized memory access patterns and SIMD instructions +//! - **5x Faster**: Sub-linear memory usage and linear-time sequence modeling +//! - **Selective State Spaces**: Advanced state selection mechanisms +//! - **Sub-5μs**: Target inference latency for HFT applications +//! - **Integer Precision**: 10,000x scaling for financial precision +//! +//! ## Architecture Improvements +//! +//! ```text +//! ┌─────────────────────────────────────────────────────────────┐ +//! │ Mamba-2 Block │ +//! ├─────────────────┬─────────────────┬─────────────────────────┤ +//! │ SSD Layer │ Hardware-Aware │ Selective State │ +//! │ │ Optimization │ Mechanism │ +//! │ • Linear Attn │ • SIMD Vectors │ • Advanced Selection │ +//! │ • Structured │ • Cache-Friendly│ • Dynamic Parameters │ +//! │ Duality │ • Prefetching │ • State Compression │ +//! └─────────────────┴─────────────────┴─────────────────────────┘ +//! ``` +//! +//! ## Performance Targets +//! +//! - Inference: <5μs per sequence step (5x faster than Mamba-1) +//! - Memory: Sub-linear growth with sequence length +//! - Throughput: >1M sequences/sec +//! - Latency: 99.9% percentile <10μs + +mod hardware_aware; +mod scan_algorithms; +pub mod selective_state; +mod ssd_layer; +pub mod loss; + +// Public exports for types used in mod.rs and by external crates +pub use hardware_aware::{HardwareCapabilities, HardwareOptimizer}; +pub use scan_algorithms::{ParallelScanEngine, ScanBenchmark, ScanOperator}; +pub use selective_state::{ + SelectiveStateConfig, SelectiveStateSpace, StateCompressor, StateImportance, +}; +pub use ssd_layer::SSDLayer; + +use std::collections::{HashMap, VecDeque}; +use std::sync::atomic::{AtomicU64, Ordering}; +use std::sync::Arc; +use std::time::{Duration, Instant, SystemTime}; + +use candle_core::{DType, Device, Tensor, Var}; +use candle_nn::Module; +use candle_nn::{Dropout, Linear, VarBuilder}; +use serde::{Deserialize, Serialize}; +use tracing::{debug, info, instrument, trace, warn}; +use uuid::Uuid; + +use ml_core::cuda_compat::layer_norm_with_fallback; +use ml_core::mixed_precision::training_dtype; +use ml_core::MLError; + +/// Optimizer type for training +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] +pub enum OptimizerType { + /// Adam optimizer with adaptive learning rates (coupled weight decay) + Adam, + /// AdamW optimizer with decoupled weight decay (recommended for SSMs) + AdamW, + /// Stochastic Gradient Descent with momentum + SGD, +} + +impl Default for OptimizerType { + fn default() -> Self { + Self::AdamW // AdamW is superior for state-space models + } +} + +/// Configuration for `MAMBA-2` state-space model +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct Mamba2Config { + /// Model dimension + pub d_model: usize, + /// State dimension + pub d_state: usize, + /// Head dimension for multi-head attention + pub d_head: usize, + /// Number of attention heads + pub num_heads: usize, + /// Expansion factor for inner dimension + pub expand: usize, + /// Number of layers + pub num_layers: usize, + /// Dropout rate + pub dropout: f64, + /// Use structured state duality + pub use_ssd: bool, + /// Use selective state mechanism + pub use_selective_state: bool, + /// Enable hardware optimizations + pub hardware_aware: bool, + /// Target latency in microseconds + pub target_latency_us: u64, + /// Maximum sequence length + pub max_seq_len: usize, + /// Learning rate + pub learning_rate: f64, + /// Weight decay + pub weight_decay: f64, + /// Gradient clipping threshold + pub grad_clip: f64, + /// Warmup steps + pub warmup_steps: usize, + /// Adam beta1 momentum parameter + pub adam_beta1: f64, + /// P1: Adam beta2 parameter (Agent 2) + pub adam_beta2: f64, + /// P1: Adam epsilon (Agent 2) + pub adam_epsilon: f64, + /// P1: Total decay steps for cosine schedule (Agent 2) + pub total_decay_steps: usize, + /// Optimizer type (Adam or SGD) + pub optimizer_type: OptimizerType, + /// SGD momentum (only used when optimizer_type = SGD) + pub sgd_momentum: f64, + /// Training batch size + pub batch_size: usize, + /// Sequence length for training + pub seq_len: usize, + /// Shuffle batches every epoch (default: false for reproducibility) + pub shuffle_batches: bool, + /// P2: Sequence stride for overlapping windows (Agent 3) + pub sequence_stride: usize, + /// P2: Normalization epsilon for layer norm (Agent 3) + pub norm_eps: f64, + /// Enable early stopping (default: true) + pub early_stopping_enabled: bool, + /// Early stopping patience (epochs without improvement) + pub early_stopping_patience: usize, + /// Early stopping threshold (minimum improvement) + pub early_stopping_min_delta: f64, + /// Minimum epochs before early stopping can trigger + pub early_stopping_min_epochs: usize, +} + +impl Default for Mamba2Config { + fn default() -> Self { + Self::emergency_safe_defaults() + } +} + +impl Mamba2Config { + /// Create Mamba2 config from central configuration system + /// + /// CRITICAL: Eliminates dangerous hardcoded defaults that could cause + /// training instability or memory issues in production + pub fn from_config_manager( + _config_manager: &config::ConfigManager, + ) -> Result> { + // Use emergency defaults since specific MAMBA configs may not be available + tracing::warn!( + "Using emergency MAMBA config defaults - MAMBA configs not available in ServiceConfig" + ); + Ok(Self::emergency_safe_defaults()) + } + + /// EMERGENCY FALLBACK: Ultra-conservative Mamba2 defaults + /// + /// WARNING: These defaults prioritize safety over performance + /// and are not suitable for production training + pub fn emergency_safe_defaults() -> Self { + tracing::error!( + "Using emergency Mamba2 defaults - check configuration system immediately!" + ); + Self { + d_model: 225, // Wave C (201) + Wave D (24) = 225 + d_state: 64, // P0 FIX: Mamba-2 official recommendation (was 16) + d_head: 16, // Small head size + num_heads: 2, // Minimal heads + expand: 1, // No expansion to minimize memory + num_layers: 1, // Single layer only + dropout: 0.5, // High dropout for safety + use_ssd: false, // Disable advanced features + use_selective_state: false, // Disable advanced features + hardware_aware: false, // Disable optimizations + target_latency_us: 1000, // Very conservative latency + max_seq_len: 128, // Short sequences only + learning_rate: 1e-6, // Extremely conservative learning rate + weight_decay: 1e-3, // High weight decay for stability + grad_clip: 0.1, // Aggressive gradient clipping + warmup_steps: 10, // Minimal warmup + adam_beta1: 0.9, // Standard Adam beta1 + adam_beta2: 0.999, // P1: Standard Adam beta2 + adam_epsilon: 1e-8, // P1: Standard Adam epsilon + total_decay_steps: 10000, // P1: Standard decay schedule + optimizer_type: OptimizerType::AdamW, // AdamW for better SSM training + sgd_momentum: 0.9, // Standard SGD momentum + batch_size: 1, // Single sample batches + seq_len: 64, // Very short sequences + shuffle_batches: false, // Deterministic by default + sequence_stride: 1, // P2: No overlapping (safe default) + norm_eps: 1e-5, // P2: Standard layer norm epsilon + early_stopping_enabled: true, // Enable early stopping by default + early_stopping_patience: 20, // 20 epochs patience (TFT default) + early_stopping_min_delta: 1e-4, // Minimum improvement threshold + early_stopping_min_epochs: 20, // Minimum 20 epochs before stopping + } + } + + /// Estimate memory usage for safety validation + fn estimate_memory_usage(config: &Mamba2Config) -> usize { + // Rough estimation: d_model * num_layers * batch_size * seq_len * 4 bytes (f32) + // Plus additional overhead for state and intermediate computations + let base_memory = + config.d_model * config.num_layers * config.batch_size * config.seq_len * 4; + let overhead_factor = 3; // Account for gradients, optimizer states, etc. + (base_memory * overhead_factor) / (1024 * 1024) // Convert to MB + } +} + +/// `MAMBA-2` state container +#[derive(Debug, Clone)] +pub struct Mamba2State { + /// Hidden states for each layer + pub hidden_states: Vec, + /// Selective state components + pub selective_state: Vec, + /// State transition matrices A, B, C + pub ssm_states: Vec, + /// Compression indices for memory efficiency + pub compression_indices: Vec, + /// Performance metrics + pub metrics: HashMap, + /// Best validation loss (for early stopping) + pub best_val_loss: f64, + /// Patience counter (epochs without improvement) + pub patience_counter: usize, + /// Early stopping triggered flag + pub stopped: bool, + /// Epoch where early stopping triggered + pub stopped_at_epoch: Option, + /// Last update timestamp + pub last_update: Instant, +} + +/// State Space Model state matrices +/// +/// Mathematical notation: A, B, C matrices follow standard SSM formulation +/// where uppercase letters represent state-space matrices as per control theory convention +#[derive(Debug, Clone)] +#[allow(non_snake_case)] +pub struct SSMState { + /// State transition matrix A (d_state × d_state) + /// Mathematical notation: uppercase A is standard in control theory and SSM literature + #[allow(non_snake_case)] + pub A: Tensor, + /// Input matrix B (d_state × d_model) + /// Mathematical notation: uppercase B is standard in control theory and SSM literature + #[allow(non_snake_case)] + pub B: Tensor, + /// Output matrix C (d_model × d_state) + /// Mathematical notation: uppercase C is standard in control theory and SSM literature + #[allow(non_snake_case)] + pub C: Tensor, + /// Discretization parameter Δ (Delta) + pub delta: Tensor, + /// Current hidden state + pub hidden: Tensor, +} + +impl SSMState { + /// Reset SSM state to zeros (call between epochs to prevent state accumulation) + /// + /// # Errors + /// + /// Returns `MLError` if tensor operations fail + pub fn reset(&mut self) -> Result<(), MLError> { + // Reset A, B, C matrices to initial random values (small initialization for stability) + // Clone device first to avoid borrow checker issues + let device = self.A.device().clone(); + let d_state = self.A.dim(0)?; + let d_inner = self.B.dim(1)?; + let d_model = self.delta.dims()[0]; + let batch_size = self.hidden.dim(0)?; + + // Re-initialize A matrix [d_state, d_state] -- use Tensor::randn to avoid temporary Vec allocation + self.A = Tensor::randn(0_f32, 0.02, (d_state, d_state), &device)?; + + // Re-initialize B matrix [d_state, d_inner] + self.B = Tensor::randn(0_f32, 0.02, (d_state, d_inner), &device)?; + + // Re-initialize C matrix [d_inner, d_state] + self.C = Tensor::randn(0_f32, 0.02, (d_inner, d_state), &device)?; + + // Reset delta to ones + self.delta = Tensor::ones((d_model,), DType::F32, &device)?; + + // Reset hidden state to zeros + self.hidden = Tensor::zeros((batch_size, d_state), DType::F32, &device)?; + + Ok(()) + } +} + +impl Mamba2State { + /// Create a zero-initialized state + /// + /// # Errors + /// + /// Returns `MLError` if: + /// - CUDA device initialization fails (falls back to CPU) + /// - Tensor allocation fails + /// - Memory allocation exceeds available resources + pub fn zeros(config: &Mamba2Config, device: &Device) -> Result { + let mut hidden_states = Vec::new(); + let mut ssm_states = Vec::new(); + let d_inner = config.d_model * config.expand; // CRITICAL: Use d_inner after input_projection + + for layer_idx in 0..config.num_layers { + // Create hidden state with proper error handling + let hidden = Tensor::zeros((config.batch_size, config.d_model), DType::F32, device) + .map_err(|e| MLError::TensorCreationError { + operation: format!("hidden state creation for layer {}", layer_idx), + reason: e.to_string(), + })?; + hidden_states.push(hidden); + + // Initialize SSM matrices with F32 dtype for GPU throughput + let A = { + let shape = (config.d_state, config.d_state); + let num_elements = shape.0 * shape.1; + let values: Vec = (0..num_elements) + .map(|_| { + use rand::Rng; + let mut rng = rand::thread_rng(); + rng.gen_range(-1.0_f32..1.0) * 0.02 // Small initialization for stability + }) + .collect(); + Tensor::from_vec(values, shape, device).map_err(|e| { + MLError::TensorCreationError { + operation: format!("SSM A matrix creation for layer {}", layer_idx), + reason: e.to_string(), + } + })? + }; + trace!( + "Layer {} A matrix initialized: shape={:?}, dtype=F32", + layer_idx, + A.dims() + ); + + // B must be [d_state, d_inner] with F32 dtype + let B = { + let shape = (config.d_state, d_inner); + let num_elements = shape.0 * shape.1; + let values: Vec = (0..num_elements) + .map(|_| { + use rand::Rng; + let mut rng = rand::thread_rng(); + rng.gen_range(-1.0_f32..1.0) * 0.02 + }) + .collect(); + Tensor::from_vec(values, shape, device).map_err(|e| { + MLError::TensorCreationError { + operation: format!("SSM B matrix creation for layer {}", layer_idx), + reason: e.to_string(), + } + })? + }; + trace!( + "Layer {} B matrix initialized: shape={:?}, dtype=F32", + layer_idx, + B.dims() + ); + + // C must be [d_inner, d_state] with F32 dtype + let C = { + let shape = (d_inner, config.d_state); + let num_elements = shape.0 * shape.1; + let values: Vec = (0..num_elements) + .map(|_| { + use rand::Rng; + let mut rng = rand::thread_rng(); + rng.gen_range(-1.0_f32..1.0) * 0.02 + }) + .collect(); + Tensor::from_vec(values, shape, device).map_err(|e| { + MLError::TensorCreationError { + operation: format!("SSM C matrix creation for layer {}", layer_idx), + reason: e.to_string(), + } + })? + }; + trace!( + "Layer {} C matrix initialized: shape={:?}, dtype=F32", + layer_idx, + C.dims() + ); + + let delta = Tensor::ones((config.d_model,), DType::F32, device).map_err(|e| { + MLError::TensorCreationError { + operation: format!("delta tensor creation for layer {}", layer_idx), + reason: e.to_string(), + } + })?; + + let ssm_hidden = Tensor::zeros((config.batch_size, config.d_state), DType::F32, device) + .map_err(|e| MLError::TensorCreationError { + operation: format!("SSM hidden state creation for layer {}", layer_idx), + reason: e.to_string(), + })?; + + ssm_states.push(SSMState { + A, + B, + C, + delta, + hidden: ssm_hidden, + }); + } + + Ok(Self { + hidden_states, + selective_state: vec![0.0; config.d_model * config.expand], + ssm_states, + compression_indices: Vec::new(), + metrics: HashMap::new(), + best_val_loss: f64::INFINITY, + patience_counter: 0, + stopped: false, + stopped_at_epoch: None, + last_update: Instant::now(), + }) + } + + /// Compress state to reduce memory usage + pub fn compress(&mut self, compression_ratio: f64) { + let target_size = (self.selective_state.len() as f64 * compression_ratio) as usize; + + // Sort by magnitude and keep top components + let mut indexed_values: Vec<(usize, f64)> = self + .selective_state + .iter() + .enumerate() + .map(|(i, &v)| (i, v.abs())) + .collect(); + + indexed_values.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal)); + + self.compression_indices.clear(); + for i in 0..target_size.min(indexed_values.len()) { + self.compression_indices.push(indexed_values[i].0); + } + + // Zero out non-selected components + for i in 0..self.selective_state.len() { + if !self.compression_indices.contains(&i) { + self.selective_state[i] = 0.0; + } + } + } +} + +/// Training metadata for `MAMBA-2` model +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct Mamba2Metadata { + pub model_id: String, + pub created_at: SystemTime, + pub version: String, + pub input_dim: usize, + pub output_dim: usize, + pub num_parameters: usize, + pub training_history: VecDeque, + pub performance_stats: HashMap, + pub last_checkpoint: Option, +} + +/// Training epoch information +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct TrainingEpoch { + pub epoch: usize, + pub loss: f64, + pub accuracy: f64, + pub learning_rate: f64, + pub duration_seconds: f64, + pub timestamp: SystemTime, +} + +/// CUDA-compatible LayerNorm wrapper for MAMBA-2 +/// +/// This wrapper uses manual CUDA implementation to avoid +/// "no cuda implementation for layer-norm" error from Candle. +#[derive(Debug, Clone)] +pub struct CudaLayerNorm { + normalized_shape: Vec, + weight: Option, + bias: Option, + eps: f64, +} + +impl CudaLayerNorm { + pub fn new(normalized_shape: usize, eps: f64, vb: VarBuilder<'_>) -> Result { + // Create learnable weight and bias parameters + let weight = vb.get(normalized_shape, "weight")?; + let bias = vb.get(normalized_shape, "bias")?; + + Ok(Self { + normalized_shape: vec![normalized_shape], + weight: Some(weight), + bias: Some(bias), + eps, + }) + } + + pub fn forward(&self, x: &Tensor) -> Result { + layer_norm_with_fallback( + x, + &self.normalized_shape, + self.weight.as_ref(), + self.bias.as_ref(), + self.eps, + ) + } +} + +/// `MAMBA-2` State-Space Model implementation +pub struct Mamba2SSM { + pub config: Mamba2Config, + pub metadata: Mamba2Metadata, + pub state: Mamba2State, + pub ssd_layers: Vec, + pub selective_state: Option, + pub hardware_optimizer: Option, + pub scan_engine: Arc, + pub is_trained: bool, + pub device: Device, + + // Model parameters + pub input_projection: Linear, + pub output_projection: Linear, + pub layer_norms: Vec, + pub dropouts: Vec, + + // Training state + pub optimizer_state: HashMap, + pub gradients: HashMap, + pub grad_scaler: f64, + pub step_count: usize, + pub current_lr: f64, + pub total_training_samples: usize, + + // Performance counters + pub total_inferences: AtomicU64, + pub total_training_steps: AtomicU64, + pub latency_histogram: VecDeque, + + // AGENT F2: VarMap for checkpoint saving (CRITICAL FIX) + // This stores all trainable parameters for safetensors serialization + pub varmap: Arc, +} + +impl std::fmt::Debug for Mamba2SSM { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("Mamba2SSM") + .field("config", &self.config) + .field("metadata", &self.metadata) + .field("is_trained", &self.is_trained) + .field("device", &self.device) + .field("step_count", &self.step_count) + .field("varmap", &"") + .finish() + } +} + +impl Mamba2SSM { + /// Create a scalar tensor with automatic dtype conversion + /// + /// Helper function to eliminate repetitive dtype matching boilerplate. + /// Automatically converts f64 values to the appropriate tensor dtype. + fn scalar_tensor(value: f64, dtype: DType, device: &Device) -> Result { + match dtype { + DType::F32 => { + Tensor::new(&[value as f32], device).map_err(|e| MLError::TensorCreationError { + operation: "scalar_tensor (F32)".to_owned(), + reason: e.to_string(), + }) + }, + DType::F64 => { + // F64 path kept for backward compatibility but should not be hit + // after the F32 migration + Tensor::new(&[value], device).map_err(|e| MLError::TensorCreationError { + operation: "scalar_tensor (F64)".to_owned(), + reason: e.to_string(), + }) + }, + DType::BF16 | DType::F16 => { + // Create as F32 scalar, then cast to the target half-precision dtype. + Tensor::new(&[value as f32], device) + .and_then(|t| t.to_dtype(dtype)) + .map_err(|e| MLError::TensorCreationError { + operation: format!("scalar_tensor ({:?} via F32)", dtype), + reason: e.to_string(), + }) + }, + DType::F8E4M3 | DType::U8 | DType::U32 | DType::I64 => { + Err(MLError::ModelError(format!( + "Unsupported dtype: {:?}", + dtype + ))) + }, + } + } + + /// Create new `MAMBA-2` model + /// + /// # Errors + /// + /// Returns `MLError` if: + /// - Variable initialization fails + /// - Linear layer creation fails + /// - Layer norm creation fails + /// - SSD layer initialization fails + pub fn new(config: Mamba2Config, device: &Device) -> Result { + if config.d_model == 0 { + return Err(MLError::ConfigError("Mamba2 requires d_model > 0".to_owned())); + } + + let vs = Arc::new(candle_nn::VarMap::new()); + let vb = VarBuilder::from_varmap(&vs, training_dtype(device), device); + + let d_inner = config.d_model * config.expand; + + let input_projection = candle_nn::linear(config.d_model, d_inner, vb.pp("input_proj"))?; + // FIXED (Agent 246): Output projection should map d_inner to 1 for regression (price prediction) + // The model performs price regression, NOT sequence-to-sequence modeling + // Output shape: [batch, seq, d_inner] → [batch, seq, 1] + let output_projection = candle_nn::linear(d_inner, 1, vb.pp("output_proj"))?; + + let mut layer_norms = Vec::new(); + let mut dropouts = Vec::new(); + let mut ssd_layers = Vec::new(); + + // Layer norm must match d_inner (d_model * expand) since input_projection expands the dimension + for i in 0..config.num_layers { + let ln = CudaLayerNorm::new(d_inner, config.norm_eps, vb.pp(format!("ln_{}", i)))?; + layer_norms.push(ln); + + let dropout = Dropout::new(config.dropout as f32); + dropouts.push(dropout); + + // CRITICAL FIX: Pass VarBuilder to SSDLayer to register parameters in parent VarMap + // Previously passed device, causing SSDLayer to create local VarMap (90% of params lost) + let ssd_layer = SSDLayer::new(&config, i, vb.clone())?; + ssd_layers.push(ssd_layer); + } + + let selective_state = config + .use_selective_state + .then(|| SelectiveStateSpace::new(&config)) + .transpose()?; + + let hardware_optimizer = config + .hardware_aware + .then(|| HardwareOptimizer::new(&config)) + .transpose()?; + + let scan_engine = Arc::new(ParallelScanEngine::new(device.clone(), 1_000_000)); + + let metadata = Mamba2Metadata { + model_id: Uuid::new_v4().to_string(), + created_at: SystemTime::now(), + version: "2.0.0".to_owned(), + input_dim: config.d_model, + output_dim: 1, // FIXED (Agent 246): Regression output (price prediction), not sequence-to-sequence + num_parameters: Self::count_parameters(&config), + training_history: VecDeque::new(), + performance_stats: HashMap::new(), + last_checkpoint: None, + }; + + let state = Mamba2State::zeros(&config, device)?; + + // Store learning_rate before moving config + let learning_rate = config.learning_rate; + + Ok(Self { + config, + metadata, + state, + ssd_layers, + selective_state, + hardware_optimizer, + scan_engine, + is_trained: false, + device: device.clone(), + input_projection, + output_projection, + layer_norms, + dropouts, + optimizer_state: HashMap::new(), + gradients: HashMap::new(), + grad_scaler: 1.0, + step_count: 0, + current_lr: learning_rate, + total_training_samples: 0, + total_inferences: AtomicU64::new(0), + total_training_steps: AtomicU64::new(0), + latency_histogram: VecDeque::new(), + varmap: vs, // AGENT F2: Store VarMap for checkpoint saving + }) + } + + /// Count total parameters in model + fn count_parameters(config: &Mamba2Config) -> usize { + let d_inner = config.d_model * config.expand; + let input_proj_params = config.d_model * d_inner; + let output_proj_params = d_inner; // FIXED (Agent 246): d_inner * 1 for regression output + let layer_params = config.num_layers + * ( + config.d_model * 3 + // Layer norm + config.d_model * config.d_state * 3 + // A, B, C matrices + config.d_model + // Delta parameters + ); + + input_proj_params + output_proj_params + layer_params + } + + /// Create HFT-optimized configuration + /// + /// # Errors + /// + /// Returns `MLError` if: + /// - Model initialization fails + /// - Hardware configuration is invalid + /// - Resource allocation fails + pub fn default_hft(device: &Device) -> Result { + let config = Mamba2Config { + d_model: 256, + d_state: 64, // P0 FIX: Mamba-2 official recommendation (was 32) + d_head: 32, + num_heads: 8, + expand: 2, + num_layers: 4, + target_latency_us: 3, + hardware_aware: true, + use_ssd: true, + use_selective_state: true, + max_seq_len: 1024, + batch_size: 16, + seq_len: 256, + ..Default::default() + }; + + Self::new(config, device) + } + + /// Forward pass through the model + /// + /// # Errors + /// + /// Returns `MLError` if: + /// - Input projection fails + /// - Layer normalization fails + /// - SSD layer processing fails + /// - Output projection fails + /// - Tensor operations fail + #[instrument(skip(self, input))] + pub fn forward(&mut self, input: &Tensor) -> Result { + let input = ml_core::mixed_precision::ensure_training_dtype(input) + .map_err(|e| MLError::ModelError(e.to_string()))?; + let start = Instant::now(); + + // OPTIMIZATION: Device affinity check (catch cross-device transfers early) + // Compare device types (CUDA vs CPU) since Device doesn't implement PartialEq + if input.device().is_cuda() != self.device.is_cuda() { + return Err(MLError::ModelError(format!( + "Input tensor on wrong device: expected {:?}, got {:?}", + self.device, + input.device() + ))); + } + + // Input projection + let mut hidden = self.input_projection.forward(&input)?; + + // Process through each layer - collect indices first to avoid borrow conflicts + let num_layers = self.ssd_layers.len(); + for layer_idx in 0..num_layers { + // Layer normalization + let normalized = self.layer_norms[layer_idx].forward(&hidden)?; + + // OPTIMIZATION: SSD layer processing - clone ssd_layer to avoid borrow conflict + // The forward_ssd_layer method requires &mut self, so we must clone the layer + let ssd_layer = self.ssd_layers[layer_idx].clone(); + let layer_output = self.forward_ssd_layer(&ssd_layer, &normalized, layer_idx)?; + + // Residual connection + hidden = (&hidden + &layer_output)?; + + // Dropout + if self.config.dropout > 0.0 { + hidden = self.dropouts[layer_idx].forward(&hidden, true)?; + } + } + + // Output projection with sigmoid activation (P0 FIX: bound output to [0,1] for normalized targets) + let output_raw = self.output_projection.forward(&hidden)?; + let output = ml_core::cuda_compat::manual_sigmoid(&output_raw)?; + + // Cast output back to F32 for API compatibility + let output = output.to_dtype(candle_core::DType::F32) + .map_err(|e| MLError::ModelError(format!("Output dtype cast failed: {}", e)))?; + + // OPTIMIZATION: Update performance metrics with VecDeque (O(1) instead of O(n)) + let inference_time = start.elapsed(); + self.total_inferences.fetch_add(1, Ordering::Relaxed); + self.latency_histogram.push_back(inference_time); + + if self.latency_histogram.len() > 10000 { + self.latency_histogram.pop_front(); // O(1) operation instead of O(n) remove(0) + } + + Ok(output) + } + + /// Forward pass through SSD layer with selective scan + #[instrument(skip(self, _ssd_layer, input))] + fn forward_ssd_layer( + &mut self, + _ssd_layer: &SSDLayer, + input: &Tensor, + layer_idx: usize, + ) -> Result { + trace!( + "forward_ssd_layer layer {}: input shape={:?}", + layer_idx, + input.dims() + ); + + // Use references to avoid unnecessary clones (Agent MAMBA-MEMORY-FIX) + let dt = &self.state.ssm_states[layer_idx].delta; + let A = &self.state.ssm_states[layer_idx].A; + let B = &self.state.ssm_states[layer_idx].B; + let C = &self.state.ssm_states[layer_idx].C; + + trace!( + "forward_ssd_layer layer {}: B shape={:?}", + layer_idx, + B.dims() + ); + + // Discretize the continuous-time SSM + let A_discrete = self.discretize_ssm(A, dt)?; + let B_discrete = self.discretize_ssm_input(B, dt)?; + trace!( + "forward_ssd_layer layer {}: B_discrete shape={:?}", + layer_idx, + B_discrete.dims() + ); + + // Selective scan algorithm + let scan_input = self.prepare_scan_input(input, &A_discrete, &B_discrete)?; + trace!( + "scan_input shape: {:?}, B shape: {:?}, C shape: {:?}", + scan_input.dims(), + B.dims(), + C.dims() + ); + let scanned_states = self + .scan_engine + .parallel_prefix_scan(&scan_input, ScanOperator::SSMScan)?; + trace!("scanned_states shape: {:?}", scanned_states.dims()); + + // Apply output transformation + trace!( + "About to matmul: scanned_states {:?} × C.t() (C is {:?})", + scanned_states.dims(), + C.dims() + ); + let batch_size = scanned_states.dim(0)?; + // Cast C to match scanned_states dtype (SSM state is F32 but computation may be BF16) + let C_cast = C.to_dtype(scanned_states.dtype())?; + let C_t = C_cast.t()?.contiguous()?; + let C_broadcasted = + C_t.unsqueeze(0)? + .broadcast_as((batch_size, C_t.dim(0)?, C_t.dim(1)?))?; + let output = scanned_states.matmul(&C_broadcasted)?; + + // Update hidden state + let _batch_size = input.dim(0)?; + let seq_len = input.dim(1)?; + if seq_len > 0 { + let last_state = scanned_states.narrow(1, seq_len - 1, 1)?.squeeze(1)?; + self.state.ssm_states[layer_idx].hidden = last_state; + } + + Ok(output) + } + + /// Discretize continuous-time SSM matrix A + /// + /// Mathematical notation: A_cont follows standard SSM notation for continuous-time state transition matrix + #[allow(non_snake_case)] + #[allow(non_snake_case)] + fn discretize_ssm(&self, A_cont: &Tensor, dt: &Tensor) -> Result { + // Keep dt on GPU — mean_all() returns a 0-D tensor, no CPU round-trip + let dt_scalar = dt.mean_all()?; + + // Bilinear (Tustin) approximation: A_disc = I + A*dt + (A*dt)^2 / 2 + // More accurate than ZOH (I + A*dt), matches ssd_layer.rs + let A_dt = A_cont.broadcast_mul(&dt_scalar)?; + let A_dt_sq = A_dt.matmul(&A_dt)?; + let half = Tensor::new(0.5_f32, A_cont.device())?; + let second_order = A_dt_sq.broadcast_mul(&half)?; + + let identity = Tensor::eye(A_cont.dim(0)?, DType::F32, A_cont.device())?; + let A_discrete = ((&identity + &A_dt)? + &second_order)?; + + Ok(A_discrete) + } + + /// Discretize continuous-time input matrix B + /// + /// Mathematical notation: B_cont follows standard SSM notation for continuous-time input matrix + #[allow(non_snake_case)] + #[allow(non_snake_case)] + fn discretize_ssm_input(&self, B_cont: &Tensor, dt: &Tensor) -> Result { + // Keep dt on GPU — mean_all() returns a 0-D tensor, no CPU round-trip + let dt_scalar = dt.mean_all()?; + let B_discrete = B_cont.broadcast_mul(&dt_scalar)?; + + Ok(B_discrete) + } + + /// Prepare input for selective scan algorithm + /// + /// Mathematical notation: Parameters _A and B follow standard SSM notation + #[allow(non_snake_case)] + fn prepare_scan_input( + &self, + input: &Tensor, + _A: &Tensor, + B: &Tensor, + ) -> Result { + // Transpose B and broadcast to match batch dimension + // input: [batch, seq, d_inner], B: [d_state, d_inner] + // B.t(): [d_inner, d_state] → broadcast to [batch, d_inner, d_state] + let batch_size = input.dim(0)?; + trace!( + "prepare_scan_input: input shape: {:?}, B shape: {:?}", + input.dims(), + B.dims() + ); + trace!( + "prepare_scan_input: d_model: {}, d_inner: {}, d_state: {}", + self.config.d_model, + self.config.d_model * self.config.expand, + self.config.d_state + ); + + // Cast B to match input dtype (SSM state is F32 but input may be BF16) + let B_cast = B.to_dtype(input.dtype())?; + let B_t = B_cast.t()?.contiguous()?; + let d_inner = B_t.dim(0)?; + let d_state = B_t.dim(1)?; + let B_broadcasted = B_t + .unsqueeze(0)? + .broadcast_as((batch_size, d_inner, d_state))?; + trace!( + "prepare_scan_input: B broadcasted shape: {:?}", + B_broadcasted.dims() + ); + + let Bu = input.matmul(&B_broadcasted)?; + trace!( + "prepare_scan_input: Bu shape: {:?}, expected [batch={}, seq={}, d_state={}]", + Bu.dims(), + input.dim(0)?, + input.dim(1)?, + self.config.d_state + ); + + Ok(Bu) + } + /// Fast single prediction for HFT + /// + /// # Errors + /// + /// Returns `MLError` if: + /// - Tensor creation fails + /// - Forward pass fails + /// - Output extraction fails + /// - Value conversion fails + pub fn predict_single_fast(&mut self, input: &[f64]) -> Result { + let start = Instant::now(); + + if input.len() != self.config.d_model { + return Err(MLError::InvalidInput(format!( + "Expected input dimension {}, got {}", + self.config.d_model, + input.len() + ))); + } + + let device = self.device(); + // Convert f64 input to f32 for F32 model dtype + let input_f32: Vec = input.iter().map(|&v| v as f32).collect(); + let input_tensor = Tensor::from_vec(input_f32, (1, input.len()), device)?; + + let output = self.forward(&input_tensor)?; + // Model uses F32 tensors — extract as f32 then widen to f64 for API compat + let result: f32 = output.to_scalar()?; + + let elapsed = start.elapsed(); + if elapsed.as_micros() > self.config.target_latency_us as u128 { + warn!( + "Prediction exceeded target latency: {}μs", + elapsed.as_micros() + ); + } + + Ok(result as f64) + } + + /// Get performance metrics + pub fn get_performance_metrics(&self) -> HashMap { + let mut metrics = HashMap::new(); + + metrics.insert( + "total_inferences".to_owned(), + self.total_inferences.load(Ordering::Relaxed) as f64, + ); + metrics.insert( + "total_training_steps".to_owned(), + self.total_training_steps.load(Ordering::Relaxed) as f64, + ); + + if !self.latency_histogram.is_empty() { + let avg_latency = self + .latency_histogram + .iter() + .map(|d| d.as_micros() as f64) + .sum::() + / self.latency_histogram.len() as f64; + metrics.insert("avg_latency_us".to_owned(), avg_latency); + + let throughput = 1_000_000.0 / avg_latency; // predictions per second + metrics.insert("throughput_pps".to_owned(), throughput); + } + + // Hardware metrics + if let Some(hw_optimizer) = &self.hardware_optimizer { + let hw_metrics = hw_optimizer.get_performance_metrics(); + for (k, v) in hw_metrics { + metrics.insert(k, v); + } + } + + // Model-specific metrics + metrics.insert( + "model_parameters".to_owned(), + self.metadata.num_parameters as f64, + ); + metrics.insert( + "compression_ratio".to_owned(), + if self.state.selective_state.len() > 0 && !self.state.compression_indices.is_empty() { + self.state.compression_indices.len() as f64 + / self.state.selective_state.len() as f64 + } else { + 1.0 + }, + ); + + let latency_target_ratio = if !self.latency_histogram.is_empty() { + let avg_latency = self + .latency_histogram + .iter() + .map(|d| d.as_micros() as f64) + .sum::() + / self.latency_histogram.len() as f64; + avg_latency / self.config.target_latency_us as f64 + } else { + 0.0 + }; + metrics.insert("latency_target_ratio".to_owned(), latency_target_ratio); + + // Additional production metrics for compatibility + metrics.insert("cache_hit_rate".to_owned(), 0.95); + metrics.insert("simd_ops_per_inference".to_owned(), 1000.0); + metrics.insert( + "state_compression_ratio".to_owned(), + metrics.get("compression_ratio").copied().unwrap_or(1.0), + ); + + metrics + } + + /// Get the device this model is on + fn device(&self) -> &Device { + &self.device + } + + /// Clear internal SSM state (call between epochs to prevent state accumulation) + /// + /// # Errors + /// + /// Returns `MLError` if: + /// - SSM state reset fails + /// - Tensor operations fail + pub fn clear_state(&mut self) -> Result<(), MLError> { + // Reset SSM state for each layer to prevent accumulation across epochs + for (layer_idx, ssm_state) in self.state.ssm_states.iter_mut().enumerate() { + ssm_state.reset()?; + trace!("Cleared MAMBA2 SSM state for layer {}", layer_idx); + } + + // Clear selective state components + self.state.selective_state.fill(0.0); + self.state.compression_indices.clear(); + + info!( + "Cleared MAMBA2 SSM state for all {} layers", + self.state.ssm_states.len() + ); + Ok(()) + } + + /// Check early stopping condition with patience (TFT pattern) + /// + /// This method implements the same early stopping logic as TFT + /// (see `ml/src/trainers/tft.rs:1702-1731`). + /// + /// # Arguments + /// + /// * `epoch` - Current epoch number (0-indexed) + /// * `val_loss` - Validation loss for current epoch + /// + /// # Returns + /// + /// `true` if training should stop early, `false` otherwise + /// + /// # Behavior + /// + /// 1. **Before min_epochs**: Always returns `false` (don't stop prematurely) + /// 2. **Improvement detected**: Resets patience counter, updates best_val_loss + /// 3. **No improvement**: Increments patience counter + /// 4. **Patience exhausted**: Sets stopped flag and returns `true` + pub fn check_early_stopping(&mut self, epoch: usize, val_loss: f64) -> bool { + // Don't stop before min_epochs + if epoch < self.config.early_stopping_min_epochs { + return false; + } + + // Check if validation loss improved by more than min_delta + if val_loss < self.state.best_val_loss - self.config.early_stopping_min_delta { + // Improvement detected - reset patience counter + self.state.best_val_loss = val_loss; + self.state.patience_counter = 0; + false + } else { + // No improvement - increment patience counter + self.state.patience_counter += 1; + + if self.state.patience_counter >= self.config.early_stopping_patience { + // Patience exhausted - trigger early stopping + self.state.stopped = true; + self.state.stopped_at_epoch = Some(epoch); + info!( + "Early stopping triggered at epoch {} (patience: {}, best val loss: {:.6})", + epoch, self.config.early_stopping_patience, self.state.best_val_loss + ); + true + } else { + debug!( + "Patience: {}/{} (best val loss: {:.6}, current: {:.6})", + self.state.patience_counter, + self.config.early_stopping_patience, + self.state.best_val_loss, + val_loss + ); + false + } + } + } + + /// Train the model with selective scan algorithm + #[instrument(skip(self, train_data, val_data, checkpoint_dir))] + pub async fn train( + &mut self, + train_data: &[(Tensor, Tensor)], + val_data: &[(Tensor, Tensor)], + epochs: usize, + checkpoint_dir: Option<&std::path::Path>, + ) -> Result, MLError> { + info!("Starting MAMBA-2 training with {} epochs", epochs); + + const MAX_TRAINING_HISTORY: usize = 100; + let mut training_history: VecDeque = VecDeque::with_capacity(MAX_TRAINING_HISTORY + 1); + let mut best_val_loss = f64::INFINITY; + + // FIXED (Agent P2): Set total training samples for accurate LR schedule + self.total_training_samples = train_data.len(); + + // Initialize optimizer + self.initialize_optimizer()?; + + for epoch in 0..epochs { + let epoch_start = Instant::now(); + + // FIXED: Do NOT clear SSM state (A, B, C parameters) - these are model weights + // that must persist across epochs to accumulate gradient updates. + // Clearing them was causing the E11 validation spike by reinitializing with random values. + + let mut epoch_loss = 0.0; + let mut batch_count = 0; + + // Create batch indices (shuffle if configured) + let mut batch_indices: Vec = (0..train_data.len()) + .step_by(self.config.batch_size) + .collect(); + + if self.config.shuffle_batches { + use rand::seq::SliceRandom; + batch_indices.shuffle(&mut rand::thread_rng()); + } + + // Training phase + for &batch_idx in &batch_indices { + let batch_end = (batch_idx + self.config.batch_size).min(train_data.len()); + let batch = &train_data[batch_idx..batch_end]; + + let batch_loss = self.train_batch(batch, epoch)?; + epoch_loss += batch_loss; + batch_count += 1; + + // Update learning rate + self.update_learning_rate(epoch, batch_idx)?; + + if batch_idx % 100 == 0 { + // FIXED (Agent P2): Log current learning rate for monitoring + let current_lr = self.get_current_learning_rate(); + debug!( + "Epoch {}, Batch {}: Loss = {:.6}, LR = {:.6}", + epoch, batch_idx, batch_loss, current_lr + ); + } + } + + epoch_loss /= batch_count as f64; + + // Validation phase + let val_loss = self.validate(val_data)?; + let epoch_accuracy = self.calculate_accuracy(val_data)?; + + // Update learning rate scheduler + let current_lr = self.get_current_learning_rate(); + + let epoch_duration = epoch_start.elapsed().as_secs_f64(); + let training_epoch = TrainingEpoch { + epoch, + loss: epoch_loss, + accuracy: epoch_accuracy, + learning_rate: current_lr, + duration_seconds: epoch_duration, + timestamp: SystemTime::now(), + }; + + training_history.push_back(training_epoch.clone()); + self.metadata.training_history.push_back(training_epoch); + + // Bound training history to prevent unbounded memory growth + while training_history.len() > MAX_TRAINING_HISTORY { + training_history.pop_front(); + } + while self.metadata.training_history.len() > MAX_TRAINING_HISTORY { + self.metadata.training_history.pop_front(); + } + + // Save checkpoint if best model + if val_loss < best_val_loss { + best_val_loss = val_loss; + + let checkpoint_path = if let Some(dir) = checkpoint_dir { + dir.join(format!("best_epoch_{}.ckpt", epoch)) + } else { + std::path::PathBuf::from(format!("best_epoch_{}.ckpt", epoch)) + }; + + let path_str = checkpoint_path.to_str().ok_or_else(|| { + MLError::ConfigError("Checkpoint path contains invalid UTF-8".to_owned()) + })?; + self.save_checkpoint(path_str).await?; + info!( + "New best validation loss: {:.6} at epoch {}", + val_loss, epoch + ); + } + + // Log epoch results + info!( + "Epoch {}/{}: Loss = {:.6}, Val Loss = {:.6}, Accuracy = {:.4}, LR = {:.2e}, Time = {:.2}s", + epoch + 1, epochs, epoch_loss, val_loss, epoch_accuracy, current_lr, epoch_duration + ); + + // Early stopping check + if self.config.early_stopping_enabled && self.check_early_stopping(epoch, val_loss) { + info!( + "Early stopping triggered at epoch {} (patience exhausted)", + epoch + ); + break; + } + } + + self.is_trained = true; + info!("Training completed with {} epochs", training_history.len()); + + Ok(training_history.into()) + } + + /// Train the model with async data loading (prefetch optimization) + /// + /// This method uses AsyncDataLoader to prefetch batches while GPU trains, + /// improving GPU utilization from ~78% to ~90-95% and reducing training + /// time by 20-30%. + /// + /// # Arguments + /// + /// * `train_data` - Training data as (feature, target) tensor pairs + /// * `val_data` - Validation data + /// * `epochs` - Number of training epochs + /// * `batch_size` - Batch size for training + /// * `prefetch_count` - Number of batches to prefetch (2-3 recommended) + /// + /// # Returns + /// + /// Training history with metrics per epoch + /// + /// # Errors + /// + /// Returns `MLError` if training fails + #[instrument(skip(self, train_data, val_data, checkpoint_dir))] + pub async fn train_async( + &mut self, + train_data: &[(Tensor, Tensor)], + val_data: &[(Tensor, Tensor)], + epochs: usize, + batch_size: usize, + prefetch_count: usize, + checkpoint_dir: Option<&std::path::Path>, + ) -> Result, MLError> { + info!( + "Starting MAMBA-2 async training with {} epochs (prefetch={})", + epochs, prefetch_count + ); + + const MAX_TRAINING_HISTORY: usize = 100; + let mut training_history: VecDeque = VecDeque::with_capacity(MAX_TRAINING_HISTORY + 1); + let mut best_val_loss = f64::INFINITY; + + // Set total training samples for accurate LR schedule + self.total_training_samples = train_data.len(); + + // Initialize optimizer + self.initialize_optimizer()?; + + for epoch in 0..epochs { + let epoch_start = Instant::now(); + + let mut epoch_loss = 0.0; + let mut batch_count = 0; + + // Batch iteration over pre-loaded training data + // Data loading from disk is handled upstream by StreamingDbnLoader + let mut batch_idx = 0; + for chunk in train_data.chunks(batch_size) { + let (inputs, targets): (Vec<_>, Vec<_>) = chunk.iter().cloned().unzip(); + let batched_input = Tensor::stack(&inputs, 0)?; + let batched_target = Tensor::stack(&targets, 0)?; + // Zero gradients + self.zero_gradients()?; + + // Forward pass with selective scan on batched input + let output = self.forward_with_gradients(&batched_input)?; + + // Extract last timestep for next-step prediction + // output: [batch, seq_len, d_model] → [batch, 1, d_model] + let seq_len = output.dim(1)?; + let output_last = output.narrow(1, seq_len - 1, 1)?; + + // Compute loss on last timestep prediction + let loss = self.compute_loss(&output_last, &batched_target)?; + + // Backward pass - compute gradients for SSM parameters + self.backward_pass(&loss, &batched_input, &batched_target)?; + + // Extract scalar AFTER backward to avoid stalling GPU pipeline + let loss_value = loss.to_scalar::()? as f64; + + // Update parameters + self.optimizer_step()?; + + epoch_loss += loss_value; + batch_count += 1; + + // Update learning rate + self.update_learning_rate(epoch, batch_idx)?; + + if batch_idx % 100 == 0 { + let current_lr = self.get_current_learning_rate(); + debug!( + "Epoch {}, Batch {}: Loss = {:.6}, LR = {:.6}", + epoch, batch_idx, loss_value, current_lr + ); + } + + batch_idx += batch_size; + } + + epoch_loss /= batch_count as f64; + + // Validation phase + let val_loss = self.validate(val_data)?; + let epoch_accuracy = self.calculate_accuracy(val_data)?; + + // Update learning rate scheduler + let current_lr = self.get_current_learning_rate(); + + let epoch_duration = epoch_start.elapsed().as_secs_f64(); + let training_epoch = TrainingEpoch { + epoch, + loss: epoch_loss, + accuracy: epoch_accuracy, + learning_rate: current_lr, + duration_seconds: epoch_duration, + timestamp: SystemTime::now(), + }; + + training_history.push_back(training_epoch.clone()); + self.metadata.training_history.push_back(training_epoch); + + // Bound training history to prevent unbounded memory growth + while training_history.len() > MAX_TRAINING_HISTORY { + training_history.pop_front(); + } + while self.metadata.training_history.len() > MAX_TRAINING_HISTORY { + self.metadata.training_history.pop_front(); + } + + // Save checkpoint if best model + if val_loss < best_val_loss { + best_val_loss = val_loss; + + let checkpoint_path = if let Some(dir) = checkpoint_dir { + dir.join(format!("best_epoch_{}.ckpt", epoch)) + } else { + std::path::PathBuf::from(format!("best_epoch_{}.ckpt", epoch)) + }; + + let path_str = checkpoint_path.to_str().ok_or_else(|| { + MLError::ConfigError("Checkpoint path contains invalid UTF-8".to_owned()) + })?; + self.save_checkpoint(path_str).await?; + info!( + "New best validation loss: {:.6} at epoch {}", + val_loss, epoch + ); + } + + // Log epoch results + info!( + "Epoch {}/{}: Loss = {:.6}, Val Loss = {:.6}, Accuracy = {:.4}, LR = {:.2e}, Time = {:.2}s", + epoch + 1, epochs, epoch_loss, val_loss, epoch_accuracy, current_lr, epoch_duration + ); + + // Early stopping check + if self.config.early_stopping_enabled && self.check_early_stopping(epoch, val_loss) { + info!( + "Early stopping triggered at epoch {} (patience exhausted)", + epoch + ); + break; + } + } + + self.is_trained = true; + info!( + "Async training completed with {} epochs", + training_history.len() + ); + + Ok(training_history.into()) + } + + /// Train a single batch with selective scan + #[instrument(skip(self, batch))] + fn train_batch(&mut self, batch: &[(Tensor, Tensor)], _epoch: usize) -> Result { + if batch.is_empty() { + return Ok(0.0); + } + + // FIXED: Batch all individual sequences together into a single batched tensor + // Individual sequences are shape [1, seq_len, d_model], we need [batch_size, seq_len, d_model] + + let actual_batch_size = batch.len(); + + // Collect all input tensors and concatenate along batch dimension + let input_tensors: Vec<&Tensor> = batch.iter().map(|(input, _)| input).collect(); + let batched_input = if actual_batch_size == 1 { + // Single sample - no concatenation needed + input_tensors[0].clone() + } else { + // Concatenate along dimension 0 (batch dimension) + Tensor::cat( + &input_tensors + .iter() + .map(|t| (*t).clone()) + .collect::>(), + 0, + )? + }; + + // Collect all target tensors and concatenate + let target_tensors: Vec<&Tensor> = batch.iter().map(|(_, target)| target).collect(); + let batched_target = if actual_batch_size == 1 { + target_tensors[0].clone() + } else { + Tensor::cat( + &target_tensors + .iter() + .map(|t| (*t).clone()) + .collect::>(), + 0, + )? + }; + + // FIXED: Ensure input and target tensors are on the model's device (GPU) + // This prevents device mismatch errors during forward pass + let batched_input = batched_input.to_device(&self.device)?; + let batched_target = batched_target.to_device(&self.device)?; + + // Zero gradients + self.zero_gradients()?; + + // Forward pass with selective scan on batched input + let output = self.forward_with_gradients(&batched_input)?; + trace!( + "Training loop: batched_input: {:?}, batched_target: {:?}, forward output: {:?}", + batched_input.dims(), + batched_target.dims(), + output.dims() + ); + + // FIXED (Agent 211): Extract last timestep for next-step prediction + // output: [batch, seq_len, d_model] → [batch, 1, d_model] + // This matches target shape [batch, 1, d_model] + let seq_len = output.dim(1)?; + let output_last = output.narrow(1, seq_len - 1, 1)?; + trace!( + "Training loop: output_last (for loss): {:?}", + output_last.dims() + ); + + // Compute loss on last timestep prediction + let loss = self.compute_loss(&output_last, &batched_target)?; + + // Backward pass - compute gradients for SSM parameters + self.backward_pass(&loss, &batched_input, &batched_target)?; + + // Extract scalar AFTER backward to avoid stalling GPU pipeline + let loss_value = loss.to_scalar::()? as f64; + + // Update parameters + self.optimizer_step()?; + + // Update selective state based on gradients (use first sample for importance scoring) + if let Some(selective_state) = &mut self.selective_state { + // Use the first sample in the batch for importance updates + let first_input = input_tensors[0]; + selective_state.update_importance_scores(first_input, &mut self.state)?; + } + + self.total_training_steps.fetch_add(1, Ordering::Relaxed); + self.step_count += 1; + + // Explicit memory cleanup to prevent GPU memory accumulation + drop(output); + drop(output_last); + drop(loss); + drop(batched_input); + drop(batched_target); + + Ok(loss_value) + } + + /// Forward pass with gradient computation enabled + pub fn forward_with_gradients(&mut self, input: &Tensor) -> Result { + // Gradient flow enabled - do not detach + let input = input; + + // Input projection with gradients + let mut hidden = self.input_projection.forward(input)?; + + // Process through each layer with SSM gradients - collect indices first to avoid borrow conflicts + let num_layers = self.ssd_layers.len(); + for layer_idx in 0..num_layers { + // Layer normalization + let normalized = self.layer_norms[layer_idx].forward(&hidden)?; + + // SSD layer processing with selective scan and gradients + let layer_output = { + let ssd_layer = self.ssd_layers[layer_idx].clone(); + self.forward_ssd_layer_with_gradients(&ssd_layer, &normalized, layer_idx)? + }; + + // Residual connection + hidden = (&hidden + &layer_output)?; + + // Dropout (enabled during training) + if self.config.dropout > 0.0 { + hidden = self.dropouts[layer_idx].forward(&hidden, true)?; + } + } + + // Output projection + trace!( + "Before output_projection: hidden shape: {:?}", + hidden.dims() + ); + // P0 FIX: Add sigmoid activation to bound output to [0,1] for normalized targets + let output_raw = self.output_projection.forward(&hidden)?; + let output = ml_core::cuda_compat::manual_sigmoid(&output_raw)?; + trace!( + "After output_projection + sigmoid: output shape: {:?}", + output.dims() + ); + + Ok(output) + } + + /// Forward pass through SSD layer with gradient tracking + fn forward_ssd_layer_with_gradients( + &mut self, + _ssd_layer: &SSDLayer, + input: &Tensor, + layer_idx: usize, + ) -> Result { + // Use references to avoid unnecessary clones (Agent MAMBA-MEMORY-FIX) + let dt = &self.state.ssm_states[layer_idx].delta; + let A = &self.state.ssm_states[layer_idx].A; + let B = &self.state.ssm_states[layer_idx].B; + let C = &self.state.ssm_states[layer_idx].C; + + // Discretize with gradient tracking + let A_discrete = self.discretize_ssm_with_gradients(A, dt)?; + let B_discrete = self.discretize_ssm_input_with_gradients(B, dt)?; + + // Selective scan with gradient computation + let scan_input = self.prepare_scan_input_with_gradients(input, &A_discrete, &B_discrete)?; + let scanned_states = self.selective_scan_with_gradients(&scan_input, &A_discrete)?; + + // Output transformation with gradients + // FIXED (Agent 207): Broadcast C correctly after transpose + let batch_size = scanned_states.dim(0)?; + trace!("C matrix broadcast: scanned_states shape: {:?}, C original shape (d_inner, d_state): {:?}", scanned_states.dims(), C.dims()); + + // For matmul: [batch, seq, d_state] × [batch, d_state, d_inner] = [batch, seq, d_inner] + // scanned_states: [32, 60, 16] + // C stored as: [d_inner, d_state] = [512, 16] + // Need: [batch, d_state, d_inner] = [32, 16, 512] + // Cast C to match scanned_states dtype (SSM state is F32 but computation may be BF16) + let C_cast = C.to_dtype(scanned_states.dtype())?; + let C_t = C_cast.t()?.contiguous()?; // [512, 16] → [16, 512] + trace!("C transposed (d_state, d_inner): {:?}", C_t.dims()); + + // Now broadcast [16, 512] to [32, 16, 512] + let d_state = C_t.dim(0)?; // 16 + let d_inner = C_t.dim(1)?; // 512 + let C_broadcasted = C_t + .unsqueeze(0)? + .broadcast_as((batch_size, d_state, d_inner))?; + trace!( + "C broadcasted shape: {:?}, expected: [batch={}, d_state={}, d_inner={}]", + C_broadcasted.dims(), + batch_size, + d_state, + d_inner + ); + + let output = scanned_states.matmul(&C_broadcasted)?; + trace!("Output shape: {:?}", output.dims()); + + // Update hidden state + let _batch_size = input.dim(0)?; + let seq_len = input.dim(1)?; + if seq_len > 0 { + let last_state = scanned_states.narrow(1, seq_len - 1, 1)?.squeeze(1)?; + self.state.ssm_states[layer_idx].hidden = last_state; + } + + Ok(output) + } + + /// Selective scan algorithm with gradient computation + /// + /// Mathematical notation: Parameter A represents the state transition matrix + #[allow(non_snake_case)] + fn selective_scan_with_gradients(&self, input: &Tensor, A: &Tensor) -> Result { + let seq_len = input.dim(1)?; + let d_state = input.dim(2)?; + let device = input.device(); + + // AGENT 176 FIX: Add shape assertions to catch dimension bugs early + tracing::debug!( + "[AGENT 176] selective_scan_with_gradients: input={:?}, A={:?}", + input.dims(), + A.dims() + ); + assert_eq!(input.dims().len(), 3, "Input must be [batch, seq, d_state]"); + assert_eq!(A.dims().len(), 2, "A must be [d_state, d_state]"); + assert_eq!( + A.dim(0)?, + d_state, + "A.dim(0) must equal input.dim(2) (d_state)" + ); + + // Initialize state sequence - pre-allocate result tensor to avoid Vec accumulation + // This prevents the 750MB memory leak from accumulating 60 tensors in Vec + let batch_size = input.dim(0)?; + let mut result = Tensor::zeros((batch_size, seq_len, d_state), input.dtype(), device)?; + let mut current_state = Tensor::zeros((batch_size, d_state), input.dtype(), device)?; + + // Cast A to match input dtype (SSM state matrices are F32 but computation may be BF16) + let A_cast = A.to_dtype(input.dtype())?; + + // Sequential scan with state transitions (maintaining gradients) + for t in 0..seq_len { + let x_t = input.narrow(1, t, 1)?.squeeze(1)?; + + // AGENT 176 FIX: Correct batch matrix multiplication + // State transition: h_t = h_{t-1} @ A^T + x_t + // current_state [batch, d_state] × A.t() [d_state, d_state] = [batch, d_state] + // This is the correct way to do batch SSM state transitions + current_state = (current_state.matmul(&A_cast.t()?)? + &x_t)?; + + // Write directly to result tensor (no Vec accumulation, no Tensor::cat doubling) + let current_unsqueezed = current_state.unsqueeze(1)?; + result = result.slice_assign( + &[0..batch_size, t..(t + 1), 0..d_state], + ¤t_unsqueezed, + )?; + } + + // AGENT 176 FIX: Verify output shape matches expected dimensions + tracing::debug!( + "[AGENT 176] selective_scan_with_gradients: output={:?}", + result.dims() + ); + assert_eq!( + result.dims(), + &[input.dim(0)?, seq_len, d_state], + "Output must be [batch, seq, d_state], got {:?}", + result.dims() + ); + + Ok(result) + } + + /// Discretize SSM with gradient tracking + /// + /// Mathematical notation: A_cont follows standard SSM notation for continuous-time state transition matrix + #[allow(non_snake_case)] + #[allow(non_snake_case)] + fn discretize_ssm_with_gradients( + &self, + A_cont: &Tensor, + dt: &Tensor, + ) -> Result { + // Keep dt on GPU — mean_all() returns a 0-D tensor, no CPU round-trip + let dt_scalar = dt.mean_all()?; + + // Scale A matrix by dt + let A_scaled = A_cont.broadcast_mul(&dt_scalar)?; + + // Matrix exponential approximation: exp(A) ≈ I + A + A²/2 + A³/6 + let identity = Tensor::eye(A_cont.dim(0)?, DType::F32, A_cont.device())?; + let A2 = A_scaled.matmul(&A_scaled)?; + let A3 = A2.matmul(&A_scaled)?; + + let half = Tensor::new(0.5_f32, A_cont.device())?; + let sixth = Tensor::new(1.0_f32 / 6.0, A_cont.device())?; + let A_discrete = (&identity + + &A_scaled + + &A2.broadcast_mul(&half)? + + &A3.broadcast_mul(&sixth)?)?; + + Ok(A_discrete) + } + + /// Discretize input matrix with gradients + /// + /// Mathematical notation: B_cont follows standard SSM notation for continuous-time input matrix + #[allow(non_snake_case)] + #[allow(non_snake_case)] + fn discretize_ssm_input_with_gradients( + &self, + B_cont: &Tensor, + dt: &Tensor, + ) -> Result { + // Keep dt on GPU — mean_all() returns a 0-D tensor, no CPU round-trip + let dt_scalar = dt.mean_all()?; + let B_discrete = B_cont.broadcast_mul(&dt_scalar)?; + Ok(B_discrete) + } + + /// Prepare scan input with gradient tracking + /// + /// Mathematical notation: Parameters _A and B follow standard SSM notation + #[allow(non_snake_case)] + fn prepare_scan_input_with_gradients( + &self, + input: &Tensor, + _A: &Tensor, + B: &Tensor, + ) -> Result { + // FIXED (Agent 248 + Agent 250): Explicit batch broadcast for B matrix + // input: [batch, seq, d_inner], B: [d_state, d_inner] + // B.t(): [d_inner, d_state] → explicit repeat to [batch, d_inner, d_state] + let batch_size = input.dim(0)?; + // Cast B to match input dtype (SSM state is F32 but input may be BF16) + let B_cast = B.to_dtype(input.dtype())?; + let B_t = B_cast.t()?.contiguous()?; // [d_state, d_inner] → [d_inner, d_state] + + // CRITICAL FIX: Use repeat/expand instead of broadcast_as for CUDA compatibility + // Create [batch, d_inner, d_state] by repeating the [d_inner, d_state] tensor + let B_expanded = B_t.unsqueeze(0)?; // [1, d_inner, d_state] + + // Repeat along batch dimension + let B_broadcasted = B_expanded.expand(&[batch_size, B_t.dim(0)?, B_t.dim(1)?])?; + + trace!( + "[Agent 250] B matrix broadcast: B_t={:?} → B_broadcasted={:?}", + B_t.dims(), + B_broadcasted.dims() + ); + + let Bu = input.matmul(&B_broadcasted)?; + trace!("[Agent 250] Bu result shape: {:?}", Bu.dims()); + Ok(Bu) + } + + /// Compute training loss + pub fn compute_loss(&self, output: &Tensor, target: &Tensor) -> Result { + // Mean Squared Error for regression + let diff = (output - target)?; + let squared_diff = (&diff * &diff)?; + let loss = squared_diff.mean_all()?; + // loss is F64 from mean_all() + Ok(loss) + } + + /// Backward pass - compute gradients for model parameters + pub fn backward_pass( + &mut self, + loss: &Tensor, + _input: &Tensor, + _target: &Tensor, + ) -> Result<(), MLError> { + // Compute gradients using automatic differentiation + // The loss tensor should already have the computational graph attached + let grads = loss.backward()?; + + // FIXED (P0): Extract REAL gradients from VarMap trainable parameters + // The trainable parameters are: input_projection, output_projection, layer_norms (weight/bias) + // SSM matrices (A, B, C, delta) are NOT trainable - they're part of the model state + trace!("[P0 FIX] Extracting real gradients from VarMap trainable parameters"); + + self.gradients.clear(); + + // Extract gradients from all VarMap parameters + let all_vars = self.varmap.all_vars(); + let mut total_grad_norm = 0.0_f64; + let mut params_with_grads = 0; + + for (idx, var) in all_vars.iter().enumerate() { + if let Some(grad) = grads.get(var) { + // Compute gradient norm for monitoring + let grad_vec = grad + .flatten_all() + .map_err(|e| MLError::TensorCreationError { + operation: "gradient flatten".to_owned(), + reason: e.to_string(), + })? + .to_vec1::() + .map_err(|e| MLError::TensorCreationError { + operation: "gradient to_vec1".to_owned(), + reason: e.to_string(), + })?; + + let grad_norm: f64 = grad_vec.iter().map(|&g| (g as f64).powi(2)).sum::().sqrt(); + + // Store gradient with descriptive key + let key = format!("varmap_param_{}", idx); + self.gradients.insert(key.clone(), grad.clone()); + + if grad_norm > 1e-12 { + params_with_grads += 1; + total_grad_norm += grad_norm; + } + + trace!("[P0 FIX] VarMap param {}: grad_norm={:.6}", idx, grad_norm); + } else { + trace!( + "[P0 FIX] VarMap param {} has no gradient (not in computational graph)", + idx + ); + } + } + + trace!( + "[P0 FIX] Extracted {} gradients from {} VarMap parameters, total_grad_norm={:.6}", + params_with_grads, + all_vars.len(), + total_grad_norm + ); + + // Verify we got non-zero gradients + if total_grad_norm < 1e-12 { + return Err(MLError::TrainingError(format!( + "Zero gradients extracted from VarMap (total_grad_norm={:.6}). \ + This indicates the loss is not connected to trainable parameters.", + total_grad_norm + ))); + } + + self.clip_gradients(self.config.grad_clip)?; + + // Gradients flow through the trainable VarMap parameters: + // 1. input_projection: Projects d_model → d_inner + // 2. output_projection: Projects d_inner → 1 (regression) + // 3. layer_norms: Normalization weights/biases for each layer + // + // SSM matrices (A, B, C, delta) are NOT trainable in standard MAMBA-2. + // They are part of the model state and are used for selective state-space computation. + + Ok(()) + } + + /// Initialize optimizer state + pub fn initialize_optimizer(&mut self) -> Result<(), MLError> { + // Initialize Adam optimizer state + self.optimizer_state.clear(); + + // Add momentum and variance terms for each parameter + // In real implementation, this would be handled by candle's optimizers + + Ok(()) + } + + /// Zero gradients + pub fn zero_gradients(&mut self) -> Result<(), MLError> { + // Clear all gradients for SSM parameters + for _ssm_state in &mut self.state.ssm_states { + // Zero gradients for A, B, C matrices and delta parameter + if let Some(grad) = self.gradients.get("A").cloned() { + self.gradients.insert("A".to_owned(), grad.zeros_like()?); + } + if let Some(grad) = self.gradients.get("B").cloned() { + self.gradients.insert("B".to_owned(), grad.zeros_like()?); + } + if let Some(grad) = self.gradients.get("C").cloned() { + self.gradients.insert("C".to_owned(), grad.zeros_like()?); + } + if let Some(grad) = self.gradients.get("delta").cloned() { + self.gradients + .insert("delta".to_owned(), grad.zeros_like()?); + } + } + + // Clear optimizer state gradients if they exist + for (param_name, tensor) in self.optimizer_state.iter_mut() { + if param_name.contains("grad") { + *tensor = tensor.zeros_like()?; + } + } + + Ok(()) + } + + /// Optimizer step - dispatches to Adam or SGD based on config + pub fn optimizer_step(&mut self) -> Result<(), MLError> { + match self.config.optimizer_type { + OptimizerType::Adam => self.optimizer_step_adam(), + OptimizerType::AdamW => self.optimizer_step_adam(), + OptimizerType::SGD => self.optimizer_step_sgd(), + } + } + + /// Adam optimizer step implementation + fn optimizer_step_adam(&mut self) -> Result<(), MLError> { + // P0: Use beta1 from config for hyperparameter optimization + let beta1: f64 = self.config.adam_beta1; + let beta2: f64 = 0.999; + let eps: f64 = 1e-8; // Standard epsilon for Adam optimizer + let lr = self.config.learning_rate; + + // Increment step counter for bias correction + let step = self + .optimizer_state + .get("step") + .and_then(|t| t.to_scalar::().ok()) + .unwrap_or(0.0) as f64 + + 1.0; + + let device = self.device(); + let step_tensor = Tensor::new(&[step as f32], device)?; + self.optimizer_state.insert("step".to_owned(), step_tensor); + + // Bias correction uses Rust-side f64 for precision + let beta1_t = beta1.powf(step); + let beta2_t = beta2.powf(step); + let bias_correction1 = 1.0 - beta1_t; + let bias_correction2 = 1.0 - beta2_t; + + // Apply Adam updates to all SSM parameters per layer + let num_layers = self.state.ssm_states.len(); + for layer_idx in 0..num_layers { + // Collect layer-specific gradients + let a_grad = self.gradients.get(&format!("A_{}", layer_idx)).cloned(); + let b_grad = self.gradients.get(&format!("B_{}", layer_idx)).cloned(); + let c_grad = self.gradients.get(&format!("C_{}", layer_idx)).cloned(); + let delta_grad = self.gradients.get(&format!("delta_{}", layer_idx)).cloned(); + + // Update A matrix (state transition matrix) + if let Some(ref A_grad) = a_grad { + trace!("[Agent 225] Updating A matrix for layer {}", layer_idx); + let mut A_param = self.state.ssm_states[layer_idx].A.clone(); + self.apply_adam_update( + &mut A_param, + A_grad, + layer_idx, + "A", + lr, + beta1, + beta2, + eps, + bias_correction1, + bias_correction2, + false, // No weight decay for A matrix (maintains stability) + )?; + self.state.ssm_states[layer_idx].A = A_param; + } + + // Update B matrix (input matrix) + if let Some(ref B_grad) = b_grad { + trace!("[Agent 225] Updating B matrix for layer {}", layer_idx); + let mut B_param = self.state.ssm_states[layer_idx].B.clone(); + self.apply_adam_update( + &mut B_param, + B_grad, + layer_idx, + "B", + lr, + beta1, + beta2, + eps, + bias_correction1, + bias_correction2, + true, // Apply weight decay to B matrix + )?; + self.state.ssm_states[layer_idx].B = B_param; + } + + // Update C matrix (output matrix) + if let Some(ref C_grad) = c_grad { + let mut C_param = self.state.ssm_states[layer_idx].C.clone(); + self.apply_adam_update( + &mut C_param, + C_grad, + layer_idx, + "C", + lr, + beta1, + beta2, + eps, + bias_correction1, + bias_correction2, + true, // Apply weight decay to C matrix + )?; + self.state.ssm_states[layer_idx].C = C_param; + } + + // Update Delta parameter (discretization parameter) + if let Some(ref delta_grad) = delta_grad { + let mut delta_param = self.state.ssm_states[layer_idx].delta.clone(); + self.apply_adam_update( + &mut delta_param, + delta_grad, + layer_idx, + "delta", + lr, + beta1, + beta2, + eps, + bias_correction1, + bias_correction2, + false, // No weight decay for Delta (maintains discretization stability) + )?; + self.state.ssm_states[layer_idx].delta = delta_param; + } + } + + // After updating A matrices, project to maintain spectral radius < 1 + self.project_ssm_matrices()?; + + Ok(()) + } + + /// AdamW optimizer step implementation with decoupled weight decay + /// + /// CRITICAL DIFFERENCE from Adam: + /// - Adam: weight_decay applied to gradients → interferes with SSM dynamics + /// - AdamW: weight_decay applied directly to parameters → preserves SSM constraints + /// + /// AdamW Formula: + /// 1. m_t = β1 * m_{t-1} + (1 - β1) * g_t + /// 2. v_t = β2 * v_{t-1} + (1 - β2) * g_t^2 + /// 3. m_hat = m_t / (1 - β1^t) + /// 4. v_hat = v_t / (1 - β2^t) + /// 5. θ_t = θ_{t-1} * (1 - λ * lr) - lr * m_hat / (√v_hat + ε) + /// ^^^^^^^^^^^^^^^^^^^^^^^^ ← DECOUPLED weight decay + /// + /// Where λ is weight_decay coefficient (independent of gradients) + fn optimizer_step_adamw(&mut self) -> Result<(), MLError> { + let beta1: f64 = self.config.adam_beta1; + let beta2: f64 = self.config.adam_beta2; + let eps: f64 = self.config.adam_epsilon; + let lr = self.config.learning_rate; + let wd = self.config.weight_decay; + + // Increment step counter for bias correction + let step = self + .optimizer_state + .get("step") + .and_then(|t| t.to_scalar::().ok()) + .unwrap_or(0.0) as f64 + + 1.0; + + let device = self.device(); + let step_tensor = Tensor::new(&[step as f32], device)?; + self.optimizer_state.insert("step".to_owned(), step_tensor); + + // Bias correction factors (Rust-side f64 for precision) + let beta1_t = beta1.powf(step); + let beta2_t = beta2.powf(step); + let bias_correction1 = 1.0 - beta1_t; + let bias_correction2 = 1.0 - beta2_t; + + // Apply AdamW updates to all SSM parameters per layer + let num_layers = self.state.ssm_states.len(); + for layer_idx in 0..num_layers { + // Collect layer-specific gradients + let a_grad = self.gradients.get(&format!("A_{}", layer_idx)).cloned(); + let b_grad = self.gradients.get(&format!("B_{}", layer_idx)).cloned(); + let c_grad = self.gradients.get(&format!("C_{}", layer_idx)).cloned(); + let delta_grad = self.gradients.get(&format!("delta_{}", layer_idx)).cloned(); + + // Update A matrix (state transition matrix) + if let Some(ref A_grad) = a_grad { + trace!("[AdamW] Updating A matrix for layer {}", layer_idx); + let mut A_param = self.state.ssm_states[layer_idx].A.clone(); + self.apply_adamw_update( + &mut A_param, + A_grad, + layer_idx, + "A", + lr, + beta1, + beta2, + eps, + bias_correction1, + bias_correction2, + 0.0, // No weight decay for A matrix (maintains stability) + )?; + self.state.ssm_states[layer_idx].A = A_param; + } + + // Update B matrix (input matrix) + if let Some(ref B_grad) = b_grad { + trace!("[AdamW] Updating B matrix for layer {}", layer_idx); + let mut B_param = self.state.ssm_states[layer_idx].B.clone(); + self.apply_adamw_update( + &mut B_param, + B_grad, + layer_idx, + "B", + lr, + beta1, + beta2, + eps, + bias_correction1, + bias_correction2, + wd, // Apply weight decay to B matrix + )?; + self.state.ssm_states[layer_idx].B = B_param; + } + + // Update C matrix (output matrix) + if let Some(ref C_grad) = c_grad { + let mut C_param = self.state.ssm_states[layer_idx].C.clone(); + self.apply_adamw_update( + &mut C_param, + C_grad, + layer_idx, + "C", + lr, + beta1, + beta2, + eps, + bias_correction1, + bias_correction2, + wd, // Apply weight decay to C matrix + )?; + self.state.ssm_states[layer_idx].C = C_param; + } + + // Update Delta parameter (discretization parameter) + if let Some(ref delta_grad) = delta_grad { + let mut delta_param = self.state.ssm_states[layer_idx].delta.clone(); + self.apply_adamw_update( + &mut delta_param, + delta_grad, + layer_idx, + "delta", + lr, + beta1, + beta2, + eps, + bias_correction1, + bias_correction2, + 0.0, // No weight decay for Delta (maintains discretization stability) + )?; + self.state.ssm_states[layer_idx].delta = delta_param; + } + } + + // After updating A matrices, project to maintain spectral radius < 1 + self.project_ssm_matrices()?; + + Ok(()) + } + + /// SGD optimizer step implementation with momentum + fn optimizer_step_sgd(&mut self) -> Result<(), MLError> { + let lr = self.config.learning_rate; + let momentum = self.config.sgd_momentum; + + // PRIORITY 2 FIX (Agent 225): Use layer-specific gradient keys + // Apply SGD updates to all SSM parameters per layer + let num_layers = self.state.ssm_states.len(); + for layer_idx in 0..num_layers { + // Collect layer-specific gradients + let a_grad = self.gradients.get(&format!("A_{}", layer_idx)).cloned(); + let b_grad = self.gradients.get(&format!("B_{}", layer_idx)).cloned(); + let c_grad = self.gradients.get(&format!("C_{}", layer_idx)).cloned(); + let delta_grad = self.gradients.get(&format!("delta_{}", layer_idx)).cloned(); + + // Update A matrix (state transition matrix) + if let Some(ref A_grad) = a_grad { + trace!("[SGD] Updating A matrix for layer {}", layer_idx); + let mut A_param = self.state.ssm_states[layer_idx].A.clone(); + self.apply_sgd_update( + &mut A_param, + A_grad, + layer_idx, + "A", + lr, + momentum, + false, // No weight decay for A matrix (maintains stability) + )?; + self.state.ssm_states[layer_idx].A = A_param; + } + + // Update B matrix (input matrix) + if let Some(ref B_grad) = b_grad { + trace!("[SGD] Updating B matrix for layer {}", layer_idx); + let mut B_param = self.state.ssm_states[layer_idx].B.clone(); + self.apply_sgd_update( + &mut B_param, + B_grad, + layer_idx, + "B", + lr, + momentum, + true, // Apply weight decay to B matrix + )?; + self.state.ssm_states[layer_idx].B = B_param; + } + + // Update C matrix (output matrix) + if let Some(ref C_grad) = c_grad { + let mut C_param = self.state.ssm_states[layer_idx].C.clone(); + self.apply_sgd_update( + &mut C_param, + C_grad, + layer_idx, + "C", + lr, + momentum, + true, // Apply weight decay to C matrix + )?; + self.state.ssm_states[layer_idx].C = C_param; + } + + // Update Delta parameter (discretization parameter) + if let Some(ref delta_grad) = delta_grad { + let mut delta_param = self.state.ssm_states[layer_idx].delta.clone(); + self.apply_sgd_update( + &mut delta_param, + delta_grad, + layer_idx, + "delta", + lr, + momentum, + false, // No weight decay for Delta (maintains discretization stability) + )?; + self.state.ssm_states[layer_idx].delta = delta_param; + } + } + + // After updating A matrices, project to maintain spectral radius < 1 + self.project_ssm_matrices()?; + + Ok(()) + } + + /// Update learning rate with warmup and decay + fn update_learning_rate(&mut self, epoch: usize, batch_idx: usize) -> Result<(), MLError> { + // FIXED (Agent P2): Use actual training data length instead of hardcoded 1000 + // Calculate total steps based on actual data size + let batches_per_epoch = if self.total_training_samples > 0 { + self.total_training_samples / self.config.batch_size + } else { + // Fallback to reasonable default if not set + 1000 / self.config.batch_size + }; + + let total_steps = epoch * batches_per_epoch + (batch_idx / self.config.batch_size); + + // FIXED (Agent P2): Remove underscore prefix - we DO use this value + let lr = if total_steps < self.config.warmup_steps { + // Linear warmup: LR increases from 0 to configured LR + self.config.learning_rate * (total_steps as f64 / self.config.warmup_steps as f64) + } else { + // Cosine decay after warmup (P0 FIX: use config value, not hardcoded) + let progress = (total_steps - self.config.warmup_steps) as f64; + let total_decay_steps = self.config.total_decay_steps as f64; + let decay_ratio = (progress / total_decay_steps).min(1.0); + self.config.learning_rate * 0.5 * (1.0 + (std::f64::consts::PI * decay_ratio).cos()) + }; + + // FIXED (Agent P2): Actually apply the computed learning rate + self.current_lr = lr; + + Ok(()) + } + + /// Get current learning rate + fn get_current_learning_rate(&self) -> f64 { + // FIXED (Agent P2): Return actual current LR, not constant config value + self.current_lr + } + + /// Validate model on validation set + fn validate(&mut self, val_data: &[(Tensor, Tensor)]) -> Result { + let mut total_loss = 0.0; + let mut count = 0; + + // Disable dropout for validation + for (input, target) in val_data { + // FIXED: Ensure input and target tensors are on the model's device (GPU) + let input = input.to_device(&self.device)?; + let target = target.to_device(&self.device)?; + + let output = self.forward(&input)?; + // FIXED (Agent 217): Extract last timestep for validation loss (same as training) + let seq_len = output.dim(1)?; + let output_last = output.narrow(1, seq_len - 1, 1)?; + let loss = self.compute_loss(&output_last, &target)?; + total_loss += loss.to_scalar::()? as f64; + count += 1; + + if count >= 100 { + // Limit validation set size for speed + break; + } + } + + Ok(total_loss / count as f64) + } + + /// Calculate accuracy metric + fn calculate_accuracy(&mut self, val_data: &[(Tensor, Tensor)]) -> Result { + if val_data.is_empty() { + return Ok(0.0); + } + + let mut correct = 0; + let mut total = 0; + + for (input, target) in val_data { + // CRITICAL FIX: Transfer tensors to device before forward pass (matches validate()) + let input = input.to_device(&self.device)?; + let target = target.to_device(&self.device)?; + + let output = self.forward(&input)?; + let seq_len = output.dims()[1]; + + // Extract last timestep predictions + let output_last = output.narrow(1, seq_len - 1, 1)?.squeeze(1)?; + + // FIX: Reshape target from [batch_size, 1, 1] to [batch_size] + // Previous double squeeze failed due to dimension index shifting + let batch_size = target.dim(0)?; + let target_squeezed = target.reshape(&[batch_size])?; + + // FIX: Use element-wise comparison instead of mean_all() + for i in 0..batch_size { + // FIX: .get(i) returns different shapes depending on input: + // - If input is [N], .get(i) returns scalar [] + // - If input is [N, 1], .get(i) returns [1] + // Check rank and squeeze conditionally + let pred_tensor = output_last.get(i)?; + let pred_value = if pred_tensor.rank() == 0 { + pred_tensor.to_scalar::()? as f64 + } else { + pred_tensor.squeeze(0)?.to_scalar::()? as f64 + }; + + let target_tensor = target_squeezed.get(i)?; + let target_value = if target_tensor.rank() == 0 { + target_tensor.to_scalar::()? as f64 + } else { + target_tensor.squeeze(0)?.to_scalar::()? as f64 + }; + + // FIX: Use absolute error (not MAPE) with 5% threshold + let abs_error = (pred_value - target_value).abs(); + + // 5% of [0,1] range = 0.05 (equivalent to ~$50 in ES price space) + if abs_error < 0.05 { + correct += 1; + } + total += 1; + } + + if total >= 100 { + break; + } + } + + Ok(correct as f64 / total as f64) + } + + /// Save model checkpoint + pub async fn save_checkpoint(&mut self, path: &str) -> Result<(), MLError> { + use std::collections::HashMap as StdHashMap; + + info!("Saving MAMBA-2 checkpoint to {}", path); + + // Update metadata + self.metadata.last_checkpoint = Some(path.to_string()); + self.metadata.performance_stats = self.get_performance_metrics(); + + // AGENT F2: CRITICAL FIX - Actually save model weights to disk using safetensors + // This replaces the stub implementation that only logged without saving + + // Add .safetensors extension if not present + let safetensors_path = if path.ends_with(".safetensors") || path.ends_with(".ckpt") { + if path.ends_with(".ckpt") { + path.replace(".ckpt", ".safetensors") + } else { + path.to_string() + } + } else { + format!("{}.safetensors", path) + }; + + // Extract all tensors from VarMap + let vars_data = self.varmap.data().lock().map_err(|e| { + MLError::LockError(format!("Failed to lock VarMap for checkpoint: {}", e)) + })?; + + // Build tensor map for safetensors serialization + let mut tensors: StdHashMap = StdHashMap::new(); + for (name, var) in vars_data.iter() { + tensors.insert(name.clone(), var.as_tensor().clone()); + } + + // Save using safetensors format (thread-safe serialization) + candle_core::safetensors::save(&tensors, &safetensors_path) + .map_err(|e| MLError::CheckpointError(format!("Failed to save safetensors: {}", e)))?; + + // Verify checkpoint was saved successfully + let metadata = std::fs::metadata(&safetensors_path).map_err(|e| { + MLError::CheckpointError(format!("Checkpoint verification failed: {}", e)) + })?; + + let file_size_mb = metadata.len() as f64 / (1024.0 * 1024.0); + + info!( + "✓ MAMBA-2 checkpoint saved successfully: {} ({:.2} MB, {} parameters)", + safetensors_path, file_size_mb, self.metadata.num_parameters + ); + + // Validate checkpoint size is reasonable (>1MB for non-trivial models) + if file_size_mb < 0.1 { + warn!( + "⚠️ Checkpoint file size is suspiciously small ({:.2} MB) - may indicate incomplete save", + file_size_mb + ); + } + + Ok(()) + } + + /// Load model checkpoint + pub async fn load_checkpoint(&mut self, path: &str) -> Result<(), MLError> { + info!("Loading MAMBA-2 checkpoint from {}", path); + + // AGENT F2: CRITICAL FIX - Actually load model weights from disk + // This replaces the stub implementation that only set flags without loading + + // Add .safetensors extension if not present + let safetensors_path = if path.ends_with(".safetensors") || path.ends_with(".ckpt") { + if path.ends_with(".ckpt") { + path.replace(".ckpt", ".safetensors") + } else { + path.to_string() + } + } else { + format!("{}.safetensors", path) + }; + + // Verify checkpoint file exists + if !std::path::Path::new(&safetensors_path).exists() { + return Err(MLError::CheckpointError(format!( + "Checkpoint file not found: {}", + safetensors_path + ))); + } + + // Load tensors from safetensors + let tensors = candle_core::safetensors::load(&safetensors_path, &self.device) + .map_err(|e| MLError::CheckpointError(format!("Failed to load safetensors: {}", e)))?; + + // Populate VarMap with loaded tensors + let mut vars_data = self.varmap.data().lock().map_err(|e| { + MLError::LockError(format!("Failed to lock VarMap for checkpoint load: {}", e)) + })?; + + for (name, tensor) in tensors.iter() { + // Create new Var from loaded tensor + let var = Var::from_tensor(tensor)?; + vars_data.insert(name.clone(), var); + } + + self.is_trained = true; + self.metadata.last_checkpoint = Some(path.to_string()); + + info!( + "✓ MAMBA-2 checkpoint loaded successfully: {} ({} tensors)", + safetensors_path, + tensors.len() + ); + + Ok(()) + } + + /// Apply gradient clipping to prevent exploding gradients + /// + /// SSM gradients (A, B, C, delta) are not directly trainable in standard MAMBA-2; + /// they flow through the VarMap and are clipped by AdamW weight_decay. + /// The previous implementation computed per-parameter norms (4N GPU syncs) + /// but discarded the clipped results. Trainable parameter gradients are + /// handled by the optimizer step. + fn clip_gradients(&mut self, _max_norm: f64) -> Result<(), MLError> { + Ok(()) + } + + /// Apply Adam optimizer update to a single parameter + fn apply_adam_update( + &mut self, + param: &mut Tensor, + grad: &Tensor, + layer_idx: usize, + param_name: &str, + lr: f64, + beta1: f64, + beta2: f64, + eps: f64, + bias_correction1: f64, + bias_correction2: f64, + apply_weight_decay: bool, + ) -> Result<(), MLError> { + // Create unique keys for momentum and variance + let m_key = format!( + "layer_{}_{}_{}_m", + layer_idx, + param_name, + param.dims().len() + ); + let v_key = format!( + "layer_{}_{}_{}_v", + layer_idx, + param_name, + param.dims().len() + ); + + // Initialize momentum and variance if not present + if !self.optimizer_state.contains_key(&m_key) { + let m_init = grad.zeros_like()?; + let v_init = grad.zeros_like()?; + self.optimizer_state.insert(m_key.clone(), m_init); + self.optimizer_state.insert(v_key.clone(), v_init); + } + + // Get momentum and variance tensors separately to avoid double borrow + let m_tensor = self + .optimizer_state + .get(&m_key) + .ok_or_else(|| { + MLError::ModelError(format!("Missing momentum tensor for key: {}", m_key)) + })? + .clone(); + let v_tensor = self + .optimizer_state + .get(&v_key) + .ok_or_else(|| { + MLError::ModelError(format!("Missing variance tensor for key: {}", v_key)) + })? + .clone(); + + // Apply weight decay if specified + // REFACTORED (Agent 234): Use scalar_tensor helper to eliminate dtype boilerplate + let device = self.device(); + let dtype = param.dtype(); + let effective_grad = if apply_weight_decay && self.config.weight_decay > 0.0 { + let weight_decay_scalar = Self::scalar_tensor(self.config.weight_decay, dtype, device)?; + let weight_decay_term = param.broadcast_mul(&weight_decay_scalar)?; + grad.add(&weight_decay_term)? + } else { + grad.clone() + }; + + // Update biased first moment estimate: m_t = β1 * m_{t-1} + (1 - β1) * g_t + // REFACTORED (Agent 234): Use scalar_tensor helper (was 87 lines of boilerplate) + let beta1_scalar = Self::scalar_tensor(beta1, dtype, device)?; + let m_scaled = m_tensor.broadcast_mul(&beta1_scalar)?; + let grad_scalar = Self::scalar_tensor(1.0 - beta1, dtype, device)?; + let grad_scaled = effective_grad.broadcast_mul(&grad_scalar)?; + let new_m = m_scaled.add(&grad_scaled)?; + + // Update biased second moment estimate: v_t = β2 * v_{t-1} + (1 - β2) * g_t^2 + let grad_squared = effective_grad.mul(&effective_grad)?; + let beta2_scalar = Self::scalar_tensor(beta2, dtype, device)?; + let v_scaled = v_tensor.broadcast_mul(&beta2_scalar)?; + let grad_squared_scalar = Self::scalar_tensor(1.0 - beta2, dtype, device)?; + let grad_squared_scaled = grad_squared.broadcast_mul(&grad_squared_scalar)?; + let new_v = v_scaled.add(&grad_squared_scaled)?; + + // Compute bias-corrected estimates + let bias_corr1_scalar = Self::scalar_tensor(1.0 / bias_correction1, dtype, device)?; + let m_hat = new_m.broadcast_mul(&bias_corr1_scalar)?; + let bias_corr2_scalar = Self::scalar_tensor(1.0 / bias_correction2, dtype, device)?; + let v_hat = new_v.broadcast_mul(&bias_corr2_scalar)?; + + // Compute parameter update: θ = θ - lr * m_hat / (√(v_hat) + ε) + let sqrt_v_hat = v_hat.sqrt()?; + let eps_scalar = Self::scalar_tensor(eps, dtype, device)?; + let denominator = sqrt_v_hat.broadcast_add(&eps_scalar)?; + let lr_scalar = Self::scalar_tensor(lr, dtype, device)?; + let update = m_hat.div(&denominator)?.broadcast_mul(&lr_scalar)?; + + // Update parameter: θ_{t+1} = θ_t - update + *param = param.sub(&update)?; + + // Store updated momentum and variance back + self.optimizer_state.insert(m_key, new_m); + self.optimizer_state.insert(v_key, new_v); + + Ok(()) + } + + /// Apply SGD optimizer update with momentum to a single parameter + /// Apply AdamW optimizer update to a single parameter with decoupled weight decay + /// + /// CRITICAL: Weight decay is applied DIRECTLY to parameters, NOT to gradients. + /// This prevents interference with SSM spectral radius constraints. + /// + /// AdamW Update Formula: + /// 1. m_t = β1 * m_{t-1} + (1 - β1) * g_t (WITHOUT weight decay in gradient) + /// 2. v_t = β2 * v_{t-1} + (1 - β2) * g_t^2 + /// 3. m_hat = m_t / (1 - β1^t), v_hat = v_t / (1 - β2^t) + /// 4. θ_t = θ_{t-1} * (1 - λ * lr) - lr * m_hat / (√v_hat + ε) + /// ^^^^^^^^^^^^^^^^^^^^^^^^ ← Weight decay applied to parameter + fn apply_adamw_update( + &mut self, + param: &mut Tensor, + grad: &Tensor, + layer_idx: usize, + param_name: &str, + lr: f64, + beta1: f64, + beta2: f64, + eps: f64, + bias_correction1: f64, + bias_correction2: f64, + weight_decay: f64, + ) -> Result<(), MLError> { + // Create unique keys for momentum and variance + let m_key = format!( + "layer_{}_{}_{}_m", + layer_idx, + param_name, + param.dims().len() + ); + let v_key = format!( + "layer_{}_{}_{}_v", + layer_idx, + param_name, + param.dims().len() + ); + + // Initialize momentum and variance if not present + if !self.optimizer_state.contains_key(&m_key) { + let m_init = grad.zeros_like()?; + let v_init = grad.zeros_like()?; + self.optimizer_state.insert(m_key.clone(), m_init); + self.optimizer_state.insert(v_key.clone(), v_init); + } + + // Get momentum and variance tensors + let m_tensor = self + .optimizer_state + .get(&m_key) + .ok_or_else(|| { + MLError::ModelError(format!("Missing momentum tensor for key: {}", m_key)) + })? + .clone(); + let v_tensor = self + .optimizer_state + .get(&v_key) + .ok_or_else(|| { + MLError::ModelError(format!("Missing variance tensor for key: {}", v_key)) + })? + .clone(); + + let device = self.device(); + let dtype = param.dtype(); + + // CRITICAL: NO weight decay applied to gradient (pure gradient) + // This is the key difference from Adam optimizer + + // Update biased first moment estimate: m_t = β1 * m_{t-1} + (1 - β1) * g_t + let beta1_scalar = Self::scalar_tensor(beta1, dtype, device)?; + let m_scaled = m_tensor.broadcast_mul(&beta1_scalar)?; + let grad_scalar = Self::scalar_tensor(1.0 - beta1, dtype, device)?; + let grad_scaled = grad.broadcast_mul(&grad_scalar)?; + let new_m = m_scaled.add(&grad_scaled)?; + + // Update biased second moment estimate: v_t = β2 * v_{t-1} + (1 - β2) * g_t^2 + let grad_squared = grad.mul(grad)?; + let beta2_scalar = Self::scalar_tensor(beta2, dtype, device)?; + let v_scaled = v_tensor.broadcast_mul(&beta2_scalar)?; + let grad_squared_scalar = Self::scalar_tensor(1.0 - beta2, dtype, device)?; + let grad_squared_scaled = grad_squared.broadcast_mul(&grad_squared_scalar)?; + let new_v = v_scaled.add(&grad_squared_scaled)?; + + // Compute bias-corrected estimates + let bias_corr1_scalar = Self::scalar_tensor(1.0 / bias_correction1, dtype, device)?; + let m_hat = new_m.broadcast_mul(&bias_corr1_scalar)?; + let bias_corr2_scalar = Self::scalar_tensor(1.0 / bias_correction2, dtype, device)?; + let v_hat = new_v.broadcast_mul(&bias_corr2_scalar)?; + + // Compute gradient update: lr * m_hat / (√(v_hat) + ε) + let sqrt_v_hat = v_hat.sqrt()?; + let eps_scalar = Self::scalar_tensor(eps, dtype, device)?; + let denominator = sqrt_v_hat.broadcast_add(&eps_scalar)?; + let lr_scalar = Self::scalar_tensor(lr, dtype, device)?; + let grad_update = m_hat.div(&denominator)?.broadcast_mul(&lr_scalar)?; + + // CRITICAL: Apply decoupled weight decay directly to parameter + // θ_t = θ_{t-1} * (1 - λ * lr) - grad_update + // This is the key innovation of AdamW vs Adam + let updated_param = if weight_decay > 0.0 { + // Apply weight decay: param * (1 - wd * lr) + let decay_factor = 1.0 - weight_decay * lr; + let decay_scalar = Self::scalar_tensor(decay_factor, dtype, device)?; + let decayed_param = param.broadcast_mul(&decay_scalar)?; + // Then subtract gradient update + decayed_param.sub(&grad_update)? + } else { + // No weight decay, just gradient update + param.sub(&grad_update)? + }; + + *param = updated_param; + + // Store updated momentum and variance + self.optimizer_state.insert(m_key, new_m); + self.optimizer_state.insert(v_key, new_v); + + Ok(()) + } + + /// Apply SGD optimizer update with momentum to a single parameter + /// + /// + /// SGD Update Formula: + /// - Momentum: v_t = μ * v_{t-1} + (1 - μ) * g_t + /// - Update: θ_t = θ_{t-1} - lr * v_t + /// + /// Where: + /// - v_t: velocity (momentum state) + /// - μ: momentum coefficient (typically 0.9) + /// - g_t: gradient (with optional weight decay) + /// - lr: learning rate + fn apply_sgd_update( + &mut self, + param: &mut Tensor, + grad: &Tensor, + layer_idx: usize, + param_name: &str, + lr: f64, + momentum: f64, + apply_weight_decay: bool, + ) -> Result<(), MLError> { + // Create unique key for velocity (momentum state) + let v_key = format!( + "layer_{}_{}_{}_velocity", + layer_idx, + param_name, + param.dims().len() + ); + + // Initialize velocity if not present (zeros) + if !self.optimizer_state.contains_key(&v_key) { + let v_init = grad.zeros_like()?; + self.optimizer_state.insert(v_key.clone(), v_init); + } + + // Get velocity tensor + let v_tensor = self + .optimizer_state + .get(&v_key) + .ok_or_else(|| { + MLError::ModelError(format!("Missing velocity tensor for key: {}", v_key)) + })? + .clone(); + + // Apply weight decay if specified (L2 regularization) + let device = self.device(); + let dtype = param.dtype(); + let effective_grad = if apply_weight_decay && self.config.weight_decay > 0.0 { + let weight_decay_scalar = Self::scalar_tensor(self.config.weight_decay, dtype, device)?; + let weight_decay_term = param.broadcast_mul(&weight_decay_scalar)?; + grad.add(&weight_decay_term)? + } else { + grad.clone() + }; + + // Update velocity: v_t = μ * v_{t-1} + (1 - μ) * g_t + let momentum_scalar = Self::scalar_tensor(momentum, dtype, device)?; + let v_scaled = v_tensor.broadcast_mul(&momentum_scalar)?; + let grad_scalar = Self::scalar_tensor(1.0 - momentum, dtype, device)?; + let grad_scaled = effective_grad.broadcast_mul(&grad_scalar)?; + let new_v = v_scaled.add(&grad_scaled)?; + + // Compute parameter update: θ_{t+1} = θ_t - lr * v_t + let lr_scalar = Self::scalar_tensor(lr, dtype, device)?; + let update = new_v.broadcast_mul(&lr_scalar)?; + *param = param.sub(&update)?; + + // Store updated velocity back + self.optimizer_state.insert(v_key, new_v); + + Ok(()) + } + + /// Project SSM matrices to maintain stability + fn project_ssm_matrices(&mut self) -> Result<(), MLError> { + // Avoid borrow checker issues by processing each state separately + for i in 0..self.state.ssm_states.len() { + // Ensure A matrix has spectral radius < 1 for stability + let spectral_radius = { + let ssm_state = &self.state.ssm_states[i]; + self.compute_spectral_radius(&ssm_state.A)? + }; + if spectral_radius >= 1.0 { + let scale_factor = 0.99 / spectral_radius; + let device = self.device(); + let scale_tensor = Tensor::new(&[scale_factor as f32], device)?; + self.state.ssm_states[i].A = + self.state.ssm_states[i].A.broadcast_mul(&scale_tensor)?; + } + + // Ensure Delta parameter stays positive and reasonable + let device = self.device(); + let delta_min = Tensor::new(&[1e-6_f32], device)?; + let delta_max = Tensor::new(&[1.0_f32], device)?; + let delta_clamped = self.state.ssm_states[i] + .delta + .broadcast_maximum(&delta_min)? + .broadcast_minimum(&delta_max)?; + self.state.ssm_states[i].delta = delta_clamped; + } + + Ok(()) + } + + /// Compute spectral radius (largest eigenvalue magnitude) of a matrix + fn compute_spectral_radius(&self, matrix: &Tensor) -> Result { + // For simplicity, use Frobenius norm as approximation + // In production, we'd compute actual eigenvalues + let frobenius_norm = matrix.powf(2.0)?.sum_all()?.to_scalar::()? as f64; + let frobenius_norm = frobenius_norm.sqrt(); + + // Frobenius norm upper bounds spectral radius + // For better approximation, we scale by sqrt of matrix size + let dims = matrix.dims(); + if dims.len() >= 2 { + let size = (dims[0].min(dims[1]) as f64).sqrt(); + Ok(frobenius_norm / size) + } else { + Ok(frobenius_norm) + } + } + + /// Get current training state from model metadata + pub fn get_current_training_state(&self) -> (Option, Option) { + if let Some(last_epoch) = self.metadata.training_history.back() { + (Some(last_epoch.epoch as u64), None) // MAMBA doesn't track steps within epochs + } else { + (None, None) + } + } + + /// Get training metrics from model performance data + pub fn get_training_metrics(&self) -> HashMap { + let mut metrics = HashMap::new(); + + if let Some(last_epoch) = self.metadata.training_history.back() { + metrics.insert("training_loss".to_owned(), last_epoch.loss); + metrics.insert("validation_loss".to_owned(), last_epoch.loss); + metrics.insert("directional_accuracy".to_owned(), last_epoch.accuracy); + metrics.insert("mae".to_owned(), last_epoch.loss); + metrics.insert("rmse".to_owned(), last_epoch.loss.sqrt()); + metrics.insert("r_squared".to_owned(), 1.0 - last_epoch.loss.min(1.0)); + } + + // Add other available metrics + let perf_metrics = self.get_performance_metrics(); + for (key, value) in perf_metrics { + if key.contains("loss") || key.contains("accuracy") { + metrics.insert(key, value); + } + } + + metrics + } + + /// Get inference performance statistics + pub fn get_inference_stats(&self) -> HashMap { + let mut stats = HashMap::new(); + + // Calculate average latency from latency histogram + let total_inferences = self + .total_inferences + .load(Ordering::Relaxed) as f64; + if total_inferences > 0.0 { + // Simulate latency calculation from internal metrics + let avg_latency = self.config.target_latency_us as f64 * 0.8; // Assume 80% of target + stats.insert("avg_latency_us".to_owned(), avg_latency); + + // Calculate throughput based on latency + let throughput_pps = if avg_latency > 0.0 { + 1_000_000.0 / avg_latency // predictions per second + } else { + 0.0 + }; + stats.insert("throughput_pps".to_owned(), throughput_pps); + } + + stats + } + + /// Extract SSD layer weights from the model + pub fn extract_ssd_weights(&self) -> Vec> { + // In a real implementation, this would extract actual SSD layer weights + // For now, return structured weight data based on model configuration + let num_layers = self.config.num_layers; + let d_model = self.config.d_model; + let expand = self.config.expand; + + let mut weights = Vec::new(); + for layer in 0..num_layers { + // Each SSD layer has weights of size [d_model * expand, d_model] + let layer_size = d_model * expand; + let mut layer_weights = Vec::with_capacity(layer_size); + + // Generate realistic weight values based on layer index + let scale = 1.0 / (layer + 1) as f32; + for i in 0..layer_size { + let weight = scale * (i as f32 / layer_size as f32 - 0.5); + layer_weights.push(weight); + } + weights.push(layer_weights); + } + + weights + } + + /// Extract input projection weights + pub fn extract_input_projection_weights(&self) -> Vec { + let d_model = self.config.d_model; + let mut weights = Vec::with_capacity(d_model); + + // Generate input projection weights + for i in 0..d_model { + let weight = (i as f32 / d_model as f32 - 0.5) * 0.1; + weights.push(weight); + } + + weights + } + + /// Extract output projection weights + pub fn extract_output_projection_weights(&self) -> Vec { + let d_model = self.config.d_model; + let mut weights = Vec::with_capacity(d_model); + + // Generate output projection weights + for i in 0..d_model { + let weight = (i as f32 / d_model as f32 - 0.5) * 0.05; + weights.push(weight); + } + + weights + } + + /// Extract layer normalization weights + pub fn extract_layer_norm_weights(&self) -> Vec> { + let num_layers = self.config.num_layers; + let d_model = self.config.d_model; + let mut weights = Vec::new(); + + for _layer in 0..num_layers { + let mut layer_norm = Vec::with_capacity(d_model); + // Layer norm weights typically start at 1.0 + for _i in 0..d_model { + layer_norm.push(1.0); + } + weights.push(layer_norm); + } + + weights + } + + /// Extract state space model matrices + pub fn extract_ssm_matrices(&self, matrix_type: &str) -> Vec> { + let num_layers = self.config.num_layers; + let d_state = self.config.d_state; + let d_model = self.config.d_model; + let mut matrices = Vec::new(); + + for layer in 0..num_layers { + let matrix_size = match matrix_type { + "A" => d_state * d_state, // A matrix is [d_state, d_state] + "B" => d_state * d_model, // B matrix is [d_state, d_model] + "C" => d_model * d_state, // C matrix is [d_model, d_state] + _ => d_state, + }; + + let mut matrix = Vec::with_capacity(matrix_size); + let scale = match matrix_type { + "A" => -0.1, // A matrices typically have negative values for stability + "B" => 0.1, + "C" => 0.1, + _ => 0.1, + }; + + for i in 0..matrix_size { + let value = scale * (i as f32 / matrix_size as f32 - 0.5) * (layer + 1) as f32; + matrix.push(value); + } + matrices.push(matrix); + } + + matrices + } + + /// Extract delta parameters for selective state space + pub fn extract_delta_params(&self) -> Vec { + let d_model = self.config.d_model; + let mut deltas = Vec::with_capacity(d_model); + + // Delta parameters control the timescale of state updates + for i in 0..d_model { + // Initialize with reasonable timescale values + let delta = 1.0 + (i as f32 / d_model as f32) * 0.1; + deltas.push(delta); + } + + deltas + } + + /// Restore SSD layer weights + pub fn restore_ssd_weights(&mut self, weights: &[Vec]) { + debug!("Restoring {} SSD layer weight matrices", weights.len()); + + // Validate weight matrix dimensions against config + let expected_layers = self.config.num_layers; + if weights.len() != expected_layers { + warn!( + "SSD weight count mismatch: expected {} layers, got {}", + expected_layers, + weights.len() + ); + } + + // Store weights in SSD layers - using actual struct fields + for (layer_idx, (_layer, layer_weights)) in + self.ssd_layers.iter_mut().zip(weights.iter()).enumerate() + { + // Update the actual layer weights (this depends on SSDLayer implementation) + // For now, we'll store in optimizer_state as a workaround + let layer_key = format!("ssd_layer_{}", layer_idx); + if let Ok(tensor) = + Tensor::from_slice(layer_weights, (layer_weights.len(),), &Device::Cpu) + { + self.optimizer_state.insert(layer_key, tensor); + } + } + + // Validate individual layer weight dimensions + for (layer_idx, layer_weights) in weights.into_iter().enumerate() { + let expected_size = self.config.d_model * self.config.d_model; // Simplified square matrix + if layer_weights.len() != expected_size { + warn!( + "Layer {} weight size mismatch: expected {}, got {}", + layer_idx, + expected_size, + layer_weights.len() + ); + } + + // Validate weight values are finite + let invalid_count = layer_weights.iter().filter(|&&w| !w.is_finite()).count(); + if invalid_count > 0 { + warn!( + "Layer {} contains {} invalid weight values", + layer_idx, invalid_count + ); + } + + debug!( + "SSD Layer {}: {} weights loaded, range [{:.4}, {:.4}]", + layer_idx, + layer_weights.len(), + layer_weights.iter().fold(f32::INFINITY, |a, &b| a.min(b)), + layer_weights + .iter() + .fold(f32::NEG_INFINITY, |a, &b| a.max(b)) + ); + } + + info!( + "Successfully restored {} SSD layer weight matrices", + weights.len() + ); + } + + /// Restore input projection weights + pub fn restore_input_projection_weights(&mut self, weights: &[f32]) { + debug!("Restoring {} input projection weights", weights.len()); + + // Validate weight dimensions (input projection typically projects from vocab_size to d_model, simplified as d_model * d_model) + let expected_size = self.config.d_model * self.config.d_model; + if weights.len() != expected_size { + warn!( + "Input projection weight size mismatch: expected {}, got {}", + expected_size, + weights.len() + ); + } + + // Validate weight values are finite + let invalid_count = weights.iter().filter(|&&w| !w.is_finite()).count(); + if invalid_count > 0 { + warn!( + "Input projection contains {} invalid weight values", + invalid_count + ); + } + + // Store input projection weights using actual struct field + // The actual input_projection is a Linear layer, store in optimizer_state as workaround + let key = "input_projection_weights".to_owned(); + if let Ok(tensor) = Tensor::from_slice(weights, (weights.len(),), &Device::Cpu) { + self.optimizer_state.insert(key, tensor); + } + + let weight_range = if !weights.is_empty() { + ( + weights.iter().fold(f32::INFINITY, |a, &b| a.min(b)), + weights.iter().fold(f32::NEG_INFINITY, |a, &b| a.max(b)), + ) + } else { + (0.0, 0.0) + }; + + info!( + "Input projection weights restored: {} parameters, range [{:.4}, {:.4}]", + weights.len(), + weight_range.0, + weight_range.1 + ); + } + + /// Restore output projection weights + pub fn restore_output_projection_weights(&mut self, weights: &[f32]) { + debug!("Restoring {} output projection weights", weights.len()); + + // Validate weight dimensions (output projection typically projects from d_model to vocab_size, simplified as d_model * d_model) + let expected_size = self.config.d_model * self.config.d_model; + if weights.len() != expected_size { + warn!( + "Output projection weight size mismatch: expected {}, got {}", + expected_size, + weights.len() + ); + } + + // Validate weight values are finite + let invalid_count = weights.iter().filter(|&&w| !w.is_finite()).count(); + if invalid_count > 0 { + warn!( + "Output projection contains {} invalid weight values", + invalid_count + ); + } + + // Store output projection weights using actual struct field + // The actual output_projection is a Linear layer, store in optimizer_state as workaround + let key = "output_projection_weights".to_owned(); + if let Ok(tensor) = Tensor::from_slice(weights, (weights.len(),), &Device::Cpu) { + self.optimizer_state.insert(key, tensor); + } + + let weight_range = if !weights.is_empty() { + ( + weights.iter().fold(f32::INFINITY, |a, &b| a.min(b)), + weights.iter().fold(f32::NEG_INFINITY, |a, &b| a.max(b)), + ) + } else { + (0.0, 0.0) + }; + + info!( + "Output projection weights restored: {} parameters, range [{:.4}, {:.4}]", + weights.len(), + weight_range.0, + weight_range.1 + ); + } + + /// Restore layer normalization weights + pub fn restore_layer_norm_weights(&mut self, weights: &[Vec]) { + debug!("Restoring {} layer norm weight matrices", weights.len()); + + // Validate layer count + let expected_layers = self.config.num_layers; + if weights.len() != expected_layers { + warn!( + "Layer norm count mismatch: expected {} layers, got {}", + expected_layers, + weights.len() + ); + } + + // Store layer norm weights using actual struct field + // The layer_norms field contains actual LayerNorm objects, store in optimizer_state as workaround + for (idx, layer_weights) in weights.into_iter().enumerate() { + let key = format!("layer_norm_weights_{}", idx); + if let Ok(tensor) = + Tensor::from_slice(layer_weights, (layer_weights.len(),), &Device::Cpu) + { + self.optimizer_state.insert(key, tensor); + } + } + + // Validate individual layer norm weights + for (layer_idx, layer_weights) in weights.into_iter().enumerate() { + let expected_size = self.config.d_model; // Layer norm has d_model parameters + if layer_weights.len() != expected_size { + warn!( + "Layer norm {} weight size mismatch: expected {}, got {}", + layer_idx, + expected_size, + layer_weights.len() + ); + } + + // Validate weight values are finite and positive (layer norm weights should be positive) + let invalid_count = layer_weights + .iter() + .filter(|&&w| !w.is_finite() || w <= 0.0) + .count(); + if invalid_count > 0 { + warn!( + "Layer norm {} contains {} invalid weight values (non-positive or non-finite)", + layer_idx, invalid_count + ); + } + + let weight_range = if !layer_weights.is_empty() { + ( + layer_weights.iter().fold(f32::INFINITY, |a, &b| a.min(b)), + layer_weights + .iter() + .fold(f32::NEG_INFINITY, |a, &b| a.max(b)), + ) + } else { + (0.0, 0.0) + }; + + debug!( + "Layer norm {}: {} weights, range [{:.4}, {:.4}]", + layer_idx, + layer_weights.len(), + weight_range.0, + weight_range.1 + ); + } + + info!( + "Successfully restored {} layer normalization weight matrices", + weights.len() + ); + } + + /// Restore state space model matrices + pub fn restore_ssm_matrices(&mut self, matrix_type: &str, matrices: &[Vec]) { + debug!("Restoring {} {} matrices", matrices.len(), matrix_type); + + // Validate matrix count + let expected_layers = self.config.num_layers; + if matrices.len() != expected_layers { + warn!( + "{} matrix count mismatch: expected {} layers, got {}", + matrix_type, + expected_layers, + matrices.len() + ); + } + + // Store matrices in appropriate fields based on type + match matrix_type { + "A" => { + let key = "ssm_A_matrices".to_owned(); + for (idx, matrix) in matrices.into_iter().enumerate() { + let matrix_key = format!("{}_{}", key, idx); + if let Ok(tensor) = Tensor::from_slice(matrix, (matrix.len(),), &Device::Cpu) { + self.optimizer_state.insert(matrix_key, tensor); + } + } + info!( + "Restored {} A matrices for state space model", + matrices.len() + ); + }, + "B" => { + let key = "ssm_B_matrices".to_owned(); + for (idx, matrix) in matrices.into_iter().enumerate() { + let matrix_key = format!("{}_{}", key, idx); + if let Ok(tensor) = Tensor::from_slice(matrix, (matrix.len(),), &Device::Cpu) { + self.optimizer_state.insert(matrix_key, tensor); + } + } + info!( + "Restored {} B matrices for state space model", + matrices.len() + ); + }, + "C" => { + let key = "ssm_C_matrices".to_owned(); + for (idx, matrix) in matrices.into_iter().enumerate() { + let matrix_key = format!("{}_{}", key, idx); + if let Ok(tensor) = Tensor::from_slice(matrix, (matrix.len(),), &Device::Cpu) { + self.optimizer_state.insert(matrix_key, tensor); + } + } + info!( + "Restored {} C matrices for state space model", + matrices.len() + ); + }, + _ => { + warn!("Unknown SSM matrix type: {}", matrix_type); + return; + }, + } + + // Validate individual matrices + for (layer_idx, matrix) in matrices.into_iter().enumerate() { + let expected_size = match matrix_type { + "A" => self.config.d_state * self.config.d_state, + "B" => self.config.d_state * self.config.d_model, + "C" => self.config.d_model * self.config.d_state, + _ => self.config.d_state, + }; + + if matrix.len() != expected_size { + warn!( + "SSM {} matrix {} size mismatch: expected {}, got {}", + matrix_type, + layer_idx, + expected_size, + matrix.len() + ); + } + + // Validate matrix values are finite + let invalid_count = matrix.iter().filter(|&&v| !v.is_finite()).count(); + if invalid_count > 0 { + warn!( + "SSM {} matrix {} contains {} invalid values", + matrix_type, layer_idx, invalid_count + ); + } + + let matrix_range = if !matrix.is_empty() { + ( + matrix.iter().fold(f32::INFINITY, |a, &b| a.min(b)), + matrix.iter().fold(f32::NEG_INFINITY, |a, &b| a.max(b)), + ) + } else { + (0.0, 0.0) + }; + + debug!( + "SSM {} matrix {}: {} values, range [{:.4}, {:.4}]", + matrix_type, + layer_idx, + matrix.len(), + matrix_range.0, + matrix_range.1 + ); + } + } + + /// Restore delta parameters + pub fn restore_delta_params(&mut self, deltas: &[f32]) { + debug!("Restoring {} delta parameters", deltas.len()); + + // Validate delta parameter count + let expected_size = self.config.d_model; + if deltas.len() != expected_size { + warn!( + "Delta parameter count mismatch: expected {}, got {}", + expected_size, + deltas.len() + ); + } + + // Validate delta values are finite and positive (deltas control timescales) + let invalid_count = deltas + .iter() + .filter(|&&d| !d.is_finite() || d <= 0.0) + .count(); + if invalid_count > 0 { + warn!( + "Delta parameters contain {} invalid values (non-positive or non-finite)", + invalid_count + ); + } + + // Store delta parameters using optimizer_state since the field doesn't exist + let key = "ssm_delta_params".to_owned(); + if let Ok(tensor) = Tensor::from_slice(deltas, (deltas.len(),), &Device::Cpu) { + self.optimizer_state.insert(key, tensor); + } + + let delta_range = if !deltas.is_empty() { + ( + deltas.iter().fold(f32::INFINITY, |a, &b| a.min(b)), + deltas.iter().fold(f32::NEG_INFINITY, |a, &b| a.max(b)), + ) + } else { + (0.0, 0.0) + }; + + info!( + "Delta parameters restored: {} values, range [{:.4}, {:.4}]", + deltas.len(), + delta_range.0, + delta_range.1 + ); + + // Log statistics for debugging + if !deltas.is_empty() { + let mean = deltas.iter().sum::() / deltas.len() as f32; + let variance = + deltas.iter().map(|&x| (x - mean).powi(2)).sum::() / deltas.len() as f32; + debug!( + "Delta parameter statistics: mean={:.4}, variance={:.4}", + mean, variance + ); + } + } +} + +impl Clone for Mamba2SSM { + /// Clone implementation for checkpoint saving + /// + /// Note: This is a shallow clone that copies configuration and metadata, + /// but shares tensor references. Use for checkpoint operations only. + fn clone(&self) -> Self { + // Create a new model with same configuration + // Note: This is a simplified clone for checkpoint operations + // Full deep cloning of all tensors would be expensive + match Mamba2SSM::new(self.config.clone(), &self.device) { + Ok(model) => model, + Err(e) => { + tracing::error!("Mamba2SSM clone failed: {}", e); + std::process::abort(); + } + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + use anyhow::Result; + + #[tokio::test] + async fn test_mamba_creation() -> Result<()> { + let config = Mamba2Config { + d_model: 8, + d_state: 4, + d_head: 4, + num_heads: 2, + ..Default::default() + }; + + let device = Device::Cpu; + let model = Mamba2SSM::new(config, &device) + .map_err(|_| anyhow::anyhow!("Failed to create MAMBA model"))?; + assert_eq!(model.metadata.input_dim, 8); + assert_eq!(model.metadata.output_dim, 1); + Ok(()) + } + + #[test] + fn test_mamba_config_default() -> Result<()> { + let config = Mamba2Config::default(); + assert!(config.d_model > 0); + assert!(config.d_state > 0); + assert!(config.num_heads > 0); + Ok(()) + } + + #[test] + fn test_mamba_state_creation() -> Result<()> { + let config = Mamba2Config { + d_model: 4, + d_state: 2, + d_head: 2, + num_heads: 2, + ..Default::default() + }; + + let device = Device::Cpu; + let state = Mamba2State::zeros(&config, &device) + .map_err(|_| anyhow::anyhow!("Failed to create MAMBA state"))?; + assert_eq!(state.ssm_states.len(), config.num_layers); + assert!(!state.selective_state.is_empty()); + Ok(()) + } + + #[test] + fn test_mamba_performance_metrics() -> Result<()> { + let config = Mamba2Config { + d_model: 4, + target_latency_us: 5, + ..Default::default() + }; + + let device = Device::Cpu; + let model = Mamba2SSM::new(config, &device) + .map_err(|_| anyhow::anyhow!("Failed to create MAMBA model"))?; + let metrics = model.get_performance_metrics(); + + assert!(metrics.contains_key("total_inferences")); + assert!(metrics.contains_key("model_parameters")); + assert!(metrics.contains_key("compression_ratio")); + Ok(()) + } + + #[test] + fn test_mamba_learning_rate_schedule() -> Result<()> { + let config = Mamba2Config { + d_model: 4, + num_layers: 1, + learning_rate: 0.001, + warmup_steps: 10, + batch_size: 2, + ..Default::default() + }; + + let device = Device::Cpu; + let mut model = Mamba2SSM::new(config.clone(), &device) + .map_err(|_| anyhow::anyhow!("Failed to create MAMBA model"))?; + + // Set total training samples (simulate 100 samples) + model.total_training_samples = 100; + + // Test warmup phase (steps 0-9) + for step in 0..10 { + let epoch = step / (model.total_training_samples / model.config.batch_size); + let batch_idx = (step % (model.total_training_samples / model.config.batch_size)) + * model.config.batch_size; + + model + .update_learning_rate(epoch, batch_idx) + .map_err(|_| anyhow::anyhow!("Failed to update LR"))?; + + let current_lr = model.get_current_learning_rate(); + let expected_lr = config.learning_rate * (step as f64 / config.warmup_steps as f64); + + // Allow small floating point error + assert!( + (current_lr - expected_lr).abs() < 1e-9, + "Warmup step {}: expected {}, got {}", + step, + expected_lr, + current_lr + ); + } + + // Test decay phase (after warmup) + let step = 15; + let epoch = step / (model.total_training_samples / model.config.batch_size); + let batch_idx = (step % (model.total_training_samples / model.config.batch_size)) + * model.config.batch_size; + + model + .update_learning_rate(epoch, batch_idx) + .map_err(|_| anyhow::anyhow!("Failed to update LR"))?; + + let decay_lr = model.get_current_learning_rate(); + // During decay, LR should be less than initial LR but greater than 0 + assert!( + decay_lr < config.learning_rate && decay_lr > 0.0, + "Decay phase: LR should be in (0, {}), got {}", + config.learning_rate, + decay_lr + ); + + Ok(()) + } + + #[test] + fn test_mamba_hft_config() -> Result<()> { + let device = Device::Cpu; + let model = Mamba2SSM::default_hft(&device) + .map_err(|_| anyhow::anyhow!("Failed to create HFT MAMBA model"))?; + assert_eq!(model.config.target_latency_us, 3); + assert!(model.config.hardware_aware); + assert!(model.config.use_ssd); + assert!(model.config.use_selective_state); + Ok(()) + } + + #[test] + fn test_mamba_shuffle_batches_deterministic() -> Result<()> { + // Test that with shuffle_batches=false, batch order is deterministic + let config = Mamba2Config { + d_model: 4, + d_state: 2, + batch_size: 2, + seq_len: 4, + shuffle_batches: false, + ..Default::default() + }; + + let device = Device::Cpu; + let model = Mamba2SSM::new(config.clone(), &device) + .map_err(|_| anyhow::anyhow!("Failed to create MAMBA model"))?; + + // Verify shuffle is disabled + assert!(!model.config.shuffle_batches); + + // Create a simple data sequence + let data_len = 10; + let batch_indices: Vec = (0..data_len).step_by(config.batch_size).collect(); + + // Verify deterministic order (should be [0, 2, 4, 6, 8]) + assert_eq!(batch_indices, vec![0, 2, 4, 6, 8]); + + Ok(()) + } + + #[test] + fn test_mamba_shuffle_batches_enabled() -> Result<()> { + // Test that with shuffle_batches=true, batches can be in different order + let config = Mamba2Config { + d_model: 4, + d_state: 2, + batch_size: 2, + seq_len: 4, + shuffle_batches: true, + ..Default::default() + }; + + let device = Device::Cpu; + let model = Mamba2SSM::new(config.clone(), &device) + .map_err(|_| anyhow::anyhow!("Failed to create MAMBA model"))?; + + // Verify shuffle is enabled + assert!(model.config.shuffle_batches); + + // Test that shuffling actually works + use rand::seq::SliceRandom; + let mut batch_indices: Vec = (0..10).step_by(2).collect(); + let _original = batch_indices.clone(); + + batch_indices.shuffle(&mut rand::thread_rng()); + + // Note: There's a small chance this could fail if shuffle happens to + // produce the same order, but probability is low (1/5! = 1/120) + // For a proper test, we'd need a deterministic RNG with a seed + + Ok(()) + } +} + +#[test] +fn test_mamba_parameter_count() -> anyhow::Result<()> { + let config = Mamba2Config { + d_model: 8, + num_layers: 2, + ..Default::default() + }; + + let param_count = Mamba2SSM::count_parameters(&config); + assert!(param_count > 0); + Ok(()) +} + +#[test] +fn test_bilinear_discretization_more_accurate_than_zoh() { + let zoh = 1.0 + (-1.0) * 0.1; + let bilinear = 1.0 + (-1.0) * 0.1 + ((-1.0) * 0.1_f64).powi(2) / 2.0; + let exact = (-0.1_f64).exp(); + + assert!((bilinear - exact).abs() < (zoh - exact).abs(), + "bilinear {} should be closer to exact {} than zoh {}", + bilinear, exact, zoh); +} diff --git a/crates/ml/src/mamba/scan_algorithms.rs b/crates/ml-supervised/src/mamba/scan_algorithms.rs similarity index 99% rename from crates/ml/src/mamba/scan_algorithms.rs rename to crates/ml-supervised/src/mamba/scan_algorithms.rs index e13752340..143de4375 100644 --- a/crates/ml/src/mamba/scan_algorithms.rs +++ b/crates/ml-supervised/src/mamba/scan_algorithms.rs @@ -13,7 +13,7 @@ use std::collections::HashMap; use std::time::Instant; use crate::liquid::FixedPoint; -use crate::MLError; // Import FixedPoint for financial precision +use ml_core::MLError; // Import FixedPoint for financial precision use candle_core::{Device, Tensor}; use tracing::{debug, instrument}; diff --git a/crates/ml/src/mamba/selective_state.rs b/crates/ml-supervised/src/mamba/selective_state.rs similarity index 99% rename from crates/ml/src/mamba/selective_state.rs rename to crates/ml-supervised/src/mamba/selective_state.rs index 2b5c62846..f2edc176a 100644 --- a/crates/ml/src/mamba/selective_state.rs +++ b/crates/ml-supervised/src/mamba/selective_state.rs @@ -24,7 +24,7 @@ use nalgebra::DVector; use tracing::{debug, instrument}; use super::{Mamba2Config, Mamba2State}; -use crate::MLError; +use ml_core::MLError; /// Configuration for selective state space mechanism #[derive(Debug, Clone)] diff --git a/crates/ml/src/mamba/ssd_layer.rs b/crates/ml-supervised/src/mamba/ssd_layer.rs similarity index 99% rename from crates/ml/src/mamba/ssd_layer.rs rename to crates/ml-supervised/src/mamba/ssd_layer.rs index 2379a5e14..5a0b5d9ae 100644 --- a/crates/ml/src/mamba/ssd_layer.rs +++ b/crates/ml-supervised/src/mamba/ssd_layer.rs @@ -24,7 +24,7 @@ use candle_nn::{Linear, Module, VarBuilder}; use tracing::instrument; use super::{Mamba2Config, Mamba2State}; -use crate::MLError; +use ml_core::MLError; /// Structured State Duality (SSD) Layer implementation #[derive(Debug)] @@ -553,7 +553,7 @@ mod tests { /// Helper to create VarBuilder for tests fn create_test_varbuilder(device: &Device) -> VarBuilder<'_> { let vs = Arc::new(candle_nn::VarMap::new()); - VarBuilder::from_varmap(&vs, crate::dqn::mixed_precision::training_dtype(device), device) + VarBuilder::from_varmap(&vs, ml_core::mixed_precision::training_dtype(device), device) } #[test] diff --git a/crates/ml/src/tft/gated_residual.rs b/crates/ml-supervised/src/tft/gated_residual.rs similarity index 99% rename from crates/ml/src/tft/gated_residual.rs rename to crates/ml-supervised/src/tft/gated_residual.rs index 906226d07..0235edf96 100644 --- a/crates/ml/src/tft/gated_residual.rs +++ b/crates/ml-supervised/src/tft/gated_residual.rs @@ -6,8 +6,8 @@ use candle_core::{Module, Tensor}; use candle_nn::{linear, Linear, VarBuilder}; -use crate::cuda_compat::{layer_norm_with_fallback, manual_sigmoid}; -use crate::MLError; +use ml_core::cuda_compat::{layer_norm_with_fallback, manual_sigmoid}; +use ml_core::MLError; /// CUDA-compatible LayerNorm wrapper /// diff --git a/crates/ml/src/tft/hft_optimizations.rs b/crates/ml-supervised/src/tft/hft_optimizations.rs similarity index 99% rename from crates/ml/src/tft/hft_optimizations.rs rename to crates/ml-supervised/src/tft/hft_optimizations.rs index 5377ea113..3e9fa6f01 100644 --- a/crates/ml/src/tft/hft_optimizations.rs +++ b/crates/ml-supervised/src/tft/hft_optimizations.rs @@ -30,7 +30,7 @@ use tracing::{info, instrument, warn}; use super::TemporalFusionTransformer; use crate::liquid::FixedPoint; -use crate::MLError; +use ml_core::MLError; use common::types::Price; // Import Price for financial predictions // Import FixedPoint for financial precision /// HFT-specific configuration for ultra-low latency inference diff --git a/crates/ml/src/tft/lstm_encoder.rs b/crates/ml-supervised/src/tft/lstm_encoder.rs similarity index 99% rename from crates/ml/src/tft/lstm_encoder.rs rename to crates/ml-supervised/src/tft/lstm_encoder.rs index af20f9bef..c77ef7513 100644 --- a/crates/ml/src/tft/lstm_encoder.rs +++ b/crates/ml-supervised/src/tft/lstm_encoder.rs @@ -19,8 +19,8 @@ use candle_core::{Module, Tensor}; use candle_nn::{linear, Linear, VarBuilder}; use std::collections::HashMap; -use crate::cuda_compat::manual_sigmoid; -use crate::MLError; +use ml_core::cuda_compat::manual_sigmoid; +use ml_core::MLError; /// Single LSTM layer with 4 gates #[derive(Debug)] diff --git a/crates/ml-supervised/src/tft/mod.rs b/crates/ml-supervised/src/tft/mod.rs new file mode 100644 index 000000000..2ade97947 --- /dev/null +++ b/crates/ml-supervised/src/tft/mod.rs @@ -0,0 +1,1410 @@ +//! # Temporal Fusion Transformer (TFT) for HFT +//! +//! State-of-the-art multi-horizon forecasting with variable selection networks, +//! temporal self-attention, gated residual networks, and uncertainty quantification. +//! +//! ## Key Features +//! +//! - Multi-horizon forecasting (1-tick to 100-tick ahead) +//! - Variable selection networks for feature importance +//! - Gated residual networks for improved gradient flow +//! - Quantile outputs for uncertainty estimation +//! - Temporal self-attention for sequential modeling +//! - Sub-50μs inference latency optimized for HFT +//! +//! ## Performance Targets +//! +//! - Inference: <50μs per prediction +//! - Accuracy improvement: +15% over baseline +//! - Memory usage: <1GB +//! - Throughput: >100K predictions/sec + +use std::collections::HashMap; +use std::num::NonZeroUsize; +use std::sync::atomic::{AtomicU64, Ordering}; +use std::sync::Arc; +use std::time::{Instant, SystemTime}; + +use candle_core::{Device, Module, Tensor}; +use candle_nn::{linear, AdamW, Linear, Optimizer, ParamsAdamW, VarBuilder, VarMap}; + +use ml_core::mixed_precision::training_dtype; +use lru::LruCache; +use ndarray::{Array1, Array2}; +use serde::{Deserialize, Serialize}; +use serde_json::Value; +use tracing::{debug, info, instrument, warn}; +use uuid::Uuid; + +use ml_core::MLError; + +// Import TFT components +pub mod gated_residual; +pub mod hft_optimizations; +pub mod lstm_encoder; +pub mod qat_tft; // Quantization-Aware Training wrapper - RE-ENABLED: Device mismatch fix applied +pub mod quantile_outputs; +pub mod quantized_attention; // Re-enabled Wave 9.12 +pub mod quantized_grn; +pub mod quantized_lstm; +pub mod quantized_tft; // Re-enabled Wave 9.12 +pub mod quantized_vsn; +pub mod temporal_attention; +pub mod variable_selection; +pub mod varmap_quantization; + +// Public exports for TFT components +pub use gated_residual::{GRNStack, GatedResidualNetwork}; +pub use lstm_encoder::LSTMEncoder; +pub use qat_tft::QATTemporalFusionTransformer; // Quantization-Aware Training wrapper - RE-ENABLED: Device mismatch fix applied +pub use quantile_outputs::QuantileLayer; +pub use quantized_attention::QuantizedTemporalAttention; // Re-enabled Wave 9.12 +pub use quantized_grn::QuantizedGatedResidualNetwork; +pub use quantized_lstm::QuantizedLSTMEncoder; +pub use quantized_tft::QuantizedTemporalFusionTransformer; // Re-enabled Wave 9.12 +pub use quantized_vsn::QuantizedVariableSelectionNetwork; +pub use temporal_attention::TemporalSelfAttention; +pub use variable_selection::VariableSelectionNetwork; +pub use varmap_quantization::{ + load_quantized_weights, quantize_varmap, quantize_varmap_parallel, save_quantized_weights, +}; + +/// `TFT` Configuration +/// TFT model variant selection (F32 vs INT8) +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] +pub enum TFTVariant { + /// Full precision (F32) model + F32, + /// INT8 quantized model (75% memory reduction) + INT8, +} + +impl Default for TFTVariant { + fn default() -> Self { + Self::F32 + } +} + +impl TFTVariant { + /// Check if variant uses quantization + pub fn is_quantized(&self) -> bool { + matches!(self, Self::INT8) + } + + /// Get expected memory reduction ratio vs F32 + pub fn memory_reduction_ratio(&self) -> f64 { + match self { + Self::F32 => 1.0, + Self::INT8 => 0.25, // 75% reduction → 25% of original + } + } +} + +/// `TFT` Configuration +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct TFTConfig { + // Model architecture + pub input_dim: usize, + pub hidden_dim: usize, + pub num_heads: usize, + pub num_layers: usize, + + // Forecasting parameters + pub prediction_horizon: usize, + pub sequence_length: usize, + pub num_quantiles: usize, + + // Feature types + pub num_static_features: usize, + pub num_known_features: usize, + pub num_unknown_features: usize, + + // Training parameters + pub learning_rate: f64, + pub batch_size: usize, + pub dropout_rate: f64, + pub l2_regularization: f64, + + // HFT optimization + pub use_flash_attention: bool, + pub mixed_precision: bool, + pub memory_efficient: bool, + + // Performance constraints + pub max_inference_latency_us: u64, + pub target_throughput_pps: u64, +} + +impl Default for TFTConfig { + fn default() -> Self { + Self { + // Wave C+D: 225 features (201 Wave C + 24 Wave D) + // Wave C: 201 features (indices 0-200) + // Wave D: 24 features (indices 201-224) + input_dim: 225, + hidden_dim: 128, + num_heads: 8, + num_layers: 3, + prediction_horizon: 10, + sequence_length: 50, + num_quantiles: 9, + // Feature split for 225 total features: + // - Static: 5 features (symbol metadata) + // - Known: 10 features (future time features) + // - Unknown: 210 features (historical OHLCV + technical + microstructure + regime) + num_static_features: 5, + num_known_features: 10, + num_unknown_features: 210, + learning_rate: 1e-3, + batch_size: 64, + dropout_rate: 0.1, + l2_regularization: 1e-4, + use_flash_attention: true, + mixed_precision: true, + memory_efficient: true, + max_inference_latency_us: 50, + target_throughput_pps: 100_000, + } + } +} + +/// `TFT` Model State for incremental processing +/// +/// **MEMORY SAFETY FIX (2025-10-25)**: +/// - Replaced unbounded HashMap with LRU cache (max 1000 entries) +/// - Prevents 3.6GB/hour memory leak in production inference +/// - Automatically evicts oldest cache entries when full +/// - Tested: 1-hour inference run with 10K predictions = stable 24MB memory +#[derive(Debug, Clone)] +pub struct TFTState { + pub hidden_state: Option, + pub attention_cache: LruCache, + pub last_update: u64, +} + +impl TFTState { + /// Maximum attention cache entries (2000 = ~48MB for TFT-225, 60% training speedup) + /// Chosen to balance: + /// - Memory safety: <100MB cache overhead (acceptable for training) + /// - Hit rate: >95% for typical 50-sequence inference + /// - Eviction overhead: <0.5% latency impact (reduced by 2x cache size) + pub const MAX_CACHE_ENTRIES: usize = 2000; + + pub fn zeros(_config: &TFTConfig) -> Result { + // MAX_CACHE_ENTRIES (2000) is non-zero by construction + let capacity = NonZeroUsize::new(Self::MAX_CACHE_ENTRIES) + .ok_or_else(|| MLError::ConfigError("MAX_CACHE_ENTRIES must be non-zero".to_owned()))?; + + Ok(Self { + hidden_state: None, + attention_cache: LruCache::new(capacity), + last_update: 0, + }) + } + + /// Clear attention cache to free memory + /// Call this after training/inference batch to prevent memory accumulation + pub fn clear_cache(&mut self) { + self.attention_cache.clear(); + self.hidden_state = None; + } +} + +/// `TFT` Model Metadata +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct TFTMetadata { + pub model_id: String, + pub version: String, + pub input_dim: usize, + pub output_dim: usize, + pub created_at: SystemTime, + pub last_trained: Option, + pub training_samples: u64, + pub performance_metrics: HashMap, +} + +/// Multi-horizon prediction result +#[derive(Debug, Clone)] +pub struct MultiHorizonPrediction { + pub predictions: Vec, // Point predictions for each horizon + pub quantiles: Vec>, // Quantile predictions [horizon][quantile] + pub uncertainty: Vec, // Uncertainty estimates + pub confidence_intervals: Vec<(f64, f64)>, // 90% confidence intervals + pub attention_weights: HashMap>, // Attention interpretability + pub feature_importance: Vec, // Variable importance scores + pub latency_us: u64, // Inference latency +} + +/// Complete Temporal Fusion Transformer +pub struct TemporalFusionTransformer { + pub config: TFTConfig, + pub metadata: TFTMetadata, + pub is_trained: bool, + + // Core TFT components (None when feature count is 0 — avoids zero-dim CUDA tensors) + static_variable_selection: Option, + historical_variable_selection: VariableSelectionNetwork, + future_variable_selection: Option, + + // Encoding layers (None when corresponding feature count is 0) + static_encoder: Option, + historical_encoder: GRNStack, + future_encoder: Option, + + // Temporal processing + lstm_encoder: Linear, // Simplified LSTM representation + lstm_decoder: Linear, + + // Attention mechanism + temporal_attention: TemporalSelfAttention, + + // Output layers + pub quantile_outputs: QuantileLayer, + + // Performance tracking + inference_count: AtomicU64, + total_latency_us: AtomicU64, + max_latency_us: AtomicU64, + + pub device: Device, + + // Variable map for checkpointing + pub varmap: Arc, +} + +impl std::fmt::Debug for TemporalFusionTransformer { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("TemporalFusionTransformer") + .field("config", &self.config) + .field("metadata", &self.metadata) + .field("is_trained", &self.is_trained) + .field( + "inference_count", + &self.inference_count.load(Ordering::Relaxed), + ) + .field( + "total_latency_us", + &self.total_latency_us.load(Ordering::Relaxed), + ) + .field( + "max_latency_us", + &self.max_latency_us.load(Ordering::Relaxed), + ) + .field("device", &format!("{:?}", self.device)) + .field("varmap", &"Arc") + .finish_non_exhaustive() + } +} + +impl TemporalFusionTransformer { + pub fn new(config: TFTConfig) -> Result { + Self::new_with_device(config, Device::cuda_if_available(0).unwrap_or(Device::Cpu)) + } + + pub fn new_with_device(config: TFTConfig, device: Device) -> Result { + // Validate configuration + let total_features = + config.num_static_features + config.num_known_features + config.num_unknown_features; + if config.num_unknown_features == 0 { + return Err(MLError::ConfigError("TFT requires num_unknown_features > 0 (temporal input)".to_owned())); + } + if total_features != config.input_dim { + return Err(MLError::ConfigError(format!( + "Feature count mismatch: static({}) + known({}) + unknown({}) = {} != input_dim({})", + config.num_static_features, + config.num_known_features, + config.num_unknown_features, + total_features, + config.input_dim + ))); + } + + // Log configuration for debugging + debug!( + "Creating TFT with {} input features (static: {}, known: {}, unknown: {})", + config.input_dim, + config.num_static_features, + config.num_known_features, + config.num_unknown_features + ); + + let varmap = Arc::new(VarMap::new()); + let vs = VarBuilder::from_varmap(&varmap, training_dtype(&device), &device); + + // Create variable selection networks (skip when feature count is 0 — CUDA + // cannot handle zero-dim tensors in linear layers) + let static_variable_selection = (config.num_static_features > 0) + .then(|| { + VariableSelectionNetwork::new( + config.num_static_features, + config.hidden_dim, + vs.pp("static_vsn"), + ) + }) + .transpose()?; + + let historical_variable_selection = VariableSelectionNetwork::new( + config.num_unknown_features, + config.hidden_dim, + vs.pp("historical_vsn"), + )?; + + let future_variable_selection = (config.num_known_features > 0) + .then(|| { + VariableSelectionNetwork::new( + config.num_known_features, + config.hidden_dim, + vs.pp("future_vsn"), + ) + }) + .transpose()?; + + // Create encoding stacks (skip when corresponding VSN is absent) + let static_encoder = (config.num_static_features > 0) + .then(|| { + GRNStack::new( + config.hidden_dim, + config.hidden_dim, + config.hidden_dim, + config.num_layers, + vs.pp("static_encoder"), + ) + }) + .transpose()?; + + let historical_encoder = GRNStack::new( + config.hidden_dim, + config.hidden_dim, + config.hidden_dim, + config.num_layers, + vs.pp("historical_encoder"), + )?; + + let future_encoder = (config.num_known_features > 0) + .then(|| { + GRNStack::new( + config.hidden_dim, + config.hidden_dim, + config.hidden_dim, + config.num_layers, + vs.pp("future_encoder"), + ) + }) + .transpose()?; + + // Simplified LSTM layers (in practice, would use proper LSTM) + let lstm_encoder = linear(config.hidden_dim, config.hidden_dim, vs.pp("lstm_encoder"))?; + let lstm_decoder = linear(config.hidden_dim, config.hidden_dim, vs.pp("lstm_decoder"))?; + + // Temporal attention + let temporal_attention = TemporalSelfAttention::new( + config.hidden_dim, + config.num_heads, + config.dropout_rate, + config.use_flash_attention, + vs.pp("temporal_attention"), + )?; + + // Quantile output layer + let quantile_outputs = QuantileLayer::new( + config.hidden_dim, + config.prediction_horizon, + config.num_quantiles, + vs.pp("quantile_outputs"), + )?; + + // Metadata + let metadata = TFTMetadata { + model_id: Uuid::new_v4().to_string(), + version: "1.0.0".to_owned(), + input_dim: config.input_dim, + output_dim: config.prediction_horizon, + created_at: SystemTime::now(), + last_trained: None, + training_samples: 0, + performance_metrics: HashMap::new(), + }; + + Ok(Self { + config, + metadata, + is_trained: false, + static_variable_selection, + historical_variable_selection, + future_variable_selection, + static_encoder, + historical_encoder, + future_encoder, + lstm_encoder, + lstm_decoder, + temporal_attention, + quantile_outputs, + inference_count: AtomicU64::new(0), + total_latency_us: AtomicU64::new(0), + max_latency_us: AtomicU64::new(0), + device, + varmap, + }) + } + + /// Get reference to the model's VarMap for checkpointing and quantization + pub fn varmap(&self) -> &Arc { + &self.varmap + } + + /// Get mutable reference to the model's VarMap for checkpoint loading + pub fn varmap_mut(&mut self) -> &mut Arc { + &mut self.varmap + } + + /// Get reference to the model's Device + pub fn device(&self) -> &Device { + &self.device + } + + /// Validate input tensor dimensions match configuration + fn validate_input_dimensions( + &self, + static_features: &Tensor, + historical_features: &Tensor, + future_features: &Tensor, + ) -> Result<(), MLError> { + // Validate static features: [batch, num_static_features] (skip if 0) + if self.config.num_static_features > 0 { + let static_dims = static_features.dims(); + if static_dims.len() != 2 { + return Err(MLError::ModelError(format!( + "Static features must be 2D [batch, features], got {} dimensions", + static_dims.len() + ))); + } + if static_dims[1] != self.config.num_static_features { + return Err(MLError::ModelError(format!( + "Static features dimension mismatch: expected {}, got {}", + self.config.num_static_features, static_dims[1] + ))); + } + } + + // Validate historical features: [batch, seq_len, num_unknown_features] + let hist_dims = historical_features.dims(); + if hist_dims.len() != 3 { + return Err(MLError::ModelError(format!( + "Historical features must be 3D [batch, seq, features], got {} dimensions", + hist_dims.len() + ))); + } + if hist_dims[2] != self.config.num_unknown_features { + return Err(MLError::ModelError(format!( + "Historical features dimension mismatch: expected {}, got {}", + self.config.num_unknown_features, + hist_dims[2] + ))); + } + + // Validate future features: [batch, horizon, num_known_features] (skip if 0) + if self.config.num_known_features > 0 { + let fut_dims = future_features.dims(); + if fut_dims.len() != 3 { + return Err(MLError::ModelError(format!( + "Future features must be 3D [batch, horizon, features], got {} dimensions", + fut_dims.len() + ))); + } + if fut_dims[2] != self.config.num_known_features { + return Err(MLError::ModelError(format!( + "Future features dimension mismatch: expected {}, got {}", + self.config.num_known_features, fut_dims[2] + ))); + } + } + + Ok(()) + } + + /// Forward pass through the complete `TFT` architecture + #[instrument(skip(self, static_features, historical_features, future_features))] + pub fn forward( + &mut self, + static_features: &Tensor, + historical_features: &Tensor, + future_features: &Tensor, + ) -> Result { + self.forward_with_checkpointing( + static_features, + historical_features, + future_features, + false, + ) + } + + /// Forward pass with optional gradient checkpointing + /// + /// When gradient checkpointing is enabled: + /// - Memory usage reduced by 30-40% (doesn't store intermediate activations) + /// - Training time increases by ~20% (recomputes activations during backprop) + /// + /// # Arguments + /// * `static_features` - Static input features [batch, num_static_features] + /// * `historical_features` - Historical features [batch, seq_len, num_unknown_features] + /// * `future_features` - Future features [batch, horizon, num_known_features] + /// * `use_checkpointing` - Whether to use gradient checkpointing + #[instrument(skip(self, static_features, historical_features, future_features))] + pub fn forward_with_checkpointing( + &mut self, + static_features: &Tensor, + historical_features: &Tensor, + future_features: &Tensor, + use_checkpointing: bool, + ) -> Result { + let static_features = ml_core::mixed_precision::ensure_training_dtype(static_features) + .map_err(|e| MLError::ModelError(e.to_string()))?; + let historical_features = ml_core::mixed_precision::ensure_training_dtype(historical_features) + .map_err(|e| MLError::ModelError(e.to_string()))?; + let future_features = ml_core::mixed_precision::ensure_training_dtype(future_features) + .map_err(|e| MLError::ModelError(e.to_string()))?; + let start_time = Instant::now(); + + // Validate input dimensions + self.validate_input_dimensions(&static_features, &historical_features, &future_features)?; + + // Log device placement for debugging + debug!("Forward pass device check:"); + debug!(" static_features: {:?}", static_features.device()); + debug!(" historical_features: {:?}", historical_features.device()); + debug!(" future_features: {:?}", future_features.device()); + debug!(" model device: {:?}", self.device); + + // 1. Variable Selection Networks (skip absent feature paths) + let static_encoded = if let Some(ref mut static_vsn) = self.static_variable_selection { + let static_selected = static_vsn.forward(&static_features, None)?; + let encoder = self.static_encoder.as_mut().ok_or_else(|| { + MLError::ModelError( + "static_encoder must exist when static_variable_selection exists".to_owned(), + ) + })?; + if use_checkpointing { + Some(encoder.forward(&static_selected.detach(), None)?) + } else { + Some(encoder.forward(&static_selected, None)?) + } + } else { + None + }; + + let historical_selected = self + .historical_variable_selection + .forward(&historical_features, None)?; + + let historical_encoded = if use_checkpointing { + self.historical_encoder + .forward(&historical_selected.detach(), None)? + } else { + self.historical_encoder + .forward(&historical_selected, None)? + }; + + let future_encoded = if let Some(ref mut future_vsn) = self.future_variable_selection { + let future_selected = future_vsn.forward(&future_features, None)?; + let encoder = self.future_encoder.as_mut().ok_or_else(|| { + MLError::ModelError( + "future_encoder must exist when future_variable_selection exists".to_owned(), + ) + })?; + if use_checkpointing { + Some(encoder.forward(&future_selected.detach(), None)?) + } else { + Some(encoder.forward(&future_selected, None)?) + } + } else { + None + }; + + // 3. Temporal Processing + let historical_temporal = if use_checkpointing { + self.lstm_encoder.forward(&historical_encoded.detach())? + } else { + self.lstm_encoder.forward(&historical_encoded)? + }; + + // 4. Combine temporal representations (skip future if absent) + let combined_temporal = if let Some(ref fut_enc) = future_encoded { + let future_temporal = if use_checkpointing { + self.lstm_decoder.forward(&fut_enc.detach())? + } else { + self.lstm_decoder.forward(fut_enc)? + }; + Tensor::cat(&[&historical_temporal, &future_temporal], 1)? + } else { + historical_temporal + }; + + // 5. Self-Attention + let attended = self.temporal_attention.forward_with_checkpointing( + &combined_temporal, + true, + use_checkpointing, + )?; + + // 6. Apply static context (skip if no static features) + let contextualized = if let Some(ref static_enc) = static_encoded { + self.apply_static_context(&attended, static_enc)? + } else { + attended + }; + + // 7. Quantile Outputs (no checkpointing on final layer) + let quantile_preds = self.quantile_outputs.forward(&contextualized)?; + + debug!(" quantile_preds: {:?}", quantile_preds.device()); + + // Cast output back to F32 for API compatibility + let quantile_preds = quantile_preds.to_dtype(candle_core::DType::F32) + .map_err(|e| MLError::ModelError(format!("Output dtype cast failed: {}", e)))?; + + // Update performance metrics + let latency = start_time.elapsed().as_micros() as u64; + self.update_performance_metrics(latency); + + Ok(quantile_preds) + } + + fn apply_static_context( + &self, + temporal: &Tensor, + static_context: &Tensor, + ) -> Result { + let (batch_size, seq_len, hidden_dim) = temporal.dims3()?; + + // Static context comes from variable selection + GRN encoding + // It has shape [batch, 1, hidden] (variable selection adds seq_len=1 dimension) + // We need to expand it to [batch, seq_len, hidden] to match temporal features + + // First, squeeze out the seq_len=1 dimension to get [batch, hidden] + let static_squeezed = static_context.squeeze(1)?; + + // Then expand to match sequence length using broadcast (zero-copy) + let static_expanded = static_squeezed + .unsqueeze(1)? // [batch, 1, hidden] + .broadcast_as((batch_size, seq_len, hidden_dim))?; // [batch, seq_len, hidden] - zero-copy broadcast + + // Add static context to temporal features + let contextualized = (temporal + &static_expanded)?; + + Ok(contextualized) + } + + /// Multi-horizon prediction interface + pub fn predict_horizons( + &mut self, + static_features: &Array1, + historical_features: &Array2, + future_features: &Array2, + ) -> Result { + if !self.is_trained { + return Err(MLError::ModelError("Model not trained".to_owned())); + } + + let start_time = Instant::now(); + + // Convert ndarray to tensors + let static_tensor = self.array_to_tensor_1d(static_features)?; + let historical_tensor = self.array_to_tensor_2d(historical_features)?; + let future_tensor = self.array_to_tensor_2d(future_features)?; + + // Add batch dimension + let static_batched = static_tensor.unsqueeze(0)?; + let historical_batched = historical_tensor.unsqueeze(0)?; + let future_batched = future_tensor.unsqueeze(0)?; + + // Forward pass + let quantile_preds = self.forward(&static_batched, &historical_batched, &future_batched)?; + + // Extract predictions and process outputs + let pred_data = quantile_preds.squeeze(0)?.to_vec2::()?; // [horizon, quantiles] + + let mut predictions = Vec::new(); + let mut quantiles = Vec::new(); + let mut uncertainty = Vec::new(); + let mut confidence_intervals = Vec::new(); + + for horizon in 0..self.config.prediction_horizon { + let horizon_quantiles = &pred_data[horizon]; + + // Point prediction (median) + let median_idx = self.config.num_quantiles / 2; + predictions.push(horizon_quantiles[median_idx] as f64); + + // All quantiles for this horizon + quantiles.push(horizon_quantiles.iter().map(|&x| x as f64).collect()); + + // Uncertainty (IQR) + let q75_idx = (self.config.num_quantiles * 3) / 4; + let q25_idx = self.config.num_quantiles / 4; + let iqr = horizon_quantiles[q75_idx] - horizon_quantiles[q25_idx]; + uncertainty.push(iqr as f64); + + // 90% confidence interval + let lower_idx = self.config.num_quantiles / 10; // ~10th percentile + let upper_idx = (self.config.num_quantiles * 9) / 10; // ~90th percentile + let ci = ( + horizon_quantiles[lower_idx] as f64, + horizon_quantiles[upper_idx] as f64, + ); + confidence_intervals.push(ci); + } + + // Get feature importance and attention weights + let feature_importance = self.static_variable_selection + .as_ref() + .map(|vsn| vsn.get_importance_scores()) + .transpose()? + .unwrap_or_default(); + let mut attention_weights = HashMap::new(); + let weights = self.temporal_attention.get_attention_weights(); + for (key, weight) in weights { + attention_weights.insert(key, vec![weight]); + } + + let latency = start_time.elapsed().as_micros() as u64; + + Ok(MultiHorizonPrediction { + predictions, + quantiles, + uncertainty, + confidence_intervals, + attention_weights, + feature_importance, + latency_us: latency, + }) + } + + fn array_to_tensor_1d(&self, arr: &Array1) -> Result { + let data: Vec = arr.iter().map(|&x| x as f32).collect(); + let tensor = Tensor::from_slice(&data, arr.len(), &self.device)?; + Ok(tensor) + } + + fn array_to_tensor_2d(&self, arr: &Array2) -> Result { + let data: Vec = arr.iter().map(|&x| x as f32).collect(); + let shape = arr.shape(); + let tensor = Tensor::from_slice(&data, (shape[0], shape[1]), &self.device)?; + Ok(tensor) + } + + fn update_performance_metrics(&self, latency_us: u64) { + self.inference_count.fetch_add(1, Ordering::Relaxed); + self.total_latency_us + .fetch_add(latency_us, Ordering::Relaxed); + + // Update max latency atomically + let mut current_max = self.max_latency_us.load(Ordering::Relaxed); + while latency_us > current_max { + match self.max_latency_us.compare_exchange_weak( + current_max, + latency_us, + Ordering::Relaxed, + Ordering::Relaxed, + ) { + Ok(_) => break, + Err(new_max) => current_max = new_max, + } + } + } + + /// Get performance metrics + pub fn get_metrics(&self) -> HashMap { + let inference_count = self.inference_count.load(Ordering::Relaxed); + let total_latency = self.total_latency_us.load(Ordering::Relaxed); + let max_latency = self.max_latency_us.load(Ordering::Relaxed); + + let avg_latency = if inference_count > 0 { + total_latency as f64 / inference_count as f64 + } else { + 0.0 + }; + + let throughput = if avg_latency > 0.0 { + 1_000_000.0 / avg_latency // predictions per second + } else { + 0.0 + }; + + let mut metrics = HashMap::new(); + metrics.insert("total_inferences".to_owned(), inference_count as f64); + metrics.insert("avg_latency_us".to_owned(), avg_latency); + metrics.insert("max_latency_us".to_owned(), max_latency as f64); + metrics.insert("throughput_pps".to_owned(), throughput); + + metrics + } + + /// Get reference to VarMap for weight extraction + pub fn get_varmap(&self) -> &Arc { + &self.varmap + } + + /// Training interface with real backward pass and optimizer + pub async fn train( + &mut self, + training_data: &[(Array1, Array2, Array2, Array1)], // (static, historical, future, targets) + validation_data: &[(Array1, Array2, Array2, Array1)], + epochs: usize, + ) -> Result<(), MLError> { + info!("Starting TFT training for {} epochs", epochs); + + // Initialize AdamW optimizer with model parameters + let params = self.varmap.all_vars(); + let num_params: usize = params.iter().map(|v| v.as_tensor().elem_count()).sum(); + let lr = 1e-3; + let mut optimizer = AdamW::new( + params, + ParamsAdamW { + lr, + beta1: 0.9, + beta2: 0.999, + eps: 1e-8, + weight_decay: 1e-4, + }, + ) + .map_err(|e| MLError::TrainingError(format!("Failed to create optimizer: {}", e)))?; + + info!( + "Initialized AdamW optimizer: lr={:.2e}, {} parameters", + lr, num_params + ); + + let mut best_val_loss = f64::MAX; + + for epoch in 0..epochs { + let mut epoch_loss = 0.0; + + for (_i, (static_feat, hist_feat, fut_feat, targets)) in + training_data.iter().enumerate() + { + // Convert to tensors + let static_tensor = self.array_to_tensor_1d(static_feat)?.unsqueeze(0)?; + let hist_tensor = self.array_to_tensor_2d(hist_feat)?.unsqueeze(0)?; + let fut_tensor = self.array_to_tensor_2d(fut_feat)?.unsqueeze(0)?; + let target_tensor = self.array_to_tensor_1d(targets)?.unsqueeze(0)?; + + // Forward pass + let predictions = self.forward(&static_tensor, &hist_tensor, &fut_tensor)?; + + // Compute quantile loss + let loss = self + .quantile_outputs + .quantile_loss(&predictions, &target_tensor)?; + epoch_loss += loss.to_vec0::()? as f64; + + // Backward pass -- compute gradients + let grads = loss.backward().map_err(|e| { + MLError::TrainingError(format!("Backward pass failed: {}", e)) + })?; + + // Check gradient health before stepping + let varmap_data = self.varmap.data().lock().map_err(|e| { + MLError::TrainingError(format!("Failed to lock VarMap: {}", e)) + })?; + let mut grad_norm_sq = 0.0_f64; + for (_name, var) in varmap_data.iter() { + if let Some(grad) = grads.get(var.as_tensor()) { + let norm_sq = grad + .sqr() + .and_then(|t| t.sum_all()) + .and_then(|t| t.to_dtype(candle_core::DType::F64)) + .and_then(|t| t.to_scalar::()) + .unwrap_or(0.0); + grad_norm_sq += norm_sq; + } + } + drop(varmap_data); + let grad_norm = grad_norm_sq.sqrt(); + + if grad_norm.is_nan() || grad_norm.is_infinite() { + warn!( + "Gradient explosion detected (norm={}), skipping update", + grad_norm + ); + continue; + } + + // Optimizer step -- update parameters + optimizer.step(&grads).map_err(|e| { + MLError::TrainingError(format!("Optimizer step failed: {}", e)) + })?; + } + + let avg_epoch_loss = epoch_loss / training_data.len().max(1) as f64; + debug!("Epoch {}: Average Loss = {:.6}", epoch, avg_epoch_loss); + + // Validation every 10 epochs + if epoch % 10 == 0 { + let val_loss = self.validate(validation_data).await?; + info!( + "Epoch {}: Train Loss = {:.6}, Val Loss = {:.6}", + epoch, avg_epoch_loss, val_loss + ); + if val_loss < best_val_loss { + best_val_loss = val_loss; + } + } + } + + self.is_trained = true; + self.metadata.last_trained = Some(SystemTime::now()); + self.metadata.training_samples = training_data.len() as u64; + + info!( + "TFT training completed: {} epochs, best val loss = {:.6}", + epochs, best_val_loss + ); + Ok(()) + } + + async fn validate( + &mut self, + validation_data: &[(Array1, Array2, Array2, Array1)], + ) -> Result { + let mut total_loss = 0.0; + + for (static_feat, hist_feat, fut_feat, targets) in validation_data { + let static_tensor = self.array_to_tensor_1d(static_feat)?.unsqueeze(0)?; + let hist_tensor = self.array_to_tensor_2d(hist_feat)?.unsqueeze(0)?; + let fut_tensor = self.array_to_tensor_2d(fut_feat)?.unsqueeze(0)?; + let target_tensor = self.array_to_tensor_1d(targets)?.unsqueeze(0)?; + + let predictions = self.forward(&static_tensor, &hist_tensor, &fut_tensor)?; + let loss = self + .quantile_outputs + .quantile_loss(&predictions, &target_tensor)?; + total_loss += loss.to_vec0::()? as f64; + } + + Ok(total_loss / validation_data.len() as f64) + } + + /// Compute quantile loss for training + pub fn compute_quantile_loss( + &self, + predictions: &Tensor, + targets: &Tensor, + ) -> Result { + self.quantile_outputs.quantile_loss(predictions, targets) + } + + /// HFT-optimized inference + pub fn predict_fast( + &mut self, + static_features: &[f32], + historical_features: &[f32], + future_features: &[f32], + ) -> Result, MLError> { + let start = Instant::now(); + + // Convert to tensors (optimized path) + let static_tensor = + Tensor::from_slice(static_features, static_features.len(), &self.device)? + .unsqueeze(0)?; + + let hist_len = self.config.sequence_length; + let hist_dim = self.config.num_unknown_features; + let historical_tensor = + Tensor::from_slice(historical_features, (hist_len, hist_dim), &self.device)? + .unsqueeze(0)?; + + let fut_len = self.config.prediction_horizon; + let fut_dim = self.config.num_known_features; + let future_tensor = + Tensor::from_slice(future_features, (fut_len, fut_dim), &self.device)?.unsqueeze(0)?; + + // Forward pass + let quantile_preds = self.forward(&static_tensor, &historical_tensor, &future_tensor)?; + + // Extract median predictions + let pred_data = quantile_preds.squeeze(0)?.to_vec2::()?; + let median_idx = self.config.num_quantiles / 2; + let predictions: Vec = pred_data + .iter() + .map(|horizon_quantiles| horizon_quantiles[median_idx]) + .collect(); + + let latency = start.elapsed().as_micros() as u64; + self.update_performance_metrics(latency); + + if latency > self.config.max_inference_latency_us { + warn!( + "Inference latency {}μs exceeds target {}μs", + latency, self.config.max_inference_latency_us + ); + } + + Ok(predictions) + } +} + +// Architecture info helpers used by the checkpoint bridge in ml crate +impl TemporalFusionTransformer { + /// Get hyperparameters for checkpoint metadata + pub fn get_hyperparameters(&self) -> HashMap { + let mut params = HashMap::new(); + // Core architecture params (Wave C+D: 225 features) + params.insert("input_dim".to_owned(), Value::from(self.config.input_dim)); + params.insert( + "hidden_dim".to_owned(), + Value::from(self.config.hidden_dim), + ); + params.insert("num_heads".to_owned(), Value::from(self.config.num_heads)); + params.insert( + "num_layers".to_owned(), + Value::from(self.config.num_layers), + ); + params.insert( + "prediction_horizon".to_owned(), + Value::from(self.config.prediction_horizon), + ); + params.insert( + "sequence_length".to_owned(), + Value::from(self.config.sequence_length), + ); + params.insert( + "num_quantiles".to_owned(), + Value::from(self.config.num_quantiles), + ); + + // Feature split (critical for Wave C+D compatibility) + params.insert( + "num_static_features".to_owned(), + Value::from(self.config.num_static_features), + ); + params.insert( + "num_known_features".to_owned(), + Value::from(self.config.num_known_features), + ); + params.insert( + "num_unknown_features".to_owned(), + Value::from(self.config.num_unknown_features), + ); + + // Training params + params.insert( + "learning_rate".to_owned(), + Value::from(self.config.learning_rate), + ); + params.insert( + "batch_size".to_owned(), + Value::from(self.config.batch_size), + ); + params.insert( + "dropout_rate".to_owned(), + Value::from(self.config.dropout_rate), + ); + params.insert( + "l2_regularization".to_owned(), + Value::from(self.config.l2_regularization), + ); + + // HFT optimization flags + params.insert( + "use_flash_attention".to_owned(), + Value::from(self.config.use_flash_attention), + ); + params.insert( + "mixed_precision".to_owned(), + Value::from(self.config.mixed_precision), + ); + params.insert( + "memory_efficient".to_owned(), + Value::from(self.config.memory_efficient), + ); + + params + } + + pub fn get_architecture_info(&self) -> HashMap { + let mut info = HashMap::new(); + info.insert("network_type".to_owned(), Value::from("TFT")); + info.insert( + "input_dim".to_owned(), + Value::from(self.metadata.input_dim), + ); + info.insert( + "output_dim".to_owned(), + Value::from(self.metadata.output_dim), + ); + info.insert( + "hidden_dim".to_owned(), + Value::from(self.config.hidden_dim), + ); + info.insert("num_heads".to_owned(), Value::from(self.config.num_heads)); + info.insert( + "num_layers".to_owned(), + Value::from(self.config.num_layers), + ); + info.insert( + "num_static_features".to_owned(), + Value::from(self.config.num_static_features), + ); + info.insert( + "num_known_features".to_owned(), + Value::from(self.config.num_known_features), + ); + info.insert( + "num_unknown_features".to_owned(), + Value::from(self.config.num_unknown_features), + ); + info + } +} + +#[cfg(test)] +mod tests { + use super::*; + use anyhow::Result; + use candle_core::DType; + + #[tokio::test] + async fn test_tft_creation() -> Result<()> { + let config = TFTConfig { + input_dim: 10, + hidden_dim: 32, + num_heads: 4, + num_quantiles: 5, + prediction_horizon: 5, + sequence_length: 20, + num_static_features: 2, + num_known_features: 3, + num_unknown_features: 5, + ..Default::default() + }; + + let tft = TemporalFusionTransformer::new(config) + .map_err(|_| anyhow::anyhow!("Failed to create TFT"))?; + assert_eq!(tft.metadata.input_dim, 10); + assert_eq!(tft.metadata.output_dim, 5); + Ok(()) + } + + #[test] + fn test_tft_225_features_default() -> Result<()> { + // Test default configuration uses 225 features (Wave C+D) + let config = TFTConfig::default(); + assert_eq!( + config.input_dim, 225, + "Default TFT config should use 225 features" + ); + assert_eq!(config.num_static_features, 5); + assert_eq!(config.num_known_features, 10); + assert_eq!(config.num_unknown_features, 210); + + let tft = TemporalFusionTransformer::new(config) + .map_err(|_| anyhow::anyhow!("Failed to create TFT with 225 features"))?; + assert_eq!(tft.metadata.input_dim, 225); + Ok(()) + } + + #[test] + fn test_tft_225_features_validation() -> Result<()> { + // Test that 225-feature TFT validates input dimensions correctly + let config = TFTConfig::default(); // 225 features + let device = Device::Cpu; + let tft = TemporalFusionTransformer::new_with_device(config.clone(), device.clone()) + .map_err(|_| anyhow::anyhow!("Failed to create TFT"))?; + + // Create valid input tensors + let batch_size = 2; + let seq_len = 50; + let horizon = 10; + + let static_features = Tensor::zeros( + (batch_size, config.num_static_features), + DType::F32, + &device, + )?; + let historical_features = Tensor::zeros( + (batch_size, seq_len, config.num_unknown_features), + DType::F32, + &device, + )?; + let future_features = Tensor::zeros( + (batch_size, horizon, config.num_known_features), + DType::F32, + &device, + )?; + + // Should validate successfully + let result = + tft.validate_input_dimensions(&static_features, &historical_features, &future_features); + assert!( + result.is_ok(), + "Valid 225-feature input should pass validation" + ); + + // Test invalid historical features dimension + let invalid_hist = Tensor::zeros((batch_size, seq_len, 50), DType::F32, &device)?; // Wrong dim: 50 instead of 210 + let result = + tft.validate_input_dimensions(&static_features, &invalid_hist, &future_features); + assert!( + result.is_err(), + "Invalid historical features should fail validation" + ); + + Ok(()) + } + + #[test] + fn test_tft_config_mismatch_detection() -> Result<()> { + // Test that mismatched feature counts are detected during construction + let invalid_config = TFTConfig { + input_dim: 225, + num_static_features: 5, + num_known_features: 10, + num_unknown_features: 100, // Wrong: should be 210 for 225 total + ..Default::default() + }; + + let result = TemporalFusionTransformer::new(invalid_config); + assert!( + result.is_err(), + "Mismatched feature counts should be rejected" + ); + + let err_msg = format!("{:?}", result.unwrap_err()); + assert!( + err_msg.contains("Feature count mismatch"), + "Error should mention feature count mismatch" + ); + + Ok(()) + } + + #[test] + fn test_tft_checkpoint_preserves_config() -> Result<()> { + // Test that checkpoint save/load preserves 225-feature configuration + let config = TFTConfig::default(); // 225 features + let tft = TemporalFusionTransformer::new(config.clone()) + .map_err(|_| anyhow::anyhow!("Failed to create TFT"))?; + + let hyperparams = tft.get_hyperparameters(); + + // Verify all critical config params are saved + assert_eq!( + hyperparams.get("input_dim").and_then(|v| v.as_u64()), + Some(225) + ); + assert_eq!( + hyperparams + .get("num_static_features") + .and_then(|v| v.as_u64()), + Some(5) + ); + assert_eq!( + hyperparams + .get("num_known_features") + .and_then(|v| v.as_u64()), + Some(10) + ); + assert_eq!( + hyperparams + .get("num_unknown_features") + .and_then(|v| v.as_u64()), + Some(210) + ); + + Ok(()) + } + + #[test] + fn test_tft_wave_c_config() -> Result<()> { + // Test Wave C configuration (201 features) + let wave_c_config = TFTConfig { + input_dim: 201, + num_static_features: 5, + num_known_features: 10, + num_unknown_features: 186, // 201 - 5 - 10 = 186 + ..Default::default() + }; + + let tft = TemporalFusionTransformer::new(wave_c_config) + .map_err(|_| anyhow::anyhow!("Failed to create TFT with 201 features"))?; + assert_eq!(tft.metadata.input_dim, 201); + + Ok(()) + } + + #[test] + fn test_tft_state_creation() -> Result<()> { + let config = TFTConfig { + hidden_dim: 32, + sequence_length: 20, + num_heads: 4, + ..Default::default() + }; + + let state = + TFTState::zeros(&config).map_err(|_| anyhow::anyhow!("Failed to create state"))?; + assert!(state.last_update == 0); + Ok(()) + } + + #[test] + fn test_tft_config_default() -> Result<()> { + let config = TFTConfig::default(); + assert!(config.input_dim > 0); + assert!(config.hidden_dim > 0); + assert!(config.num_heads > 0); + Ok(()) + } + + #[test] + fn test_tft_performance_metrics() -> Result<()> { + let config = TFTConfig { + input_dim: 30, + hidden_dim: 32, + num_static_features: 5, + num_known_features: 10, + num_unknown_features: 15, // 30 - 5 - 10 = 15 + ..Default::default() + }; + + let tft = TemporalFusionTransformer::new(config) + .map_err(|_| anyhow::anyhow!("Failed to create TFT"))?; + let metrics = tft.get_metrics(); + + assert!(metrics.contains_key("total_inferences")); + assert!(metrics.contains_key("avg_latency_us")); + assert!(metrics.contains_key("max_latency_us")); + assert!(metrics.contains_key("throughput_pps")); + Ok(()) + } + + #[test] + fn test_tft_training_state() -> Result<()> { + let config = TFTConfig::default(); + let mut tft = TemporalFusionTransformer::new(config) + .map_err(|_| anyhow::anyhow!("Failed to create TFT"))?; + + assert!(!tft.is_trained); + tft.is_trained = true; + assert!(tft.is_trained); + Ok(()) + } + + #[test] + fn test_tft_metadata() -> Result<()> { + let config = TFTConfig { + input_dim: 30, + prediction_horizon: 12, + num_static_features: 5, + num_known_features: 10, + num_unknown_features: 15, // 30 - 5 - 10 = 15 + ..Default::default() + }; + + let tft = TemporalFusionTransformer::new(config) + .map_err(|_| anyhow::anyhow!("Failed to create TFT"))?; + assert_eq!(tft.metadata.input_dim, 30); + assert_eq!(tft.metadata.output_dim, 12); + Ok(()) + } +} diff --git a/crates/ml/src/tft/qat_tft.rs b/crates/ml-supervised/src/tft/qat_tft.rs similarity index 99% rename from crates/ml/src/tft/qat_tft.rs rename to crates/ml-supervised/src/tft/qat_tft.rs index 8ec7165d4..f0170596c 100644 --- a/crates/ml/src/tft/qat_tft.rs +++ b/crates/ml-supervised/src/tft/qat_tft.rs @@ -43,7 +43,7 @@ //! ``` use crate::tft::{QuantizedTemporalFusionTransformer, TemporalFusionTransformer}; -use crate::MLError; +use ml_core::MLError; use candle_core::DeviceLocation; use candle_core::{Device, Tensor}; use std::collections::HashMap; diff --git a/crates/ml/src/tft/quantile_outputs.rs b/crates/ml-supervised/src/tft/quantile_outputs.rs similarity index 99% rename from crates/ml/src/tft/quantile_outputs.rs rename to crates/ml-supervised/src/tft/quantile_outputs.rs index a63a43634..9987e1dab 100644 --- a/crates/ml/src/tft/quantile_outputs.rs +++ b/crates/ml-supervised/src/tft/quantile_outputs.rs @@ -6,7 +6,7 @@ use candle_core::{Module, Tensor}; use candle_nn::{linear, Linear, VarBuilder}; -use crate::MLError; +use ml_core::MLError; /// Quantile output layer for uncertainty estimation #[derive(Debug, Clone)] diff --git a/crates/ml/src/tft/quantized_attention.rs b/crates/ml-supervised/src/tft/quantized_attention.rs similarity index 99% rename from crates/ml/src/tft/quantized_attention.rs rename to crates/ml-supervised/src/tft/quantized_attention.rs index c03c225d1..e029d67cb 100644 --- a/crates/ml/src/tft/quantized_attention.rs +++ b/crates/ml-supervised/src/tft/quantized_attention.rs @@ -5,10 +5,10 @@ //! Supports optional causal masking for autoregressive prediction. //! Provides optional weight caching (4x memory for 2-3x speed improvement). -use crate::memory_optimization::quantization::{ +use ml_core::memory_optimization::quantization::{ QuantizationConfig, QuantizationType, QuantizedTensor, Quantizer, }; -use crate::MLError; +use ml_core::MLError; use candle_core::{Device, Tensor}; use candle_nn::VarBuilder; use std::collections::HashMap; @@ -409,7 +409,7 @@ impl QuantizedTemporalAttention { #[cfg(test)] mod tests { use super::*; - use crate::dqn::mixed_precision::training_dtype; + use ml_core::mixed_precision::training_dtype; fn create_test_attention() -> QuantizedTemporalAttention { let device = Device::Cpu; diff --git a/crates/ml/src/tft/quantized_grn.rs b/crates/ml-supervised/src/tft/quantized_grn.rs similarity index 97% rename from crates/ml/src/tft/quantized_grn.rs rename to crates/ml-supervised/src/tft/quantized_grn.rs index c6298f6e8..cdcbf345f 100644 --- a/crates/ml/src/tft/quantized_grn.rs +++ b/crates/ml-supervised/src/tft/quantized_grn.rs @@ -6,10 +6,10 @@ use candle_core::{Device, Tensor}; use tracing::debug; -use crate::cuda_compat::manual_sigmoid; -use crate::memory_optimization::quantization::{QuantizationType, QuantizedTensor, Quantizer}; +use ml_core::cuda_compat::manual_sigmoid; +use ml_core::memory_optimization::quantization::{QuantizationType, QuantizedTensor, Quantizer}; use crate::tft::gated_residual::GatedResidualNetwork; -use crate::MLError; +use ml_core::MLError; /// Quantized Gated Residual Network /// @@ -227,7 +227,7 @@ impl QuantizedGatedResidualNetwork { /// Apply layer normalization (kept in F32) fn apply_layer_norm(&self, x: &Tensor) -> Result { if let Some(ln) = &self.layer_norm { - crate::cuda_compat::layer_norm_with_fallback( + ml_core::cuda_compat::layer_norm_with_fallback( x, &ln.normalized_shape, ln.weight.as_ref(), @@ -283,8 +283,8 @@ impl QuantizedGatedResidualNetwork { #[cfg(test)] mod tests { use super::*; - use crate::dqn::mixed_precision::training_dtype; - use crate::memory_optimization::quantization::QuantizationConfig; + use ml_core::mixed_precision::training_dtype; + use ml_core::memory_optimization::quantization::QuantizationConfig; use candle_nn::{VarBuilder, VarMap}; use std::sync::Arc; diff --git a/crates/ml/src/tft/quantized_lstm.rs b/crates/ml-supervised/src/tft/quantized_lstm.rs similarity index 98% rename from crates/ml/src/tft/quantized_lstm.rs rename to crates/ml-supervised/src/tft/quantized_lstm.rs index 31caada32..33270d855 100644 --- a/crates/ml/src/tft/quantized_lstm.rs +++ b/crates/ml-supervised/src/tft/quantized_lstm.rs @@ -18,9 +18,9 @@ use candle_core::{Device, Tensor}; use std::collections::HashMap; -use crate::cuda_compat::manual_sigmoid; -use crate::memory_optimization::quantization::{QuantizationConfig, QuantizedTensor, Quantizer}; -use crate::MLError; +use ml_core::cuda_compat::manual_sigmoid; +use ml_core::memory_optimization::quantization::{QuantizationConfig, QuantizedTensor, Quantizer}; +use ml_core::MLError; use super::lstm_encoder::LSTMEncoder; @@ -405,11 +405,11 @@ impl QuantizedLSTMEncoder { #[cfg(test)] mod tests { use super::*; - use crate::memory_optimization::quantization::QuantizationType; + use ml_core::memory_optimization::quantization::QuantizationType; #[test] fn test_quantized_lstm_creation() -> anyhow::Result<()> { - use crate::dqn::mixed_precision::training_dtype; + use ml_core::mixed_precision::training_dtype; use candle_nn::{VarBuilder, VarMap}; let device = Device::Cpu; @@ -434,7 +434,7 @@ mod tests { #[test] fn test_memory_reduction() -> anyhow::Result<()> { - use crate::dqn::mixed_precision::training_dtype; + use ml_core::mixed_precision::training_dtype; use candle_nn::{VarBuilder, VarMap}; let device = Device::Cpu; diff --git a/crates/ml/src/tft/quantized_tft.rs b/crates/ml-supervised/src/tft/quantized_tft.rs similarity index 99% rename from crates/ml/src/tft/quantized_tft.rs rename to crates/ml-supervised/src/tft/quantized_tft.rs index 2daa21f86..fed71fd25 100644 --- a/crates/ml/src/tft/quantized_tft.rs +++ b/crates/ml-supervised/src/tft/quantized_tft.rs @@ -4,12 +4,12 @@ //! Currently returns zero-initialized tensors for compatibility. //! Full quantization logic planned for future optimization (Wave 9.12+). -use crate::cuda_compat::manual_sigmoid; -use crate::memory_optimization::quantization::{ +use ml_core::cuda_compat::manual_sigmoid; +use ml_core::memory_optimization::quantization::{ QuantizationConfig, QuantizationType, QuantizedTensor, Quantizer, }; use crate::tft::TFTConfig; -use crate::MLError; +use ml_core::MLError; use candle_core::{Device, Tensor}; use candle_nn::VarMap; use std::collections::HashMap; diff --git a/crates/ml/src/tft/quantized_vsn.rs b/crates/ml-supervised/src/tft/quantized_vsn.rs similarity index 97% rename from crates/ml/src/tft/quantized_vsn.rs rename to crates/ml-supervised/src/tft/quantized_vsn.rs index 2058e13ac..f61e999ce 100644 --- a/crates/ml/src/tft/quantized_vsn.rs +++ b/crates/ml-supervised/src/tft/quantized_vsn.rs @@ -12,14 +12,14 @@ use candle_nn::{VarBuilder, VarMap}; use tracing::{debug, info}; use super::variable_selection::VariableSelectionNetwork; -use crate::dqn::mixed_precision::training_dtype; +use ml_core::mixed_precision::training_dtype; #[cfg(test)] -use crate::memory_optimization::quantization::{ +use ml_core::memory_optimization::quantization::{ QuantizationConfig, QuantizationType, QuantizedTensor, Quantizer, }; #[cfg(not(test))] -use crate::memory_optimization::quantization::{QuantizationConfig, QuantizedTensor, Quantizer}; -use crate::MLError; +use ml_core::memory_optimization::quantization::{QuantizationConfig, QuantizedTensor, Quantizer}; +use ml_core::MLError; /// Quantized Variable Selection Network /// @@ -242,7 +242,7 @@ impl QuantizedVariableSelectionNetwork { #[cfg(test)] mod tests { use super::*; - use crate::dqn::mixed_precision::training_dtype; + use ml_core::mixed_precision::training_dtype; #[test] fn test_quantized_vsn_creation() -> Result<(), MLError> { diff --git a/crates/ml/src/tft/temporal_attention.rs b/crates/ml-supervised/src/tft/temporal_attention.rs similarity index 99% rename from crates/ml/src/tft/temporal_attention.rs rename to crates/ml-supervised/src/tft/temporal_attention.rs index 725cb0331..8fe6a883f 100644 --- a/crates/ml/src/tft/temporal_attention.rs +++ b/crates/ml-supervised/src/tft/temporal_attention.rs @@ -20,8 +20,8 @@ use candle_core::{Device, Module, Tensor}; use candle_nn::{linear, Dropout, Linear, VarBuilder}; use tracing::{instrument, warn}; -use crate::cuda_compat::layer_norm_with_fallback; -use crate::MLError; +use ml_core::cuda_compat::layer_norm_with_fallback; +use ml_core::MLError; /// CUDA-compatible LayerNorm wrapper for TFT #[derive(Debug, Clone)] diff --git a/crates/ml/src/tft/variable_selection.rs b/crates/ml-supervised/src/tft/variable_selection.rs similarity index 99% rename from crates/ml/src/tft/variable_selection.rs rename to crates/ml-supervised/src/tft/variable_selection.rs index 671e511e5..f5e621155 100644 --- a/crates/ml/src/tft/variable_selection.rs +++ b/crates/ml-supervised/src/tft/variable_selection.rs @@ -9,7 +9,7 @@ use candle_core::{Device, Module, Tensor}; use candle_nn::{linear, Linear, VarBuilder}; use super::GatedResidualNetwork; -use crate::MLError; +use ml_core::MLError; /// Variable Selection Network for feature importance learning #[derive(Debug, Clone)] diff --git a/crates/ml/src/tft/varmap_quantization.rs b/crates/ml-supervised/src/tft/varmap_quantization.rs similarity index 99% rename from crates/ml/src/tft/varmap_quantization.rs rename to crates/ml-supervised/src/tft/varmap_quantization.rs index f93607cd3..ce8c6fee0 100644 --- a/crates/ml/src/tft/varmap_quantization.rs +++ b/crates/ml-supervised/src/tft/varmap_quantization.rs @@ -13,8 +13,8 @@ //! - Special case handling: small tensors, bias terms, LayerNorm params //! - Parallel processing with Rayon (optional) -use crate::memory_optimization::quantization::{QuantizationType, QuantizedTensor, Quantizer}; -use crate::MLError; +use ml_core::memory_optimization::quantization::{QuantizationType, QuantizedTensor, Quantizer}; +use ml_core::MLError; use candle_core::{DType, Device, Tensor}; use candle_nn::VarMap; use rayon::prelude::*; @@ -319,7 +319,7 @@ pub fn quantize_varmap_parallel( // Create thread-safe quantizer (one per thread via Arc) let quantizer = Arc::new(Mutex::new(Quantizer::new( - crate::memory_optimization::quantization::QuantizationConfig { + ml_core::memory_optimization::quantization::QuantizationConfig { quant_type: QuantizationType::Int8, symmetric: true, per_channel: false, @@ -662,8 +662,8 @@ pub fn load_quantized_weights( #[cfg(test)] mod tests { use super::*; - use crate::dqn::mixed_precision::training_dtype; - use crate::memory_optimization::quantization::{QuantizationConfig, Quantizer}; + use ml_core::mixed_precision::training_dtype; + use ml_core::memory_optimization::quantization::{QuantizationConfig, Quantizer}; use candle_core::Var; use candle_nn::VarBuilder; diff --git a/crates/ml/src/tgnn/gating.rs b/crates/ml-supervised/src/tgnn/gating.rs similarity index 99% rename from crates/ml/src/tgnn/gating.rs rename to crates/ml-supervised/src/tgnn/gating.rs index bb0b3a499..8d52cad64 100644 --- a/crates/ml/src/tgnn/gating.rs +++ b/crates/ml-supervised/src/tgnn/gating.rs @@ -5,7 +5,7 @@ use ndarray::{s, Array1, Array2, Axis}; use serde::{Deserialize, Serialize}; -use crate::MLError; +use ml_core::MLError; use rand::prelude::*; // Replace common::rng with standard rand /// Gradients for attention mechanism components diff --git a/crates/ml/src/tgnn/graph.rs b/crates/ml-supervised/src/tgnn/graph.rs similarity index 99% rename from crates/ml/src/tgnn/graph.rs rename to crates/ml-supervised/src/tgnn/graph.rs index ddcdf3d8b..82acad5b4 100644 --- a/crates/ml/src/tgnn/graph.rs +++ b/crates/ml-supervised/src/tgnn/graph.rs @@ -12,7 +12,7 @@ use serde::{Deserialize, Serialize}; use tracing::debug; use super::{MarketEdge, NodeId, NodeType}; -use crate::{MLError, PRECISION_FACTOR}; +use ml_core::{MLError, PRECISION_FACTOR}; /// Graph statistics #[derive(Debug, Clone, Default, Serialize, Deserialize)] diff --git a/crates/ml/src/tgnn/message_passing.rs b/crates/ml-supervised/src/tgnn/message_passing.rs similarity index 99% rename from crates/ml/src/tgnn/message_passing.rs rename to crates/ml-supervised/src/tgnn/message_passing.rs index 099772c47..b71a5d462 100644 --- a/crates/ml/src/tgnn/message_passing.rs +++ b/crates/ml-supervised/src/tgnn/message_passing.rs @@ -5,7 +5,7 @@ use ndarray::{s, Array1, Array2}; use serde::{Deserialize, Serialize}; -use crate::MLError; +use ml_core::MLError; use rand::prelude::*; // Replace common::rng with standard rand /// Cache for forward pass computations needed for backpropagation diff --git a/crates/ml-supervised/src/tgnn/mod.rs b/crates/ml-supervised/src/tgnn/mod.rs new file mode 100644 index 000000000..2844b701f --- /dev/null +++ b/crates/ml-supervised/src/tgnn/mod.rs @@ -0,0 +1,1310 @@ +//! # Temporal Graph Gated Networks (TGNN) for HFT +//! +//! Ultra-low latency implementation of TGNN for market microstructure analysis. +//! +//! ## Key Features +//! +//! - Sub-1μs graph neural network inference +//! - Real-time order book graph construction +//! - Market maker and liquidity flow modeling +//! - Cache-friendly graph operations +//! - Integer arithmetic for precision +//! +//! ## Performance Targets +//! +//! - Graph construction: <500ns from order book +//! - GNN inference: <1μs per prediction +//! - Node updates: <100ns per update +//! - Memory: Minimal allocations + +// Module imports +pub mod gating; +pub mod graph; +pub mod message_passing; +pub mod traits; +pub mod types; + +// DO NOT RE-EXPORT - Use explicit imports at usage sites + +// Import types from main crate - this fixes the circular dependency +use ml_core::{InferenceResult, MLError, ModelMetadata, ModelType, PRECISION_FACTOR}; +// Import traits from the local traits module +use traits::MLModel; +// Import types from this module +use types::{TrainingMetrics, ValidationMetrics}; + +// Import TGNN component types from submodules +use gating::GatingMechanism; +use graph::MarketGraph; +use message_passing::MessagePassing; + +use std::collections::HashMap; +use std::sync::atomic::{AtomicU64, Ordering}; +use std::time::{Instant, SystemTime}; + +// Import RNG utilities from types crate +use rand::prelude::*; // Replace common::rng with standard rand + +use async_trait::async_trait; +use dashmap::DashMap; +use ndarray::{s, Array1, Array2}; +use serde::{Deserialize, Serialize}; +use tracing::{debug, info, warn}; + +/// Node types in market microstructure graph +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)] +pub enum NodeType { + /// Price level in order book + PriceLevel, + /// Market maker entity + MarketMaker, + /// Liquidity pool + LiquidityPool, + /// Order cluster + OrderCluster, +} + +/// Edge types representing market relationships +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)] +pub enum EdgeType { + /// Price proximity relationship + PriceProximity, + /// Liquidity flow + LiquidityFlow, + /// Market maker connection + MarketMaking, + /// Order correlation + OrderCorrelation, +} + +/// Market node identifier +#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)] +pub struct NodeId { + pub node_type: NodeType, + pub id: String, +} + +impl NodeId { + pub fn price_level(price: i64) -> Self { + Self { + node_type: NodeType::PriceLevel, + id: format!("price_{}", price), + } + } + + pub fn market_maker>(name: S) -> Self { + Self { + node_type: NodeType::MarketMaker, + id: name.into(), + } + } +} + +/// Market edge with temporal properties +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct MarketEdge { + pub edge_type: EdgeType, + pub weight: i64, + pub strength: f64, + pub timestamp: u64, + pub decay_factor: f64, +} + +impl MarketEdge { + pub fn new(edge_type: EdgeType, weight: i64, strength: f64) -> Self { + Self { + edge_type, + weight, + strength, + timestamp: SystemTime::now() + .duration_since(SystemTime::UNIX_EPOCH) + .unwrap_or_default() + .as_nanos() as u64, + decay_factor: 0.99, + } + } + + pub fn apply_temporal_decay(&mut self, current_time: u64) { + let age = current_time.saturating_sub(self.timestamp); + let decay = self.decay_factor.powf(age as f64 / 1_000_000_000.0); // per second + self.strength *= decay; + self.weight = (self.weight as f64 * decay) as i64; + } +} + +/// TGGN model configuration +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct TGGNConfig { + /// Maximum number of nodes + pub max_nodes: usize, + + /// Maximum number of edges + pub max_edges: usize, + + /// Node feature dimension + pub node_dim: usize, + + /// Edge feature dimension + pub edge_dim: usize, + + /// Hidden dimension for GNN layers + pub hidden_dim: usize, + + /// Number of message passing layers + pub num_layers: usize, + + /// Temporal decay factor + pub temporal_decay: f64, + + /// Graph update frequency (nanoseconds) + pub update_frequency_ns: u64, + + /// Enable SIMD optimizations + pub use_simd: bool, +} + +impl Default for TGGNConfig { + fn default() -> Self { + Self { + max_nodes: 1000, + max_edges: 10000, + node_dim: 32, + edge_dim: 16, + hidden_dim: 64, + num_layers: 3, + temporal_decay: 0.99, + update_frequency_ns: 1_000_000, // 1ms + use_simd: true, + } + } +} + +/// Temporal Graph Gated Networks for market microstructure +#[derive(Debug)] +pub struct TGGN { + /// Model configuration + config: TGGNConfig, + + /// Model metadata + pub metadata: ModelMetadata, + + /// Market graph structure + graph: MarketGraph, + + /// Gating mechanism + gating: GatingMechanism, + + /// Message passing layers + message_passing: Vec, + + /// Node embeddings cache + node_embeddings: DashMap>, + + /// Edge embeddings cache + edge_embeddings: DashMap<(NodeId, NodeId), Array1>, + + /// Whether model is trained + is_trained: bool, + + /// Performance counters + inference_count: AtomicU64, + total_latency_ns: AtomicU64, + max_latency_ns: AtomicU64, + graph_updates: AtomicU64, + + /// Last update timestamp + last_update: AtomicU64, +} + +impl TGGN { + /// Create new TGGN model + pub fn new(config: TGGNConfig) -> Result { + let mut metadata = ModelMetadata::new( + ModelType::TGNN, + "1.0.0".to_owned(), + config.node_dim, + 1.0, // Single output for prediction + ); + metadata.add_metadata("max_nodes", config.max_nodes.to_string()); + metadata.add_metadata("max_edges", config.max_edges.to_string()); + metadata.add_metadata("hidden_dim", config.hidden_dim.to_string()); + metadata.add_metadata("num_layers", config.num_layers.to_string()); + + let graph = MarketGraph::new(config.max_nodes, config.max_edges)?; + let gating = GatingMechanism::new(config.hidden_dim)?; + + // Initialize message passing layers + let mut message_passing = Vec::with_capacity(config.num_layers); + for layer in 0..config.num_layers { + let input_dim = if layer == 0 { + config.node_dim + } else { + config.hidden_dim + }; + message_passing.push(MessagePassing::new(input_dim, config.hidden_dim)?); + } + + info!( + "Initialized TGGN with {} nodes, {} layers", + config.max_nodes, config.num_layers + ); + + Ok(Self { + config, + metadata, + graph, + gating, + message_passing, + node_embeddings: DashMap::new(), + edge_embeddings: DashMap::new(), + is_trained: false, + inference_count: AtomicU64::new(0), + total_latency_ns: AtomicU64::new(0), + max_latency_ns: AtomicU64::new(0), + graph_updates: AtomicU64::new(0), + last_update: AtomicU64::new(0), + }) + } + + /// Create with default configuration + pub fn default() -> Result { + Self::new(TGGNConfig::default()) + } + + /// Update graph from order book data + pub fn update_from_order_book( + &mut self, + bids: &[(i64, i64)], // (price, volume) pairs + asks: &[(i64, i64)], + timestamp: u64, + ) -> Result<(), MLError> { + let start = Instant::now(); + + // Clear old nodes and edges + self.graph + .clear_temporal_data(timestamp, self.config.temporal_decay)?; + + // Add price level nodes for bids + for (i, &(price, volume)) in bids.iter().enumerate() { + let node_id = NodeId::price_level(price); + let features = self.create_price_level_features(price, volume, true, i)?; + self.graph.add_node(node_id.clone(), features.to_vec())?; + self.node_embeddings.insert(node_id, features); + } + + // Add price level nodes for asks + for (i, &(price, volume)) in asks.iter().enumerate() { + let node_id = NodeId::price_level(price); + let features = self.create_price_level_features(price, volume, false, i)?; + self.graph.add_node(node_id.clone(), features.to_vec())?; + self.node_embeddings.insert(node_id, features); + } + + // Create edges between nearby price levels + self.create_proximity_edges(bids, asks)?; + + // Create liquidity flow edges + self.create_liquidity_edges(bids, asks)?; + + let elapsed = start.elapsed(); + self.graph_updates.fetch_add(1, Ordering::Relaxed); + self.last_update.store(timestamp, Ordering::Relaxed); + + debug!( + "Updated graph in {}ns: {} nodes, {} edges", + elapsed.as_nanos(), + self.graph.node_count(), + self.graph.edge_count() + ); + + // Check latency target + let latency_ns = elapsed.as_nanos() as u64; + if latency_ns > 500 { + // 500ns target + warn!("Graph update {}ns exceeds target 500ns", latency_ns); + } + + Ok(()) + } + + /// Perform graph neural network inference + pub fn gnn_inference( + &mut self, + target_nodes: &[NodeId], + ) -> Result, MLError> { + let start = Instant::now(); + + let mut predictions = HashMap::new(); + + // Get current node embeddings + let mut node_features = HashMap::new(); + for node_id in target_nodes { + if let Some(embedding) = self.node_embeddings.get(node_id) { + node_features.insert(node_id.clone(), embedding.clone()); + } else { + // Create default features if node not found + let default_features = Array1::zeros(self.config.node_dim); + node_features.insert(node_id.clone(), default_features); + } + } + + // Apply message passing layers + for (layer_idx, layer) in self.message_passing.iter().enumerate() { + let layer_start = Instant::now(); + + // Collect messages from neighbors + for node_id in target_nodes { + if let Some(neighbors) = self.graph.get_neighbors(node_id) { + let messages = self.collect_messages(node_id, &neighbors, &node_features)?; + + // Apply gating mechanism + let gated_messages = self.gating.apply(&messages)?; + + // Update node features with gated messages + if let Some(current_features) = node_features.get_mut(node_id) { + let updated = layer.forward(current_features, &gated_messages)?; + *current_features = updated; + } + } + } + + debug!( + "Layer {} completed in {}ns", + layer_idx, + layer_start.elapsed().as_nanos() + ); + } + + // Generate predictions from final node features + for node_id in target_nodes { + if let Some(features) = node_features.get(node_id) { + // Simple prediction: weighted sum of features + let prediction = features.sum() / features.len() as f64; + predictions.insert(node_id.clone(), prediction); + } + } + + let elapsed = start.elapsed(); + self.inference_count.fetch_add(1, Ordering::Relaxed); + self.total_latency_ns + .fetch_add(elapsed.as_nanos() as u64, Ordering::Relaxed); + + let latency_ns = elapsed.as_nanos() as u64; + let current_max = self.max_latency_ns.load(Ordering::Relaxed); + if latency_ns > current_max { + self.max_latency_ns + .compare_exchange_weak( + current_max, + latency_ns, + Ordering::Relaxed, + Ordering::Relaxed, + ) + .ok(); + } + + // Check sub-1μs target + if latency_ns > 1000 { + // 1μs = 1000ns + warn!("GNN inference {}ns exceeds target 1000ns", latency_ns); + } else { + debug!("GNN inference completed in {}ns", latency_ns); + } + + Ok(predictions) + } + + /// Create features for price level nodes + fn create_price_level_features( + &self, + price: i64, + volume: i64, + is_bid: bool, + depth_level: usize, + ) -> Result, MLError> { + let mut features = Array1::zeros(self.config.node_dim); + + // Normalize price and volume + let price_norm = (price as f64) / PRECISION_FACTOR as f64; + let volume_norm = (volume as f64) / PRECISION_FACTOR as f64; + + // Feature 0-3: Basic price/volume info + if features.len() > 0 { + features[0] = price_norm; + } + if features.len() > 1 { + features[1] = volume_norm; + } + if features.len() > 2 { + features[2] = if is_bid { 1.0 } else { -1.0 }; + } + if features.len() > 3 { + features[3] = depth_level as f64 / 10.0; + } + + // Feature 4-7: Statistical features + if features.len() > 4 { + features[4] = price_norm.ln(); + } // Log price + if features.len() > 5 { + features[5] = volume_norm.sqrt(); + } // Sqrt volume + if features.len() > 6 { + features[6] = price_norm * volume_norm; + } // Price * volume + if features.len() > 7 { + features[7] = volume_norm / (price_norm + 1e-8); + } // Volume/price ratio + + // Feature 8-15: Technical indicators (simplified) + for i in 8..features.len().min(16) { + let phase = (i as f64 * std::f64::consts::PI) / 8.0; + features[i] = (price_norm * phase.cos() + volume_norm * phase.sin()) / 10.0; + } + + // Feature 16+: Reserved for market microstructure + for i in 16..features.len() { + features[i] = thread_rng().gen::() * 0.01; // Small random noise + } + + Ok(features) + } + + /// Create edges between nearby price levels + fn create_proximity_edges( + &mut self, + bids: &[(i64, i64)], + asks: &[(i64, i64)], + ) -> Result<(), MLError> { + // Connect adjacent price levels within same side + for window in bids.windows(2) { + let node1 = NodeId::price_level(window[0].0); + let node2 = NodeId::price_level(window[1].0); + let weight = ((window[0].1 + window[1].1) / 2) as i64; // Avg volume + let edge = MarketEdge::new(EdgeType::PriceProximity, weight, 0.8); + self.graph.add_edge(&node1, &node2, edge)?; + } + + for window in asks.windows(2) { + let node1 = NodeId::price_level(window[0].0); + let node2 = NodeId::price_level(window[1].0); + let weight = ((window[0].1 + window[1].1) / 2) as i64; + let edge = MarketEdge::new(EdgeType::PriceProximity, weight, 0.8); + self.graph.add_edge(&node1, &node2, edge)?; + } + + // Connect best bid and ask + if !bids.is_empty() && !asks.is_empty() { + let best_bid = NodeId::price_level(bids[0].0); + let best_ask = NodeId::price_level(asks[0].0); + let spread_weight = (asks[0].0 - bids[0].0).abs(); + let edge = MarketEdge::new(EdgeType::PriceProximity, spread_weight, 0.9); + self.graph.add_edge(&best_bid, &best_ask, edge)?; + } + + Ok(()) + } + + /// Create liquidity flow edges + fn create_liquidity_edges( + &mut self, + bids: &[(i64, i64)], + asks: &[(i64, i64)], + ) -> Result<(), MLError> { + // Create flow edges based on volume imbalance + let total_bid_volume: i64 = bids.iter().map(|(_, v)| v).sum(); + let total_ask_volume: i64 = asks.iter().map(|(_, v)| v).sum(); + + let imbalance = total_bid_volume - total_ask_volume; + let flow_strength = + (imbalance.abs() as f64) / (total_bid_volume + total_ask_volume + 1) as f64; + + // Connect high-volume levels with flow edges + for &(price, volume) in bids.into_iter().take(3) { + for &(ask_price, ask_volume) in asks.into_iter().take(3) { + if volume > total_bid_volume / 10 && ask_volume > total_ask_volume / 10 { + let node1 = NodeId::price_level(price); + let node2 = NodeId::price_level(ask_price); + let weight = (volume.min(ask_volume)) as i64; + let edge = MarketEdge::new(EdgeType::LiquidityFlow, weight, flow_strength); + self.graph.add_edge(&node1, &node2, edge)?; + } + } + } + + Ok(()) + } + + /// Collect messages from neighboring nodes + fn collect_messages( + &self, + node_id: &NodeId, + neighbors: &[NodeId], + node_features: &HashMap>, + ) -> Result>, MLError> { + let mut messages = Vec::new(); + + for neighbor in neighbors { + if let Some(neighbor_features) = node_features.get(neighbor) { + // Get edge weight if available + let edge_weight = self.graph.get_edge_weight(node_id, neighbor).unwrap_or(1.0); + + // Weight neighbor features by edge strength + let weighted_message = neighbor_features.mapv(|x| x * edge_weight); + messages.push(weighted_message); + } + } + + Ok(messages) + } + + /// Get performance statistics + pub fn get_performance_stats(&self) -> HashMap { + let mut stats = HashMap::new(); + + let inference_count = self.inference_count.load(Ordering::Relaxed); + let total_latency = self.total_latency_ns.load(Ordering::Relaxed); + let max_latency = self.max_latency_ns.load(Ordering::Relaxed); + let graph_updates = self.graph_updates.load(Ordering::Relaxed); + + stats.insert("inference_count".to_owned(), inference_count as f64); + stats.insert("graph_updates".to_owned(), graph_updates as f64); + stats.insert("max_latency_ns".to_owned(), max_latency as f64); + + if inference_count > 0 { + stats.insert( + "avg_latency_ns".to_owned(), + total_latency as f64 / inference_count as f64, + ); + } + + stats.insert("node_count".to_owned(), self.graph.node_count() as f64); + stats.insert("edge_count".to_owned(), self.graph.edge_count() as f64); + + stats + } + + /// Public getters for checkpoint operations + pub fn node_embeddings(&self) -> &DashMap> { + &self.node_embeddings + } + + pub fn edge_embeddings(&self) -> &DashMap<(NodeId, NodeId), Array1> { + &self.edge_embeddings + } + + pub fn config(&self) -> &TGGNConfig { + &self.config + } + + pub fn inference_count(&self) -> &AtomicU64 { + &self.inference_count + } + + pub fn graph_updates(&self) -> &AtomicU64 { + &self.graph_updates + } + + pub fn is_trained(&self) -> bool { + self.is_trained + } + + /// Get graph statistics for monitoring and checkpointing + pub fn get_graph_stats(&self) -> (usize, usize) { + (self.graph.node_count(), self.graph.edge_count()) + } + + /// Restore node embeddings from checkpoint state + pub fn restore_node_embeddings( + &mut self, + embeddings: &HashMap>, + ) -> Result<(), MLError> { + for (node_id_str, embedding) in embeddings { + // Convert f32 to f64 + let embedding_f64: Vec = embedding.iter().map(|&x| x as f64).collect(); + let array = Array1::from_vec(embedding_f64); + // Parse the node ID string to create proper NodeId + let node_id = if node_id_str.starts_with("price_") { + let price = node_id_str + .strip_prefix("price_") + .and_then(|s| s.parse::().ok()) + .unwrap_or(0); + NodeId::price_level(price) + } else if node_id_str.starts_with("mm_") { + NodeId::market_maker(node_id_str.get(3..).unwrap_or_default()) + } else { + // Default case - use the string as-is with generic type + NodeId { + node_type: NodeType::PriceLevel, + id: node_id_str.clone(), + } + }; + self.node_embeddings.insert(node_id, array); + } + Ok(()) + } + + /// Restore edge embeddings from checkpoint state + pub fn restore_edge_embeddings( + &mut self, + embeddings: &HashMap>, + ) -> Result<(), MLError> { + for (edge_key, embedding) in embeddings { + // Convert f32 to f64 + let embedding_f64: Vec = embedding.iter().map(|&x| x as f64).collect(); + let array = Array1::from_vec(embedding_f64); + + // Parse edge key (assuming format like "from_id->to_id") + if let Some((from_str, to_str)) = edge_key.split_once("->") { + // Parse from node + let from_node = if from_str.starts_with("price_") { + let price = from_str + .strip_prefix("price_") + .and_then(|s| s.parse::().ok()) + .unwrap_or(0); + NodeId::price_level(price) + } else if from_str.starts_with("mm_") { + NodeId::market_maker(from_str.get(3..).unwrap_or_default()) + } else { + NodeId { + node_type: NodeType::PriceLevel, + id: from_str.to_string(), + } + }; + + // Parse to node + let to_node = if to_str.starts_with("price_") { + let price = to_str + .strip_prefix("price_") + .and_then(|s| s.parse::().ok()) + .unwrap_or(0); + NodeId::price_level(price) + } else if to_str.starts_with("mm_") { + NodeId::market_maker(to_str.get(3..).unwrap_or_default()) + } else { + NodeId { + node_type: NodeType::PriceLevel, + id: to_str.to_string(), + } + }; + + self.edge_embeddings.insert((from_node, to_node), array); + } + } + Ok(()) + } + + /// Restore graph statistics from checkpoint state + pub fn restore_graph_statistics( + &mut self, + _stats: &HashMap, + ) -> Result<(), MLError> { + // Production implementation for now - graph statistics would be restored here + Ok(()) + } + + /// Restore message passing weights from checkpoint state + pub fn restore_message_passing_weights( + &mut self, + _weights: &Vec>, + ) -> Result<(), MLError> { + // Production implementation for now - message passing weights would be restored here + Ok(()) + } +} + +#[async_trait] +impl MLModel for TGGN { + type Config = serde_json::Value; + + fn metadata(&self) -> &ModelMetadata { + &self.metadata + } + + fn is_ready(&self) -> bool { + self.is_trained + } + + async fn train( + &mut self, + features: &Array2, + targets: &Array2, + ) -> Result { + info!("Starting TGGN training with {} samples", features.nrows()); + + let start = Instant::now(); + let _n_samples = features.nrows(); + + // For TGGN, training involves learning message passing weights with real gradients + let learning_rate = 0.001; + + // Prepare batch data for message passing training (moved outside loop) + let mut node_features_batch = Vec::new(); + let mut neighbor_messages_batch = Vec::new(); + let mut targets_batch = Vec::new(); + + // Convert features to node features and targets for each layer + // Start with input features (32-dim), progressively transform through layers + let mut current_features = features.clone(); + let num_layers = self.message_passing.len(); + + for (layer_idx, layer) in self.message_passing.iter_mut().enumerate() { + info!("Training layer {} with real backpropagation", layer_idx); + + // Clear batch data for this layer + node_features_batch.clear(); + neighbor_messages_batch.clear(); + targets_batch.clear(); + + for sample_idx in 0..current_features.nrows().min(targets.nrows()) { + let node_features = current_features.row(sample_idx).to_owned(); + // Create target with correct dimension (hidden_dim, not 1) + // Replicate the single target value across all hidden dimensions + let target_value = targets[[sample_idx, 0]]; + let target = Array1::from_elem(self.config.hidden_dim, target_value); + + // For training, create synthetic neighbor messages from nearby samples + let mut neighbor_messages = Vec::new(); + for neighbor_idx in 0..3.min(current_features.nrows()) { + // Use up to 3 neighbors + if neighbor_idx != sample_idx { + let neighbor_features = current_features.row(neighbor_idx).to_owned(); + neighbor_messages.push(neighbor_features); + } + } + + node_features_batch.push(node_features); + neighbor_messages_batch.push(neighbor_messages); + targets_batch.push(target); + } + + // Train layer with proper backpropagation + layer + .train_weights( + &node_features_batch, + &neighbor_messages_batch, + &targets_batch, + learning_rate, + ) + .map_err(|e| MLError::TrainingError(format!("Layer training failed: {}", e)))?; + + // Transform features through this layer for next layer's training + // This ensures layer 1 gets 64-dim inputs, not 32-dim + if layer_idx < num_layers - 1 { + let mut transformed_features = Vec::new(); + for sample_idx in 0..current_features.nrows() { + let node_features = current_features.row(sample_idx).to_owned(); + + // Get neighbor messages for transformation + let mut neighbor_messages = Vec::new(); + for neighbor_idx in 0..3.min(current_features.nrows()) { + if neighbor_idx != sample_idx { + let neighbor_features = current_features.row(neighbor_idx).to_owned(); + neighbor_messages.push(neighbor_features); + } + } + + // Transform through layer + let transformed = layer + .forward(&node_features, &neighbor_messages) + .unwrap_or_else(|_| Array1::zeros(self.config.hidden_dim)); + transformed_features.push(transformed); + } + + // Convert to Array2 for next layer + let n_samples = transformed_features.len(); + let feature_dim = self.config.hidden_dim; + let flat_len = n_samples * feature_dim; + let flat: Vec = transformed_features + .into_iter() + .flat_map(|arr| arr.to_vec()) + .collect(); + current_features = + Array2::from_shape_vec((n_samples, feature_dim), flat).map_err(|_| { + MLError::DimensionMismatch { + expected: flat_len, + actual: flat_len, + } + })?; + } + } + + // Update gating mechanism with real gradients + if !node_features_batch.is_empty() { + // Transform node features to hidden_dim for gating mechanism + // The gating mechanism expects hidden_dim inputs and outputs + let mut transformed_inputs = Vec::new(); + for (node_features, neighbor_messages) in node_features_batch + .into_iter() + .zip(neighbor_messages_batch.into_iter()) + { + // Use first layer to transform node features to hidden_dim + if let Some(first_layer) = self.message_passing.first() { + let transformed = first_layer + .forward(&node_features, &neighbor_messages) + .unwrap_or_else(|_| Array1::zeros(self.config.hidden_dim)); + transformed_inputs.push(transformed); + } + } + + // Only train gating if we have transformed inputs + if !transformed_inputs.is_empty() { + // Gating mechanism uses GLU which halves the dimension + // So we need to adjust targets to match the output dimension (hidden_dim/2) + let glu_output_dim = self.config.hidden_dim / 2; + let mut gating_targets = Vec::new(); + for target in &targets_batch { + // Truncate or pad targets to match GLU output dimension + let adjusted_target = if target.len() >= glu_output_dim { + target.slice(s![..glu_output_dim]).to_owned() + } else { + let mut padded = Array1::zeros(glu_output_dim); + padded.slice_mut(s![..target.len()]).assign(target); + padded + }; + gating_targets.push(adjusted_target); + } + + self.gating + .update_weights(&transformed_inputs, &gating_targets, learning_rate) + .map_err(|e| { + MLError::TrainingError(format!("Gating training failed: {}", e)) + })?; + } + } + + self.is_trained = true; + self.metadata.mark_trained(); + + let training_time = start.elapsed().as_secs_f64(); + + info!("TGGN training completed in {:.2}s", training_time); + + Ok(TrainingMetrics { + loss: 0.1, + accuracy: 0.9, + precision: 0.88, + recall: 0.85, + f1_score: 0.865, + training_time_seconds: training_time, + epochs_trained: 1, + convergence_achieved: true, + additional_metrics: HashMap::new(), + }) + } + + async fn predict(&self, features: &[f64]) -> Result { + if !self.is_trained { + return Err(MLError::NotTrained("TGGN not trained".to_owned())); + } + + let start = Instant::now(); + + // Simple prediction based on features + let prediction = features.iter().sum::() / features.len() as f64; + let confidence = 0.9; // High confidence for graph-based predictions + + let result = InferenceResult::new( + "tgnn_1.0".to_owned(), + prediction, + confidence, + start.elapsed().as_micros() as u64, + start.elapsed().as_nanos() as u64, + self.metadata.clone(), + ); + + Ok(result) + } + + async fn validate( + &self, + features: &Array2, + targets: &Array2, + ) -> Result { + let mut total_error = 0.0; + let mut correct_predictions = 0; + + for i in 0..features.nrows() { + let row_features: Vec = features.row(i).to_vec(); + let prediction_result = self.predict(&row_features).await?; + let prediction = prediction_result.prediction_as_float(); + + let target = targets[[i, 0]]; + let error = (prediction - target).abs(); + total_error += error; + + if error < 0.1 { + // Threshold for "correct" + correct_predictions += 1; + } + } + + let mse = total_error / features.nrows() as f64; + let accuracy = correct_predictions as f64 / features.nrows() as f64; + + Ok(ValidationMetrics { + validation_loss: mse, + validation_accuracy: accuracy, + validation_precision: accuracy * 0.95, + validation_recall: accuracy * 0.93, + validation_f1_score: accuracy * 0.94, + samples_validated: features.nrows(), + additional_metrics: HashMap::new(), + }) + } + + async fn update( + &mut self, + features: &Array2, + targets: &Array2, + ) -> Result<(), MLError> { + // Online learning for TGGN + self.train(features, targets).await?; + Ok(()) + } + + async fn save(&self, path: &str) -> Result<(), MLError> { + let data = serde_json::json!({ + "config": self.config, + "metadata": self.metadata, + "is_trained": self.is_trained, + "performance_stats": self.get_performance_stats(), + }); + + let serialized = + serde_json::to_string_pretty(&data).map_err(|e| MLError::SerializationError { + reason: e.to_string(), + })?; + tokio::fs::write(path, serialized) + .await + .map_err(|e| MLError::SerializationError { + reason: e.to_string(), + })?; + + info!("Saved TGGN model to {}", path); + Ok(()) + } + + async fn load(&mut self, path: &str) -> Result<(), MLError> { + let content = + tokio::fs::read_to_string(path) + .await + .map_err(|e| MLError::SerializationError { + reason: e.to_string(), + })?; + + let data: serde_json::Value = + serde_json::from_str(&content).map_err(|e| MLError::SerializationError { + reason: e.to_string(), + })?; + + self.config = serde_json::from_value(data["config"].clone()).map_err(|e| { + MLError::SerializationError { + reason: e.to_string(), + } + })?; + + self.metadata = serde_json::from_value(data["metadata"].clone()).map_err(|e| { + MLError::SerializationError { + reason: e.to_string(), + } + })?; + + self.is_trained = data["is_trained"].as_bool().unwrap_or(false); + + info!("Loaded TGGN model from {}", path); + Ok(()) + } + + fn config(&self) -> Self::Config { + serde_json::to_value(&self.config).unwrap_or_default() + } + + fn set_config(&mut self, config: Self::Config) -> Result<(), MLError> { + self.config = serde_json::from_value(config).map_err(|e| MLError::ConfigError(e.to_string()))?; + Ok(()) + } +} + +/// Training pipeline for TGGN with order book data +#[derive(Debug)] +pub struct TGGNTrainingPipeline { + pub model: TGGN, + pub training_data: Vec<(Vec<(i64, i64)>, Vec<(i64, i64)>, f64)>, // (bids, asks, target) +} + +impl TGGNTrainingPipeline { + pub fn new(config: TGGNConfig) -> Result { + let model = TGGN::new(config)?; + Ok(Self { + model, + training_data: Vec::new(), + }) + } + + pub fn add_training_sample( + &mut self, + bids: Vec<(i64, i64)>, + asks: Vec<(i64, i64)>, + target: f64, + ) { + self.training_data.push((bids, asks, target)); + } + + pub async fn train_from_order_book_data(&mut self) -> Result { + info!( + "Training TGGN from {} order book samples", + self.training_data.len() + ); + + // Convert order book data to feature matrices + let mut features_vec = Vec::new(); + let mut targets_vec = Vec::new(); + + for (bids, asks, target) in &self.training_data { + // Update graph with order book data + let timestamp = SystemTime::now() + .duration_since(SystemTime::UNIX_EPOCH) + .unwrap_or_default() + .as_nanos() as u64; + + self.model.update_from_order_book(bids, asks, timestamp)?; + + // Extract features from updated graph + let graph_features = self.extract_graph_features()?; + features_vec.push(graph_features); + targets_vec.push(vec![*target]); + } + + // Convert to ndarray format + let features = Array2::from_shape_vec( + (features_vec.len(), features_vec[0].len()), + features_vec.into_iter().flatten().collect(), + ) + .map_err(|e| MLError::DimensionMismatch { + expected: self.model.config.node_dim, + actual: e.to_string().len(), + })?; + + let targets = Array2::from_shape_vec( + (targets_vec.len(), 1), + targets_vec.into_iter().flatten().collect(), + ) + .map_err(|e| MLError::DimensionMismatch { + expected: 1, + actual: e.to_string().len(), + })?; + + // Train the model (convert types::MLError to MLError) + self.model + .train(&features, &targets) + .await + .map_err(|e| MLError::TrainingError(format!("TGNN training failed: {}", e))) + } + + fn extract_graph_features(&self) -> Result, MLError> { + let stats = self.model.graph.get_stats(); + + // Graph topology features + let mut features = vec![ + stats.node_count as f64, + stats.edge_count as f64, + stats.density, + stats.average_degree, + ]; + + // Fill remaining features with zeros if needed + while features.len() < self.model.config.node_dim { + features.push(0.0); + } + + Ok(features) + } +} + +// Checkpoint helper methods used by the checkpoint bridge in ml crate +impl TGGN { + /// Extract graph statistics for checkpoint state + pub fn extract_graph_statistics(&self) -> HashMap { + let mut stats = HashMap::new(); + + let (node_count, edge_count) = self.get_graph_stats(); + let node_count_f64 = node_count as f64; + let edge_count_f64 = edge_count as f64; + + stats.insert("node_count".to_owned(), node_count_f64); + stats.insert("max_nodes".to_owned(), self.config().max_nodes as f64); + stats.insert( + "node_utilization".to_owned(), + node_count_f64 / self.config().max_nodes as f64, + ); + + stats.insert("edge_count".to_owned(), edge_count_f64); + stats.insert("max_edges".to_owned(), self.config().max_edges as f64); + stats.insert( + "edge_utilization".to_owned(), + edge_count_f64 / self.config().max_edges as f64, + ); + + if node_count_f64 > 1.0 { + let max_edges = node_count_f64 * (node_count_f64 - 1.0) / 2.0; + stats.insert("graph_density".to_owned(), edge_count_f64 / max_edges); + } else { + stats.insert("graph_density".to_owned(), 0.0); + } + + if node_count_f64 > 0.0 { + stats.insert( + "avg_degree".to_owned(), + (2.0 * edge_count_f64) / node_count_f64, + ); + } else { + stats.insert("avg_degree".to_owned(), 0.0); + } + + stats + } + + /// Extract message passing weights from the model layers + pub fn extract_message_passing_weights(&self) -> Vec { + let mut weights = Vec::new(); + + for layer_idx in 0..self.config().num_layers { + let layer_size = self.config().node_dim * self.config().edge_dim; + for i in 0..layer_size { + weights.push((layer_idx as f32 + (i as f32).rem_euclid(10.0)) * 0.1); + } + } + + weights + } + + /// Extract gating mechanism weights + pub fn extract_gating_weights(&self) -> Vec { + let mut gating_weights = Vec::new(); + + let num_gates = self.config().num_layers * 3; + for gate_idx in 0..num_gates { + let gate_size = self.config().node_dim; + for i in 0..gate_size { + gating_weights.push(((gate_idx * gate_size + i) as f32).rem_euclid(100.0) * 0.01); + } + } + + gating_weights + } + + /// Calculate average inference latency from internal statistics + pub fn calculate_avg_inference_latency(&self) -> u64 { + let total_inferences = self + .inference_count() + .load(std::sync::atomic::Ordering::Relaxed); + + if total_inferences > 0 { + let base_latency_ns = match self.config().num_layers { + 1..=3 => 1_000_000, + 4..=6 => 5_000_000, + _ => 10_000_000, + }; + + let graph_complexity = (100 + 200) as u64; + let complexity_factor = (graph_complexity / 1000).max(1); + + base_latency_ns * complexity_factor + } else { + 0 + } + } + + /// Calculate average graph update latency + pub fn calculate_avg_graph_update_latency(&self) -> u64 { + let total_updates = self + .graph_updates() + .load(std::sync::atomic::Ordering::Relaxed); + + if total_updates > 0 { + let base_update_latency_ns = 500_000; + let graph_size_factor = ((100 + 200) / 100) as u64; + + base_update_latency_ns * graph_size_factor + } else { + 0 + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[tokio::test] + async fn test_tggn_creation() -> Result<(), MLError> { + let config = TGGNConfig::default(); + let model = TGGN::new(config)?; + + assert_eq!(model.config.max_nodes, 1000); + assert_eq!(model.config.num_layers, 3); + assert!(!model.is_trained); + Ok(()) + } + + #[tokio::test] + async fn test_order_book_update() -> Result<(), MLError> { + let config = TGGNConfig::default(); + let mut model = TGGN::new(config)?; + + let bids = vec![(100_00000000, 1000_00000000), (99_00000000, 500_00000000)]; + let asks = vec![(101_00000000, 800_00000000), (102_00000000, 600_00000000)]; + let timestamp = 1234567890; + + let result = model.update_from_order_book(&bids, &asks, timestamp); + assert!(result.is_ok()); + + assert_eq!(model.graph.node_count(), 4); // 2 bids + 2 asks + assert!(model.graph.edge_count() > 0); + Ok(()) + } + + #[tokio::test] + async fn test_gnn_inference() -> Result<(), MLError> { + let config = TGGNConfig::default(); + let mut model = TGGN::new(config)?; + + // Setup graph with some nodes + let bids = vec![(100_00000000, 1000_00000000)]; + let asks = vec![(101_00000000, 800_00000000)]; + let timestamp = 1234567890; + + model.update_from_order_book(&bids, &asks, timestamp)?; + + let target_nodes = vec![NodeId::price_level(100_00000000)]; + let predictions = model.gnn_inference(&target_nodes)?; + + assert_eq!(predictions.len(), 1); + assert!(predictions.contains_key(&NodeId::price_level(100_00000000))); + Ok(()) + } + + #[tokio::test] + async fn test_training_pipeline() -> Result<(), MLError> { + let config = TGGNConfig::default(); + let mut pipeline = TGGNTrainingPipeline::new(config)?; + + // Add some training samples + pipeline.add_training_sample( + vec![(100_00000000, 1000_00000000)], + vec![(101_00000000, 800_00000000)], + 0.5, + ); + + pipeline.add_training_sample( + vec![(99_00000000, 1200_00000000)], + vec![(100_00000000, 900_00000000)], + -0.3, + ); + + let metrics = pipeline.train_from_order_book_data().await?; + assert!(metrics.training_time_seconds > 0.0); + assert!(pipeline.model.is_trained); + Ok(()) + } +} diff --git a/crates/ml/src/tgnn/traits.rs b/crates/ml-supervised/src/tgnn/traits.rs similarity index 97% rename from crates/ml/src/tgnn/traits.rs rename to crates/ml-supervised/src/tgnn/traits.rs index 0de5d3a99..048670a75 100644 --- a/crates/ml/src/tgnn/traits.rs +++ b/crates/ml-supervised/src/tgnn/traits.rs @@ -1,7 +1,7 @@ //! Traits for TGNN implementation use super::types::*; -use crate::{InferenceResult, MLError, ModelMetadata}; +use ml_core::{InferenceResult, MLError, ModelMetadata}; use async_trait::async_trait; use ndarray::Array2; diff --git a/crates/ml/src/tgnn/types.rs b/crates/ml-supervised/src/tgnn/types.rs similarity index 100% rename from crates/ml/src/tgnn/types.rs rename to crates/ml-supervised/src/tgnn/types.rs diff --git a/crates/ml/src/tlob/analytics.rs b/crates/ml-supervised/src/tlob/analytics.rs similarity index 100% rename from crates/ml/src/tlob/analytics.rs rename to crates/ml-supervised/src/tlob/analytics.rs diff --git a/crates/ml/src/tlob/features.rs b/crates/ml-supervised/src/tlob/features.rs similarity index 99% rename from crates/ml/src/tlob/features.rs rename to crates/ml-supervised/src/tlob/features.rs index 9a51e7643..f49acff18 100644 --- a/crates/ml/src/tlob/features.rs +++ b/crates/ml-supervised/src/tlob/features.rs @@ -22,7 +22,7 @@ use anyhow::Result; use serde::{Deserialize, Serialize}; use tracing::{instrument, warn}; -use crate::MLError; +use ml_core::MLError; /// Total number of TLOB features extracted pub const TLOB_FEATURE_COUNT: usize = 51; diff --git a/crates/ml/src/tlob/mbp10_feature_extractor.rs b/crates/ml-supervised/src/tlob/mbp10_feature_extractor.rs similarity index 99% rename from crates/ml/src/tlob/mbp10_feature_extractor.rs rename to crates/ml-supervised/src/tlob/mbp10_feature_extractor.rs index 001d31d56..6e459e516 100644 --- a/crates/ml/src/tlob/mbp10_feature_extractor.rs +++ b/crates/ml-supervised/src/tlob/mbp10_feature_extractor.rs @@ -4,7 +4,7 @@ //! for transformer-based limit order book prediction. use crate::tlob::features::{FeatureVector, TLOBFeatureExtractor, TLOBFeatures}; -use crate::MLError; +use ml_core::MLError; use anyhow::Result; use data::providers::databento::mbp10::Mbp10Snapshot; use tracing::{debug, instrument}; diff --git a/crates/ml-supervised/src/tlob/mod.rs b/crates/ml-supervised/src/tlob/mod.rs new file mode 100644 index 000000000..96d342064 --- /dev/null +++ b/crates/ml-supervised/src/tlob/mod.rs @@ -0,0 +1,23 @@ +//! Time Limit Order Book (TLOB) Transformer +//! +//! High-performance TLOB analysis for HFT systems with sub-50μs latency requirements. +//! Based on advanced order flow analytics from institutional trading systems. + +pub mod analytics; +pub mod features; +pub mod mbp10_feature_extractor; // MBP-10 to TLOB feature extraction +pub mod performance; +pub mod transformer; + +// Re-export key types for external use +pub use features::{ + ExtractionMetrics, + FeatureVector as TLOBFeatureVector, // Rename to avoid conflict with main FeatureVector from lib.rs + TLOBFeatureExtractor, + TLOBFeatures as TLOBInputFeatures, + TLOB_FEATURE_COUNT, +}; +pub use transformer::{TLOBConfig, TLOBMetrics, TLOBTransformer}; + +// Re-export transformer-specific TLOBFeatures with a different name to avoid conflicts +pub use transformer::TLOBFeatures as TLOBPredictionFeatures; diff --git a/crates/ml/src/tlob/performance.rs b/crates/ml-supervised/src/tlob/performance.rs similarity index 100% rename from crates/ml/src/tlob/performance.rs rename to crates/ml-supervised/src/tlob/performance.rs diff --git a/crates/ml/src/tlob/transformer.rs b/crates/ml-supervised/src/tlob/transformer.rs similarity index 99% rename from crates/ml/src/tlob/transformer.rs rename to crates/ml-supervised/src/tlob/transformer.rs index 37a9503a1..de22d1041 100644 --- a/crates/ml/src/tlob/transformer.rs +++ b/crates/ml-supervised/src/tlob/transformer.rs @@ -6,8 +6,8 @@ use std::sync::Arc; use std::time::Instant; -use crate::FeatureVector; -use crate::MLError; +use ml_core::FeatureVector; +use ml_core::MLError; use anyhow::Result; use candle_core::Device; // ONNX Runtime removed - keeping interface for compatibility diff --git a/crates/ml/src/xlstm/block.rs b/crates/ml-supervised/src/xlstm/block.rs similarity index 98% rename from crates/ml/src/xlstm/block.rs rename to crates/ml-supervised/src/xlstm/block.rs index ff648bcec..98b2b7055 100644 --- a/crates/ml/src/xlstm/block.rs +++ b/crates/ml-supervised/src/xlstm/block.rs @@ -7,7 +7,7 @@ use candle_core::Tensor; use candle_nn::{layer_norm, LayerNorm, LayerNormConfig, Module, VarBuilder}; -use crate::MLError; +use ml_core::MLError; use super::slstm::SLSTMCell; use super::mlstm::MLSTMCell; @@ -120,7 +120,7 @@ mod tests { use super::*; use candle_core::Device; use candle_nn::VarMap; - use crate::dqn::mixed_precision::training_dtype; + use ml_core::mixed_precision::training_dtype; #[test] fn test_slstm_block_forward() { diff --git a/crates/ml/src/xlstm/config.rs b/crates/ml-supervised/src/xlstm/config.rs similarity index 100% rename from crates/ml/src/xlstm/config.rs rename to crates/ml-supervised/src/xlstm/config.rs diff --git a/crates/ml/src/xlstm/mlstm.rs b/crates/ml-supervised/src/xlstm/mlstm.rs similarity index 98% rename from crates/ml/src/xlstm/mlstm.rs rename to crates/ml-supervised/src/xlstm/mlstm.rs index 1ddd3f477..18fe47527 100644 --- a/crates/ml/src/xlstm/mlstm.rs +++ b/crates/ml-supervised/src/xlstm/mlstm.rs @@ -14,7 +14,7 @@ use candle_core::Tensor; use candle_nn::{linear, Linear, Module, VarBuilder}; -use crate::MLError; +use ml_core::MLError; /// mLSTM cell with matrix memory. #[derive(Debug)] @@ -91,7 +91,7 @@ impl MLSTMCell { let (h_prev, c_flat_prev) = match state { Some((h, c)) => (h.clone(), c.clone()), None => { - let dtype = crate::dqn::mixed_precision::training_dtype(dev); + let dtype = ml_core::mixed_precision::training_dtype(dev); let h_zeros = Tensor::zeros(&[batch, self.hidden_dim], dtype, dev) .map_err(|e| MLError::ModelError(format!("mLSTM h_zeros: {e}")))?; let c_zeros = Tensor::zeros(&[batch, c_flat_size], dtype, dev) @@ -221,7 +221,7 @@ mod tests { use super::*; use candle_core::Device; use candle_nn::VarMap; - use crate::dqn::mixed_precision::training_dtype; + use ml_core::mixed_precision::training_dtype; #[test] fn test_mlstm_output_shape() { diff --git a/crates/ml-supervised/src/xlstm/mod.rs b/crates/ml-supervised/src/xlstm/mod.rs new file mode 100644 index 000000000..b4c3fccb5 --- /dev/null +++ b/crates/ml-supervised/src/xlstm/mod.rs @@ -0,0 +1,17 @@ +//! xLSTM (Extended Long Short-Term Memory). +//! +//! Combines two cell types: +//! - **sLSTM**: Exponential gating with scalar memory (better for sequential patterns) +//! - **mLSTM**: Matrix memory with key-value association (higher memory capacity) +//! +//! OOM notes: mLSTM matrix memory is (num_heads, head_dim, head_dim) per sample. +//! Keep hidden_dim/num_heads ratio reasonable (head_dim ≤ 32 recommended). + +pub mod block; +pub mod config; +pub mod mlstm; +pub mod network; +pub mod slstm; + +pub use config::XLSTMConfig; +pub use network::XLSTMNetwork; diff --git a/crates/ml/src/xlstm/network.rs b/crates/ml-supervised/src/xlstm/network.rs similarity index 97% rename from crates/ml/src/xlstm/network.rs rename to crates/ml-supervised/src/xlstm/network.rs index d29c7664f..f7fc5d1b9 100644 --- a/crates/ml/src/xlstm/network.rs +++ b/crates/ml-supervised/src/xlstm/network.rs @@ -8,7 +8,7 @@ use candle_core::Tensor; use candle_nn::{linear, Linear, Module, VarBuilder}; -use crate::MLError; +use ml_core::MLError; use super::block::XLSTMBlock; use super::config::XLSTMConfig; @@ -81,7 +81,7 @@ impl XLSTMNetwork { /// /// Returns `(batch, output_dim)`. pub fn forward(&self, input: &Tensor) -> Result { - let input = crate::dqn::mixed_precision::ensure_training_dtype(input) + let input = ml_core::mixed_precision::ensure_training_dtype(input) .map_err(|e| MLError::ModelError(e.to_string()))?; let dims = input.dims(); let (batch, seq_len, _features) = match dims.len() { @@ -101,7 +101,7 @@ impl XLSTMNetwork { let mut final_h = Tensor::zeros( &[batch, self.hidden_dim], - crate::dqn::mixed_precision::training_dtype(input.device()), + ml_core::mixed_precision::training_dtype(input.device()), input.device(), ) .map_err(|e| MLError::ModelError(format!("xLSTM zeros: {e}")))?; @@ -154,7 +154,7 @@ mod tests { use super::*; use candle_core::Device; use candle_nn::VarMap; - use crate::dqn::mixed_precision::training_dtype; + use ml_core::mixed_precision::training_dtype; fn small_config() -> XLSTMConfig { XLSTMConfig { diff --git a/crates/ml/src/xlstm/slstm.rs b/crates/ml-supervised/src/xlstm/slstm.rs similarity index 98% rename from crates/ml/src/xlstm/slstm.rs rename to crates/ml-supervised/src/xlstm/slstm.rs index 903810720..a19735673 100644 --- a/crates/ml/src/xlstm/slstm.rs +++ b/crates/ml-supervised/src/xlstm/slstm.rs @@ -7,7 +7,7 @@ use candle_core::Tensor; use candle_nn::{linear, Linear, Module, VarBuilder}; -use crate::MLError; +use ml_core::MLError; /// sLSTM cell with exponential gating and scalar memory. #[derive(Debug)] @@ -59,7 +59,7 @@ impl SLSTMCell { None => { let zeros = Tensor::zeros( &[batch, self.hidden_dim], - crate::dqn::mixed_precision::training_dtype(dev), + ml_core::mixed_precision::training_dtype(dev), dev, ) .map_err(|e| MLError::ModelError(format!("sLSTM init zeros: {e}")))?; @@ -133,7 +133,7 @@ mod tests { use super::*; use candle_core::Device; use candle_nn::VarMap; - use crate::dqn::mixed_precision::training_dtype; + use ml_core::mixed_precision::training_dtype; #[test] fn test_slstm_output_shape() { diff --git a/crates/ml/Cargo.toml b/crates/ml/Cargo.toml index 6993947af..7d675bf16 100644 --- a/crates/ml/Cargo.toml +++ b/crates/ml/Cargo.toml @@ -71,6 +71,7 @@ colored = "2.1" # Terminal color output for evaluation reports ml-core.workspace = true ml-dqn.workspace = true ml-ppo.workspace = true +ml-supervised.workspace = true config.workspace = true common = { workspace = true, features = ["questdb"] } risk = { path = "../risk" } diff --git a/crates/ml/src/checkpoint/model_implementations.rs b/crates/ml/src/checkpoint/model_implementations.rs index 507ff96bf..88599a1b1 100644 --- a/crates/ml/src/checkpoint/model_implementations.rs +++ b/crates/ml/src/checkpoint/model_implementations.rs @@ -403,591 +403,6 @@ impl Checkpointable for Mamba2SSM { } } -impl Mamba2SSM { - /// Get current training state from model metadata - fn get_current_training_state(&self) -> (Option, Option) { - if let Some(last_epoch) = self.metadata.training_history.back() { - (Some(last_epoch.epoch as u64), None) // MAMBA doesn't track steps within epochs - } else { - (None, None) - } - } - - /// Get training metrics from model performance data - fn get_training_metrics(&self) -> HashMap { - let mut metrics = HashMap::new(); - - if let Some(last_epoch) = self.metadata.training_history.back() { - metrics.insert("training_loss".to_owned(), last_epoch.loss); - metrics.insert("validation_loss".to_owned(), last_epoch.loss); - metrics.insert("directional_accuracy".to_owned(), last_epoch.accuracy); - metrics.insert("mae".to_owned(), last_epoch.loss); - metrics.insert("rmse".to_owned(), last_epoch.loss.sqrt()); - metrics.insert("r_squared".to_owned(), 1.0 - last_epoch.loss.min(1.0)); - } - - // Add other available metrics - let perf_metrics = self.get_performance_metrics(); - for (key, value) in perf_metrics { - if key.contains("loss") || key.contains("accuracy") { - metrics.insert(key, value); - } - } - - metrics - } - - /// Get inference performance statistics - fn get_inference_stats(&self) -> HashMap { - let mut stats = HashMap::new(); - - // Calculate average latency from latency histogram - let total_inferences = self - .total_inferences - .load(std::sync::atomic::Ordering::Relaxed) as f64; - if total_inferences > 0.0 { - // Simulate latency calculation from internal metrics - let avg_latency = self.config.target_latency_us as f64 * 0.8; // Assume 80% of target - stats.insert("avg_latency_us".to_owned(), avg_latency); - - // Calculate throughput based on latency - let throughput_pps = if avg_latency > 0.0 { - 1_000_000.0 / avg_latency // predictions per second - } else { - 0.0 - }; - stats.insert("throughput_pps".to_owned(), throughput_pps); - } - - stats - } - - /// Extract SSD layer weights from the model - fn extract_ssd_weights(&self) -> Vec> { - // In a real implementation, this would extract actual SSD layer weights - // For now, return structured weight data based on model configuration - let num_layers = self.config.num_layers; - let d_model = self.config.d_model; - let expand = self.config.expand; - - let mut weights = Vec::new(); - for layer in 0..num_layers { - // Each SSD layer has weights of size [d_model * expand, d_model] - let layer_size = d_model * expand; - let mut layer_weights = Vec::with_capacity(layer_size); - - // Generate realistic weight values based on layer index - let scale = 1.0 / (layer + 1) as f32; - for i in 0..layer_size { - let weight = scale * (i as f32 / layer_size as f32 - 0.5); - layer_weights.push(weight); - } - weights.push(layer_weights); - } - - weights - } - - /// Extract input projection weights - fn extract_input_projection_weights(&self) -> Vec { - let d_model = self.config.d_model; - let mut weights = Vec::with_capacity(d_model); - - // Generate input projection weights - for i in 0..d_model { - let weight = (i as f32 / d_model as f32 - 0.5) * 0.1; - weights.push(weight); - } - - weights - } - - /// Extract output projection weights - fn extract_output_projection_weights(&self) -> Vec { - let d_model = self.config.d_model; - let mut weights = Vec::with_capacity(d_model); - - // Generate output projection weights - for i in 0..d_model { - let weight = (i as f32 / d_model as f32 - 0.5) * 0.05; - weights.push(weight); - } - - weights - } - - /// Extract layer normalization weights - fn extract_layer_norm_weights(&self) -> Vec> { - let num_layers = self.config.num_layers; - let d_model = self.config.d_model; - let mut weights = Vec::new(); - - for _layer in 0..num_layers { - let mut layer_norm = Vec::with_capacity(d_model); - // Layer norm weights typically start at 1.0 - for _i in 0..d_model { - layer_norm.push(1.0); - } - weights.push(layer_norm); - } - - weights - } - - /// Extract state space model matrices - fn extract_ssm_matrices(&self, matrix_type: &str) -> Vec> { - let num_layers = self.config.num_layers; - let d_state = self.config.d_state; - let d_model = self.config.d_model; - let mut matrices = Vec::new(); - - for layer in 0..num_layers { - let matrix_size = match matrix_type { - "A" => d_state * d_state, // A matrix is [d_state, d_state] - "B" => d_state * d_model, // B matrix is [d_state, d_model] - "C" => d_model * d_state, // C matrix is [d_model, d_state] - _ => d_state, - }; - - let mut matrix = Vec::with_capacity(matrix_size); - let scale = match matrix_type { - "A" => -0.1, // A matrices typically have negative values for stability - "B" => 0.1, - "C" => 0.1, - _ => 0.1, - }; - - for i in 0..matrix_size { - let value = scale * (i as f32 / matrix_size as f32 - 0.5) * (layer + 1) as f32; - matrix.push(value); - } - matrices.push(matrix); - } - - matrices - } - - /// Extract delta parameters for selective state space - fn extract_delta_params(&self) -> Vec { - let d_model = self.config.d_model; - let mut deltas = Vec::with_capacity(d_model); - - // Delta parameters control the timescale of state updates - for i in 0..d_model { - // Initialize with reasonable timescale values - let delta = 1.0 + (i as f32 / d_model as f32) * 0.1; - deltas.push(delta); - } - - deltas - } - - /// Restore SSD layer weights - fn restore_ssd_weights(&mut self, weights: &[Vec]) { - debug!("Restoring {} SSD layer weight matrices", weights.len()); - - // Validate weight matrix dimensions against config - let expected_layers = self.config.num_layers; - if weights.len() != expected_layers { - warn!( - "SSD weight count mismatch: expected {} layers, got {}", - expected_layers, - weights.len() - ); - } - - // Store weights in SSD layers - using actual struct fields - for (layer_idx, (_layer, layer_weights)) in - self.ssd_layers.iter_mut().zip(weights.iter()).enumerate() - { - // Update the actual layer weights (this depends on SSDLayer implementation) - // For now, we'll store in optimizer_state as a workaround - let layer_key = format!("ssd_layer_{}", layer_idx); - if let Ok(tensor) = - Tensor::from_slice(layer_weights, (layer_weights.len(),), &Device::Cpu) - { - self.optimizer_state.insert(layer_key, tensor); - } - } - - // Validate individual layer weight dimensions - for (layer_idx, layer_weights) in weights.into_iter().enumerate() { - let expected_size = self.config.d_model * self.config.d_model; // Simplified square matrix - if layer_weights.len() != expected_size { - warn!( - "Layer {} weight size mismatch: expected {}, got {}", - layer_idx, - expected_size, - layer_weights.len() - ); - } - - // Validate weight values are finite - let invalid_count = layer_weights.iter().filter(|&&w| !w.is_finite()).count(); - if invalid_count > 0 { - warn!( - "Layer {} contains {} invalid weight values", - layer_idx, invalid_count - ); - } - - debug!( - "SSD Layer {}: {} weights loaded, range [{:.4}, {:.4}]", - layer_idx, - layer_weights.len(), - layer_weights.iter().fold(f32::INFINITY, |a, &b| a.min(b)), - layer_weights - .iter() - .fold(f32::NEG_INFINITY, |a, &b| a.max(b)) - ); - } - - info!( - "Successfully restored {} SSD layer weight matrices", - weights.len() - ); - } - - /// Restore input projection weights - fn restore_input_projection_weights(&mut self, weights: &[f32]) { - debug!("Restoring {} input projection weights", weights.len()); - - // Validate weight dimensions (input projection typically projects from vocab_size to d_model, simplified as d_model * d_model) - let expected_size = self.config.d_model * self.config.d_model; - if weights.len() != expected_size { - warn!( - "Input projection weight size mismatch: expected {}, got {}", - expected_size, - weights.len() - ); - } - - // Validate weight values are finite - let invalid_count = weights.iter().filter(|&&w| !w.is_finite()).count(); - if invalid_count > 0 { - warn!( - "Input projection contains {} invalid weight values", - invalid_count - ); - } - - // Store input projection weights using actual struct field - // The actual input_projection is a Linear layer, store in optimizer_state as workaround - let key = "input_projection_weights".to_owned(); - if let Ok(tensor) = Tensor::from_slice(weights, (weights.len(),), &Device::Cpu) { - self.optimizer_state.insert(key, tensor); - } - - let weight_range = if !weights.is_empty() { - ( - weights.iter().fold(f32::INFINITY, |a, &b| a.min(b)), - weights.iter().fold(f32::NEG_INFINITY, |a, &b| a.max(b)), - ) - } else { - (0.0, 0.0) - }; - - info!( - "Input projection weights restored: {} parameters, range [{:.4}, {:.4}]", - weights.len(), - weight_range.0, - weight_range.1 - ); - } - - /// Restore output projection weights - fn restore_output_projection_weights(&mut self, weights: &[f32]) { - debug!("Restoring {} output projection weights", weights.len()); - - // Validate weight dimensions (output projection typically projects from d_model to vocab_size, simplified as d_model * d_model) - let expected_size = self.config.d_model * self.config.d_model; - if weights.len() != expected_size { - warn!( - "Output projection weight size mismatch: expected {}, got {}", - expected_size, - weights.len() - ); - } - - // Validate weight values are finite - let invalid_count = weights.iter().filter(|&&w| !w.is_finite()).count(); - if invalid_count > 0 { - warn!( - "Output projection contains {} invalid weight values", - invalid_count - ); - } - - // Store output projection weights using actual struct field - // The actual output_projection is a Linear layer, store in optimizer_state as workaround - let key = "output_projection_weights".to_owned(); - if let Ok(tensor) = Tensor::from_slice(weights, (weights.len(),), &Device::Cpu) { - self.optimizer_state.insert(key, tensor); - } - - let weight_range = if !weights.is_empty() { - ( - weights.iter().fold(f32::INFINITY, |a, &b| a.min(b)), - weights.iter().fold(f32::NEG_INFINITY, |a, &b| a.max(b)), - ) - } else { - (0.0, 0.0) - }; - - info!( - "Output projection weights restored: {} parameters, range [{:.4}, {:.4}]", - weights.len(), - weight_range.0, - weight_range.1 - ); - } - - /// Restore layer normalization weights - fn restore_layer_norm_weights(&mut self, weights: &[Vec]) { - debug!("Restoring {} layer norm weight matrices", weights.len()); - - // Validate layer count - let expected_layers = self.config.num_layers; - if weights.len() != expected_layers { - warn!( - "Layer norm count mismatch: expected {} layers, got {}", - expected_layers, - weights.len() - ); - } - - // Store layer norm weights using actual struct field - // The layer_norms field contains actual LayerNorm objects, store in optimizer_state as workaround - for (idx, layer_weights) in weights.into_iter().enumerate() { - let key = format!("layer_norm_weights_{}", idx); - if let Ok(tensor) = - Tensor::from_slice(layer_weights, (layer_weights.len(),), &Device::Cpu) - { - self.optimizer_state.insert(key, tensor); - } - } - - // Validate individual layer norm weights - for (layer_idx, layer_weights) in weights.into_iter().enumerate() { - let expected_size = self.config.d_model; // Layer norm has d_model parameters - if layer_weights.len() != expected_size { - warn!( - "Layer norm {} weight size mismatch: expected {}, got {}", - layer_idx, - expected_size, - layer_weights.len() - ); - } - - // Validate weight values are finite and positive (layer norm weights should be positive) - let invalid_count = layer_weights - .iter() - .filter(|&&w| !w.is_finite() || w <= 0.0) - .count(); - if invalid_count > 0 { - warn!( - "Layer norm {} contains {} invalid weight values (non-positive or non-finite)", - layer_idx, invalid_count - ); - } - - let weight_range = if !layer_weights.is_empty() { - ( - layer_weights.iter().fold(f32::INFINITY, |a, &b| a.min(b)), - layer_weights - .iter() - .fold(f32::NEG_INFINITY, |a, &b| a.max(b)), - ) - } else { - (0.0, 0.0) - }; - - debug!( - "Layer norm {}: {} weights, range [{:.4}, {:.4}]", - layer_idx, - layer_weights.len(), - weight_range.0, - weight_range.1 - ); - } - - info!( - "Successfully restored {} layer normalization weight matrices", - weights.len() - ); - } - - /// Restore state space model matrices - fn restore_ssm_matrices(&mut self, matrix_type: &str, matrices: &[Vec]) { - debug!("Restoring {} {} matrices", matrices.len(), matrix_type); - - // Validate matrix count - let expected_layers = self.config.num_layers; - if matrices.len() != expected_layers { - warn!( - "{} matrix count mismatch: expected {} layers, got {}", - matrix_type, - expected_layers, - matrices.len() - ); - } - - // Store matrices in appropriate fields based on type - match matrix_type { - "A" => { - let key = "ssm_A_matrices".to_owned(); - for (idx, matrix) in matrices.into_iter().enumerate() { - let matrix_key = format!("{}_{}", key, idx); - if let Ok(tensor) = Tensor::from_slice(matrix, (matrix.len(),), &Device::Cpu) { - self.optimizer_state.insert(matrix_key, tensor); - } - } - info!( - "Restored {} A matrices for state space model", - matrices.len() - ); - }, - "B" => { - let key = "ssm_B_matrices".to_owned(); - for (idx, matrix) in matrices.into_iter().enumerate() { - let matrix_key = format!("{}_{}", key, idx); - if let Ok(tensor) = Tensor::from_slice(matrix, (matrix.len(),), &Device::Cpu) { - self.optimizer_state.insert(matrix_key, tensor); - } - } - info!( - "Restored {} B matrices for state space model", - matrices.len() - ); - }, - "C" => { - let key = "ssm_C_matrices".to_owned(); - for (idx, matrix) in matrices.into_iter().enumerate() { - let matrix_key = format!("{}_{}", key, idx); - if let Ok(tensor) = Tensor::from_slice(matrix, (matrix.len(),), &Device::Cpu) { - self.optimizer_state.insert(matrix_key, tensor); - } - } - info!( - "Restored {} C matrices for state space model", - matrices.len() - ); - }, - _ => { - warn!("Unknown SSM matrix type: {}", matrix_type); - return; - }, - } - - // Validate individual matrices - for (layer_idx, matrix) in matrices.into_iter().enumerate() { - let expected_size = match matrix_type { - "A" => self.config.d_state * self.config.d_state, - "B" => self.config.d_state * self.config.d_model, - "C" => self.config.d_model * self.config.d_state, - _ => self.config.d_state, - }; - - if matrix.len() != expected_size { - warn!( - "SSM {} matrix {} size mismatch: expected {}, got {}", - matrix_type, - layer_idx, - expected_size, - matrix.len() - ); - } - - // Validate matrix values are finite - let invalid_count = matrix.iter().filter(|&&v| !v.is_finite()).count(); - if invalid_count > 0 { - warn!( - "SSM {} matrix {} contains {} invalid values", - matrix_type, layer_idx, invalid_count - ); - } - - let matrix_range = if !matrix.is_empty() { - ( - matrix.iter().fold(f32::INFINITY, |a, &b| a.min(b)), - matrix.iter().fold(f32::NEG_INFINITY, |a, &b| a.max(b)), - ) - } else { - (0.0, 0.0) - }; - - debug!( - "SSM {} matrix {}: {} values, range [{:.4}, {:.4}]", - matrix_type, - layer_idx, - matrix.len(), - matrix_range.0, - matrix_range.1 - ); - } - } - - /// Restore delta parameters - fn restore_delta_params(&mut self, deltas: &[f32]) { - debug!("Restoring {} delta parameters", deltas.len()); - - // Validate delta parameter count - let expected_size = self.config.d_model; - if deltas.len() != expected_size { - warn!( - "Delta parameter count mismatch: expected {}, got {}", - expected_size, - deltas.len() - ); - } - - // Validate delta values are finite and positive (deltas control timescales) - let invalid_count = deltas - .iter() - .filter(|&&d| !d.is_finite() || d <= 0.0) - .count(); - if invalid_count > 0 { - warn!( - "Delta parameters contain {} invalid values (non-positive or non-finite)", - invalid_count - ); - } - - // Store delta parameters using optimizer_state since the field doesn't exist - let key = "ssm_delta_params".to_owned(); - if let Ok(tensor) = Tensor::from_slice(deltas, (deltas.len(),), &Device::Cpu) { - self.optimizer_state.insert(key, tensor); - } - - let delta_range = if !deltas.is_empty() { - ( - deltas.iter().fold(f32::INFINITY, |a, &b| a.min(b)), - deltas.iter().fold(f32::NEG_INFINITY, |a, &b| a.max(b)), - ) - } else { - (0.0, 0.0) - }; - - info!( - "Delta parameters restored: {} values, range [{:.4}, {:.4}]", - deltas.len(), - delta_range.0, - delta_range.1 - ); - - // Log statistics for debugging - if !deltas.is_empty() { - let mean = deltas.iter().sum::() / deltas.len() as f32; - let variance = - deltas.iter().map(|&x| (x - mean).powi(2)).sum::() / deltas.len() as f32; - debug!( - "Delta parameter statistics: mean={:.4}, variance={:.4}", - mean, variance - ); - } - } -} - /// Serializable state for `TFT` model #[derive(Debug, Clone, Serialize, Deserialize)] pub struct TFTCheckpointState { @@ -1161,132 +576,6 @@ impl Checkpointable for TGGN { } } -impl TGGN { - /// Extract graph statistics for checkpoint state - fn extract_graph_statistics(&self) -> HashMap { - let mut stats = HashMap::new(); - - // Node statistics - now using proper graph access method - let (node_count, edge_count) = self.get_graph_stats(); - let node_count_f64 = node_count as f64; - let edge_count_f64 = edge_count as f64; - - stats.insert("node_count".to_owned(), node_count_f64); - stats.insert("max_nodes".to_owned(), self.config().max_nodes as f64); - stats.insert( - "node_utilization".to_owned(), - node_count_f64 / self.config().max_nodes as f64, - ); - - // Edge statistics - stats.insert("edge_count".to_owned(), edge_count_f64); - stats.insert("max_edges".to_owned(), self.config().max_edges as f64); - stats.insert( - "edge_utilization".to_owned(), - edge_count_f64 / self.config().max_edges as f64, - ); - - // Graph density - if node_count_f64 > 1.0 { - let max_edges = node_count_f64 * (node_count_f64 - 1.0) / 2.0; - stats.insert("graph_density".to_owned(), edge_count_f64 / max_edges); - } else { - stats.insert("graph_density".to_owned(), 0.0); - } - - // Average degree - if node_count_f64 > 0.0 { - stats.insert( - "avg_degree".to_owned(), - (2.0 * edge_count_f64) / node_count_f64, - ); - } else { - stats.insert("avg_degree".to_owned(), 0.0); - } - - stats - } - - /// Extract message passing weights from the model layers - fn extract_message_passing_weights(&self) -> Vec { - // Extract weights from the message passing layers - // This would depend on the actual implementation of the TGGN layers - let mut weights = Vec::new(); - - // Simulate extracting weights from different layers - for layer_idx in 0..self.config().num_layers { - // Add simulated layer weights (in real implementation, extract from actual layers) - let layer_size = self.config().node_dim * self.config().edge_dim; - for i in 0..layer_size { - weights.push((layer_idx as f32 + (i as f32).rem_euclid(10.0)) * 0.1); - } - } - - weights - } - - /// Extract gating mechanism weights - fn extract_gating_weights(&self) -> Vec { - // Extract weights from gating mechanisms - let mut gating_weights = Vec::new(); - - // Simulate extracting gating weights (in real implementation, extract from actual gates) - let num_gates = self.config().num_layers * 3; // Assume 3 gates per layer - for gate_idx in 0..num_gates { - let gate_size = self.config().node_dim; - for i in 0..gate_size { - gating_weights.push(((gate_idx * gate_size + i) as f32).rem_euclid(100.0) * 0.01); - } - } - - gating_weights - } - - /// Calculate average inference latency from internal statistics - fn calculate_avg_inference_latency(&self) -> u64 { - let total_inferences = self - .inference_count() - .load(std::sync::atomic::Ordering::Relaxed); - - if total_inferences > 0 { - // Simulate calculation from internal timing statistics - // In real implementation, this would use actual timing data - let base_latency_ns = match self.config().num_layers { - 1..=3 => 1_000_000, // 1ms for small models - 4..=6 => 5_000_000, // 5ms for medium models - _ => 10_000_000, // 10ms for large models - }; - - // Add complexity factor based on graph size - let graph_complexity = (100 /* production */ + 200/* production */) as u64; - let complexity_factor = (graph_complexity / 1000).max(1); - - base_latency_ns * complexity_factor - } else { - 0 - } - } - - /// Calculate average graph update latency - fn calculate_avg_graph_update_latency(&self) -> u64 { - let total_updates = self - .graph_updates() - .load(std::sync::atomic::Ordering::Relaxed); - - if total_updates > 0 { - // Simulate calculation from internal timing statistics - let base_update_latency_ns = 500_000; // 0.5ms base - - // Scale with graph size - let graph_size_factor = - ((100 /* production */ + 200/* production */) / 100) as u64; - - base_update_latency_ns * graph_size_factor - } else { - 0 - } - } -} /// Serializable state for Liquid Neural Network #[derive(Debug, Clone, Serialize, Deserialize)] diff --git a/crates/ml/src/diffusion/mod.rs b/crates/ml/src/diffusion/mod.rs index d4c8f6d7d..110b3bf9f 100644 --- a/crates/ml/src/diffusion/mod.rs +++ b/crates/ml/src/diffusion/mod.rs @@ -1,19 +1,13 @@ -//! Diffusion model (DDPM/DDIM) for price path generation. +//! Diffusion model (DDPM/DDIM) for price path generation //! -//! Implements a denoising diffusion probabilistic model with: -//! - Cosine/linear noise schedules -//! - Fully-connected denoiser with sinusoidal time embedding -//! - DDIM sampling for fast inference -//! - UnifiedTrainable adapter for the training pipeline +//! This module re-exports the `ml-supervised` Diffusion implementation and keeps bridge +//! modules that depend on types from both `ml-supervised` and the parent `ml` crate. -pub mod config; -pub mod denoiser; -pub mod noise; -pub mod sampler; +// Re-export everything from the ml-supervised diffusion module +pub use ml_supervised::diffusion::*; + +// Bridge modules that depend on ml-internal types (UnifiedTrainable) pub mod trainable; -pub use config::{DiffusionConfig, NoiseSchedule}; -pub use denoiser::Denoiser; -pub use noise::NoiseScheduler; -pub use sampler::DDIMSampler; +// Re-export bridge types pub use trainable::DiffusionTrainableAdapter; diff --git a/crates/ml/src/hyperopt/adapters/async_data_loader.rs b/crates/ml/src/hyperopt/adapters/async_data_loader.rs deleted file mode 100644 index 1b92b915d..000000000 --- a/crates/ml/src/hyperopt/adapters/async_data_loader.rs +++ /dev/null @@ -1,646 +0,0 @@ -//! Async Data Loader for GPU Training Optimization -//! -//! This module implements prefetch-based async data loading to improve GPU utilization -//! by overlapping data preparation with GPU computation. Key features: -//! -//! - Prefetch 2-3 batches ahead while GPU trains -//! - Thread-safe channel-based communication -//! - Graceful shutdown and error handling -//! - Zero-copy tensor transfer where possible -//! -//! ## Performance Impact -//! -//! - CPU utilization: 7% → 30-40% -//! - GPU utilization: 78% → 90-95% -//! - Training speedup: 20-30% -//! -//! ## Usage Example -//! -//! ```rust,no_run -//! use ml::hyperopt::adapters::async_data_loader::AsyncDataLoader; -//! use candle_core::Device; -//! -//! # fn example() -> anyhow::Result<()> { -//! let device = Device::cuda_if_available(0)?; -//! let data = vec![/* training data */]; -//! let batch_size = 32; -//! let prefetch_count = 3; -//! -//! let mut loader = AsyncDataLoader::new(data, batch_size, prefetch_count, &device)?; -//! -//! while let Some((features, targets)) = loader.next_batch() { -//! // GPU trains on current batch while next batches are being prepared -//! } -//! # Ok(()) -//! # } -//! ``` - -use anyhow::Result; -use candle_core::{Device, Tensor}; -use std::sync::mpsc::{sync_channel, Receiver, RecvTimeoutError, SyncSender, TryRecvError}; -use std::thread::{self, JoinHandle}; -use std::time::Duration; -use tracing::{debug, warn}; - -use crate::MLError; - -/// Async data loader with prefetching for GPU training optimization -/// -/// This loader runs a background thread that prepares batches on CPU and transfers -/// them to GPU ahead of time. The training loop can then consume batches without -/// waiting for data preparation, maximizing GPU utilization. -/// -/// ## Architecture -/// -/// ```text -/// ┌─────────────────────────────────────────────────────────────┐ -/// │ AsyncDataLoader │ -/// │ │ -/// │ ┌─────────────────┐ ┌──────────────────────────┐ │ -/// │ │ Prefetch Thread │ ────> │ Bounded Channel (size=3) │ │ -/// │ │ (CPU prep) │ │ (GPU tensors ready) │ │ -/// │ └─────────────────┘ └──────────────────────────┘ │ -/// │ │ │ -/// └────────────────────────────────────────┼─────────────────────┘ -/// │ -/// ▼ -/// Training Loop -/// (GPU compute) -/// ``` -/// -/// ## Thread Safety -/// -/// - Channel-based communication (thread-safe by design) -/// - No shared mutable state -/// - Graceful shutdown on drop or error -#[derive(Debug)] -pub struct AsyncDataLoader { - /// Receiver for prefetched batches - receiver: Receiver>, - /// Background prefetch thread - prefetch_thread: Option>, - /// Total number of batches - total_batches: usize, - /// Current batch index - current_batch: usize, - /// Device for error reporting - device: Device, -} - -impl AsyncDataLoader { - /// Create a new async data loader with prefetching - /// - /// # Arguments - /// - /// * `data` - Training data as (feature_tensor, target_tensor) pairs - /// * `batch_size` - Number of samples per batch - /// * `prefetch_count` - Number of batches to prefetch (typically 2-3) - /// * `device` - GPU device for tensor transfers - /// - /// # Returns - /// - /// AsyncDataLoader ready to stream batches - /// - /// # Errors - /// - /// Returns error if: - /// - Data is empty - /// - Device clone fails - /// - Thread spawn fails - pub fn new( - data: Vec<(Tensor, Tensor)>, - batch_size: usize, - prefetch_count: usize, - device: &Device, - ) -> Result { - if data.is_empty() { - return Err( - MLError::InvalidInput("Cannot create loader with empty data".to_owned()).into(), - ); - } - - if batch_size == 0 { - return Err(MLError::InvalidInput("Batch size must be > 0".to_owned()).into()); - } - - let total_batches = (data.len() + batch_size - 1) / batch_size; - - // CRITICAL FIX: Explicitly clone device for storage - // Issue: device.clone() on &Device might not properly clone CUDA devices - // Solution: Use Device::clone() explicitly to ensure proper cloning - let device_owned = (*device).clone(); - let device_clone = device_owned.clone(); - - debug!( - "Creating AsyncDataLoader: {} samples, batch_size={}, prefetch={}, batches={}, device={:?}", - data.len(), - batch_size, - prefetch_count, - total_batches, - device_owned - ); - - // Create bounded channel - blocks if full (backpressure) - let (sender, receiver) = sync_channel(prefetch_count); - - // Spawn prefetch thread - let prefetch_thread = thread::spawn(move || { - Self::prefetch_worker(data, batch_size, sender, device_clone); - }); - - Ok(Self { - receiver, - prefetch_thread: Some(prefetch_thread), - total_batches, - current_batch: 0, - device: device_owned, - }) - } - - /// Background worker that prefetches and prepares batches - /// - /// This runs in a separate thread and: - /// 1. Chunks data into batches - /// 2. Concatenates tensors for each batch ON CPU - /// 3. Sends via channel to training loop - /// - /// NOTE: GPU transfer happens in main thread to avoid CUDA context issues. - /// CUDA contexts are thread-local, so transferring in background thread can fail. - /// - /// Stops when: - /// - All batches processed - /// - Channel receiver drops (training stopped) - /// - Error occurs - fn prefetch_worker( - data: Vec<(Tensor, Tensor)>, - batch_size: usize, - sender: SyncSender>, - _device: Device, // Unused - kept for API compatibility - ) { - debug!("Prefetch worker started: {} samples", data.len()); - - for (batch_idx, batch_data) in data.chunks(batch_size).enumerate() { - // Note: SyncSender doesn't have is_disconnected(), we'll rely on send() error instead - - // Prepare batch on CPU ONLY (no GPU transfer in background thread) - let batch_result = Self::prepare_batch_cpu(batch_data); - - // Send to training loop (blocks if channel full) - if let Err(e) = sender.send(batch_result) { - warn!("Prefetch worker failed to send batch {}: {}", batch_idx, e); - break; - } - - if batch_idx % 20 == 0 { - debug!("Prefetch worker: prepared batch {}", batch_idx); - } - } - - debug!("Prefetch worker finished"); - } - - /// Prepare a single batch on CPU: concatenate tensors - /// - /// This is the CPU-intensive operation we want to overlap with GPU training. - /// GPU transfer happens in the main thread to avoid CUDA context issues. - /// - /// # Arguments - /// - /// * `batch_data` - Slice of (feature, target) tensor pairs - /// - /// # Returns - /// - /// Batched tensors on CPU, or error if preparation fails - fn prepare_batch_cpu(batch_data: &[(Tensor, Tensor)]) -> Result<(Tensor, Tensor), MLError> { - if batch_data.is_empty() { - return Err(MLError::InvalidInput("Empty batch".to_owned())); - } - - let actual_batch_size = batch_data.len(); - - // Concatenate feature tensors along batch dimension - let feature_tensors: Vec<&Tensor> = batch_data.iter().map(|(f, _)| f).collect(); - let batched_features = if actual_batch_size == 1 { - feature_tensors[0].clone() - } else { - Tensor::cat( - &feature_tensors - .iter() - .map(|t| (*t).clone()) - .collect::>(), - 0, - ) - .map_err(|e| MLError::TensorCreationError { - operation: "concatenate features".to_owned(), - reason: e.to_string(), - })? - }; - - // Concatenate target tensors along batch dimension - let target_tensors: Vec<&Tensor> = batch_data.iter().map(|(_, t)| t).collect(); - let batched_targets = if actual_batch_size == 1 { - target_tensors[0].clone() - } else { - Tensor::cat( - &target_tensors - .iter() - .map(|t| (*t).clone()) - .collect::>(), - 0, - ) - .map_err(|e| MLError::TensorCreationError { - operation: "concatenate targets".to_owned(), - reason: e.to_string(), - })? - }; - - // Return CPU tensors - GPU transfer happens in main thread - Ok((batched_features, batched_targets)) - } - - /// Get the next batch (non-blocking) - /// - /// Returns `None` when all batches consumed or on error. - /// - /// NOTE: This method transfers tensors from CPU to GPU in the main thread - /// to avoid CUDA context issues. The background thread only prepares batches on CPU. - /// - /// # Returns - /// - /// - `Some((features, targets))` - Next batch ready on GPU - /// - `None` - No more batches or error occurred - pub fn next_batch(&mut self) -> Option<(Tensor, Tensor)> { - debug!( - "next_batch() called: current_batch={}, total_batches={}", - self.current_batch, self.total_batches - ); - - if self.current_batch >= self.total_batches { - debug!("next_batch() returning None: reached total_batches"); - return None; - } - - debug!("next_batch() waiting for receiver.recv_timeout(30s)..."); - match self.receiver.recv_timeout(Duration::from_secs(30)) { - Ok(Ok((features, targets))) => { - self.current_batch += 1; - debug!( - "Received batch {} from prefetch worker (features: {:?}, targets: {:?})", - self.current_batch, - features.device(), - targets.device() - ); - - // Transfer to GPU in main thread (CUDA context is valid here) - debug!( - "Batch {} - Before transfer: features device={:?}, targets device={:?}, target device={:?}", - self.current_batch, - features.device(), - targets.device(), - self.device - ); - - let features_gpu = match features.to_device(&self.device) { - Ok(t) => { - // CRITICAL FIX: Verify transfer actually happened - if t.device().is_cuda() != self.device.is_cuda() { - warn!( - "Device transfer failed: expected {:?}, got {:?} (to_device returned wrong device)", - self.device, t.device() - ); - return None; - } - debug!( - "Batch {} - After transfer: features device={:?}", - self.current_batch, - t.device() - ); - t - }, - Err(e) => { - warn!( - "Failed to transfer features to GPU at batch {}: {}", - self.current_batch - 1, - e - ); - return None; - }, - }; - - let targets_gpu = match targets.to_device(&self.device) { - Ok(t) => { - // CRITICAL FIX: Verify transfer actually happened - if t.device().is_cuda() != self.device.is_cuda() { - warn!( - "Device transfer failed: expected {:?}, got {:?} (to_device returned wrong device)", - self.device, t.device() - ); - return None; - } - debug!( - "Batch {} - After transfer: targets device={:?}", - self.current_batch, - t.device() - ); - t - }, - Err(e) => { - warn!( - "Failed to transfer targets to GPU at batch {}: {}", - self.current_batch - 1, - e - ); - return None; - }, - }; - - Some((features_gpu, targets_gpu)) - }, - Ok(Err(e)) => { - warn!( - "Batch preparation error at batch {}: {}", - self.current_batch, e - ); - None - }, - Err(RecvTimeoutError::Timeout) => { - warn!( - "Prefetch timeout after 30s at batch {} (expected {} total batches)", - self.current_batch, self.total_batches - ); - None - }, - Err(RecvTimeoutError::Disconnected) => { - // Channel closed (worker finished or crashed) - debug!("Prefetch channel closed at batch {}", self.current_batch); - None - }, - } - } - - /// Try to get the next batch without blocking - /// - /// Useful for checking if data is ready without waiting. - /// - /// NOTE: This method transfers tensors from CPU to GPU in the main thread - /// to avoid CUDA context issues. - /// - /// # Returns - /// - /// - `Some((features, targets))` - Batch ready immediately on GPU - /// - `None` - No batch ready yet (try again later) or stream ended - pub fn try_next_batch(&mut self) -> Option<(Tensor, Tensor)> { - if self.current_batch >= self.total_batches { - return None; - } - - match self.receiver.try_recv() { - Ok(Ok((features, targets))) => { - self.current_batch += 1; - - // Transfer to GPU in main thread (CUDA context is valid here) - debug!( - "try_next_batch - Batch {} - Before transfer: features device={:?}, targets device={:?}, target device={:?}", - self.current_batch, - features.device(), - targets.device(), - self.device - ); - - let features_gpu = match features.to_device(&self.device) { - Ok(t) => { - debug!( - "try_next_batch - Batch {} - After transfer: features device={:?}", - self.current_batch, - t.device() - ); - t - }, - Err(e) => { - warn!("Failed to transfer features to GPU: {}", e); - return None; - }, - }; - - let targets_gpu = match targets.to_device(&self.device) { - Ok(t) => { - debug!( - "try_next_batch - Batch {} - After transfer: targets device={:?}", - self.current_batch, - t.device() - ); - t - }, - Err(e) => { - warn!("Failed to transfer targets to GPU: {}", e); - return None; - }, - }; - - Some((features_gpu, targets_gpu)) - }, - Ok(Err(e)) => { - warn!("Batch preparation error: {}", e); - None - }, - Err(TryRecvError::Empty) => None, // No batch ready yet - Err(TryRecvError::Disconnected) => None, // Worker finished - } - } - - /// Get progress: (current_batch, total_batches) - pub fn progress(&self) -> (usize, usize) { - (self.current_batch, self.total_batches) - } - - /// Check if all batches have been consumed - pub fn is_complete(&self) -> bool { - self.current_batch >= self.total_batches - } -} - -impl Drop for AsyncDataLoader { - /// Ensure prefetch thread is joined on drop - fn drop(&mut self) { - if let Some(handle) = self.prefetch_thread.take() { - // CRITICAL FIX: Explicitly drop receiver BEFORE join() to signal worker to stop - // Without this, thread may block on sender.send() waiting for receiver, - // while join() waits for thread - creating a 60s deadlock - drop(std::mem::replace( - &mut self.receiver, - sync_channel(1).1, // Dummy receiver to satisfy type - )); - - // Wait for worker to finish (should be quick now that receiver is dropped) - if let Err(e) = handle.join() { - warn!("Prefetch thread panicked during join: {:?}", e); - } - } - } -} - -#[cfg(test)] -mod tests { - use super::*; - - fn create_test_data(count: usize, device: &Device) -> Result> { - let mut data = Vec::new(); - for i in 0..count { - let features = Tensor::new(&[i as f64; 10], device)?.reshape((1, 10, 1))?; - let target = Tensor::new(&[i as f64], device)?.reshape((1, 1, 1))?; - data.push((features, target)); - } - Ok(data) - } - - #[test] - fn test_async_loader_basic() -> Result<()> { - let device = Device::Cpu; - let data = create_test_data(100, &device)?; - let batch_size = 10; - let prefetch = 2; - - let mut loader = AsyncDataLoader::new(data, batch_size, prefetch, &device)?; - - let mut batch_count = 0; - while let Some(_batch) = loader.next_batch() { - batch_count += 1; - } - - assert_eq!(batch_count, 10, "Should get 10 batches (100 / 10)"); - assert!(loader.is_complete()); - - Ok(()) - } - - #[test] - fn test_async_loader_partial_batch() -> Result<()> { - let device = Device::Cpu; - let data = create_test_data(95, &device)?; - let batch_size = 10; - let prefetch = 2; - - let mut loader = AsyncDataLoader::new(data, batch_size, prefetch, &device)?; - - let mut batch_count = 0; - while let Some(_batch) = loader.next_batch() { - batch_count += 1; - } - - assert_eq!( - batch_count, 10, - "Should get 10 batches (95 / 10 = 9.5 -> 10)" - ); - - Ok(()) - } - - #[test] - fn test_async_loader_progress() -> Result<()> { - let device = Device::Cpu; - let data = create_test_data(50, &device)?; - let batch_size = 10; - let prefetch = 2; - - let mut loader = AsyncDataLoader::new(data, batch_size, prefetch, &device)?; - - assert_eq!(loader.progress(), (0, 5)); - - loader.next_batch(); - assert_eq!(loader.progress(), (1, 5)); - - loader.next_batch(); - loader.next_batch(); - assert_eq!(loader.progress(), (3, 5)); - - Ok(()) - } - - #[test] - fn test_async_loader_empty_data() { - let device = Device::Cpu; - let data: Vec<(Tensor, Tensor)> = vec![]; - let result = AsyncDataLoader::new(data, 10, 2, &device); - - assert!(result.is_err()); - assert!(result.unwrap_err().to_string().contains("empty")); - } - - #[test] - fn test_async_loader_zero_batch_size() -> Result<()> { - let device = Device::Cpu; - let data = create_test_data(10, &device)?; - let result = AsyncDataLoader::new(data, 0, 2, &device); - - assert!(result.is_err()); - assert!(result.unwrap_err().to_string().contains("Batch size")); - - Ok(()) - } - - #[test] - fn test_early_termination() -> Result<()> { - let device = Device::Cpu; - let data = create_test_data(100, &device)?; - let batch_size = 10; - let prefetch = 2; - - let mut loader = AsyncDataLoader::new(data, batch_size, prefetch, &device)?; - - // Consume only 3 batches - loader.next_batch(); - loader.next_batch(); - loader.next_batch(); - - // Drop loader (should cleanly shut down prefetch thread) - drop(loader); - - Ok(()) - } - - #[test] - fn test_try_next_batch() -> Result<()> { - let device = Device::Cpu; - let data = create_test_data(20, &device)?; - let batch_size = 10; - let prefetch = 2; - - let mut loader = AsyncDataLoader::new(data, batch_size, prefetch, &device)?; - - // First try should succeed (prefetch filled channel) - std::thread::sleep(std::time::Duration::from_millis(100)); - assert!(loader.try_next_batch().is_some()); - - Ok(()) - } - - #[test] - fn test_batch_tensor_shapes() -> Result<()> { - let device = Device::Cpu; - let data = create_test_data(25, &device)?; - let batch_size = 10; - let prefetch = 2; - - let mut loader = AsyncDataLoader::new(data, batch_size, prefetch, &device)?; - - // First batch: 10 samples - if let Some((features, targets)) = loader.next_batch() { - assert_eq!(features.dims()[0], 10, "Batch size should be 10"); - assert_eq!(targets.dims()[0], 10, "Batch size should be 10"); - } - - // Second batch: 10 samples - if let Some((features, targets)) = loader.next_batch() { - assert_eq!(features.dims()[0], 10, "Batch size should be 10"); - assert_eq!(targets.dims()[0], 10, "Batch size should be 10"); - } - - // Third batch: 5 samples (partial) - if let Some((features, targets)) = loader.next_batch() { - assert_eq!(features.dims()[0], 5, "Last batch should be 5"); - assert_eq!(targets.dims()[0], 5, "Last batch should be 5"); - } - - Ok(()) - } -} diff --git a/crates/ml/src/hyperopt/adapters/mod.rs b/crates/ml/src/hyperopt/adapters/mod.rs index ae48566ba..69cc738fc 100644 --- a/crates/ml/src/hyperopt/adapters/mod.rs +++ b/crates/ml/src/hyperopt/adapters/mod.rs @@ -52,7 +52,6 @@ pub mod dbn_loader; // Active adapters (production-ready) -pub mod async_data_loader; pub mod continuous_ppo; pub mod dqn; pub mod ensemble; @@ -67,7 +66,6 @@ pub mod diffusion; pub mod xlstm; // Re-export adapters for convenience -pub use async_data_loader::AsyncDataLoader; pub use kan::{KANMetrics, KANParams, KANTrainer}; pub use liquid::{LiquidMetrics, LiquidParams, LiquidTrainer}; pub use continuous_ppo::{ContinuousPPOMetrics, ContinuousPPOParams, ContinuousPPOTrainer}; diff --git a/crates/ml/src/kan/mod.rs b/crates/ml/src/kan/mod.rs index 3bb8168bf..6382994bb 100644 --- a/crates/ml/src/kan/mod.rs +++ b/crates/ml/src/kan/mod.rs @@ -1,15 +1,13 @@ -//! KAN (Kolmogorov-Arnold Network) module. +//! KAN (Kolmogorov-Arnold Network) //! -//! Implements KAN with learnable B-spline activation functions on each edge -//! of the network, replacing fixed activations (ReLU) with data-driven -//! non-linearities. +//! This module re-exports the `ml-supervised` KAN implementation and keeps bridge +//! modules that depend on types from both `ml-supervised` and the parent `ml` crate. -pub mod config; -pub mod layer; -pub mod network; -pub mod spline; +// Re-export everything from the ml-supervised kan module +pub use ml_supervised::kan::*; + +// Bridge modules that depend on ml-internal types (UnifiedTrainable) pub mod trainable; -pub use config::KANConfig; -pub use network::KANNetwork; +// Re-export bridge types pub use trainable::KANTrainableAdapter; diff --git a/crates/ml/src/liquid/cuda/liquid_kernels.cu b/crates/ml/src/liquid/cuda/liquid_kernels.cu deleted file mode 100644 index 956113ce1..000000000 --- a/crates/ml/src/liquid/cuda/liquid_kernels.cu +++ /dev/null @@ -1,513 +0,0 @@ -/** - * CUDA Kernels for Liquid Neural Networks - * - * GPU-accelerated implementation of Liquid Time-constant (LTC) and - * Closed-form Continuous-time (CfC) neural networks for ultra-low - * latency inference in HFT applications. - */ - -#include -#include -#include -#include - -// Fixed-point precision for ultra-low latency operations -#define PRECISION 100000000L // 8 decimal places -#define PRECISION_F 100000000.0f - -/** - * Convert float to fixed-point representation - */ -__device__ __forceinline__ long long float_to_fixed(float value) { - return (long long)(value * PRECISION_F); -} - -/** - * Convert fixed-point to float representation - */ -__device__ __forceinline__ float fixed_to_float(long long value) { - return (float)value / PRECISION_F; -} - -/** - * Fixed-point multiplication with overflow protection - */ -__device__ __forceinline__ long long fixed_mul(long long a, long long b) { - return ((long long)a * (long long)b) / PRECISION; -} - -/** - * Fixed-point division with zero protection - */ -__device__ __forceinline__ long long fixed_div(long long a, long long b) { - if (b == 0) return 0; - return ((long long)a * PRECISION) / (long long)b; -} - -/** - * Activation functions for liquid networks - */ -__device__ __forceinline__ float activation_tanh(float x) { - return tanhf(x); -} - -__device__ __forceinline__ float activation_sigmoid(float x) { - return 1.0f / (1.0f + expf(-x)); -} - -__device__ __forceinline__ float activation_relu(float x) { - return fmaxf(0.0f, x); -} - -__device__ __forceinline__ float activation_linear(float x) { - return x; -} - -/** - * Apply activation function based on type - * 0=Linear, 1=ReLU, 2=Sigmoid, 3=Tanh - */ -__device__ __forceinline__ float apply_activation(float x, int activation_type) { - switch (activation_type) { - case 0: return activation_linear(x); - case 1: return activation_relu(x); - case 2: return activation_sigmoid(x); - case 3: return activation_tanh(x); - default: return activation_tanh(x); - } -} - -/** - * Fused kernel for LTC cell forward pass - * - * Computes multiple LTC neurons in parallel with fused operations: - * 1. Input transformation - * 2. Recurrent computation - * 3. Time constant adaptation - * 4. ODE integration (Euler method) - * 5. Activation application - * - * @param input Input tensor [batch_size, input_size] - * @param hidden_state Current hidden state [batch_size, hidden_size] - * @param input_weights Input weight matrix [hidden_size, input_size] - * @param recurrent_weights Recurrent weight matrix [hidden_size, hidden_size] - * @param bias Bias vector [hidden_size] - * @param time_constants Time constants [hidden_size] - * @param new_hidden_state Output hidden state [batch_size, hidden_size] - * @param dt Integration time step - * @param volatility Market volatility for adaptation - * @param activation_type Activation function type - * @param batch_size Number of samples in batch - * @param input_size Input dimension - * @param hidden_size Hidden dimension - */ -__global__ void fused_ltc_forward( - const float* input, - const float* hidden_state, - const float* input_weights, - const float* recurrent_weights, - const float* bias, - float* time_constants, - float* new_hidden_state, - float dt, - float volatility, - int activation_type, - int batch_size, - int input_size, - int hidden_size -) { - int batch_idx = blockIdx.y * blockDim.y + threadIdx.y; - int neuron_idx = blockIdx.x * blockDim.x + threadIdx.x; - - if (batch_idx >= batch_size || neuron_idx >= hidden_size) return; - - // Shared memory for efficient memory access - extern __shared__ float shared_mem[]; - float* shared_input = shared_mem; - float* shared_hidden = shared_input + input_size; - - // Load input and hidden state to shared memory - if (threadIdx.y == 0 && threadIdx.x < input_size) { - shared_input[threadIdx.x] = input[batch_idx * input_size + threadIdx.x]; - } - if (threadIdx.y == 0 && threadIdx.x < hidden_size) { - shared_hidden[threadIdx.x] = hidden_state[batch_idx * hidden_size + threadIdx.x]; - } - __syncthreads(); - - // Compute input contribution - float input_sum = 0.0f; - for (int i = 0; i < input_size; i++) { - input_sum += input_weights[neuron_idx * input_size + i] * shared_input[i]; - } - - // Compute recurrent contribution - float recurrent_sum = 0.0f; - for (int i = 0; i < hidden_size; i++) { - recurrent_sum += recurrent_weights[neuron_idx * hidden_size + i] * shared_hidden[i]; - } - - // Add bias - float total_input = input_sum + recurrent_sum + bias[neuron_idx]; - - // Apply activation - float activated = apply_activation(total_input, activation_type); - - // Adaptive time constant based on volatility - float base_tau = time_constants[neuron_idx]; - float adapted_tau = base_tau * (1.0f + 0.1f * volatility); // Simple adaptation - adapted_tau = fmaxf(0.01f, fminf(1.0f, adapted_tau)); // Clamp to reasonable range - - // Update time constant - time_constants[neuron_idx] = adapted_tau; - - // LTC dynamics: dx/dt = (1/tau) * (-x + activated) - float current_x = shared_hidden[neuron_idx]; - float dx_dt = (1.0f / adapted_tau) * (-current_x + activated); - - // Euler integration - float new_x = current_x + dt * dx_dt; - - // Store result - new_hidden_state[batch_idx * hidden_size + neuron_idx] = new_x; -} - -/** - * Fused kernel for CfC cell forward pass with backbone network - * - * @param input Input tensor [batch_size, input_size] - * @param hidden_state Current hidden state [batch_size, hidden_size] - * @param backbone_weights Backbone network weights [num_layers][max_layer_size][max_input_size] - * @param backbone_bias Backbone network bias [num_layers][max_layer_size] - * @param layer_sizes Size of each backbone layer [num_layers] - * @param output_weights Final output weights [hidden_size] - * @param new_hidden_state Output hidden state [batch_size, hidden_size] - * @param dt Integration time step - * @param batch_size Number of samples in batch - * @param input_size Input dimension - * @param hidden_size Hidden dimension - * @param num_layers Number of backbone layers - * @param max_layer_size Maximum layer size in backbone - */ -__global__ void fused_cfc_forward( - const float* input, - const float* hidden_state, - const float* backbone_weights, - const float* backbone_bias, - const int* layer_sizes, - const float* output_weights, - float* new_hidden_state, - float dt, - int batch_size, - int input_size, - int hidden_size, - int num_layers, - int max_layer_size -) { - int batch_idx = blockIdx.y * blockDim.y + threadIdx.y; - int neuron_idx = blockIdx.x * blockDim.x + threadIdx.x; - - if (batch_idx >= batch_size || neuron_idx >= hidden_size) return; - - // Shared memory for backbone computation - extern __shared__ float shared_backbone[]; - float* current_layer = shared_backbone; - float* next_layer = shared_backbone + max_layer_size; - - // Initialize first layer with concatenated input and hidden state - if (threadIdx.x < input_size && threadIdx.y == 0) { - current_layer[threadIdx.x] = input[batch_idx * input_size + threadIdx.x]; - } - if (threadIdx.x < hidden_size && threadIdx.y == 0) { - current_layer[input_size + threadIdx.x] = hidden_state[batch_idx * hidden_size + threadIdx.x]; - } - __syncthreads(); - - int current_size = input_size + hidden_size; - - // Forward through backbone layers - for (int layer = 0; layer < num_layers; layer++) { - int layer_size = layer_sizes[layer]; - - if (threadIdx.x < layer_size && threadIdx.y == 0) { - float sum = 0.0f; - - // Compute weighted sum for this neuron - for (int i = 0; i < current_size; i++) { - int weight_idx = layer * max_layer_size * max_layer_size + - threadIdx.x * max_layer_size + i; - sum += backbone_weights[weight_idx] * current_layer[i]; - } - - // Add bias and apply tanh activation - sum += backbone_bias[layer * max_layer_size + threadIdx.x]; - next_layer[threadIdx.x] = tanhf(sum); - } - __syncthreads(); - - // Swap layers - float* temp = current_layer; - current_layer = next_layer; - next_layer = temp; - current_size = layer_sizes[layer]; - __syncthreads(); - } - - // Use backbone output to compute CfC dynamics - if (neuron_idx < hidden_size) { - float current_x = hidden_state[batch_idx * hidden_size + neuron_idx]; - - // Simple CfC dynamics using backbone modulation - float backbone_modulation = (neuron_idx < current_size) ? current_layer[neuron_idx] : 0.0f; - float target = tanhf(backbone_modulation + output_weights[neuron_idx] * current_x); - - // Simple integration step - float dx_dt = target - current_x; - float new_x = current_x + dt * dx_dt; - - new_hidden_state[batch_idx * hidden_size + neuron_idx] = new_x; - } -} - -/** - * Fused kernel for liquid network output layer computation - * - * @param hidden_states Hidden states from all layers [batch_size, total_hidden_size] - * @param output_weights Output layer weights [output_size, total_hidden_size] - * @param output_bias Output bias [output_size] - * @param outputs Final outputs [batch_size, output_size] - * @param activation_type Output activation type - * @param batch_size Number of samples - * @param total_hidden_size Total hidden dimension - * @param output_size Output dimension - */ -__global__ void fused_liquid_output( - const float* hidden_states, - const float* output_weights, - const float* output_bias, - float* outputs, - int activation_type, - int batch_size, - int total_hidden_size, - int output_size -) { - int batch_idx = blockIdx.y * blockDim.y + threadIdx.y; - int output_idx = blockIdx.x * blockDim.x + threadIdx.x; - - if (batch_idx >= batch_size || output_idx >= output_size) return; - - float sum = output_bias[output_idx]; - - // Compute weighted sum - for (int i = 0; i < total_hidden_size; i++) { - sum += output_weights[output_idx * total_hidden_size + i] * - hidden_states[batch_idx * total_hidden_size + i]; - } - - // Apply output activation - outputs[batch_idx * output_size + output_idx] = apply_activation(sum, activation_type); -} - -/** - * Kernel for market regime adaptation - * - * Updates time constants and other parameters based on market volatility - * - * @param time_constants Time constants to update [hidden_size] - * @param base_time_constants Base time constants [hidden_size] - * @param volatility Current market volatility - * @param tau_min Minimum allowed time constant - * @param tau_max Maximum allowed time constant - * @param hidden_size Number of neurons - */ -__global__ void adapt_time_constants( - float* time_constants, - const float* base_time_constants, - float volatility, - float tau_min, - float tau_max, - int hidden_size -) { - int idx = blockIdx.x * blockDim.x + threadIdx.x; - - if (idx >= hidden_size) return; - - float base_tau = base_time_constants[idx]; - - // Adaptive time constant based on volatility - // High volatility -> faster adaptation (smaller tau) - // Low volatility -> slower adaptation (larger tau) - float adaptation_factor = 1.0f / (1.0f + volatility); - float adapted_tau = base_tau * adaptation_factor; - - // Clamp to valid range - adapted_tau = fmaxf(tau_min, fminf(tau_max, adapted_tau)); - - time_constants[idx] = adapted_tau; -} - -// Host function declarations for Rust FFI -extern "C" { - void launch_fused_ltc_forward( - const float* input, - const float* hidden_state, - const float* input_weights, - const float* recurrent_weights, - const float* bias, - float* time_constants, - float* new_hidden_state, - float dt, - float volatility, - int activation_type, - int batch_size, - int input_size, - int hidden_size, - cudaStream_t stream - ); - - void launch_fused_cfc_forward( - const float* input, - const float* hidden_state, - const float* backbone_weights, - const float* backbone_bias, - const int* layer_sizes, - const float* output_weights, - float* new_hidden_state, - float dt, - int batch_size, - int input_size, - int hidden_size, - int num_layers, - int max_layer_size, - cudaStream_t stream - ); - - void launch_fused_liquid_output( - const float* hidden_states, - const float* output_weights, - const float* output_bias, - float* outputs, - int activation_type, - int batch_size, - int total_hidden_size, - int output_size, - cudaStream_t stream - ); - - void launch_adapt_time_constants( - float* time_constants, - const float* base_time_constants, - float volatility, - float tau_min, - float tau_max, - int hidden_size, - cudaStream_t stream - ); -} - -/** - * Host function implementations - */ - -void launch_fused_ltc_forward( - const float* input, - const float* hidden_state, - const float* input_weights, - const float* recurrent_weights, - const float* bias, - float* time_constants, - float* new_hidden_state, - float dt, - float volatility, - int activation_type, - int batch_size, - int input_size, - int hidden_size, - cudaStream_t stream -) { - dim3 block_size(16, 16); - dim3 grid_size( - (hidden_size + block_size.x - 1) / block_size.x, - (batch_size + block_size.y - 1) / block_size.y - ); - - size_t shared_mem_size = (input_size + hidden_size) * sizeof(float); - - fused_ltc_forward<<>>( - input, hidden_state, input_weights, recurrent_weights, bias, - time_constants, new_hidden_state, dt, volatility, activation_type, - batch_size, input_size, hidden_size - ); -} - -void launch_fused_cfc_forward( - const float* input, - const float* hidden_state, - const float* backbone_weights, - const float* backbone_bias, - const int* layer_sizes, - const float* output_weights, - float* new_hidden_state, - float dt, - int batch_size, - int input_size, - int hidden_size, - int num_layers, - int max_layer_size, - cudaStream_t stream -) { - dim3 block_size(16, 16); - dim3 grid_size( - (hidden_size + block_size.x - 1) / block_size.x, - (batch_size + block_size.y - 1) / block_size.y - ); - - size_t shared_mem_size = 2 * max_layer_size * sizeof(float); - - fused_cfc_forward<<>>( - input, hidden_state, backbone_weights, backbone_bias, layer_sizes, - output_weights, new_hidden_state, dt, batch_size, input_size, - hidden_size, num_layers, max_layer_size - ); -} - -void launch_fused_liquid_output( - const float* hidden_states, - const float* output_weights, - const float* output_bias, - float* outputs, - int activation_type, - int batch_size, - int total_hidden_size, - int output_size, - cudaStream_t stream -) { - dim3 block_size(16, 16); - dim3 grid_size( - (output_size + block_size.x - 1) / block_size.x, - (batch_size + block_size.y - 1) / block_size.y - ); - - fused_liquid_output<<>>( - hidden_states, output_weights, output_bias, outputs, - activation_type, batch_size, total_hidden_size, output_size - ); -} - -void launch_adapt_time_constants( - float* time_constants, - const float* base_time_constants, - float volatility, - float tau_min, - float tau_max, - int hidden_size, - cudaStream_t stream -) { - const int block_size = 256; - const int grid_size = (hidden_size + block_size - 1) / block_size; - - adapt_time_constants<<>>( - time_constants, base_time_constants, volatility, - tau_min, tau_max, hidden_size - ); -} \ No newline at end of file diff --git a/crates/ml/src/liquid/cuda/memory.rs b/crates/ml/src/liquid/cuda/memory.rs deleted file mode 100644 index 2caed25bf..000000000 --- a/crates/ml/src/liquid/cuda/memory.rs +++ /dev/null @@ -1,319 +0,0 @@ -//! GPU Memory Management for Liquid Networks -//! -//! Efficient memory allocation and management for CUDA-accelerated Liquid Networks, -//! optimized for minimal allocation overhead in high-frequency trading scenarios. - -use std::collections::HashMap; -use std::sync::{Arc, Mutex}; - -use cudarc::driver::{CudaDevice, CudaSlice, DevicePtr}; - -use super::{LiquidError, Result}; - -/// `GPU` memory pool for efficient allocation/deallocation -#[derive(Debug)] -pub struct GpuMemoryPool { - device: Arc, - free_blocks: HashMap>>, - allocated_blocks: HashMap, usize>, - total_allocated: usize, - max_pool_size: usize, -} - -impl GpuMemoryPool { - /// Create a new `GPU` memory pool - pub fn new(device: Arc, max_pool_size: usize) -> Result { - Ok(Self { - device, - free_blocks: HashMap::new(), - allocated_blocks: HashMap::new(), - total_allocated: 0, - max_pool_size, - }) - } - - /// Allocate memory from the pool - pub fn allocate(&mut self, size: usize) -> Result> { - // Round up to nearest power of 2 for better reuse - let aligned_size = size.next_power_of_two(); - - // Try to reuse existing block - if let Some(blocks) = self.free_blocks.get_mut(&aligned_size) { - if let Some(ptr) = blocks.pop() { - self.allocated_blocks.insert(ptr, aligned_size); - return Ok(ptr); - } - } - - // Check if we have room for new allocation - if self.total_allocated + aligned_size > self.max_pool_size { - return Err(LiquidError::InferenceError( - "GPU memory pool exhausted".to_owned(), - )); - } - - // Allocate new block - let ptr = self.device.alloc_zeros::(aligned_size) - .map_err(|e| LiquidError::InferenceError(format!("GPU allocation failed: {}", e)))? - .device_ptr(); - - self.allocated_blocks.insert(ptr, aligned_size); - self.total_allocated += aligned_size; - - Ok(ptr) - } - - /// Deallocate memory back to the pool - pub fn deallocate(&mut self, ptr: DevicePtr) -> Result<()> { - if let Some(size) = self.allocated_blocks.remove(&ptr) { - self.free_blocks.entry(size).or_insert_with(Vec::new).push(ptr); - Ok(()) - } else { - Err(LiquidError::InferenceError( - "Attempted to deallocate untracked pointer".to_owned(), - )) - } - } - - /// Get memory statistics - pub fn get_stats(&self) -> MemoryStats { - let free_memory = self.free_blocks.values() - .map(|blocks| blocks.len()) - .sum::(); - let allocated_memory = self.allocated_blocks.len(); - - MemoryStats { - total_allocated_bytes: self.total_allocated, - free_blocks: free_memory, - allocated_blocks: allocated_memory, - max_pool_size_bytes: self.max_pool_size, - fragmentation_ratio: if self.total_allocated > 0 { - (self.total_allocated - allocated_memory) as f64 / self.total_allocated as f64 - } else { - 0.0 - }, - } - } - - /// Clear all free blocks to reclaim memory - pub fn compact(&mut self) -> Result<()> { - for blocks in self.free_blocks.values() { - for &ptr in blocks { - // In a real implementation, we would free the GPU memory here - // For now, we just track it - } - } - - let freed_bytes: usize = self.free_blocks.iter() - .map(|(&size, blocks)| size * blocks.len()) - .sum(); - - self.free_blocks.clear(); - self.total_allocated -= freed_bytes; - - Ok(()) - } -} - -/// Memory statistics for monitoring -#[derive(Debug, Clone)] -pub struct MemoryStats { - pub total_allocated_bytes: usize, - pub free_blocks: usize, - pub allocated_blocks: usize, - pub max_pool_size_bytes: usize, - pub fragmentation_ratio: f64, -} - -/// Thread-safe `GPU` memory manager -#[derive(Debug)] -pub struct GpuMemoryManager { - pool: Arc>, - device: Arc, -} - -impl GpuMemoryManager { - /// Create a new `GPU` memory manager - pub fn new(device: Arc, max_pool_size: usize) -> Result { - let pool = GpuMemoryPool::new(device.clone(), max_pool_size)?; - - Ok(Self { - pool: Arc::new(Mutex::new(pool)), - device, - }) - } - - /// Allocate typed memory slice - pub fn allocate_slice(&self, count: usize) -> Result> - where - T: Clone + Default + cudarc::driver::DeviceRepr, - { - let size_bytes = count * std::mem::size_of::(); - - // For simplicity, use device allocation directly - // In production, would use the memory pool - self.device.alloc_zeros::(count) - .map_err(|e| LiquidError::InferenceError(format!("Failed to allocate GPU memory: {}", e))) - } - - /// Get memory statistics - pub fn get_stats(&self) -> Result { - let pool = self.pool.lock() - .map_err(|_| LiquidError::InferenceError("Failed to lock memory pool".to_owned()))?; - Ok(pool.get_stats()) - } - - /// Compact memory pool - pub fn compact(&self) -> Result<()> { - let mut pool = self.pool.lock() - .map_err(|_| LiquidError::InferenceError("Failed to lock memory pool".to_owned()))?; - pool.compact() - } - - /// Get device info - pub fn device_info(&self) -> DeviceInfo { - // In a real implementation, would query actual device properties - DeviceInfo { - name: "CUDA Device".to_owned(), - total_memory_mb: 8192, // 8GB default - free_memory_mb: 4096, // 4GB default - compute_capability: (7, 5), // Turing architecture - max_threads_per_block: 1024, - max_blocks_per_grid: 65535, - warp_size: 32, - } - } -} - -/// `CUDA` device information -#[derive(Debug, Clone)] -pub struct DeviceInfo { - pub name: String, - pub total_memory_mb: usize, - pub free_memory_mb: usize, - pub compute_capability: (u32, u32), - pub max_threads_per_block: u32, - pub max_blocks_per_grid: u32, - pub warp_size: u32, -} - -/// Specialized allocator for liquid network tensors -#[derive(Debug)] -pub struct LiquidTensorAllocator { - memory_manager: Arc, - preallocated_buffers: HashMap>>>>, -} - -impl LiquidTensorAllocator { - /// Create a new tensor allocator - pub fn new(memory_manager: Arc) -> Self { - Self { - memory_manager, - preallocated_buffers: HashMap::new(), - } - } - - /// Preallocate buffers for common tensor sizes - pub fn preallocate_buffers(&mut self, common_sizes: &[(String, usize, usize)]) -> Result<()> { - for (name, size, count) in common_sizes { - let mut buffers = Vec::new(); - - for _ in 0..*count { - let buffer = self.memory_manager.allocate_slice::(*size)?; - buffers.push(buffer); - } - - self.preallocated_buffers.insert( - name.clone(), - Arc::new(Mutex::new(buffers)), - ); - } - - Ok(()) - } - - /// Get a preallocated buffer - pub fn get_buffer(&self, name: &str) -> Result>> { - if let Some(buffers) = self.preallocated_buffers.get(name) { - let mut buffers = buffers.lock() - .map_err(|_| LiquidError::InferenceError("Failed to lock buffer pool".to_owned()))?; - - Ok(buffers.pop()) - } else { - Ok(None) - } - } - - /// Return a buffer to the pool - pub fn return_buffer(&self, name: &str, buffer: CudaSlice) -> Result<()> { - if let Some(buffers) = self.preallocated_buffers.get(name) { - let mut buffers = buffers.lock() - .map_err(|_| LiquidError::InferenceError("Failed to lock buffer pool".to_owned()))?; - - buffers.push(buffer); - } - - Ok(()) - } - - /// Allocate a new tensor - pub fn allocate_tensor(&self, size: usize) -> Result> { - self.memory_manager.allocate_slice::(size) - } -} - -impl Clone for GpuMemoryManager { - fn clone(&self) -> Self { - Self { - pool: self.pool.clone(), - device: self.device.clone(), - } - } -} - -#[cfg(test)] -mod tests { - use super::*; - - #[test] - fn test_memory_pool_creation() { - // This test would require actual CUDA device - // For now, just test the structure - let max_size = 1024 * 1024; // 1MB - - // Test would create device and pool - assert!(max_size > 0); - } - - #[test] - fn test_memory_stats() { - let stats = MemoryStats { - total_allocated_bytes: 1024, - free_blocks: 2, - allocated_blocks: 3, - max_pool_size_bytes: 2048, - fragmentation_ratio: 0.1, - }; - - assert_eq!(stats.total_allocated_bytes, 1024); - assert_eq!(stats.free_blocks, 2); - assert!(stats.fragmentation_ratio < 1.0); - } - - #[test] - fn test_device_info() { - let info = DeviceInfo { - name: "Test GPU".to_owned(), - total_memory_mb: 8192, - free_memory_mb: 4096, - compute_capability: (7, 5), - max_threads_per_block: 1024, - max_blocks_per_grid: 65535, - warp_size: 32, - }; - - assert_eq!(info.name, "Test GPU"); - assert_eq!(info.warp_size, 32); - assert!(info.total_memory_mb > info.free_memory_mb); - } -} \ No newline at end of file diff --git a/crates/ml/src/liquid/cuda/mod.rs b/crates/ml/src/liquid/cuda/mod.rs deleted file mode 100644 index 568ac163b..000000000 --- a/crates/ml/src/liquid/cuda/mod.rs +++ /dev/null @@ -1,640 +0,0 @@ -#![allow(unsafe_code)] // Intentional unsafe for CUDA operations - -//! CUDA-accelerated Liquid Neural Networks -//! -//! GPU implementation of Liquid Time-constant (LTC) and Closed-form Continuous-time (CfC) -//! neural networks with optimized CUDA kernels for ultra-low latency inference. - -use std::collections::HashMap; -use std::ffi::c_void; -use std::ptr; -use std::sync::Arc; - -use cudarc::driver::{CudaDevice, CudaSlice, DevicePtr, LaunchAsync, LaunchConfig}; -use cudarc::nvrtc::Ptx; -use serde::{Deserialize, Serialize}; - -use super::{FixedPoint, LiquidError, MarketRegime, NetworkType, PerformanceMetrics, Result}; -use crate::{MLError, MLResult}; - -pub mod bindings; -pub mod memory; -pub mod stream_manager; - -// DO NOT RE-EXPORT - Use explicit imports at usage sites - -/// `CUDA`-accelerated Liquid Network configuration -#[derive(Debug, Clone, Serialize, Deserialize)] -pub struct CudaLiquidConfig { - pub device_id: usize, - pub max_batch_size: usize, - pub use_shared_memory: bool, - pub stream_count: usize, - pub memory_pool_size_mb: usize, - pub enable_profiling: bool, -} - -impl Default for CudaLiquidConfig { - fn default() -> Self { - let caps = crate::gpu::capabilities::cached_capabilities(); - // Scale batch size to GPU VRAM: - // ≤4GB: 256, 5-12GB: 512, 13-25GB: 1024, 26+GB: 2048 - let max_batch_size = match caps.free_vram_mb as u64 { - 0..=4096 => 256, - 4097..=12288 => 512, - 12289..=25600 => 1024, - _ => 2048, - }; - // Scale memory pool to GPU VRAM (10% of free VRAM, min 256MB) - let memory_pool_size_mb = ((caps.free_vram_mb * 0.10) as usize).max(256); - Self { - device_id: 0, - max_batch_size, - use_shared_memory: true, - stream_count: 4, - memory_pool_size_mb, - enable_profiling: false, - } - } -} - -/// `GPU` memory buffers for Liquid Networks -#[derive(Debug)] -pub struct CudaBuffers { - // Input/Output buffers - pub input: CudaSlice, - pub hidden_state: CudaSlice, - pub new_hidden_state: CudaSlice, - pub output: CudaSlice, - - // Weight buffers - pub input_weights: CudaSlice, - pub recurrent_weights: CudaSlice, - pub output_weights: CudaSlice, - pub bias: CudaSlice, - pub output_bias: CudaSlice, - - // Dynamic parameters - pub time_constants: CudaSlice, - pub base_time_constants: CudaSlice, - - // CfC-specific buffers - pub backbone_weights: Option>, - pub backbone_bias: Option>, - pub layer_sizes: Option>, -} - -/// `CUDA`-accelerated Liquid Neural Network -#[derive(Debug)] -pub struct CudaLiquidNetwork { - pub config: CudaLiquidConfig, - pub device: Arc, - pub buffers: CudaBuffers, - pub stream_manager: cudarc::driver::CudaStreamManager, - pub memory_manager: GpuMemoryManager, - - // Network parameters - pub network_type: NetworkType, - pub input_size: usize, - pub hidden_size: usize, - pub output_size: usize, - pub batch_size: usize, - - // Performance tracking - pub performance_metrics: PerformanceMetrics, - pub current_regime: MarketRegime, - - // CUDA function handles - ltc_forward_fn: cudarc::driver::CudaFunction, - cfc_forward_fn: Option, - output_fn: cudarc::driver::CudaFunction, - adapt_tau_fn: cudarc::driver::CudaFunction, -} - -impl CudaLiquidNetwork { - /// Create a new `CUDA`-accelerated Liquid Network - pub fn new( - network_type: NetworkType, - input_size: usize, - hidden_size: usize, - output_size: usize, - config: CudaLiquidConfig, - ) -> Result { - // Initialize CUDA device - let device = CudaDevice::new(config.device_id) - .map_err(|e| LiquidError::InferenceError(format!("Failed to initialize CUDA device: {}", e)))?; - let device = Arc::new(device); - - // Load CUDA kernels - let ptx = compile_liquid_kernels()?; - device.load_ptx(ptx, "liquid_kernels", &[ - "fused_ltc_forward", - "fused_cfc_forward", - "fused_liquid_output", - "adapt_time_constants" - ]).map_err(|e| LiquidError::InferenceError(format!("Failed to load CUDA kernels: {}", e)))?; - - // Get kernel functions - let ltc_forward_fn = device.get_func("liquid_kernels", "fused_ltc_forward") - .map_err(|e| LiquidError::InferenceError(format!("Failed to get LTC kernel: {}", e)))?; - let cfc_forward_fn = if matches!(network_type, NetworkType::CfC | NetworkType::Mixed) { - Some(device.get_func("liquid_kernels", "fused_cfc_forward") - .map_err(|e| LiquidError::InferenceError(format!("Failed to get CfC kernel: {}", e)))?) - } else { - None - }; - let output_fn = device.get_func("liquid_kernels", "fused_liquid_output") - .map_err(|e| LiquidError::InferenceError(format!("Failed to get output kernel: {}", e)))?; - let adapt_tau_fn = device.get_func("liquid_kernels", "adapt_time_constants") - .map_err(|e| LiquidError::InferenceError(format!("Failed to get adaptation kernel: {}", e)))?; - - // Initialize memory manager - let memory_manager = GpuMemoryManager::new( - device.clone(), - config.memory_pool_size_mb * 1024 * 1024, - )?; - - // Initialize stream manager - let stream_manager = cudarc::driver::CudaStreamManager::new(device.clone(), config.stream_count)?; - - // Allocate GPU buffers - let batch_size = config.max_batch_size; - let buffers = Self::allocate_buffers( - &device, - &memory_manager, - batch_size, - input_size, - hidden_size, - output_size, - &network_type, - )?; - - let performance_metrics = PerformanceMetrics { - total_inferences: 0, - average_inference_time_ns: 0, - average_inference_time_us: 0.0, - total_parameters: Self::calculate_parameter_count(input_size, hidden_size, output_size), - current_regime: MarketRegime::Normal, - regime_switches: 0, - last_adaptation_time: None, - }; - - Ok(Self { - config, - device, - buffers, - stream_manager, - memory_manager, - network_type, - input_size, - hidden_size, - output_size, - batch_size, - performance_metrics, - current_regime: MarketRegime::Normal, - ltc_forward_fn, - cfc_forward_fn, - output_fn, - adapt_tau_fn, - }) - } - - /// Allocate `GPU` memory buffers - fn allocate_buffers( - device: &CudaDevice, - memory_manager: &GpuMemoryManager, - batch_size: usize, - input_size: usize, - hidden_size: usize, - output_size: usize, - network_type: &NetworkType, - ) -> Result { - // Input/Output buffers - let input = device.alloc_zeros::(batch_size * input_size) - .map_err(|e| LiquidError::InferenceError(format!("Failed to allocate input buffer: {}", e)))?; - let hidden_state = device.alloc_zeros::(batch_size * hidden_size) - .map_err(|e| LiquidError::InferenceError(format!("Failed to allocate hidden state buffer: {}", e)))?; - let new_hidden_state = device.alloc_zeros::(batch_size * hidden_size) - .map_err(|e| LiquidError::InferenceError(format!("Failed to allocate new hidden state buffer: {}", e)))?; - let output = device.alloc_zeros::(batch_size * output_size) - .map_err(|e| LiquidError::InferenceError(format!("Failed to allocate output buffer: {}", e)))?; - - // Weight buffers - let input_weights = device.alloc_zeros::(hidden_size * input_size) - .map_err(|e| LiquidError::InferenceError(format!("Failed to allocate input weights: {}", e)))?; - let recurrent_weights = device.alloc_zeros::(hidden_size * hidden_size) - .map_err(|e| LiquidError::InferenceError(format!("Failed to allocate recurrent weights: {}", e)))?; - let output_weights = device.alloc_zeros::(output_size * hidden_size) - .map_err(|e| LiquidError::InferenceError(format!("Failed to allocate output weights: {}", e)))?; - let bias = device.alloc_zeros::(hidden_size) - .map_err(|e| LiquidError::InferenceError(format!("Failed to allocate bias: {}", e)))?; - let output_bias = device.alloc_zeros::(output_size) - .map_err(|e| LiquidError::InferenceError(format!("Failed to allocate output bias: {}", e)))?; - - // Dynamic parameters - let time_constants = device.alloc_zeros::(hidden_size) - .map_err(|e| LiquidError::InferenceError(format!("Failed to allocate time constants: {}", e)))?; - let base_time_constants = device.alloc_zeros::(hidden_size) - .map_err(|e| LiquidError::InferenceError(format!("Failed to allocate base time constants: {}", e)))?; - - // CfC-specific buffers - let (backbone_weights, backbone_bias, layer_sizes) = match network_type { - NetworkType::CfC | NetworkType::Mixed => { - // For now, allocate simple backbone (2 layers of size hidden_size each) - let backbone_layers = 2; - let max_layer_size = hidden_size; - let total_backbone_weights = backbone_layers * max_layer_size * (input_size + hidden_size); - - let backbone_weights = device.alloc_zeros::(total_backbone_weights) - .map_err(|e| LiquidError::InferenceError(format!("Failed to allocate backbone weights: {}", e)))?; - let backbone_bias = device.alloc_zeros::(backbone_layers * max_layer_size) - .map_err(|e| LiquidError::InferenceError(format!("Failed to allocate backbone bias: {}", e)))?; - let layer_sizes = device.alloc_zeros::(backbone_layers) - .map_err(|e| LiquidError::InferenceError(format!("Failed to allocate layer sizes: {}", e)))?; - - (Some(backbone_weights), Some(backbone_bias), Some(layer_sizes)) - } - _ => (None, None, None), - }; - - Ok(CudaBuffers { - input, - hidden_state, - new_hidden_state, - output, - input_weights, - recurrent_weights, - output_weights, - bias, - output_bias, - time_constants, - base_time_constants, - backbone_weights, - backbone_bias, - layer_sizes, - }) - } - - /// Forward pass through the `CUDA`-accelerated network - pub fn forward_gpu(&mut self, input: &[f32], batch_size: usize) -> Result> { - if batch_size > self.config.max_batch_size { - return Err(LiquidError::InvalidInput(format!( - "Batch size {} exceeds maximum {}", - batch_size, self.config.max_batch_size - ))); - } - - if input.len() != batch_size * self.input_size { - return Err(LiquidError::InvalidInput(format!( - "Input size mismatch: expected {}, got {}", - batch_size * self.input_size, - input.len() - ))); - } - - let start_time = std::time::Instant::now(); - - // Get a stream for this operation - let stream = self.stream_manager.get_stream()?; - - // Copy input to GPU - self.device.htod_sync_copy_into(input, &mut self.buffers.input) - .map_err(|e| LiquidError::InferenceError(format!("Failed to copy input to GPU: {}", e)))?; - - // Launch appropriate forward kernel based on network type - match self.network_type { - NetworkType::LTC => { - self.launch_ltc_forward(batch_size, stream)?; - } - NetworkType::CfC => { - self.launch_cfc_forward(batch_size, stream)?; - } - NetworkType::Mixed => { - // Run both LTC and CfC in parallel on different parts of hidden state - self.launch_ltc_forward(batch_size, stream)?; - // Wait for LTC to complete, then run CfC - self.device.synchronize() - .map_err(|e| LiquidError::InferenceError(format!("CUDA sync failed: {}", e)))?; - self.launch_cfc_forward(batch_size, stream)?; - } - } - - // Launch output layer kernel - self.launch_output_layer(batch_size, stream)?; - - // Copy result back to CPU - let mut output = vec![0.0_f32; batch_size * self.output_size]; - self.device.dtoh_sync_copy_into(&self.buffers.output, &mut output) - .map_err(|e| LiquidError::InferenceError(format!("Failed to copy output from GPU: {}", e)))?; - - // Update performance metrics - let elapsed = start_time.elapsed(); - self.performance_metrics.total_inferences += 1; - let total_time_ns = self.performance_metrics.average_inference_time_ns - * (self.performance_metrics.total_inferences - 1) - + elapsed.as_nanos() as u64; - self.performance_metrics.average_inference_time_ns = - total_time_ns / self.performance_metrics.total_inferences; - self.performance_metrics.average_inference_time_us = - self.performance_metrics.average_inference_time_ns as f64 / 1000.0; - - Ok(output) - } - - /// Launch LTC forward kernel - fn launch_ltc_forward(&self, batch_size: usize, stream: &cudarc::driver::CudaStream) -> Result<()> { - let grid_x = (self.hidden_size + 15) / 16; - let grid_y = (batch_size + 15) / 16; - let grid_z = 1; - - let block_x = 16; - let block_y = 16; - let block_z = 1; - - let shared_mem_size = (self.input_size + self.hidden_size) * 4; // 4 bytes per f32 - - let config = LaunchConfig { - grid_dim: (grid_x as u32, grid_y as u32, grid_z as u32), - block_dim: (block_x as u32, block_y as u32, block_z as u32), - shared_mem_bytes: shared_mem_size as u32, - }; - - let params = ( - &self.buffers.input, - &self.buffers.hidden_state, - &self.buffers.input_weights, - &self.buffers.recurrent_weights, - &self.buffers.bias, - &self.buffers.time_constants, - &self.buffers.new_hidden_state, - 0.01_f32, // dt - 0.5_f32, // volatility - 3_i32, // activation_type (Tanh) - batch_size as i32, - self.input_size as i32, - self.hidden_size as i32, - ); - - // SAFETY: CUDA kernel launch for LTC forward pass - // - Invariant 1: All device buffers allocated and sized correctly - // - Invariant 2: Grid/block dimensions calculated to cover batch_size × hidden_size - // - Invariant 3: Stream handle valid from CudaDevice::fork_default_stream - // - Verified: Buffer allocation checked in launch_ltc_forward caller - // - Risk: HIGH - GPU kernel execution, incorrect params cause device errors - // SAFETY: Unsafe operation validated - invariants maintained by surrounding code - unsafe { - self.ltc_forward_fn.launch_async(config, params, stream) - .map_err(|e| LiquidError::InferenceError(format!("LTC kernel launch failed: {}", e)))?; - } - - Ok(()) - } - - /// Launch CfC forward kernel - fn launch_cfc_forward(&self, batch_size: usize, stream: &cudarc::driver::CudaStream) -> Result<()> { - let cfc_fn = self.cfc_forward_fn.as_ref() - .ok_or_else(|| LiquidError::InferenceError("CfC kernel not available".to_owned()))?; - - let backbone_weights = self.buffers.backbone_weights.as_ref() - .ok_or_else(|| LiquidError::InferenceError("Backbone weights not allocated".to_owned()))?; - let backbone_bias = self.buffers.backbone_bias.as_ref() - .ok_or_else(|| LiquidError::InferenceError("Backbone bias not allocated".to_owned()))?; - let layer_sizes = self.buffers.layer_sizes.as_ref() - .ok_or_else(|| LiquidError::InferenceError("Layer sizes not allocated".to_owned()))?; - - let grid_x = (self.hidden_size + 15) / 16; - let grid_y = (batch_size + 15) / 16; - let grid_z = 1; - - let block_x = 16; - let block_y = 16; - let block_z = 1; - - let shared_mem_size = 2 * self.hidden_size * 4; // Two layers in shared memory - - let config = LaunchConfig { - grid_dim: (grid_x as u32, grid_y as u32, grid_z as u32), - block_dim: (block_x as u32, block_y as u32, block_z as u32), - shared_mem_bytes: shared_mem_size as u32, - }; - - let params = ( - &self.buffers.input, - &self.buffers.hidden_state, - backbone_weights, - backbone_bias, - layer_sizes, - &self.buffers.output_weights, - &self.buffers.new_hidden_state, - 0.01_f32, // dt - batch_size as i32, - self.input_size as i32, - self.hidden_size as i32, - 2_i32, // num_layers - self.hidden_size as i32, // max_layer_size - ); - - // SAFETY: CUDA kernel launch for CfC forward pass - // - Invariant 1: CfC kernel function verified present via ok_or_else - // - Invariant 2: Backbone weights/bias buffers checked for allocation - // - Invariant 3: Grid dimensions cover hidden_size with 16×16 thread blocks - // - Verified: All buffer allocation validated before kernel params creation - // - Risk: HIGH - Complex kernel with multiple buffers, size mismatches critical - // SAFETY: Unsafe operation validated - invariants maintained by surrounding code - unsafe { - cfc_fn.launch_async(config, params, stream) - .map_err(|e| LiquidError::InferenceError(format!("CfC kernel launch failed: {}", e)))?; - } - - Ok(()) - } - - /// Launch output layer kernel - fn launch_output_layer(&self, batch_size: usize, stream: &cudarc::driver::CudaStream) -> Result<()> { - let grid_x = (self.output_size + 15) / 16; - let grid_y = (batch_size + 15) / 16; - let grid_z = 1; - - let block_x = 16; - let block_y = 16; - let block_z = 1; - - let config = LaunchConfig { - grid_dim: (grid_x as u32, grid_y as u32, grid_z as u32), - block_dim: (block_x as u32, block_y as u32, block_z as u32), - shared_mem_bytes: 0, - }; - - let params = ( - &self.buffers.new_hidden_state, - &self.buffers.output_weights, - &self.buffers.output_bias, - &self.buffers.output, - 0_i32, // activation_type (Linear) - batch_size as i32, - self.hidden_size as i32, - self.output_size as i32, - ); - - // SAFETY: CUDA kernel launch for output layer - // - Invariant 1: Output buffers allocated with correct dimensions - // - Invariant 2: Grid covers batch_size × output_size matrix - // - Invariant 3: Hidden state from previous layer available - // - Verified: Output buffer size matches output_size parameter - // - Risk: MEDIUM - Final layer, errors visible in incorrect predictions - // SAFETY: Unsafe operation validated - invariants maintained by surrounding code - unsafe { - self.output_fn.launch_async(config, params, stream) - .map_err(|e| LiquidError::InferenceError(format!("Output kernel launch failed: {}", e)))?; - } - - Ok(()) - } - - /// Update market volatility and adapt network parameters - pub fn update_market_volatility_gpu(&mut self, volatility: f32) -> Result<()> { - let stream = self.stream_manager.get_stream()?; - - let grid_size = (self.hidden_size + 255) / 256; - let block_size = 256; - - let config = LaunchConfig { - grid_dim: (grid_size as u32, 1, 1), - block_dim: (block_size as u32, 1, 1), - shared_mem_bytes: 0, - }; - - let params = ( - &self.buffers.time_constants, - &self.buffers.base_time_constants, - volatility, - 0.01_f32, // tau_min - 1.0_f32, // tau_max - self.hidden_size as i32, - ); - - // SAFETY: CUDA kernel launch for time constant adaptation - // - Invariant 1: Tau parameters buffer sized for hidden_size elements - // - Invariant 2: Volatility value validated as finite positive float - // - Invariant 3: Grid dimensions match tau buffer layout - // - Verified: Volatility check ensures valid kernel input - // - Risk: MEDIUM - Adaptation kernel, invalid params affect convergence - // SAFETY: Unsafe operation validated - invariants maintained by surrounding code - unsafe { - self.adapt_tau_fn.launch_async(config, params, stream) - .map_err(|e| LiquidError::InferenceError(format!("Adaptation kernel launch failed: {}", e)))?; - } - - // Update regime based on volatility - let new_regime = if volatility < 0.2 { - MarketRegime::Normal - } else if volatility < 1.0 { - MarketRegime::Sideways - } else if volatility < 3.0 { - MarketRegime::Trending - } else if volatility < 5.0 { - MarketRegime::Bull - } else { - MarketRegime::Crisis - }; - - if new_regime != self.current_regime { - self.current_regime = new_regime; - self.performance_metrics.current_regime = new_regime; - self.performance_metrics.regime_switches += 1; - self.performance_metrics.last_adaptation_time = Some( - std::time::SystemTime::now() - .duration_since(std::time::UNIX_EPOCH) - .map_err(|e| LiquidError::InferenceError(format!("System time error: {}", e)))? - .as_millis() as u64, - ); - } - - Ok(()) - } - - /// Initialize network weights from CPU network - pub fn load_weights_from_cpu( - &mut self, - input_weights: &[f32], - recurrent_weights: &[f32], - output_weights: &[f32], - bias: &[f32], - output_bias: &[f32], - time_constants: &[f32], - ) -> Result<()> { - // Copy weights to GPU - self.device.htod_sync_copy_into(input_weights, &mut self.buffers.input_weights) - .map_err(|e| LiquidError::InferenceError(format!("Failed to copy input weights: {}", e)))?; - self.device.htod_sync_copy_into(recurrent_weights, &mut self.buffers.recurrent_weights) - .map_err(|e| LiquidError::InferenceError(format!("Failed to copy recurrent weights: {}", e)))?; - self.device.htod_sync_copy_into(output_weights, &mut self.buffers.output_weights) - .map_err(|e| LiquidError::InferenceError(format!("Failed to copy output weights: {}", e)))?; - self.device.htod_sync_copy_into(bias, &mut self.buffers.bias) - .map_err(|e| LiquidError::InferenceError(format!("Failed to copy bias: {}", e)))?; - self.device.htod_sync_copy_into(output_bias, &mut self.buffers.output_bias) - .map_err(|e| LiquidError::InferenceError(format!("Failed to copy output bias: {}", e)))?; - self.device.htod_sync_copy_into(time_constants, &mut self.buffers.time_constants) - .map_err(|e| LiquidError::InferenceError(format!("Failed to copy time constants: {}", e)))?; - self.device.htod_sync_copy_into(time_constants, &mut self.buffers.base_time_constants) - .map_err(|e| LiquidError::InferenceError(format!("Failed to copy base time constants: {}", e)))?; - - Ok(()) - } - - /// Get performance metrics - pub fn get_performance_metrics(&self) -> &PerformanceMetrics { - &self.performance_metrics - } - - /// Calculate total parameter count - fn calculate_parameter_count(input_size: usize, hidden_size: usize, output_size: usize) -> usize { - let input_params = hidden_size * input_size; - let recurrent_params = hidden_size * hidden_size; - let output_params = output_size * hidden_size; - let bias_params = hidden_size + output_size; - let tau_params = hidden_size; - - input_params + recurrent_params + output_params + bias_params + tau_params - } -} - -/// Compile `CUDA` kernels from source -fn compile_liquid_kernels() -> Result { - let kernel_src = include_str!("liquid_kernels.cu"); - cudarc::nvrtc::compile_ptx(kernel_src).map_err(|e| { - LiquidError::InferenceError(format!( - "CUDA kernel compilation failed. Ensure NVCC is installed: {}", - e - )) - }) -} - -// TECHNICAL DEBT ELIMINATED - Use cudarc types directly - -#[cfg(test)] -mod tests { - use super::*; - - #[test] - fn test_cuda_liquid_config_default_scales_to_vram() { - let config = CudaLiquidConfig::default(); - // Should be at least 256 on any system (was hardcoded 32) - assert!( - config.max_batch_size >= 256, - "max_batch_size should be at least 256, got {}", - config.max_batch_size - ); - } - - #[test] - fn test_cuda_liquid_config_memory_pool_scales() { - let config = CudaLiquidConfig::default(); - // Memory pool should be at least 256MB (minimum floor) - assert!( - config.memory_pool_size_mb >= 256, - "memory_pool_size_mb should be at least 256, got {}", - config.memory_pool_size_mb - ); - } -} \ No newline at end of file diff --git a/crates/ml/src/liquid/mod.rs b/crates/ml/src/liquid/mod.rs index 15e6bd0e3..fdf42d29a 100644 --- a/crates/ml/src/liquid/mod.rs +++ b/crates/ml/src/liquid/mod.rs @@ -1,212 +1,13 @@ //! Liquid Neural Networks for Ultra-Low Latency HFT //! -//! Implementation of Liquid Time-constant (LTC) and Closed-form Continuous-time (CfC) -//! neural networks with fixed-point arithmetic for sub-100μs inference. +//! This module re-exports the `ml-supervised` Liquid implementation and keeps bridge +//! modules that depend on types from both `ml-supervised` and the parent `ml` crate. -use std::error::Error; -use std::fmt; +// Re-export everything from the ml-supervised liquid module +pub use ml_supervised::liquid::*; -use serde::{Deserialize, Serialize}; - -// Import MarketRegime from core types to avoid type conflicts -use crate::MLError; -use common::trading::MarketRegime; - -pub mod activation; +// Bridge modules that depend on ml-internal types (UnifiedTrainable) pub mod adapter; -pub mod candle_cfc; -pub mod cells; -pub mod network; -pub mod ode_solvers; -pub mod training; -#[cfg(test)] -mod tests; - -// Re-export main types for external usage -pub use activation::ActivationType; +// Re-export bridge types pub use adapter::LiquidTrainableAdapter; -pub use candle_cfc::{BackboneMLP, CandleCfCNetwork, CfCCell, CfCTrainConfig, DeviceConfig}; -pub use cells::{CfCConfig, LTCConfig}; -pub use network::{LayerConfig, LiquidNetwork, LiquidNetworkConfig, OutputLayerConfig}; -pub use ode_solvers::SolverType; -pub use training::{ - CandleCfCTrainer, CfCTrainerConfig, LiquidTrainer, LiquidTrainingConfig, TrainingBatch, - TrainingMetrics, TrainingSample, TrainingUtils, -}; - -/// Fixed-point arithmetic for ultra-low latency inference -pub const PRECISION: i64 = 100_000_000; // 8 decimal places - -#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize, Default)] -pub struct FixedPoint(pub i64); - -impl FixedPoint { - pub fn from_f64(value: f64) -> Self { - FixedPoint((value * PRECISION as f64) as i64) - } - - pub fn to_f64(self) -> f64 { - self.0 as f64 / PRECISION as f64 - } - - pub fn zero() -> Self { - FixedPoint(0) - } - - pub fn one() -> Self { - FixedPoint(PRECISION) - } - - pub fn is_finite(&self) -> bool { - self.0.abs() < i64::MAX / 2 - } -} - -impl std::ops::Add for FixedPoint { - type Output = Result; - - fn add(self, rhs: FixedPoint) -> Self::Output { - self.0 - .checked_add(rhs.0) - .map(FixedPoint) - .ok_or(LiquidError::Overflow("Addition overflow".to_owned())) - } -} - -impl std::ops::Sub for FixedPoint { - type Output = Result; - - fn sub(self, rhs: FixedPoint) -> Self::Output { - self.0 - .checked_sub(rhs.0) - .map(FixedPoint) - .ok_or(LiquidError::Overflow("Subtraction overflow".to_owned())) - } -} - -impl std::ops::Mul for FixedPoint { - type Output = Result; - - fn mul(self, rhs: FixedPoint) -> Self::Output { - let result = ((self.0 as i128) * (rhs.0 as i128)) / (PRECISION as i128); - if result > i64::MAX as i128 || result < i64::MIN as i128 { - Err(LiquidError::Overflow("Multiplication overflow".to_owned())) - } else { - Ok(FixedPoint(result as i64)) - } - } -} - -impl std::ops::Div for FixedPoint { - type Output = Result; - - fn div(self, rhs: FixedPoint) -> Self::Output { - if rhs.0 == 0 { - return Err(LiquidError::DivisionByZero); - } - let result = ((self.0 as i128) * (PRECISION as i128)) / (rhs.0 as i128); - if result > i64::MAX as i128 || result < i64::MIN as i128 { - Err(LiquidError::Overflow("Division overflow".to_owned())) - } else { - Ok(FixedPoint(result as i64)) - } - } -} - -/// Liquid Neural Network specific errors -#[derive(Debug, Clone, Serialize, Deserialize)] -pub enum LiquidError { - InvalidConfiguration(String), - InvalidInput(String), - Overflow(String), - DivisionByZero, - InferenceError(String), - TrainingError(String), - SolverError(String), -} - -impl fmt::Display for LiquidError { - fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { - match self { - LiquidError::InvalidConfiguration(msg) => write!(f, "Invalid configuration: {}", msg), - LiquidError::InvalidInput(msg) => write!(f, "Invalid input: {}", msg), - LiquidError::Overflow(msg) => write!(f, "Overflow error: {}", msg), - LiquidError::DivisionByZero => write!(f, "Division by zero"), - LiquidError::InferenceError(msg) => write!(f, "Inference error: {}", msg), - LiquidError::TrainingError(msg) => write!(f, "Training error: {}", msg), - LiquidError::SolverError(msg) => write!(f, "ODE solver error: {}", msg), - } - } -} - -impl Error for LiquidError {} - -impl From for MLError { - fn from(err: LiquidError) -> Self { - match err { - LiquidError::InvalidConfiguration(msg) => MLError::ConfigError(msg), - LiquidError::InvalidInput(msg) => MLError::InvalidInput(msg), - LiquidError::InferenceError(msg) => MLError::InferenceError(msg), - LiquidError::TrainingError(msg) => MLError::TrainingError(msg), - LiquidError::Overflow(_) - | LiquidError::DivisionByZero - | LiquidError::SolverError(_) => MLError::ModelError(err.to_string()), - } - } -} - -impl From for LiquidError { - fn from(err: MLError) -> Self { - match err { - MLError::ConfigError(msg) => { - LiquidError::InvalidConfiguration(msg) - }, - MLError::InvalidInput(msg) => LiquidError::InvalidInput(msg), - MLError::InferenceError(msg) => LiquidError::InferenceError(msg), - MLError::TrainingError(msg) => LiquidError::TrainingError(msg), - MLError::DimensionMismatch { .. } - | MLError::GraphError { .. } - | MLError::ResourceLimit { .. } - | MLError::SerializationError { .. } - | MLError::ValidationError { .. } - | MLError::ConcurrencyError { .. } - | MLError::InitializationError { .. } - | MLError::ModelError(_) - | MLError::NotTrained(_) - | MLError::AnyhowError(_) - | MLError::TensorCreationError { .. } - | MLError::TensorOperationError(_) - | MLError::LockError(_) - | MLError::ModelNotFound(_) - | MLError::InsufficientData(_) - | MLError::CheckpointError(_) - | MLError::DeviceError(_) => LiquidError::InferenceError(err.to_string()), - } - } -} - -pub type Result = std::result::Result; - -/// Network type for liquid neural networks -#[derive(Debug, Clone, Serialize, Deserialize)] -pub enum NetworkType { - LTC, // Liquid Time-constant - CfC, // Closed-form Continuous-time - Mixed, // Combination of LTC and CfC layers -} - -// REMOVED: MarketRegime enum - now using common::MarketRegime instead -// This eliminates the type conflict and ensures consistency across the entire system - -/// Performance metrics for liquid networks -#[derive(Debug, Clone, Serialize, Deserialize)] -pub struct PerformanceMetrics { - pub total_inferences: u64, - pub average_inference_time_ns: u64, - pub average_inference_time_us: f64, - pub total_parameters: usize, - pub current_regime: MarketRegime, // Now uses core MarketRegime enum - pub regime_switches: u32, - pub last_adaptation_time: Option, // Store as timestamp millis instead of Instant -} diff --git a/crates/ml/src/mamba/mod.rs b/crates/ml/src/mamba/mod.rs index 0cad442e0..0dbfd7c36 100644 --- a/crates/ml/src/mamba/mod.rs +++ b/crates/ml/src/mamba/mod.rs @@ -1,3194 +1,10 @@ -//! # Mamba-2 State-Space Model for HFT +//! Mamba-2 State-Space Model for HFT //! -//! Next-generation Mamba-2 implementation with Structured State Duality (SSD), -//! hardware-aware algorithms, and 5x performance improvements over Mamba-1. -//! -//! **Note**: This module uses mathematical notation (A, B, C for state-space matrices). -//! Non-snake-case warnings are allowed for mathematical clarity. +//! This module re-exports the `ml-supervised` Mamba-2 implementation and keeps bridge +//! modules that depend on types from both `ml-supervised` and the parent `ml` crate. -#![allow(non_snake_case)] -//! -//! ## Key Features -//! -//! - **SSD Layers**: Structured State Duality for linear attention mechanisms -//! - **Hardware-aware**: Optimized memory access patterns and SIMD instructions -//! - **5x Faster**: Sub-linear memory usage and linear-time sequence modeling -//! - **Selective State Spaces**: Advanced state selection mechanisms -//! - **Sub-5μs**: Target inference latency for HFT applications -//! - **Integer Precision**: 10,000x scaling for financial precision -//! -//! ## Architecture Improvements -//! -//! ```text -//! ┌─────────────────────────────────────────────────────────────┐ -//! │ Mamba-2 Block │ -//! ├─────────────────┬─────────────────┬─────────────────────────┤ -//! │ SSD Layer │ Hardware-Aware │ Selective State │ -//! │ │ Optimization │ Mechanism │ -//! │ • Linear Attn │ • SIMD Vectors │ • Advanced Selection │ -//! │ • Structured │ • Cache-Friendly│ • Dynamic Parameters │ -//! │ Duality │ • Prefetching │ • State Compression │ -//! └─────────────────┴─────────────────┴─────────────────────────┘ -//! ``` -//! -//! ## Performance Targets -//! -//! - Inference: <5μs per sequence step (5x faster than Mamba-1) -//! - Memory: Sub-linear growth with sequence length -//! - Throughput: >1M sequences/sec -//! - Latency: 99.9% percentile <10μs +// Re-export everything from the ml-supervised mamba module +pub use ml_supervised::mamba::*; -mod hardware_aware; -mod scan_algorithms; -pub mod selective_state; -mod ssd_layer; -pub mod loss; +// Bridge modules that depend on ml-internal types (UnifiedTrainable) pub mod trainable_adapter; - -// Public exports for types used in mod.rs and by external crates -pub use hardware_aware::{HardwareCapabilities, HardwareOptimizer}; -pub use scan_algorithms::{ParallelScanEngine, ScanBenchmark, ScanOperator}; -pub use selective_state::{ - SelectiveStateConfig, SelectiveStateSpace, StateCompressor, StateImportance, -}; -pub use ssd_layer::SSDLayer; - -use std::collections::{HashMap, VecDeque}; -use std::sync::atomic::{AtomicU64, Ordering}; -use std::sync::Arc; -use std::time::{Duration, Instant, SystemTime}; - -use candle_core::{DType, Device, Tensor, Var}; -use candle_nn::Module; -use candle_nn::{Dropout, Linear, VarBuilder}; -use serde::{Deserialize, Serialize}; -use tracing::{debug, info, instrument, trace, warn}; -use uuid::Uuid; - -use crate::cuda_compat::layer_norm_with_fallback; -use crate::dqn::mixed_precision::training_dtype; -use crate::MLError; - -/// Optimizer type for training -#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] -pub enum OptimizerType { - /// Adam optimizer with adaptive learning rates (coupled weight decay) - Adam, - /// AdamW optimizer with decoupled weight decay (recommended for SSMs) - AdamW, - /// Stochastic Gradient Descent with momentum - SGD, -} - -impl Default for OptimizerType { - fn default() -> Self { - Self::AdamW // AdamW is superior for state-space models - } -} - -/// Configuration for `MAMBA-2` state-space model -#[derive(Debug, Clone, Serialize, Deserialize)] -pub struct Mamba2Config { - /// Model dimension - pub d_model: usize, - /// State dimension - pub d_state: usize, - /// Head dimension for multi-head attention - pub d_head: usize, - /// Number of attention heads - pub num_heads: usize, - /// Expansion factor for inner dimension - pub expand: usize, - /// Number of layers - pub num_layers: usize, - /// Dropout rate - pub dropout: f64, - /// Use structured state duality - pub use_ssd: bool, - /// Use selective state mechanism - pub use_selective_state: bool, - /// Enable hardware optimizations - pub hardware_aware: bool, - /// Target latency in microseconds - pub target_latency_us: u64, - /// Maximum sequence length - pub max_seq_len: usize, - /// Learning rate - pub learning_rate: f64, - /// Weight decay - pub weight_decay: f64, - /// Gradient clipping threshold - pub grad_clip: f64, - /// Warmup steps - pub warmup_steps: usize, - /// Adam beta1 momentum parameter - pub adam_beta1: f64, - /// P1: Adam beta2 parameter (Agent 2) - pub adam_beta2: f64, - /// P1: Adam epsilon (Agent 2) - pub adam_epsilon: f64, - /// P1: Total decay steps for cosine schedule (Agent 2) - pub total_decay_steps: usize, - /// Optimizer type (Adam or SGD) - pub optimizer_type: OptimizerType, - /// SGD momentum (only used when optimizer_type = SGD) - pub sgd_momentum: f64, - /// Training batch size - pub batch_size: usize, - /// Sequence length for training - pub seq_len: usize, - /// Shuffle batches every epoch (default: false for reproducibility) - pub shuffle_batches: bool, - /// P2: Sequence stride for overlapping windows (Agent 3) - pub sequence_stride: usize, - /// P2: Normalization epsilon for layer norm (Agent 3) - pub norm_eps: f64, - /// Enable early stopping (default: true) - pub early_stopping_enabled: bool, - /// Early stopping patience (epochs without improvement) - pub early_stopping_patience: usize, - /// Early stopping threshold (minimum improvement) - pub early_stopping_min_delta: f64, - /// Minimum epochs before early stopping can trigger - pub early_stopping_min_epochs: usize, -} - -impl Default for Mamba2Config { - fn default() -> Self { - Self::emergency_safe_defaults() - } -} - -impl Mamba2Config { - /// Create Mamba2 config from central configuration system - /// - /// CRITICAL: Eliminates dangerous hardcoded defaults that could cause - /// training instability or memory issues in production - pub fn from_config_manager( - _config_manager: &config::ConfigManager, - ) -> Result> { - // Use emergency defaults since specific MAMBA configs may not be available - tracing::warn!( - "Using emergency MAMBA config defaults - MAMBA configs not available in ServiceConfig" - ); - Ok(Self::emergency_safe_defaults()) - } - - /// EMERGENCY FALLBACK: Ultra-conservative Mamba2 defaults - /// - /// WARNING: These defaults prioritize safety over performance - /// and are not suitable for production training - pub fn emergency_safe_defaults() -> Self { - tracing::error!( - "Using emergency Mamba2 defaults - check configuration system immediately!" - ); - Self { - d_model: 225, // Wave C (201) + Wave D (24) = 225 - d_state: 64, // P0 FIX: Mamba-2 official recommendation (was 16) - d_head: 16, // Small head size - num_heads: 2, // Minimal heads - expand: 1, // No expansion to minimize memory - num_layers: 1, // Single layer only - dropout: 0.5, // High dropout for safety - use_ssd: false, // Disable advanced features - use_selective_state: false, // Disable advanced features - hardware_aware: false, // Disable optimizations - target_latency_us: 1000, // Very conservative latency - max_seq_len: 128, // Short sequences only - learning_rate: 1e-6, // Extremely conservative learning rate - weight_decay: 1e-3, // High weight decay for stability - grad_clip: 0.1, // Aggressive gradient clipping - warmup_steps: 10, // Minimal warmup - adam_beta1: 0.9, // Standard Adam beta1 - adam_beta2: 0.999, // P1: Standard Adam beta2 - adam_epsilon: 1e-8, // P1: Standard Adam epsilon - total_decay_steps: 10000, // P1: Standard decay schedule - optimizer_type: OptimizerType::AdamW, // AdamW for better SSM training - sgd_momentum: 0.9, // Standard SGD momentum - batch_size: 1, // Single sample batches - seq_len: 64, // Very short sequences - shuffle_batches: false, // Deterministic by default - sequence_stride: 1, // P2: No overlapping (safe default) - norm_eps: 1e-5, // P2: Standard layer norm epsilon - early_stopping_enabled: true, // Enable early stopping by default - early_stopping_patience: 20, // 20 epochs patience (TFT default) - early_stopping_min_delta: 1e-4, // Minimum improvement threshold - early_stopping_min_epochs: 20, // Minimum 20 epochs before stopping - } - } - - /// Estimate memory usage for safety validation - fn estimate_memory_usage(config: &config::Mamba2Config) -> usize { - // Rough estimation: d_model * num_layers * batch_size * seq_len * 4 bytes (f32) - // Plus additional overhead for state and intermediate computations - let base_memory = - config.d_model * config.num_layers * config.batch_size * config.seq_len * 4; - let overhead_factor = 3; // Account for gradients, optimizer states, etc. - (base_memory * overhead_factor) / (1024 * 1024) // Convert to MB - } -} - -/// `MAMBA-2` state container -#[derive(Debug, Clone)] -pub struct Mamba2State { - /// Hidden states for each layer - pub hidden_states: Vec, - /// Selective state components - pub selective_state: Vec, - /// State transition matrices A, B, C - pub ssm_states: Vec, - /// Compression indices for memory efficiency - pub compression_indices: Vec, - /// Performance metrics - pub metrics: HashMap, - /// Best validation loss (for early stopping) - pub best_val_loss: f64, - /// Patience counter (epochs without improvement) - pub patience_counter: usize, - /// Early stopping triggered flag - pub stopped: bool, - /// Epoch where early stopping triggered - pub stopped_at_epoch: Option, - /// Last update timestamp - pub last_update: Instant, -} - -/// State Space Model state matrices -/// -/// Mathematical notation: A, B, C matrices follow standard SSM formulation -/// where uppercase letters represent state-space matrices as per control theory convention -#[derive(Debug, Clone)] -#[allow(non_snake_case)] -pub struct SSMState { - /// State transition matrix A (d_state × d_state) - /// Mathematical notation: uppercase A is standard in control theory and SSM literature - #[allow(non_snake_case)] - pub A: Tensor, - /// Input matrix B (d_state × d_model) - /// Mathematical notation: uppercase B is standard in control theory and SSM literature - #[allow(non_snake_case)] - pub B: Tensor, - /// Output matrix C (d_model × d_state) - /// Mathematical notation: uppercase C is standard in control theory and SSM literature - #[allow(non_snake_case)] - pub C: Tensor, - /// Discretization parameter Δ (Delta) - pub delta: Tensor, - /// Current hidden state - pub hidden: Tensor, -} - -impl SSMState { - /// Reset SSM state to zeros (call between epochs to prevent state accumulation) - /// - /// # Errors - /// - /// Returns `MLError` if tensor operations fail - pub fn reset(&mut self) -> Result<(), MLError> { - // Reset A, B, C matrices to initial random values (small initialization for stability) - // Clone device first to avoid borrow checker issues - let device = self.A.device().clone(); - let d_state = self.A.dim(0)?; - let d_inner = self.B.dim(1)?; - let d_model = self.delta.dims()[0]; - let batch_size = self.hidden.dim(0)?; - - // Re-initialize A matrix [d_state, d_state] -- use Tensor::randn to avoid temporary Vec allocation - self.A = Tensor::randn(0_f32, 0.02, (d_state, d_state), &device)?; - - // Re-initialize B matrix [d_state, d_inner] - self.B = Tensor::randn(0_f32, 0.02, (d_state, d_inner), &device)?; - - // Re-initialize C matrix [d_inner, d_state] - self.C = Tensor::randn(0_f32, 0.02, (d_inner, d_state), &device)?; - - // Reset delta to ones - self.delta = Tensor::ones((d_model,), DType::F32, &device)?; - - // Reset hidden state to zeros - self.hidden = Tensor::zeros((batch_size, d_state), DType::F32, &device)?; - - Ok(()) - } -} - -impl Mamba2State { - /// Create a zero-initialized state - /// - /// # Errors - /// - /// Returns `MLError` if: - /// - CUDA device initialization fails (falls back to CPU) - /// - Tensor allocation fails - /// - Memory allocation exceeds available resources - pub fn zeros(config: &Mamba2Config, device: &Device) -> Result { - let mut hidden_states = Vec::new(); - let mut ssm_states = Vec::new(); - let d_inner = config.d_model * config.expand; // CRITICAL: Use d_inner after input_projection - - for layer_idx in 0..config.num_layers { - // Create hidden state with proper error handling - let hidden = Tensor::zeros((config.batch_size, config.d_model), DType::F32, device) - .map_err(|e| MLError::TensorCreationError { - operation: format!("hidden state creation for layer {}", layer_idx), - reason: e.to_string(), - })?; - hidden_states.push(hidden); - - // Initialize SSM matrices with F32 dtype for GPU throughput - let A = { - let shape = (config.d_state, config.d_state); - let num_elements = shape.0 * shape.1; - let values: Vec = (0..num_elements) - .map(|_| { - use rand::Rng; - let mut rng = rand::thread_rng(); - rng.gen_range(-1.0_f32..1.0) * 0.02 // Small initialization for stability - }) - .collect(); - Tensor::from_vec(values, shape, device).map_err(|e| { - MLError::TensorCreationError { - operation: format!("SSM A matrix creation for layer {}", layer_idx), - reason: e.to_string(), - } - })? - }; - trace!( - "Layer {} A matrix initialized: shape={:?}, dtype=F32", - layer_idx, - A.dims() - ); - - // B must be [d_state, d_inner] with F32 dtype - let B = { - let shape = (config.d_state, d_inner); - let num_elements = shape.0 * shape.1; - let values: Vec = (0..num_elements) - .map(|_| { - use rand::Rng; - let mut rng = rand::thread_rng(); - rng.gen_range(-1.0_f32..1.0) * 0.02 - }) - .collect(); - Tensor::from_vec(values, shape, device).map_err(|e| { - MLError::TensorCreationError { - operation: format!("SSM B matrix creation for layer {}", layer_idx), - reason: e.to_string(), - } - })? - }; - trace!( - "Layer {} B matrix initialized: shape={:?}, dtype=F32", - layer_idx, - B.dims() - ); - - // C must be [d_inner, d_state] with F32 dtype - let C = { - let shape = (d_inner, config.d_state); - let num_elements = shape.0 * shape.1; - let values: Vec = (0..num_elements) - .map(|_| { - use rand::Rng; - let mut rng = rand::thread_rng(); - rng.gen_range(-1.0_f32..1.0) * 0.02 - }) - .collect(); - Tensor::from_vec(values, shape, device).map_err(|e| { - MLError::TensorCreationError { - operation: format!("SSM C matrix creation for layer {}", layer_idx), - reason: e.to_string(), - } - })? - }; - trace!( - "Layer {} C matrix initialized: shape={:?}, dtype=F32", - layer_idx, - C.dims() - ); - - let delta = Tensor::ones((config.d_model,), DType::F32, device).map_err(|e| { - MLError::TensorCreationError { - operation: format!("delta tensor creation for layer {}", layer_idx), - reason: e.to_string(), - } - })?; - - let ssm_hidden = Tensor::zeros((config.batch_size, config.d_state), DType::F32, device) - .map_err(|e| MLError::TensorCreationError { - operation: format!("SSM hidden state creation for layer {}", layer_idx), - reason: e.to_string(), - })?; - - ssm_states.push(SSMState { - A, - B, - C, - delta, - hidden: ssm_hidden, - }); - } - - Ok(Self { - hidden_states, - selective_state: vec![0.0; config.d_model * config.expand], - ssm_states, - compression_indices: Vec::new(), - metrics: HashMap::new(), - best_val_loss: f64::INFINITY, - patience_counter: 0, - stopped: false, - stopped_at_epoch: None, - last_update: Instant::now(), - }) - } - - /// Compress state to reduce memory usage - pub fn compress(&mut self, compression_ratio: f64) { - let target_size = (self.selective_state.len() as f64 * compression_ratio) as usize; - - // Sort by magnitude and keep top components - let mut indexed_values: Vec<(usize, f64)> = self - .selective_state - .iter() - .enumerate() - .map(|(i, &v)| (i, v.abs())) - .collect(); - - indexed_values.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal)); - - self.compression_indices.clear(); - for i in 0..target_size.min(indexed_values.len()) { - self.compression_indices.push(indexed_values[i].0); - } - - // Zero out non-selected components - for i in 0..self.selective_state.len() { - if !self.compression_indices.contains(&i) { - self.selective_state[i] = 0.0; - } - } - } -} - -/// Training metadata for `MAMBA-2` model -#[derive(Debug, Clone, Serialize, Deserialize)] -pub struct Mamba2Metadata { - pub model_id: String, - pub created_at: SystemTime, - pub version: String, - pub input_dim: usize, - pub output_dim: usize, - pub num_parameters: usize, - pub training_history: VecDeque, - pub performance_stats: HashMap, - pub last_checkpoint: Option, -} - -/// Training epoch information -#[derive(Debug, Clone, Serialize, Deserialize)] -pub struct TrainingEpoch { - pub epoch: usize, - pub loss: f64, - pub accuracy: f64, - pub learning_rate: f64, - pub duration_seconds: f64, - pub timestamp: SystemTime, -} - -/// CUDA-compatible LayerNorm wrapper for MAMBA-2 -/// -/// This wrapper uses manual CUDA implementation to avoid -/// "no cuda implementation for layer-norm" error from Candle. -#[derive(Debug, Clone)] -pub struct CudaLayerNorm { - normalized_shape: Vec, - weight: Option, - bias: Option, - eps: f64, -} - -impl CudaLayerNorm { - pub fn new(normalized_shape: usize, eps: f64, vb: VarBuilder<'_>) -> Result { - // Create learnable weight and bias parameters - let weight = vb.get(normalized_shape, "weight")?; - let bias = vb.get(normalized_shape, "bias")?; - - Ok(Self { - normalized_shape: vec![normalized_shape], - weight: Some(weight), - bias: Some(bias), - eps, - }) - } - - pub fn forward(&self, x: &Tensor) -> Result { - layer_norm_with_fallback( - x, - &self.normalized_shape, - self.weight.as_ref(), - self.bias.as_ref(), - self.eps, - ) - } -} - -/// `MAMBA-2` State-Space Model implementation -pub struct Mamba2SSM { - pub config: Mamba2Config, - pub metadata: Mamba2Metadata, - pub state: Mamba2State, - pub ssd_layers: Vec, - pub selective_state: Option, - pub hardware_optimizer: Option, - pub scan_engine: Arc, - pub is_trained: bool, - pub device: Device, - - // Model parameters - pub input_projection: Linear, - pub output_projection: Linear, - pub layer_norms: Vec, - pub dropouts: Vec, - - // Training state - pub optimizer_state: HashMap, - pub gradients: HashMap, - pub grad_scaler: f64, - pub step_count: usize, - pub current_lr: f64, - pub total_training_samples: usize, - - // Performance counters - pub total_inferences: AtomicU64, - pub total_training_steps: AtomicU64, - pub latency_histogram: VecDeque, - - // AGENT F2: VarMap for checkpoint saving (CRITICAL FIX) - // This stores all trainable parameters for safetensors serialization - pub varmap: Arc, -} - -impl std::fmt::Debug for Mamba2SSM { - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - f.debug_struct("Mamba2SSM") - .field("config", &self.config) - .field("metadata", &self.metadata) - .field("is_trained", &self.is_trained) - .field("device", &self.device) - .field("step_count", &self.step_count) - .field("varmap", &"") - .finish() - } -} - -impl Mamba2SSM { - /// Create a scalar tensor with automatic dtype conversion - /// - /// Helper function to eliminate repetitive dtype matching boilerplate. - /// Automatically converts f64 values to the appropriate tensor dtype. - fn scalar_tensor(value: f64, dtype: DType, device: &Device) -> Result { - match dtype { - DType::F32 => { - Tensor::new(&[value as f32], device).map_err(|e| MLError::TensorCreationError { - operation: "scalar_tensor (F32)".to_owned(), - reason: e.to_string(), - }) - }, - DType::F64 => { - // F64 path kept for backward compatibility but should not be hit - // after the F32 migration - Tensor::new(&[value], device).map_err(|e| MLError::TensorCreationError { - operation: "scalar_tensor (F64)".to_owned(), - reason: e.to_string(), - }) - }, - DType::BF16 | DType::F16 => { - // Create as F32 scalar, then cast to the target half-precision dtype. - Tensor::new(&[value as f32], device) - .and_then(|t| t.to_dtype(dtype)) - .map_err(|e| MLError::TensorCreationError { - operation: format!("scalar_tensor ({:?} via F32)", dtype), - reason: e.to_string(), - }) - }, - DType::F8E4M3 | DType::U8 | DType::U32 | DType::I64 => { - Err(MLError::ModelError(format!( - "Unsupported dtype: {:?}", - dtype - ))) - }, - } - } - - /// Create new `MAMBA-2` model - /// - /// # Errors - /// - /// Returns `MLError` if: - /// - Variable initialization fails - /// - Linear layer creation fails - /// - Layer norm creation fails - /// - SSD layer initialization fails - pub fn new(config: Mamba2Config, device: &Device) -> Result { - if config.d_model == 0 { - return Err(MLError::ConfigError("Mamba2 requires d_model > 0".to_owned())); - } - - let vs = Arc::new(candle_nn::VarMap::new()); - let vb = VarBuilder::from_varmap(&vs, training_dtype(device), device); - - let d_inner = config.d_model * config.expand; - - let input_projection = candle_nn::linear(config.d_model, d_inner, vb.pp("input_proj"))?; - // FIXED (Agent 246): Output projection should map d_inner to 1 for regression (price prediction) - // The model performs price regression, NOT sequence-to-sequence modeling - // Output shape: [batch, seq, d_inner] → [batch, seq, 1] - let output_projection = candle_nn::linear(d_inner, 1, vb.pp("output_proj"))?; - - let mut layer_norms = Vec::new(); - let mut dropouts = Vec::new(); - let mut ssd_layers = Vec::new(); - - // Layer norm must match d_inner (d_model * expand) since input_projection expands the dimension - for i in 0..config.num_layers { - let ln = CudaLayerNorm::new(d_inner, config.norm_eps, vb.pp(format!("ln_{}", i)))?; - layer_norms.push(ln); - - let dropout = Dropout::new(config.dropout as f32); - dropouts.push(dropout); - - // CRITICAL FIX: Pass VarBuilder to SSDLayer to register parameters in parent VarMap - // Previously passed device, causing SSDLayer to create local VarMap (90% of params lost) - let ssd_layer = SSDLayer::new(&config, i, vb.clone())?; - ssd_layers.push(ssd_layer); - } - - let selective_state = config - .use_selective_state - .then(|| SelectiveStateSpace::new(&config)) - .transpose()?; - - let hardware_optimizer = config - .hardware_aware - .then(|| HardwareOptimizer::new(&config)) - .transpose()?; - - let scan_engine = Arc::new(ParallelScanEngine::new(device.clone(), 1_000_000)); - - let metadata = Mamba2Metadata { - model_id: Uuid::new_v4().to_string(), - created_at: SystemTime::now(), - version: "2.0.0".to_owned(), - input_dim: config.d_model, - output_dim: 1, // FIXED (Agent 246): Regression output (price prediction), not sequence-to-sequence - num_parameters: Self::count_parameters(&config), - training_history: VecDeque::new(), - performance_stats: HashMap::new(), - last_checkpoint: None, - }; - - let state = Mamba2State::zeros(&config, device)?; - - // Store learning_rate before moving config - let learning_rate = config.learning_rate; - - Ok(Self { - config, - metadata, - state, - ssd_layers, - selective_state, - hardware_optimizer, - scan_engine, - is_trained: false, - device: device.clone(), - input_projection, - output_projection, - layer_norms, - dropouts, - optimizer_state: HashMap::new(), - gradients: HashMap::new(), - grad_scaler: 1.0, - step_count: 0, - current_lr: learning_rate, - total_training_samples: 0, - total_inferences: AtomicU64::new(0), - total_training_steps: AtomicU64::new(0), - latency_histogram: VecDeque::new(), - varmap: vs, // AGENT F2: Store VarMap for checkpoint saving - }) - } - - /// Count total parameters in model - fn count_parameters(config: &Mamba2Config) -> usize { - let d_inner = config.d_model * config.expand; - let input_proj_params = config.d_model * d_inner; - let output_proj_params = d_inner; // FIXED (Agent 246): d_inner * 1 for regression output - let layer_params = config.num_layers - * ( - config.d_model * 3 + // Layer norm - config.d_model * config.d_state * 3 + // A, B, C matrices - config.d_model - // Delta parameters - ); - - input_proj_params + output_proj_params + layer_params - } - - /// Create HFT-optimized configuration - /// - /// # Errors - /// - /// Returns `MLError` if: - /// - Model initialization fails - /// - Hardware configuration is invalid - /// - Resource allocation fails - pub fn default_hft(device: &Device) -> Result { - let config = Mamba2Config { - d_model: 256, - d_state: 64, // P0 FIX: Mamba-2 official recommendation (was 32) - d_head: 32, - num_heads: 8, - expand: 2, - num_layers: 4, - target_latency_us: 3, - hardware_aware: true, - use_ssd: true, - use_selective_state: true, - max_seq_len: 1024, - batch_size: 16, - seq_len: 256, - ..Default::default() - }; - - Self::new(config, device) - } - - /// Forward pass through the model - /// - /// # Errors - /// - /// Returns `MLError` if: - /// - Input projection fails - /// - Layer normalization fails - /// - SSD layer processing fails - /// - Output projection fails - /// - Tensor operations fail - #[instrument(skip(self, input))] - pub fn forward(&mut self, input: &Tensor) -> Result { - let input = crate::dqn::mixed_precision::ensure_training_dtype(input) - .map_err(|e| MLError::ModelError(e.to_string()))?; - let start = Instant::now(); - - // OPTIMIZATION: Device affinity check (catch cross-device transfers early) - // Compare device types (CUDA vs CPU) since Device doesn't implement PartialEq - if input.device().is_cuda() != self.device.is_cuda() { - return Err(MLError::ModelError(format!( - "Input tensor on wrong device: expected {:?}, got {:?}", - self.device, - input.device() - ))); - } - - // Input projection - let mut hidden = self.input_projection.forward(&input)?; - - // Process through each layer - collect indices first to avoid borrow conflicts - let num_layers = self.ssd_layers.len(); - for layer_idx in 0..num_layers { - // Layer normalization - let normalized = self.layer_norms[layer_idx].forward(&hidden)?; - - // OPTIMIZATION: SSD layer processing - clone ssd_layer to avoid borrow conflict - // The forward_ssd_layer method requires &mut self, so we must clone the layer - let ssd_layer = self.ssd_layers[layer_idx].clone(); - let layer_output = self.forward_ssd_layer(&ssd_layer, &normalized, layer_idx)?; - - // Residual connection - hidden = (&hidden + &layer_output)?; - - // Dropout - if self.config.dropout > 0.0 { - hidden = self.dropouts[layer_idx].forward(&hidden, true)?; - } - } - - // Output projection with sigmoid activation (P0 FIX: bound output to [0,1] for normalized targets) - let output_raw = self.output_projection.forward(&hidden)?; - let output = crate::cuda_compat::manual_sigmoid(&output_raw)?; - - // Cast output back to F32 for API compatibility - let output = output.to_dtype(candle_core::DType::F32) - .map_err(|e| MLError::ModelError(format!("Output dtype cast failed: {}", e)))?; - - // OPTIMIZATION: Update performance metrics with VecDeque (O(1) instead of O(n)) - let inference_time = start.elapsed(); - self.total_inferences.fetch_add(1, Ordering::Relaxed); - self.latency_histogram.push_back(inference_time); - - if self.latency_histogram.len() > 10000 { - self.latency_histogram.pop_front(); // O(1) operation instead of O(n) remove(0) - } - - Ok(output) - } - - /// Forward pass through SSD layer with selective scan - #[instrument(skip(self, _ssd_layer, input))] - fn forward_ssd_layer( - &mut self, - _ssd_layer: &SSDLayer, - input: &Tensor, - layer_idx: usize, - ) -> Result { - trace!( - "forward_ssd_layer layer {}: input shape={:?}", - layer_idx, - input.dims() - ); - - // Use references to avoid unnecessary clones (Agent MAMBA-MEMORY-FIX) - let dt = &self.state.ssm_states[layer_idx].delta; - let A = &self.state.ssm_states[layer_idx].A; - let B = &self.state.ssm_states[layer_idx].B; - let C = &self.state.ssm_states[layer_idx].C; - - trace!( - "forward_ssd_layer layer {}: B shape={:?}", - layer_idx, - B.dims() - ); - - // Discretize the continuous-time SSM - let A_discrete = self.discretize_ssm(A, dt)?; - let B_discrete = self.discretize_ssm_input(B, dt)?; - trace!( - "forward_ssd_layer layer {}: B_discrete shape={:?}", - layer_idx, - B_discrete.dims() - ); - - // Selective scan algorithm - let scan_input = self.prepare_scan_input(input, &A_discrete, &B_discrete)?; - trace!( - "scan_input shape: {:?}, B shape: {:?}, C shape: {:?}", - scan_input.dims(), - B.dims(), - C.dims() - ); - let scanned_states = self - .scan_engine - .parallel_prefix_scan(&scan_input, ScanOperator::SSMScan)?; - trace!("scanned_states shape: {:?}", scanned_states.dims()); - - // Apply output transformation - trace!( - "About to matmul: scanned_states {:?} × C.t() (C is {:?})", - scanned_states.dims(), - C.dims() - ); - let batch_size = scanned_states.dim(0)?; - // Cast C to match scanned_states dtype (SSM state is F32 but computation may be BF16) - let C_cast = C.to_dtype(scanned_states.dtype())?; - let C_t = C_cast.t()?.contiguous()?; - let C_broadcasted = - C_t.unsqueeze(0)? - .broadcast_as((batch_size, C_t.dim(0)?, C_t.dim(1)?))?; - let output = scanned_states.matmul(&C_broadcasted)?; - - // Update hidden state - let _batch_size = input.dim(0)?; - let seq_len = input.dim(1)?; - if seq_len > 0 { - let last_state = scanned_states.narrow(1, seq_len - 1, 1)?.squeeze(1)?; - self.state.ssm_states[layer_idx].hidden = last_state; - } - - Ok(output) - } - - /// Discretize continuous-time SSM matrix A - /// - /// Mathematical notation: A_cont follows standard SSM notation for continuous-time state transition matrix - #[allow(non_snake_case)] - #[allow(non_snake_case)] - fn discretize_ssm(&self, A_cont: &Tensor, dt: &Tensor) -> Result { - // Keep dt on GPU — mean_all() returns a 0-D tensor, no CPU round-trip - let dt_scalar = dt.mean_all()?; - - // Bilinear (Tustin) approximation: A_disc = I + A*dt + (A*dt)^2 / 2 - // More accurate than ZOH (I + A*dt), matches ssd_layer.rs - let A_dt = A_cont.broadcast_mul(&dt_scalar)?; - let A_dt_sq = A_dt.matmul(&A_dt)?; - let half = Tensor::new(0.5_f32, A_cont.device())?; - let second_order = A_dt_sq.broadcast_mul(&half)?; - - let identity = Tensor::eye(A_cont.dim(0)?, DType::F32, A_cont.device())?; - let A_discrete = ((&identity + &A_dt)? + &second_order)?; - - Ok(A_discrete) - } - - /// Discretize continuous-time input matrix B - /// - /// Mathematical notation: B_cont follows standard SSM notation for continuous-time input matrix - #[allow(non_snake_case)] - #[allow(non_snake_case)] - fn discretize_ssm_input(&self, B_cont: &Tensor, dt: &Tensor) -> Result { - // Keep dt on GPU — mean_all() returns a 0-D tensor, no CPU round-trip - let dt_scalar = dt.mean_all()?; - let B_discrete = B_cont.broadcast_mul(&dt_scalar)?; - - Ok(B_discrete) - } - - /// Prepare input for selective scan algorithm - /// - /// Mathematical notation: Parameters _A and B follow standard SSM notation - #[allow(non_snake_case)] - fn prepare_scan_input( - &self, - input: &Tensor, - _A: &Tensor, - B: &Tensor, - ) -> Result { - // Transpose B and broadcast to match batch dimension - // input: [batch, seq, d_inner], B: [d_state, d_inner] - // B.t(): [d_inner, d_state] → broadcast to [batch, d_inner, d_state] - let batch_size = input.dim(0)?; - trace!( - "prepare_scan_input: input shape: {:?}, B shape: {:?}", - input.dims(), - B.dims() - ); - trace!( - "prepare_scan_input: d_model: {}, d_inner: {}, d_state: {}", - self.config.d_model, - self.config.d_model * self.config.expand, - self.config.d_state - ); - - // Cast B to match input dtype (SSM state is F32 but input may be BF16) - let B_cast = B.to_dtype(input.dtype())?; - let B_t = B_cast.t()?.contiguous()?; - let d_inner = B_t.dim(0)?; - let d_state = B_t.dim(1)?; - let B_broadcasted = B_t - .unsqueeze(0)? - .broadcast_as((batch_size, d_inner, d_state))?; - trace!( - "prepare_scan_input: B broadcasted shape: {:?}", - B_broadcasted.dims() - ); - - let Bu = input.matmul(&B_broadcasted)?; - trace!( - "prepare_scan_input: Bu shape: {:?}, expected [batch={}, seq={}, d_state={}]", - Bu.dims(), - input.dim(0)?, - input.dim(1)?, - self.config.d_state - ); - - Ok(Bu) - } - /// Fast single prediction for HFT - /// - /// # Errors - /// - /// Returns `MLError` if: - /// - Tensor creation fails - /// - Forward pass fails - /// - Output extraction fails - /// - Value conversion fails - pub fn predict_single_fast(&mut self, input: &[f64]) -> Result { - let start = Instant::now(); - - if input.len() != self.config.d_model { - return Err(MLError::InvalidInput(format!( - "Expected input dimension {}, got {}", - self.config.d_model, - input.len() - ))); - } - - let device = self.device(); - // Convert f64 input to f32 for F32 model dtype - let input_f32: Vec = input.iter().map(|&v| v as f32).collect(); - let input_tensor = Tensor::from_vec(input_f32, (1, input.len()), device)?; - - let output = self.forward(&input_tensor)?; - // Model uses F32 tensors — extract as f32 then widen to f64 for API compat - let result: f32 = output.to_scalar()?; - - let elapsed = start.elapsed(); - if elapsed.as_micros() > self.config.target_latency_us as u128 { - warn!( - "Prediction exceeded target latency: {}μs", - elapsed.as_micros() - ); - } - - Ok(result as f64) - } - - /// Get performance metrics - pub fn get_performance_metrics(&self) -> HashMap { - let mut metrics = HashMap::new(); - - metrics.insert( - "total_inferences".to_owned(), - self.total_inferences.load(Ordering::Relaxed) as f64, - ); - metrics.insert( - "total_training_steps".to_owned(), - self.total_training_steps.load(Ordering::Relaxed) as f64, - ); - - if !self.latency_histogram.is_empty() { - let avg_latency = self - .latency_histogram - .iter() - .map(|d| d.as_micros() as f64) - .sum::() - / self.latency_histogram.len() as f64; - metrics.insert("avg_latency_us".to_owned(), avg_latency); - - let throughput = 1_000_000.0 / avg_latency; // predictions per second - metrics.insert("throughput_pps".to_owned(), throughput); - } - - // Hardware metrics - if let Some(hw_optimizer) = &self.hardware_optimizer { - let hw_metrics = hw_optimizer.get_performance_metrics(); - for (k, v) in hw_metrics { - metrics.insert(k, v); - } - } - - // Model-specific metrics - metrics.insert( - "model_parameters".to_owned(), - self.metadata.num_parameters as f64, - ); - metrics.insert( - "compression_ratio".to_owned(), - if self.state.selective_state.len() > 0 && !self.state.compression_indices.is_empty() { - self.state.compression_indices.len() as f64 - / self.state.selective_state.len() as f64 - } else { - 1.0 - }, - ); - - let latency_target_ratio = if !self.latency_histogram.is_empty() { - let avg_latency = self - .latency_histogram - .iter() - .map(|d| d.as_micros() as f64) - .sum::() - / self.latency_histogram.len() as f64; - avg_latency / self.config.target_latency_us as f64 - } else { - 0.0 - }; - metrics.insert("latency_target_ratio".to_owned(), latency_target_ratio); - - // Additional production metrics for compatibility - metrics.insert("cache_hit_rate".to_owned(), 0.95); - metrics.insert("simd_ops_per_inference".to_owned(), 1000.0); - metrics.insert( - "state_compression_ratio".to_owned(), - metrics.get("compression_ratio").copied().unwrap_or(1.0), - ); - - metrics - } - - /// Get the device this model is on - fn device(&self) -> &Device { - &self.device - } - - /// Clear internal SSM state (call between epochs to prevent state accumulation) - /// - /// # Errors - /// - /// Returns `MLError` if: - /// - SSM state reset fails - /// - Tensor operations fail - pub fn clear_state(&mut self) -> Result<(), MLError> { - // Reset SSM state for each layer to prevent accumulation across epochs - for (layer_idx, ssm_state) in self.state.ssm_states.iter_mut().enumerate() { - ssm_state.reset()?; - trace!("Cleared MAMBA2 SSM state for layer {}", layer_idx); - } - - // Clear selective state components - self.state.selective_state.fill(0.0); - self.state.compression_indices.clear(); - - info!( - "Cleared MAMBA2 SSM state for all {} layers", - self.state.ssm_states.len() - ); - Ok(()) - } - - /// Check early stopping condition with patience (TFT pattern) - /// - /// This method implements the same early stopping logic as TFT - /// (see `ml/src/trainers/tft.rs:1702-1731`). - /// - /// # Arguments - /// - /// * `epoch` - Current epoch number (0-indexed) - /// * `val_loss` - Validation loss for current epoch - /// - /// # Returns - /// - /// `true` if training should stop early, `false` otherwise - /// - /// # Behavior - /// - /// 1. **Before min_epochs**: Always returns `false` (don't stop prematurely) - /// 2. **Improvement detected**: Resets patience counter, updates best_val_loss - /// 3. **No improvement**: Increments patience counter - /// 4. **Patience exhausted**: Sets stopped flag and returns `true` - pub fn check_early_stopping(&mut self, epoch: usize, val_loss: f64) -> bool { - // Don't stop before min_epochs - if epoch < self.config.early_stopping_min_epochs { - return false; - } - - // Check if validation loss improved by more than min_delta - if val_loss < self.state.best_val_loss - self.config.early_stopping_min_delta { - // Improvement detected - reset patience counter - self.state.best_val_loss = val_loss; - self.state.patience_counter = 0; - false - } else { - // No improvement - increment patience counter - self.state.patience_counter += 1; - - if self.state.patience_counter >= self.config.early_stopping_patience { - // Patience exhausted - trigger early stopping - self.state.stopped = true; - self.state.stopped_at_epoch = Some(epoch); - info!( - "Early stopping triggered at epoch {} (patience: {}, best val loss: {:.6})", - epoch, self.config.early_stopping_patience, self.state.best_val_loss - ); - true - } else { - debug!( - "Patience: {}/{} (best val loss: {:.6}, current: {:.6})", - self.state.patience_counter, - self.config.early_stopping_patience, - self.state.best_val_loss, - val_loss - ); - false - } - } - } - - /// Train the model with selective scan algorithm - #[instrument(skip(self, train_data, val_data, checkpoint_dir))] - pub async fn train( - &mut self, - train_data: &[(Tensor, Tensor)], - val_data: &[(Tensor, Tensor)], - epochs: usize, - checkpoint_dir: Option<&std::path::Path>, - ) -> Result, MLError> { - info!("Starting MAMBA-2 training with {} epochs", epochs); - - const MAX_TRAINING_HISTORY: usize = 100; - let mut training_history: VecDeque = VecDeque::with_capacity(MAX_TRAINING_HISTORY + 1); - let mut best_val_loss = f64::INFINITY; - - // FIXED (Agent P2): Set total training samples for accurate LR schedule - self.total_training_samples = train_data.len(); - - // Initialize optimizer - self.initialize_optimizer()?; - - for epoch in 0..epochs { - let epoch_start = Instant::now(); - - // FIXED: Do NOT clear SSM state (A, B, C parameters) - these are model weights - // that must persist across epochs to accumulate gradient updates. - // Clearing them was causing the E11 validation spike by reinitializing with random values. - - let mut epoch_loss = 0.0; - let mut batch_count = 0; - - // Create batch indices (shuffle if configured) - let mut batch_indices: Vec = (0..train_data.len()) - .step_by(self.config.batch_size) - .collect(); - - if self.config.shuffle_batches { - use rand::seq::SliceRandom; - batch_indices.shuffle(&mut rand::thread_rng()); - } - - // Training phase - for &batch_idx in &batch_indices { - let batch_end = (batch_idx + self.config.batch_size).min(train_data.len()); - let batch = &train_data[batch_idx..batch_end]; - - let batch_loss = self.train_batch(batch, epoch)?; - epoch_loss += batch_loss; - batch_count += 1; - - // Update learning rate - self.update_learning_rate(epoch, batch_idx)?; - - if batch_idx % 100 == 0 { - // FIXED (Agent P2): Log current learning rate for monitoring - let current_lr = self.get_current_learning_rate(); - debug!( - "Epoch {}, Batch {}: Loss = {:.6}, LR = {:.6}", - epoch, batch_idx, batch_loss, current_lr - ); - } - } - - epoch_loss /= batch_count as f64; - - // Validation phase - let val_loss = self.validate(val_data)?; - let epoch_accuracy = self.calculate_accuracy(val_data)?; - - // Update learning rate scheduler - let current_lr = self.get_current_learning_rate(); - - let epoch_duration = epoch_start.elapsed().as_secs_f64(); - let training_epoch = TrainingEpoch { - epoch, - loss: epoch_loss, - accuracy: epoch_accuracy, - learning_rate: current_lr, - duration_seconds: epoch_duration, - timestamp: SystemTime::now(), - }; - - training_history.push_back(training_epoch.clone()); - self.metadata.training_history.push_back(training_epoch); - - // Bound training history to prevent unbounded memory growth - while training_history.len() > MAX_TRAINING_HISTORY { - training_history.pop_front(); - } - while self.metadata.training_history.len() > MAX_TRAINING_HISTORY { - self.metadata.training_history.pop_front(); - } - - // Save checkpoint if best model - if val_loss < best_val_loss { - best_val_loss = val_loss; - - let checkpoint_path = if let Some(dir) = checkpoint_dir { - dir.join(format!("best_epoch_{}.ckpt", epoch)) - } else { - std::path::PathBuf::from(format!("best_epoch_{}.ckpt", epoch)) - }; - - let path_str = checkpoint_path.to_str().ok_or_else(|| { - MLError::ConfigError("Checkpoint path contains invalid UTF-8".to_owned()) - })?; - self.save_checkpoint(path_str).await?; - info!( - "New best validation loss: {:.6} at epoch {}", - val_loss, epoch - ); - } - - // Log epoch results - info!( - "Epoch {}/{}: Loss = {:.6}, Val Loss = {:.6}, Accuracy = {:.4}, LR = {:.2e}, Time = {:.2}s", - epoch + 1, epochs, epoch_loss, val_loss, epoch_accuracy, current_lr, epoch_duration - ); - - // Early stopping check - if self.config.early_stopping_enabled && self.check_early_stopping(epoch, val_loss) { - info!( - "Early stopping triggered at epoch {} (patience exhausted)", - epoch - ); - break; - } - } - - self.is_trained = true; - info!("Training completed with {} epochs", training_history.len()); - - Ok(training_history.into()) - } - - /// Train the model with async data loading (prefetch optimization) - /// - /// This method uses AsyncDataLoader to prefetch batches while GPU trains, - /// improving GPU utilization from ~78% to ~90-95% and reducing training - /// time by 20-30%. - /// - /// # Arguments - /// - /// * `train_data` - Training data as (feature, target) tensor pairs - /// * `val_data` - Validation data - /// * `epochs` - Number of training epochs - /// * `batch_size` - Batch size for training - /// * `prefetch_count` - Number of batches to prefetch (2-3 recommended) - /// - /// # Returns - /// - /// Training history with metrics per epoch - /// - /// # Errors - /// - /// Returns `MLError` if training fails - #[instrument(skip(self, train_data, val_data, checkpoint_dir))] - pub async fn train_async( - &mut self, - train_data: &[(Tensor, Tensor)], - val_data: &[(Tensor, Tensor)], - epochs: usize, - batch_size: usize, - prefetch_count: usize, - checkpoint_dir: Option<&std::path::Path>, - ) -> Result, MLError> { - info!( - "Starting MAMBA-2 async training with {} epochs (prefetch={})", - epochs, prefetch_count - ); - - const MAX_TRAINING_HISTORY: usize = 100; - let mut training_history: VecDeque = VecDeque::with_capacity(MAX_TRAINING_HISTORY + 1); - let mut best_val_loss = f64::INFINITY; - - // Set total training samples for accurate LR schedule - self.total_training_samples = train_data.len(); - - // Initialize optimizer - self.initialize_optimizer()?; - - for epoch in 0..epochs { - let epoch_start = Instant::now(); - - let mut epoch_loss = 0.0; - let mut batch_count = 0; - - // Create AsyncDataLoader for this epoch - // CRITICAL FIX: Use actual device from model parameters, not self.device field - // self.device field can become stale or incorrect, causing device mismatch errors - let actual_device = self.input_projection.weight().device(); - - debug!("AsyncDataLoader device verification:"); - debug!(" Model self.device field: {:?}", self.device); - debug!(" Actual parameter device: {:?}", actual_device); - - let mut loader = crate::hyperopt::adapters::async_data_loader::AsyncDataLoader::new( - train_data.to_vec(), - batch_size, - prefetch_count, - actual_device, - ) - .map_err(|e| MLError::TrainingError(format!("Failed to create async loader: {}", e)))?; - - // Training phase with async prefetch - let mut batch_idx = 0; - while let Some((batched_input, batched_target)) = loader.next_batch() { - // Zero gradients - self.zero_gradients()?; - - // Forward pass with selective scan on batched input - let output = self.forward_with_gradients(&batched_input)?; - - // Extract last timestep for next-step prediction - // output: [batch, seq_len, d_model] → [batch, 1, d_model] - let seq_len = output.dim(1)?; - let output_last = output.narrow(1, seq_len - 1, 1)?; - - // Compute loss on last timestep prediction - let loss = self.compute_loss(&output_last, &batched_target)?; - - // Backward pass - compute gradients for SSM parameters - self.backward_pass(&loss, &batched_input, &batched_target)?; - - // Extract scalar AFTER backward to avoid stalling GPU pipeline - let loss_value = loss.to_scalar::()? as f64; - - // Update parameters - self.optimizer_step()?; - - epoch_loss += loss_value; - batch_count += 1; - - // Update learning rate - self.update_learning_rate(epoch, batch_idx)?; - - if batch_idx % 100 == 0 { - let current_lr = self.get_current_learning_rate(); - debug!( - "Epoch {}, Batch {}: Loss = {:.6}, LR = {:.6}", - epoch, batch_idx, loss_value, current_lr - ); - } - - batch_idx += batch_size; - } - - epoch_loss /= batch_count as f64; - - // Validation phase - let val_loss = self.validate(val_data)?; - let epoch_accuracy = self.calculate_accuracy(val_data)?; - - // Update learning rate scheduler - let current_lr = self.get_current_learning_rate(); - - let epoch_duration = epoch_start.elapsed().as_secs_f64(); - let training_epoch = TrainingEpoch { - epoch, - loss: epoch_loss, - accuracy: epoch_accuracy, - learning_rate: current_lr, - duration_seconds: epoch_duration, - timestamp: SystemTime::now(), - }; - - training_history.push_back(training_epoch.clone()); - self.metadata.training_history.push_back(training_epoch); - - // Bound training history to prevent unbounded memory growth - while training_history.len() > MAX_TRAINING_HISTORY { - training_history.pop_front(); - } - while self.metadata.training_history.len() > MAX_TRAINING_HISTORY { - self.metadata.training_history.pop_front(); - } - - // Save checkpoint if best model - if val_loss < best_val_loss { - best_val_loss = val_loss; - - let checkpoint_path = if let Some(dir) = checkpoint_dir { - dir.join(format!("best_epoch_{}.ckpt", epoch)) - } else { - std::path::PathBuf::from(format!("best_epoch_{}.ckpt", epoch)) - }; - - let path_str = checkpoint_path.to_str().ok_or_else(|| { - MLError::ConfigError("Checkpoint path contains invalid UTF-8".to_owned()) - })?; - self.save_checkpoint(path_str).await?; - info!( - "New best validation loss: {:.6} at epoch {}", - val_loss, epoch - ); - } - - // Log epoch results - info!( - "Epoch {}/{}: Loss = {:.6}, Val Loss = {:.6}, Accuracy = {:.4}, LR = {:.2e}, Time = {:.2}s", - epoch + 1, epochs, epoch_loss, val_loss, epoch_accuracy, current_lr, epoch_duration - ); - - // Early stopping check - if self.config.early_stopping_enabled && self.check_early_stopping(epoch, val_loss) { - info!( - "Early stopping triggered at epoch {} (patience exhausted)", - epoch - ); - break; - } - } - - self.is_trained = true; - info!( - "Async training completed with {} epochs", - training_history.len() - ); - - Ok(training_history.into()) - } - - /// Train a single batch with selective scan - #[instrument(skip(self, batch))] - fn train_batch(&mut self, batch: &[(Tensor, Tensor)], _epoch: usize) -> Result { - if batch.is_empty() { - return Ok(0.0); - } - - // FIXED: Batch all individual sequences together into a single batched tensor - // Individual sequences are shape [1, seq_len, d_model], we need [batch_size, seq_len, d_model] - - let actual_batch_size = batch.len(); - - // Collect all input tensors and concatenate along batch dimension - let input_tensors: Vec<&Tensor> = batch.iter().map(|(input, _)| input).collect(); - let batched_input = if actual_batch_size == 1 { - // Single sample - no concatenation needed - input_tensors[0].clone() - } else { - // Concatenate along dimension 0 (batch dimension) - Tensor::cat( - &input_tensors - .iter() - .map(|t| (*t).clone()) - .collect::>(), - 0, - )? - }; - - // Collect all target tensors and concatenate - let target_tensors: Vec<&Tensor> = batch.iter().map(|(_, target)| target).collect(); - let batched_target = if actual_batch_size == 1 { - target_tensors[0].clone() - } else { - Tensor::cat( - &target_tensors - .iter() - .map(|t| (*t).clone()) - .collect::>(), - 0, - )? - }; - - // FIXED: Ensure input and target tensors are on the model's device (GPU) - // This prevents device mismatch errors during forward pass - let batched_input = batched_input.to_device(&self.device)?; - let batched_target = batched_target.to_device(&self.device)?; - - // Zero gradients - self.zero_gradients()?; - - // Forward pass with selective scan on batched input - let output = self.forward_with_gradients(&batched_input)?; - trace!( - "Training loop: batched_input: {:?}, batched_target: {:?}, forward output: {:?}", - batched_input.dims(), - batched_target.dims(), - output.dims() - ); - - // FIXED (Agent 211): Extract last timestep for next-step prediction - // output: [batch, seq_len, d_model] → [batch, 1, d_model] - // This matches target shape [batch, 1, d_model] - let seq_len = output.dim(1)?; - let output_last = output.narrow(1, seq_len - 1, 1)?; - trace!( - "Training loop: output_last (for loss): {:?}", - output_last.dims() - ); - - // Compute loss on last timestep prediction - let loss = self.compute_loss(&output_last, &batched_target)?; - - // Backward pass - compute gradients for SSM parameters - self.backward_pass(&loss, &batched_input, &batched_target)?; - - // Extract scalar AFTER backward to avoid stalling GPU pipeline - let loss_value = loss.to_scalar::()? as f64; - - // Update parameters - self.optimizer_step()?; - - // Update selective state based on gradients (use first sample for importance scoring) - if let Some(selective_state) = &mut self.selective_state { - // Use the first sample in the batch for importance updates - let first_input = input_tensors[0]; - selective_state.update_importance_scores(first_input, &mut self.state)?; - } - - self.total_training_steps.fetch_add(1, Ordering::Relaxed); - self.step_count += 1; - - // Explicit memory cleanup to prevent GPU memory accumulation - drop(output); - drop(output_last); - drop(loss); - drop(batched_input); - drop(batched_target); - - Ok(loss_value) - } - - /// Forward pass with gradient computation enabled - pub fn forward_with_gradients(&mut self, input: &Tensor) -> Result { - // Gradient flow enabled - do not detach - let input = input; - - // Input projection with gradients - let mut hidden = self.input_projection.forward(input)?; - - // Process through each layer with SSM gradients - collect indices first to avoid borrow conflicts - let num_layers = self.ssd_layers.len(); - for layer_idx in 0..num_layers { - // Layer normalization - let normalized = self.layer_norms[layer_idx].forward(&hidden)?; - - // SSD layer processing with selective scan and gradients - let layer_output = { - let ssd_layer = self.ssd_layers[layer_idx].clone(); - self.forward_ssd_layer_with_gradients(&ssd_layer, &normalized, layer_idx)? - }; - - // Residual connection - hidden = (&hidden + &layer_output)?; - - // Dropout (enabled during training) - if self.config.dropout > 0.0 { - hidden = self.dropouts[layer_idx].forward(&hidden, true)?; - } - } - - // Output projection - trace!( - "Before output_projection: hidden shape: {:?}", - hidden.dims() - ); - // P0 FIX: Add sigmoid activation to bound output to [0,1] for normalized targets - let output_raw = self.output_projection.forward(&hidden)?; - let output = crate::cuda_compat::manual_sigmoid(&output_raw)?; - trace!( - "After output_projection + sigmoid: output shape: {:?}", - output.dims() - ); - - Ok(output) - } - - /// Forward pass through SSD layer with gradient tracking - fn forward_ssd_layer_with_gradients( - &mut self, - _ssd_layer: &SSDLayer, - input: &Tensor, - layer_idx: usize, - ) -> Result { - // Use references to avoid unnecessary clones (Agent MAMBA-MEMORY-FIX) - let dt = &self.state.ssm_states[layer_idx].delta; - let A = &self.state.ssm_states[layer_idx].A; - let B = &self.state.ssm_states[layer_idx].B; - let C = &self.state.ssm_states[layer_idx].C; - - // Discretize with gradient tracking - let A_discrete = self.discretize_ssm_with_gradients(A, dt)?; - let B_discrete = self.discretize_ssm_input_with_gradients(B, dt)?; - - // Selective scan with gradient computation - let scan_input = self.prepare_scan_input_with_gradients(input, &A_discrete, &B_discrete)?; - let scanned_states = self.selective_scan_with_gradients(&scan_input, &A_discrete)?; - - // Output transformation with gradients - // FIXED (Agent 207): Broadcast C correctly after transpose - let batch_size = scanned_states.dim(0)?; - trace!("C matrix broadcast: scanned_states shape: {:?}, C original shape (d_inner, d_state): {:?}", scanned_states.dims(), C.dims()); - - // For matmul: [batch, seq, d_state] × [batch, d_state, d_inner] = [batch, seq, d_inner] - // scanned_states: [32, 60, 16] - // C stored as: [d_inner, d_state] = [512, 16] - // Need: [batch, d_state, d_inner] = [32, 16, 512] - // Cast C to match scanned_states dtype (SSM state is F32 but computation may be BF16) - let C_cast = C.to_dtype(scanned_states.dtype())?; - let C_t = C_cast.t()?.contiguous()?; // [512, 16] → [16, 512] - trace!("C transposed (d_state, d_inner): {:?}", C_t.dims()); - - // Now broadcast [16, 512] to [32, 16, 512] - let d_state = C_t.dim(0)?; // 16 - let d_inner = C_t.dim(1)?; // 512 - let C_broadcasted = C_t - .unsqueeze(0)? - .broadcast_as((batch_size, d_state, d_inner))?; - trace!( - "C broadcasted shape: {:?}, expected: [batch={}, d_state={}, d_inner={}]", - C_broadcasted.dims(), - batch_size, - d_state, - d_inner - ); - - let output = scanned_states.matmul(&C_broadcasted)?; - trace!("Output shape: {:?}", output.dims()); - - // Update hidden state - let _batch_size = input.dim(0)?; - let seq_len = input.dim(1)?; - if seq_len > 0 { - let last_state = scanned_states.narrow(1, seq_len - 1, 1)?.squeeze(1)?; - self.state.ssm_states[layer_idx].hidden = last_state; - } - - Ok(output) - } - - /// Selective scan algorithm with gradient computation - /// - /// Mathematical notation: Parameter A represents the state transition matrix - #[allow(non_snake_case)] - fn selective_scan_with_gradients(&self, input: &Tensor, A: &Tensor) -> Result { - let seq_len = input.dim(1)?; - let d_state = input.dim(2)?; - let device = input.device(); - - // AGENT 176 FIX: Add shape assertions to catch dimension bugs early - tracing::debug!( - "[AGENT 176] selective_scan_with_gradients: input={:?}, A={:?}", - input.dims(), - A.dims() - ); - assert_eq!(input.dims().len(), 3, "Input must be [batch, seq, d_state]"); - assert_eq!(A.dims().len(), 2, "A must be [d_state, d_state]"); - assert_eq!( - A.dim(0)?, - d_state, - "A.dim(0) must equal input.dim(2) (d_state)" - ); - - // Initialize state sequence - pre-allocate result tensor to avoid Vec accumulation - // This prevents the 750MB memory leak from accumulating 60 tensors in Vec - let batch_size = input.dim(0)?; - let mut result = Tensor::zeros((batch_size, seq_len, d_state), input.dtype(), device)?; - let mut current_state = Tensor::zeros((batch_size, d_state), input.dtype(), device)?; - - // Cast A to match input dtype (SSM state matrices are F32 but computation may be BF16) - let A_cast = A.to_dtype(input.dtype())?; - - // Sequential scan with state transitions (maintaining gradients) - for t in 0..seq_len { - let x_t = input.narrow(1, t, 1)?.squeeze(1)?; - - // AGENT 176 FIX: Correct batch matrix multiplication - // State transition: h_t = h_{t-1} @ A^T + x_t - // current_state [batch, d_state] × A.t() [d_state, d_state] = [batch, d_state] - // This is the correct way to do batch SSM state transitions - current_state = (current_state.matmul(&A_cast.t()?)? + &x_t)?; - - // Write directly to result tensor (no Vec accumulation, no Tensor::cat doubling) - let current_unsqueezed = current_state.unsqueeze(1)?; - result = result.slice_assign( - &[0..batch_size, t..(t + 1), 0..d_state], - ¤t_unsqueezed, - )?; - } - - // AGENT 176 FIX: Verify output shape matches expected dimensions - tracing::debug!( - "[AGENT 176] selective_scan_with_gradients: output={:?}", - result.dims() - ); - assert_eq!( - result.dims(), - &[input.dim(0)?, seq_len, d_state], - "Output must be [batch, seq, d_state], got {:?}", - result.dims() - ); - - Ok(result) - } - - /// Discretize SSM with gradient tracking - /// - /// Mathematical notation: A_cont follows standard SSM notation for continuous-time state transition matrix - #[allow(non_snake_case)] - #[allow(non_snake_case)] - fn discretize_ssm_with_gradients( - &self, - A_cont: &Tensor, - dt: &Tensor, - ) -> Result { - // Keep dt on GPU — mean_all() returns a 0-D tensor, no CPU round-trip - let dt_scalar = dt.mean_all()?; - - // Scale A matrix by dt - let A_scaled = A_cont.broadcast_mul(&dt_scalar)?; - - // Matrix exponential approximation: exp(A) ≈ I + A + A²/2 + A³/6 - let identity = Tensor::eye(A_cont.dim(0)?, DType::F32, A_cont.device())?; - let A2 = A_scaled.matmul(&A_scaled)?; - let A3 = A2.matmul(&A_scaled)?; - - let half = Tensor::new(0.5_f32, A_cont.device())?; - let sixth = Tensor::new(1.0_f32 / 6.0, A_cont.device())?; - let A_discrete = (&identity - + &A_scaled - + &A2.broadcast_mul(&half)? - + &A3.broadcast_mul(&sixth)?)?; - - Ok(A_discrete) - } - - /// Discretize input matrix with gradients - /// - /// Mathematical notation: B_cont follows standard SSM notation for continuous-time input matrix - #[allow(non_snake_case)] - #[allow(non_snake_case)] - fn discretize_ssm_input_with_gradients( - &self, - B_cont: &Tensor, - dt: &Tensor, - ) -> Result { - // Keep dt on GPU — mean_all() returns a 0-D tensor, no CPU round-trip - let dt_scalar = dt.mean_all()?; - let B_discrete = B_cont.broadcast_mul(&dt_scalar)?; - Ok(B_discrete) - } - - /// Prepare scan input with gradient tracking - /// - /// Mathematical notation: Parameters _A and B follow standard SSM notation - #[allow(non_snake_case)] - fn prepare_scan_input_with_gradients( - &self, - input: &Tensor, - _A: &Tensor, - B: &Tensor, - ) -> Result { - // FIXED (Agent 248 + Agent 250): Explicit batch broadcast for B matrix - // input: [batch, seq, d_inner], B: [d_state, d_inner] - // B.t(): [d_inner, d_state] → explicit repeat to [batch, d_inner, d_state] - let batch_size = input.dim(0)?; - // Cast B to match input dtype (SSM state is F32 but input may be BF16) - let B_cast = B.to_dtype(input.dtype())?; - let B_t = B_cast.t()?.contiguous()?; // [d_state, d_inner] → [d_inner, d_state] - - // CRITICAL FIX: Use repeat/expand instead of broadcast_as for CUDA compatibility - // Create [batch, d_inner, d_state] by repeating the [d_inner, d_state] tensor - let B_expanded = B_t.unsqueeze(0)?; // [1, d_inner, d_state] - - // Repeat along batch dimension - let B_broadcasted = B_expanded.expand(&[batch_size, B_t.dim(0)?, B_t.dim(1)?])?; - - trace!( - "[Agent 250] B matrix broadcast: B_t={:?} → B_broadcasted={:?}", - B_t.dims(), - B_broadcasted.dims() - ); - - let Bu = input.matmul(&B_broadcasted)?; - trace!("[Agent 250] Bu result shape: {:?}", Bu.dims()); - Ok(Bu) - } - - /// Compute training loss - pub fn compute_loss(&self, output: &Tensor, target: &Tensor) -> Result { - // Mean Squared Error for regression - let diff = (output - target)?; - let squared_diff = (&diff * &diff)?; - let loss = squared_diff.mean_all()?; - // loss is F64 from mean_all() - Ok(loss) - } - - /// Backward pass - compute gradients for model parameters - pub fn backward_pass( - &mut self, - loss: &Tensor, - _input: &Tensor, - _target: &Tensor, - ) -> Result<(), MLError> { - // Compute gradients using automatic differentiation - // The loss tensor should already have the computational graph attached - let grads = loss.backward()?; - - // FIXED (P0): Extract REAL gradients from VarMap trainable parameters - // The trainable parameters are: input_projection, output_projection, layer_norms (weight/bias) - // SSM matrices (A, B, C, delta) are NOT trainable - they're part of the model state - trace!("[P0 FIX] Extracting real gradients from VarMap trainable parameters"); - - self.gradients.clear(); - - // Extract gradients from all VarMap parameters - let all_vars = self.varmap.all_vars(); - let mut total_grad_norm = 0.0_f64; - let mut params_with_grads = 0; - - for (idx, var) in all_vars.iter().enumerate() { - if let Some(grad) = grads.get(var) { - // Compute gradient norm for monitoring - let grad_vec = grad - .flatten_all() - .map_err(|e| MLError::TensorCreationError { - operation: "gradient flatten".to_owned(), - reason: e.to_string(), - })? - .to_vec1::() - .map_err(|e| MLError::TensorCreationError { - operation: "gradient to_vec1".to_owned(), - reason: e.to_string(), - })?; - - let grad_norm: f64 = grad_vec.iter().map(|&g| (g as f64).powi(2)).sum::().sqrt(); - - // Store gradient with descriptive key - let key = format!("varmap_param_{}", idx); - self.gradients.insert(key.clone(), grad.clone()); - - if grad_norm > 1e-12 { - params_with_grads += 1; - total_grad_norm += grad_norm; - } - - trace!("[P0 FIX] VarMap param {}: grad_norm={:.6}", idx, grad_norm); - } else { - trace!( - "[P0 FIX] VarMap param {} has no gradient (not in computational graph)", - idx - ); - } - } - - trace!( - "[P0 FIX] Extracted {} gradients from {} VarMap parameters, total_grad_norm={:.6}", - params_with_grads, - all_vars.len(), - total_grad_norm - ); - - // Verify we got non-zero gradients - if total_grad_norm < 1e-12 { - return Err(MLError::TrainingError(format!( - "Zero gradients extracted from VarMap (total_grad_norm={:.6}). \ - This indicates the loss is not connected to trainable parameters.", - total_grad_norm - ))); - } - - self.clip_gradients(self.config.grad_clip)?; - - // Gradients flow through the trainable VarMap parameters: - // 1. input_projection: Projects d_model → d_inner - // 2. output_projection: Projects d_inner → 1 (regression) - // 3. layer_norms: Normalization weights/biases for each layer - // - // SSM matrices (A, B, C, delta) are NOT trainable in standard MAMBA-2. - // They are part of the model state and are used for selective state-space computation. - - Ok(()) - } - - /// Initialize optimizer state - pub fn initialize_optimizer(&mut self) -> Result<(), MLError> { - // Initialize Adam optimizer state - self.optimizer_state.clear(); - - // Add momentum and variance terms for each parameter - // In real implementation, this would be handled by candle's optimizers - - Ok(()) - } - - /// Zero gradients - pub fn zero_gradients(&mut self) -> Result<(), MLError> { - // Clear all gradients for SSM parameters - for _ssm_state in &mut self.state.ssm_states { - // Zero gradients for A, B, C matrices and delta parameter - if let Some(grad) = self.gradients.get("A").cloned() { - self.gradients.insert("A".to_owned(), grad.zeros_like()?); - } - if let Some(grad) = self.gradients.get("B").cloned() { - self.gradients.insert("B".to_owned(), grad.zeros_like()?); - } - if let Some(grad) = self.gradients.get("C").cloned() { - self.gradients.insert("C".to_owned(), grad.zeros_like()?); - } - if let Some(grad) = self.gradients.get("delta").cloned() { - self.gradients - .insert("delta".to_owned(), grad.zeros_like()?); - } - } - - // Clear optimizer state gradients if they exist - for (param_name, tensor) in self.optimizer_state.iter_mut() { - if param_name.contains("grad") { - *tensor = tensor.zeros_like()?; - } - } - - Ok(()) - } - - /// Optimizer step - dispatches to Adam or SGD based on config - pub fn optimizer_step(&mut self) -> Result<(), MLError> { - match self.config.optimizer_type { - OptimizerType::Adam => self.optimizer_step_adam(), - OptimizerType::AdamW => self.optimizer_step_adam(), - OptimizerType::SGD => self.optimizer_step_sgd(), - } - } - - /// Adam optimizer step implementation - fn optimizer_step_adam(&mut self) -> Result<(), MLError> { - // P0: Use beta1 from config for hyperparameter optimization - let beta1: f64 = self.config.adam_beta1; - let beta2: f64 = 0.999; - let eps: f64 = 1e-8; // Standard epsilon for Adam optimizer - let lr = self.config.learning_rate; - - // Increment step counter for bias correction - let step = self - .optimizer_state - .get("step") - .and_then(|t| t.to_scalar::().ok()) - .unwrap_or(0.0) as f64 - + 1.0; - - let device = self.device(); - let step_tensor = Tensor::new(&[step as f32], device)?; - self.optimizer_state.insert("step".to_owned(), step_tensor); - - // Bias correction uses Rust-side f64 for precision - let beta1_t = beta1.powf(step); - let beta2_t = beta2.powf(step); - let bias_correction1 = 1.0 - beta1_t; - let bias_correction2 = 1.0 - beta2_t; - - // Apply Adam updates to all SSM parameters per layer - let num_layers = self.state.ssm_states.len(); - for layer_idx in 0..num_layers { - // Collect layer-specific gradients - let a_grad = self.gradients.get(&format!("A_{}", layer_idx)).cloned(); - let b_grad = self.gradients.get(&format!("B_{}", layer_idx)).cloned(); - let c_grad = self.gradients.get(&format!("C_{}", layer_idx)).cloned(); - let delta_grad = self.gradients.get(&format!("delta_{}", layer_idx)).cloned(); - - // Update A matrix (state transition matrix) - if let Some(ref A_grad) = a_grad { - trace!("[Agent 225] Updating A matrix for layer {}", layer_idx); - let mut A_param = self.state.ssm_states[layer_idx].A.clone(); - self.apply_adam_update( - &mut A_param, - A_grad, - layer_idx, - "A", - lr, - beta1, - beta2, - eps, - bias_correction1, - bias_correction2, - false, // No weight decay for A matrix (maintains stability) - )?; - self.state.ssm_states[layer_idx].A = A_param; - } - - // Update B matrix (input matrix) - if let Some(ref B_grad) = b_grad { - trace!("[Agent 225] Updating B matrix for layer {}", layer_idx); - let mut B_param = self.state.ssm_states[layer_idx].B.clone(); - self.apply_adam_update( - &mut B_param, - B_grad, - layer_idx, - "B", - lr, - beta1, - beta2, - eps, - bias_correction1, - bias_correction2, - true, // Apply weight decay to B matrix - )?; - self.state.ssm_states[layer_idx].B = B_param; - } - - // Update C matrix (output matrix) - if let Some(ref C_grad) = c_grad { - let mut C_param = self.state.ssm_states[layer_idx].C.clone(); - self.apply_adam_update( - &mut C_param, - C_grad, - layer_idx, - "C", - lr, - beta1, - beta2, - eps, - bias_correction1, - bias_correction2, - true, // Apply weight decay to C matrix - )?; - self.state.ssm_states[layer_idx].C = C_param; - } - - // Update Delta parameter (discretization parameter) - if let Some(ref delta_grad) = delta_grad { - let mut delta_param = self.state.ssm_states[layer_idx].delta.clone(); - self.apply_adam_update( - &mut delta_param, - delta_grad, - layer_idx, - "delta", - lr, - beta1, - beta2, - eps, - bias_correction1, - bias_correction2, - false, // No weight decay for Delta (maintains discretization stability) - )?; - self.state.ssm_states[layer_idx].delta = delta_param; - } - } - - // After updating A matrices, project to maintain spectral radius < 1 - self.project_ssm_matrices()?; - - Ok(()) - } - - /// AdamW optimizer step implementation with decoupled weight decay - /// - /// CRITICAL DIFFERENCE from Adam: - /// - Adam: weight_decay applied to gradients → interferes with SSM dynamics - /// - AdamW: weight_decay applied directly to parameters → preserves SSM constraints - /// - /// AdamW Formula: - /// 1. m_t = β1 * m_{t-1} + (1 - β1) * g_t - /// 2. v_t = β2 * v_{t-1} + (1 - β2) * g_t^2 - /// 3. m_hat = m_t / (1 - β1^t) - /// 4. v_hat = v_t / (1 - β2^t) - /// 5. θ_t = θ_{t-1} * (1 - λ * lr) - lr * m_hat / (√v_hat + ε) - /// ^^^^^^^^^^^^^^^^^^^^^^^^ ← DECOUPLED weight decay - /// - /// Where λ is weight_decay coefficient (independent of gradients) - fn optimizer_step_adamw(&mut self) -> Result<(), MLError> { - let beta1: f64 = self.config.adam_beta1; - let beta2: f64 = self.config.adam_beta2; - let eps: f64 = self.config.adam_epsilon; - let lr = self.config.learning_rate; - let wd = self.config.weight_decay; - - // Increment step counter for bias correction - let step = self - .optimizer_state - .get("step") - .and_then(|t| t.to_scalar::().ok()) - .unwrap_or(0.0) as f64 - + 1.0; - - let device = self.device(); - let step_tensor = Tensor::new(&[step as f32], device)?; - self.optimizer_state.insert("step".to_owned(), step_tensor); - - // Bias correction factors (Rust-side f64 for precision) - let beta1_t = beta1.powf(step); - let beta2_t = beta2.powf(step); - let bias_correction1 = 1.0 - beta1_t; - let bias_correction2 = 1.0 - beta2_t; - - // Apply AdamW updates to all SSM parameters per layer - let num_layers = self.state.ssm_states.len(); - for layer_idx in 0..num_layers { - // Collect layer-specific gradients - let a_grad = self.gradients.get(&format!("A_{}", layer_idx)).cloned(); - let b_grad = self.gradients.get(&format!("B_{}", layer_idx)).cloned(); - let c_grad = self.gradients.get(&format!("C_{}", layer_idx)).cloned(); - let delta_grad = self.gradients.get(&format!("delta_{}", layer_idx)).cloned(); - - // Update A matrix (state transition matrix) - if let Some(ref A_grad) = a_grad { - trace!("[AdamW] Updating A matrix for layer {}", layer_idx); - let mut A_param = self.state.ssm_states[layer_idx].A.clone(); - self.apply_adamw_update( - &mut A_param, - A_grad, - layer_idx, - "A", - lr, - beta1, - beta2, - eps, - bias_correction1, - bias_correction2, - 0.0, // No weight decay for A matrix (maintains stability) - )?; - self.state.ssm_states[layer_idx].A = A_param; - } - - // Update B matrix (input matrix) - if let Some(ref B_grad) = b_grad { - trace!("[AdamW] Updating B matrix for layer {}", layer_idx); - let mut B_param = self.state.ssm_states[layer_idx].B.clone(); - self.apply_adamw_update( - &mut B_param, - B_grad, - layer_idx, - "B", - lr, - beta1, - beta2, - eps, - bias_correction1, - bias_correction2, - wd, // Apply weight decay to B matrix - )?; - self.state.ssm_states[layer_idx].B = B_param; - } - - // Update C matrix (output matrix) - if let Some(ref C_grad) = c_grad { - let mut C_param = self.state.ssm_states[layer_idx].C.clone(); - self.apply_adamw_update( - &mut C_param, - C_grad, - layer_idx, - "C", - lr, - beta1, - beta2, - eps, - bias_correction1, - bias_correction2, - wd, // Apply weight decay to C matrix - )?; - self.state.ssm_states[layer_idx].C = C_param; - } - - // Update Delta parameter (discretization parameter) - if let Some(ref delta_grad) = delta_grad { - let mut delta_param = self.state.ssm_states[layer_idx].delta.clone(); - self.apply_adamw_update( - &mut delta_param, - delta_grad, - layer_idx, - "delta", - lr, - beta1, - beta2, - eps, - bias_correction1, - bias_correction2, - 0.0, // No weight decay for Delta (maintains discretization stability) - )?; - self.state.ssm_states[layer_idx].delta = delta_param; - } - } - - // After updating A matrices, project to maintain spectral radius < 1 - self.project_ssm_matrices()?; - - Ok(()) - } - - /// SGD optimizer step implementation with momentum - fn optimizer_step_sgd(&mut self) -> Result<(), MLError> { - let lr = self.config.learning_rate; - let momentum = self.config.sgd_momentum; - - // PRIORITY 2 FIX (Agent 225): Use layer-specific gradient keys - // Apply SGD updates to all SSM parameters per layer - let num_layers = self.state.ssm_states.len(); - for layer_idx in 0..num_layers { - // Collect layer-specific gradients - let a_grad = self.gradients.get(&format!("A_{}", layer_idx)).cloned(); - let b_grad = self.gradients.get(&format!("B_{}", layer_idx)).cloned(); - let c_grad = self.gradients.get(&format!("C_{}", layer_idx)).cloned(); - let delta_grad = self.gradients.get(&format!("delta_{}", layer_idx)).cloned(); - - // Update A matrix (state transition matrix) - if let Some(ref A_grad) = a_grad { - trace!("[SGD] Updating A matrix for layer {}", layer_idx); - let mut A_param = self.state.ssm_states[layer_idx].A.clone(); - self.apply_sgd_update( - &mut A_param, - A_grad, - layer_idx, - "A", - lr, - momentum, - false, // No weight decay for A matrix (maintains stability) - )?; - self.state.ssm_states[layer_idx].A = A_param; - } - - // Update B matrix (input matrix) - if let Some(ref B_grad) = b_grad { - trace!("[SGD] Updating B matrix for layer {}", layer_idx); - let mut B_param = self.state.ssm_states[layer_idx].B.clone(); - self.apply_sgd_update( - &mut B_param, - B_grad, - layer_idx, - "B", - lr, - momentum, - true, // Apply weight decay to B matrix - )?; - self.state.ssm_states[layer_idx].B = B_param; - } - - // Update C matrix (output matrix) - if let Some(ref C_grad) = c_grad { - let mut C_param = self.state.ssm_states[layer_idx].C.clone(); - self.apply_sgd_update( - &mut C_param, - C_grad, - layer_idx, - "C", - lr, - momentum, - true, // Apply weight decay to C matrix - )?; - self.state.ssm_states[layer_idx].C = C_param; - } - - // Update Delta parameter (discretization parameter) - if let Some(ref delta_grad) = delta_grad { - let mut delta_param = self.state.ssm_states[layer_idx].delta.clone(); - self.apply_sgd_update( - &mut delta_param, - delta_grad, - layer_idx, - "delta", - lr, - momentum, - false, // No weight decay for Delta (maintains discretization stability) - )?; - self.state.ssm_states[layer_idx].delta = delta_param; - } - } - - // After updating A matrices, project to maintain spectral radius < 1 - self.project_ssm_matrices()?; - - Ok(()) - } - - /// Update learning rate with warmup and decay - fn update_learning_rate(&mut self, epoch: usize, batch_idx: usize) -> Result<(), MLError> { - // FIXED (Agent P2): Use actual training data length instead of hardcoded 1000 - // Calculate total steps based on actual data size - let batches_per_epoch = if self.total_training_samples > 0 { - self.total_training_samples / self.config.batch_size - } else { - // Fallback to reasonable default if not set - 1000 / self.config.batch_size - }; - - let total_steps = epoch * batches_per_epoch + (batch_idx / self.config.batch_size); - - // FIXED (Agent P2): Remove underscore prefix - we DO use this value - let lr = if total_steps < self.config.warmup_steps { - // Linear warmup: LR increases from 0 to configured LR - self.config.learning_rate * (total_steps as f64 / self.config.warmup_steps as f64) - } else { - // Cosine decay after warmup (P0 FIX: use config value, not hardcoded) - let progress = (total_steps - self.config.warmup_steps) as f64; - let total_decay_steps = self.config.total_decay_steps as f64; - let decay_ratio = (progress / total_decay_steps).min(1.0); - self.config.learning_rate * 0.5 * (1.0 + (std::f64::consts::PI * decay_ratio).cos()) - }; - - // FIXED (Agent P2): Actually apply the computed learning rate - self.current_lr = lr; - - Ok(()) - } - - /// Get current learning rate - fn get_current_learning_rate(&self) -> f64 { - // FIXED (Agent P2): Return actual current LR, not constant config value - self.current_lr - } - - /// Validate model on validation set - fn validate(&mut self, val_data: &[(Tensor, Tensor)]) -> Result { - let mut total_loss = 0.0; - let mut count = 0; - - // Disable dropout for validation - for (input, target) in val_data { - // FIXED: Ensure input and target tensors are on the model's device (GPU) - let input = input.to_device(&self.device)?; - let target = target.to_device(&self.device)?; - - let output = self.forward(&input)?; - // FIXED (Agent 217): Extract last timestep for validation loss (same as training) - let seq_len = output.dim(1)?; - let output_last = output.narrow(1, seq_len - 1, 1)?; - let loss = self.compute_loss(&output_last, &target)?; - total_loss += loss.to_scalar::()? as f64; - count += 1; - - if count >= 100 { - // Limit validation set size for speed - break; - } - } - - Ok(total_loss / count as f64) - } - - /// Calculate accuracy metric - fn calculate_accuracy(&mut self, val_data: &[(Tensor, Tensor)]) -> Result { - if val_data.is_empty() { - return Ok(0.0); - } - - let mut correct = 0; - let mut total = 0; - - for (input, target) in val_data { - // CRITICAL FIX: Transfer tensors to device before forward pass (matches validate()) - let input = input.to_device(&self.device)?; - let target = target.to_device(&self.device)?; - - let output = self.forward(&input)?; - let seq_len = output.dims()[1]; - - // Extract last timestep predictions - let output_last = output.narrow(1, seq_len - 1, 1)?.squeeze(1)?; - - // FIX: Reshape target from [batch_size, 1, 1] to [batch_size] - // Previous double squeeze failed due to dimension index shifting - let batch_size = target.dim(0)?; - let target_squeezed = target.reshape(&[batch_size])?; - - // FIX: Use element-wise comparison instead of mean_all() - for i in 0..batch_size { - // FIX: .get(i) returns different shapes depending on input: - // - If input is [N], .get(i) returns scalar [] - // - If input is [N, 1], .get(i) returns [1] - // Check rank and squeeze conditionally - let pred_tensor = output_last.get(i)?; - let pred_value = if pred_tensor.rank() == 0 { - pred_tensor.to_scalar::()? as f64 - } else { - pred_tensor.squeeze(0)?.to_scalar::()? as f64 - }; - - let target_tensor = target_squeezed.get(i)?; - let target_value = if target_tensor.rank() == 0 { - target_tensor.to_scalar::()? as f64 - } else { - target_tensor.squeeze(0)?.to_scalar::()? as f64 - }; - - // FIX: Use absolute error (not MAPE) with 5% threshold - let abs_error = (pred_value - target_value).abs(); - - // 5% of [0,1] range = 0.05 (equivalent to ~$50 in ES price space) - if abs_error < 0.05 { - correct += 1; - } - total += 1; - } - - if total >= 100 { - break; - } - } - - Ok(correct as f64 / total as f64) - } - - /// Save model checkpoint - pub async fn save_checkpoint(&mut self, path: &str) -> Result<(), MLError> { - use std::collections::HashMap as StdHashMap; - - info!("Saving MAMBA-2 checkpoint to {}", path); - - // Update metadata - self.metadata.last_checkpoint = Some(path.to_string()); - self.metadata.performance_stats = self.get_performance_metrics(); - - // AGENT F2: CRITICAL FIX - Actually save model weights to disk using safetensors - // This replaces the stub implementation that only logged without saving - - // Add .safetensors extension if not present - let safetensors_path = if path.ends_with(".safetensors") || path.ends_with(".ckpt") { - if path.ends_with(".ckpt") { - path.replace(".ckpt", ".safetensors") - } else { - path.to_string() - } - } else { - format!("{}.safetensors", path) - }; - - // Extract all tensors from VarMap - let vars_data = self.varmap.data().lock().map_err(|e| { - MLError::LockError(format!("Failed to lock VarMap for checkpoint: {}", e)) - })?; - - // Build tensor map for safetensors serialization - let mut tensors: StdHashMap = StdHashMap::new(); - for (name, var) in vars_data.iter() { - tensors.insert(name.clone(), var.as_tensor().clone()); - } - - // Save using safetensors format (thread-safe serialization) - candle_core::safetensors::save(&tensors, &safetensors_path) - .map_err(|e| MLError::CheckpointError(format!("Failed to save safetensors: {}", e)))?; - - // Verify checkpoint was saved successfully - let metadata = std::fs::metadata(&safetensors_path).map_err(|e| { - MLError::CheckpointError(format!("Checkpoint verification failed: {}", e)) - })?; - - let file_size_mb = metadata.len() as f64 / (1024.0 * 1024.0); - - info!( - "✓ MAMBA-2 checkpoint saved successfully: {} ({:.2} MB, {} parameters)", - safetensors_path, file_size_mb, self.metadata.num_parameters - ); - - // Validate checkpoint size is reasonable (>1MB for non-trivial models) - if file_size_mb < 0.1 { - warn!( - "⚠️ Checkpoint file size is suspiciously small ({:.2} MB) - may indicate incomplete save", - file_size_mb - ); - } - - Ok(()) - } - - /// Load model checkpoint - pub async fn load_checkpoint(&mut self, path: &str) -> Result<(), MLError> { - info!("Loading MAMBA-2 checkpoint from {}", path); - - // AGENT F2: CRITICAL FIX - Actually load model weights from disk - // This replaces the stub implementation that only set flags without loading - - // Add .safetensors extension if not present - let safetensors_path = if path.ends_with(".safetensors") || path.ends_with(".ckpt") { - if path.ends_with(".ckpt") { - path.replace(".ckpt", ".safetensors") - } else { - path.to_string() - } - } else { - format!("{}.safetensors", path) - }; - - // Verify checkpoint file exists - if !std::path::Path::new(&safetensors_path).exists() { - return Err(MLError::CheckpointError(format!( - "Checkpoint file not found: {}", - safetensors_path - ))); - } - - // Load tensors from safetensors - let tensors = candle_core::safetensors::load(&safetensors_path, &self.device) - .map_err(|e| MLError::CheckpointError(format!("Failed to load safetensors: {}", e)))?; - - // Populate VarMap with loaded tensors - let mut vars_data = self.varmap.data().lock().map_err(|e| { - MLError::LockError(format!("Failed to lock VarMap for checkpoint load: {}", e)) - })?; - - for (name, tensor) in tensors.iter() { - // Create new Var from loaded tensor - let var = Var::from_tensor(tensor)?; - vars_data.insert(name.clone(), var); - } - - self.is_trained = true; - self.metadata.last_checkpoint = Some(path.to_string()); - - info!( - "✓ MAMBA-2 checkpoint loaded successfully: {} ({} tensors)", - safetensors_path, - tensors.len() - ); - - Ok(()) - } - - /// Apply gradient clipping to prevent exploding gradients - /// - /// SSM gradients (A, B, C, delta) are not directly trainable in standard MAMBA-2; - /// they flow through the VarMap and are clipped by AdamW weight_decay. - /// The previous implementation computed per-parameter norms (4N GPU syncs) - /// but discarded the clipped results. Trainable parameter gradients are - /// handled by the optimizer step. - fn clip_gradients(&mut self, _max_norm: f64) -> Result<(), MLError> { - Ok(()) - } - - /// Apply Adam optimizer update to a single parameter - fn apply_adam_update( - &mut self, - param: &mut Tensor, - grad: &Tensor, - layer_idx: usize, - param_name: &str, - lr: f64, - beta1: f64, - beta2: f64, - eps: f64, - bias_correction1: f64, - bias_correction2: f64, - apply_weight_decay: bool, - ) -> Result<(), MLError> { - // Create unique keys for momentum and variance - let m_key = format!( - "layer_{}_{}_{}_m", - layer_idx, - param_name, - param.dims().len() - ); - let v_key = format!( - "layer_{}_{}_{}_v", - layer_idx, - param_name, - param.dims().len() - ); - - // Initialize momentum and variance if not present - if !self.optimizer_state.contains_key(&m_key) { - let m_init = grad.zeros_like()?; - let v_init = grad.zeros_like()?; - self.optimizer_state.insert(m_key.clone(), m_init); - self.optimizer_state.insert(v_key.clone(), v_init); - } - - // Get momentum and variance tensors separately to avoid double borrow - let m_tensor = self - .optimizer_state - .get(&m_key) - .ok_or_else(|| { - MLError::ModelError(format!("Missing momentum tensor for key: {}", m_key)) - })? - .clone(); - let v_tensor = self - .optimizer_state - .get(&v_key) - .ok_or_else(|| { - MLError::ModelError(format!("Missing variance tensor for key: {}", v_key)) - })? - .clone(); - - // Apply weight decay if specified - // REFACTORED (Agent 234): Use scalar_tensor helper to eliminate dtype boilerplate - let device = self.device(); - let dtype = param.dtype(); - let effective_grad = if apply_weight_decay && self.config.weight_decay > 0.0 { - let weight_decay_scalar = Self::scalar_tensor(self.config.weight_decay, dtype, device)?; - let weight_decay_term = param.broadcast_mul(&weight_decay_scalar)?; - grad.add(&weight_decay_term)? - } else { - grad.clone() - }; - - // Update biased first moment estimate: m_t = β1 * m_{t-1} + (1 - β1) * g_t - // REFACTORED (Agent 234): Use scalar_tensor helper (was 87 lines of boilerplate) - let beta1_scalar = Self::scalar_tensor(beta1, dtype, device)?; - let m_scaled = m_tensor.broadcast_mul(&beta1_scalar)?; - let grad_scalar = Self::scalar_tensor(1.0 - beta1, dtype, device)?; - let grad_scaled = effective_grad.broadcast_mul(&grad_scalar)?; - let new_m = m_scaled.add(&grad_scaled)?; - - // Update biased second moment estimate: v_t = β2 * v_{t-1} + (1 - β2) * g_t^2 - let grad_squared = effective_grad.mul(&effective_grad)?; - let beta2_scalar = Self::scalar_tensor(beta2, dtype, device)?; - let v_scaled = v_tensor.broadcast_mul(&beta2_scalar)?; - let grad_squared_scalar = Self::scalar_tensor(1.0 - beta2, dtype, device)?; - let grad_squared_scaled = grad_squared.broadcast_mul(&grad_squared_scalar)?; - let new_v = v_scaled.add(&grad_squared_scaled)?; - - // Compute bias-corrected estimates - let bias_corr1_scalar = Self::scalar_tensor(1.0 / bias_correction1, dtype, device)?; - let m_hat = new_m.broadcast_mul(&bias_corr1_scalar)?; - let bias_corr2_scalar = Self::scalar_tensor(1.0 / bias_correction2, dtype, device)?; - let v_hat = new_v.broadcast_mul(&bias_corr2_scalar)?; - - // Compute parameter update: θ = θ - lr * m_hat / (√(v_hat) + ε) - let sqrt_v_hat = v_hat.sqrt()?; - let eps_scalar = Self::scalar_tensor(eps, dtype, device)?; - let denominator = sqrt_v_hat.broadcast_add(&eps_scalar)?; - let lr_scalar = Self::scalar_tensor(lr, dtype, device)?; - let update = m_hat.div(&denominator)?.broadcast_mul(&lr_scalar)?; - - // Update parameter: θ_{t+1} = θ_t - update - *param = param.sub(&update)?; - - // Store updated momentum and variance back - self.optimizer_state.insert(m_key, new_m); - self.optimizer_state.insert(v_key, new_v); - - Ok(()) - } - - /// Apply SGD optimizer update with momentum to a single parameter - /// Apply AdamW optimizer update to a single parameter with decoupled weight decay - /// - /// CRITICAL: Weight decay is applied DIRECTLY to parameters, NOT to gradients. - /// This prevents interference with SSM spectral radius constraints. - /// - /// AdamW Update Formula: - /// 1. m_t = β1 * m_{t-1} + (1 - β1) * g_t (WITHOUT weight decay in gradient) - /// 2. v_t = β2 * v_{t-1} + (1 - β2) * g_t^2 - /// 3. m_hat = m_t / (1 - β1^t), v_hat = v_t / (1 - β2^t) - /// 4. θ_t = θ_{t-1} * (1 - λ * lr) - lr * m_hat / (√v_hat + ε) - /// ^^^^^^^^^^^^^^^^^^^^^^^^ ← Weight decay applied to parameter - fn apply_adamw_update( - &mut self, - param: &mut Tensor, - grad: &Tensor, - layer_idx: usize, - param_name: &str, - lr: f64, - beta1: f64, - beta2: f64, - eps: f64, - bias_correction1: f64, - bias_correction2: f64, - weight_decay: f64, - ) -> Result<(), MLError> { - // Create unique keys for momentum and variance - let m_key = format!( - "layer_{}_{}_{}_m", - layer_idx, - param_name, - param.dims().len() - ); - let v_key = format!( - "layer_{}_{}_{}_v", - layer_idx, - param_name, - param.dims().len() - ); - - // Initialize momentum and variance if not present - if !self.optimizer_state.contains_key(&m_key) { - let m_init = grad.zeros_like()?; - let v_init = grad.zeros_like()?; - self.optimizer_state.insert(m_key.clone(), m_init); - self.optimizer_state.insert(v_key.clone(), v_init); - } - - // Get momentum and variance tensors - let m_tensor = self - .optimizer_state - .get(&m_key) - .ok_or_else(|| { - MLError::ModelError(format!("Missing momentum tensor for key: {}", m_key)) - })? - .clone(); - let v_tensor = self - .optimizer_state - .get(&v_key) - .ok_or_else(|| { - MLError::ModelError(format!("Missing variance tensor for key: {}", v_key)) - })? - .clone(); - - let device = self.device(); - let dtype = param.dtype(); - - // CRITICAL: NO weight decay applied to gradient (pure gradient) - // This is the key difference from Adam optimizer - - // Update biased first moment estimate: m_t = β1 * m_{t-1} + (1 - β1) * g_t - let beta1_scalar = Self::scalar_tensor(beta1, dtype, device)?; - let m_scaled = m_tensor.broadcast_mul(&beta1_scalar)?; - let grad_scalar = Self::scalar_tensor(1.0 - beta1, dtype, device)?; - let grad_scaled = grad.broadcast_mul(&grad_scalar)?; - let new_m = m_scaled.add(&grad_scaled)?; - - // Update biased second moment estimate: v_t = β2 * v_{t-1} + (1 - β2) * g_t^2 - let grad_squared = grad.mul(grad)?; - let beta2_scalar = Self::scalar_tensor(beta2, dtype, device)?; - let v_scaled = v_tensor.broadcast_mul(&beta2_scalar)?; - let grad_squared_scalar = Self::scalar_tensor(1.0 - beta2, dtype, device)?; - let grad_squared_scaled = grad_squared.broadcast_mul(&grad_squared_scalar)?; - let new_v = v_scaled.add(&grad_squared_scaled)?; - - // Compute bias-corrected estimates - let bias_corr1_scalar = Self::scalar_tensor(1.0 / bias_correction1, dtype, device)?; - let m_hat = new_m.broadcast_mul(&bias_corr1_scalar)?; - let bias_corr2_scalar = Self::scalar_tensor(1.0 / bias_correction2, dtype, device)?; - let v_hat = new_v.broadcast_mul(&bias_corr2_scalar)?; - - // Compute gradient update: lr * m_hat / (√(v_hat) + ε) - let sqrt_v_hat = v_hat.sqrt()?; - let eps_scalar = Self::scalar_tensor(eps, dtype, device)?; - let denominator = sqrt_v_hat.broadcast_add(&eps_scalar)?; - let lr_scalar = Self::scalar_tensor(lr, dtype, device)?; - let grad_update = m_hat.div(&denominator)?.broadcast_mul(&lr_scalar)?; - - // CRITICAL: Apply decoupled weight decay directly to parameter - // θ_t = θ_{t-1} * (1 - λ * lr) - grad_update - // This is the key innovation of AdamW vs Adam - let updated_param = if weight_decay > 0.0 { - // Apply weight decay: param * (1 - wd * lr) - let decay_factor = 1.0 - weight_decay * lr; - let decay_scalar = Self::scalar_tensor(decay_factor, dtype, device)?; - let decayed_param = param.broadcast_mul(&decay_scalar)?; - // Then subtract gradient update - decayed_param.sub(&grad_update)? - } else { - // No weight decay, just gradient update - param.sub(&grad_update)? - }; - - *param = updated_param; - - // Store updated momentum and variance - self.optimizer_state.insert(m_key, new_m); - self.optimizer_state.insert(v_key, new_v); - - Ok(()) - } - - /// Apply SGD optimizer update with momentum to a single parameter - /// - /// - /// SGD Update Formula: - /// - Momentum: v_t = μ * v_{t-1} + (1 - μ) * g_t - /// - Update: θ_t = θ_{t-1} - lr * v_t - /// - /// Where: - /// - v_t: velocity (momentum state) - /// - μ: momentum coefficient (typically 0.9) - /// - g_t: gradient (with optional weight decay) - /// - lr: learning rate - fn apply_sgd_update( - &mut self, - param: &mut Tensor, - grad: &Tensor, - layer_idx: usize, - param_name: &str, - lr: f64, - momentum: f64, - apply_weight_decay: bool, - ) -> Result<(), MLError> { - // Create unique key for velocity (momentum state) - let v_key = format!( - "layer_{}_{}_{}_velocity", - layer_idx, - param_name, - param.dims().len() - ); - - // Initialize velocity if not present (zeros) - if !self.optimizer_state.contains_key(&v_key) { - let v_init = grad.zeros_like()?; - self.optimizer_state.insert(v_key.clone(), v_init); - } - - // Get velocity tensor - let v_tensor = self - .optimizer_state - .get(&v_key) - .ok_or_else(|| { - MLError::ModelError(format!("Missing velocity tensor for key: {}", v_key)) - })? - .clone(); - - // Apply weight decay if specified (L2 regularization) - let device = self.device(); - let dtype = param.dtype(); - let effective_grad = if apply_weight_decay && self.config.weight_decay > 0.0 { - let weight_decay_scalar = Self::scalar_tensor(self.config.weight_decay, dtype, device)?; - let weight_decay_term = param.broadcast_mul(&weight_decay_scalar)?; - grad.add(&weight_decay_term)? - } else { - grad.clone() - }; - - // Update velocity: v_t = μ * v_{t-1} + (1 - μ) * g_t - let momentum_scalar = Self::scalar_tensor(momentum, dtype, device)?; - let v_scaled = v_tensor.broadcast_mul(&momentum_scalar)?; - let grad_scalar = Self::scalar_tensor(1.0 - momentum, dtype, device)?; - let grad_scaled = effective_grad.broadcast_mul(&grad_scalar)?; - let new_v = v_scaled.add(&grad_scaled)?; - - // Compute parameter update: θ_{t+1} = θ_t - lr * v_t - let lr_scalar = Self::scalar_tensor(lr, dtype, device)?; - let update = new_v.broadcast_mul(&lr_scalar)?; - *param = param.sub(&update)?; - - // Store updated velocity back - self.optimizer_state.insert(v_key, new_v); - - Ok(()) - } - - /// Project SSM matrices to maintain stability - fn project_ssm_matrices(&mut self) -> Result<(), MLError> { - // Avoid borrow checker issues by processing each state separately - for i in 0..self.state.ssm_states.len() { - // Ensure A matrix has spectral radius < 1 for stability - let spectral_radius = { - let ssm_state = &self.state.ssm_states[i]; - self.compute_spectral_radius(&ssm_state.A)? - }; - if spectral_radius >= 1.0 { - let scale_factor = 0.99 / spectral_radius; - let device = self.device(); - let scale_tensor = Tensor::new(&[scale_factor as f32], device)?; - self.state.ssm_states[i].A = - self.state.ssm_states[i].A.broadcast_mul(&scale_tensor)?; - } - - // Ensure Delta parameter stays positive and reasonable - let device = self.device(); - let delta_min = Tensor::new(&[1e-6_f32], device)?; - let delta_max = Tensor::new(&[1.0_f32], device)?; - let delta_clamped = self.state.ssm_states[i] - .delta - .broadcast_maximum(&delta_min)? - .broadcast_minimum(&delta_max)?; - self.state.ssm_states[i].delta = delta_clamped; - } - - Ok(()) - } - - /// Compute spectral radius (largest eigenvalue magnitude) of a matrix - fn compute_spectral_radius(&self, matrix: &Tensor) -> Result { - // For simplicity, use Frobenius norm as approximation - // In production, we'd compute actual eigenvalues - let frobenius_norm = matrix.powf(2.0)?.sum_all()?.to_scalar::()? as f64; - let frobenius_norm = frobenius_norm.sqrt(); - - // Frobenius norm upper bounds spectral radius - // For better approximation, we scale by sqrt of matrix size - let dims = matrix.dims(); - if dims.len() >= 2 { - let size = (dims[0].min(dims[1]) as f64).sqrt(); - Ok(frobenius_norm / size) - } else { - Ok(frobenius_norm) - } - } -} - -#[cfg(test)] -mod tests { - use super::*; - use anyhow::Result; - - #[tokio::test] - async fn test_mamba_creation() -> Result<()> { - let config = Mamba2Config { - d_model: 8, - d_state: 4, - d_head: 4, - num_heads: 2, - ..Default::default() - }; - - let device = Device::Cpu; - let model = Mamba2SSM::new(config, &device) - .map_err(|_| anyhow::anyhow!("Failed to create MAMBA model"))?; - assert_eq!(model.metadata.input_dim, 8); - assert_eq!(model.metadata.output_dim, 1); - Ok(()) - } - - #[test] - fn test_mamba_config_default() -> Result<()> { - let config = Mamba2Config::default(); - assert!(config.d_model > 0); - assert!(config.d_state > 0); - assert!(config.num_heads > 0); - Ok(()) - } - - #[test] - fn test_mamba_state_creation() -> Result<()> { - let config = Mamba2Config { - d_model: 4, - d_state: 2, - d_head: 2, - num_heads: 2, - ..Default::default() - }; - - let device = Device::Cpu; - let state = Mamba2State::zeros(&config, &device) - .map_err(|_| anyhow::anyhow!("Failed to create MAMBA state"))?; - assert_eq!(state.ssm_states.len(), config.num_layers); - assert!(!state.selective_state.is_empty()); - Ok(()) - } - - #[test] - fn test_mamba_performance_metrics() -> Result<()> { - let config = Mamba2Config { - d_model: 4, - target_latency_us: 5, - ..Default::default() - }; - - let device = Device::Cpu; - let model = Mamba2SSM::new(config, &device) - .map_err(|_| anyhow::anyhow!("Failed to create MAMBA model"))?; - let metrics = model.get_performance_metrics(); - - assert!(metrics.contains_key("total_inferences")); - assert!(metrics.contains_key("model_parameters")); - assert!(metrics.contains_key("compression_ratio")); - Ok(()) - } - - #[test] - fn test_mamba_learning_rate_schedule() -> Result<()> { - let config = Mamba2Config { - d_model: 4, - num_layers: 1, - learning_rate: 0.001, - warmup_steps: 10, - batch_size: 2, - ..Default::default() - }; - - let device = Device::Cpu; - let mut model = Mamba2SSM::new(config.clone(), &device) - .map_err(|_| anyhow::anyhow!("Failed to create MAMBA model"))?; - - // Set total training samples (simulate 100 samples) - model.total_training_samples = 100; - - // Test warmup phase (steps 0-9) - for step in 0..10 { - let epoch = step / (model.total_training_samples / model.config.batch_size); - let batch_idx = (step % (model.total_training_samples / model.config.batch_size)) - * model.config.batch_size; - - model - .update_learning_rate(epoch, batch_idx) - .map_err(|_| anyhow::anyhow!("Failed to update LR"))?; - - let current_lr = model.get_current_learning_rate(); - let expected_lr = config.learning_rate * (step as f64 / config.warmup_steps as f64); - - // Allow small floating point error - assert!( - (current_lr - expected_lr).abs() < 1e-9, - "Warmup step {}: expected {}, got {}", - step, - expected_lr, - current_lr - ); - } - - // Test decay phase (after warmup) - let step = 15; - let epoch = step / (model.total_training_samples / model.config.batch_size); - let batch_idx = (step % (model.total_training_samples / model.config.batch_size)) - * model.config.batch_size; - - model - .update_learning_rate(epoch, batch_idx) - .map_err(|_| anyhow::anyhow!("Failed to update LR"))?; - - let decay_lr = model.get_current_learning_rate(); - // During decay, LR should be less than initial LR but greater than 0 - assert!( - decay_lr < config.learning_rate && decay_lr > 0.0, - "Decay phase: LR should be in (0, {}), got {}", - config.learning_rate, - decay_lr - ); - - Ok(()) - } - - #[test] - fn test_mamba_hft_config() -> Result<()> { - let device = Device::Cpu; - let model = Mamba2SSM::default_hft(&device) - .map_err(|_| anyhow::anyhow!("Failed to create HFT MAMBA model"))?; - assert_eq!(model.config.target_latency_us, 3); - assert!(model.config.hardware_aware); - assert!(model.config.use_ssd); - assert!(model.config.use_selective_state); - Ok(()) - } - - #[test] - fn test_mamba_shuffle_batches_deterministic() -> Result<()> { - // Test that with shuffle_batches=false, batch order is deterministic - let config = Mamba2Config { - d_model: 4, - d_state: 2, - batch_size: 2, - seq_len: 4, - shuffle_batches: false, - ..Default::default() - }; - - let device = Device::Cpu; - let model = Mamba2SSM::new(config.clone(), &device) - .map_err(|_| anyhow::anyhow!("Failed to create MAMBA model"))?; - - // Verify shuffle is disabled - assert!(!model.config.shuffle_batches); - - // Create a simple data sequence - let data_len = 10; - let batch_indices: Vec = (0..data_len).step_by(config.batch_size).collect(); - - // Verify deterministic order (should be [0, 2, 4, 6, 8]) - assert_eq!(batch_indices, vec![0, 2, 4, 6, 8]); - - Ok(()) - } - - #[test] - fn test_mamba_shuffle_batches_enabled() -> Result<()> { - // Test that with shuffle_batches=true, batches can be in different order - let config = Mamba2Config { - d_model: 4, - d_state: 2, - batch_size: 2, - seq_len: 4, - shuffle_batches: true, - ..Default::default() - }; - - let device = Device::Cpu; - let model = Mamba2SSM::new(config.clone(), &device) - .map_err(|_| anyhow::anyhow!("Failed to create MAMBA model"))?; - - // Verify shuffle is enabled - assert!(model.config.shuffle_batches); - - // Test that shuffling actually works - use rand::seq::SliceRandom; - let mut batch_indices: Vec = (0..10).step_by(2).collect(); - let _original = batch_indices.clone(); - - batch_indices.shuffle(&mut rand::thread_rng()); - - // Note: There's a small chance this could fail if shuffle happens to - // produce the same order, but probability is low (1/5! = 1/120) - // For a proper test, we'd need a deterministic RNG with a seed - - Ok(()) - } -} - -#[test] -fn test_mamba_parameter_count() -> anyhow::Result<()> { - let config = Mamba2Config { - d_model: 8, - num_layers: 2, - ..Default::default() - }; - - let param_count = Mamba2SSM::count_parameters(&config); - assert!(param_count > 0); - Ok(()) -} - -#[test] -fn test_bilinear_discretization_more_accurate_than_zoh() { - let zoh = 1.0 + (-1.0) * 0.1; - let bilinear = 1.0 + (-1.0) * 0.1 + ((-1.0) * 0.1_f64).powi(2) / 2.0; - let exact = (-0.1_f64).exp(); - - assert!((bilinear - exact).abs() < (zoh - exact).abs(), - "bilinear {} should be closer to exact {} than zoh {}", - bilinear, exact, zoh); -} diff --git a/crates/ml/src/mamba/trainable_adapter.rs b/crates/ml/src/mamba/trainable_adapter.rs index 110bb35f3..e03dd189c 100644 --- a/crates/ml/src/mamba/trainable_adapter.rs +++ b/crates/ml/src/mamba/trainable_adapter.rs @@ -334,25 +334,6 @@ impl UnifiedTrainable for Mamba2SSM { } } -impl Clone for Mamba2SSM { - /// Clone implementation for checkpoint saving - /// - /// Note: This is a shallow clone that copies configuration and metadata, - /// but shares tensor references. Use for checkpoint operations only. - fn clone(&self) -> Self { - // Create a new model with same configuration - // Note: This is a simplified clone for checkpoint operations - // Full deep cloning of all tensors would be expensive - match Mamba2SSM::new(self.config.clone(), &self.device) { - Ok(model) => model, - Err(e) => { - tracing::error!("Mamba2SSM clone failed: {}", e); - std::process::abort(); - } - } - } -} - #[cfg(test)] mod tests { use super::super::Mamba2Config; diff --git a/crates/ml/src/tft/mod.rs b/crates/ml/src/tft/mod.rs index 04c9903d6..48d6b8af0 100644 --- a/crates/ml/src/tft/mod.rs +++ b/crates/ml/src/tft/mod.rs @@ -1,1525 +1,14 @@ -//! # Temporal Fusion Transformer (TFT) for HFT +//! Temporal Fusion Transformer (TFT) for HFT //! -//! State-of-the-art multi-horizon forecasting with variable selection networks, -//! temporal self-attention, gated residual networks, and uncertainty quantification. -//! -//! ## Key Features -//! -//! - Multi-horizon forecasting (1-tick to 100-tick ahead) -//! - Variable selection networks for feature importance -//! - Gated residual networks for improved gradient flow -//! - Quantile outputs for uncertainty estimation -//! - Temporal self-attention for sequential modeling -//! - Sub-50μs inference latency optimized for HFT -//! -//! ## Performance Targets -//! -//! - Inference: <50μs per prediction -//! - Accuracy improvement: +15% over baseline -//! - Memory usage: <1GB -//! - Throughput: >100K predictions/sec +//! This module re-exports the `ml-supervised` TFT implementation and keeps bridge +//! modules that depend on types from both `ml-supervised` and the parent `ml` crate. -use std::collections::HashMap; -use std::num::NonZeroUsize; -use std::sync::atomic::{AtomicU64, Ordering}; -use std::sync::Arc; -use std::time::{Instant, SystemTime}; +// Re-export everything from the ml-supervised tft module +pub use ml_supervised::tft::*; -use async_trait::async_trait; -use candle_core::{Device, Module, Tensor}; -use candle_nn::{linear, AdamW, Linear, Optimizer, ParamsAdamW, VarBuilder, VarMap}; - -use crate::dqn::mixed_precision::training_dtype; -use lru::LruCache; -use ndarray::{Array1, Array2}; -use serde::{Deserialize, Serialize}; -use serde_json::Value; -use tracing::{debug, info, instrument, warn}; -use uuid::Uuid; - -use crate::checkpoint::Checkpointable; -use crate::{MLError, ModelType}; - -// Import TFT components -pub mod gated_residual; -pub mod hft_optimizations; -pub mod lstm_encoder; -pub mod qat_tft; // Quantization-Aware Training wrapper - RE-ENABLED: Device mismatch fix applied -pub mod quantile_outputs; -pub mod quantized_attention; // Re-enabled Wave 9.12 -pub mod quantized_grn; -pub mod quantized_lstm; -pub mod quantized_tft; // Re-enabled Wave 9.12 -pub mod quantized_vsn; -pub mod temporal_attention; +// Bridge modules that depend on ml-internal types (UnifiedTrainable, Checkpointable) pub mod trainable_adapter; pub mod training; -pub mod variable_selection; -pub mod varmap_quantization; -// Public exports for TFT components -pub use gated_residual::{GRNStack, GatedResidualNetwork}; -pub use lstm_encoder::LSTMEncoder; -pub use qat_tft::QATTemporalFusionTransformer; // Quantization-Aware Training wrapper - RE-ENABLED: Device mismatch fix applied -pub use quantile_outputs::QuantileLayer; -pub use quantized_attention::QuantizedTemporalAttention; // Re-enabled Wave 9.12 -pub use quantized_grn::QuantizedGatedResidualNetwork; -pub use quantized_lstm::QuantizedLSTMEncoder; -pub use quantized_tft::QuantizedTemporalFusionTransformer; // Re-enabled Wave 9.12 -pub use quantized_vsn::QuantizedVariableSelectionNetwork; -pub use temporal_attention::TemporalSelfAttention; +// Re-export bridge types pub use trainable_adapter::TrainableTFT; -pub use variable_selection::VariableSelectionNetwork; -pub use varmap_quantization::{ - load_quantized_weights, quantize_varmap, quantize_varmap_parallel, save_quantized_weights, -}; - -/// `TFT` Configuration -/// TFT model variant selection (F32 vs INT8) -#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] -pub enum TFTVariant { - /// Full precision (F32) model - F32, - /// INT8 quantized model (75% memory reduction) - INT8, -} - -impl Default for TFTVariant { - fn default() -> Self { - Self::F32 - } -} - -impl TFTVariant { - /// Check if variant uses quantization - pub fn is_quantized(&self) -> bool { - matches!(self, Self::INT8) - } - - /// Get expected memory reduction ratio vs F32 - pub fn memory_reduction_ratio(&self) -> f64 { - match self { - Self::F32 => 1.0, - Self::INT8 => 0.25, // 75% reduction → 25% of original - } - } -} - -/// `TFT` Configuration -#[derive(Debug, Clone, Serialize, Deserialize)] -pub struct TFTConfig { - // Model architecture - pub input_dim: usize, - pub hidden_dim: usize, - pub num_heads: usize, - pub num_layers: usize, - - // Forecasting parameters - pub prediction_horizon: usize, - pub sequence_length: usize, - pub num_quantiles: usize, - - // Feature types - pub num_static_features: usize, - pub num_known_features: usize, - pub num_unknown_features: usize, - - // Training parameters - pub learning_rate: f64, - pub batch_size: usize, - pub dropout_rate: f64, - pub l2_regularization: f64, - - // HFT optimization - pub use_flash_attention: bool, - pub mixed_precision: bool, - pub memory_efficient: bool, - - // Performance constraints - pub max_inference_latency_us: u64, - pub target_throughput_pps: u64, -} - -impl Default for TFTConfig { - fn default() -> Self { - Self { - // Wave C+D: 225 features (201 Wave C + 24 Wave D) - // Wave C: 201 features (indices 0-200) - // Wave D: 24 features (indices 201-224) - input_dim: 225, - hidden_dim: 128, - num_heads: 8, - num_layers: 3, - prediction_horizon: 10, - sequence_length: 50, - num_quantiles: 9, - // Feature split for 225 total features: - // - Static: 5 features (symbol metadata) - // - Known: 10 features (future time features) - // - Unknown: 210 features (historical OHLCV + technical + microstructure + regime) - num_static_features: 5, - num_known_features: 10, - num_unknown_features: 210, - learning_rate: 1e-3, - batch_size: 64, - dropout_rate: 0.1, - l2_regularization: 1e-4, - use_flash_attention: true, - mixed_precision: true, - memory_efficient: true, - max_inference_latency_us: 50, - target_throughput_pps: 100_000, - } - } -} - -/// `TFT` Model State for incremental processing -/// -/// **MEMORY SAFETY FIX (2025-10-25)**: -/// - Replaced unbounded HashMap with LRU cache (max 1000 entries) -/// - Prevents 3.6GB/hour memory leak in production inference -/// - Automatically evicts oldest cache entries when full -/// - Tested: 1-hour inference run with 10K predictions = stable 24MB memory -#[derive(Debug, Clone)] -pub struct TFTState { - pub hidden_state: Option, - pub attention_cache: LruCache, - pub last_update: u64, -} - -impl TFTState { - /// Maximum attention cache entries (2000 = ~48MB for TFT-225, 60% training speedup) - /// Chosen to balance: - /// - Memory safety: <100MB cache overhead (acceptable for training) - /// - Hit rate: >95% for typical 50-sequence inference - /// - Eviction overhead: <0.5% latency impact (reduced by 2x cache size) - pub const MAX_CACHE_ENTRIES: usize = 2000; - - pub fn zeros(_config: &TFTConfig) -> Result { - // MAX_CACHE_ENTRIES (2000) is non-zero by construction - let capacity = NonZeroUsize::new(Self::MAX_CACHE_ENTRIES) - .ok_or_else(|| MLError::ConfigError("MAX_CACHE_ENTRIES must be non-zero".to_owned()))?; - - Ok(Self { - hidden_state: None, - attention_cache: LruCache::new(capacity), - last_update: 0, - }) - } - - /// Clear attention cache to free memory - /// Call this after training/inference batch to prevent memory accumulation - pub fn clear_cache(&mut self) { - self.attention_cache.clear(); - self.hidden_state = None; - } -} - -/// `TFT` Model Metadata -#[derive(Debug, Clone, Serialize, Deserialize)] -pub struct TFTMetadata { - pub model_id: String, - pub version: String, - pub input_dim: usize, - pub output_dim: usize, - pub created_at: SystemTime, - pub last_trained: Option, - pub training_samples: u64, - pub performance_metrics: HashMap, -} - -/// Multi-horizon prediction result -#[derive(Debug, Clone)] -pub struct MultiHorizonPrediction { - pub predictions: Vec, // Point predictions for each horizon - pub quantiles: Vec>, // Quantile predictions [horizon][quantile] - pub uncertainty: Vec, // Uncertainty estimates - pub confidence_intervals: Vec<(f64, f64)>, // 90% confidence intervals - pub attention_weights: HashMap>, // Attention interpretability - pub feature_importance: Vec, // Variable importance scores - pub latency_us: u64, // Inference latency -} - -/// Complete Temporal Fusion Transformer -pub struct TemporalFusionTransformer { - pub config: TFTConfig, - pub metadata: TFTMetadata, - pub is_trained: bool, - - // Core TFT components (None when feature count is 0 — avoids zero-dim CUDA tensors) - static_variable_selection: Option, - historical_variable_selection: VariableSelectionNetwork, - future_variable_selection: Option, - - // Encoding layers (None when corresponding feature count is 0) - static_encoder: Option, - historical_encoder: GRNStack, - future_encoder: Option, - - // Temporal processing - lstm_encoder: Linear, // Simplified LSTM representation - lstm_decoder: Linear, - - // Attention mechanism - temporal_attention: TemporalSelfAttention, - - // Output layers - quantile_outputs: QuantileLayer, - - // Performance tracking - inference_count: AtomicU64, - total_latency_us: AtomicU64, - max_latency_us: AtomicU64, - - device: Device, - - // Variable map for checkpointing - varmap: Arc, -} - -impl std::fmt::Debug for TemporalFusionTransformer { - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - f.debug_struct("TemporalFusionTransformer") - .field("config", &self.config) - .field("metadata", &self.metadata) - .field("is_trained", &self.is_trained) - .field( - "inference_count", - &self.inference_count.load(Ordering::Relaxed), - ) - .field( - "total_latency_us", - &self.total_latency_us.load(Ordering::Relaxed), - ) - .field( - "max_latency_us", - &self.max_latency_us.load(Ordering::Relaxed), - ) - .field("device", &format!("{:?}", self.device)) - .field("varmap", &"Arc") - .finish_non_exhaustive() - } -} - -impl TemporalFusionTransformer { - pub fn new(config: TFTConfig) -> Result { - Self::new_with_device(config, Device::cuda_if_available(0).unwrap_or(Device::Cpu)) - } - - pub fn new_with_device(config: TFTConfig, device: Device) -> Result { - // Validate configuration - let total_features = - config.num_static_features + config.num_known_features + config.num_unknown_features; - if config.num_unknown_features == 0 { - return Err(MLError::ConfigError("TFT requires num_unknown_features > 0 (temporal input)".to_owned())); - } - if total_features != config.input_dim { - return Err(MLError::ConfigError(format!( - "Feature count mismatch: static({}) + known({}) + unknown({}) = {} != input_dim({})", - config.num_static_features, - config.num_known_features, - config.num_unknown_features, - total_features, - config.input_dim - ))); - } - - // Log configuration for debugging - debug!( - "Creating TFT with {} input features (static: {}, known: {}, unknown: {})", - config.input_dim, - config.num_static_features, - config.num_known_features, - config.num_unknown_features - ); - - let varmap = Arc::new(VarMap::new()); - let vs = VarBuilder::from_varmap(&varmap, training_dtype(&device), &device); - - // Create variable selection networks (skip when feature count is 0 — CUDA - // cannot handle zero-dim tensors in linear layers) - let static_variable_selection = (config.num_static_features > 0) - .then(|| { - VariableSelectionNetwork::new( - config.num_static_features, - config.hidden_dim, - vs.pp("static_vsn"), - ) - }) - .transpose()?; - - let historical_variable_selection = VariableSelectionNetwork::new( - config.num_unknown_features, - config.hidden_dim, - vs.pp("historical_vsn"), - )?; - - let future_variable_selection = (config.num_known_features > 0) - .then(|| { - VariableSelectionNetwork::new( - config.num_known_features, - config.hidden_dim, - vs.pp("future_vsn"), - ) - }) - .transpose()?; - - // Create encoding stacks (skip when corresponding VSN is absent) - let static_encoder = (config.num_static_features > 0) - .then(|| { - GRNStack::new( - config.hidden_dim, - config.hidden_dim, - config.hidden_dim, - config.num_layers, - vs.pp("static_encoder"), - ) - }) - .transpose()?; - - let historical_encoder = GRNStack::new( - config.hidden_dim, - config.hidden_dim, - config.hidden_dim, - config.num_layers, - vs.pp("historical_encoder"), - )?; - - let future_encoder = (config.num_known_features > 0) - .then(|| { - GRNStack::new( - config.hidden_dim, - config.hidden_dim, - config.hidden_dim, - config.num_layers, - vs.pp("future_encoder"), - ) - }) - .transpose()?; - - // Simplified LSTM layers (in practice, would use proper LSTM) - let lstm_encoder = linear(config.hidden_dim, config.hidden_dim, vs.pp("lstm_encoder"))?; - let lstm_decoder = linear(config.hidden_dim, config.hidden_dim, vs.pp("lstm_decoder"))?; - - // Temporal attention - let temporal_attention = TemporalSelfAttention::new( - config.hidden_dim, - config.num_heads, - config.dropout_rate, - config.use_flash_attention, - vs.pp("temporal_attention"), - )?; - - // Quantile output layer - let quantile_outputs = QuantileLayer::new( - config.hidden_dim, - config.prediction_horizon, - config.num_quantiles, - vs.pp("quantile_outputs"), - )?; - - // Metadata - let metadata = TFTMetadata { - model_id: Uuid::new_v4().to_string(), - version: "1.0.0".to_owned(), - input_dim: config.input_dim, - output_dim: config.prediction_horizon, - created_at: SystemTime::now(), - last_trained: None, - training_samples: 0, - performance_metrics: HashMap::new(), - }; - - Ok(Self { - config, - metadata, - is_trained: false, - static_variable_selection, - historical_variable_selection, - future_variable_selection, - static_encoder, - historical_encoder, - future_encoder, - lstm_encoder, - lstm_decoder, - temporal_attention, - quantile_outputs, - inference_count: AtomicU64::new(0), - total_latency_us: AtomicU64::new(0), - max_latency_us: AtomicU64::new(0), - device, - varmap, - }) - } - - /// Get reference to the model's VarMap for checkpointing and quantization - pub fn varmap(&self) -> &Arc { - &self.varmap - } - - /// Get mutable reference to the model's VarMap for checkpoint loading - pub fn varmap_mut(&mut self) -> &mut Arc { - &mut self.varmap - } - - /// Get reference to the model's Device - pub fn device(&self) -> &Device { - &self.device - } - - /// Validate input tensor dimensions match configuration - fn validate_input_dimensions( - &self, - static_features: &Tensor, - historical_features: &Tensor, - future_features: &Tensor, - ) -> Result<(), MLError> { - // Validate static features: [batch, num_static_features] (skip if 0) - if self.config.num_static_features > 0 { - let static_dims = static_features.dims(); - if static_dims.len() != 2 { - return Err(MLError::ModelError(format!( - "Static features must be 2D [batch, features], got {} dimensions", - static_dims.len() - ))); - } - if static_dims[1] != self.config.num_static_features { - return Err(MLError::ModelError(format!( - "Static features dimension mismatch: expected {}, got {}", - self.config.num_static_features, static_dims[1] - ))); - } - } - - // Validate historical features: [batch, seq_len, num_unknown_features] - let hist_dims = historical_features.dims(); - if hist_dims.len() != 3 { - return Err(MLError::ModelError(format!( - "Historical features must be 3D [batch, seq, features], got {} dimensions", - hist_dims.len() - ))); - } - if hist_dims[2] != self.config.num_unknown_features { - return Err(MLError::ModelError(format!( - "Historical features dimension mismatch: expected {}, got {}", - self.config.num_unknown_features, - hist_dims[2] - ))); - } - - // Validate future features: [batch, horizon, num_known_features] (skip if 0) - if self.config.num_known_features > 0 { - let fut_dims = future_features.dims(); - if fut_dims.len() != 3 { - return Err(MLError::ModelError(format!( - "Future features must be 3D [batch, horizon, features], got {} dimensions", - fut_dims.len() - ))); - } - if fut_dims[2] != self.config.num_known_features { - return Err(MLError::ModelError(format!( - "Future features dimension mismatch: expected {}, got {}", - self.config.num_known_features, fut_dims[2] - ))); - } - } - - Ok(()) - } - - /// Forward pass through the complete `TFT` architecture - #[instrument(skip(self, static_features, historical_features, future_features))] - pub fn forward( - &mut self, - static_features: &Tensor, - historical_features: &Tensor, - future_features: &Tensor, - ) -> Result { - self.forward_with_checkpointing( - static_features, - historical_features, - future_features, - false, - ) - } - - /// Forward pass with optional gradient checkpointing - /// - /// When gradient checkpointing is enabled: - /// - Memory usage reduced by 30-40% (doesn't store intermediate activations) - /// - Training time increases by ~20% (recomputes activations during backprop) - /// - /// # Arguments - /// * `static_features` - Static input features [batch, num_static_features] - /// * `historical_features` - Historical features [batch, seq_len, num_unknown_features] - /// * `future_features` - Future features [batch, horizon, num_known_features] - /// * `use_checkpointing` - Whether to use gradient checkpointing - #[instrument(skip(self, static_features, historical_features, future_features))] - pub fn forward_with_checkpointing( - &mut self, - static_features: &Tensor, - historical_features: &Tensor, - future_features: &Tensor, - use_checkpointing: bool, - ) -> Result { - let static_features = crate::dqn::mixed_precision::ensure_training_dtype(static_features) - .map_err(|e| MLError::ModelError(e.to_string()))?; - let historical_features = crate::dqn::mixed_precision::ensure_training_dtype(historical_features) - .map_err(|e| MLError::ModelError(e.to_string()))?; - let future_features = crate::dqn::mixed_precision::ensure_training_dtype(future_features) - .map_err(|e| MLError::ModelError(e.to_string()))?; - let start_time = Instant::now(); - - // Validate input dimensions - self.validate_input_dimensions(&static_features, &historical_features, &future_features)?; - - // Log device placement for debugging - debug!("Forward pass device check:"); - debug!(" static_features: {:?}", static_features.device()); - debug!(" historical_features: {:?}", historical_features.device()); - debug!(" future_features: {:?}", future_features.device()); - debug!(" model device: {:?}", self.device); - - // 1. Variable Selection Networks (skip absent feature paths) - let static_encoded = if let Some(ref mut static_vsn) = self.static_variable_selection { - let static_selected = static_vsn.forward(&static_features, None)?; - let encoder = self.static_encoder.as_mut().ok_or_else(|| { - MLError::ModelError( - "static_encoder must exist when static_variable_selection exists".to_owned(), - ) - })?; - if use_checkpointing { - Some(encoder.forward(&static_selected.detach(), None)?) - } else { - Some(encoder.forward(&static_selected, None)?) - } - } else { - None - }; - - let historical_selected = self - .historical_variable_selection - .forward(&historical_features, None)?; - - let historical_encoded = if use_checkpointing { - self.historical_encoder - .forward(&historical_selected.detach(), None)? - } else { - self.historical_encoder - .forward(&historical_selected, None)? - }; - - let future_encoded = if let Some(ref mut future_vsn) = self.future_variable_selection { - let future_selected = future_vsn.forward(&future_features, None)?; - let encoder = self.future_encoder.as_mut().ok_or_else(|| { - MLError::ModelError( - "future_encoder must exist when future_variable_selection exists".to_owned(), - ) - })?; - if use_checkpointing { - Some(encoder.forward(&future_selected.detach(), None)?) - } else { - Some(encoder.forward(&future_selected, None)?) - } - } else { - None - }; - - // 3. Temporal Processing - let historical_temporal = if use_checkpointing { - self.lstm_encoder.forward(&historical_encoded.detach())? - } else { - self.lstm_encoder.forward(&historical_encoded)? - }; - - // 4. Combine temporal representations (skip future if absent) - let combined_temporal = if let Some(ref fut_enc) = future_encoded { - let future_temporal = if use_checkpointing { - self.lstm_decoder.forward(&fut_enc.detach())? - } else { - self.lstm_decoder.forward(fut_enc)? - }; - Tensor::cat(&[&historical_temporal, &future_temporal], 1)? - } else { - historical_temporal - }; - - // 5. Self-Attention - let attended = self.temporal_attention.forward_with_checkpointing( - &combined_temporal, - true, - use_checkpointing, - )?; - - // 6. Apply static context (skip if no static features) - let contextualized = if let Some(ref static_enc) = static_encoded { - self.apply_static_context(&attended, static_enc)? - } else { - attended - }; - - // 7. Quantile Outputs (no checkpointing on final layer) - let quantile_preds = self.quantile_outputs.forward(&contextualized)?; - - debug!(" quantile_preds: {:?}", quantile_preds.device()); - - // Cast output back to F32 for API compatibility - let quantile_preds = quantile_preds.to_dtype(candle_core::DType::F32) - .map_err(|e| MLError::ModelError(format!("Output dtype cast failed: {}", e)))?; - - // Update performance metrics - let latency = start_time.elapsed().as_micros() as u64; - self.update_performance_metrics(latency); - - Ok(quantile_preds) - } - - fn apply_static_context( - &self, - temporal: &Tensor, - static_context: &Tensor, - ) -> Result { - let (batch_size, seq_len, hidden_dim) = temporal.dims3()?; - - // Static context comes from variable selection + GRN encoding - // It has shape [batch, 1, hidden] (variable selection adds seq_len=1 dimension) - // We need to expand it to [batch, seq_len, hidden] to match temporal features - - // First, squeeze out the seq_len=1 dimension to get [batch, hidden] - let static_squeezed = static_context.squeeze(1)?; - - // Then expand to match sequence length using broadcast (zero-copy) - let static_expanded = static_squeezed - .unsqueeze(1)? // [batch, 1, hidden] - .broadcast_as((batch_size, seq_len, hidden_dim))?; // [batch, seq_len, hidden] - zero-copy broadcast - - // Add static context to temporal features - let contextualized = (temporal + &static_expanded)?; - - Ok(contextualized) - } - - /// Multi-horizon prediction interface - pub fn predict_horizons( - &mut self, - static_features: &Array1, - historical_features: &Array2, - future_features: &Array2, - ) -> Result { - if !self.is_trained { - return Err(MLError::ModelError("Model not trained".to_owned())); - } - - let start_time = Instant::now(); - - // Convert ndarray to tensors - let static_tensor = self.array_to_tensor_1d(static_features)?; - let historical_tensor = self.array_to_tensor_2d(historical_features)?; - let future_tensor = self.array_to_tensor_2d(future_features)?; - - // Add batch dimension - let static_batched = static_tensor.unsqueeze(0)?; - let historical_batched = historical_tensor.unsqueeze(0)?; - let future_batched = future_tensor.unsqueeze(0)?; - - // Forward pass - let quantile_preds = self.forward(&static_batched, &historical_batched, &future_batched)?; - - // Extract predictions and process outputs - let pred_data = quantile_preds.squeeze(0)?.to_vec2::()?; // [horizon, quantiles] - - let mut predictions = Vec::new(); - let mut quantiles = Vec::new(); - let mut uncertainty = Vec::new(); - let mut confidence_intervals = Vec::new(); - - for horizon in 0..self.config.prediction_horizon { - let horizon_quantiles = &pred_data[horizon]; - - // Point prediction (median) - let median_idx = self.config.num_quantiles / 2; - predictions.push(horizon_quantiles[median_idx] as f64); - - // All quantiles for this horizon - quantiles.push(horizon_quantiles.iter().map(|&x| x as f64).collect()); - - // Uncertainty (IQR) - let q75_idx = (self.config.num_quantiles * 3) / 4; - let q25_idx = self.config.num_quantiles / 4; - let iqr = horizon_quantiles[q75_idx] - horizon_quantiles[q25_idx]; - uncertainty.push(iqr as f64); - - // 90% confidence interval - let lower_idx = self.config.num_quantiles / 10; // ~10th percentile - let upper_idx = (self.config.num_quantiles * 9) / 10; // ~90th percentile - let ci = ( - horizon_quantiles[lower_idx] as f64, - horizon_quantiles[upper_idx] as f64, - ); - confidence_intervals.push(ci); - } - - // Get feature importance and attention weights - let feature_importance = self.static_variable_selection - .as_ref() - .map(|vsn| vsn.get_importance_scores()) - .transpose()? - .unwrap_or_default(); - let mut attention_weights = HashMap::new(); - let weights = self.temporal_attention.get_attention_weights(); - for (key, weight) in weights { - attention_weights.insert(key, vec![weight]); - } - - let latency = start_time.elapsed().as_micros() as u64; - - Ok(MultiHorizonPrediction { - predictions, - quantiles, - uncertainty, - confidence_intervals, - attention_weights, - feature_importance, - latency_us: latency, - }) - } - - fn array_to_tensor_1d(&self, arr: &Array1) -> Result { - let data: Vec = arr.iter().map(|&x| x as f32).collect(); - let tensor = Tensor::from_slice(&data, arr.len(), &self.device)?; - Ok(tensor) - } - - fn array_to_tensor_2d(&self, arr: &Array2) -> Result { - let data: Vec = arr.iter().map(|&x| x as f32).collect(); - let shape = arr.shape(); - let tensor = Tensor::from_slice(&data, (shape[0], shape[1]), &self.device)?; - Ok(tensor) - } - - fn update_performance_metrics(&self, latency_us: u64) { - self.inference_count.fetch_add(1, Ordering::Relaxed); - self.total_latency_us - .fetch_add(latency_us, Ordering::Relaxed); - - // Update max latency atomically - let mut current_max = self.max_latency_us.load(Ordering::Relaxed); - while latency_us > current_max { - match self.max_latency_us.compare_exchange_weak( - current_max, - latency_us, - Ordering::Relaxed, - Ordering::Relaxed, - ) { - Ok(_) => break, - Err(new_max) => current_max = new_max, - } - } - } - - /// Get performance metrics - pub fn get_metrics(&self) -> HashMap { - let inference_count = self.inference_count.load(Ordering::Relaxed); - let total_latency = self.total_latency_us.load(Ordering::Relaxed); - let max_latency = self.max_latency_us.load(Ordering::Relaxed); - - let avg_latency = if inference_count > 0 { - total_latency as f64 / inference_count as f64 - } else { - 0.0 - }; - - let throughput = if avg_latency > 0.0 { - 1_000_000.0 / avg_latency // predictions per second - } else { - 0.0 - }; - - let mut metrics = HashMap::new(); - metrics.insert("total_inferences".to_owned(), inference_count as f64); - metrics.insert("avg_latency_us".to_owned(), avg_latency); - metrics.insert("max_latency_us".to_owned(), max_latency as f64); - metrics.insert("throughput_pps".to_owned(), throughput); - - metrics - } - - /// Get reference to VarMap for weight extraction - pub fn get_varmap(&self) -> &Arc { - &self.varmap - } - - /// Training interface with real backward pass and optimizer - pub async fn train( - &mut self, - training_data: &[(Array1, Array2, Array2, Array1)], // (static, historical, future, targets) - validation_data: &[(Array1, Array2, Array2, Array1)], - epochs: usize, - ) -> Result<(), MLError> { - info!("Starting TFT training for {} epochs", epochs); - - // Initialize AdamW optimizer with model parameters - let params = self.varmap.all_vars(); - let num_params: usize = params.iter().map(|v| v.as_tensor().elem_count()).sum(); - let lr = 1e-3; - let mut optimizer = AdamW::new( - params, - ParamsAdamW { - lr, - beta1: 0.9, - beta2: 0.999, - eps: 1e-8, - weight_decay: 1e-4, - }, - ) - .map_err(|e| MLError::TrainingError(format!("Failed to create optimizer: {}", e)))?; - - info!( - "Initialized AdamW optimizer: lr={:.2e}, {} parameters", - lr, num_params - ); - - let mut best_val_loss = f64::MAX; - - for epoch in 0..epochs { - let mut epoch_loss = 0.0; - - for (_i, (static_feat, hist_feat, fut_feat, targets)) in - training_data.iter().enumerate() - { - // Convert to tensors - let static_tensor = self.array_to_tensor_1d(static_feat)?.unsqueeze(0)?; - let hist_tensor = self.array_to_tensor_2d(hist_feat)?.unsqueeze(0)?; - let fut_tensor = self.array_to_tensor_2d(fut_feat)?.unsqueeze(0)?; - let target_tensor = self.array_to_tensor_1d(targets)?.unsqueeze(0)?; - - // Forward pass - let predictions = self.forward(&static_tensor, &hist_tensor, &fut_tensor)?; - - // Compute quantile loss - let loss = self - .quantile_outputs - .quantile_loss(&predictions, &target_tensor)?; - epoch_loss += loss.to_vec0::()? as f64; - - // Backward pass -- compute gradients - let grads = loss.backward().map_err(|e| { - MLError::TrainingError(format!("Backward pass failed: {}", e)) - })?; - - // Check gradient health before stepping - let varmap_data = self.varmap.data().lock().map_err(|e| { - MLError::TrainingError(format!("Failed to lock VarMap: {}", e)) - })?; - let mut grad_norm_sq = 0.0_f64; - for (_name, var) in varmap_data.iter() { - if let Some(grad) = grads.get(var.as_tensor()) { - let norm_sq = grad - .sqr() - .and_then(|t| t.sum_all()) - .and_then(|t| t.to_dtype(candle_core::DType::F64)) - .and_then(|t| t.to_scalar::()) - .unwrap_or(0.0); - grad_norm_sq += norm_sq; - } - } - drop(varmap_data); - let grad_norm = grad_norm_sq.sqrt(); - - if grad_norm.is_nan() || grad_norm.is_infinite() { - warn!( - "Gradient explosion detected (norm={}), skipping update", - grad_norm - ); - continue; - } - - // Optimizer step -- update parameters - optimizer.step(&grads).map_err(|e| { - MLError::TrainingError(format!("Optimizer step failed: {}", e)) - })?; - } - - let avg_epoch_loss = epoch_loss / training_data.len().max(1) as f64; - debug!("Epoch {}: Average Loss = {:.6}", epoch, avg_epoch_loss); - - // Validation every 10 epochs - if epoch % 10 == 0 { - let val_loss = self.validate(validation_data).await?; - info!( - "Epoch {}: Train Loss = {:.6}, Val Loss = {:.6}", - epoch, avg_epoch_loss, val_loss - ); - if val_loss < best_val_loss { - best_val_loss = val_loss; - } - } - } - - self.is_trained = true; - self.metadata.last_trained = Some(SystemTime::now()); - self.metadata.training_samples = training_data.len() as u64; - - info!( - "TFT training completed: {} epochs, best val loss = {:.6}", - epochs, best_val_loss - ); - Ok(()) - } - - async fn validate( - &mut self, - validation_data: &[(Array1, Array2, Array2, Array1)], - ) -> Result { - let mut total_loss = 0.0; - - for (static_feat, hist_feat, fut_feat, targets) in validation_data { - let static_tensor = self.array_to_tensor_1d(static_feat)?.unsqueeze(0)?; - let hist_tensor = self.array_to_tensor_2d(hist_feat)?.unsqueeze(0)?; - let fut_tensor = self.array_to_tensor_2d(fut_feat)?.unsqueeze(0)?; - let target_tensor = self.array_to_tensor_1d(targets)?.unsqueeze(0)?; - - let predictions = self.forward(&static_tensor, &hist_tensor, &fut_tensor)?; - let loss = self - .quantile_outputs - .quantile_loss(&predictions, &target_tensor)?; - total_loss += loss.to_vec0::()? as f64; - } - - Ok(total_loss / validation_data.len() as f64) - } - - /// Compute quantile loss for training - pub fn compute_quantile_loss( - &self, - predictions: &Tensor, - targets: &Tensor, - ) -> Result { - self.quantile_outputs.quantile_loss(predictions, targets) - } - - /// HFT-optimized inference - pub fn predict_fast( - &mut self, - static_features: &[f32], - historical_features: &[f32], - future_features: &[f32], - ) -> Result, MLError> { - let start = Instant::now(); - - // Convert to tensors (optimized path) - let static_tensor = - Tensor::from_slice(static_features, static_features.len(), &self.device)? - .unsqueeze(0)?; - - let hist_len = self.config.sequence_length; - let hist_dim = self.config.num_unknown_features; - let historical_tensor = - Tensor::from_slice(historical_features, (hist_len, hist_dim), &self.device)? - .unsqueeze(0)?; - - let fut_len = self.config.prediction_horizon; - let fut_dim = self.config.num_known_features; - let future_tensor = - Tensor::from_slice(future_features, (fut_len, fut_dim), &self.device)?.unsqueeze(0)?; - - // Forward pass - let quantile_preds = self.forward(&static_tensor, &historical_tensor, &future_tensor)?; - - // Extract median predictions - let pred_data = quantile_preds.squeeze(0)?.to_vec2::()?; - let median_idx = self.config.num_quantiles / 2; - let predictions: Vec = pred_data - .iter() - .map(|horizon_quantiles| horizon_quantiles[median_idx]) - .collect(); - - let latency = start.elapsed().as_micros() as u64; - self.update_performance_metrics(latency); - - if latency > self.config.max_inference_latency_us { - warn!( - "Inference latency {}μs exceeds target {}μs", - latency, self.config.max_inference_latency_us - ); - } - - Ok(predictions) - } -} - -/// Implement Checkpointable trait for TFT -#[async_trait] -impl Checkpointable for TemporalFusionTransformer { - fn model_type(&self) -> ModelType { - ModelType::TFT - } - - fn model_name(&self) -> &str { - &self.metadata.model_id - } - - fn model_version(&self) -> &str { - &self.metadata.version - } - - async fn serialize_state(&self) -> Result, MLError> { - // Save VarMap to temporary file, then read as bytes - // VarMap.save() requires a Path, not a writer - let temp_dir = std::env::temp_dir(); - let temp_path = temp_dir.join(format!("tft_checkpoint_{}.safetensors", Uuid::new_v4())); - - // Convert temp_path to string for VarMap::save() - let temp_path_str = temp_path - .to_str() - .ok_or_else(|| MLError::ModelError("Invalid temp path".to_owned()))?; - - self.varmap - .save(temp_path_str) - .map_err(|e| MLError::ModelError(format!("Failed to serialize TFT state: {}", e)))?; - - // Read the file into bytes - let buffer = std::fs::read(&temp_path) - .map_err(|e| MLError::ModelError(format!("Failed to read checkpoint file: {}", e)))?; - - // Clean up temp file - drop(std::fs::remove_file(&temp_path)); - - debug!("Serialized TFT state: {} bytes", buffer.len()); - Ok(buffer) - } - - async fn deserialize_state(&mut self, data: &[u8]) -> Result<(), MLError> { - // Write bytes to temporary file, then load VarMap - let temp_dir = std::env::temp_dir(); - let temp_path = temp_dir.join(format!("tft_restore_{}.safetensors", Uuid::new_v4())); - - std::fs::write(&temp_path, data) - .map_err(|e| MLError::ModelError(format!("Failed to write temp checkpoint: {}", e)))?; - - // Convert temp_path to string for VarMap::load() - let temp_path_str = temp_path - .to_str() - .ok_or_else(|| MLError::ModelError("Invalid temp path".to_owned()))?; - - // Try to get mutable access to the VarMap through Arc - let varmap_mut = Arc::get_mut(&mut self.varmap).ok_or_else(|| { - MLError::ModelError( - "Cannot load checkpoint: VarMap has multiple references. \ - This indicates the model is being shared across threads. \ - Clone the model before loading checkpoint." - .to_string(), - ) - })?; - - // Load the checkpoint into the VarMap - varmap_mut - .load(temp_path_str) - .map_err(|e| MLError::ModelError(format!("Failed to load TFT state: {}", e)))?; - - // Clean up temp file - drop(std::fs::remove_file(&temp_path)); - - debug!("Deserialized TFT state from {} bytes", data.len()); - Ok(()) - } - - fn get_training_state(&self) -> (Option, Option, Option, Option) { - // TFT doesn't track epochs/steps in the current implementation - // Return metadata-based info if available - ( - None, // epoch - None, // step - None, // loss - None, // accuracy - ) - } - - fn get_hyperparameters(&self) -> HashMap { - let mut params = HashMap::new(); - // Core architecture params (Wave C+D: 225 features) - params.insert("input_dim".to_owned(), Value::from(self.config.input_dim)); - params.insert( - "hidden_dim".to_owned(), - Value::from(self.config.hidden_dim), - ); - params.insert("num_heads".to_owned(), Value::from(self.config.num_heads)); - params.insert( - "num_layers".to_owned(), - Value::from(self.config.num_layers), - ); - params.insert( - "prediction_horizon".to_owned(), - Value::from(self.config.prediction_horizon), - ); - params.insert( - "sequence_length".to_owned(), - Value::from(self.config.sequence_length), - ); - params.insert( - "num_quantiles".to_owned(), - Value::from(self.config.num_quantiles), - ); - - // Feature split (critical for Wave C+D compatibility) - params.insert( - "num_static_features".to_owned(), - Value::from(self.config.num_static_features), - ); - params.insert( - "num_known_features".to_owned(), - Value::from(self.config.num_known_features), - ); - params.insert( - "num_unknown_features".to_owned(), - Value::from(self.config.num_unknown_features), - ); - - // Training params - params.insert( - "learning_rate".to_owned(), - Value::from(self.config.learning_rate), - ); - params.insert( - "batch_size".to_owned(), - Value::from(self.config.batch_size), - ); - params.insert( - "dropout_rate".to_owned(), - Value::from(self.config.dropout_rate), - ); - params.insert( - "l2_regularization".to_owned(), - Value::from(self.config.l2_regularization), - ); - - // HFT optimization flags - params.insert( - "use_flash_attention".to_owned(), - Value::from(self.config.use_flash_attention), - ); - params.insert( - "mixed_precision".to_owned(), - Value::from(self.config.mixed_precision), - ); - params.insert( - "memory_efficient".to_owned(), - Value::from(self.config.memory_efficient), - ); - - params - } - - fn get_metrics(&self) -> HashMap { - // Call the existing get_metrics method from TemporalFusionTransformer - let inference_count = self.inference_count.load(Ordering::Relaxed); - let total_latency = self.total_latency_us.load(Ordering::Relaxed); - let max_latency = self.max_latency_us.load(Ordering::Relaxed); - - let avg_latency = if inference_count > 0 { - total_latency as f64 / inference_count as f64 - } else { - 0.0 - }; - - let throughput = if avg_latency > 0.0 { - 1_000_000.0 / avg_latency - } else { - 0.0 - }; - - let mut metrics = HashMap::new(); - metrics.insert("total_inferences".to_owned(), inference_count as f64); - metrics.insert("avg_latency_us".to_owned(), avg_latency); - metrics.insert("max_latency_us".to_owned(), max_latency as f64); - metrics.insert("throughput_pps".to_owned(), throughput); - metrics - } - - fn get_architecture_info(&self) -> HashMap { - let mut info = HashMap::new(); - info.insert("network_type".to_owned(), Value::from("TFT")); - info.insert( - "input_dim".to_owned(), - Value::from(self.metadata.input_dim), - ); - info.insert( - "output_dim".to_owned(), - Value::from(self.metadata.output_dim), - ); - info.insert( - "hidden_dim".to_owned(), - Value::from(self.config.hidden_dim), - ); - info.insert("num_heads".to_owned(), Value::from(self.config.num_heads)); - info.insert( - "num_layers".to_owned(), - Value::from(self.config.num_layers), - ); - info.insert( - "num_static_features".to_owned(), - Value::from(self.config.num_static_features), - ); - info.insert( - "num_known_features".to_owned(), - Value::from(self.config.num_known_features), - ); - info.insert( - "num_unknown_features".to_owned(), - Value::from(self.config.num_unknown_features), - ); - info - } -} - -#[cfg(test)] -mod tests { - use super::*; - use anyhow::Result; - use candle_core::DType; - - #[tokio::test] - async fn test_tft_creation() -> Result<()> { - let config = TFTConfig { - input_dim: 10, - hidden_dim: 32, - num_heads: 4, - num_quantiles: 5, - prediction_horizon: 5, - sequence_length: 20, - num_static_features: 2, - num_known_features: 3, - num_unknown_features: 5, - ..Default::default() - }; - - let tft = TemporalFusionTransformer::new(config) - .map_err(|_| anyhow::anyhow!("Failed to create TFT"))?; - assert_eq!(tft.metadata.input_dim, 10); - assert_eq!(tft.metadata.output_dim, 5); - Ok(()) - } - - #[test] - fn test_tft_225_features_default() -> Result<()> { - // Test default configuration uses 225 features (Wave C+D) - let config = TFTConfig::default(); - assert_eq!( - config.input_dim, 225, - "Default TFT config should use 225 features" - ); - assert_eq!(config.num_static_features, 5); - assert_eq!(config.num_known_features, 10); - assert_eq!(config.num_unknown_features, 210); - - let tft = TemporalFusionTransformer::new(config) - .map_err(|_| anyhow::anyhow!("Failed to create TFT with 225 features"))?; - assert_eq!(tft.metadata.input_dim, 225); - Ok(()) - } - - #[test] - fn test_tft_225_features_validation() -> Result<()> { - // Test that 225-feature TFT validates input dimensions correctly - let config = TFTConfig::default(); // 225 features - let device = Device::Cpu; - let tft = TemporalFusionTransformer::new_with_device(config.clone(), device.clone()) - .map_err(|_| anyhow::anyhow!("Failed to create TFT"))?; - - // Create valid input tensors - let batch_size = 2; - let seq_len = 50; - let horizon = 10; - - let static_features = Tensor::zeros( - (batch_size, config.num_static_features), - DType::F32, - &device, - )?; - let historical_features = Tensor::zeros( - (batch_size, seq_len, config.num_unknown_features), - DType::F32, - &device, - )?; - let future_features = Tensor::zeros( - (batch_size, horizon, config.num_known_features), - DType::F32, - &device, - )?; - - // Should validate successfully - let result = - tft.validate_input_dimensions(&static_features, &historical_features, &future_features); - assert!( - result.is_ok(), - "Valid 225-feature input should pass validation" - ); - - // Test invalid historical features dimension - let invalid_hist = Tensor::zeros((batch_size, seq_len, 50), DType::F32, &device)?; // Wrong dim: 50 instead of 210 - let result = - tft.validate_input_dimensions(&static_features, &invalid_hist, &future_features); - assert!( - result.is_err(), - "Invalid historical features should fail validation" - ); - - Ok(()) - } - - #[test] - fn test_tft_config_mismatch_detection() -> Result<()> { - // Test that mismatched feature counts are detected during construction - let invalid_config = TFTConfig { - input_dim: 225, - num_static_features: 5, - num_known_features: 10, - num_unknown_features: 100, // Wrong: should be 210 for 225 total - ..Default::default() - }; - - let result = TemporalFusionTransformer::new(invalid_config); - assert!( - result.is_err(), - "Mismatched feature counts should be rejected" - ); - - let err_msg = format!("{:?}", result.unwrap_err()); - assert!( - err_msg.contains("Feature count mismatch"), - "Error should mention feature count mismatch" - ); - - Ok(()) - } - - #[test] - fn test_tft_checkpoint_preserves_config() -> Result<()> { - // Test that checkpoint save/load preserves 225-feature configuration - let config = TFTConfig::default(); // 225 features - let tft = TemporalFusionTransformer::new(config.clone()) - .map_err(|_| anyhow::anyhow!("Failed to create TFT"))?; - - let hyperparams = tft.get_hyperparameters(); - - // Verify all critical config params are saved - assert_eq!( - hyperparams.get("input_dim").and_then(|v| v.as_u64()), - Some(225) - ); - assert_eq!( - hyperparams - .get("num_static_features") - .and_then(|v| v.as_u64()), - Some(5) - ); - assert_eq!( - hyperparams - .get("num_known_features") - .and_then(|v| v.as_u64()), - Some(10) - ); - assert_eq!( - hyperparams - .get("num_unknown_features") - .and_then(|v| v.as_u64()), - Some(210) - ); - - Ok(()) - } - - #[test] - fn test_tft_wave_c_config() -> Result<()> { - // Test Wave C configuration (201 features) - let wave_c_config = TFTConfig { - input_dim: 201, - num_static_features: 5, - num_known_features: 10, - num_unknown_features: 186, // 201 - 5 - 10 = 186 - ..Default::default() - }; - - let tft = TemporalFusionTransformer::new(wave_c_config) - .map_err(|_| anyhow::anyhow!("Failed to create TFT with 201 features"))?; - assert_eq!(tft.metadata.input_dim, 201); - - Ok(()) - } - - #[test] - fn test_tft_state_creation() -> Result<()> { - let config = TFTConfig { - hidden_dim: 32, - sequence_length: 20, - num_heads: 4, - ..Default::default() - }; - - let state = - TFTState::zeros(&config).map_err(|_| anyhow::anyhow!("Failed to create state"))?; - assert!(state.last_update == 0); - Ok(()) - } - - #[test] - fn test_tft_config_default() -> Result<()> { - let config = TFTConfig::default(); - assert!(config.input_dim > 0); - assert!(config.hidden_dim > 0); - assert!(config.num_heads > 0); - Ok(()) - } - - #[test] - fn test_tft_performance_metrics() -> Result<()> { - let config = TFTConfig { - input_dim: 30, - hidden_dim: 32, - num_static_features: 5, - num_known_features: 10, - num_unknown_features: 15, // 30 - 5 - 10 = 15 - ..Default::default() - }; - - let tft = TemporalFusionTransformer::new(config) - .map_err(|_| anyhow::anyhow!("Failed to create TFT"))?; - let metrics = tft.get_metrics(); - - assert!(metrics.contains_key("total_inferences")); - assert!(metrics.contains_key("avg_latency_us")); - assert!(metrics.contains_key("max_latency_us")); - assert!(metrics.contains_key("throughput_pps")); - Ok(()) - } - - #[test] - fn test_tft_training_state() -> Result<()> { - let config = TFTConfig::default(); - let mut tft = TemporalFusionTransformer::new(config) - .map_err(|_| anyhow::anyhow!("Failed to create TFT"))?; - - assert!(!tft.is_trained); - tft.is_trained = true; - assert!(tft.is_trained); - Ok(()) - } - - #[test] - fn test_tft_metadata() -> Result<()> { - let config = TFTConfig { - input_dim: 30, - prediction_horizon: 12, - num_static_features: 5, - num_known_features: 10, - num_unknown_features: 15, // 30 - 5 - 10 = 15 - ..Default::default() - }; - - let tft = TemporalFusionTransformer::new(config) - .map_err(|_| anyhow::anyhow!("Failed to create TFT"))?; - assert_eq!(tft.metadata.input_dim, 30); - assert_eq!(tft.metadata.output_dim, 12); - Ok(()) - } -} diff --git a/crates/ml/src/tft/quantized_tft_forward.rs b/crates/ml/src/tft/quantized_tft_forward.rs deleted file mode 100644 index ecfebf28a..000000000 --- a/crates/ml/src/tft/quantized_tft_forward.rs +++ /dev/null @@ -1,311 +0,0 @@ -// This file contains the complete forward pass implementation for QuantizedTFT -// To be integrated into quantized_tft.rs - -use super::*; - -impl QuantizedTemporalFusionTransformer { - /// Validate input tensor dimensions and device placement - fn validate_inputs( - &self, - static_features: &Tensor, - historical_features: &Tensor, - future_features: &Tensor, - ) -> Result<(), MLError> { - // Check static features: [batch, num_static_features] - let static_dims = static_features.dims(); - if static_dims.len() != 2 { - return Err(MLError::ModelError(format!( - "Static features must be 2D [batch, features], got {:?}", - static_dims - ))); - } - if static_dims[1] != self.config.num_static_features { - return Err(MLError::ModelError(format!( - "Static features dimension mismatch: expected {}, got {}", - self.config.num_static_features, static_dims[1] - ))); - } - - // Check historical features: [batch, seq_len, num_unknown_features] - let hist_dims = historical_features.dims(); - if hist_dims.len() != 3 { - return Err(MLError::ModelError(format!( - "Historical features must be 3D [batch, seq, features], got {:?}", - hist_dims - ))); - } - if hist_dims[2] != self.config.num_unknown_features { - return Err(MLError::ModelError(format!( - "Historical features dimension mismatch: expected {}, got {}", - self.config.num_unknown_features, hist_dims[2] - ))); - } - - // Check future features: [batch, horizon, num_known_features] - let fut_dims = future_features.dims(); - if fut_dims.len() != 3 { - return Err(MLError::ModelError(format!( - "Future features must be 3D [batch, horizon, features], got {:?}", - fut_dims - ))); - } - if fut_dims[2] != self.config.num_known_features { - return Err(MLError::ModelError(format!( - "Future features dimension mismatch: expected {}, got {}", - self.config.num_known_features, fut_dims[2] - ))); - } - - // Verify device consistency (compare CUDA vs CPU since Device doesn't implement PartialEq) - let target_device = &self.device; - if static_features.device().is_cuda() != target_device.is_cuda() { - return Err(MLError::ModelError(format!( - "Static features on wrong device: expected {:?}, got {:?}", - target_device, - static_features.device() - ))); - } - if historical_features.device().is_cuda() != target_device.is_cuda() { - return Err(MLError::ModelError(format!( - "Historical features on wrong device: expected {:?}, got {:?}", - target_device, - historical_features.device() - ))); - } - if future_features.device().is_cuda() != target_device.is_cuda() { - return Err(MLError::ModelError(format!( - "Future features on wrong device: expected {:?}, got {:?}", - target_device, - future_features.device() - ))); - } - - Ok(()) - } - - /// Step 1: Static Variable Selection Network (placeholder) - fn forward_static_vsn(&self, static_features: &Tensor) -> Result { - // Placeholder: Project static features to hidden_dim - // [batch, num_static_features] -> [batch, 1, hidden_dim] - let batch_size = static_features.dims()[0]; - - // Simple linear projection (in production, use quantized VSN) - let hidden = Tensor::zeros( - (batch_size, 1, self.config.hidden_dim), - DType::F32, - &self.device, - )?; - - Ok(hidden) - } - - /// Step 2: Historical Encoder (LSTM with INT8 weights) - fn forward_historical_encoder(&self, historical_features: &Tensor) -> Result { - let (batch_size, seq_len, _input_dim) = historical_features.dims3()?; - - if self.lstm_weights.is_empty() { - // Fallback: return zeros if LSTM not initialized - let output = Tensor::zeros( - (batch_size, seq_len, self.config.hidden_dim), - DType::F32, - &self.device, - )?; - return Ok(output); - } - - // Initialize hidden and cell states - let hidden_dim = self.config.hidden_dim; - let num_layers = self.lstm_weights.len(); - - let mut h = Tensor::zeros((num_layers, batch_size, hidden_dim), DType::F32, &self.device)?; - let mut c = Tensor::zeros((num_layers, batch_size, hidden_dim), DType::F32, &self.device)?; - - // Process through quantized LSTM layers - let mut layer_input = historical_features.clone(); - - for (layer_idx, layer_weights) in self.lstm_weights.iter().enumerate() { - let (layer_output, h_new, c_new) = self.forward_lstm_layer( - &layer_input, - &h.narrow(0, layer_idx, 1)?.squeeze(0)?, - &c.narrow(0, layer_idx, 1)?.squeeze(0)?, - layer_weights, - )?; - - layer_input = layer_output; - - // Update hidden and cell states - h = h.slice_assign(&[layer_idx..layer_idx + 1], &h_new.unsqueeze(0)?)?; - c = c.slice_assign(&[layer_idx..layer_idx + 1], &c_new.unsqueeze(0)?)?; - } - - Ok(layer_input) - } - - /// Forward pass through a single LSTM layer with quantized weights - fn forward_lstm_layer( - &self, - input: &Tensor, - h_prev: &Tensor, - c_prev: &Tensor, - layer_weights: &HashMap, - ) -> Result<(Tensor, Tensor, Tensor), MLError> { - let (batch_size, seq_len, _input_dim) = input.dims3()?; - let _hidden_dim = self.config.hidden_dim; - - // Dequantize all LSTM weights - let w_ii = self.quantizer.dequantize_tensor( - layer_weights.get("W_ii").ok_or_else(|| MLError::ModelError("Missing W_ii".to_owned()))? - )?; - let w_if = self.quantizer.dequantize_tensor( - layer_weights.get("W_if").ok_or_else(|| MLError::ModelError("Missing W_if".to_owned()))? - )?; - let w_ig = self.quantizer.dequantize_tensor( - layer_weights.get("W_ig").ok_or_else(|| MLError::ModelError("Missing W_ig".to_owned()))? - )?; - let w_io = self.quantizer.dequantize_tensor( - layer_weights.get("W_io").ok_or_else(|| MLError::ModelError("Missing W_io".to_owned()))? - )?; - - let w_hi = self.quantizer.dequantize_tensor( - layer_weights.get("W_hi").ok_or_else(|| MLError::ModelError("Missing W_hi".to_owned()))? - )?; - let w_hf = self.quantizer.dequantize_tensor( - layer_weights.get("W_hf").ok_or_else(|| MLError::ModelError("Missing W_hf".to_owned()))? - )?; - let w_hg = self.quantizer.dequantize_tensor( - layer_weights.get("W_hg").ok_or_else(|| MLError::ModelError("Missing W_hg".to_owned()))? - )?; - let w_ho = self.quantizer.dequantize_tensor( - layer_weights.get("W_ho").ok_or_else(|| MLError::ModelError("Missing W_ho".to_owned()))? - )?; - - let mut outputs = Vec::new(); - let mut h_t = h_prev.clone(); - let mut c_t = c_prev.clone(); - - // Process each time step - for t in 0..seq_len { - let x_t = input.narrow(1, t, 1)?.squeeze(1)?; // [batch, input_dim] - - // Input gate: i_t = sigmoid(W_ii @ x_t + W_hi @ h_t) - let i_t = manual_sigmoid(&(x_t.matmul(&w_ii)? + h_t.matmul(&w_hi)?)?)?; - - // Forget gate: f_t = sigmoid(W_if @ x_t + W_hf @ h_t) - let f_t = manual_sigmoid(&(x_t.matmul(&w_if)? + h_t.matmul(&w_hf)?)?)?; - - // Cell gate: g_t = tanh(W_ig @ x_t + W_hg @ h_t) - let g_t = (x_t.matmul(&w_ig)? + h_t.matmul(&w_hg)?)?.tanh()?; - - // Output gate: o_t = sigmoid(W_io @ x_t + W_ho @ h_t) - let o_t = manual_sigmoid(&(x_t.matmul(&w_io)? + h_t.matmul(&w_ho)?)?)?; - - // Cell state: c_t = f_t * c_t + i_t * g_t - c_t = (f_t.mul(&c_t)? + i_t.mul(&g_t)?)?; - - // Hidden state: h_t = o_t * tanh(c_t) - h_t = o_t.mul(&c_t.tanh()?)?; - - outputs.push(h_t.unsqueeze(1)?); - } - - // Concatenate outputs along time dimension - let output = Tensor::cat(&outputs, 1)?; // [batch, seq_len, hidden_dim] - - Ok((output, h_t, c_t)) - } - - /// Step 3: Future Decoder (simplified - uses same LSTM architecture) - fn forward_future_decoder(&self, future_features: &Tensor) -> Result { - // For simplicity, treat future decoder similarly to historical encoder - // In production TFT, this would be a separate decoder with different weights - self.forward_historical_encoder(future_features) - } - - /// Step 5: Combine encodings (static, attention output, future) - fn combine_contexts( - &self, - static_encoding: &Tensor, - attention_output: &Tensor, - future_encoding: &Tensor, - ) -> Result { - let (batch_size, hist_seq_len, _hidden_dim) = attention_output.dims3()?; - let (_, fut_seq_len, _) = future_encoding.dims3()?; - - // Expand static context to match sequence length - let total_seq = hist_seq_len + fut_seq_len; - let static_expanded = static_encoding - .squeeze(1)? // [batch, hidden_dim] - .unsqueeze(1)? // [batch, 1, hidden_dim] - .repeat(&[1, total_seq, 1])?; // [batch, total_seq, hidden_dim] - - // Concatenate historical and future along time dimension - let temporal_combined = Tensor::cat(&[attention_output, future_encoding], 1)?; - - // Add static context - let combined = (&temporal_combined + &static_expanded)?; - - // Pool over time dimension to get fixed-size representation - // Use mean pooling: [batch, seq, hidden] -> [batch, hidden] - let pooled = combined.mean(1)?; - - Ok(pooled) - } - - /// Step 6: Quantile Output Layer (placeholder) - fn forward_quantile_output(&self, combined: &Tensor) -> Result { - let batch_size = combined.dims()[0]; - - // In production, apply GRN + linear layer for quantile predictions - // For now, return zeros with correct shape - let output = Tensor::zeros( - (batch_size, self.config.prediction_horizon, self.config.num_quantiles), - DType::F32, - &self.device, - )?; - - Ok(output) - } - - /// Complete end-to-end INT8 forward pass - pub fn forward_integrated( - &self, - static_features: &Tensor, - historical_features: &Tensor, - future_features: &Tensor, - ) -> Result { - // Validate inputs - self.validate_inputs(static_features, historical_features, future_features)?; - - let batch_size = static_features.dims()[0]; - - // Step 1: Static Variable Selection (FWD-01) - let static_encoding = self.forward_static_vsn(static_features)?; - - // Step 2: Historical Encoder (FWD-02) - let historical_encoding = self.forward_historical_encoder(historical_features)?; - - // Step 3: Future Decoder (FWD-03) - let future_encoding = self.forward_future_decoder(future_features)?; - - // Step 4: Temporal Attention (FWD-04) - let attention_output = self.forward_temporal_attention(&historical_encoding, false)?; - - // Step 5: Combine encodings - let combined = self.combine_contexts(&static_encoding, &attention_output, &future_encoding)?; - - // Step 6: Quantile Output (FWD-05) - let predictions = self.forward_quantile_output(&combined)?; - - // Verify output shape - let expected_shape = vec![batch_size, self.config.prediction_horizon, self.config.num_quantiles]; - if predictions.dims() != expected_shape { - return Err(MLError::ModelError(format!( - "Output shape mismatch: expected {:?}, got {:?}", - expected_shape, - predictions.dims() - ))); - } - - Ok(predictions) - } -} diff --git a/crates/ml/src/tgnn/mod.rs b/crates/ml/src/tgnn/mod.rs index ad2ebf141..77b9ddb09 100644 --- a/crates/ml/src/tgnn/mod.rs +++ b/crates/ml/src/tgnn/mod.rs @@ -1,1200 +1,10 @@ -//! # Temporal Graph Gated Networks (TGNN) for HFT +//! Temporal Graph Gated Networks (TGNN) for HFT //! -//! Ultra-low latency implementation of TGNN for market microstructure analysis. -//! -//! ## Key Features -//! -//! - Sub-1μs graph neural network inference -//! - Real-time order book graph construction -//! - Market maker and liquidity flow modeling -//! - Cache-friendly graph operations -//! - Integer arithmetic for precision -//! -//! ## Performance Targets -//! -//! - Graph construction: <500ns from order book -//! - GNN inference: <1μs per prediction -//! - Node updates: <100ns per update -//! - Memory: Minimal allocations +//! This module re-exports the `ml-supervised` TGNN implementation and keeps bridge +//! modules that depend on types from both `ml-supervised` and the parent `ml` crate. -// Module imports -pub mod gating; -pub mod graph; -pub mod message_passing; +// Re-export everything from the ml-supervised tgnn module +pub use ml_supervised::tgnn::*; + +// Bridge modules that depend on ml-internal types (UnifiedTrainable) pub mod trainable_adapter; -pub mod traits; -pub mod types; - -// DO NOT RE-EXPORT - Use explicit imports at usage sites - -// Import types from main crate - this fixes the circular dependency -use crate::{InferenceResult, MLError, ModelMetadata, ModelType, PRECISION_FACTOR}; -// Import traits from the local traits module -use traits::MLModel; -// Import types from this module -use types::{TrainingMetrics, ValidationMetrics}; - -// Import TGNN component types from submodules -use gating::GatingMechanism; -use graph::MarketGraph; -use message_passing::MessagePassing; - -use std::collections::HashMap; -use std::sync::atomic::{AtomicU64, Ordering}; -use std::time::{Instant, SystemTime}; - -// Import RNG utilities from types crate -use rand::prelude::*; // Replace common::rng with standard rand - -use async_trait::async_trait; -use dashmap::DashMap; -use ndarray::{s, Array1, Array2}; -use serde::{Deserialize, Serialize}; -use tracing::{debug, info, warn}; - -/// Node types in market microstructure graph -#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)] -pub enum NodeType { - /// Price level in order book - PriceLevel, - /// Market maker entity - MarketMaker, - /// Liquidity pool - LiquidityPool, - /// Order cluster - OrderCluster, -} - -/// Edge types representing market relationships -#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)] -pub enum EdgeType { - /// Price proximity relationship - PriceProximity, - /// Liquidity flow - LiquidityFlow, - /// Market maker connection - MarketMaking, - /// Order correlation - OrderCorrelation, -} - -/// Market node identifier -#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)] -pub struct NodeId { - pub node_type: NodeType, - pub id: String, -} - -impl NodeId { - pub fn price_level(price: i64) -> Self { - Self { - node_type: NodeType::PriceLevel, - id: format!("price_{}", price), - } - } - - pub fn market_maker>(name: S) -> Self { - Self { - node_type: NodeType::MarketMaker, - id: name.into(), - } - } -} - -/// Market edge with temporal properties -#[derive(Debug, Clone, Serialize, Deserialize)] -pub struct MarketEdge { - pub edge_type: EdgeType, - pub weight: i64, - pub strength: f64, - pub timestamp: u64, - pub decay_factor: f64, -} - -impl MarketEdge { - pub fn new(edge_type: EdgeType, weight: i64, strength: f64) -> Self { - Self { - edge_type, - weight, - strength, - timestamp: SystemTime::now() - .duration_since(SystemTime::UNIX_EPOCH) - .unwrap_or_default() - .as_nanos() as u64, - decay_factor: 0.99, - } - } - - pub fn apply_temporal_decay(&mut self, current_time: u64) { - let age = current_time.saturating_sub(self.timestamp); - let decay = self.decay_factor.powf(age as f64 / 1_000_000_000.0); // per second - self.strength *= decay; - self.weight = (self.weight as f64 * decay) as i64; - } -} - -/// TGGN model configuration -#[derive(Debug, Clone, Serialize, Deserialize)] -pub struct TGGNConfig { - /// Maximum number of nodes - pub max_nodes: usize, - - /// Maximum number of edges - pub max_edges: usize, - - /// Node feature dimension - pub node_dim: usize, - - /// Edge feature dimension - pub edge_dim: usize, - - /// Hidden dimension for GNN layers - pub hidden_dim: usize, - - /// Number of message passing layers - pub num_layers: usize, - - /// Temporal decay factor - pub temporal_decay: f64, - - /// Graph update frequency (nanoseconds) - pub update_frequency_ns: u64, - - /// Enable SIMD optimizations - pub use_simd: bool, -} - -impl Default for TGGNConfig { - fn default() -> Self { - Self { - max_nodes: 1000, - max_edges: 10000, - node_dim: 32, - edge_dim: 16, - hidden_dim: 64, - num_layers: 3, - temporal_decay: 0.99, - update_frequency_ns: 1_000_000, // 1ms - use_simd: true, - } - } -} - -/// Temporal Graph Gated Networks for market microstructure -#[derive(Debug)] -pub struct TGGN { - /// Model configuration - config: TGGNConfig, - - /// Model metadata - pub metadata: ModelMetadata, - - /// Market graph structure - graph: MarketGraph, - - /// Gating mechanism - gating: GatingMechanism, - - /// Message passing layers - message_passing: Vec, - - /// Node embeddings cache - node_embeddings: DashMap>, - - /// Edge embeddings cache - edge_embeddings: DashMap<(NodeId, NodeId), Array1>, - - /// Whether model is trained - is_trained: bool, - - /// Performance counters - inference_count: AtomicU64, - total_latency_ns: AtomicU64, - max_latency_ns: AtomicU64, - graph_updates: AtomicU64, - - /// Last update timestamp - last_update: AtomicU64, -} - -impl TGGN { - /// Create new TGGN model - pub fn new(config: TGGNConfig) -> Result { - let mut metadata = ModelMetadata::new( - ModelType::TGNN, - "1.0.0".to_owned(), - config.node_dim, - 1.0, // Single output for prediction - ); - metadata.add_metadata("max_nodes", config.max_nodes.to_string()); - metadata.add_metadata("max_edges", config.max_edges.to_string()); - metadata.add_metadata("hidden_dim", config.hidden_dim.to_string()); - metadata.add_metadata("num_layers", config.num_layers.to_string()); - - let graph = MarketGraph::new(config.max_nodes, config.max_edges)?; - let gating = GatingMechanism::new(config.hidden_dim)?; - - // Initialize message passing layers - let mut message_passing = Vec::with_capacity(config.num_layers); - for layer in 0..config.num_layers { - let input_dim = if layer == 0 { - config.node_dim - } else { - config.hidden_dim - }; - message_passing.push(MessagePassing::new(input_dim, config.hidden_dim)?); - } - - info!( - "Initialized TGGN with {} nodes, {} layers", - config.max_nodes, config.num_layers - ); - - Ok(Self { - config, - metadata, - graph, - gating, - message_passing, - node_embeddings: DashMap::new(), - edge_embeddings: DashMap::new(), - is_trained: false, - inference_count: AtomicU64::new(0), - total_latency_ns: AtomicU64::new(0), - max_latency_ns: AtomicU64::new(0), - graph_updates: AtomicU64::new(0), - last_update: AtomicU64::new(0), - }) - } - - /// Create with default configuration - pub fn default() -> Result { - Self::new(TGGNConfig::default()) - } - - /// Update graph from order book data - pub fn update_from_order_book( - &mut self, - bids: &[(i64, i64)], // (price, volume) pairs - asks: &[(i64, i64)], - timestamp: u64, - ) -> Result<(), MLError> { - let start = Instant::now(); - - // Clear old nodes and edges - self.graph - .clear_temporal_data(timestamp, self.config.temporal_decay)?; - - // Add price level nodes for bids - for (i, &(price, volume)) in bids.iter().enumerate() { - let node_id = NodeId::price_level(price); - let features = self.create_price_level_features(price, volume, true, i)?; - self.graph.add_node(node_id.clone(), features.to_vec())?; - self.node_embeddings.insert(node_id, features); - } - - // Add price level nodes for asks - for (i, &(price, volume)) in asks.iter().enumerate() { - let node_id = NodeId::price_level(price); - let features = self.create_price_level_features(price, volume, false, i)?; - self.graph.add_node(node_id.clone(), features.to_vec())?; - self.node_embeddings.insert(node_id, features); - } - - // Create edges between nearby price levels - self.create_proximity_edges(bids, asks)?; - - // Create liquidity flow edges - self.create_liquidity_edges(bids, asks)?; - - let elapsed = start.elapsed(); - self.graph_updates.fetch_add(1, Ordering::Relaxed); - self.last_update.store(timestamp, Ordering::Relaxed); - - debug!( - "Updated graph in {}ns: {} nodes, {} edges", - elapsed.as_nanos(), - self.graph.node_count(), - self.graph.edge_count() - ); - - // Check latency target - let latency_ns = elapsed.as_nanos() as u64; - if latency_ns > 500 { - // 500ns target - warn!("Graph update {}ns exceeds target 500ns", latency_ns); - } - - Ok(()) - } - - /// Perform graph neural network inference - pub fn gnn_inference( - &mut self, - target_nodes: &[NodeId], - ) -> Result, MLError> { - let start = Instant::now(); - - let mut predictions = HashMap::new(); - - // Get current node embeddings - let mut node_features = HashMap::new(); - for node_id in target_nodes { - if let Some(embedding) = self.node_embeddings.get(node_id) { - node_features.insert(node_id.clone(), embedding.clone()); - } else { - // Create default features if node not found - let default_features = Array1::zeros(self.config.node_dim); - node_features.insert(node_id.clone(), default_features); - } - } - - // Apply message passing layers - for (layer_idx, layer) in self.message_passing.iter().enumerate() { - let layer_start = Instant::now(); - - // Collect messages from neighbors - for node_id in target_nodes { - if let Some(neighbors) = self.graph.get_neighbors(node_id) { - let messages = self.collect_messages(node_id, &neighbors, &node_features)?; - - // Apply gating mechanism - let gated_messages = self.gating.apply(&messages)?; - - // Update node features with gated messages - if let Some(current_features) = node_features.get_mut(node_id) { - let updated = layer.forward(current_features, &gated_messages)?; - *current_features = updated; - } - } - } - - debug!( - "Layer {} completed in {}ns", - layer_idx, - layer_start.elapsed().as_nanos() - ); - } - - // Generate predictions from final node features - for node_id in target_nodes { - if let Some(features) = node_features.get(node_id) { - // Simple prediction: weighted sum of features - let prediction = features.sum() / features.len() as f64; - predictions.insert(node_id.clone(), prediction); - } - } - - let elapsed = start.elapsed(); - self.inference_count.fetch_add(1, Ordering::Relaxed); - self.total_latency_ns - .fetch_add(elapsed.as_nanos() as u64, Ordering::Relaxed); - - let latency_ns = elapsed.as_nanos() as u64; - let current_max = self.max_latency_ns.load(Ordering::Relaxed); - if latency_ns > current_max { - self.max_latency_ns - .compare_exchange_weak( - current_max, - latency_ns, - Ordering::Relaxed, - Ordering::Relaxed, - ) - .ok(); - } - - // Check sub-1μs target - if latency_ns > 1000 { - // 1μs = 1000ns - warn!("GNN inference {}ns exceeds target 1000ns", latency_ns); - } else { - debug!("GNN inference completed in {}ns", latency_ns); - } - - Ok(predictions) - } - - /// Create features for price level nodes - fn create_price_level_features( - &self, - price: i64, - volume: i64, - is_bid: bool, - depth_level: usize, - ) -> Result, MLError> { - let mut features = Array1::zeros(self.config.node_dim); - - // Normalize price and volume - let price_norm = (price as f64) / PRECISION_FACTOR as f64; - let volume_norm = (volume as f64) / PRECISION_FACTOR as f64; - - // Feature 0-3: Basic price/volume info - if features.len() > 0 { - features[0] = price_norm; - } - if features.len() > 1 { - features[1] = volume_norm; - } - if features.len() > 2 { - features[2] = if is_bid { 1.0 } else { -1.0 }; - } - if features.len() > 3 { - features[3] = depth_level as f64 / 10.0; - } - - // Feature 4-7: Statistical features - if features.len() > 4 { - features[4] = price_norm.ln(); - } // Log price - if features.len() > 5 { - features[5] = volume_norm.sqrt(); - } // Sqrt volume - if features.len() > 6 { - features[6] = price_norm * volume_norm; - } // Price * volume - if features.len() > 7 { - features[7] = volume_norm / (price_norm + 1e-8); - } // Volume/price ratio - - // Feature 8-15: Technical indicators (simplified) - for i in 8..features.len().min(16) { - let phase = (i as f64 * std::f64::consts::PI) / 8.0; - features[i] = (price_norm * phase.cos() + volume_norm * phase.sin()) / 10.0; - } - - // Feature 16+: Reserved for market microstructure - for i in 16..features.len() { - features[i] = thread_rng().gen::() * 0.01; // Small random noise - } - - Ok(features) - } - - /// Create edges between nearby price levels - fn create_proximity_edges( - &mut self, - bids: &[(i64, i64)], - asks: &[(i64, i64)], - ) -> Result<(), MLError> { - // Connect adjacent price levels within same side - for window in bids.windows(2) { - let node1 = NodeId::price_level(window[0].0); - let node2 = NodeId::price_level(window[1].0); - let weight = ((window[0].1 + window[1].1) / 2) as i64; // Avg volume - let edge = MarketEdge::new(EdgeType::PriceProximity, weight, 0.8); - self.graph.add_edge(&node1, &node2, edge)?; - } - - for window in asks.windows(2) { - let node1 = NodeId::price_level(window[0].0); - let node2 = NodeId::price_level(window[1].0); - let weight = ((window[0].1 + window[1].1) / 2) as i64; - let edge = MarketEdge::new(EdgeType::PriceProximity, weight, 0.8); - self.graph.add_edge(&node1, &node2, edge)?; - } - - // Connect best bid and ask - if !bids.is_empty() && !asks.is_empty() { - let best_bid = NodeId::price_level(bids[0].0); - let best_ask = NodeId::price_level(asks[0].0); - let spread_weight = (asks[0].0 - bids[0].0).abs(); - let edge = MarketEdge::new(EdgeType::PriceProximity, spread_weight, 0.9); - self.graph.add_edge(&best_bid, &best_ask, edge)?; - } - - Ok(()) - } - - /// Create liquidity flow edges - fn create_liquidity_edges( - &mut self, - bids: &[(i64, i64)], - asks: &[(i64, i64)], - ) -> Result<(), MLError> { - // Create flow edges based on volume imbalance - let total_bid_volume: i64 = bids.iter().map(|(_, v)| v).sum(); - let total_ask_volume: i64 = asks.iter().map(|(_, v)| v).sum(); - - let imbalance = total_bid_volume - total_ask_volume; - let flow_strength = - (imbalance.abs() as f64) / (total_bid_volume + total_ask_volume + 1) as f64; - - // Connect high-volume levels with flow edges - for &(price, volume) in bids.into_iter().take(3) { - for &(ask_price, ask_volume) in asks.into_iter().take(3) { - if volume > total_bid_volume / 10 && ask_volume > total_ask_volume / 10 { - let node1 = NodeId::price_level(price); - let node2 = NodeId::price_level(ask_price); - let weight = (volume.min(ask_volume)) as i64; - let edge = MarketEdge::new(EdgeType::LiquidityFlow, weight, flow_strength); - self.graph.add_edge(&node1, &node2, edge)?; - } - } - } - - Ok(()) - } - - /// Collect messages from neighboring nodes - fn collect_messages( - &self, - node_id: &NodeId, - neighbors: &[NodeId], - node_features: &HashMap>, - ) -> Result>, MLError> { - let mut messages = Vec::new(); - - for neighbor in neighbors { - if let Some(neighbor_features) = node_features.get(neighbor) { - // Get edge weight if available - let edge_weight = self.graph.get_edge_weight(node_id, neighbor).unwrap_or(1.0); - - // Weight neighbor features by edge strength - let weighted_message = neighbor_features.mapv(|x| x * edge_weight); - messages.push(weighted_message); - } - } - - Ok(messages) - } - - /// Get performance statistics - pub fn get_performance_stats(&self) -> HashMap { - let mut stats = HashMap::new(); - - let inference_count = self.inference_count.load(Ordering::Relaxed); - let total_latency = self.total_latency_ns.load(Ordering::Relaxed); - let max_latency = self.max_latency_ns.load(Ordering::Relaxed); - let graph_updates = self.graph_updates.load(Ordering::Relaxed); - - stats.insert("inference_count".to_owned(), inference_count as f64); - stats.insert("graph_updates".to_owned(), graph_updates as f64); - stats.insert("max_latency_ns".to_owned(), max_latency as f64); - - if inference_count > 0 { - stats.insert( - "avg_latency_ns".to_owned(), - total_latency as f64 / inference_count as f64, - ); - } - - stats.insert("node_count".to_owned(), self.graph.node_count() as f64); - stats.insert("edge_count".to_owned(), self.graph.edge_count() as f64); - - stats - } - - /// Public getters for checkpoint operations - pub fn node_embeddings(&self) -> &DashMap> { - &self.node_embeddings - } - - pub fn edge_embeddings(&self) -> &DashMap<(NodeId, NodeId), Array1> { - &self.edge_embeddings - } - - pub fn config(&self) -> &TGGNConfig { - &self.config - } - - pub fn inference_count(&self) -> &AtomicU64 { - &self.inference_count - } - - pub fn graph_updates(&self) -> &AtomicU64 { - &self.graph_updates - } - - pub fn is_trained(&self) -> bool { - self.is_trained - } - - /// Get graph statistics for monitoring and checkpointing - pub fn get_graph_stats(&self) -> (usize, usize) { - (self.graph.node_count(), self.graph.edge_count()) - } - - /// Restore node embeddings from checkpoint state - pub fn restore_node_embeddings( - &mut self, - embeddings: &HashMap>, - ) -> Result<(), MLError> { - for (node_id_str, embedding) in embeddings { - // Convert f32 to f64 - let embedding_f64: Vec = embedding.iter().map(|&x| x as f64).collect(); - let array = Array1::from_vec(embedding_f64); - // Parse the node ID string to create proper NodeId - let node_id = if node_id_str.starts_with("price_") { - let price = node_id_str - .strip_prefix("price_") - .and_then(|s| s.parse::().ok()) - .unwrap_or(0); - NodeId::price_level(price) - } else if node_id_str.starts_with("mm_") { - NodeId::market_maker(node_id_str.get(3..).unwrap_or_default()) - } else { - // Default case - use the string as-is with generic type - NodeId { - node_type: NodeType::PriceLevel, - id: node_id_str.clone(), - } - }; - self.node_embeddings.insert(node_id, array); - } - Ok(()) - } - - /// Restore edge embeddings from checkpoint state - pub fn restore_edge_embeddings( - &mut self, - embeddings: &HashMap>, - ) -> Result<(), MLError> { - for (edge_key, embedding) in embeddings { - // Convert f32 to f64 - let embedding_f64: Vec = embedding.iter().map(|&x| x as f64).collect(); - let array = Array1::from_vec(embedding_f64); - - // Parse edge key (assuming format like "from_id->to_id") - if let Some((from_str, to_str)) = edge_key.split_once("->") { - // Parse from node - let from_node = if from_str.starts_with("price_") { - let price = from_str - .strip_prefix("price_") - .and_then(|s| s.parse::().ok()) - .unwrap_or(0); - NodeId::price_level(price) - } else if from_str.starts_with("mm_") { - NodeId::market_maker(from_str.get(3..).unwrap_or_default()) - } else { - NodeId { - node_type: NodeType::PriceLevel, - id: from_str.to_string(), - } - }; - - // Parse to node - let to_node = if to_str.starts_with("price_") { - let price = to_str - .strip_prefix("price_") - .and_then(|s| s.parse::().ok()) - .unwrap_or(0); - NodeId::price_level(price) - } else if to_str.starts_with("mm_") { - NodeId::market_maker(to_str.get(3..).unwrap_or_default()) - } else { - NodeId { - node_type: NodeType::PriceLevel, - id: to_str.to_string(), - } - }; - - self.edge_embeddings.insert((from_node, to_node), array); - } - } - Ok(()) - } - - /// Restore graph statistics from checkpoint state - pub fn restore_graph_statistics( - &mut self, - _stats: &HashMap, - ) -> Result<(), MLError> { - // Production implementation for now - graph statistics would be restored here - Ok(()) - } - - /// Restore message passing weights from checkpoint state - pub fn restore_message_passing_weights( - &mut self, - _weights: &Vec>, - ) -> Result<(), MLError> { - // Production implementation for now - message passing weights would be restored here - Ok(()) - } -} - -#[async_trait] -impl MLModel for TGGN { - type Config = serde_json::Value; - - fn metadata(&self) -> &ModelMetadata { - &self.metadata - } - - fn is_ready(&self) -> bool { - self.is_trained - } - - async fn train( - &mut self, - features: &Array2, - targets: &Array2, - ) -> Result { - info!("Starting TGGN training with {} samples", features.nrows()); - - let start = Instant::now(); - let _n_samples = features.nrows(); - - // For TGGN, training involves learning message passing weights with real gradients - let learning_rate = 0.001; - - // Prepare batch data for message passing training (moved outside loop) - let mut node_features_batch = Vec::new(); - let mut neighbor_messages_batch = Vec::new(); - let mut targets_batch = Vec::new(); - - // Convert features to node features and targets for each layer - // Start with input features (32-dim), progressively transform through layers - let mut current_features = features.clone(); - let num_layers = self.message_passing.len(); - - for (layer_idx, layer) in self.message_passing.iter_mut().enumerate() { - info!("Training layer {} with real backpropagation", layer_idx); - - // Clear batch data for this layer - node_features_batch.clear(); - neighbor_messages_batch.clear(); - targets_batch.clear(); - - for sample_idx in 0..current_features.nrows().min(targets.nrows()) { - let node_features = current_features.row(sample_idx).to_owned(); - // Create target with correct dimension (hidden_dim, not 1) - // Replicate the single target value across all hidden dimensions - let target_value = targets[[sample_idx, 0]]; - let target = Array1::from_elem(self.config.hidden_dim, target_value); - - // For training, create synthetic neighbor messages from nearby samples - let mut neighbor_messages = Vec::new(); - for neighbor_idx in 0..3.min(current_features.nrows()) { - // Use up to 3 neighbors - if neighbor_idx != sample_idx { - let neighbor_features = current_features.row(neighbor_idx).to_owned(); - neighbor_messages.push(neighbor_features); - } - } - - node_features_batch.push(node_features); - neighbor_messages_batch.push(neighbor_messages); - targets_batch.push(target); - } - - // Train layer with proper backpropagation - layer - .train_weights( - &node_features_batch, - &neighbor_messages_batch, - &targets_batch, - learning_rate, - ) - .map_err(|e| MLError::TrainingError(format!("Layer training failed: {}", e)))?; - - // Transform features through this layer for next layer's training - // This ensures layer 1 gets 64-dim inputs, not 32-dim - if layer_idx < num_layers - 1 { - let mut transformed_features = Vec::new(); - for sample_idx in 0..current_features.nrows() { - let node_features = current_features.row(sample_idx).to_owned(); - - // Get neighbor messages for transformation - let mut neighbor_messages = Vec::new(); - for neighbor_idx in 0..3.min(current_features.nrows()) { - if neighbor_idx != sample_idx { - let neighbor_features = current_features.row(neighbor_idx).to_owned(); - neighbor_messages.push(neighbor_features); - } - } - - // Transform through layer - let transformed = layer - .forward(&node_features, &neighbor_messages) - .unwrap_or_else(|_| Array1::zeros(self.config.hidden_dim)); - transformed_features.push(transformed); - } - - // Convert to Array2 for next layer - let n_samples = transformed_features.len(); - let feature_dim = self.config.hidden_dim; - let flat_len = n_samples * feature_dim; - let flat: Vec = transformed_features - .into_iter() - .flat_map(|arr| arr.to_vec()) - .collect(); - current_features = - Array2::from_shape_vec((n_samples, feature_dim), flat).map_err(|_| { - MLError::DimensionMismatch { - expected: flat_len, - actual: flat_len, - } - })?; - } - } - - // Update gating mechanism with real gradients - if !node_features_batch.is_empty() { - // Transform node features to hidden_dim for gating mechanism - // The gating mechanism expects hidden_dim inputs and outputs - let mut transformed_inputs = Vec::new(); - for (node_features, neighbor_messages) in node_features_batch - .into_iter() - .zip(neighbor_messages_batch.into_iter()) - { - // Use first layer to transform node features to hidden_dim - if let Some(first_layer) = self.message_passing.first() { - let transformed = first_layer - .forward(&node_features, &neighbor_messages) - .unwrap_or_else(|_| Array1::zeros(self.config.hidden_dim)); - transformed_inputs.push(transformed); - } - } - - // Only train gating if we have transformed inputs - if !transformed_inputs.is_empty() { - // Gating mechanism uses GLU which halves the dimension - // So we need to adjust targets to match the output dimension (hidden_dim/2) - let glu_output_dim = self.config.hidden_dim / 2; - let mut gating_targets = Vec::new(); - for target in &targets_batch { - // Truncate or pad targets to match GLU output dimension - let adjusted_target = if target.len() >= glu_output_dim { - target.slice(s![..glu_output_dim]).to_owned() - } else { - let mut padded = Array1::zeros(glu_output_dim); - padded.slice_mut(s![..target.len()]).assign(target); - padded - }; - gating_targets.push(adjusted_target); - } - - self.gating - .update_weights(&transformed_inputs, &gating_targets, learning_rate) - .map_err(|e| { - MLError::TrainingError(format!("Gating training failed: {}", e)) - })?; - } - } - - self.is_trained = true; - self.metadata.mark_trained(); - - let training_time = start.elapsed().as_secs_f64(); - - info!("TGGN training completed in {:.2}s", training_time); - - Ok(TrainingMetrics { - loss: 0.1, - accuracy: 0.9, - precision: 0.88, - recall: 0.85, - f1_score: 0.865, - training_time_seconds: training_time, - epochs_trained: 1, - convergence_achieved: true, - additional_metrics: HashMap::new(), - }) - } - - async fn predict(&self, features: &[f64]) -> Result { - if !self.is_trained { - return Err(MLError::NotTrained("TGGN not trained".to_owned())); - } - - let start = Instant::now(); - - // Simple prediction based on features - let prediction = features.iter().sum::() / features.len() as f64; - let confidence = 0.9; // High confidence for graph-based predictions - - let result = InferenceResult::new( - "tgnn_1.0".to_owned(), - prediction, - confidence, - start.elapsed().as_micros() as u64, - start.elapsed().as_nanos() as u64, - self.metadata.clone(), - ); - - Ok(result) - } - - async fn validate( - &self, - features: &Array2, - targets: &Array2, - ) -> Result { - let mut total_error = 0.0; - let mut correct_predictions = 0; - - for i in 0..features.nrows() { - let row_features: Vec = features.row(i).to_vec(); - let prediction_result = self.predict(&row_features).await?; - let prediction = prediction_result.prediction_as_float(); - - let target = targets[[i, 0]]; - let error = (prediction - target).abs(); - total_error += error; - - if error < 0.1 { - // Threshold for "correct" - correct_predictions += 1; - } - } - - let mse = total_error / features.nrows() as f64; - let accuracy = correct_predictions as f64 / features.nrows() as f64; - - Ok(ValidationMetrics { - validation_loss: mse, - validation_accuracy: accuracy, - validation_precision: accuracy * 0.95, - validation_recall: accuracy * 0.93, - validation_f1_score: accuracy * 0.94, - samples_validated: features.nrows(), - additional_metrics: HashMap::new(), - }) - } - - async fn update( - &mut self, - features: &Array2, - targets: &Array2, - ) -> Result<(), MLError> { - // Online learning for TGGN - self.train(features, targets).await?; - Ok(()) - } - - async fn save(&self, path: &str) -> Result<(), MLError> { - let data = serde_json::json!({ - "config": self.config, - "metadata": self.metadata, - "is_trained": self.is_trained, - "performance_stats": self.get_performance_stats(), - }); - - let serialized = - serde_json::to_string_pretty(&data).map_err(|e| MLError::SerializationError { - reason: e.to_string(), - })?; - tokio::fs::write(path, serialized) - .await - .map_err(|e| MLError::SerializationError { - reason: e.to_string(), - })?; - - info!("Saved TGGN model to {}", path); - Ok(()) - } - - async fn load(&mut self, path: &str) -> Result<(), MLError> { - let content = - tokio::fs::read_to_string(path) - .await - .map_err(|e| MLError::SerializationError { - reason: e.to_string(), - })?; - - let data: serde_json::Value = - serde_json::from_str(&content).map_err(|e| MLError::SerializationError { - reason: e.to_string(), - })?; - - self.config = serde_json::from_value(data["config"].clone()).map_err(|e| { - MLError::SerializationError { - reason: e.to_string(), - } - })?; - - self.metadata = serde_json::from_value(data["metadata"].clone()).map_err(|e| { - MLError::SerializationError { - reason: e.to_string(), - } - })?; - - self.is_trained = data["is_trained"].as_bool().unwrap_or(false); - - info!("Loaded TGGN model from {}", path); - Ok(()) - } - - fn config(&self) -> Self::Config { - serde_json::to_value(&self.config).unwrap_or_default() - } - - fn set_config(&mut self, config: Self::Config) -> Result<(), MLError> { - self.config = serde_json::from_value(config).map_err(|e| MLError::ConfigError(e.to_string()))?; - Ok(()) - } -} - -/// Training pipeline for TGGN with order book data -#[derive(Debug)] -pub struct TGGNTrainingPipeline { - pub model: TGGN, - pub training_data: Vec<(Vec<(i64, i64)>, Vec<(i64, i64)>, f64)>, // (bids, asks, target) -} - -impl TGGNTrainingPipeline { - pub fn new(config: TGGNConfig) -> Result { - let model = TGGN::new(config)?; - Ok(Self { - model, - training_data: Vec::new(), - }) - } - - pub fn add_training_sample( - &mut self, - bids: Vec<(i64, i64)>, - asks: Vec<(i64, i64)>, - target: f64, - ) { - self.training_data.push((bids, asks, target)); - } - - pub async fn train_from_order_book_data(&mut self) -> Result { - info!( - "Training TGGN from {} order book samples", - self.training_data.len() - ); - - // Convert order book data to feature matrices - let mut features_vec = Vec::new(); - let mut targets_vec = Vec::new(); - - for (bids, asks, target) in &self.training_data { - // Update graph with order book data - let timestamp = SystemTime::now() - .duration_since(SystemTime::UNIX_EPOCH) - .unwrap_or_default() - .as_nanos() as u64; - - self.model.update_from_order_book(bids, asks, timestamp)?; - - // Extract features from updated graph - let graph_features = self.extract_graph_features()?; - features_vec.push(graph_features); - targets_vec.push(vec![*target]); - } - - // Convert to ndarray format - let features = Array2::from_shape_vec( - (features_vec.len(), features_vec[0].len()), - features_vec.into_iter().flatten().collect(), - ) - .map_err(|e| MLError::DimensionMismatch { - expected: self.model.config.node_dim, - actual: e.to_string().len(), - })?; - - let targets = Array2::from_shape_vec( - (targets_vec.len(), 1), - targets_vec.into_iter().flatten().collect(), - ) - .map_err(|e| MLError::DimensionMismatch { - expected: 1, - actual: e.to_string().len(), - })?; - - // Train the model (convert types::MLError to MLError) - self.model - .train(&features, &targets) - .await - .map_err(|e| MLError::TrainingError(format!("TGNN training failed: {}", e))) - } - - fn extract_graph_features(&self) -> Result, MLError> { - let stats = self.model.graph.get_stats(); - - // Graph topology features - let mut features = vec![ - stats.node_count as f64, - stats.edge_count as f64, - stats.density, - stats.average_degree, - ]; - - // Fill remaining features with zeros if needed - while features.len() < self.model.config.node_dim { - features.push(0.0); - } - - Ok(features) - } -} - -#[cfg(test)] -mod tests { - use super::*; - - #[tokio::test] - async fn test_tggn_creation() -> Result<(), MLError> { - let config = TGGNConfig::default(); - let model = TGGN::new(config)?; - - assert_eq!(model.config.max_nodes, 1000); - assert_eq!(model.config.num_layers, 3); - assert!(!model.is_trained); - Ok(()) - } - - #[tokio::test] - async fn test_order_book_update() -> Result<(), MLError> { - let config = TGGNConfig::default(); - let mut model = TGGN::new(config)?; - - let bids = vec![(100_00000000, 1000_00000000), (99_00000000, 500_00000000)]; - let asks = vec![(101_00000000, 800_00000000), (102_00000000, 600_00000000)]; - let timestamp = 1234567890; - - let result = model.update_from_order_book(&bids, &asks, timestamp); - assert!(result.is_ok()); - - assert_eq!(model.graph.node_count(), 4); // 2 bids + 2 asks - assert!(model.graph.edge_count() > 0); - Ok(()) - } - - #[tokio::test] - async fn test_gnn_inference() -> Result<(), MLError> { - let config = TGGNConfig::default(); - let mut model = TGGN::new(config)?; - - // Setup graph with some nodes - let bids = vec![(100_00000000, 1000_00000000)]; - let asks = vec![(101_00000000, 800_00000000)]; - let timestamp = 1234567890; - - model.update_from_order_book(&bids, &asks, timestamp)?; - - let target_nodes = vec![NodeId::price_level(100_00000000)]; - let predictions = model.gnn_inference(&target_nodes)?; - - assert_eq!(predictions.len(), 1); - assert!(predictions.contains_key(&NodeId::price_level(100_00000000))); - Ok(()) - } - - #[tokio::test] - async fn test_training_pipeline() -> Result<(), MLError> { - let config = TGGNConfig::default(); - let mut pipeline = TGGNTrainingPipeline::new(config)?; - - // Add some training samples - pipeline.add_training_sample( - vec![(100_00000000, 1000_00000000)], - vec![(101_00000000, 800_00000000)], - 0.5, - ); - - pipeline.add_training_sample( - vec![(99_00000000, 1200_00000000)], - vec![(100_00000000, 900_00000000)], - -0.3, - ); - - let metrics = pipeline.train_from_order_book_data().await?; - assert!(metrics.training_time_seconds > 0.0); - assert!(pipeline.model.is_trained); - Ok(()) - } -} diff --git a/crates/ml/src/tlob/mod.rs b/crates/ml/src/tlob/mod.rs index 59495bce4..9124b3a22 100644 --- a/crates/ml/src/tlob/mod.rs +++ b/crates/ml/src/tlob/mod.rs @@ -1,25 +1,13 @@ //! Time Limit Order Book (TLOB) Transformer //! -//! High-performance TLOB analysis for HFT systems with sub-50μs latency requirements. -//! Based on advanced order flow analytics from institutional trading systems. +//! This module re-exports the `ml-supervised` TLOB implementation and keeps bridge +//! modules that depend on types from both `ml-supervised` and the parent `ml` crate. -pub mod analytics; -pub mod features; -pub mod mbp10_feature_extractor; // MBP-10 to TLOB feature extraction -pub mod performance; +// Re-export everything from the ml-supervised tlob module +pub use ml_supervised::tlob::*; + +// Bridge modules that depend on ml-internal types (UnifiedTrainable) pub mod trainable_adapter; -pub mod transformer; -// Re-export key types for external use -pub use features::{ - ExtractionMetrics, - FeatureVector as TLOBFeatureVector, // Rename to avoid conflict with main FeatureVector from lib.rs - TLOBFeatureExtractor, - TLOBFeatures as TLOBInputFeatures, - TLOB_FEATURE_COUNT, -}; +// Re-export bridge types pub use trainable_adapter::{TLOBAdapterConfig, TLOBTrainableAdapter}; -pub use transformer::{TLOBConfig, TLOBMetrics, TLOBTransformer}; - -// Re-export transformer-specific TLOBFeatures with a different name to avoid conflicts -pub use transformer::TLOBFeatures as TLOBPredictionFeatures; diff --git a/crates/ml/src/xlstm/mod.rs b/crates/ml/src/xlstm/mod.rs index e492d1389..a1ca25601 100644 --- a/crates/ml/src/xlstm/mod.rs +++ b/crates/ml/src/xlstm/mod.rs @@ -1,19 +1,13 @@ -//! xLSTM (Extended Long Short-Term Memory). +//! xLSTM (Extended Long Short-Term Memory) //! -//! Combines two cell types: -//! - **sLSTM**: Exponential gating with scalar memory (better for sequential patterns) -//! - **mLSTM**: Matrix memory with key-value association (higher memory capacity) -//! -//! OOM notes: mLSTM matrix memory is (num_heads, head_dim, head_dim) per sample. -//! Keep hidden_dim/num_heads ratio reasonable (head_dim ≤ 32 recommended). +//! This module re-exports the `ml-supervised` xLSTM implementation and keeps bridge +//! modules that depend on types from both `ml-supervised` and the parent `ml` crate. -pub mod block; -pub mod config; -pub mod mlstm; -pub mod network; -pub mod slstm; +// Re-export everything from the ml-supervised xlstm module +pub use ml_supervised::xlstm::*; + +// Bridge modules that depend on ml-internal types (UnifiedTrainable) pub mod trainable; -pub use config::XLSTMConfig; -pub use network::XLSTMNetwork; +// Re-export bridge types pub use trainable::XLSTMTrainableAdapter; diff --git a/crates/ml/tests/async_data_loading_benchmark.rs b/crates/ml/tests/async_data_loading_benchmark.rs deleted file mode 100644 index 13ca4c8d8..000000000 --- a/crates/ml/tests/async_data_loading_benchmark.rs +++ /dev/null @@ -1,265 +0,0 @@ -//! Benchmark: Async Data Loading vs Synchronous Loading -//! -//! This test compares training time with and without async data loading -//! to validate the 20-30% speedup claim. -//! -//! Expected results: -//! - Sync loading: ~100% baseline -//! - Async loading: ~70-80% (20-30% speedup) -//! - CPU utilization: 7% → 30-40% -//! - GPU utilization: 78% → 90-95% - -use anyhow::Result; -use candle_core::{Device, Tensor}; -use ml::hyperopt::adapters::async_data_loader::AsyncDataLoader; -use std::time::Instant; - -/// Create mock training data -fn create_mock_data( - count: usize, - d_model: usize, - seq_len: usize, - device: &Device, -) -> Result> { - let mut data = Vec::new(); - for i in 0..count { - let features: Vec = (0..seq_len * d_model) - .map(|j| (i as f64 + j as f64) / 1000.0) - .collect(); - - let features_tensor = - Tensor::new(features.as_slice(), device)?.reshape((1, seq_len, d_model))?; - - let target_tensor = Tensor::new(&[i as f64 / 1000.0], device)?.reshape((1, 1, 1))?; - - data.push((features_tensor, target_tensor)); - } - Ok(data) -} - -/// Simulate GPU training on a batch (just tensor operations) -fn simulate_gpu_training(features: &Tensor, targets: &Tensor) -> Result { - // Simulate forward pass: matrix multiply + activation - let batch_size = features.dim(0)?; - let seq_len = features.dim(1)?; - let d_model = features.dim(2)?; - - // Flatten for matmul - let features_flat = features.reshape((batch_size * seq_len, d_model))?; - - // Create weight matrix - let weights = Tensor::randn(0.0, 1.0, (d_model, 1), features.device())?; - - // Forward pass - let output = features_flat.matmul(&weights)?; - - // Simulate loss - let predicted = output.mean_all()?.to_scalar::()?; - let target_val = targets.mean_all()?.to_scalar::()?; - let loss = (predicted - target_val).abs(); - - Ok(loss) -} - -/// Test synchronous data loading -fn test_sync_loading( - data: Vec<(Tensor, Tensor)>, - batch_size: usize, - device: &Device, -) -> Result { - let start = Instant::now(); - - let mut total_loss = 0.0; - let mut batch_count = 0; - - // Process batches synchronously (CPU prepares, then GPU trains) - for batch_data in data.chunks(batch_size) { - // CPU: Concatenate batch - let features: Vec<&Tensor> = batch_data.iter().map(|(f, _)| f).collect(); - let batched_features = if batch_data.len() == 1 { - features[0].clone() - } else { - Tensor::cat( - &features.iter().map(|t| (*t).clone()).collect::>(), - 0, - )? - }; - - let targets: Vec<&Tensor> = batch_data.iter().map(|(_, t)| t).collect(); - let batched_targets = if batch_data.len() == 1 { - targets[0].clone() - } else { - Tensor::cat(&targets.iter().map(|t| (*t).clone()).collect::>(), 0)? - }; - - // CPU: Transfer to GPU - let batched_features = batched_features.to_device(device)?; - let batched_targets = batched_targets.to_device(device)?; - - // GPU: Train (simulated) - let loss = simulate_gpu_training(&batched_features, &batched_targets)?; - total_loss += loss; - batch_count += 1; - } - - let elapsed = start.elapsed(); - println!( - "Sync loading: {:.2}s, avg loss: {:.6}, batches: {}", - elapsed.as_secs_f64(), - total_loss / batch_count as f64, - batch_count - ); - - Ok(elapsed) -} - -/// Test asynchronous data loading -fn test_async_loading( - data: Vec<(Tensor, Tensor)>, - batch_size: usize, - prefetch_count: usize, - device: &Device, -) -> Result { - let start = Instant::now(); - - let mut loader = AsyncDataLoader::new(data, batch_size, prefetch_count, device)?; - - let mut total_loss = 0.0; - let mut batch_count = 0; - - // Process batches asynchronously (CPU prefetches while GPU trains) - while let Some((batched_features, batched_targets)) = loader.next_batch() { - // GPU: Train (simulated) - CPU prefetches next batch in parallel - let loss = simulate_gpu_training(&batched_features, &batched_targets)?; - total_loss += loss; - batch_count += 1; - } - - let elapsed = start.elapsed(); - println!( - "Async loading: {:.2}s, avg loss: {:.6}, batches: {}", - elapsed.as_secs_f64(), - total_loss / batch_count as f64, - batch_count - ); - - Ok(elapsed) -} - -#[test] -fn benchmark_sync_vs_async_loading() -> Result<()> { - println!("\n=== Async Data Loading Benchmark ===\n"); - - let device = Device::cuda_if_available(0)?; - println!("Device: {:?}", device); - - // Configuration - let num_samples = 1000; - let batch_size = 32; - let prefetch_count = 3; - let d_model = 54; // State dimension (updated to 54) - let seq_len = 60; - - println!("Samples: {}", num_samples); - println!("Batch size: {}", batch_size); - println!("Prefetch: {}", prefetch_count); - println!("Feature dim: {} x {}", seq_len, d_model); - println!(); - - // Create test data - println!("Creating mock data..."); - let data = create_mock_data(num_samples, d_model, seq_len, &device)?; - - // Test sync loading - println!("\n[1/3] Testing synchronous loading..."); - let sync_time = test_sync_loading(data.clone(), batch_size, &device)?; - - // Small delay to let GPU settle - std::thread::sleep(std::time::Duration::from_millis(500)); - - // Test async loading - println!("\n[2/3] Testing asynchronous loading..."); - let async_time = test_async_loading(data.clone(), batch_size, prefetch_count, &device)?; - - // Test async loading again (warm cache) - println!("\n[3/3] Testing asynchronous loading (warm cache)..."); - let async_time_warm = test_async_loading(data, batch_size, prefetch_count, &device)?; - - // Results - println!("\n=== Results ==="); - println!("Sync time: {:.3}s (100%)", sync_time.as_secs_f64()); - println!( - "Async time: {:.3}s ({:.1}%)", - async_time.as_secs_f64(), - (async_time.as_secs_f64() / sync_time.as_secs_f64()) * 100.0 - ); - println!( - "Async time (warm): {:.3}s ({:.1}%)", - async_time_warm.as_secs_f64(), - (async_time_warm.as_secs_f64() / sync_time.as_secs_f64()) * 100.0 - ); - - let speedup = (sync_time.as_secs_f64() / async_time.as_secs_f64() - 1.0) * 100.0; - let speedup_warm = (sync_time.as_secs_f64() / async_time_warm.as_secs_f64() - 1.0) * 100.0; - - println!("\nSpeedup: {:.1}%", speedup); - println!("Speedup (warm): {:.1}%", speedup_warm); - - // Assertions - println!("\n=== Validation ==="); - - // Async should be faster (or at least not significantly slower) - // Allow 10% margin for test variability - if async_time_warm.as_secs_f64() <= sync_time.as_secs_f64() * 1.1 { - println!("✓ Async loading is faster or comparable"); - } else { - println!("✗ Async loading is slower than expected"); - println!(" This may indicate CPU bottleneck or insufficient prefetch buffer"); - } - - // Check if we achieved target speedup (15-30% range) - if speedup_warm >= 10.0 { - println!("✓ Achieved significant speedup ({:.1}%)", speedup_warm); - } else { - println!("⚠ Speedup lower than expected ({:.1}% < 15%)", speedup_warm); - println!(" This is expected for small datasets or CPU workloads"); - } - - Ok(()) -} - -#[test] -fn benchmark_different_prefetch_counts() -> Result<()> { - println!("\n=== Prefetch Count Impact ===\n"); - - let device = Device::cuda_if_available(0)?; - let num_samples = 500; - let batch_size = 32; - let d_model = 54; - let seq_len = 60; - - let data = create_mock_data(num_samples, d_model, seq_len, &device)?; - - // Test different prefetch counts - for prefetch in [2, 3, 5, 10] { - println!("Prefetch count: {}", prefetch); - let start = Instant::now(); - - let mut loader = AsyncDataLoader::new(data.clone(), batch_size, prefetch, &device)?; - let mut batch_count = 0; - - while let Some((features, targets)) = loader.next_batch() { - let _loss = simulate_gpu_training(&features, &targets)?; - batch_count += 1; - } - - let elapsed = start.elapsed(); - println!( - " Time: {:.3}s, batches: {}\n", - elapsed.as_secs_f64(), - batch_count - ); - } - - Ok(()) -}