# 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

```rust
/// 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**

```rust
/// 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.

```rust
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**

```bash
SQLX_OFFLINE=true cargo check -p ml
```

- [ ] **Step 4: Commit**

```bash
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**

```rust
/// 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`:

```rust
/// 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**

```bash
SQLX_OFFLINE=true cargo check -p ml
```

- [ ] **Step 4: Commit**

```bash
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**

```rust
/// 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**

```rust
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**

```bash
SQLX_OFFLINE=true cargo check -p ml
```

- [ ] **Step 5: Commit**

```bash
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:

```rust
// 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:
```rust
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:
```rust
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**

```bash
SQLX_OFFLINE=true cargo check -p ml --example train_baseline_rl
```

- [ ] **Step 4: Commit**

```bash
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`:

```rust
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**

```bash
SQLX_OFFLINE=true cargo check -p ml --example train_baseline_rl
```

- [ ] **Step 7: Commit**

```bash
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**

```rust
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**

```bash
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**

```bash
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)**

```bash
argo submit -n foxhunt --from workflowtemplate/precompute-features
```

- [ ] **Step 2: Submit H100 training**

```bash
./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 |
