From b47b2fabfb26511632c69099a5d1a1b14096acb2 Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Tue, 19 May 2026 01:19:13 +0200 Subject: [PATCH] fix(ml-core): deterministic GPU weight init via scoped_init_seed MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ml_core::cuda_autograd::init::generate_uniform (backing xavier_uniform, kaiming_uniform, bias_uniform, near_zero_xavier) defaulted to seeding from SystemTime::now() + thread_id, producing non-reproducible weights across processes. Mamba2 stacks initialise via OwnedGpuLinear::xavier, which routes through this helper — so PerceptionTrainer.evaluate output diverged 5-30% across fresh-process runs with identical cfg.seed. Fix: thread-local seedable RNG override. New API: let _g = ml_core::cuda_autograd::init::scoped_init_seed(seed); // ... all xavier/kaiming/bias/near_zero calls draw from // StdRng::seed_from_u64(seed) chain while _g is alive ... // _g dropped here -> restores default time-based seeding PerceptionTrainer::new now installs the guard before any Mamba2Block construction, so the trainer is reproducible from cfg.seed end-to-end. CfC/VSN/heads already used explicit ChaCha8Rng::seed_from_u64 — only Mamba2 was affected. Production behavior unchanged when no guard is set. ml-core: 306 tests pass, ml-alpha: 34 lib tests pass. Regression test: crates/ml-alpha/tests/perception_forward_golden.rs captures bit-exact PerceptionTrainer.evaluate output (loss + 160 probs on a deterministic seed=42 fixture) into a 644-byte golden file. Three consecutive runs now produce max_abs_diff=0; pre-fix runs varied by 0.1-0.3 absolute on individual probs. .gitignore: added exception for crates/ml-alpha/tests/fixtures/*.bin so deterministic test fixtures land in repo. Per pearl_scoped_init_seed_for_reproducibility in project memory. Co-Authored-By: Claude Opus 4.7 (1M context) --- .gitignore | 2 + crates/ml-alpha/src/trainer/perception.rs | 14 ++ .../fixtures/perception_forward_golden.bin | Bin 0 -> 644 bytes .../tests/perception_forward_golden.rs | 168 ++++++++++++++++++ crates/ml-core/src/cuda_autograd/init.rs | 98 ++++++++-- 5 files changed, 271 insertions(+), 11 deletions(-) create mode 100644 crates/ml-alpha/tests/fixtures/perception_forward_golden.bin create mode 100644 crates/ml-alpha/tests/perception_forward_golden.rs 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 0000000000000000000000000000000000000000..5c2dcfe4eb4adb739c412423793c87615421a19b GIT binary patch literal 644 zcmV~$do0ue00(e)=p0Uze&=^RM9f;xKU7QG_uIqHWVMl*O0BWwvDA8~C0UJ%HifO# zTG4B_d9PxIc_nf#TTEk@DCaeK{Df|E4}CtNV+jZEd;0QCZUw#yrMMd6!0Us7j7t_7 zv9B3n)g`cIE9rN|pE;L(ctqO)H)RR5CTCXp`?7IVV12UzxfP}8&y+K7VLc=31)4Ly zAv3rFBX1qpUo3K_SfF=$A2wH4VW_|k`wKO9O$i)J96-#T54en5_AaCR(HaJvQS-cN1U=ccu(U73@3a^Hb*mYlG>YDg zI)vP`pfk#ohb?NxC5*$C&;Uvn{7hNRyD>ftYM4O5&PEhXTVOrv!8>{%=1)$d=UWrB z&#myNQc-L2;bPDfT!j|QH(7CU*qwSuk&1{vkUwd~P_hk8f$l8ZCNk~SG@N=q!Q5_x z|1~!rKPS={FpF5(XT+YhW6I*;zCLU9ALaE`yS_v(WJF!yZz{*7z-uNxDzDPpHQYW5y$XMFu!dedv z%X%bq3wrF{D&thN3*(kVI-4bo?bpK?CS%rG7oIYR6y_w*+KF;+8T0BDw6u%NSIwjF zo*oy{WgLB=VE=28FJ~o)Iz9I4WK2G(V9K&c-8%`v=t8_%#?(U!26}1OkuM>&r3-0s pa?Tu3@PtWZxJH7quM2g#avD_%%5RB$UpNPwlK~qm<&2zJ#s9#U7$X1x literal 0 HcmV?d00001 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); }