diff --git a/docs/plans/2026-03-01-gpu-full-sweep-design.md b/docs/plans/2026-03-01-gpu-full-sweep-design.md new file mode 100644 index 000000000..215d87930 --- /dev/null +++ b/docs/plans/2026-03-01-gpu-full-sweep-design.md @@ -0,0 +1,118 @@ +# GPU Optimization Full Sweep — Design Document + +## Problem Statement + +Phases 1–2 built GPU optimization infrastructure (mixed precision, dynamic batching, gradient accumulation, DoubleBufferedLoader, EpochPrefetcher, GpuBufferPool, tensor core alignment). However, audit reveals: + +1. **Training binary bypasses all of it** — `train_baseline_rl.rs` uses raw `DQN::new()`/`PPO::new()`, never `DQNTrainer`/`PpoTrainer` +2. **Training binary explicitly disables Rainbow DQN** — hardcodes `use_per: false, use_dueling: false, use_distributional: false, use_noisy_nets: false, use_cql: false, use_iqn: false` +3. **PPO trainer lacks parity** — no mixed precision auto-detect, gradient accumulation is no-op +4. **Supervised trainers have correctness bugs** — Liquid/TLOB eval paths lack `detach()`, TLOB gradient clipping is a stub +5. **Mamba2 hardcoded to 4GB RTX 3050 Ti** — ignores L4/H100 capacity +6. **No multi-GPU support** — all trainers hardcoded to device 0 +7. **Ensemble inference is sequential** — 10 models one-by-one, no parallelism + +## Architecture + +### Section 1: Training Binary Refactor (P0) + +**Goal**: Connect `train_baseline_rl.rs` to `DQNTrainer`/`PpoTrainer`. + +**Current flow**: Binary → `DQN::new()` → manual per-bar loop → `dqn.train_step()` (bypasses all Trainer infrastructure) + +**New flow**: Binary → `DQNTrainer::new(config)` → `trainer.train_with_preloaded_data(features, val, callback)` per fold + +- Walk-forward fold iteration stays in the binary (data splitting logic) +- Per-fold training delegates to Trainers +- All GPU optimizations + Rainbow components activate automatically +- `DqnTrainerConfig` built from CLI args + hyperopt JSON, using `..DqnTrainerConfig::default()` for Rainbow defaults +- Same pattern for PPO: `PpoTrainer::new()` → `trainer.train(data)` + +**Key files**: +- Modify: `crates/ml/examples/train_baseline_rl.rs` +- Read: `crates/ml/src/trainers/dqn/trainer.rs` (train_with_preloaded_data interface) +- Read: `crates/ml/src/trainers/ppo.rs` (train interface) + +### Section 2: PPO Trainer Parity (P1) + +**Goal**: Bring PPO trainer to feature parity with DQN trainer. + +**Changes to `trainers/ppo.rs`**: + +1. **Mixed precision auto-detection**: Add `MixedPrecisionConfig::detect_from_gpu_name()` in `PpoTrainer::new()`, matching DQN trainer lines 327-330 +2. **Gradient accumulation loop**: In `update_policy()`, divide loss by `accumulation_steps`, call optimizer only every N mini-batches (mirroring DQN trainer lines 3137-3370) +3. **DoubleBufferedLoader**: Create `PpoDoubleBufferedLoader` or generalize `DoubleBufferedLoader` with a trait, so PPO fold transitions don't stall GPU + +### Section 3: Supervised Trainer Fixes (P1-P2) + +| Trainer | File | Fix | Details | +|---|---|---|---| +| **Liquid** | `trainers/liquid.rs:315-337` | `detach()` in eval | Add `.detach()` to forward pass outputs in `evaluate()` to prevent gradient graph accumulation | +| **TLOB** | `trainers/tlob.rs:398-427` | `detach()` in eval | Add `.detach()` in `validate_epoch()` | +| **TLOB** | `trainers/tlob.rs:474` | Gradient clipping | Replace stub `clip_gradients()` with real VarMap iteration + L2 norm clamping | +| **TLOB** | `trainers/tlob.rs:481` | Gradient norm | Replace stub `calculate_gradient_norm()` with real L2 norm computation over GradStore | +| **Mamba2** | `trainers/mamba2.rs:72-119` | HardwareBudget | Replace hardcoded 4GB constraints with `HardwareBudget::detect()` dynamic scaling | +| **All supervised** | Various config constructors | Tensor core alignment | Apply `align_to_tensor_cores()` to hidden layer widths | + +### Section 4: CUDA Pipeline Completion (P2) + +1. **Wire EpochPrefetcher**: Connect `EpochPrefetcher::spawn()` to DQN trainer fold transitions so disk I/O overlaps with GPU training +2. **Wire GpuBufferPool**: Use in DQN trainer's walk-forward loop instead of per-fold fresh allocation via `upload_dqn()` +3. **Align hidden dims**: Apply `align_to_tensor_cores()` to DQN/PPO hidden layer configs at trainer init (not just CUDA kernel sizing) + +### Section 5: Inference + Benchmarks (P3) + +1. **Ensemble parallel inference**: Replace sequential model iteration with `rayon::par_iter()` over ensemble models for sub-ms multi-model inference +2. **Benchmark BF16 variants**: Add mixed-precision benchmark configs to `dqn_benchmark.rs:473` and `tft_benchmark.rs:519,547` + +### Section 6: Multi-GPU with NCCL (P2) + +**Goal**: Native multi-GPU data parallelism via NCCL all-reduce on H100 SXM nodes (NVLink = 900 GB/s). + +**Dependencies**: `cudarc` (transitive via Candle) with `nccl` feature. + +**New file**: `crates/ml/src/cuda_pipeline/multi_gpu.rs` + +**Components**: + +1. **`NcclGradientSync`**: Wraps `cudarc::nccl::Comm`. After `loss.backward()`, iterates VarMap parameters, extracts `CudaSlice`, calls `all_reduce_in_place(Sum)`, divides by world_size. Native GPU-to-GPU transfer. + +2. **`MultiGpuConfig`**: + ``` + devices: Vec // enumerated from CUDA_VISIBLE_DEVICES + sync_every_n_steps: usize // default 1 + ``` + +3. **Device enumeration**: `MultiGpuConfig::detect()` reads `CUDA_VISIBLE_DEVICES` or `--gpu-ids` CLI arg, creates `Vec`. + +4. **Data sharding**: Mini-batch split evenly across devices. Each device processes `batch_size / n_gpus` samples. + +5. **Model replication**: Clone VarMap to each device at init. Gradients stay synchronized via NCCL all-reduce. + +6. **Trainer integration**: Optional `multi_gpu: Option` field. Single-GPU remains default, zero overhead when not configured. + +**Scope**: Single-node data parallelism only. Cross-node distributed training is out of scope. + +## Out of Scope + +- **Cross-node distributed training**: Requires process management (MPI/torchrun equivalent) +- **Pinned memory**: Candle doesn't expose `cudaMallocHost`. GpuBufferPool approximates with reusable staging buffers +- **Model parallelism / tensor sharding**: Models fit in single-GPU memory even on L4 (24GB) + +## Implementation Order + +1. P0: Training binary refactor (unlocks everything else) +2. P1: PPO parity + Liquid/TLOB correctness fixes +3. P2: Mamba2 HardwareBudget + CUDA pipeline wiring + Multi-GPU NCCL +4. P3: Ensemble parallelism + benchmark updates + +## Success Criteria + +- `train_baseline_rl.rs` uses `DQNTrainer`/`PpoTrainer` — all GPU optimizations activate +- Rainbow DQN (6 components + IQN + CQL) enabled in production training +- PPO has mixed precision, gradient accumulation, DoubleBufferedLoader +- Liquid/TLOB eval paths use `detach()` — no gradient graph in validation +- TLOB gradient clipping is real (not a stub) +- Mamba2 scales batch size to GPU capacity (not capped at 16) +- Multi-GPU training works on 2+ GPU nodes via NCCL +- 2,437+ ml tests pass, 0 failures