diff --git a/.gitignore b/.gitignore index d2fd7f2b7..2dbd9510f 100644 --- a/.gitignore +++ b/.gitignore @@ -189,6 +189,8 @@ test_data/databento/samples/ *.pt *.pth *.bin +# Exception: deterministic test fixtures (golden bytes for bit-equivalence) +!crates/ml-alpha/tests/fixtures/*.bin # Load test results services/*/load_tests/results/ services/*/load_tests/*.json diff --git a/crates/ml-alpha/src/trainer/perception.rs b/crates/ml-alpha/src/trainer/perception.rs index b624c13de..1488ce6dc 100644 --- a/crates/ml-alpha/src/trainer/perception.rs +++ b/crates/ml-alpha/src/trainer/perception.rs @@ -497,6 +497,20 @@ impl PerceptionTrainer { pub fn new(dev: &MlDevice, cfg: &PerceptionTrainerConfig) -> Result { anyhow::ensure!(cfg.seq_len >= 2, "Mamba2 requires seq_len >= 2"); + // Seed GPU-side weight init (Mamba2 stacks initialise their weights via + // `OwnedGpuLinear::xavier`, which delegates to + // `ml_core::cuda_autograd::init::generate_uniform`). Without the + // thread-local guard, that helper seeds from SystemTime+thread-id and + // produces non-deterministic weights — every process gets fresh values. + // CfC/VSN/heads below use an explicit `ChaCha8Rng::seed_from_u64(cfg.seed)` + // chain already; the guard closes the remaining init paths so the + // whole trainer is reproducible from `cfg.seed`. + // + // Guard is dropped at the end of `new`, restoring default time-based + // seeding for any later xavier callers outside the trainer (e.g., + // dqn / ppo crates that also use OwnedGpuLinear::xavier). + let _seed_guard = ml_core::cuda_autograd::init::scoped_init_seed(cfg.seed); + let stream = dev.cuda_stream().context("trainer stream")?.clone(); let ctx = dev.cuda_context().context("trainer ctx")?; // Disable cudarc's automatic per-allocation read/write event diff --git a/crates/ml-alpha/tests/fixtures/perception_forward_golden.bin b/crates/ml-alpha/tests/fixtures/perception_forward_golden.bin new file mode 100644 index 000000000..5c2dcfe4e Binary files /dev/null and b/crates/ml-alpha/tests/fixtures/perception_forward_golden.bin differ diff --git a/crates/ml-alpha/tests/perception_forward_golden.rs b/crates/ml-alpha/tests/perception_forward_golden.rs new file mode 100644 index 000000000..939ca8d85 --- /dev/null +++ b/crates/ml-alpha/tests/perception_forward_golden.rs @@ -0,0 +1,168 @@ +//! Bit-equivalence gate for the v2 trunk-grows refactor. +//! +//! Captures a deterministic snapshot of `PerceptionTrainer.evaluate` +//! output on a 32-snapshot fixture with seed=42. Every refactor commit +//! X1..X11 must reproduce these bytes to `max_abs_diff < 1e-4` +//! (cross-GPU tolerance per spec §2.4). +//! +//! Runtime invocation (GPU required): +//! +//! SQLX_OFFLINE=true cargo test -p ml-alpha \ +//! --test perception_forward_golden -- --ignored --nocapture +//! +//! First run: writes the golden file at +//! crates/ml-alpha/tests/fixtures/perception_forward_golden.bin +//! Subsequent runs: re-runs the forward path and asserts bit-equivalence +//! against the golden. +//! +//! Golden bytes were captured on RTX 3050 (sm_86). L40S runs use the +//! same tight-tolerance comparison; cross-GPU drift within 1e-4 is +//! accepted per spec §2.4. + +use std::fs; +use std::path::PathBuf; + +use anyhow::{Context, Result}; +use ml_alpha::cfc::snap_features::{Mbp10RawInput, REGIME_DIM}; +use ml_alpha::heads::N_HORIZONS; +use ml_alpha::trainer::perception::{PerceptionTrainer, PerceptionTrainerConfig}; +use ml_core::device::MlDevice; +use rand::{Rng, SeedableRng}; + +const SEQ_LEN: usize = 32; +const SEED: u64 = 42; +const GOLDEN_PATH: &str = "tests/fixtures/perception_forward_golden.bin"; + +/// Build a deterministic 32-snapshot fixture using a fixed PRNG seed. +/// Output bytes are stable across runs on the same GPU model. +fn fixture_snapshots() -> Vec { + let mut rng = rand::rngs::StdRng::seed_from_u64(SEED); + let mut prev_mid = 100.0_f32; + let mut prev_ts_ns = 1_000_000_000_u64; + (0..SEQ_LEN) + .map(|i| { + let mid_jitter: f32 = rng.gen_range(-0.5_f32..0.5_f32); + let mid = 100.0 + mid_jitter; + let ts_ns = prev_ts_ns + 1_000_000; + let bid_px: [f32; 10] = std::array::from_fn(|j| mid - 0.125 - (j as f32) * 0.25); + let ask_px: [f32; 10] = std::array::from_fn(|j| mid + 0.125 + (j as f32) * 0.25); + let bid_sz: [f32; 10] = [10.0; 10]; + let ask_sz: [f32; 10] = [10.0; 10]; + let regime: [f32; REGIME_DIM] = std::array::from_fn(|_| rng.gen_range(-1.0_f32..1.0_f32)); + let trade_signed_vol: f32 = rng.gen_range(-5.0_f32..5.0_f32); + let trade_count: u32 = rng.gen_range(0_u32..50_u32); + let snap = Mbp10RawInput { + bid_px, + bid_sz, + ask_px, + ask_sz, + prev_mid, + trade_signed_vol, + trade_count, + ts_ns, + prev_ts_ns, + regime, + }; + prev_mid = mid; + prev_ts_ns = ts_ns; + let _ = i; + snap + }) + .collect() +} + +/// Deterministic labels — random in [0, 1] per (position, horizon). +fn fixture_labels() -> Vec<[f32; N_HORIZONS]> { + let mut rng = rand::rngs::StdRng::seed_from_u64(SEED.wrapping_add(1)); + (0..SEQ_LEN) + .map(|_| { + let mut row = [0.0_f32; N_HORIZONS]; + for x in &mut row { + *x = rng.gen_range(0.0_f32..1.0_f32); + } + row + }) + .collect() +} + +fn run_forward() -> Result<(f32, Vec)> { + let dev = MlDevice::cuda(0).context("init MlDevice")?; + let cfg = PerceptionTrainerConfig { + seq_len: SEQ_LEN, + n_batch: 1, + mamba2_state_dim: 16, + seed: SEED, + ..Default::default() + }; + let mut trainer = PerceptionTrainer::new(&dev, &cfg).context("trainer init")?; + let snapshots = fixture_snapshots(); + let labels = fixture_labels(); + let (loss, probs) = trainer.evaluate(&snapshots, &labels)?; + Ok((loss, probs)) +} + +fn golden_file_path() -> PathBuf { + PathBuf::from(env!("CARGO_MANIFEST_DIR")).join(GOLDEN_PATH) +} + +#[test] +#[ignore = "GPU required; gates the v2 trunk-grows refactor commits"] +fn golden_matches() -> Result<()> { + let (loss, probs) = run_forward()?; + let path = golden_file_path(); + if let Some(parent) = path.parent() { + fs::create_dir_all(parent) + .with_context(|| format!("create fixture dir {}", parent.display()))?; + } + if !path.exists() { + // First run: capture the golden. + let mut bytes = Vec::with_capacity(4 + probs.len() * 4); + bytes.extend_from_slice(&loss.to_le_bytes()); + for p in &probs { + bytes.extend_from_slice(&p.to_le_bytes()); + } + fs::write(&path, &bytes) + .with_context(|| format!("write golden {}", path.display()))?; + eprintln!( + "WROTE golden: {} bytes, {} probs at {}", + bytes.len(), + probs.len(), + path.display() + ); + return Ok(()); + } + let golden = fs::read(&path).context("read golden")?; + let expected_size = 4 + probs.len() * 4; + anyhow::ensure!( + golden.len() == expected_size, + "golden size {} != expected {} (probs len changed?)", + golden.len(), + expected_size + ); + let golden_loss = f32::from_le_bytes(golden[0..4].try_into().unwrap()); + let mut max_diff: f32 = (loss - golden_loss).abs(); + let mut max_diff_idx: usize = usize::MAX; + for (i, p) in probs.iter().enumerate() { + let off = 4 + i * 4; + let g = f32::from_le_bytes(golden[off..off + 4].try_into().unwrap()); + let d = (p - g).abs(); + if d > max_diff { + max_diff = d; + max_diff_idx = i; + } + } + let tol = 1e-4_f32; + eprintln!( + "golden loss={:.6e} run loss={:.6e} max_diff={:.6e} (idx={}) tol={:.0e}", + golden_loss, loss, max_diff, max_diff_idx, tol + ); + anyhow::ensure!( + max_diff < tol, + "max_abs_diff {:.6e} >= tolerance {:.0e} — refactor regression vs golden \ + (largest diff at prob[{}])", + max_diff, + tol, + max_diff_idx + ); + Ok(()) +} diff --git a/crates/ml-core/src/cuda_autograd/init.rs b/crates/ml-core/src/cuda_autograd/init.rs index b91e0d924..3bb378dbf 100644 --- a/crates/ml-core/src/cuda_autograd/init.rs +++ b/crates/ml-core/src/cuda_autograd/init.rs @@ -4,13 +4,69 @@ //! and upload to GPU. For typical ML weight sizes (thousands to millions of //! elements) the CPU generation + HtoD transfer is negligible compared to //! a full training step, and avoids the complexity of a GPU RNG kernel. +//! +//! ## Deterministic init for reproducibility +//! +//! By default, `generate_uniform` seeds its PRNG from `SystemTime::now()` +//! + thread id, giving each process fresh weights. Callers that need +//! reproducibility (bit-equivalence tests, regression suites, the +//! `PerceptionTrainer::new(cfg)` path with a fixed `cfg.seed`) opt in via +//! [`scoped_init_seed(seed)`], which installs a thread-local seedable +//! RNG that overrides the default for the lifetime of the returned +//! guard. Drop the guard to restore the default behavior. +use std::cell::RefCell; use std::sync::Arc; use cudarc::driver::{CudaSlice, CudaStream}; +use rand::rngs::StdRng; +use rand::{Rng, SeedableRng}; use crate::MLError; +thread_local! { + /// When `Some`, all `generate_uniform` calls on this thread draw + /// from this seeded RNG instead of the time-based default. The RNG + /// chain advances across calls — multiple init-tensor draws share + /// one chain, so a single seed determines the full init for the + /// caller's session. + static INIT_RNG: RefCell> = const { RefCell::new(None) }; +} + +/// Install a thread-local seeded RNG that overrides +/// `generate_uniform`'s default time-based seed. Returns a guard that +/// clears the thread-local on drop — use `let _g = scoped_init_seed(s);` +/// at the top of any function that needs deterministic init for the +/// duration of its scope. +/// +/// **Why a guard:** RAII ensures the deterministic seed never leaks +/// out of the function that installed it. Nested calls within the +/// same thread share the chain (advancing it), which is what callers +/// want (every init pulls from the same seeded stream in deterministic +/// order). +#[must_use = "scoped_init_seed returns a guard; drop it to restore default seeding"] +pub fn scoped_init_seed(seed: u64) -> InitSeedGuard { + INIT_RNG.with(|cell| { + *cell.borrow_mut() = Some(StdRng::seed_from_u64(seed)); + }); + InitSeedGuard { _private: () } +} + +/// RAII guard returned by [`scoped_init_seed`]. Restores the default +/// time-based seeding on drop. +#[derive(Debug)] +pub struct InitSeedGuard { + _private: (), +} + +impl Drop for InitSeedGuard { + fn drop(&mut self) { + INIT_RNG.with(|cell| { + *cell.borrow_mut() = None; + }); + } +} + /// Initialize a `CudaSlice` with Xavier uniform values. /// /// Samples from U(-limit, limit) where `limit = sqrt(6 / (fan_in + fan_out))`. @@ -81,28 +137,49 @@ pub fn near_zero_xavier( /// Generate `n` uniform random f32 values in `[lo, hi)`. /// -/// Uses a simple xoshiro256++ PRNG seeded from thread_rng for speed. -/// This runs on CPU — acceptable for one-time initialization. +/// If a [`scoped_init_seed`] guard is active on this thread, draws from +/// the thread-local seeded RNG (reproducible across processes given the +/// same seed). Otherwise falls back to a time+thread-id-seeded +/// xoshiro256++ chain (the original behavior — fresh weights per process). fn generate_uniform(n: usize, lo: f64, hi: f64) -> Vec { + let range = hi - lo; + let mut out = Vec::with_capacity(n); + + let consumed_from_thread_local = INIT_RNG.with(|cell| { + let mut maybe_rng = cell.borrow_mut(); + if let Some(rng) = maybe_rng.as_mut() { + for _ in 0..n { + let u: f64 = rng.gen_range(0.0_f64..1.0_f64); + out.push((lo + u * range) as f32); + } + true + } else { + false + } + }); + if consumed_from_thread_local { + return out; + } + + // Default path: time+thread-id-seeded xoshiro256++ (unchanged from + // pre-deterministic-opt-in behavior). use std::collections::hash_map::DefaultHasher; use std::hash::{Hash, Hasher}; use std::time::SystemTime; - // Seed from time + thread id for reasonable uniqueness let mut hasher = DefaultHasher::new(); SystemTime::now().hash(&mut hasher); std::thread::current().id().hash(&mut hasher); let seed = hasher.finish(); - // Simple xoshiro256++ state seeded from the hash - let mut s = [seed, seed.wrapping_mul(6364136223846793005).wrapping_add(1), - seed.wrapping_mul(1442695040888963407).wrapping_add(3), - seed.wrapping_mul(2891336453).wrapping_add(7)]; + let mut s = [ + seed, + seed.wrapping_mul(6364136223846793005).wrapping_add(1), + seed.wrapping_mul(1442695040888963407).wrapping_add(3), + seed.wrapping_mul(2891336453).wrapping_add(7), + ]; - let range = hi - lo; - let mut out = Vec::with_capacity(n); for _ in 0..n { - // xoshiro256++ step let result = s[0].wrapping_add(s[3]).rotate_left(23).wrapping_add(s[0]); let t = s[1] << 17; s[2] ^= s[0]; @@ -112,7 +189,6 @@ fn generate_uniform(n: usize, lo: f64, hi: f64) -> Vec { s[2] ^= t; s[3] = s[3].rotate_left(45); - // Convert to f64 in [0, 1) then scale to [lo, hi) let f = (result >> 11) as f64 / (1_u64 << 53) as f64; out.push((lo + f * range) as f32); }