diff --git a/docs/plans/2026-03-01-io-pipeline-wiring-design.md b/docs/plans/2026-03-01-io-pipeline-wiring-design.md new file mode 100644 index 000000000..878b15fb3 --- /dev/null +++ b/docs/plans/2026-03-01-io-pipeline-wiring-design.md @@ -0,0 +1,71 @@ +# I/O Pipeline Wiring — Design Document + +## Problem Statement + +GPU optimization infrastructure is fully built but 3 I/O pipeline components are unused: + +1. **EpochPrefetcher** (`cuda_pipeline/prefetch.rs`) — background thread loads next fold's data. DQNTrainer has `set_prefetcher()`/`take_prefetched_data()` methods but `train_baseline_rl.rs` never calls them. + +2. **GpuBufferPool** (`cuda_pipeline/mod.rs:278-367`) — pre-allocated CPU staging buffers for zero-alloc GPU uploads. DQNTrainer initializes it in `new()` but never references `self.buffer_pool`. + +3. **DoubleBufferedLoader** (`cuda_pipeline/double_buffer.rs`) — ping-pong GPU tensor slots for zero-downtime fold transitions. Not used by any trainer. + +All core GPU optimizations (mixed precision, dynamic batching, gradient accumulation, tensor core alignment, Rainbow DQN) are already wired and active. This work is the final I/O throughput piece. + +## Architecture + +### Fold-Level Pipeline (train_baseline_rl.rs) + +``` +for fold_idx in 0..num_folds: + // 1. Get data: from prefetcher (fold 1+) or fresh load (fold 0) + let fold_data = if let Some(prefetched) = trainer.take_prefetched_data() { + prefetched? + } else { + load_fold_data(fold_idx) // first fold, no prefetch available + }; + + // 2. Kick off prefetch for NEXT fold (if not last) + if fold_idx + 1 < num_folds { + let next_path = fold_paths[fold_idx + 1].clone(); + trainer.set_prefetcher(EpochPrefetcher::spawn(move || { + load_fold_data_from_path(&next_path) + })); + } + + // 3. Train this fold (buffer_pool + double_buffer used internally) + trainer.train_with_preloaded_data(fold_data, val_data, callback).await?; +``` + +### Trainer-Internal Pipeline (DQNTrainer) + +Inside `train_with_preloaded_data()`: + +1. **GpuBufferPool**: Use `self.buffer_pool.upload_dqn()` instead of `DqnGpuData::upload()` to avoid per-fold heap allocation. The pool pre-allocates staging buffers for 100k bars once. + +2. **DoubleBufferedLoader**: Add `double_buffer: Option` field. In the fold transition path, upload next fold's data to staging slot while current fold trains. `swap()` at fold boundary is O(1). + +### PPO Path + +PPO uses `PpoTrainer` which has a different data flow (trajectory collection, not preloaded features). The prefetcher and buffer pool don't directly apply. PPO already has gradient accumulation and mixed precision wired. No changes needed. + +## Files to Modify + +| File | Change | +|---|---| +| `crates/ml/examples/train_baseline_rl.rs` | Wire `set_prefetcher()`/`take_prefetched_data()` in DQN fold loop | +| `crates/ml/src/trainers/dqn/trainer.rs` | Use `self.buffer_pool.upload_dqn()` in `train_with_preloaded_data()`, add `DoubleBufferedLoader` field, integrate in fold transitions | + +## Out of Scope + +- PPO I/O pipeline (different data flow — trajectory-based, not preloaded) +- Multi-node distributed training (single-node multi-GPU kept as-is) +- Pinned memory (Candle doesn't expose `cudaMallocHost`) + +## Success Criteria + +- `train_baseline_rl.rs` DQN path uses prefetcher for fold N+1 while fold N trains +- DQNTrainer uses `buffer_pool.upload_dqn()` instead of direct `DqnGpuData::upload()` +- DoubleBufferedLoader manages active/staging GPU slots during fold transitions +- All 2444+ ml tests pass +- Training binary compiles clean