# 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<T>` 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<Device>       // 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<Device>`.

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<MultiGpuConfig>` 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
