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

21 KiB
Raw Blame History

Zero-Copy FxCache Training Pipeline — 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 all per-fold CPU waste by building the entire training pipeline on FxCacheData (features[42] + targets[4] + OFI[8] + timestamps). Load once, upload to GPU once, slice by index per fold. Walk-forward uses full MBP-10/trades data via precomputed OFI features.

Architecture: FxCacheData replaces OHLCVBar as the data backbone. Walk-forward generates index ranges from timestamps. Trainer accepts (&[[f64; 42]], &[[f64; 4]], &[[f64; 8]]) slices instead of Vec<(FeatureVector, Vec<f64>)>. Data uploaded to GPU once; per-fold transitions set index bounds only.

Tech Stack: Rust 1.85, cudarc 0.17.3, CUDA


Current Data Flow (per fold)

fxcache ─→ (features[42], targets[4], OFI[8], timestamps)     # 470MB, loaded once
                    │
                    ▼
    OHLCVBars reconstructed from targets (close prices)        # WASTE: fxcache has everything
                    │
                    ▼
    generate_walk_forward_windows(bars)                         # Clones bars into Vec per fold
                    │
                    ▼
    prepare_fold_data(window)                                   # RE-EXTRACTS features from bars (!)
                    │
                    ▼
    features_to_trainer_format(features, bars)                  # 4M × Vec<f64> heap allocs
                    │
                    ▼
    DQNTrainer::new() per fold                                  # Rebuilds GPU pipeline
                    │
                    ▼
    init_gpu_data → flatten + upload to GPU                     # Unpacks tuples, flattens, HtoD
                    │
                    ▼
    Training epochs (GPU)                                       # The only useful work

Target Data Flow

fxcache ─→ FxCacheData { features[42], targets[4], OFI[8], timestamps }
                    │
                    ▼ (one-time)
    Upload to GPU: features_cuda, targets_cuda, ofi_cuda
                    │
                    ▼ (one-time)
    generate_walk_forward_indices(timestamps) → Vec<FoldRange>
                    │
                    ▼ (one-time)
    DQNTrainer::new()
                    │
                    ▼ (per fold, ~0.1s)
    trainer.set_training_range(range.train_start, range.train_end, ...)
    trainer.reset_for_fold()
    NormStats from &features[train_start..train_end]            # Slice, no copy
                    │
                    ▼
    Training epochs (GPU reads from offset in pre-uploaded arrays)

File Structure

Modified Files

File Changes
crates/ml/src/walk_forward.rs generate_walk_forward_indices_from_timestamps(&[i64], config) → Vec<FoldRange> — date slicing on nanosecond timestamps from fxcache, no OHLCVBar dependency
crates/ml/src/trainers/dqn/trainer/mod.rs New train_with_fxcache_data() API accepting &FxCacheData. Remove Vec<(FeatureVector, Vec<f64>)> from the hot path. Add upload_fxcache_to_gpu() for one-time upload
crates/ml/src/trainers/dqn/trainer/training_loop.rs init_gpu_raw_buffers_from_slices() accepting (&[[f64; 42]], &[[f64; 4]], &[[f64; 8]]). Remove per-element flattening loops
crates/ml/src/cuda_pipeline/mod.rs DqnGpuData::upload_slices() accepting contiguous &[[f64; 42]] + &[[f64; 4]] — no tuple unpacking
crates/ml/examples/train_baseline_rl.rs Complete rewrite of fold loop: load fxcache → upload GPU → index ranges → per-fold: set range, reset, train. Delete prepare_fold_data, FoldData, features_to_trainer_format, double-buffer, prefetch thread

Task 1: Walk-forward from timestamps (already partially done)

Files:

  • Modify: crates/ml/src/walk_forward.rs

generate_walk_forward_indices (added in previous commit) uses &[OHLCVBar] for date slicing. Add a variant that works directly on &[i64] timestamps from fxcache — no OHLCVBar needed.

  • Step 1: Add generate_walk_forward_indices_from_timestamps
