Files
foxhunt/docs/superpowers/plans/2026-04-02-zero-copy-walk-forward.md
2026-04-02 23:29:04 +02:00

15 KiB
Raw Blame History

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_indices function

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_fold method
/// 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_range method
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_fold to accept &mut DQNTrainer instead 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_data function entirely

  • Step 3: Delete FoldData type 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