gitadmin/foxhunt
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

fix(rl-cli): build B×K snapshot tensor per step at b_size>1

Browse Source 
Crash in alpha-rl-ljn8k (commit 9c6c280bd) at step 0:
  forward_only: expected 512 snapshots (B=16 * K=32); got 32

Root cause: CLI binary called next_sequence_pair() once per step
and passed the resulting K-snapshot window to step_with_lobsim. At
b_size=1 the encoder's B×K=32 contract matched; at b_size=16 it
silently expected B×K=512 and bailed.

Fix: sample n_backtests INDEPENDENT pairs per step and concat into
B×K row-major layout. This is the "proper" per-batch market
diversity that delivers the gradient-variance reduction promised
by `pearl_b_size_1_signal_starvation_blocks_q_learning` —
tiling one window B times would be bit-identical encoder input
across batch slots and contribute zero encoder gradient diversity.

Bumps loader cap from n_steps×2 to n_steps×n_backtests×2 so the
extra pair-sampling doesn't EOF mid-run. Eval loop fixed
identically.

Staging buffers preallocated outside the step loop to avoid
per-step Vec alloc thrash at the hot path.

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-05-24 11:38:37 +02:00
parent 9c6c280bd8
commit a776fab31f
1 changed files with 72 additions and 31 deletions
Download Patch File Download Diff File Expand all files Collapse all files

103
crates/ml-alpha/examples/alpha_rl_train.rs
View File

