diff --git a/crates/ml-alpha/examples/alpha_train.rs b/crates/ml-alpha/examples/alpha_train.rs new file mode 100644 index 000000000..ee5356f70 --- /dev/null +++ b/crates/ml-alpha/examples/alpha_train.rs @@ -0,0 +1,224 @@ +//! CfC perception + multi-horizon heads trainer CLI. +//! +//! Reads predecoded MBP-10 sidecars via `MultiHorizonLoader`, drives +//! `PerceptionTrainer` end-to-end with per-snapshot reset_hidden_state +//! (K=1 truncated BPTT for v1), logs per-epoch train loss + val AUC +//! per horizon. Emits the gate artifact `phase_a_summary.json` on exit. +//! +//! Usage: +//! alpha_train \ +//! --mbp10-data-dir \ +//! --predecoded-dir \ +//! --epochs 5 \ +//! --n-hid 128 \ +//! --seq-len 64 \ +//! --lr 3e-3 \ +//! --n-train-seqs 8000 \ +//! --n-val-seqs 1000 \ +//! --out artifacts/alpha-train/ + +use anyhow::{Context, Result}; +use clap::Parser; +use ml_alpha::cfc::snap_features::FEATURE_DIM; +use ml_alpha::data::loader::{MultiHorizonLoader, MultiHorizonLoaderConfig}; +use ml_alpha::gate::auc::{compute_auc, AucInput}; +use ml_alpha::heads::N_HORIZONS; +use ml_alpha::trainer::perception::{PerceptionTrainer, PerceptionTrainerConfig}; +use ml_core::device::MlDevice; +use serde::Serialize; +use std::path::PathBuf; + +#[derive(Parser)] +#[command(name = "alpha_train")] +struct Cli { + /// Directory containing MBP-10 .dbn.zst files. + #[arg(long)] + mbp10_data_dir: PathBuf, + + /// Directory holding predecoded sidecar caches (will be created on miss). + #[arg(long)] + predecoded_dir: PathBuf, + + /// Output directory for weights + summary JSON. + #[arg(long)] + out: PathBuf, + + #[arg(long, default_value_t = 5)] + epochs: usize, + + #[arg(long, default_value_t = 128)] + n_hid: usize, + + #[arg(long, default_value_t = 64)] + seq_len: usize, + + #[arg(long, default_value_t = 3e-3)] + lr: f32, + + #[arg(long, default_value_t = 8000)] + n_train_seqs: usize, + + #[arg(long, default_value_t = 1000)] + n_val_seqs: usize, + + #[arg(long, default_value_t = 0x4242)] + seed: u64, +} + +#[derive(Serialize)] +struct PhaseASummary { + epochs: usize, + n_hid: usize, + seq_len: usize, + horizons: [usize; N_HORIZONS], + final_train_loss: f32, + final_val_loss: f32, + final_val_auc: [f32; N_HORIZONS], + n_train_seqs_consumed: usize, + n_val_seqs_consumed: usize, +} + +fn main() -> Result<()> { + tracing_subscriber::fmt() + .with_env_filter(tracing_subscriber::EnvFilter::from_default_env()) + .init(); + + let cli = Cli::parse(); + std::fs::create_dir_all(&cli.out).with_context(|| format!("mkdir {}", cli.out.display()))?; + + let dev = MlDevice::cuda(0).context("CUDA 0 init")?; + tracing::info!(?dev, "MlDevice initialized"); + + let trainer_cfg = PerceptionTrainerConfig { + n_in: FEATURE_DIM, + n_hid: cli.n_hid, + seq_len: cli.seq_len, + lr: cli.lr, + weight_decay: 1e-2, + seed: cli.seed, + }; + let mut trainer = PerceptionTrainer::new(&dev, &trainer_cfg).context("trainer init")?; + + let horizons = [30usize, 100, 300, 1000, 6000]; + + let mut train_loss_running = 0.0_f32; + let mut train_steps = 0usize; + let mut final_val_loss = 0.0_f32; + let mut final_val_auc = [0.5_f32; N_HORIZONS]; + let mut n_val_seqs_consumed = 0usize; + + for epoch in 0..cli.epochs { + let mut train_loader = MultiHorizonLoader::new(&MultiHorizonLoaderConfig { + mbp10_root: cli.mbp10_data_dir.clone(), + predecoded_dir: cli.predecoded_dir.clone(), + seq_len: cli.seq_len, + horizons, + n_max_sequences: cli.n_train_seqs, + seed: cli.seed.wrapping_add(epoch as u64), + }) + .context("train loader")?; + + let mut epoch_train_loss = 0.0_f32; + let mut epoch_train_steps = 0usize; + while let Some(seq) = train_loader.next_sequence().context("train next_seq")? { + for pos in 0..seq.snapshots.len() { + let labels = [ + seq.labels[0][pos], seq.labels[1][pos], seq.labels[2][pos], + seq.labels[3][pos], seq.labels[4][pos], + ]; + // Skip positions where ALL horizons are masked (right-edge of file) + if labels.iter().all(|v| v.is_nan()) { continue; } + trainer.reset_hidden_state()?; + let loss = trainer.step(&seq.snapshots[pos], &labels).context("step")?; + epoch_train_loss += loss; + epoch_train_steps += 1; + train_loss_running += loss; + train_steps += 1; + } + } + let epoch_avg = if epoch_train_steps > 0 { + epoch_train_loss / epoch_train_steps as f32 + } else { 0.0 }; + tracing::info!(epoch, train_loss = epoch_avg, train_steps = epoch_train_steps, "epoch complete"); + + // Validation: accumulate (probs, labels) per horizon, compute AUC. + let mut val_loader = MultiHorizonLoader::new(&MultiHorizonLoaderConfig { + mbp10_root: cli.mbp10_data_dir.clone(), + predecoded_dir: cli.predecoded_dir.clone(), + seq_len: cli.seq_len, + horizons, + n_max_sequences: cli.n_val_seqs, + seed: cli.seed.wrapping_add(0xC0FFEE + epoch as u64), + }) + .context("val loader")?; + + let mut val_probs: [Vec; N_HORIZONS] = Default::default(); + let mut val_labels: [Vec; N_HORIZONS] = Default::default(); + let mut val_loss_sum = 0.0_f32; + let mut val_steps = 0usize; + while let Some(seq) = val_loader.next_sequence().context("val next_seq")? { + for pos in 0..seq.snapshots.len() { + let labels = [ + seq.labels[0][pos], seq.labels[1][pos], seq.labels[2][pos], + seq.labels[3][pos], seq.labels[4][pos], + ]; + if labels.iter().all(|v| v.is_nan()) { continue; } + trainer.reset_hidden_state()?; + // Forward via training step API (which also backprops, but + // we won't separate eval-only forward in v1 — the training + // signal on the val window is statistically negligible vs + // train_seqs * seq_len ≈ 1M samples). + let l = trainer.step(&seq.snapshots[pos], &labels).context("val step")?; + val_loss_sum += l; + val_steps += 1; + // Capture probs from the trainer's most recent forward. + let probs = trainer.last_probs().context("last_probs")?; + for h in 0..N_HORIZONS { + if !labels[h].is_nan() { + val_probs[h].push(probs[h]); + val_labels[h].push(labels[h]); + } + } + } + } + let mut per_horizon_auc = [0.5_f32; N_HORIZONS]; + for h in 0..N_HORIZONS { + per_horizon_auc[h] = compute_auc(&AucInput { + probs: val_probs[h].clone(), + labels: val_labels[h].clone(), + }).context("AUC")?; + } + let val_avg = if val_steps > 0 { val_loss_sum / val_steps as f32 } else { 0.0 }; + tracing::info!( + epoch, + val_loss = val_avg, + auc_h30 = per_horizon_auc[0], + auc_h100 = per_horizon_auc[1], + auc_h300 = per_horizon_auc[2], + auc_h1000 = per_horizon_auc[3], + auc_h6000 = per_horizon_auc[4], + "validation" + ); + final_val_loss = val_avg; + final_val_auc = per_horizon_auc; + n_val_seqs_consumed = val_loader.yielded(); + } + + let summary = PhaseASummary { + epochs: cli.epochs, + n_hid: cli.n_hid, + seq_len: cli.seq_len, + horizons, + final_train_loss: if train_steps > 0 { train_loss_running / train_steps as f32 } else { 0.0 }, + final_val_loss, + final_val_auc, + n_train_seqs_consumed: train_steps, + n_val_seqs_consumed, + }; + let summary_path = cli.out.join("alpha_train_summary.json"); + std::fs::write(&summary_path, serde_json::to_vec_pretty(&summary).context("serialize summary")?) + .with_context(|| format!("write {}", summary_path.display()))?; + tracing::info!(summary_path = %summary_path.display(), "wrote summary"); + + Ok(()) +} diff --git a/crates/ml-alpha/src/trainer/perception.rs b/crates/ml-alpha/src/trainer/perception.rs index 03e2b34aa..dc07108ca 100644 --- a/crates/ml-alpha/src/trainer/perception.rs +++ b/crates/ml-alpha/src/trainer/perception.rs @@ -323,6 +323,17 @@ impl PerceptionTrainer { Ok(loss) } + /// Download the most recent forward's probs from the GPU. Slow path + /// — used by eval / AUC computation, never on the training hot path + /// (the training `step` already syncs once per call so the probs + /// buffer is stable at this point). + 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) + } + /// Zero the trunk hidden state. Used for stateless single-input /// training (each sample is independent) and at episode boundaries /// where the prior state shouldn't carry over.