15 KiB
Zero-Copy Walk-Forward Training — Implementation Plan
For agentic workers: REQUIRED SUB-SKILL: Use superpowers:subagent-driven-development (recommended) or superpowers:executing-plans to implement this plan task-by-task. Steps use checkbox (
- [ ]) syntax for tracking.
Goal: Eliminate per-fold CPU waste: no bar cloning, no feature re-extraction, no trainer re-creation. Data loaded to GPU once, sliced by index per fold.
Architecture: Load fxcache once → upload features+targets to GPU once → generate walk-forward index ranges (not data copies) → reuse DQNTrainer across folds, changing only the data range. The experience collector and training loop operate on index-bounded slices of the GPU-resident arrays.
Tech Stack: Rust 1.85, cudarc 0.17.3, CUDA
Current Flow (per fold)
1. generate_walk_forward_windows → clones bars into Vec<OHLCVBar> per fold
2. prepare_fold_data → extract_ml_features on CPU (REDUNDANT — fxcache already has them)
3. NormStats::from_features → CPU normalization stats
4. normalize_batch → CPU normalization
5. train_dqn_fold → creates NEW DQNTrainer (GPU init, kernel compile, graph capture)
6. train_with_preloaded_data → uploads Vec<(FeatureVector, Vec<f64>)> to GPU
7. init_gpu_raw_buffers → uploads features_raw + targets_raw to GPU
8. Per-epoch: collect_gpu_experiences → runs on GPU data
9. Per-epoch: run_training_steps → GPU training
Waste:
- Step 1: clones 4M OHLCVBars per fold (~320MB copied)
- Step 2: re-extracts 42-dim features on CPU (~30s on 4M bars)
- Step 3-4: CPU normalization on 4M × 42 floats
- Step 5: rebuilds ENTIRE GPU pipeline per fold (~2-3s on H100)
- Step 6: CPU→GPU upload of features+targets per fold
Target Flow
1. Load fxcache ONCE → all_features[N, 42], all_targets[N, 4], all_timestamps[N]
2. Upload features+targets to GPU ONCE → features_raw_cuda, targets_raw_cuda
3. Generate walk-forward INDEX RANGES → [(train_start, train_end, val_start, val_end), ...]
4. Create DQNTrainer ONCE
5. Per fold:
a. Set index range on trainer → trainer.set_fold_range(train_start, train_end, val_start, val_end)
b. Compute NormStats from GPU → GPU reduction kernel (mean, std per feature)
c. Normalize features on GPU → GPU element-wise kernel
d. collect_gpu_experiences → experience collector uses fold range
e. run_training_steps → training loop unchanged
f. Reset trainer state → clear replay buffer, reset optimizer, keep CUDA graphs
Eliminated:
- Zero bar cloning (index ranges only)
- Zero CPU feature extraction (fxcache → GPU directly)
- Zero trainer re-creation (reuse across folds)
- Zero CPU→GPU upload per fold (data stays on GPU)
- CPU normalization → GPU normalization kernel
File Structure
Modified Files
| File | Changes |
|---|---|
crates/ml/examples/train_baseline_rl.rs |
Remove prepare_fold_data, extract_ml_features calls. Load fxcache once, upload to GPU once, generate index ranges, reuse trainer across folds |
crates/ml/src/walk_forward.rs |
Add generate_walk_forward_indices() that returns Vec<FoldRange> (start/end indices) instead of cloning bars |
crates/ml/src/trainers/dqn/trainer/mod.rs |
Add set_fold_range() to reuse trainer across folds, reset_for_fold() to clear per-fold state |
crates/ml/src/trainers/dqn/trainer/training_loop.rs |
Support index-bounded training (experience collector uses fold range) |
New Types
/// Walk-forward fold defined by index ranges into pre-loaded arrays.
/// Zero data copying — just indices.
pub struct FoldRange {
pub fold: usize,
pub train_start: usize,
pub train_end: usize, // exclusive
pub val_start: usize,
pub val_end: usize, // exclusive
pub difficulty_score: f64,
}
Task 1: Add FoldRange and generate_walk_forward_indices
Files:
-
Modify:
crates/ml/src/walk_forward.rs -
Step 1: Add FoldRange struct
/// Walk-forward fold as index ranges into pre-loaded arrays.
/// Zero data copying — just start/end indices.
#[derive(Debug, Clone)]
pub struct FoldRange {
pub fold: usize,
pub train_start: usize,
pub train_end: usize,
pub val_start: usize,
pub val_end: usize,
pub difficulty_score: f64,
}
- Step 2: Add
generate_walk_forward_indicesfunction
Same date-boundary logic as generate_walk_forward_windows but returns indices instead of cloned Vecs. Uses partition_point for O(log n) index lookup instead of filtering all bars.
pub fn generate_walk_forward_indices(
bars: &[OHLCVBar],
config: &WalkForwardConfig,
) -> Vec<FoldRange> {
// Same date arithmetic as generate_walk_forward_windows
// But instead of:
// bars.iter().filter(...).copied().collect()
// Use:
// bars.partition_point(|b| b.timestamp.date_naive() < train_end)
// to get index boundaries in O(log n)
}
- Step 3: Compile and test
SQLX_OFFLINE=true cargo check -p ml
- Step 4: Commit
git commit -m "feat: FoldRange + generate_walk_forward_indices — zero-copy walk-forward"
Task 2: Add reset_for_fold to DQNTrainer
Files:
- Modify:
crates/ml/src/trainers/dqn/trainer/mod.rs
The trainer currently holds per-fold state (replay buffer contents, optimizer momentum, epoch counters). To reuse across folds, we need a reset that clears per-fold state but keeps GPU infrastructure (CUDA context, compiled kernels, CUDA graphs, cuBLAS handles).
- Step 1: Add
reset_for_foldmethod
/// Reset per-fold training state while keeping GPU infrastructure.
/// Clears: replay buffer, optimizer momentum, epoch counters, loss history.
/// Keeps: CUDA context, compiled kernels, CUDA graphs, cuBLAS handles,
/// GPU data buffers (features_raw_cuda, targets_raw_cuda).
pub async fn reset_for_fold(&mut self) -> Result<()> {
// Clear replay buffer
let agent = self.agent.write().await;
agent.memory().clear()?;
drop(agent);
// Reset epoch counters
self.current_epoch = 0;
self.gradient_logging_step = 0;
self.steps_since_varmap_sync = 0;
// Reset loss/Q-value history
self.loss_history.clear();
self.q_value_history.clear();
self.val_loss_history.clear();
self.sharpe_history.clear();
self.safety_loss_history.clear();
// Reset fused training context state (keep graphs, reset step counters)
if let Some(ref mut fused) = self.fused_ctx {
fused.reset_for_fold()?;
}
// Reset training guard accumulators
if let Some(ref mut guard) = self.training_guard {
guard.reset_accumulators()?;
}
Ok(())
}
- Step 2: Add
FusedTrainingCtx::reset_for_fold
In crates/ml/src/trainers/dqn/fused_training.rs:
/// Reset per-fold state. Keeps GPU infrastructure + CUDA graphs.
pub(crate) fn reset_for_fold(&mut self) -> Result<()> {
self.steps_since_varmap_sync = 0;
self.last_combined_norm = 0.0;
// Invalidate graph_aux — fold data changed, graph topology might differ
self.graph_aux = None;
Ok(())
}
- Step 3: Compile and test
SQLX_OFFLINE=true cargo check -p ml
- Step 4: Commit
git commit -m "feat: DQNTrainer::reset_for_fold — reuse GPU pipeline across folds"
Task 3: Add set_training_range for index-bounded training
Files:
- Modify:
crates/ml/src/trainers/dqn/trainer/training_loop.rs
The experience collector and training loop need to know which slice of the GPU data to use for this fold. Instead of re-uploading, set index bounds that the collector uses.
- Step 1: Add fold range fields to DQNTrainer
/// Current fold's training data range in the GPU-resident arrays.
/// Set by set_training_range(), used by collect_gpu_experiences().
pub(crate) fold_train_start: usize,
pub(crate) fold_train_end: usize,
pub(crate) fold_val_start: usize,
pub(crate) fold_val_end: usize,
- Step 2: Add
set_training_rangemethod
pub fn set_training_range(&mut self, train_start: usize, train_end: usize, val_start: usize, val_end: usize) {
self.fold_train_start = train_start;
self.fold_train_end = train_end;
self.fold_val_start = val_start;
self.fold_val_end = val_end;
// Clear gpu_data to force re-init with new range
// (init_gpu_data checks is_some() — clearing forces re-run)
self.gpu_data = None;
}
- Step 3: Update experience collector to use fold range
In collect_gpu_experiences, replace training_data.len() with self.fold_train_end - self.fold_train_start. Pass the offset to the experience collector so it reads from the correct slice of the GPU arrays.
- Step 4: Compile and test
SQLX_OFFLINE=true cargo check -p ml
- Step 5: Commit
git commit -m "feat: index-bounded training — experience collector uses fold range"
Task 4a: Move DQN trainer creation before the fold loop
Files:
- Modify:
crates/ml/examples/train_baseline_rl.rs
Currently train_dqn_fold() creates a new DQNTrainer::new(hyperparams) per fold (line 576). Move this to BEFORE the fold loop so the trainer (and all its GPU infrastructure) is created once.
- Step 1: Create DQN trainer once before fold loop
After the walk-forward windows are generated (line 910), create the trainer:
// Create DQN trainer ONCE — GPU infra persists across folds
let hp = load_hyperopt_params(&args.hyperopt_params, "dqn");
let dqn_hyperparams = build_dqn_hyperparams(args, &hp, &gpu_profile);
let mut dqn_trainer = if train_dqn {
Some(DQNTrainer::new(dqn_hyperparams)?)
} else {
None
};
This requires extracting the hyperparams construction from train_dqn_fold into a separate function build_dqn_hyperparams().
- Step 2: Change
train_dqn_foldto accept&mut DQNTrainerinstead of creating one
Change signature from:
fn train_dqn_fold(fold, train_features, val_features, train_bars, val_bars, args, output_dir, hp, pre_uploaded_gpu_data, prefix) -> Result<f64>
To:
fn train_dqn_fold(trainer: &mut DQNTrainer, fold, train_features, val_features, train_bars, val_bars, args, output_dir, prefix) -> Result<f64>
Remove the DQNTrainer::new() call inside the function. The trainer is now passed in.
- Step 3: Compile
SQLX_OFFLINE=true cargo check -p ml --example train_baseline_rl
- Step 4: Commit
git commit -m "refactor: move DQN trainer creation before fold loop — create once, reuse across folds"
Task 4b: Replace prepare_fold_data with index-based feature slicing
Files:
- Modify:
crates/ml/examples/train_baseline_rl.rs
prepare_fold_data calls extract_ml_features(&window.train) which re-computes 42-dim features from raw OHLCV bars on CPU. This is redundant — all_features already has them from the fxcache load.
- Step 1: Replace fold data prep in the fold loop
Instead of calling prepare_fold_data(window, &args.output_dir), use the pre-loaded all_features array and slice by index using FoldRange:
let fold_ranges = generate_walk_forward_indices(aligned_bars, &wf_config);
for range in &fold_ranges {
// Slice pre-loaded features by index — zero CPU re-extraction
let train_feat = &all_features[range.train_start..range.train_end];
let val_feat = &all_features[range.val_start..range.val_end];
let train_bars = &aligned_bars[range.train_start..range.train_end];
let val_bars = &aligned_bars[range.val_start..range.val_end];
// Normalize (still CPU for now — GPU normalization is a future task)
let norm_stats = NormStats::from_features(train_feat);
let train_norm = norm_stats.normalize_batch(train_feat);
let val_norm = norm_stats.normalize_batch(val_feat);
// Save NormStats
let norm_path = args.output_dir.join(format!("norm_stats_fold{}.json", range.fold));
std::fs::write(&norm_path, serde_json::to_string_pretty(&norm_stats)?)?;
// Train
train_dqn_fold(&mut trainer, range.fold, &train_norm, &val_norm, train_bars, val_bars, ...)?;
}
-
Step 2: Delete
prepare_fold_datafunction entirely -
Step 3: Delete
FoldDatatype alias -
Step 4: Delete fold prefetch thread (lines 954-973)
The background prefetch thread ran prepare_fold_data on the next fold. With zero-copy slicing, fold prep is instant — no prefetch needed.
- Step 5: Delete double-buffer GPU upload (lines 1111-1131)
The DqnGpuData::upload + dqn_gpu_staged double-buffer pre-uploaded next fold's data to GPU between folds. With data staying on GPU, this is dead code.
- Step 6: Compile
SQLX_OFFLINE=true cargo check -p ml --example train_baseline_rl
- Step 7: Commit
git commit -m "perf: zero-copy fold slicing — pre-loaded features sliced by index, no CPU re-extraction"
Task 4c: Wire reset_for_fold into the fold loop
Files:
- Modify:
crates/ml/examples/train_baseline_rl.rs
Now that the trainer persists across folds, each fold must reset per-fold state (replay buffer, optimizer momentum, epoch counters) while keeping GPU infrastructure.
- Step 1: Add reset_for_fold call at start of each fold
for range in &fold_ranges {
// Reset per-fold state (keep GPU infra)
if let Some(ref mut trainer) = dqn_trainer {
trainer.reset_for_fold().await?;
}
// ... rest of fold training
}
- Step 2: Compile and run locally
SQLX_OFFLINE=true FOXHUNT_TEST_DATA=test_data/futures-baseline cargo run --release -p ml --example train_baseline_rl -- --model dqn --epochs 3 --data-dir test_data/futures-baseline
Expected: 3 folds × 3 epochs in <30s (was 159s).
- Step 3: Commit
git commit -m "perf: zero-copy walk-forward complete — load once, slice by index, reset between folds"
Task 5: Validate on H100
- Step 1: Run precompute (if needed)
argo submit -n foxhunt --from workflowtemplate/precompute-features
- Step 2: Submit H100 training
./scripts/argo-train.sh dqn --epochs 50 --trials 0 --gpu-pool ci-training-h100
Expected:
- Data load: <5s (fxcache hit)
- Per-fold init: <0.5s (reset only, no re-creation)
- Per-epoch: <5s (batch=8192, 488 steps)
- Total 50 epochs × 3 folds: <15 minutes
Expected Performance Impact
| Metric | Before | After |
|---|---|---|
| Data load per fold | ~30s CPU (feature extraction) | 0s (already on GPU) |
| GPU init per fold | ~3s (kernel compile + graph capture) | ~0.1s (reset counters only) |
| Data upload per fold | ~1s (CPU→GPU DMA) | 0s (data stays on GPU) |
| Memory per fold | ~640MB (cloned bars + features + targets) | ~0 (index ranges only) |
| Total fold transition | ~34s | ~0.1s |
| 3-fold overhead | ~102s | ~0.3s |
Risks
| Risk | Impact | Mitigation |
|---|---|---|
| Optimizer momentum carries across folds | Training quality regression | reset_for_fold clears Adam m/v buffers |
| CUDA graphs invalid after range change | Graph replay crash | Invalidate graph_aux on range change, re-capture on first step |
| NormStats differ per fold | Feature distribution shift | Compute NormStats per fold from GPU slice (GPU reduction kernel, future task) |
| Replay buffer from previous fold bleeds | Wrong training distribution | buffer.clear() in reset_for_fold |