/// Generate walk-forward index ranges from nanosecond timestamps.
/// Works directly on fxcache timestamps — no OHLCVBar conversion needed.
pub fn generate_walk_forward_indices_from_timestamps(
    timestamps_ns: &[i64],
    config: &WalkForwardConfig,
) -> Vec<FoldRange> {
    if timestamps_ns.is_empty() {
        return Vec::new();
    }
    let ts_to_date = |ts: i64| -> NaiveDate {
        chrono::DateTime::from_timestamp_nanos(ts).date_naive()
    };
    let data_start = ts_to_date(timestamps_ns[0]);
    let data_end = ts_to_date(*timestamps_ns.last().unwrap());
    // ... same fold logic as generate_walk_forward_indices
    // but uses partition_point on timestamps instead of bars
}
  • Step 2: Add compute_difficulty_from_features that reads ADX from features[40]

The current compute_difficulty(&[OHLCVBar]) computes ADX from raw bars. With fxcache, ADX is already at feature index 40. Add:

pub fn compute_difficulty_from_features(features: &[[f64; 42]]) -> f64 {
    if features.is_empty() { return 0.0; }
    let adx_idx = 40;
    let sum: f64 = features.iter().map(|f| f[adx_idx]).sum();
    sum / features.len() as f64
}
  • Step 3: Compile
SQLX_OFFLINE=true cargo check -p ml
  • Step 4: Commit

Task 2: DqnGpuData::upload_slices — contiguous array upload

Files:

  • Modify: crates/ml/src/cuda_pipeline/mod.rs

The current upload() takes &[([f64; 42], Vec<f64>)] and iterates element-by-element to flatten. Add upload_slices() that accepts pre-contiguous arrays from fxcache.

  • Step 1: Add upload_slices method
/// Upload contiguous feature/target slices to GPU (zero per-element iteration).
/// Accepts fxcache arrays directly — no tuple unpacking, no Vec<f64> allocation.
pub fn upload_slices(
    features: &[[f64; 42]],
    targets: &[[f64; 4]],
    ofi: &[[f64; 8]],
    stream: &Arc<CudaStream>,
) -> Result<Self, MLError> {
    let num_bars = features.len();
    let feature_dim = 42;

    // Cast f64 → f32 in one contiguous pass (no per-element Vec push)
    let flat_features: Vec<f32> = features.iter()
        .flat_map(|f| f.iter().map(|&v| v as f32))
        .collect();
    let flat_targets: Vec<f32> = targets.iter()
        .flat_map(|t| t.iter().map(|&v| v as f32))
        .collect();

    let features_gpu = clone_htod_f32_to_bf16(stream, &flat_features)?;
    let targets_gpu = clone_htod_f32_to_bf16(stream, &flat_targets)?;

    let mut data = Self {
        features: features_gpu,
        targets: targets_gpu,
        ofi_features: None,
        num_bars,
        feature_dim,
        aligned_state_dim: None,
    };

    // Upload OFI if present (non-zero)
    if !ofi.is_empty() && ofi.iter().any(|o| o.iter().any(|&v| v != 0.0)) {
        let flat_ofi: Vec<f32> = ofi.iter()
            .flat_map(|o| o.iter().map(|&v| v as f32))
            .collect();
        let ofi_gpu = clone_htod_f32_to_bf16(stream, &flat_ofi)?;
        data.ofi_features = Some(ofi_gpu);
    }

    Ok(data)
}
  • Step 2: Compile
SQLX_OFFLINE=true cargo check -p ml
  • Step 3: Commit

Task 3: Trainer API — init_from_fxcache

Files:

  • Modify: crates/ml/src/trainers/dqn/trainer/mod.rs
  • Modify: crates/ml/src/trainers/dqn/trainer/training_loop.rs

Add init_from_fxcache() that uploads the full fxcache dataset to GPU ONCE. This replaces the per-fold init_gpu_data + init_gpu_raw_buffers pair.

  • Step 1: Add init_from_fxcache on DQNTrainer
/// Upload full fxcache dataset to GPU ONCE. Call before the fold loop.
/// Replaces per-fold init_gpu_data + init_gpu_raw_buffers.
pub async fn init_from_fxcache(
    &mut self,
    features: &[[f64; 42]],
    targets: &[[f64; 4]],
    ofi: &[[f64; 8]],
) -> Result<()> {
    let stream = self.cuda_stream.as_ref()
        .ok_or_else(|| anyhow::anyhow!("CUDA stream required"))?;

    // Upload via DqnGpuData::upload_slices (contiguous, zero tuple unpacking)
    let gpu_data = DqnGpuData::upload_slices(features, targets, ofi, stream)?;
    let ofi_enabled = gpu_data.ofi_features.is_some();
    let raw_dim = if ofi_enabled { 53 } else { 45 };
    let aligned_dim = (raw_dim + 7) & !7;
    info!("init_from_fxcache: {} bars uploaded to GPU ({:.1} MB)",
        features.len(),
        (features.len() * (42 + 4 + 8) * 4) as f64 / 1_048_576.0);
    self.gpu_data = Some(gpu_data);

    // Also upload raw buffers for experience collector
    self.init_gpu_raw_buffers_from_slices(features, targets).await?;

    Ok(())
}
  • Step 2: Add init_gpu_raw_buffers_from_slices in training_loop.rs
