diff --git a/crates/ml-alpha/src/trainer/mod.rs b/crates/ml-alpha/src/trainer/mod.rs index bfac77e52..076273484 100644 --- a/crates/ml-alpha/src/trainer/mod.rs +++ b/crates/ml-alpha/src/trainer/mod.rs @@ -6,3 +6,4 @@ pub mod loss; pub mod optim; pub mod perception; +pub mod stacked; diff --git a/crates/ml-alpha/src/trainer/stacked.rs b/crates/ml-alpha/src/trainer/stacked.rs new file mode 100644 index 000000000..e94fc67d0 --- /dev/null +++ b/crates/ml-alpha/src/trainer/stacked.rs @@ -0,0 +1,494 @@ +//! Stacked Mamba2 -> CfC -> Heads trainer (per spec 2026-05-16 amendment). +//! +//! Topology per training step: +//! snap_features × seq_len → Mamba2.forward_train → h_enriched [hidden_dim] +//! → cfc_step (h_old=0) → h_new [hidden_dim] +//! → heads → probs [5] +//! → BCE(probs, labels) → loss +//! +//! Backward chain: +//! grad_probs → heads_backward → grad_h_new + grad_W_heads/b_heads +//! → cfc_step_backward → grad_W_in/W_rec/b + grad_x (=grad_h_enriched) +//! → Mamba2.backward_from_h_enriched → full Mamba2 grad set +//! +//! Optimizers: +//! - 5 CfC AdamWs (W_in, W_rec, b, heads_w, heads_b) — same as PerceptionTrainer +//! - 1 Mamba2AdamW for all 9 Mamba2 parameter tensors +//! +//! Architectural note: CfC runs with h_old=0 each step (no inter-step +//! recurrence in v1). With h_old=0, the cfc_step degenerates to a +//! per-cell-tau scaled tanh-FC layer feeding the heads. Inter-step CfC +//! state (where h_old carries between calls) is a v2 extension. + +use std::sync::Arc; + +use anyhow::{Context, Result}; +use cudarc::driver::{ + CudaFunction, CudaModule, CudaSlice, CudaStream, DevicePtr, DevicePtrMut, LaunchConfig, + PushKernelArg, +}; +use ml_core::cuda_autograd::gpu_tensor::GpuTensor; +use ml_core::device::MlDevice; +use rand::{Rng, SeedableRng}; +use rand_chacha::ChaCha8Rng; + +use crate::cfc::snap_features::{Mbp10RawInput, ES_TICK_SIZE, FEATURE_DIM}; +use crate::heads::{HIDDEN_DIM, N_HORIZONS}; +use crate::mamba2_block::{ + Mamba2AdamW, Mamba2AdamWConfig, Mamba2Block, Mamba2BlockConfig, +}; +use crate::pinned_mem::MappedF32Buffer; +use crate::trainer::loss::{bce_multi_horizon_loss_and_grad_gpu, BceInput}; +use crate::trainer::optim::AdamW; + +const SNAP_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/snap_feature_assemble.cubin")); +const STEP_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/cfc_step.cubin")); +const HEADS_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/multi_horizon_heads.cubin")); + +#[derive(Clone, Debug)] +pub struct Mamba2CfcTrainerConfig { + pub seq_len: usize, + pub mamba2_state_dim: usize, + pub lr_cfc: f32, + pub lr_mamba2: f32, + pub seed: u64, +} + +impl Default for Mamba2CfcTrainerConfig { + fn default() -> Self { + Self { + seq_len: 32, + mamba2_state_dim: 16, + lr_cfc: 3e-3, + lr_mamba2: 1e-3, + seed: 0x4242, + } + } +} + +pub struct Mamba2CfcTrainer { + cfg: Mamba2CfcTrainerConfig, + dev: MlDevice, + stream: Arc, + + // Modules + cached function handles + _snap_module: Arc, + _step_module: Arc, + _heads_module: Arc, + snap_fn: CudaFunction, + step_fn: CudaFunction, + step_bwd_fn: CudaFunction, + heads_fn: CudaFunction, + heads_bwd_fn: CudaFunction, + + // Mamba2 encoder block + its optimizer + pub mamba2: Mamba2Block, + pub mamba2_adamw: Mamba2AdamW, + + // CfC + heads weights + their AdamWs (5 groups) + pub w_in_d: CudaSlice, + pub w_rec_d: CudaSlice, + pub b_d: CudaSlice, + pub tau_d: CudaSlice, + pub heads_w_d: CudaSlice, + pub heads_b_d: CudaSlice, + opt_w_in: AdamW, + opt_w_rec: AdamW, + opt_b: AdamW, + opt_heads_w: AdamW, + opt_heads_b: AdamW, + + // CfC hidden ping-pong (h_old reset to zero per step in v1) + h_old_d: CudaSlice, + h_new_d: CudaSlice, + + // Per-step scratch + bid_px_d: CudaSlice, + bid_sz_d: CudaSlice, + ask_px_d: CudaSlice, + ask_sz_d: CudaSlice, + prev_bid_sz_d: CudaSlice, + prev_ask_sz_d: CudaSlice, + snap_feat_d: CudaSlice, // single-snapshot scratch (32 floats) + + // Forward intermediates + probs_d: CudaSlice, + + // Backward scratch + grad_w_in_d: CudaSlice, + grad_w_rec_d: CudaSlice, + grad_b_d: CudaSlice, + grad_heads_w_d: CudaSlice, + grad_heads_b_d: CudaSlice, + grad_h_old_d: CudaSlice, + grad_h_new_d: CudaSlice, + grad_x_d: CudaSlice, // gradient on cfc_step's x = h_enriched + + // Pre-allocated mapped-pinned staging for snap_feature uploads + stg_bid_px: MappedF32Buffer, + stg_bid_sz: MappedF32Buffer, + stg_ask_px: MappedF32Buffer, + stg_ask_sz: MappedF32Buffer, +} + +impl Mamba2CfcTrainer { + pub fn new(dev: &MlDevice, cfg: &Mamba2CfcTrainerConfig) -> Result { + anyhow::ensure!(cfg.seq_len >= 2, "Mamba2 requires seq_len >= 2"); + + let stream = dev.cuda_stream().context("trainer stream")?.clone(); + let ctx = dev.cuda_context().context("trainer ctx")?; + let snap_module = ctx.load_cubin(SNAP_CUBIN.to_vec()).context("snap cubin")?; + let step_module = ctx.load_cubin(STEP_CUBIN.to_vec()).context("step cubin")?; + let heads_module = ctx.load_cubin(HEADS_CUBIN.to_vec()).context("heads cubin")?; + let snap_fn = snap_module.load_function("snap_feature_assemble")?; + let step_fn = step_module.load_function("cfc_step")?; + let step_bwd_fn = step_module.load_function("cfc_step_backward")?; + let heads_fn = heads_module.load_function("multi_horizon_heads")?; + let heads_bwd_fn = heads_module.load_function("multi_horizon_heads_backward")?; + + // Mamba2 encoder: in_dim=FEATURE_DIM (32), hidden=HIDDEN_DIM (128). + let mamba2 = Mamba2Block::new( + Mamba2BlockConfig { + in_dim: FEATURE_DIM, + hidden_dim: HIDDEN_DIM, + state_dim: cfg.mamba2_state_dim, + seq_len: cfg.seq_len, + }, + stream.clone(), + ) + .context("Mamba2Block::new")?; + let mamba2_adamw = Mamba2AdamW::new( + &mamba2, + Mamba2AdamWConfig { + lr: cfg.lr_mamba2, + grad_clip_max_norm: Some(1.0), + ..Default::default() + }, + ) + .context("Mamba2AdamW::new")?; + + // CfC weights (input = h_enriched [HIDDEN_DIM], output = [HIDDEN_DIM]) + let mut r = ChaCha8Rng::seed_from_u64(cfg.seed); + let n_in = HIDDEN_DIM; + let n_hid = HIDDEN_DIM; + let scale_in = (1.0_f32 / n_in as f32).sqrt(); + let scale_rec = (1.0_f32 / n_hid as f32).sqrt(); + let w_in: Vec = (0..n_hid * n_in).map(|_| r.gen_range(-scale_in..scale_in)).collect(); + let w_rec: Vec = (0..n_hid * n_hid).map(|_| r.gen_range(-scale_rec..scale_rec)).collect(); + let b: Vec = vec![0.0; n_hid]; + let tau: Vec = (0..n_hid) + .map(|_| { + let u: f32 = r.gen_range(0.0..1.0); + (-4.6 + u * 11.5).exp() + }) + .collect(); + let head_scale = scale_rec; + let heads_w: Vec = (0..N_HORIZONS * n_hid).map(|_| r.gen_range(-head_scale..head_scale)).collect(); + let heads_b: Vec = vec![0.0; N_HORIZONS]; + + let w_in_d = upload(&stream, &w_in)?; + let w_rec_d = upload(&stream, &w_rec)?; + let b_d = upload(&stream, &b)?; + let tau_d = upload(&stream, &tau)?; + let heads_w_d = upload(&stream, &heads_w)?; + let heads_b_d = upload(&stream, &heads_b)?; + + let opt_w_in = AdamW::new(dev, n_hid * n_in, cfg.lr_cfc)?; + let opt_w_rec = AdamW::new(dev, n_hid * n_hid, cfg.lr_cfc)?; + let mut opt_b = AdamW::new(dev, n_hid, cfg.lr_cfc)?; + opt_b.wd = 0.0; + let opt_heads_w = AdamW::new(dev, N_HORIZONS * n_hid, cfg.lr_cfc)?; + let mut opt_heads_b = AdamW::new(dev, N_HORIZONS, cfg.lr_cfc)?; + opt_heads_b.wd = 0.0; + + Ok(Self { + cfg: cfg.clone(), + dev: dev.clone(), + h_old_d: stream.alloc_zeros::(n_hid)?, + h_new_d: stream.alloc_zeros::(n_hid)?, + bid_px_d: stream.alloc_zeros::(10)?, + bid_sz_d: stream.alloc_zeros::(10)?, + ask_px_d: stream.alloc_zeros::(10)?, + ask_sz_d: stream.alloc_zeros::(10)?, + prev_bid_sz_d: stream.alloc_zeros::(10)?, + prev_ask_sz_d: stream.alloc_zeros::(10)?, + snap_feat_d: stream.alloc_zeros::(FEATURE_DIM)?, + probs_d: stream.alloc_zeros::(N_HORIZONS)?, + grad_w_in_d: stream.alloc_zeros::(n_hid * n_in)?, + grad_w_rec_d: stream.alloc_zeros::(n_hid * n_hid)?, + grad_b_d: stream.alloc_zeros::(n_hid)?, + grad_heads_w_d: stream.alloc_zeros::(N_HORIZONS * n_hid)?, + grad_heads_b_d: stream.alloc_zeros::(N_HORIZONS)?, + grad_h_old_d: stream.alloc_zeros::(n_hid)?, + grad_h_new_d: stream.alloc_zeros::(n_hid)?, + grad_x_d: stream.alloc_zeros::(n_in)?, + stg_bid_px: unsafe { MappedF32Buffer::new(10) }.map_err(|e| anyhow::anyhow!("stg: {e}"))?, + stg_bid_sz: unsafe { MappedF32Buffer::new(10) }.map_err(|e| anyhow::anyhow!("stg: {e}"))?, + stg_ask_px: unsafe { MappedF32Buffer::new(10) }.map_err(|e| anyhow::anyhow!("stg: {e}"))?, + stg_ask_sz: unsafe { MappedF32Buffer::new(10) }.map_err(|e| anyhow::anyhow!("stg: {e}"))?, + w_in_d, + w_rec_d, + b_d, + tau_d, + heads_w_d, + heads_b_d, + stream, + _snap_module: snap_module, + _step_module: step_module, + _heads_module: heads_module, + snap_fn, + step_fn, + step_bwd_fn, + heads_fn, + heads_bwd_fn, + mamba2, + mamba2_adamw, + opt_w_in, + opt_w_rec, + opt_b, + opt_heads_w, + opt_heads_b, + }) + } + + /// One training step on a sequence of `seq_len` snapshots with one + /// label set (per-horizon, applied at the last position). Returns + /// the BCE loss. + pub fn step( + &mut self, + snapshots: &[Mbp10RawInput], + labels: &[f32; N_HORIZONS], + ) -> Result { + anyhow::ensure!( + snapshots.len() == self.cfg.seq_len, + "snapshots.len()={} != seq_len={}", + snapshots.len(), + self.cfg.seq_len + ); + + // 1. Build the window tensor: pack seq_len snap_features. + let mut window_tensor = GpuTensor::zeros( + &[1, self.cfg.seq_len, FEATURE_DIM], + &self.stream, + ) + .map_err(|e| anyhow::anyhow!("window alloc: {e}"))?; + + for (k, snap) in snapshots.iter().enumerate() { + // Upload raw input slots. + upload_into(&self.stream, &snap.bid_px, &self.stg_bid_px, &mut self.bid_px_d)?; + upload_into(&self.stream, &snap.bid_sz, &self.stg_bid_sz, &mut self.bid_sz_d)?; + upload_into(&self.stream, &snap.ask_px, &self.stg_ask_px, &mut self.ask_px_d)?; + upload_into(&self.stream, &snap.ask_sz, &self.stg_ask_sz, &mut self.ask_sz_d)?; + let prev_mid = snap.prev_mid; + let trade_signed_vol = snap.trade_signed_vol; + let trade_count = snap.trade_count as i32; + let ts_ns = snap.ts_ns as i64; + let prev_ts_ns = snap.prev_ts_ns as i64; + let tick_size = ES_TICK_SIZE; + let cfg1 = LaunchConfig { + grid_dim: (1, 1, 1), + block_dim: (1, 1, 1), + shared_mem_bytes: 0, + }; + { + let mut launch = self.stream.launch_builder(&self.snap_fn); + launch + .arg(&self.bid_px_d).arg(&self.bid_sz_d).arg(&self.ask_px_d).arg(&self.ask_sz_d) + .arg(&self.prev_bid_sz_d).arg(&self.prev_ask_sz_d) + .arg(&prev_mid).arg(&trade_signed_vol).arg(&trade_count) + .arg(&ts_ns).arg(&prev_ts_ns).arg(&tick_size) + .arg(&mut self.snap_feat_d); + unsafe { launch.launch(cfg1).context("snap fwd")?; } + } + // Copy snap_feat_d into window_tensor at offset k*FEATURE_DIM (DtoD). + let nbytes = FEATURE_DIM * std::mem::size_of::(); + unsafe { + let (src_ptr, _g1) = self.snap_feat_d.device_ptr(&self.stream); + let (dst_base, _g2) = window_tensor.data_mut().device_ptr_mut(&self.stream); + let dst_offset_ptr = dst_base + (k * FEATURE_DIM * std::mem::size_of::()) as u64; + cudarc::driver::result::memcpy_dtod_async( + dst_offset_ptr, + src_ptr, + nbytes, + self.stream.cu_stream(), + ) + .context("window pack DtoD")?; + } + } + self.stream.synchronize().context("window pack sync")?; + + // 2. Mamba2 forward — emits (logit, cache); we use cache.h_enriched. + let (_logit, cache) = self.mamba2.forward_train(&window_tensor).context("mamba2 fwd")?; + // h_enriched: [1, HIDDEN_DIM]. Copy its slice into CfC's input + // buffer h_old or directly use as `x` for cfc_step. + + // 3. cfc_step: x = h_enriched, h_old = zeros. + // Zero h_old first. + self.stream + .memset_zeros(&mut self.h_old_d) + .map_err(|e| anyhow::anyhow!("zero h_old: {e}"))?; + let dt_s = 1.0_f32; // unit dt for stacked v1 (CfC's role is per-cell scaled tanh) + let n_in_i = HIDDEN_DIM as i32; + let n_hid_i = HIDDEN_DIM as i32; + let block_dim = 128u32; + let grid_dim = ((HIDDEN_DIM as u32) + block_dim - 1) / block_dim; + let cfg2 = LaunchConfig { + grid_dim: (grid_dim, 1, 1), + block_dim: (block_dim, 1, 1), + shared_mem_bytes: 0, + }; + { + let mut launch = self.stream.launch_builder(&self.step_fn); + launch + .arg(&self.w_in_d).arg(&self.w_rec_d).arg(&self.b_d).arg(&self.tau_d) + .arg(cache.h_enriched.cuda_data()).arg(&self.h_old_d) + .arg(&dt_s).arg(&n_in_i).arg(&n_hid_i) + .arg(&mut self.h_new_d); + unsafe { launch.launch(cfg2).context("cfc fwd")?; } + } + + // 4. heads forward. + let cfg3 = LaunchConfig { + grid_dim: (1, 1, 1), + block_dim: (N_HORIZONS as u32, 1, 1), + shared_mem_bytes: 0, + }; + { + let mut launch = self.stream.launch_builder(&self.heads_fn); + launch.arg(&self.heads_w_d).arg(&self.heads_b_d).arg(&self.h_new_d).arg(&mut self.probs_d); + unsafe { launch.launch(cfg3).context("heads fwd")?; } + } + self.stream.synchronize().context("fwd sync")?; + + // 5. BCE forward + grad_probs. + let probs_host = download(&self.stream, &self.probs_d)?; + let bce_out = bce_multi_horizon_loss_and_grad_gpu( + &self.dev, + &BceInput { probs: probs_host.clone(), labels: labels.to_vec(), n_horizons: N_HORIZONS, n_pos: 1 }, + )?; + let loss = bce_out.loss; + + // 6. heads backward → grad_h_new + grad_heads_* + let grad_probs_d = upload(&self.stream, &bce_out.grad_probs)?; + let probs_d_alias = upload(&self.stream, &probs_host)?; + { + let mut launch = self.stream.launch_builder(&self.heads_bwd_fn); + launch + .arg(&self.heads_w_d).arg(&probs_d_alias).arg(&self.h_new_d).arg(&grad_probs_d) + .arg(&mut self.grad_heads_w_d).arg(&mut self.grad_heads_b_d).arg(&mut self.grad_h_new_d); + unsafe { launch.launch(cfg3).context("heads bwd")?; } + } + + // 7. cfc_step_backward → grad_W_in/W_rec/b + grad_h_old (discard) + grad_x (=grad_h_enriched) + let shared_mem = (2 * HIDDEN_DIM * std::mem::size_of::()) as u32; + let cfg_bwd = LaunchConfig { + grid_dim: (grid_dim, 1, 1), + block_dim: (block_dim, 1, 1), + shared_mem_bytes: shared_mem, + }; + { + let mut launch = self.stream.launch_builder(&self.step_bwd_fn); + launch + .arg(&self.w_in_d).arg(&self.w_rec_d).arg(&self.b_d).arg(&self.tau_d) + .arg(cache.h_enriched.cuda_data()).arg(&self.h_old_d).arg(&self.grad_h_new_d) + .arg(&dt_s).arg(&n_in_i).arg(&n_hid_i) + .arg(&mut self.grad_w_in_d).arg(&mut self.grad_w_rec_d) + .arg(&mut self.grad_b_d).arg(&mut self.grad_h_old_d) + .arg(&mut self.grad_x_d); + unsafe { launch.launch(cfg_bwd).context("cfc bwd")?; } + } + self.stream.synchronize().context("cfc bwd sync")?; + + // 8. Mamba2 backward: wrap grad_x_d as GpuTensor [1, HIDDEN_DIM] for backward_from_h_enriched. + // Need to construct a fresh GpuTensor from our CudaSlice. The + // backward API takes a borrow; build a temporary GpuTensor + // copy. (For v2, allocate this once and reuse.) + let mut grad_h_enriched_slice = self.stream + .alloc_zeros::(HIDDEN_DIM) + .context("grad_h_enriched alloc")?; + let nbytes = HIDDEN_DIM * std::mem::size_of::(); + unsafe { + let (src_ptr, _g1) = self.grad_x_d.device_ptr(&self.stream); + let (dst_ptr, _g2) = grad_h_enriched_slice.device_ptr_mut(&self.stream); + cudarc::driver::result::memcpy_dtod_async(dst_ptr, src_ptr, nbytes, self.stream.cu_stream()) + .context("grad_h_enriched copy")?; + } + let grad_h_enriched_tensor = + GpuTensor::new(grad_h_enriched_slice, vec![1, HIDDEN_DIM]) + .map_err(|e| anyhow::anyhow!("grad_h_enriched as GpuTensor: {e}"))?; + let mamba2_grads = self + .mamba2 + .backward_from_h_enriched(&cache, &grad_h_enriched_tensor) + .context("mamba2 backward_from_h_enriched")?; + + // 9. Apply AdamW updates on all 6 param groups. + self.opt_w_in.step(&mut self.w_in_d, &self.grad_w_in_d)?; + self.opt_w_rec.step(&mut self.w_rec_d, &self.grad_w_rec_d)?; + self.opt_b.step(&mut self.b_d, &self.grad_b_d)?; + self.opt_heads_w.step(&mut self.heads_w_d, &self.grad_heads_w_d)?; + self.opt_heads_b.step(&mut self.heads_b_d, &self.grad_heads_b_d)?; + self.mamba2_adamw + .step(&mut self.mamba2, &mamba2_grads) + .context("mamba2 AdamW step")?; + + Ok(loss) + } + + pub fn last_probs(&self) -> Result<[f32; N_HORIZONS]> { + let v = download(&self.stream, &self.probs_d)?; + let mut out = [0f32; N_HORIZONS]; + out.copy_from_slice(&v); + Ok(out) + } +} + +fn upload(stream: &Arc, host: &[f32]) -> Result> { + let n = host.len(); + let staging = unsafe { MappedF32Buffer::new(n) } + .map_err(|e| anyhow::anyhow!("stacked upload staging: {e}"))?; + staging.write_from_slice(host); + let mut dst = stream.alloc_zeros::(n).context("stacked upload alloc")?; + if n > 0 { + let nbytes = n * std::mem::size_of::(); + unsafe { + let (dst_ptr, _g) = dst.device_ptr_mut(stream); + cudarc::driver::result::memcpy_dtod_async(dst_ptr, staging.dev_ptr, nbytes, stream.cu_stream()) + .context("stacked upload DtoD")?; + } + } + Ok(dst) +} + +fn upload_into( + stream: &Arc, + host: &[f32], + staging: &MappedF32Buffer, + dst: &mut CudaSlice, +) -> Result<()> { + let n = host.len(); + assert_eq!(n, dst.len()); + staging.write_from_slice(host); + if n > 0 { + let nbytes = n * std::mem::size_of::(); + unsafe { + let (dst_ptr, _g) = dst.device_ptr_mut(stream); + cudarc::driver::result::memcpy_dtod_async(dst_ptr, staging.dev_ptr, nbytes, stream.cu_stream()) + .context("upload_into DtoD")?; + } + } + Ok(()) +} + +fn download(stream: &Arc, src: &CudaSlice) -> Result> { + let n = src.len(); + let staging = unsafe { MappedF32Buffer::new(n) } + .map_err(|e| anyhow::anyhow!("stacked download staging: {e}"))?; + let nbytes = n * std::mem::size_of::(); + unsafe { + let (src_ptr, _g) = src.device_ptr(stream); + cudarc::driver::result::memcpy_dtod_async(staging.dev_ptr, src_ptr, nbytes, stream.cu_stream()) + .context("stacked download DtoD")?; + } + stream.synchronize().context("stacked download sync")?; + Ok(staging.read_all()) +} diff --git a/crates/ml-alpha/tests/stacked_overfit.rs b/crates/ml-alpha/tests/stacked_overfit.rs new file mode 100644 index 000000000..de232891e --- /dev/null +++ b/crates/ml-alpha/tests/stacked_overfit.rs @@ -0,0 +1,122 @@ +//! Stacked Mamba2 -> CfC -> Heads synthetic-overfit smoke. +//! +//! Mirrors `perception_overfit.rs` but on the stacked trainer. The +//! signal is constant direction=+1, label=[1; 5]. Asserts the BCE loss +//! shrinks at least 40% over the training budget — proves the +//! full forward + backward (Mamba2 + CfC + heads) wires up correctly +//! end-to-end and that all 6 AdamW optimizers actually move weights. + +use ml_alpha::cfc::snap_features::Mbp10RawInput; +use ml_alpha::trainer::stacked::{Mamba2CfcTrainer, Mamba2CfcTrainerConfig}; +use ml_core::device::MlDevice; + +fn test_device() -> MlDevice { + MlDevice::cuda(0).expect("CUDA 0 required for ml-alpha tests") +} + +fn synthetic_seq(seq_len: usize, mut prev_mid: f32, mut ts_ns: u64) -> (Vec, [f32; 5]) { + let mut out = Vec::with_capacity(seq_len); + for k in 0..seq_len { + let next_mid = prev_mid + 0.25; + let mut bid_px = [0.0f32; 10]; + let mut bid_sz = [0.0f32; 10]; + let mut ask_px = [0.0f32; 10]; + let mut ask_sz = [0.0f32; 10]; + for i in 0..10 { + bid_px[i] = next_mid - 0.125 - 0.25 * i as f32; + ask_px[i] = next_mid + 0.125 + 0.25 * i as f32; + bid_sz[i] = 10.0; + ask_sz[i] = 10.0; + } + let prev_ts = ts_ns; + ts_ns += 20_000_000; + out.push(Mbp10RawInput { + bid_px, + bid_sz, + ask_px, + ask_sz, + prev_mid, + trade_signed_vol: 1.0, + trade_count: 1, + ts_ns, + prev_ts_ns: prev_ts, + }); + prev_mid = next_mid; + let _ = k; + } + (out, [1.0, 1.0, 1.0, 1.0, 1.0]) +} + +#[test] +fn stacked_trainer_constructs_cleanly() { + let dev = test_device(); + let cfg = Mamba2CfcTrainerConfig::default(); + let t = Mamba2CfcTrainer::new(&dev, &cfg).expect("init"); + drop(t); +} + +#[test] +fn stacked_trainer_loss_shrinks_on_constant_signal() { + let dev = test_device(); + let cfg = Mamba2CfcTrainerConfig { + seq_len: 16, // smaller for smoke speed; Mamba2 needs >=2 + mamba2_state_dim: 8, + lr_cfc: 3e-3, + lr_mamba2: 1e-3, + seed: 0x4242, + }; + let mut trainer = Mamba2CfcTrainer::new(&dev, &cfg).expect("init"); + + // Initial loss over 8 batches. + let mut initial_total = 0.0_f32; + let mut ts = 1_000_000u64; + let mut prev_mid = 5500.0_f32; + for _ in 0..8 { + let (seq, labels) = synthetic_seq(cfg.seq_len, prev_mid, ts); + let l = trainer.step(&seq, &labels).expect("step warm"); + initial_total += l; + prev_mid = 0.5 * (seq.last().unwrap().bid_px[0] + seq.last().unwrap().ask_px[0]); + ts = seq.last().unwrap().ts_ns; + } + let initial_avg = initial_total / 8.0; + eprintln!("initial_avg = {initial_avg:.4}"); + + // Train 250 steps. Print every 50. + let mut window_loss = 0.0_f32; + let mut window_count = 0usize; + for step_idx in 0..250 { + let (seq, labels) = synthetic_seq(cfg.seq_len, prev_mid, ts); + let l = trainer.step(&seq, &labels).expect("train step"); + window_loss += l; + window_count += 1; + if step_idx % 50 == 49 { + eprintln!( + " step {}: window_avg_loss={:.4}", + step_idx + 1, + window_loss / window_count as f32 + ); + window_loss = 0.0; + window_count = 0; + } + prev_mid = 0.5 * (seq.last().unwrap().bid_px[0] + seq.last().unwrap().ask_px[0]); + ts = seq.last().unwrap().ts_ns; + } + + // Final loss over 8 batches. + let mut final_total = 0.0_f32; + for _ in 0..8 { + let (seq, labels) = synthetic_seq(cfg.seq_len, prev_mid, ts); + let l = trainer.step(&seq, &labels).expect("step final"); + final_total += l; + prev_mid = 0.5 * (seq.last().unwrap().bid_px[0] + seq.last().unwrap().ask_px[0]); + ts = seq.last().unwrap().ts_ns; + } + let final_avg = final_total / 8.0; + eprintln!("final_avg = {final_avg:.4}"); + + assert!( + final_avg < 0.6 * initial_avg || final_avg < 0.5, + "Mamba2CfcTrainer failed to overfit constant signal: \ + start={initial_avg:.4}, end={final_avg:.4}" + ); +}