0f5d5c7b4a
CB3+CB4 shipped the kernels with a holding-pattern (grad-zero) call site.
CB5 wires the real calls + per-loss ISV signals.
perception.rs:
- step_batched body: pos_weight computed host-side from
self.last_pos_fraction (n_neg/n_pos clamped to [1.0, 50.0] per E3),
uploaded mapped-pinned to stg_aux_pos_weight_{long,short}
- 4 kernel calls per step (slab mode, K*B*N_AUX_HORIZONS):
* aux_bce_loss_gpu × 2 (prof_long, prof_short) — class-weighted
* aux_huber_masked_loss_gpu × 2 (size_long, size_short) — NaN-mask
from CB1's y_size=NaN at y_prof=0 gives conditional-Huber for free
- Per-loss EMAs replace single aux_huber_ema:
* aux_prof_bce_ema_per_h (BCE EMA per horizon)
* aux_size_huber_ema_per_h (Huber EMA per horizon)
* aux_dir_acc_ema_per_h (unchanged)
- Stop-grad lift condition: aux_prof_bce_ema < 0.4 AND aux_dir_acc > 0.85
for ALL horizons (uses BCE not Huber per E3 — size Huber is
observability-only since the regression scale varies more than the
binary classification quality)
- aux_lift_huber_threshold renamed to aux_lift_prof_bce_threshold
alpha_train.rs:
- AlphaTrainSummary: final_aux_huber_ema_per_h split into
final_aux_prof_bce_ema_per_h + final_aux_size_huber_ema_per_h
- Per-epoch tracing: aux_prof_bce_h{100,300,1000} + aux_size_huber_h{...}
+ aux_dir_acc_h{...} + stop_grad_aux_to_encoder
perception_overfit.rs: synthetic test asserts BCE finite + below ln(2)
chance baseline, size Huber finite, dir_acc >= 0.5.
Synthetic test on RTX 3050:
aux_prof_bce_ema = [0.0142, 0.0142, 0.0142] (30× below threshold)
aux_size_huber_ema = [0.1213, 0.1213, 0.1213] (small)
aux_dir_acc_ema = [1.0, 1.0, 1.0] (perfect)
stop_grad lifted = false (lift fired)
Known limitation (follow-up CB6): both kernels return single joint scalar
over the [K × B × N_AUX_HORIZONS] slab — all per-horizon EMA entries
carry the broadcast joint mean. Per-h gating would require kernel split.
cargo check --workspace --all-targets clean.
cargo test -p ml-alpha --lib: 43 passed.
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
780 lines
30 KiB
Rust
780 lines
30 KiB
Rust
//! PerceptionTrainer synthetic-overfit smoke.
|
||
//!
|
||
//! Mirrors `perception_overfit.rs` but on the stacked trainer. The
|
||
//! signal is constant direction=+1, label=[1; N_HORIZONS]. 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::aux_heads::N_AUX_HORIZONS;
|
||
use ml_alpha::cfc::snap_features::Mbp10RawInput;
|
||
use ml_alpha::heads::N_HORIZONS;
|
||
use ml_alpha::trainer::perception::{PerceptionTrainer, PerceptionTrainerConfig};
|
||
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>, Vec<[f32; N_HORIZONS]>) {
|
||
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,
|
||
regime: [0.0; 6],
|
||
});
|
||
prev_mid = next_mid;
|
||
let _ = k;
|
||
}
|
||
// Per-position labels: every position knows the next K snapshots all
|
||
// move up (synthetic monotone ramp), so label = 1.0 for every horizon
|
||
// at every position. Drives the trainer to learn "always predict 1".
|
||
let labels = vec![[1.0; N_HORIZONS]; seq_len];
|
||
(out, labels)
|
||
}
|
||
|
||
/// SDD-3 CB2: synthetic aux-outcome generator matching the +0.25 price
|
||
/// ramp from `synthetic_seq`. For a constant up-trend the A+B paired
|
||
/// targets are:
|
||
/// y_prof_long = 1.0 (every long position is profitable after 2×cost)
|
||
/// y_prof_short = 0.0 (every short position loses)
|
||
/// y_size_long ≈ +ratio (positive σ-normalized signed return)
|
||
/// y_size_short = NaN (unprofitable side: size target masked)
|
||
/// Returns five vectors of `[N_AUX_HORIZONS]` rows + the per-(dir,horizon)
|
||
/// positive-class fraction expected by the trainer signature.
|
||
fn synthetic_aux_outcomes(
|
||
seq_len: usize,
|
||
) -> AuxOutcomes {
|
||
let prof_long_row = std::array::from_fn::<f32, N_AUX_HORIZONS, _>(|_| 1.0);
|
||
let prof_short_row = std::array::from_fn::<f32, N_AUX_HORIZONS, _>(|_| 0.0);
|
||
// Per-horizon magnitude — exact values don't matter for the smoke;
|
||
// the invariant is sign-correctness so the dir_acc gate measures.
|
||
let size_long_per_h: [f32; N_AUX_HORIZONS] = [0.5, 1.5, 5.0];
|
||
let size_long_row = std::array::from_fn::<f32, N_AUX_HORIZONS, _>(|h| size_long_per_h[h]);
|
||
// Short is unprofitable → size masked with NaN (loader convention).
|
||
let size_short_row = std::array::from_fn::<f32, N_AUX_HORIZONS, _>(|_| f32::NAN);
|
||
// pos_fraction is 1.0 for long (all positive) and 0.0 for short.
|
||
let mut pos_fraction = vec![0.0_f32; 2 * N_AUX_HORIZONS];
|
||
for h in 0..N_AUX_HORIZONS { pos_fraction[h] = 1.0; }
|
||
AuxOutcomes {
|
||
prof_long: vec![prof_long_row; seq_len],
|
||
prof_short: vec![prof_short_row; seq_len],
|
||
size_long: vec![size_long_row; seq_len],
|
||
size_short: vec![size_short_row; seq_len],
|
||
pos_fraction,
|
||
}
|
||
}
|
||
|
||
/// CB2: bundle for the A+B paired aux-outcome targets returned by
|
||
/// `synthetic_aux_outcomes` and consumed by `PerceptionTrainer::step`.
|
||
struct AuxOutcomes {
|
||
prof_long: Vec<[f32; N_AUX_HORIZONS]>,
|
||
prof_short: Vec<[f32; N_AUX_HORIZONS]>,
|
||
size_long: Vec<[f32; N_AUX_HORIZONS]>,
|
||
size_short: Vec<[f32; N_AUX_HORIZONS]>,
|
||
pos_fraction: Vec<f32>,
|
||
}
|
||
|
||
#[test]
|
||
fn stacked_trainer_constructs_cleanly() {
|
||
let dev = test_device();
|
||
let cfg = PerceptionTrainerConfig::default();
|
||
let t = PerceptionTrainer::new(&dev, &cfg).expect("init");
|
||
drop(t);
|
||
}
|
||
|
||
#[test]
|
||
fn stacked_trainer_loss_shrinks_on_constant_signal() {
|
||
let dev = test_device();
|
||
let cfg = PerceptionTrainerConfig {
|
||
seq_len: 16, // smaller for smoke speed; Mamba2 needs >=2
|
||
mamba2_state_dim: 8,
|
||
lr_cfc: 3e-3,
|
||
lr_mamba2: 1e-3,
|
||
seed: 0x4242,
|
||
horizon_weights: [1.0; N_HORIZONS],
|
||
n_batch: 1,
|
||
smoothness_base_lambda: 0.0,
|
||
kernel_step_trace_path: None,
|
||
bucket_warmup_cap_override: None,
|
||
lr_aux: 3e-3,
|
||
};
|
||
let mut trainer = PerceptionTrainer::new(&dev, &cfg).expect("init");
|
||
let aux = synthetic_aux_outcomes(cfg.seq_len);
|
||
|
||
// 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.as_slice(),
|
||
aux.prof_long.as_slice(),
|
||
aux.prof_short.as_slice(),
|
||
aux.size_long.as_slice(),
|
||
aux.size_short.as_slice(),
|
||
&aux.pos_fraction,
|
||
).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.as_slice(),
|
||
aux.prof_long.as_slice(),
|
||
aux.prof_short.as_slice(),
|
||
aux.size_long.as_slice(),
|
||
aux.size_short.as_slice(),
|
||
&aux.pos_fraction,
|
||
).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.as_slice(),
|
||
aux.prof_long.as_slice(),
|
||
aux.prof_short.as_slice(),
|
||
aux.size_long.as_slice(),
|
||
aux.size_short.as_slice(),
|
||
&aux.pos_fraction,
|
||
).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,
|
||
"PerceptionTrainer failed to overfit constant signal: \
|
||
start={initial_avg:.4}, end={final_avg:.4}"
|
||
);
|
||
}
|
||
|
||
/// Verifies the trainer still converges with a distinct seed and lr.
|
||
/// A1: decision_stride removed; dt_s is always 1.0 (event-rate).
|
||
/// The test previously verified stride=4 dt_s dynamics; those are
|
||
/// now unconditionally event-rate, so this becomes a second convergence
|
||
/// smoke with different hyperparameters.
|
||
#[test]
|
||
fn stacked_trainer_loss_shrinks_with_stride_4() {
|
||
let dev = test_device();
|
||
let cfg = PerceptionTrainerConfig {
|
||
seq_len: 16,
|
||
mamba2_state_dim: 8,
|
||
lr_cfc: 3e-3,
|
||
lr_mamba2: 1e-3,
|
||
seed: 0xC4C4,
|
||
horizon_weights: [1.0; N_HORIZONS],
|
||
n_batch: 1,
|
||
smoothness_base_lambda: 0.0,
|
||
kernel_step_trace_path: None,
|
||
bucket_warmup_cap_override: None,
|
||
lr_aux: 3e-3,
|
||
};
|
||
let mut trainer = PerceptionTrainer::new(&dev, &cfg).expect("init");
|
||
let aux = synthetic_aux_outcomes(cfg.seq_len);
|
||
|
||
let mut initial = 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.as_slice(),
|
||
aux.prof_long.as_slice(),
|
||
aux.prof_short.as_slice(),
|
||
aux.size_long.as_slice(),
|
||
aux.size_short.as_slice(),
|
||
&aux.pos_fraction,
|
||
).expect("step");
|
||
initial += l;
|
||
prev_mid = 0.5 * (seq.last().unwrap().bid_px[0] + seq.last().unwrap().ask_px[0]);
|
||
ts = seq.last().unwrap().ts_ns;
|
||
}
|
||
initial /= 8.0;
|
||
eprintln!("stride=4 trainer: initial={initial:.4}");
|
||
|
||
for _ in 0..200 {
|
||
let (seq, labels) = synthetic_seq(cfg.seq_len, prev_mid, ts);
|
||
trainer.step(
|
||
&seq,
|
||
labels.as_slice(),
|
||
aux.prof_long.as_slice(),
|
||
aux.prof_short.as_slice(),
|
||
aux.size_long.as_slice(),
|
||
aux.size_short.as_slice(),
|
||
&aux.pos_fraction,
|
||
).expect("step");
|
||
prev_mid = 0.5 * (seq.last().unwrap().bid_px[0] + seq.last().unwrap().ask_px[0]);
|
||
ts = seq.last().unwrap().ts_ns;
|
||
}
|
||
|
||
let mut final_loss = 0.0_f32;
|
||
for _ in 0..8 {
|
||
let (seq, labels) = synthetic_seq(cfg.seq_len, prev_mid, ts);
|
||
final_loss += trainer.step(
|
||
&seq,
|
||
labels.as_slice(),
|
||
aux.prof_long.as_slice(),
|
||
aux.prof_short.as_slice(),
|
||
aux.size_long.as_slice(),
|
||
aux.size_short.as_slice(),
|
||
&aux.pos_fraction,
|
||
).expect("step");
|
||
prev_mid = 0.5 * (seq.last().unwrap().bid_px[0] + seq.last().unwrap().ask_px[0]);
|
||
ts = seq.last().unwrap().ts_ns;
|
||
}
|
||
final_loss /= 8.0;
|
||
eprintln!("stride=4 trainer: final={final_loss:.4}");
|
||
assert!(
|
||
final_loss < 0.6 * initial || final_loss < 0.5,
|
||
"stride=4 trainer failed to converge: {initial:.4} → {final_loss:.4}"
|
||
);
|
||
}
|
||
|
||
/// Verifies the trainer converges at n_batch=32 — the FIRST test that
|
||
/// exercises the cross-batch reducer code path. Existing
|
||
/// `stacked_trainer_loss_shrinks_*` tests all use n_batch=1 so the new
|
||
/// per-batch scratch + reducer logic was never previously hit.
|
||
#[test]
|
||
fn stacked_trainer_loss_shrinks_at_batch_32() {
|
||
let dev = test_device();
|
||
let cfg = PerceptionTrainerConfig {
|
||
seq_len: 16,
|
||
mamba2_state_dim: 8,
|
||
lr_cfc: 3e-3,
|
||
lr_mamba2: 1e-3,
|
||
seed: 0xB32B,
|
||
horizon_weights: [1.0; N_HORIZONS],
|
||
n_batch: 32,
|
||
smoothness_base_lambda: 0.0,
|
||
kernel_step_trace_path: None,
|
||
bucket_warmup_cap_override: None,
|
||
lr_aux: 3e-3,
|
||
};
|
||
let mut trainer = PerceptionTrainer::new(&dev, &cfg).expect("init");
|
||
|
||
let mut ts_base = 1_000_000u64;
|
||
let mut prev_mid = 5500.0_f32;
|
||
|
||
let make_batch = |prev_mid: &mut f32, ts_base: &mut u64,
|
||
cfg: &PerceptionTrainerConfig|
|
||
-> (Vec<Vec<Mbp10RawInput>>, Vec<Vec<[f32; N_HORIZONS]>>) {
|
||
let mut seqs: Vec<Vec<Mbp10RawInput>> = Vec::with_capacity(cfg.n_batch);
|
||
let mut labels: Vec<Vec<[f32; N_HORIZONS]>> = Vec::with_capacity(cfg.n_batch);
|
||
for _ in 0..cfg.n_batch {
|
||
let (seq, lbl) = synthetic_seq(cfg.seq_len, *prev_mid, *ts_base);
|
||
*prev_mid = 0.5 * (seq.last().unwrap().bid_px[0] + seq.last().unwrap().ask_px[0]);
|
||
*ts_base = seq.last().unwrap().ts_ns;
|
||
seqs.push(seq);
|
||
labels.push(lbl);
|
||
}
|
||
(seqs, labels)
|
||
};
|
||
// SDD-3 CB2: synthetic aux outcomes are seq_len-determined and
|
||
// direction-asymmetric; the shared per-sample rows for both prof-binary
|
||
// and size targets are broadcast across all B samples per step.
|
||
let aux_seq = synthetic_aux_outcomes(cfg.seq_len);
|
||
let prof_long_batch: Vec<Vec<[f32; N_AUX_HORIZONS]>> =
|
||
(0..cfg.n_batch).map(|_| aux_seq.prof_long.clone()).collect();
|
||
let prof_short_batch: Vec<Vec<[f32; N_AUX_HORIZONS]>> =
|
||
(0..cfg.n_batch).map(|_| aux_seq.prof_short.clone()).collect();
|
||
let size_long_batch: Vec<Vec<[f32; N_AUX_HORIZONS]>> =
|
||
(0..cfg.n_batch).map(|_| aux_seq.size_long.clone()).collect();
|
||
let size_short_batch: Vec<Vec<[f32; N_AUX_HORIZONS]>> =
|
||
(0..cfg.n_batch).map(|_| aux_seq.size_short.clone()).collect();
|
||
let prof_long_refs: Vec<&[[f32; N_AUX_HORIZONS]]> =
|
||
prof_long_batch.iter().map(|v| v.as_slice()).collect();
|
||
let prof_short_refs: Vec<&[[f32; N_AUX_HORIZONS]]> =
|
||
prof_short_batch.iter().map(|v| v.as_slice()).collect();
|
||
let size_long_refs: Vec<&[[f32; N_AUX_HORIZONS]]> =
|
||
size_long_batch.iter().map(|v| v.as_slice()).collect();
|
||
let size_short_refs: Vec<&[[f32; N_AUX_HORIZONS]]> =
|
||
size_short_batch.iter().map(|v| v.as_slice()).collect();
|
||
|
||
let mut initial = 0.0_f32;
|
||
for warmup in 0..4 {
|
||
let (seqs, labels) = make_batch(&mut prev_mid, &mut ts_base, &cfg);
|
||
let seq_refs: Vec<&[Mbp10RawInput]> = seqs.iter().map(|s| s.as_slice()).collect();
|
||
let lbl_refs: Vec<&[[f32; N_HORIZONS]]> = labels.iter().map(|l| l.as_slice()).collect();
|
||
let l = trainer
|
||
.step_batched(
|
||
&seq_refs,
|
||
&lbl_refs,
|
||
&prof_long_refs,
|
||
&prof_short_refs,
|
||
&size_long_refs,
|
||
&size_short_refs,
|
||
&aux_seq.pos_fraction,
|
||
)
|
||
.expect("warm step");
|
||
if warmup >= 2 {
|
||
initial += l;
|
||
}
|
||
}
|
||
initial /= 2.0;
|
||
eprintln!("B=32 trainer: initial={initial:.4}");
|
||
|
||
for _ in 0..200 {
|
||
let (seqs, labels) = make_batch(&mut prev_mid, &mut ts_base, &cfg);
|
||
let seq_refs: Vec<&[Mbp10RawInput]> = seqs.iter().map(|s| s.as_slice()).collect();
|
||
let lbl_refs: Vec<&[[f32; N_HORIZONS]]> = labels.iter().map(|l| l.as_slice()).collect();
|
||
trainer
|
||
.step_batched(
|
||
&seq_refs,
|
||
&lbl_refs,
|
||
&prof_long_refs,
|
||
&prof_short_refs,
|
||
&size_long_refs,
|
||
&size_short_refs,
|
||
&aux_seq.pos_fraction,
|
||
)
|
||
.expect("train step");
|
||
}
|
||
|
||
let mut final_loss = 0.0_f32;
|
||
for _ in 0..4 {
|
||
let (seqs, labels) = make_batch(&mut prev_mid, &mut ts_base, &cfg);
|
||
let seq_refs: Vec<&[Mbp10RawInput]> = seqs.iter().map(|s| s.as_slice()).collect();
|
||
let lbl_refs: Vec<&[[f32; N_HORIZONS]]> = labels.iter().map(|l| l.as_slice()).collect();
|
||
final_loss += trainer
|
||
.step_batched(
|
||
&seq_refs,
|
||
&lbl_refs,
|
||
&prof_long_refs,
|
||
&prof_short_refs,
|
||
&size_long_refs,
|
||
&size_short_refs,
|
||
&aux_seq.pos_fraction,
|
||
)
|
||
.expect("final step");
|
||
}
|
||
final_loss /= 4.0;
|
||
eprintln!("B=32 trainer: final={final_loss:.4}");
|
||
assert!(
|
||
final_loss < 0.6 * initial || final_loss < 0.1,
|
||
"B=32 trainer failed to converge: {initial:.4} → {final_loss:.4}"
|
||
);
|
||
}
|
||
|
||
// NOTE on scratch-clears testing:
|
||
//
|
||
// A direct "scratch is zero between steps" test is structurally hard to
|
||
// write against this trainer. The lifecycle is:
|
||
// step 1: uncaptured warmup execute → scratch ends with grad
|
||
// step 2: begin_capture/end_capture → records kernels but does
|
||
// NOT execute them, scratch
|
||
// unchanged from step 1
|
||
// step 3+: graph.launch_replay → executes the captured graph
|
||
//
|
||
// A post-step snapshot read on step 2 returns step 1's residual,
|
||
// which is bit-identical between any two captures regardless of input
|
||
// data. Two replay steps would work but require either three+ trainer
|
||
// calls or capture-internal scratch readback, both of which require
|
||
// substantial test infrastructure.
|
||
//
|
||
// Instead we rely on the `stacked_trainer_loss_shrinks_at_batch_32`
|
||
// smoke as the implicit scratch-zero validation: if the K-loop scratch
|
||
// wasn't being zeroed at step start, gradients from step N would
|
||
// pollute step N+1, training would diverge after a handful of steps,
|
||
// and the convergence assertion would fail. The B=32 smoke explicitly
|
||
// runs 200 training steps + measures final loss, so the absence of
|
||
// divergence IS the scratch-zero guarantee.
|
||
|
||
|
||
/// Eval alone must work — proves the eval path is fine WITHOUT any
|
||
/// prior captured-graph training. If this passes but
|
||
/// `evaluate_works_after_captured_training_step` fails, the bug is in
|
||
/// the train→eval transition (captured graph leaving state hostile to
|
||
/// direct kernel launches).
|
||
#[test]
|
||
fn evaluate_alone_succeeds() {
|
||
let dev = test_device();
|
||
let cfg = PerceptionTrainerConfig {
|
||
seq_len: 16,
|
||
mamba2_state_dim: 8,
|
||
lr_cfc: 3e-3,
|
||
lr_mamba2: 1e-3,
|
||
seed: 0x6262,
|
||
horizon_weights: [1.0; N_HORIZONS],
|
||
n_batch: 1,
|
||
smoothness_base_lambda: 0.0,
|
||
kernel_step_trace_path: None,
|
||
bucket_warmup_cap_override: None,
|
||
lr_aux: 3e-3,
|
||
};
|
||
let mut trainer = PerceptionTrainer::new(&dev, &cfg).expect("init");
|
||
let ts = 1_000_000u64;
|
||
let prev_mid = 5500.0_f32;
|
||
let (seq, labels) = synthetic_seq(cfg.seq_len, prev_mid, ts);
|
||
let (loss, probs) = trainer.evaluate(&seq, labels.as_slice()).expect("eval alone");
|
||
assert!(loss.is_finite(), "eval loss must be finite, got {loss}");
|
||
assert_eq!(probs.len(), cfg.seq_len * N_HORIZONS);
|
||
}
|
||
|
||
/// Regression test for z2w9w cluster run: training step (which captures
|
||
/// a CUDA Graph) must NOT break the subsequent `evaluate()` call.
|
||
/// z2w9w hit CUDA_ERROR_INVALID_VALUE at "eval snap_batched fwd" the
|
||
/// very first time eval ran after training; the captured graph or some
|
||
/// of its launch state left the stream in a state hostile to the
|
||
/// non-captured eval path.
|
||
#[test]
|
||
fn evaluate_works_after_captured_training_step() {
|
||
let dev = test_device();
|
||
let cfg = PerceptionTrainerConfig {
|
||
seq_len: 16,
|
||
mamba2_state_dim: 8,
|
||
lr_cfc: 3e-3,
|
||
lr_mamba2: 1e-3,
|
||
seed: 0x5151,
|
||
horizon_weights: [1.0; N_HORIZONS],
|
||
n_batch: 1,
|
||
smoothness_base_lambda: 0.0,
|
||
kernel_step_trace_path: None,
|
||
bucket_warmup_cap_override: None,
|
||
lr_aux: 3e-3,
|
||
};
|
||
let mut trainer = PerceptionTrainer::new(&dev, &cfg).expect("init");
|
||
let aux = synthetic_aux_outcomes(cfg.seq_len);
|
||
|
||
// Drive enough training steps to exercise warmup → capture → replay.
|
||
let mut ts = 1_000_000u64;
|
||
let mut prev_mid = 5500.0_f32;
|
||
for _ in 0..5 {
|
||
let (seq, labels) = synthetic_seq(cfg.seq_len, prev_mid, ts);
|
||
trainer.step(
|
||
&seq,
|
||
labels.as_slice(),
|
||
aux.prof_long.as_slice(),
|
||
aux.prof_short.as_slice(),
|
||
aux.size_long.as_slice(),
|
||
aux.size_short.as_slice(),
|
||
&aux.pos_fraction,
|
||
).expect("train step");
|
||
prev_mid = 0.5 * (seq.last().unwrap().bid_px[0] + seq.last().unwrap().ask_px[0]);
|
||
ts = seq.last().unwrap().ts_ns;
|
||
}
|
||
|
||
// Now call evaluate — must NOT fail with CUDA_ERROR_INVALID_VALUE.
|
||
let (seq, labels) = synthetic_seq(cfg.seq_len, prev_mid, ts);
|
||
let (loss, probs) = trainer.evaluate(&seq, labels.as_slice()).expect("evaluate after train");
|
||
assert!(loss.is_finite(), "eval loss must be finite, got {loss}");
|
||
assert_eq!(probs.len(), cfg.seq_len * N_HORIZONS, "eval probs must be [K, N_HORIZONS] flat");
|
||
assert!(probs.iter().all(|p| p.is_finite()), "eval probs must be finite");
|
||
}
|
||
|
||
/// 2 steps = warmup + capture (no replay yet). Does eval fail right
|
||
/// after capture but BEFORE the first graph.launch?
|
||
#[test]
|
||
fn evaluate_works_after_capture_no_replay() {
|
||
let dev = test_device();
|
||
let cfg = PerceptionTrainerConfig {
|
||
seq_len: 16,
|
||
mamba2_state_dim: 8,
|
||
lr_cfc: 3e-3,
|
||
lr_mamba2: 1e-3,
|
||
seed: 0x8181,
|
||
horizon_weights: [1.0; N_HORIZONS],
|
||
n_batch: 1,
|
||
smoothness_base_lambda: 0.0,
|
||
kernel_step_trace_path: None,
|
||
bucket_warmup_cap_override: None,
|
||
lr_aux: 3e-3,
|
||
};
|
||
let mut trainer = PerceptionTrainer::new(&dev, &cfg).expect("init");
|
||
let aux = synthetic_aux_outcomes(cfg.seq_len);
|
||
let ts = 1_000_000u64;
|
||
let prev_mid = 5500.0_f32;
|
||
let (seq, labels) = synthetic_seq(cfg.seq_len, prev_mid, ts);
|
||
trainer.step(
|
||
&seq,
|
||
labels.as_slice(),
|
||
aux.prof_long.as_slice(),
|
||
aux.prof_short.as_slice(),
|
||
aux.size_long.as_slice(),
|
||
aux.size_short.as_slice(),
|
||
&aux.pos_fraction,
|
||
).expect("warmup step");
|
||
trainer.step(
|
||
&seq,
|
||
labels.as_slice(),
|
||
aux.prof_long.as_slice(),
|
||
aux.prof_short.as_slice(),
|
||
aux.size_long.as_slice(),
|
||
aux.size_short.as_slice(),
|
||
&aux.pos_fraction,
|
||
).expect("capture step");
|
||
// Eval right after capture, no replay.
|
||
let (loss, _probs) = trainer.evaluate(&seq, labels.as_slice()).expect("eval after capture");
|
||
assert!(loss.is_finite(), "eval loss must be finite, got {loss}");
|
||
}
|
||
|
||
/// ISV-driven per-horizon EMA + lambda — after a few training steps,
|
||
/// `loss_ema` should contain positive BCE values for every horizon
|
||
/// (proves the BCE kernel writes the per-horizon channel correctly),
|
||
/// and `lambda` should land inside the ASYMMETRIC clamp [1.0, 2.0]
|
||
/// (boost-only — never demote below uniform, per
|
||
/// `pearl_audit_unboundedness_for_implicit_asymmetry.md`).
|
||
#[test]
|
||
fn horizon_ema_and_lambda_track_after_training() {
|
||
let dev = test_device();
|
||
let cfg = PerceptionTrainerConfig {
|
||
seq_len: 16,
|
||
mamba2_state_dim: 8,
|
||
lr_cfc: 3e-3,
|
||
lr_mamba2: 1e-3,
|
||
seed: 0x9292,
|
||
horizon_weights: [1.0; N_HORIZONS],
|
||
n_batch: 1,
|
||
smoothness_base_lambda: 0.0,
|
||
kernel_step_trace_path: None,
|
||
bucket_warmup_cap_override: None,
|
||
lr_aux: 3e-3,
|
||
};
|
||
let mut trainer = PerceptionTrainer::new(&dev, &cfg).expect("init");
|
||
let aux = synthetic_aux_outcomes(cfg.seq_len);
|
||
|
||
// EMA / lambda before any training step.
|
||
let ema0 = trainer.loss_ema_snapshot().expect("ema initial");
|
||
let lam0 = trainer.lambda_snapshot().expect("lambda initial");
|
||
assert!(ema0.iter().all(|v| *v == 0.0), "ema should be zero-init sentinel, got {ema0:?}");
|
||
assert!(lam0.iter().all(|v| *v == 0.0), "lambda should be zero before first kernel run, got {lam0:?}");
|
||
|
||
let mut ts = 1_000_000u64;
|
||
let mut prev_mid = 5500.0_f32;
|
||
for _ in 0..5 {
|
||
let (seq, labels) = synthetic_seq(cfg.seq_len, prev_mid, ts);
|
||
trainer.step(
|
||
&seq,
|
||
labels.as_slice(),
|
||
aux.prof_long.as_slice(),
|
||
aux.prof_short.as_slice(),
|
||
aux.size_long.as_slice(),
|
||
aux.size_short.as_slice(),
|
||
&aux.pos_fraction,
|
||
).expect("train step");
|
||
prev_mid = 0.5 * (seq.last().unwrap().bid_px[0] + seq.last().unwrap().ask_px[0]);
|
||
ts = seq.last().unwrap().ts_ns;
|
||
}
|
||
|
||
let ema = trainer.loss_ema_snapshot().expect("ema after train");
|
||
let lam = trainer.lambda_snapshot().expect("lambda after train");
|
||
eprintln!("loss_ema = {ema:?}");
|
||
eprintln!("lambda = {lam:?}");
|
||
// Each horizon must have a positive finite EMA (BCE > 0 unless model is perfect).
|
||
for (h, &v) in ema.iter().enumerate() {
|
||
assert!(v.is_finite(), "ema[{h}] not finite: {v}");
|
||
assert!(v > 0.0, "ema[{h}] should be positive, got {v}");
|
||
}
|
||
// Lambda asymmetric clamp: every entry in [1.0, 2.0] (boost-only).
|
||
for (h, &v) in lam.iter().enumerate() {
|
||
assert!(v.is_finite(), "lambda[{h}] not finite: {v}");
|
||
assert!(v >= 1.0 - 1e-6 && v <= 2.0 + 1e-6,
|
||
"lambda[{h}] = {v} outside asymmetric clamp [1.0, 2.0]");
|
||
}
|
||
// Mean of lambda must be >= 1.0 (boost-only, never demote).
|
||
// Upper-bounded by 2.0 (ceiling clamp). On synthetic-constant data
|
||
// the per-horizon EMAs converge, ratios approach 1.0, mean → 1.0.
|
||
let lam_mean: f32 = lam.iter().sum::<f32>() / lam.len() as f32;
|
||
assert!(lam_mean >= 1.0 - 1e-6 && lam_mean <= 2.0 + 1e-6,
|
||
"lambda mean {lam_mean} outside [1.0, 2.0] envelope");
|
||
}
|
||
|
||
/// SDD-3 Layer B5: verifies the aux supervision Huber loss converges on
|
||
/// the same synthetic constant-direction up-ramp the BCE smoke uses.
|
||
/// Asserts:
|
||
/// * Aux per-horizon Huber EMA is finite and bounded after training.
|
||
/// * Aux per-horizon directional accuracy EMA is high (the synthetic
|
||
/// long > short separation is unambiguous).
|
||
/// * Asymmetric stop-grad to encoder either stays masked (Phase 1
|
||
/// auditor behaviour) OR lifts after the gates are satisfied —
|
||
/// either is a valid outcome on the small synthetic budget.
|
||
#[test]
|
||
fn stacked_trainer_aux_supervision_converges_on_constant_signal() {
|
||
let dev = test_device();
|
||
let cfg = PerceptionTrainerConfig {
|
||
seq_len: 16,
|
||
mamba2_state_dim: 8,
|
||
lr_cfc: 3e-3,
|
||
lr_mamba2: 1e-3,
|
||
seed: 0xA1B5,
|
||
horizon_weights: [1.0; N_HORIZONS],
|
||
n_batch: 1,
|
||
smoothness_base_lambda: 0.0,
|
||
kernel_step_trace_path: None,
|
||
bucket_warmup_cap_override: None,
|
||
lr_aux: 3e-3,
|
||
};
|
||
let mut trainer = PerceptionTrainer::new(&dev, &cfg).expect("init");
|
||
let aux = synthetic_aux_outcomes(cfg.seq_len);
|
||
|
||
let mut ts = 1_000_000u64;
|
||
let mut prev_mid = 5500.0_f32;
|
||
// Train enough steps for the aux Huber + dir_acc EMAs to settle.
|
||
for _ in 0..400 {
|
||
let (seq, labels) = synthetic_seq(cfg.seq_len, prev_mid, ts);
|
||
trainer
|
||
.step(
|
||
&seq,
|
||
labels.as_slice(),
|
||
aux.prof_long.as_slice(),
|
||
aux.prof_short.as_slice(),
|
||
aux.size_long.as_slice(),
|
||
aux.size_short.as_slice(),
|
||
&aux.pos_fraction,
|
||
)
|
||
.expect("train step");
|
||
prev_mid = 0.5 * (seq.last().unwrap().bid_px[0] + seq.last().unwrap().ask_px[0]);
|
||
ts = seq.last().unwrap().ts_ns;
|
||
}
|
||
|
||
// After ≥400 steps (CB5):
|
||
// * Aux prof BCE EMA must be finite + non-negative on every horizon
|
||
// and STRICTLY BELOW the random-baseline ln(2) ≈ 0.693 — the
|
||
// constant-up signal is perfectly classifiable so even a small
|
||
// budget should beat random. (The CB5 kernel returns a joint
|
||
// scalar that we broadcast to every horizon, so the same value
|
||
// appears at h=0..N — but that's still load-bearing as an
|
||
// invariant check.)
|
||
// * Aux size Huber EMA must be finite + non-negative (no upper
|
||
// bound assertion because the σ-normalised size target can be
|
||
// anywhere in ℝ and we're not testing the regression accuracy,
|
||
// only that the loss is well-defined and the kernel ran).
|
||
// * Per-horizon dir_acc EMA must exceed chance (0.5) on every
|
||
// horizon — the synthetic long/short separation is unambiguous,
|
||
// so a well-wired aux head should clear it easily.
|
||
const RANDOM_BCE: f32 = std::f32::consts::LN_2;
|
||
for (h, &v) in trainer.aux_prof_bce_ema_per_h.iter().enumerate() {
|
||
assert!(v.is_finite(), "aux_prof_bce_ema_per_h[{h}] not finite: {v}");
|
||
assert!(v >= 0.0, "aux_prof_bce_ema_per_h[{h}] negative: {v}");
|
||
assert!(
|
||
v < RANDOM_BCE,
|
||
"aux_prof_bce_ema_per_h[{h}] = {v} >= random baseline ln(2)={RANDOM_BCE}; \
|
||
aux prof head should beat random on unambiguous signal"
|
||
);
|
||
}
|
||
for (h, &v) in trainer.aux_size_huber_ema_per_h.iter().enumerate() {
|
||
assert!(v.is_finite(), "aux_size_huber_ema_per_h[{h}] not finite: {v}");
|
||
assert!(v >= 0.0, "aux_size_huber_ema_per_h[{h}] negative: {v}");
|
||
}
|
||
for (h, &v) in trainer.aux_dir_acc_ema_per_h.iter().enumerate() {
|
||
assert!(v.is_finite(), "aux_dir_acc_ema_per_h[{h}] not finite: {v}");
|
||
assert!(
|
||
(0.0..=1.0).contains(&v),
|
||
"aux_dir_acc_ema_per_h[{h}] out of [0, 1]: {v}"
|
||
);
|
||
assert!(
|
||
v >= 0.5,
|
||
"aux_dir_acc_ema_per_h[{h}] = {v} below chance on unambiguous signal"
|
||
);
|
||
}
|
||
eprintln!(
|
||
"aux_prof_bce_ema_per_h = {:?}",
|
||
trainer.aux_prof_bce_ema_per_h
|
||
);
|
||
eprintln!(
|
||
"aux_size_huber_ema_per_h = {:?}",
|
||
trainer.aux_size_huber_ema_per_h
|
||
);
|
||
eprintln!(
|
||
"aux_dir_acc_ema_per_h = {:?}",
|
||
trainer.aux_dir_acc_ema_per_h
|
||
);
|
||
eprintln!(
|
||
"stop_grad_aux_to_encoder = {}",
|
||
trainer.stop_grad_aux_to_encoder
|
||
);
|
||
}
|
||
|
||
/// Does the bug surface after just ONE step (warmup only, no capture)?
|
||
/// If yes, the warmup dispatch path itself breaks subsequent eval.
|
||
/// If no, the captured graph instantiation / launch is what breaks eval.
|
||
#[test]
|
||
fn evaluate_works_after_warmup_only() {
|
||
let dev = test_device();
|
||
let cfg = PerceptionTrainerConfig {
|
||
seq_len: 16,
|
||
mamba2_state_dim: 8,
|
||
lr_cfc: 3e-3,
|
||
lr_mamba2: 1e-3,
|
||
seed: 0x7171,
|
||
horizon_weights: [1.0; N_HORIZONS],
|
||
n_batch: 1,
|
||
smoothness_base_lambda: 0.0,
|
||
kernel_step_trace_path: None,
|
||
bucket_warmup_cap_override: None,
|
||
lr_aux: 3e-3,
|
||
};
|
||
let mut trainer = PerceptionTrainer::new(&dev, &cfg).expect("init");
|
||
let aux = synthetic_aux_outcomes(cfg.seq_len);
|
||
let ts = 1_000_000u64;
|
||
let prev_mid = 5500.0_f32;
|
||
let (seq, labels) = synthetic_seq(cfg.seq_len, prev_mid, ts);
|
||
// ONLY one step — warmup, no capture yet.
|
||
trainer.step(
|
||
&seq,
|
||
labels.as_slice(),
|
||
aux.prof_long.as_slice(),
|
||
aux.prof_short.as_slice(),
|
||
aux.size_long.as_slice(),
|
||
aux.size_short.as_slice(),
|
||
&aux.pos_fraction,
|
||
).expect("warmup step");
|
||
// Now eval.
|
||
let (loss, _probs) = trainer.evaluate(&seq, labels.as_slice()).expect("eval after warmup");
|
||
assert!(loss.is_finite(), "eval loss must be finite, got {loss}");
|
||
}
|