foxhunt/docs/plans/2026-03-01-gpu-full-sweep-design.md

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