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

perf(ml-alpha): preload all MBP-10 files into RAM, reset loaders per epoch

Browse Source 
cnjfl wall-time analysis: 80% of training time was disk IO. The
MultiHorizonLoader was constructed fresh per epoch in the CLI loop,
forcing 9 file deserializations from bincode (~50s/file × 9 = 7-8
min) per epoch. For a 6-epoch run that was ~45 min of pure IO out
of ~50 min total.

Now loaders preload ALL files into RAM at construction (one-time
~7 min startup) and expose a `reset(seed: u64)` method that
re-seeds anchor sampling per epoch — no disk IO between epochs.

Memory cost: ~13-15 GB for the 9-quarter ES.FUT dataset (45M
snapshots × ~280 bytes). Well under the training pod's 64 GB
limit; current 16 GB request remains sufficient since the resident
set fits.

Expected wall-time impact at K=96, B=8, 16K seqs:
  6 epochs:  ~50 min → ~13 min  (~3.7×)
  15 epochs: ~120 min → ~23 min (~5×)

The internal LoadedFile-cache-with-cycling logic is gone — the
loader now holds Vec<LoadedFile> with all files resident. Per-call
`next_sequence` picks a uniformly random file + uniformly random
anchor inside it (instead of cycling files with a per-file budget).
Distribution is equivalent: each file contributes ~n_max_sequences /
n_files samples per epoch in expectation.

77 ml-alpha tests pass.

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-05-17 11:57:07 +02:00
parent 26d91a816c
commit 3f089210f3
2 changed files with 99 additions and 88 deletions
Download Patch File Download Diff File Expand all files Collapse all files

50
crates/ml-alpha/examples/alpha_train.rs
View File

