From 41440e9ae7583d97eecbcb59eec822491852bb0a Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Tue, 10 Mar 2026 13:50:25 +0100 Subject: [PATCH] docs: add zero-CPU DQN training hot path design spec MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Defines architecture for eliminating all 10 GPU→CPU sync barriers from the training loop via 4 new GPU components: Training Guard (pinned memory predicates), Q-Value Monitor (on-device accumulator), GPU-resident action selection, and async experience collector readback. Co-Authored-By: Claude Opus 4.6 --- .../2026-03-10-zero-cpu-hotpath-design.md | 260 ++++++++++++++++++ 1 file changed, 260 insertions(+) create mode 100644 docs/superpowers/specs/2026-03-10-zero-cpu-hotpath-design.md diff --git a/docs/superpowers/specs/2026-03-10-zero-cpu-hotpath-design.md b/docs/superpowers/specs/2026-03-10-zero-cpu-hotpath-design.md new file mode 100644 index 000000000..c0ae948a5 --- /dev/null +++ b/docs/superpowers/specs/2026-03-10-zero-cpu-hotpath-design.md @@ -0,0 +1,260 @@ +# Zero-CPU DQN Training Hot Path — Design Spec + +**Goal:** Eliminate all GPU→CPU synchronization barriers (`to_scalar()`, `to_vec1()`, `to_vec2()`, `memcpy_dtoh`) from the DQN training hot path so the GPU pipeline never stalls between `backward()` and `optimizer.step()`. + +**Architecture:** Four new GPU components replace 10 CPU readback sites with on-device predicates, pinned-memory polling, and GPU-resident accumulators. All monitoring/safety checks (NaN detection, loss clipping, gradient collapse, Q-value divergence) execute as CUDA kernels writing boolean halt flags to pinned host memory — zero `cudaStreamSynchronize` barriers per training step. + +**Tech Stack:** CUDA (NVRTC compiled), cudarc 0.17, Candle v0.9.1 (Rust ML framework), pinned host memory via `cudarc::driver::CudaHostBuf` + +--- + +## Readback Inventory (10 sites to eliminate) + +| # | File:Line | Pattern | Frequency | Replacement | +|---|-----------|---------|-----------|-------------| +| 1 | trainer.rs:4531 | `stacked.to_vec1::()` | Per train step | GpuTrainingGuard::check() | +| 2 | gradient_utils.rs:108 | `total_sum.to_scalar::()` | Per accumulation boundary | GpuTrainingGuard NaN predicate | +| 3 | trainer.rs:4739 | `mean_all().to_scalar()` | Per accumulation boundary | GpuTrainingGuard::read_accumulators() | +| 4 | trainer.rs:4751 | `mean_all().to_scalar()` | Per accumulation boundary | GpuTrainingGuard::read_accumulators() | +| 5 | trainer.rs:4860 | `mean_all().to_scalar::()` | Every 50 steps | GPU Q-value accumulator | +| 6 | dqn.rs:3618 | `to_vec2::()` | Every 50 steps | GPU Q-value stats reduction | +| 7 | trainer.rs:4113 | `to_vec1::()` | Per experience batch | GPU-resident action tensor | +| 8 | trainer.rs:4302-4303 | `to_vec1::()` | Per experience batch | GPU-resident action tensor | +| 9 | trainer.rs:4334-4335 | `to_vec1::()` | CPU fallback (non-cuda) | N/A (CPU-only path) | +| 10 | gpu_experience_collector.rs:769-772 | `memcpy_dtoh` | Per GPU collection | Async deferred readback | + +**Note:** Site #9 is the CPU-only fallback path (no GPU selector). This only runs when `#[cfg(not(feature = "cuda"))]` or when GPU selector init fails. We leave this path as-is since it's not a GPU hot path. + +--- + +## Component A: GPU Training Guard + +**New files:** +- `crates/ml/src/cuda_pipeline/training_guard_kernel.cu` +- `crates/ml/src/cuda_pipeline/gpu_training_guard.rs` + +### CUDA Kernel: `training_guard_check` + +Single-thread kernel (1 block × 1 thread) that reads GPU-resident loss and grad_norm scalars and writes results to pinned host memory. + +**Inputs (device pointers):** +- `loss_val: *const f32` — scalar loss from train_step (GPU-resident) +- `grad_norm_val: *const f32` — scalar grad_norm from clip_grad_norm (GPU-resident) +- `loss_clip_threshold: f32` — 1e6 (constant) +- `grad_collapse_threshold: f32` — LR × multiplier (passed from Rust) +- `warmup_active: i32` — 1 if within warmup period (skip grad collapse check) + +**Outputs (pinned host memory, 7 × f32 = 28 bytes):** +- `[0]`: halt_nan (1.0 if loss or grad_norm is NaN/Inf, else 0.0) +- `[1]`: halt_loss_clip (1.0 if loss > threshold, else 0.0) +- `[2]`: halt_grad_collapse (1.0 if grad_norm < threshold && !warmup, else 0.0) +- `[3]`: clipped_loss (min(loss, threshold)) +- `[4]`: raw_loss (unclipped, for logging) +- `[5]`: raw_grad_norm (for logging) +- `[6]`: reserved (padding to 32 bytes) + +**On-device accumulators** (persistent CudaSlice, not pinned): +- `loss_sum: CudaSlice` — running sum of losses across steps +- `grad_norm_sum: CudaSlice` — running sum of grad norms +- `step_count: CudaSlice` — number of steps accumulated + +### CUDA Kernel: `training_guard_accumulate` + +Adds current loss and grad_norm to running sums, increments counter. + +### Rust Wrapper: `GpuTrainingGuard` + +```rust +pub struct GpuTrainingGuard { + check_func: CudaFunction, + accumulate_func: CudaFunction, + pinned_output: CudaHostBuf, // 7 × f32 pinned + loss_sum: CudaSlice, + grad_norm_sum: CudaSlice, + step_count: CudaSlice, + device: Device, +} +``` + +**Methods:** +- `new(device) -> Result` — compile PTX, allocate pinned + device buffers +- `check(loss_gpu: &Tensor, grad_norm_gpu: &Tensor, threshold: f32, warmup: bool) -> Result` — launches kernel, reads pinned memory (no cudaStreamSync) +- `accumulate(loss_gpu: &Tensor, grad_norm_gpu: &Tensor) -> Result<()>` — adds to running sums +- `read_accumulators() -> Result<(f64, f64)>` — single DtoH of mean loss + mean grad_norm (epoch boundary only) +- `reset_accumulators() -> Result<()>` — zeros device buffers + +```rust +pub struct GuardResult { + pub halt_nan: bool, + pub halt_loss_clip: bool, + pub halt_grad_collapse: bool, + pub clipped_loss: f32, + pub raw_grad_norm: f32, +} +``` + +### Wiring (trainer.rs) + +**train_step_optimized** (replaces sites #1): +``` +// BEFORE: stacked.to_vec1::() — GPU sync barrier +// AFTER: +let guard_result = self.training_guard.check( + &gpu_result.loss_gpu, + &gpu_result.grad_norm_gpu, + grad_collapse_threshold, + !past_warmup, +)?; +if guard_result.halt_nan { return Err(...) } +if guard_result.halt_loss_clip { log warning } +if guard_result.halt_grad_collapse { increment counter, check patience } +self.training_guard.accumulate(&gpu_result.loss_gpu, &gpu_result.grad_norm_gpu)?; +``` + +**train_step_with_accumulation** (replaces sites #3, #4): +``` +// BEFORE: gpu_loss_tensors concat → mean_all().to_scalar() +// AFTER: accumulate() called per sub-step, read_accumulators() at boundary +let (avg_loss, avg_grad_norm) = self.training_guard.read_accumulators()?; +self.training_guard.reset_accumulators()?; +``` + +**check_gradients_finite** (replaces site #2): +Add `#[cfg(feature = "cuda")]` bypass: when CUDA is active, the training guard kernel already checks NaN/Inf on the loss and grad_norm scalars. The per-parameter NaN check in `check_gradients_finite` is redundant with the post-clip norm check — if any gradient was NaN, the norm would be NaN, which the guard catches. + +--- + +## Component B: GPU Q-Value Monitor + +Extends `training_guard_kernel.cu` with a Q-value accumulation kernel. + +### CUDA Kernel: `qvalue_accumulate` + +**Inputs:** +- `q_values: *const f32` — [batch × actions] Q-value matrix +- `batch_size: i32, num_actions: i32` +- `qvalue_sum: *mut f32` — device accumulator (running sum of max-Q per sample) +- `qvalue_count: *mut i32` — device counter + +**Operation:** Each thread handles one sample. Finds max Q across actions, atomicAdd to qvalue_sum. + +### CUDA Kernel: `qvalue_divergence_check` + +**Inputs:** +- `q_values: *const f32` — [1, actions] single-sample Q-values +- `num_actions: i32` +- `divergence_threshold: f32` — 10000.0 + +**Output (pinned, 4 × f32):** +- `[0]`: q_min +- `[1]`: q_max +- `[2]`: q_mean +- `[3]`: divergence_detected (1.0 if |q_min| or |q_max| > threshold) + +### Wiring + +**estimate_avg_q_value_with_early_stopping** (replaces site #5): +- Forward pass stays as-is (Candle tensor) +- Replace `max(1) → mean_all() → to_scalar()` with `qvalue_accumulate` kernel call +- Read accumulator at epoch boundary, not per 50-step check + +**log_q_values** (replaces site #6): +- Replace `to_vec2::()` with `qvalue_divergence_check` kernel +- Pinned readback: 4 floats instead of entire Q-matrix +- Prometheus gauges set from pinned values + +--- + +## Component C: GPU-Resident Action Selection + +No new kernels needed — `select_actions()`, `select_actions_routed()`, and `select_actions_branching()` already return GPU-resident `Tensor`. The readback happens in the **caller** (trainer.rs) which does `to_vec1::()` to iterate actions on CPU. + +### Solution: Keep actions on GPU, defer routing to experience storage + +**select_actions_batch_gpu** (replaces sites #7, #8): + +Currently at trainer.rs:4302: +```rust +let action_indices = action_indices_tensor.to_vec1::()?; // GPU sync! +for &idx in &action_indices { ... route_action(exposure, spread) ... } +``` + +Replace with: +1. Keep `action_indices_tensor` on GPU +2. When `use_routed` is true: actions are already factored indices (0-44), store tensor directly +3. When `use_routed` is false AND `use_branching` is false: actions are exposure indices (0-4). Port `route_action()` to a batch CUDA kernel that converts exposure → factored index on-device. +4. When `use_branching` is true: actions are already factored (0-44) from branching kernel + +**New CUDA kernel: `batch_route_exposure_to_factored`** + +Add to `epsilon_greedy_kernel.cu`: +```cuda +extern "C" __global__ void batch_route_exposure_to_factored( + const unsigned int* exposure_actions, // [N] exposure indices 0-4 + unsigned int* factored_actions, // [N] output factored indices 0-44 + float spread, float median_spread, + float volatility, float median_volatility, + int N +) { + int i = blockIdx.x * blockDim.x + threadIdx.x; + if (i >= N) return; + + // Port of OrderRouter::route() + int order_type; + if (median_spread <= 0.0f) order_type = 0; // Market + else if (spread < median_spread) order_type = 1; // LimitMaker + else if (spread > 2.0f * median_spread) order_type = 0; // Market + else order_type = 2; // IoC + + int urgency; + if (median_volatility <= 0.0f) urgency = 1; // Normal + else if (volatility > 1.5f * median_volatility) urgency = 2; // Aggressive + else if (volatility < 0.5f * median_volatility) urgency = 0; // Patient + else urgency = 1; // Normal + + factored_actions[i] = exposure_actions[i] * 9 + order_type * 3 + urgency; +} +``` + +**Downstream consumer change:** The experience loop at trainer.rs:2184+ currently uses `Vec` from `select_actions_batch_gpu()`. With GPU-resident actions, we need to either: +- (a) Change the experience loop to accept GPU tensor indices + batch DtoH at the end, or +- (b) Keep `actions` as a `Vec` but derive it from a single batch readback at the loop boundary instead of per-batch + +**Chosen approach:** (b) — Single batch readback. The experience loop needs CPU-side `FactoredAction` values for reward calculation (`calculate_reward_for_action`) and portfolio tracking. These are unavoidable CPU operations. The readback happens ONCE per experience collection (not per training step), so it's outside the critical train_step hot path. The key win is eliminating the per-step readbacks in the training loop (#1-#6). + +--- + +## Component D: Experience Collector Async Readback + +**gpu_experience_collector.rs:769-772** — `memcpy_dtoh` for rewards_cpu/actions_cpu. + +These are ~0.5MB downloads needed for `pnl_history` (Sharpe ratio) and `monitor.track_reward/track_action`. They happen once per experience collection (every epoch), not per training step. + +### Solution: Async stream + deferred poll + +1. Launch `memcpy_dtoh_async` on a separate CUDA stream +2. Return `GpuExperienceBatch` immediately (with empty rewards_cpu/actions_cpu) +3. Add `poll_monitoring_data(&mut self) -> Option<(Vec, Vec)>` that checks stream completion +4. Trainer calls `poll_monitoring_data()` at epoch boundary, updates pnl_history lazily + +**Alternative (simpler):** Keep synchronous but move to AFTER `insert_batch_tensors()`. The DtoD copies for replay buffer insertion don't depend on the monitoring downloads. This reorders operations so the GPU pipeline isn't stalled during the insert. + +**Chosen approach:** Reorder — simpler, same effect. The `stream.synchronize()` at line 777 already waits for DtoD copies. Move the monitoring downloads AFTER the sync, or use a separate stream for monitoring downloads that doesn't block the main pipeline. + +--- + +## Success Criteria + +1. `grep -rn 'to_scalar\|to_vec0\|to_vec1\|to_vec2' trainer.rs` → 0 matches in `#[cfg(feature = "cuda")]` blocks of `train_step_optimized` and `train_step_with_accumulation` +2. `check_gradients_finite` bypassed under `#[cfg(feature = "cuda")]` (guard catches NaN) +3. `log_q_values` uses GPU reduction instead of `to_vec2` +4. All existing tests pass: `SQLX_OFFLINE=true cargo test -p ml -p ml-dqn -p ml-core --lib` +5. 0 clippy warnings in modified files +6. Training semantics unchanged (loss values, action distributions identical) + +## Non-Goals + +- CPU-only fallback paths (`#[cfg(not(feature = "cuda"))]`) are unchanged +- CPU argmax fallback in `select_actions_batch_gpu` (line 4334) is for when GPU selector fails to init — not a hot path +- `select_actions_batch` (the older, non-GPU-tensor path at line 4090+) is the CPU experience fallback — readback there is unavoidable since it feeds CPU route_action loop +- Experience collector monitoring downloads are low-frequency (once per epoch) and can tolerate ~1ms latency