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

feat(ml-alpha): Mamba2CfcTrainer — stacked Mamba2 -> CfC -> heads

Browse Source 
Realizes the 2026-05-16 spec amendment merging Mamba2 + CfC into one
stacked architecture (vs the original "compete via gate" framing).

Forward chain:
  snap_features × seq_len
   -> window pack [1, seq_len, FEATURE_DIM]
   -> Mamba2Block.forward_train -> (logit, cache.h_enriched [1, hidden_dim])
   -> cfc_step(x=h_enriched, h_old=0) -> h_new
   -> heads -> probs [5]
   -> BCE(probs, labels)

Backward chain:
  BCE -> grad_probs
   -> heads_backward -> grad_h_new + grad_W_heads, grad_b_heads
   -> cfc_step_backward -> grad_W_in, grad_W_rec, grad_b + grad_x (=grad_h_enriched)
   -> Mamba2.backward_from_h_enriched(&cache, &grad_h_enriched_tensor)
      -> Mamba2BackwardGrads (full 9-tensor gradient set)

Optimizers (6 total):
  - 5 CfC AdamWs (W_in, W_rec, b, heads_w, heads_b) — reused from
    PerceptionTrainer's per-param-group pattern
  - 1 Mamba2AdamW for all 9 Mamba2 parameter tensors (existing
    implementation in mamba2_block.rs)

Synthetic-overfit on constant +1 direction (seq_len=16, state_dim=8,
lr_cfc=3e-3, lr_mamba2=1e-3, 250 steps):
  initial_avg=0.5951 → final_avg=0.1917 (68% drop, well past 40% gate).
  Monotone descent at all 5 progress checkpoints.

Architectural note (v1): CfC runs with h_old=0 each step (no inter-
step recurrence). With h_old=0, the CfC layer is effectively per-cell
tau-scaled tanh FC. Inter-step CfC state (h_old carrying between
calls) is a v2 extension once the cluster gate validates the v1
foundation.

The cluster gate (Task 18) now has the actual stacked production
trainer to deploy, not a CfC-alone-vs-Mamba2-alone bench.

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-05-16 23:03:55 +02:00
parent deed15b34e
commit 45b203e31e
3 changed files with 617 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

1
crates/ml-alpha/src/trainer/mod.rs
View File

@@ -6,3 +6,4 @@
pub mod loss;
pub mod optim;
pub mod perception;
pub mod stacked;

494
crates/ml-alpha/src/trainer/stacked.rs Normal file
View File

@@ -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<CudaStream>,
// Modules + cached function handles
_snap_module: Arc<CudaModule>,
_step_module: Arc<CudaModule>,
_heads_module: Arc<CudaModule>,
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<f32>,
pub w_rec_d: CudaSlice<f32>,
pub b_d: CudaSlice<f32>,
pub tau_d: CudaSlice<f32>,
pub heads_w_d: CudaSlice<f32>,
pub heads_b_d: CudaSlice<f32>,
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<f32>,
h_new_d: CudaSlice<f32>,
// Per-step scratch
bid_px_d: CudaSlice<f32>,
bid_sz_d: CudaSlice<f32>,
ask_px_d: CudaSlice<f32>,
ask_sz_d: CudaSlice<f32>,
prev_bid_sz_d: CudaSlice<f32>,
prev_ask_sz_d: CudaSlice<f32>,
snap_feat_d: CudaSlice<f32>, // single-snapshot scratch (32 floats)
// Forward intermediates
probs_d: CudaSlice<f32>,
// Backward scratch
grad_w_in_d: CudaSlice<f32>,
grad_w_rec_d: CudaSlice<f32>,
grad_b_d: CudaSlice<f32>,
grad_heads_w_d: CudaSlice<f32>,
grad_heads_b_d: CudaSlice<f32>,
grad_h_old_d: CudaSlice<f32>,
grad_h_new_d: CudaSlice<f32>,
grad_x_d: CudaSlice<f32>, // 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<Self> {
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<f32> = (0..n_hid * n_in).map(|_| r.gen_range(-scale_in..scale_in)).collect();
let w_rec: Vec<f32> = (0..n_hid * n_hid).map(|_| r.gen_range(-scale_rec..scale_rec)).collect();
let b: Vec<f32> = vec![0.0; n_hid];
let tau: Vec<f32> = (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<f32> = (0..N_HORIZONS * n_hid).map(|_| r.gen_range(-head_scale..head_scale)).collect();
let heads_b: Vec<f32> = 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::<f32>(n_hid)?,
h_new_d: stream.alloc_zeros::<f32>(n_hid)?,
bid_px_d: stream.alloc_zeros::<f32>(10)?,
bid_sz_d: stream.alloc_zeros::<f32>(10)?,
ask_px_d: stream.alloc_zeros::<f32>(10)?,
ask_sz_d: stream.alloc_zeros::<f32>(10)?,
prev_bid_sz_d: stream.alloc_zeros::<f32>(10)?,
prev_ask_sz_d: stream.alloc_zeros::<f32>(10)?,
snap_feat_d: stream.alloc_zeros::<f32>(FEATURE_DIM)?,
probs_d: stream.alloc_zeros::<f32>(N_HORIZONS)?,
grad_w_in_d: stream.alloc_zeros::<f32>(n_hid * n_in)?,
grad_w_rec_d: stream.alloc_zeros::<f32>(n_hid * n_hid)?,
grad_b_d: stream.alloc_zeros::<f32>(n_hid)?,
grad_heads_w_d: stream.alloc_zeros::<f32>(N_HORIZONS * n_hid)?,
grad_heads_b_d: stream.alloc_zeros::<f32>(N_HORIZONS)?,
grad_h_old_d: stream.alloc_zeros::<f32>(n_hid)?,
grad_h_new_d: stream.alloc_zeros::<f32>(n_hid)?,
grad_x_d: stream.alloc_zeros::<f32>(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<f32> {
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::<f32>();
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::<f32>()) 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::<f32>()) 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::<f32>(HIDDEN_DIM)
.context("grad_h_enriched alloc")?;
let nbytes = HIDDEN_DIM * std::mem::size_of::<f32>();
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<CudaStream>, host: &[f32]) -> Result<CudaSlice<f32>> {
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::<f32>(n).context("stacked upload alloc")?;
if n > 0 {
let nbytes = n * std::mem::size_of::<f32>();
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<CudaStream>,
host: &[f32],
staging: &MappedF32Buffer,
dst: &mut CudaSlice<f32>,
) -> 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::<f32>();
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<CudaStream>, src: &CudaSlice<f32>) -> Result<Vec<f32>> {
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::<f32>();
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())
}

122
crates/ml-alpha/tests/stacked_overfit.rs Normal file
View File

@@ -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<Mbp10RawInput>, [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}"
);
}