From a38dd64a712a1d6bc1e446c30e27b7469c3bbff1 Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Sat, 28 Feb 2026 12:59:45 +0100 Subject: [PATCH] =?UTF-8?q?docs:=20add=20PPO=20CUDA=20pipeline=20Phase=202?= =?UTF-8?q?c=20design=20=E2=80=94=20zero-roundtrip=20experience=20collecti?= =?UTF-8?q?on?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Single monolithic kernel with in-kernel GAE, full portfolio simulation, shared device functions header, 5-layer critic matching trainer exactly. Co-Authored-By: Claude Opus 4.6 --- .../2026-02-28-ppo-cuda-pipeline-design.md | 201 ++++++++++++++++++ 1 file changed, 201 insertions(+) create mode 100644 docs/plans/2026-02-28-ppo-cuda-pipeline-design.md diff --git a/docs/plans/2026-02-28-ppo-cuda-pipeline-design.md b/docs/plans/2026-02-28-ppo-cuda-pipeline-design.md new file mode 100644 index 000000000..680c22c7f --- /dev/null +++ b/docs/plans/2026-02-28-ppo-cuda-pipeline-design.md @@ -0,0 +1,201 @@ +# PPO CUDA Pipeline Design — Zero-Roundtrip Experience Collection (Phase 2c) + +## Goal + +Build a single monolithic CUDA kernel that runs the entire PPO experience collection pipeline — actor forward pass, softmax categorical sampling, critic forward pass, portfolio simulation, barrier tracking, diversity entropy, curiosity inference, reward combination, and GAE advantage computation — with zero CPU-GPU roundtrips per timestep. 128 parallel episodes × 500 timesteps = 64,000 training-ready experience tuples per kernel launch. + +## Architecture + +Single entry point `ppo_full_experience_kernel` in `ppo_experience_kernel.cu`. Shared device functions (RNG, matmul, portfolio, barriers, diversity, curiosity) extracted from the DQN kernel into `common_device_functions.cuh` and reused by both kernels via source concatenation at NVRTC compile time. + +Thread mapping: 1 thread = 1 episode, N=128 threads default. Each thread processes L=500 timesteps sequentially in two phases: forward rollout (collect experiences) then backward GAE scan (compute advantages and returns). + +## Kernel Phases + +### Phase A — Forward Rollout (L timesteps per thread) + +For each timestep t: + +1. Read 51 market features from global memory +2. Compute 3 portfolio features from simulation state (normalized value, normalized position, spread) +3. **Actor forward:** state[54] → h1[128] (ReLU) → h2[64] (ReLU) → logits[45] +4. **Softmax + categorical sample:** logits → exp → normalize → CDF scan → LCG random draw → action index + log(p[action]) +5. **Critic forward:** state[54] → h1[512] (ReLU) → h2[384] (ReLU) → h3[256] (ReLU) → h4[128] (ReLU) → h5[64] (ReLU) → value[1] +6. **Portfolio simulation** (reused from DQN): action → exposure mapping, transaction costs, cash/position tracking, reversals +7. **Barrier tracking** (reused): profit-take, stop-loss, time-based exit detection +8. **Diversity entropy** (reused): action window penalty when entropy < 1.0 +9. **Curiosity inference** (reused): forward dynamics model prediction error +10. **Reward combination:** PnL × (1 + barrier_scale) + diversity_penalty + curiosity_bonus - risk_penalty +11. Store per-timestep: state[54], action, log_prob, value, reward, done +12. Episode reset on done (reinitialize portfolio, barriers, diversity window) + +### Phase B — GAE Backward Scan (after all L timesteps) + +Each thread scans backward over its own episode's stored rewards, values, and dones: + +``` +advantages[L-1] = delta[L-1] // where delta = reward + gamma * V(next) * (1-done) - V(current) +for t = L-2 down to 0: + delta[t] = rewards[t] + gamma * values[t+1] * (1 - dones[t]) - values[t] + advantages[t] = delta[t] + gamma * lambda * (1 - dones[t]) * advantages[t+1] +returns[t] = advantages[t] + values[t] +``` + +Write advantages and returns to output buffers. Values buffer is NOT downloaded — only needed for GAE. + +## Shared Device Functions Header + +Extract from `dqn_experience_kernel.cu` into `common_device_functions.cuh`: + +| Function | Purpose | +|----------|---------| +| `gpu_random` | LCG RNG, returns float in [0, 1) | +| `leaky_relu` | Activation with configurable alpha | +| `matvec_leaky_relu` | Matrix-vector multiply + bias + activation (use_leaky=0 for ReLU) | +| `action_to_exposure` | 45-action index → position exposure float | +| `action_to_tx_cost` | 45-action index → transaction cost rate | +| `barrier_init/check/reset` | Triple barrier state management | +| `diversity_entropy` | Action diversity penalty computation | +| `curiosity_inference` | Forward dynamics model inference | + +New PPO-specific functions in `ppo_experience_kernel.cu`: + +| Function | Purpose | +|----------|---------| +| `ppo_actor_forward` | MLP: 54→128(ReLU)→64(ReLU)→45 raw logits | +| `softmax_sample` | Stable softmax → CDF scan → categorical sample → (action, log_prob) | +| `ppo_critic_forward` | MLP: 54→512(ReLU)→384(ReLU)→256(ReLU)→128(ReLU)→64(ReLU)→1 | +| `compute_gae_backward` | Backward scan: rewards + values + dones → advantages + returns | + +NVRTC compilation via source concatenation in Rust: +```rust +let common = include_str!("common_device_functions.cuh"); +let kernel = include_str!("ppo_experience_kernel.cu"); +let full_source = format!("{}\n{}", common, kernel); +``` + +DQN kernel refactored to use the same pattern. + +## Weight Layout + +### Actor (6 tensors, 18,093 parameters) + +| VarMap Key | Shape | CUDA name | +|-----------|-------|-----------| +| `policy_layer_0.weight` | [128, 54] | `pw1` | +| `policy_layer_0.bias` | [128] | `pb1` | +| `policy_layer_1.weight` | [64, 128] | `pw2` | +| `policy_layer_1.bias` | [64] | `pb2` | +| `policy_output.weight` | [45, 64] | `pw3` | +| `policy_output.bias` | [45] | `pb3` | + +### Critic (12 tensors, 346,369 parameters) + +| VarMap Key | Shape | CUDA name | +|-----------|-------|-----------| +| `value_layer_0.weight` | [512, 54] | `vw1` | +| `value_layer_0.bias` | [512] | `vb1` | +| `value_layer_1.weight` | [384, 512] | `vw2` | +| `value_layer_1.bias` | [384] | `vb2` | +| `value_layer_2.weight` | [256, 384] | `vw3` | +| `value_layer_2.bias` | [256] | `vb3` | +| `value_layer_3.weight` | [128, 256] | `vw4` | +| `value_layer_3.bias` | [128] | `vb4` | +| `value_layer_4.weight` | [64, 128] | `vw5` | +| `value_layer_4.bias` | [64] | `vb5` | +| `value_output.weight` | [1, 64] | `vw6` | +| `value_output.bias` | [1] | `vb6` | + +### Curiosity (4 tensors, reuses existing CuriosityWeightSet) + +| VarMap Key | Shape | CUDA name | +|-----------|-------|-----------| +| `fc1.weight` | [64, 35] | `cw1` | +| `fc1.bias` | [64] | `cb1` | +| `fc2.weight` | [32, 64] | `cw2` | +| `fc2.bias` | [32] | `cb2` | + +## Per-Thread Memory Budget + +| Array | Size | Bytes | +|-------|------|-------| +| `state[54]` | 54 floats | 216 | +| `logits[45]` + `probs[45]` | 90 floats | 360 | +| Actor scratch: `h1[128]`, `h2[64]` | 192 floats | 768 | +| Critic scratch: `scratch_a[512]`, `scratch_b[512]` (ping-pong) | 1,024 floats | 4,096 | +| `next_state[54]` | 54 floats | 216 | +| Curiosity scratch: `cur_h[64]` | 64 floats | 256 | +| Portfolio state: 8 floats | 8 floats | 32 | +| Barrier state: 5 floats | 5 floats | 20 | +| Diversity window: 100 ints | 100 ints | 400 | +| GAE buffers: `values_buf[L]`, `rewards_buf[L]`, `dones_buf[L]` | 3 × 500 | 6,000 | +| **Total per thread** | | **~12.4 KB** | + +At 128 threads: ~1.6 MB total local memory. Well within H100/L40S capacity. + +Critic uses ping-pong pattern: two 512-wide scratch buffers alternate between layers, avoiding 5 separate arrays. + +## Output Buffers + +| Buffer | Shape | Type | Downloaded | +|--------|-------|------|-----------| +| `states` | [N × L × 54] | f32 | Yes | +| `actions` | [N × L] | i32 | Yes | +| `old_log_probs` | [N × L] | f32 | Yes | +| `advantages` | [N × L] | f32 | Yes | +| `returns` | [N × L] | f32 | Yes | +| `dones` | [N × L] | i32 | Yes | +| `values` | [N × L] | f32 | No (only used for in-kernel GAE) | + +Total download per launch: N × L × (54 + 1 + 1 + 1 + 1 + 1) × 4 bytes = 128 × 500 × 59 × 4 = ~14.4 MB. + +## Trainer Integration + +1. Add `gpu_ppo_collector: Option` to `PpoTrainer` behind `#[cfg(feature = "cuda")]` +2. Initialize when CUDA device present and curiosity module available +3. GPU path in `collect_rollouts()`: kernel launch → download → build `TrajectoryBatch` +4. CPU fallback: existing code unchanged +5. Weight sync after `update_mlp()`: actor + critic + curiosity + +Training flow: +``` +GPU kernel launch (N=128 episodes × L=500 timesteps) + ↓ +Download 6 buffers (~14.4 MB) + ↓ +Build TrajectoryBatch from flat arrays + ↓ +update_mlp() [existing Candle code, unchanged]: + For each PPO epoch (10): + For each mini-batch (512): + Recompute log_probs via Candle actor forward + ratio = exp(new_log_probs - old_log_probs) + PPO clipped surrogate loss + critic MSE loss + Backward + optimizer step + ↓ +Sync actor + critic + curiosity weights → GPU +``` + +## File Layout + +``` +crates/ml/src/cuda_pipeline/ +├── common_device_functions.cuh # NEW: extracted shared device functions +├── dqn_experience_kernel.cu # MODIFIED: uses common header via concatenation +├── ppo_experience_kernel.cu # NEW: PPO kernel entry point + actor/critic/GAE +├── gpu_weights.rs # MODIFIED: add PpoActorWeightSet, PpoCriticWeightSet +├── gpu_experience_collector.rs # EXISTING: DQN collector (unchanged) +├── gpu_ppo_collector.rs # NEW: PPO collector wrapper +├── gpu_portfolio.rs # EXISTING (unchanged) +└── mod.rs # MODIFIED: add new module declarations + +crates/ml/src/trainers/ppo.rs # MODIFIED: GPU collector integration +crates/ml/src/dqn/curiosity.rs # ALREADY DONE: forward_model_vars() exists from Phase 2b +``` + +## Testing Strategy + +- CPU-only unit tests: actor/critic VarMap key paths, param counts, kernel source verification, config defaults +- DQN regression: extract shared header must not break existing DQN kernel compilation or 14 cuda_pipeline tests +- PPO trainer regression: 125 existing PPO tests pass with CPU fallback +- Full ML suite: 2396+ tests pass