@@ -247,16 +247,38 @@ fn main() -> Result<()> {
// sequence yields one optimizer step). Used by cosine decay.
let total_steps_budget = cli.epochs * cli.n_train_seqs;
// PRELOAD: construct loaders ONCE (each preloads all files into RAM).
// Per-epoch we call reset(seed) to re-shuffle anchor sampling — no
// disk IO. Cuts wall time ~5× on a 6-epoch run, ~7× on 15 epochs.
tracing::info!("preloading train + val MBP-10 files into RAM (slow startup, fast per-epoch)…");
let preload_start = std::time::Instant::now();
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,
})
.context("train loader")?;
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),
})
.context("val loader")?;
tracing::info!(
elapsed_s = preload_start.elapsed().as_secs(),
"preload complete",
);
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")?;
// Re-seed loaders for this epoch (anchor sampling differs per epoch).
train_loader.reset(cli.seed.wrapping_add(epoch as u64));
val_loader.reset(cli.seed.wrapping_add(0xC0FFEE + epoch as u64));
let mut epoch_train_loss = 0.0_f32;
let mut epoch_train_steps = 0usize;
@@ -303,15 +325,7 @@ fn main() -> Result<()> {
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")?;
// val_loader was preloaded above and reset at top of loop.
let mut val_probs: [Vec<f32>; N_HORIZONS] = Default::default();
let mut val_labels: [Vec<f32>; N_HORIZONS] = Default::default();

137
crates/ml-alpha/src/data/loader.rs
View File

@@ -128,22 +128,23 @@ struct LoadedFile {
pub struct MultiHorizonLoader {
cfg: MultiHorizonLoaderConfig,
files: Vec<PathBuf>,
rng: ChaCha8Rng,
file_idx: usize,
yielded: usize,
/// Currently-loaded file. Cleared and replaced when we advance.
current: Option<LoadedFile>,
/// Sequences yielded from `current` since it was loaded.
yielded_from_current: usize,
/// How many sequences to yield from each file before advancing. Set
/// from `n_max_sequences / n_files` (rounded up) at construction, so
/// each file is loaded exactly once per "pass" through the loader.
sequences_per_file: usize,
/// All MBP-10 files preloaded into memory at construction. Each
/// `LoadedFile` carries its deserialized snapshots + precomputed
/// per-horizon labels + regime features. Per-epoch iteration is
/// pure memory access — no disk IO after the initial preload.
///
/// Memory cost at full ES.FUT dataset (9 quarters, ~45M snapshots
/// total): ~13-15 GB. Worth it: per-epoch wall time drops from
/// ~8 min (file-IO-bound) to ~1 min (compute-bound) on L40S, a
/// ~5-7× speedup that compounds for longer runs.
files_loaded: Vec<LoadedFile>,
}
impl MultiHorizonLoader {
/// Construct + preload all MBP-10 files. Slow (~50s × N_files),
/// done once. Subsequent `next_sequence` calls are pure memory.
pub fn new(cfg: &MultiHorizonLoaderConfig) -> Result<Self> {
let mut files: Vec<PathBuf> = std::fs::read_dir(&cfg.mbp10_root)
.with_context(|| format!("read mbp10 root {}", cfg.mbp10_root.display()))?
@@ -151,7 +152,6 @@ impl MultiHorizonLoader {
.map(|e| e.path())
.filter(|p| {
let s = p.to_string_lossy();
// Source DBN files only; not the predecoded sidecars themselves.
s.ends_with(".dbn.zst") || s.ends_with(".dbn")
})
.collect();
@@ -162,25 +162,63 @@ impl MultiHorizonLoader {
);
let mut rng = ChaCha8Rng::seed_from_u64(cfg.seed);
files.shuffle(&mut rng);
let n_files = files.len();
// Yield (n_max_sequences / n_files) sequences per file, rounded
// up so we always honor n_max_sequences. With 8000 max and 9
// files: 8000/9 = 889/file, total 9 loads instead of 8000.
let sequences_per_file = cfg.n_max_sequences.div_ceil(n_files).max(1);
let max_horizon = *cfg.horizons.iter().max().expect("non-empty horizons");
let min_size = cfg.seq_len + max_horizon + 1;
let mut files_loaded: Vec<LoadedFile> = Vec::with_capacity(files.len());
for path in &files {
let snapshots = load_or_predecode_mbp10(path, &cfg.predecoded_dir)
.with_context(|| format!("load mbp10 {}", path.display()))?;
if snapshots.len() < min_size {
tracing::warn!(
path = %path.display(),
snapshots = snapshots.len(),
min_size,
"skipping file (too few snapshots for seq_len + max_horizon)",
);
continue;
}
let prices: Vec<f32> = snapshots.iter().map(mid_price_f32).collect();
let mut labels_full: [Vec<f32>; 5] = Default::default();
for (h_idx, &h) in cfg.horizons.iter().enumerate() {
let mut full = vec![f32::NAN; snapshots.len()];
let raw = generate_labels(&prices, h);
for (i, &t) in raw.valid_indices.iter().enumerate() {
full[t] = raw.labels[i];
}
labels_full[h_idx] = full;
}
let regime_full = compute_regime_features(&snapshots);
files_loaded.push(LoadedFile { snapshots, labels_full, regime_full });
}
anyhow::ensure!(
!files_loaded.is_empty(),
"no files had enough snapshots ({}+ required) under {}",
min_size, cfg.mbp10_root.display()
);
tracing::info!(
n_files = files_loaded.len(),
total_snapshots = files_loaded.iter().map(|lf| lf.snapshots.len()).sum::<usize>(),
"MultiHorizonLoader preloaded all files",
);
Ok(Self {
cfg: cfg.clone(),
files,
rng,
file_idx: 0,
yielded: 0,
current: None,
yielded_from_current: 0,
sequences_per_file,
files_loaded,
})
}
/// Reset per-epoch state: re-seed RNG (so anchor sampling differs
/// each epoch) and zero the yielded counter. The preloaded file
/// inventory is preserved — no disk IO.
pub fn reset(&mut self, seed: u64) {
self.rng = ChaCha8Rng::seed_from_u64(seed);
self.yielded = 0;
}
pub fn n_files(&self) -> usize {
self.files.len()
self.files_loaded.len()
}
pub fn yielded(&self) -> usize {
@@ -193,52 +231,12 @@ impl MultiHorizonLoader {
}
let max_horizon = *self.cfg.horizons.iter().max().expect("non-empty horizons");
// Advance to a new file if we have none cached or have yielded our
// budget from the current one.
let need_new = self.current.is_none()
|| self.yielded_from_current >= self.sequences_per_file;
if need_new {
// Try files until one is big enough; advance file_idx each time.
let mut found: Option<LoadedFile> = None;
for _ in 0..self.files.len() * 2 {
if self.file_idx >= self.files.len() {
self.files.shuffle(&mut self.rng);
self.file_idx = 0;
}
let path = &self.files[self.file_idx];
let snapshots = load_or_predecode_mbp10(path, &self.cfg.predecoded_dir)
.with_context(|| format!("load mbp10 {}", path.display()))?;
self.file_idx += 1;
if snapshots.len() < self.cfg.seq_len + max_horizon + 1 {
continue;
}
// Precompute mid prices + per-horizon labels + regime
// features for the WHOLE file. We pay this once per file
// load, then slice cheaply per anchor on subsequent calls.
let prices: Vec<f32> = snapshots.iter().map(mid_price_f32).collect();
let mut labels_full: [Vec<f32>; 5] = Default::default();
for (h_idx, &h) in self.cfg.horizons.iter().enumerate() {
let mut full = vec![f32::NAN; snapshots.len()];
let raw = generate_labels(&prices, h);
for (i, &t) in raw.valid_indices.iter().enumerate() {
full[t] = raw.labels[i];
}
labels_full[h_idx] = full;
}
let regime_full = compute_regime_features(&snapshots);
found = Some(LoadedFile { snapshots, labels_full, regime_full });
break;
}
self.current = found;
self.yielded_from_current = 0;
if self.current.is_none() {
return Ok(None); // no file in the inventory has enough snapshots
}
}
// Sample an anchor inside the current file and build the sequence.
let lf = self.current.as_ref().expect("set above");
// Pick a random file from the preloaded inventory, then a random
// anchor within that file. Uniform-across-files sampling — each
// file contributes ~n_max_sequences / n_files samples per epoch.
let n_files = self.files_loaded.len();
let file_idx = self.rng.gen_range(0..n_files);
let lf = &self.files_loaded[file_idx];
let max_anchor = lf.snapshots.len() - self.cfg.seq_len - max_horizon;
let anchor: usize = self.rng.gen_range(0..max_anchor);
@@ -256,7 +254,6 @@ impl MultiHorizonLoader {
}
self.yielded += 1;
self.yielded_from_current += 1;
Ok(Some(LabeledSequence { snapshots: sequence, labels }))
}
}