docs: add zero-CPU DQN training hot path design spec
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 <noreply@anthropic.com>
This commit is contained in:
jgrusewski
260
docs/superpowers/specs/2026-03-10-zero-cpu-hotpath-design.md
Normal file
260
docs/superpowers/specs/2026-03-10-zero-cpu-hotpath-design.md
Normal file
@@ -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::<f32>()` | Per train step | GpuTrainingGuard::check() |
|
||||
| 2 | gradient_utils.rs:108 | `total_sum.to_scalar::<f32>()` | 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::<f32>()` | Every 50 steps | GPU Q-value accumulator |
|
||||
| 6 | dqn.rs:3618 | `to_vec2::<f32>()` | Every 50 steps | GPU Q-value stats reduction |
|
||||
| 7 | trainer.rs:4113 | `to_vec1::<u32>()` | Per experience batch | GPU-resident action tensor |
|
||||
| 8 | trainer.rs:4302-4303 | `to_vec1::<u32>()` | Per experience batch | GPU-resident action tensor |
|
||||
| 9 | trainer.rs:4334-4335 | `to_vec1::<u32>()` | 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<f32>` — running sum of losses across steps
|
||||
- `grad_norm_sum: CudaSlice<f32>` — running sum of grad norms
|
||||
- `step_count: CudaSlice<i32>` — 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<f32>, // 7 × f32 pinned
|
||||
loss_sum: CudaSlice<f32>,
|
||||
grad_norm_sum: CudaSlice<f32>,
|
||||
step_count: CudaSlice<i32>,
|
||||
device: Device,
|
||||
}
|
||||
```
|
||||
|
||||
**Methods:**
|
||||
- `new(device) -> Result<Self>` — compile PTX, allocate pinned + device buffers
|
||||
- `check(loss_gpu: &Tensor, grad_norm_gpu: &Tensor, threshold: f32, warmup: bool) -> Result<GuardResult>` — 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::<f32>() — 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::<f32>()` 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::<u32>()` 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::<u32>()?; // 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<FactoredAction>` 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<FactoredAction>` 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<f32>, Vec<i32>)>` 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
|
||||
Reference in New Issue
Block a user