/// Upload raw features + targets from contiguous slices (no tuple unpacking).
pub(crate) async fn init_gpu_raw_buffers_from_slices(
    &mut self,
    features: &[[f64; 42]],
    targets: &[[f64; 4]],
) -> Result<()> {
    if self.targets_raw_cuda.is_some() {
        return Ok(());
    }
    let stream = match self.cuda_stream {
        Some(ref s) => Arc::clone(s),
        None => return Ok(()),
    };

    let num_bars = features.len();
    let flat_targets: Vec<f32> = targets.iter()
        .flat_map(|t| t.iter().map(|&v| v as f32))
        .collect();
    let flat_features: Vec<f32> = features.iter()
        .flat_map(|f| f.iter().map(|&v| v as f32))
        .collect();

    self.targets_raw_cuda = Some(
        crate::cuda_pipeline::clone_htod_f32_to_bf16(&stream, &flat_targets)?
    );
    self.features_raw_cuda = Some(
        crate::cuda_pipeline::clone_htod_f32_to_bf16(&stream, &flat_features)?
    );

    info!("CUDA raw buffers uploaded: {} bars x 42 features + 4 targets", num_bars);
    Ok(())
}
  • Step 3: Add set_val_data_from_slices for validation

The trainer needs val_data for epoch-end Sharpe computation. Instead of Vec<(FeatureVector, Vec<f64>)>, accept slices:

/// Set validation data from contiguous slices (zero Vec<f64> allocation).
pub fn set_val_data_from_slices(
    &mut self,
    features: &[[f64; 42]],
    targets: &[[f64; 4]],
    ofi_val_offset: usize,
) {
    // Convert to the legacy format for now — val set is small (~50K bars)
    // Full GPU-native validation is a future task
    self.val_data = features.iter().zip(targets.iter())
        .map(|(f, t)| (*f, t.to_vec()))
        .collect();
    self.ofi_val_offset = ofi_val_offset;
    self.val_features_gpu = None;
    self.val_closes_gpu = None;
    self.val_ofi_gpu = None;
}
  • Step 4: Compile
SQLX_OFFLINE=true cargo check -p ml
  • Step 5: Commit

Task 4: Rewrite train_baseline_rl fold loop

Files:

  • Modify: crates/ml/examples/train_baseline_rl.rs

This is the integration task. Replace the entire fold loop to:

  1. Load fxcache ONCE
  2. Create DQNTrainer ONCE
  3. Upload to GPU ONCE via init_from_fxcache
  4. Generate index ranges from timestamps
  5. Per fold: slice features/targets, normalize, set_training_range, reset_for_fold, train
  • Step 1: Replace data loading section (lines ~798-891)

Replace the fxcache + DBN loading + OHLCVBar reconstruction with:

// Load fxcache — contains features[42] + targets[4] + OFI[8] + timestamps
// MBP-10 orderbook and trades data are pre-embedded in OFI features
let fxcache = load_fxcache_data(&args)?;
info!("Loaded fxcache: {} bars, features[42] + targets[4] + OFI[8]",
    fxcache.bar_count);

Where load_fxcache_data is a helper that does the cache lookup (same as current code but returns FxCacheData directly, no OHLCVBar conversion).

  • Step 2: Replace walk-forward window generation
let fold_ranges = generate_walk_forward_indices_from_timestamps(
    &fxcache.timestamps, &wf_config,
);
info!("Generated {} walk-forward folds", fold_ranges.len());
  • Step 3: Create DQN trainer + upload to GPU ONCE
let hp = load_hyperopt_params(&args.hyperopt_params, "dqn");
let hyperparams = build_dqn_hyperparams(args, &hp, &gpu_profile);
let mut trainer = DQNTrainer::new(hyperparams)?;

// One-time GPU upload of full dataset
let rt = tokio::runtime::Builder::new_current_thread()
    .enable_all().build()?;