@@ -32,6 +32,7 @@
use anyhow::{Context, Result};
use clap::Parser;
use data::providers::databento::dbn_parser::InstrumentFilter;
use ml_alpha::cfc::snap_features::Mbp10RawInput;
use ml_alpha::data::loader::{
discover_mbp10_files_sorted, MultiHorizonLoader, MultiHorizonLoaderConfig,
DEFAULT_OUTCOME_LABEL_COST_ES,
@@ -353,10 +354,16 @@ fn main() -> Result<()> {
predecoded_dir: cli.predecoded_dir.clone(),
multi_resolution: multi_resolution.clone(),
horizons: HORIZONS,
// Soft cap above n_steps × 2 so next_sequence_pair never
// returns Ok(None) mid-run. Loader's per-file random anchor
// sampling reuses files indefinitely up to this cap.
n_max_sequences: cli.n_steps.saturating_mul(2).max(64),
// Soft cap above n_steps × n_backtests × 2 so next_sequence_pair
// never returns Ok(None) mid-run. Each step samples n_backtests
// pairs (one per batch slot) to build the B×K snapshot tensor
// that forward_encoder consumes — per-batch market diversity is
// what gives the gradient-variance reduction promised by
// `pearl_b_size_1_signal_starvation_blocks_q_learning`.
n_max_sequences: cli.n_steps
.saturating_mul(cli.n_backtests.max(1))
.saturating_mul(2)
.max(64),
seed: cli.seed,
inference_only: false,
outcome_label_cost: cli.outcome_label_cost,
@@ -423,23 +430,47 @@ fn main() -> Result<()> {
let mut windowed_act_hist: [f32; 9] = [0.0; 9];
const WINDOWED_ACT_ALPHA: f32 = 1.0 / 1000.0;
// Per-step staging buffers for the B×K snapshot tensor that
// forward_encoder consumes. Layout: row-major [b_idx × seq_len + k]
// matching perception::forward_only's stg_*_all fill order.
// Pre-allocate once and clear+extend per step.
let n_batch = cli.n_backtests.max(1);
let bk = n_batch.saturating_mul(cli.seq_len);
let mut s_t_bk: Vec<Mbp10RawInput> = Vec::with_capacity(bk);
let mut s_tp1_bk: Vec<Mbp10RawInput> = Vec::with_capacity(bk);
let t_start = std::time::Instant::now();
for step in 0..cli.n_steps {
let pair = loader
.next_sequence_pair()
.with_context(|| format!("next_sequence_pair at step {step}"))?;
let (s_t, s_tp1) = match pair {
Some(p) => p,
None => {
eprintln!(
"loader EOF at step {step} (n_max_sequences exhausted) — exiting early"
);
break;
}
};
'outer: for step in 0..cli.n_steps {
// Sample n_backtests INDEPENDENT pairs per step. Each pair
// becomes one batch slot's K-window — distinct market segments
// per batch slot give true gradient-variance reduction (vs
// tiling one window B times, which is bit-identical encoder
// input and contributes no encoder gradient diversity).
s_t_bk.clear();
s_tp1_bk.clear();
for b_idx in 0..n_batch {
let pair = loader
.next_sequence_pair()
.with_context(|| format!("next_sequence_pair at step {step} batch {b_idx}"))?;
let (s_t, s_tp1) = match pair {
Some(p) => p,
None => {
eprintln!(
"loader EOF at step {step} batch {b_idx} (n_max_sequences exhausted) — exiting early"
);
break 'outer;
}
};
debug_assert_eq!(s_t.snapshots.len(), cli.seq_len);
debug_assert_eq!(s_tp1.snapshots.len(), cli.seq_len);
s_t_bk.extend_from_slice(&s_t.snapshots);
s_tp1_bk.extend_from_slice(&s_tp1.snapshots);
}
debug_assert_eq!(s_t_bk.len(), bk);
debug_assert_eq!(s_tp1_bk.len(), bk);
let stats = trainer
.step_with_lobsim(&s_t.snapshots, &s_tp1.snapshots, &mut sim)
.step_with_lobsim(&s_t_bk, &s_tp1_bk, &mut sim)
.with_context(|| format!("step_with_lobsim at step {step}"))?;
// Gate G8: NaN abort. Per `feedback_stop_on_anomaly` +
@@ -824,7 +855,10 @@ fn main() -> Result<()> {
predecoded_dir: cli.predecoded_dir.clone(),
multi_resolution,
horizons: HORIZONS,
n_max_sequences: cli.n_eval_steps.saturating_mul(2).max(64),
n_max_sequences: cli.n_eval_steps
.saturating_mul(cli.n_backtests.max(1))
.saturating_mul(2)
.max(64),
// Different seed for eval so sequence sampling is independent
// of the train seed (deterministic given the eval seed).
seed: cli.seed.wrapping_add(0xE7AE),
@@ -835,19 +869,26 @@ fn main() -> Result<()> {
let mut eval_loader = MultiHorizonLoader::new(&eval_loader_cfg)
.context("MultiHorizonLoader::new (eval)")?;
for eval_step in 0..cli.n_eval_steps {
let pair = eval_loader
.next_sequence_pair()
.with_context(|| format!("eval next_sequence_pair at step {eval_step}"))?;
let (s_t, s_tp1) = match pair {
Some(p) => p,
None => {
eprintln!("eval loader EOF at step {eval_step}");
break;
}
};
// Reuse staging buffers from train phase (B×K snapshot layout).
'eval_outer: for eval_step in 0..cli.n_eval_steps {
s_t_bk.clear();
s_tp1_bk.clear();
for b_idx in 0..n_batch {
let pair = eval_loader
.next_sequence_pair()
.with_context(|| format!("eval next_sequence_pair at step {eval_step} batch {b_idx}"))?;
let (s_t, s_tp1) = match pair {
Some(p) => p,
None => {
eprintln!("eval loader EOF at step {eval_step} batch {b_idx}");
break 'eval_outer;
}
};
s_t_bk.extend_from_slice(&s_t.snapshots);
s_tp1_bk.extend_from_slice(&s_tp1.snapshots);
}
let stats = trainer
.step_with_lobsim(&s_t.snapshots, &s_tp1.snapshots, &mut sim)
.step_with_lobsim(&s_t_bk, &s_tp1_bk, &mut sim)
.with_context(|| format!("step_with_lobsim eval step {eval_step}"))?;
if !stats.l_total.is_finite() {
eprintln!("G8 NaN ABORT during eval at step {eval_step}");