7cb8c305df
Audit found training binary bypasses all DQNTrainer/PpoTrainer infrastructure and explicitly disables Rainbow DQN. Design covers 6 sections: - P0: Wire train_baseline_rl.rs to Trainers - P1: PPO mixed precision + gradient accumulation + Liquid/TLOB detach fixes - P2: Mamba2 HardwareBudget + CUDA pipeline wiring + NCCL multi-GPU - P3: Ensemble parallel inference + BF16 benchmarks Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
119 lines
6.5 KiB
Markdown
119 lines
6.5 KiB
Markdown
# 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
|