diff --git a/crates/ml-alpha/examples/phase1d_mamba.rs b/crates/ml-alpha/examples/phase1d_mamba.rs new file mode 100644 index 000000000..9e33e8c14 --- /dev/null +++ b/crates/ml-alpha/examples/phase1d_mamba.rs @@ -0,0 +1,270 @@ +//! Phase 1d.1 — Stateful Mamba2 sequence encoder smoke. +//! +//! Trains the from-scratch GPU-pure Mamba2 block on the snapshot-stream +//! fxcache. Compares val AUC against the MLP baseline (Phase 1c snapshot: +//! AUC=0.685 at K=100 stateless). +//! +//! Falsification gate (per the implementation plan): +//! - AUC > 0.72 → sequence-state lifts beyond stateless MLP. Proceed +//! to Phase 1d.2 (multi-minute horizon). +//! - AUC ≤ 0.72 → stateless features have a hard ceiling; richer +//! features (not richer model) are the next move. +//! +//! Sequence semantics: each training example is a contiguous block of +//! `seq_len` snapshots ending at bar `t`; the model predicts the binary +//! direction label at bar `t` (sign of `mid_price[t+H] - mid_price[t]`, +//! computed by `prepare_phase1a_data`). +//! +//! Memory model: feature matrix + labels stay CPU-resident (640 MB +//! f32 + 1.5 MB labels for the 1.97 M-row snapshot fxcache). Per batch, +//! we gather contiguous seq_len-bar slices into a pinned host buffer, +//! DMA to GPU once, then do forward + backward + AdamW step. + +use anyhow::{Context, Result}; +use clap::Parser; +use cudarc::driver::CudaContext; +use rand::{Rng, SeedableRng}; +use rand_chacha::ChaCha8Rng; +use std::sync::Arc; +use tracing_subscriber::EnvFilter; + +use ml_alpha::eval::{accuracy_from_logits, auc_from_logits}; +use ml_alpha::mamba2_block::{ + Mamba2AdamW, Mamba2AdamWConfig, Mamba2Block, Mamba2BlockConfig, +}; +use ml_alpha::fxcache_reader::FxCacheReader; +use ml_alpha::purged_split::PurgedSplit; +use ml_alpha::training::{prepare_phase1a_data, Phase1aConfig}; + +use ml_core::cuda_autograd::gpu_tensor::GpuTensor; + +#[derive(Parser, Debug)] +#[command(name = "phase1d_mamba", about = "FoxhuntQ-Δ Phase 1d.1 Mamba2 sequence smoke")] +struct Cli { + #[arg(long)] fxcache_path: String, + #[arg(long, default_value_t = 100)] horizon: usize, + #[arg(long, default_value_t = 32)] seq_len: usize, + #[arg(long, default_value_t = 64)] hidden_dim: usize, + #[arg(long, default_value_t = 16)] state_dim: usize, + #[arg(long, default_value_t = 5)] epochs: usize, + #[arg(long, default_value_t = 256)] batch_size: usize, + #[arg(long, default_value_t = 1e-3)] lr: f32, + #[arg(long, default_value_t = 42)] seed: u64, + /// Subsample stride for train sequences (1 = use every starting position; + /// larger = subsample). Memory + time both scale linearly with this. + #[arg(long, default_value_t = 4)] train_stride: usize, +} + +fn main() -> Result<()> { + tracing_subscriber::fmt() + .with_env_filter(EnvFilter::try_from_default_env().unwrap_or_else(|_| EnvFilter::new("warn"))) + .init(); + let cli = Cli::parse(); + let ctx = CudaContext::new(0).context("init CUDA")?; + let stream = ctx.default_stream(); + + // ── Load fxcache, build features+labels (mirrors phase1a) ────────── + let mut cfg_p1a = Phase1aConfig::default(); + cfg_p1a.fxcache_path = cli.fxcache_path.clone(); + cfg_p1a.horizon = cli.horizon; + let reader = FxCacheReader::open(&cfg_p1a.fxcache_path)?; + let alpha_dim = reader.alpha_feature_dim() + .ok_or_else(|| anyhow::anyhow!("fxcache lacks alpha column"))?; + cfg_p1a.mlp.in_dim = alpha_dim; + let split = PurgedSplit::new(reader.bar_count(), cfg_p1a.train_frac, + cfg_p1a.horizon, cfg_p1a.embargo_bars)?.split(); + let data = prepare_phase1a_data(&reader, &split, &cfg_p1a)?; + println!( + "fxcache: {} rows, alpha_dim={}, n_train={}, n_val={}", + reader.bar_count(), alpha_dim, data.train_indices.len(), data.val_indices.len() + ); + + // ── Build sequence start indices ────────────────────────────────── + // A sequence ending at bar t requires t-seq_len+1 ≥ 0 and t < bar_count. + // We filter train/val to bars where the preceding seq_len-1 bars are + // within the same train (or val) range. For simplicity we use bars + // where `idx ≥ seq_len - 1`. + let seq_len = cli.seq_len; + let make_seq_starts = |indices: &[usize], stride: usize| -> Vec { + indices.iter() + .filter(|&&i| i + 1 >= seq_len) + .step_by(stride) + .map(|&i| i + 1 - seq_len) // start bar of the seq_len-bar window + .collect() + }; + let train_starts = make_seq_starts(&data.train_indices, cli.train_stride); + let val_starts = make_seq_starts(&data.val_indices, 1); + println!("train sequences (stride={}): {}", cli.train_stride, train_starts.len()); + println!("val sequences: {}", val_starts.len()); + + // ── Build Mamba2 + AdamW ────────────────────────────────────────── + let block_cfg = Mamba2BlockConfig { + in_dim: alpha_dim, + hidden_dim: cli.hidden_dim, + state_dim: cli.state_dim, + seq_len, + }; + let mut block = Mamba2Block::new(block_cfg.clone(), Arc::clone(&stream))?; + let mut opt = Mamba2AdamW::new(&block, Mamba2AdamWConfig { + lr: cli.lr, + ..Default::default() + })?; + println!("Mamba2Block params: {}", block.param_count()); + + // ── Training loop ───────────────────────────────────────────────── + let mut rng = ChaCha8Rng::seed_from_u64(cli.seed); + let train_len = train_starts.len(); + if train_len < cli.batch_size { + anyhow::bail!("n_train_sequences ({}) < batch_size ({})", train_len, cli.batch_size); + } + let batches_per_epoch = train_len / cli.batch_size; + + let bce_loss = |logits: &[f32], labels: &[f32]| -> f32 { + let mut s = 0.0_f32; + let eps = 1e-7_f32; + for (&z, &y) in logits.iter().zip(labels.iter()) { + let z = z.clamp(-50.0, 50.0); + let p = (1.0 / (1.0 + (-z).exp())).clamp(eps, 1.0 - eps); + s += -(y * p.ln() + (1.0 - y) * (1.0 - p).ln()); + } + s / logits.len() as f32 + }; + + for epoch in 0..cli.epochs { + let mut perm: Vec = (0..train_len).collect(); + for i in (1..train_len).rev() { + let j = rng.gen_range(0..=i); + perm.swap(i, j); + } + let mut loss_sum = 0.0_f32; + for batch_idx in 0..batches_per_epoch { + let sel = &perm[batch_idx * cli.batch_size .. (batch_idx + 1) * cli.batch_size]; + // Pull labels from train_labels via end-bar lookup. We rebuilt gather_batch + // to read train_labels for sequence END bars (sequence labels = last-position). + let (input, labels) = gather_batch_for_train( + &train_starts, sel, seq_len, alpha_dim, + &data.feature_matrix, &data.train_labels, &data.train_indices, &stream + )?; + let (logit, cache) = block.forward_train(&input)?; + let logit_host = logit.to_host(&stream)?; + let loss = bce_loss(&logit_host, &labels); + loss_sum += loss; + + // d_logit = (sigmoid(z) - y) / N for mean-BCE-with-logits. + let n_b = labels.len() as f32; + let d_logit_host: Vec = logit_host.iter().zip(labels.iter()) + .map(|(&z, &y)| { + let p = 1.0 / (1.0 + (-z.clamp(-50.0, 50.0)).exp()); + (p - y) / n_b + }) + .collect(); + let d_logit_dev = stream.clone_htod(&d_logit_host)?; + let d_logit = GpuTensor::new(d_logit_dev, vec![labels.len(), 1])?; + let grads = block.backward(&cache, &d_logit)?; + opt.step(&mut block, &grads)?; + } + let mean_loss = loss_sum / batches_per_epoch as f32; + println!("epoch {epoch:2} mean_train_bce={mean_loss:.5} n_batches={batches_per_epoch}"); + } + + // ── Validation pass ─────────────────────────────────────────────── + let val_chunk = cli.batch_size; + let mut val_logits: Vec = Vec::with_capacity(val_starts.len()); + let mut val_labels_y: Vec = Vec::with_capacity(val_starts.len()); + let mut i = 0; + while i < val_starts.len() { + let this = val_chunk.min(val_starts.len() - i); + let sel: Vec = (i..i + this).collect(); + let (input, labels) = gather_batch_for_val( + &val_starts, &sel, seq_len, alpha_dim, + &data.feature_matrix, &data.val_labels, &data.val_indices, &stream + )?; + let logit = block.forward(&input)?; + let logit_host = logit.to_host(&stream)?; + val_logits.extend_from_slice(&logit_host); + val_labels_y.extend_from_slice(&labels); + i += this; + } + let labels_u8: Vec = val_labels_y.iter().map(|&y| if y > 0.5 { 1 } else { 0 }).collect(); + let acc = accuracy_from_logits(&val_logits, &labels_u8); + let auc = auc_from_logits(&val_logits, &labels_u8); + println!(); + println!("==========================================================="); + println!("PHASE 1d.1 — MAMBA2 SMOKE"); + println!("==========================================================="); + println!("val sequences: {}", val_logits.len()); + println!("val accuracy: {:.4}", acc); + println!("val AUC: {:.4}", auc); + println!("MLP baseline (Phase 1c K=100): AUC=0.6849"); + if auc > 0.72 { + println!("GATE PASS: sequence-state lift; proceed to Phase 1d.2 (multi-minute horizon)."); + } else if auc > 0.6849 { + println!("MARGINAL: above MLP baseline ({:.4} > 0.6849) but below gate (0.72). Inspect.", auc); + } else { + println!("GATE FAIL: AUC ≤ MLP baseline 0.6849; no sequence-state lift."); + } + Ok(()) +} + +/// Gather a training batch. Sequence labels come from `train_labels`, which +/// is indexed *positionally* (the i-th train_label corresponds to the i-th +/// train_indices entry). We map start → end_bar → position in train_indices. +fn gather_batch_for_train( + starts: &[usize], + sel: &[usize], + seq_len: usize, + alpha_dim: usize, + feature_matrix: &[f32], + train_labels: &[f32], + train_indices: &[usize], + stream: &Arc, +) -> Result<(GpuTensor, Vec)> { + let b = sel.len(); + let mut host = Vec::with_capacity(b * seq_len * alpha_dim); + let mut labels = Vec::with_capacity(b); + for &k in sel { + let s = starts[k]; + let end_bar = s + seq_len - 1; + for t in 0..seq_len { + let off = (s + t) * alpha_dim; + host.extend_from_slice(&feature_matrix[off..off + alpha_dim]); + } + // Find end_bar in train_indices via binary search (train_indices is + // strictly ascending in the prepare_phase1a_data implementation). + let pos = train_indices.binary_search(&end_bar).ok(); + labels.push(pos.map(|p| train_labels[p]).unwrap_or(0.5)); + } + let dev = stream.clone_htod(&host)?; + let t = GpuTensor::new(dev, vec![b, seq_len, alpha_dim]) + .map_err(|e| anyhow::anyhow!("batch tensor: {e}"))?; + Ok((t, labels)) +} + +fn gather_batch_for_val( + starts: &[usize], + sel: &[usize], + seq_len: usize, + alpha_dim: usize, + feature_matrix: &[f32], + val_labels: &[f32], + val_indices: &[usize], + stream: &Arc, +) -> Result<(GpuTensor, Vec)> { + let b = sel.len(); + let mut host = Vec::with_capacity(b * seq_len * alpha_dim); + let mut labels = Vec::with_capacity(b); + for &k in sel { + let s = starts[k]; + let end_bar = s + seq_len - 1; + for t in 0..seq_len { + let off = (s + t) * alpha_dim; + host.extend_from_slice(&feature_matrix[off..off + alpha_dim]); + } + let pos = val_indices.binary_search(&end_bar).ok(); + labels.push(pos.map(|p| val_labels[p]).unwrap_or(0.5)); + } + let dev = stream.clone_htod(&host)?; + let t = GpuTensor::new(dev, vec![b, seq_len, alpha_dim]) + .map_err(|e| anyhow::anyhow!("batch tensor: {e}"))?; + Ok((t, labels)) +}