rt.block_on(trainer.init_from_fxcache(
    &fxcache.features, &fxcache.targets, &fxcache.ofi,
))?;
  • Step 4: Replace fold loop body
for range in &fold_ranges {
    info!("--- Fold {} ---", range.fold);

    // Slice features for normalization (zero copy — just index bounds)
    let train_feat = &fxcache.features[range.train_start..range.train_end];
    let val_feat = &fxcache.features[range.val_start..range.val_end];
    let train_targets = &fxcache.targets[range.train_start..range.train_end];
    let val_targets = &fxcache.targets[range.val_start..range.val_end];

    // NormStats from training slice (CPU, ~10ms for 500K bars)
    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);

    // Set fold range on trainer (GPU arrays stay put, just index bounds change)
    trainer.set_training_range(
        range.train_start, range.train_end,
        range.val_start, range.val_end,
    );

    // Set validation data for epoch-end Sharpe
    trainer.set_val_data_from_slices(val_feat, val_targets, range.train_end - range.train_start);

    // Reset per-fold state (keep GPU infra)
    rt.block_on(trainer.reset_for_fold())?;

    // Checkpoint callback
    let output_dir = args.output_dir.clone();
    let fold = range.fold;
    let checkpoint_callback = move |epoch: usize, data: Vec<u8>, is_best: bool| -> Result<String> {
        let suffix = if is_best { "best" } else { &format!("epoch{}", epoch) };
        let path = output_dir.join(format!("dqn_fold{}_{}.safetensors", fold, suffix));
        std::fs::write(&path, &data)?;
        Ok(path.to_string_lossy().into_owned())
    };

    // Train — trainer reads from GPU arrays at fold range offset
    // NOTE: still passes train_norm as Vec<(FeatureVector, Vec<f64>)> for now.
    // Full elimination of this conversion requires changing train_with_data_full_loop
    // to read directly from GPU buffers — tracked as follow-up.
    let training_data = features_to_trainer_format_fast(&train_norm, train_targets);
    let val_data_fmt = features_to_trainer_format_fast(&val_norm, val_targets);

    match rt.block_on(trainer.train_with_preloaded_data(training_data, val_data_fmt, checkpoint_callback)) {
        Ok(metrics) => {
            info!("Fold {} complete: loss={:.6}, epochs={}", fold, metrics.loss, metrics.epochs_trained);
            dqn_results.push((fold, metrics.loss));
        }
        Err(e) => error!("Fold {} failed: {:#}", fold, e),
    }
}
  • Step 5: Add features_to_trainer_format_fast helper

Faster version that avoids per-bar vec![open, high, low, close]:

fn features_to_trainer_format_fast(
    features: &[[f64; 42]],
    targets: &[[f64; 4]],
) -> Vec<(FeatureVector, Vec<f64>)> {
    features.iter().zip(targets.iter())
        .map(|(f, t)| (*f, t.to_vec()))
        .collect()
}

This still allocates Vec<f64> per bar (4 elements), but eliminates:

  • Bar cloning (320MB)
  • Feature re-extraction (30s CPU)
  • Trainer re-creation (3s GPU init)
  • Data re-upload (1s HtoD)

The Vec<f64> allocation (~32 bytes per bar × 500K bars = 16MB) is the last remaining waste. Eliminating it requires changing train_with_data_full_loop to accept &[[f64; 4]] — tracked as Task 6.

  • Step 6: Delete dead code

Remove:

  • prepare_fold_data function

  • FoldData type alias

  • features_to_trainer_format function (replaced by _fast variant)

  • prefetched_data and fold prefetch thread

  • dqn_gpu_staged double-buffer

  • DqnGpuData::upload from double-buffer path

  • Step 7: Compile and test 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 2>&1 | tail -20

Expected: 3 folds × 3 epochs in <10s (was 159s).

  • Step 8: Commit

Task 5: Eliminate last Vec allocation (trainer API)

Files:

  • Modify: crates/ml/src/trainers/dqn/trainer/training_loop.rs
  • Modify: crates/ml/src/trainers/dqn/trainer/mod.rs

Change train_with_data_full_loop to accept (&[[f64; 42]], &[[f64; 4]]) instead of &[(FeatureVector, Vec<f64>)]. This eliminates the 4M × Vec<f64> allocation.

  • Step 1: Add train_fold_from_slices on DQNTrainer
