docs: add PPO CUDA pipeline Phase 2c design — zero-roundtrip experience collection

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 <noreply@anthropic.com>
2026-02-28 12:59:45 +01:00
parent e7a6814045
commit a38dd64a71
1 changed files with 201 additions and 0 deletions
--- a/docs/plans/2026-02-28-ppo-cuda-pipeline-design.md
+++ 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<GpuPpoExperienceCollector>` 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