/// Train one fold from contiguous feature/target slices.
/// Zero Vec<f64> allocation — reads targets directly as &[[f64; 4]].
pub async fn train_fold_from_slices<F>(
    &mut self,
    features: &[[f64; 42]],
    targets: &[[f64; 4]],
    checkpoint_callback: F,
) -> Result<TrainingMetrics>
where
    F: FnMut(usize, Vec<u8>, bool) -> Result<String> + Send,
{
    // GPU data already uploaded via init_from_fxcache — skip re-upload
    // Training loop uses self.features_raw_cuda / self.targets_raw_cuda
    // which are the FULL dataset on GPU. Fold range selects the slice.

    let num_bars = features.len();
    info!("train_fold_from_slices: {} bars (GPU range {}..{})",
        num_bars, self.fold_train_start, self.fold_train_end);

    // The training loop still needs training_data for:
    // 1. .len() → use fold range
    // 2. Curriculum ADX filter → read features[bar][40]
    // 3. Vol normalizer → read targets for returns
    // Build a lightweight wrapper that avoids Vec<f64>
    let training_data_ref: Vec<([f64; 42], [f64; 4])> = features.iter()
        .zip(targets.iter())
        .map(|(f, t)| (*f, *t))
        .collect();

    self.train_with_data_full_loop_slices(&training_data_ref, checkpoint_callback).await
}
  • Step 2: Add train_with_data_full_loop_slices in training_loop.rs

Mirror of train_with_data_full_loop but accepts &[([f64; 42], [f64; 4])]. The internal code only needs .len(), features[bar][adx_idx], and targets[bar] for vol normalizer — all work with fixed-size arrays.

  • Step 3: Update fold loop in train_baseline_rl to use new API

Replace features_to_trainer_format_fast + train_with_preloaded_data with:

match rt.block_on(trainer.train_fold_from_slices(&train_norm, train_targets, checkpoint_callback)) {
  • Step 4: Compile and test

  • Step 5: Commit


Task 6: Validate locally and on H100

  • Step 1: Local validation
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 2>&1 | tail -20

Expected: <10s for 3 folds × 3 epochs.

  • Step 2: Smoketest
SQLX_OFFLINE=true FOXHUNT_TEST_DATA=test_data/futures-baseline cargo test -p ml --lib -- smoke_tests::training_stability::test_gpu_collector_auto_initializes --ignored --nocapture 2>&1 | tail -5

Expected: PASS, <60s.

  • Step 3: H100 submission
./scripts/argo-train.sh dqn --epochs 50 --trials 0 --gpu-pool ci-training-h100

Expected:

  • fxcache hit (precomputed on PVC)
  • One-time GPU upload: <2s
  • Per-fold transition: <0.5s
  • Per-epoch (batch=8192, 488 steps): <5s
  • Total 50 epochs × 3 folds: <15 minutes

Performance Summary

Operation Before After
Data load (fxcache hit) ~5s ~5s (unchanged)
Walk-forward window gen ~200ms + 320MB cloned bars ~1ms (index ranges only)
Feature extraction per fold ~30s CPU 0s (pre-computed in fxcache)
Trainer creation per fold ~3s GPU init 0s (reused)
Data upload per fold ~1s HtoD 0s (stays on GPU)
Vec allocs per fold ~16MB (4M × 4 elements) 0 (Task 5: fixed-size arrays)
NormStats per fold ~10ms (unchanged) ~10ms (slice, no copy)
Total fold transition ~34s ~10ms
3-fold overhead ~102s ~30ms

Data Quality Improvement

The walk-forward now uses the FULL precomputed feature set from fxcache:

  • 42 market features: OHLCV derivatives + technical indicators + ADX + CUSUM
  • 8 OFI features: from MBP-10 order book (OFI L1/L5, depth imbalance, VPIN, Kyle's Lambda, bid/ask slope, trade imbalance)
  • 4 targets: raw close prices for portfolio simulation

This matches production: the model trains on exactly the same feature set it will see in real-time trading, including orderbook microstructure from MBP-10 and trade flow data.

Risks

Risk Impact Mitigation
NormStats computed per fold on CPU ~10ms — negligible Future: GPU reduction kernel
val_data still uses legacy format Small alloc (~50K × Vec) set_val_data_from_slices minimizes, full GPU val in follow-up
OHLCVBar removal breaks PPO path PPO uses same fold loop Keep PPO using legacy path for now, separate refactor
train_with_data_full_loop signature change Touches many callers Task 5 adds new method, doesn't change existing