docs: H100 GPU optimization implementation plan — 10 tasks
Phase 1: CUDA events, dynamic batch size, async q-stats, phase overlap Phase 2: Multi-stream branches, target/online parallelism Phase 3: CUDA Graph experience collection, lock removal Validation: H100 benchmark with epoch time + SM utilization targets Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
621
docs/superpowers/plans/2026-04-01-h100-gpu-optimization.md
Normal file
621
docs/superpowers/plans/2026-04-01-h100-gpu-optimization.md
Normal file
@@ -0,0 +1,621 @@
|
||||
# H100 GPU Pipeline Optimization Plan
|
||||
|
||||
> **For agentic workers:** REQUIRED SUB-SKILL: Use superpowers:subagent-driven-development (recommended) or superpowers:executing-plans to implement this plan task-by-task. Steps use checkbox (`- [ ]`) syntax for tracking.
|
||||
|
||||
**Goal:** Reduce DQN epoch time from 2 min to <10s on H100 by eliminating GPU idle time from blocking synchronize calls, maximizing batch size, and enabling multi-stream parallelism.
|
||||
|
||||
**Architecture:** Replace `stream.synchronize()` with CUDA events for async scalar readback. Let `AutoBatchSizer` drive batch size (removes static cap). Capture experience collection as a CUDA Graph. Fork child streams for independent branch/target forward passes.
|
||||
|
||||
**Tech Stack:** Rust, cudarc 0.17 (CudaEvent, CudaStream::fork), CUDA 12.4, H100 PCIe 80GB.
|
||||
|
||||
---
|
||||
|
||||
## File Structure
|
||||
|
||||
| File | Action | Change |
|
||||
|------|--------|--------|
|
||||
| `crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs` | Modify | CUDA events in `replay_adam_and_readback`, async q-stats, multi-stream forward |
|
||||
| `crates/ml/src/cuda_pipeline/batched_forward.rs` | Modify | Multi-stream branch dispatch |
|
||||
| `crates/ml/src/cuda_pipeline/gpu_experience_collector.rs` | Modify | CUDA Graph capture for timestep loop |
|
||||
| `crates/ml/src/trainers/dqn/trainer/training_loop.rs` | Modify | Async training loop, event-based phase transitions |
|
||||
| `config/gpu/h100.toml` | Modify | Remove static batch_size cap |
|
||||
| `crates/ml/src/trainers/dqn/trainer/constructor.rs` | Modify | Let AutoBatchSizer drive batch_size |
|
||||
| `infra/k8s/argo/precompute-features-template.yaml` | Done | Already committed (data-dir/SYMBOL alignment) |
|
||||
|
||||
---
|
||||
|
||||
## Phase 1: Eliminate GPU Idle Time (biggest impact)
|
||||
|
||||
### Task 1: CUDA Events replace per-step stream.synchronize()
|
||||
|
||||
**Files:**
|
||||
- Modify: `crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs:2865-2893`
|
||||
|
||||
The `replay_adam_and_readback()` method calls `stream.synchronize()` on every training step (341×/epoch). This is the primary bottleneck — H100 finishes both CUDA graphs in microseconds then sits idle.
|
||||
|
||||
- [ ] **Step 1: Add event and async readback state to GpuDqnTrainer struct**
|
||||
|
||||
In `gpu_dqn_trainer.rs`, add fields to the `GpuDqnTrainer` struct (near line 479):
|
||||
|
||||
```rust
|
||||
/// CUDA event for async scalar readback (replaces stream.synchronize per step)
|
||||
readback_event: Option<Arc<CudaEvent>>,
|
||||
/// Host-side buffers for deferred readback (loss, grad_norm from previous step)
|
||||
pending_loss: f32,
|
||||
pending_grad_norm: f32,
|
||||
/// Whether there's an in-flight readback to collect
|
||||
readback_pending: bool,
|
||||
```
|
||||
|
||||
Initialize in `new()`:
|
||||
```rust
|
||||
readback_event: None,
|
||||
pending_loss: 0.0,
|
||||
pending_grad_norm: 0.0,
|
||||
readback_pending: false,
|
||||
```
|
||||
|
||||
- [ ] **Step 2: Rewrite replay_adam_and_readback to use events**
|
||||
|
||||
Replace the synchronous `replay_adam_and_readback` method:
|
||||
|
||||
```rust
|
||||
/// Collect scalars from the PREVIOUS step (non-blocking poll),
|
||||
/// then launch grad_norm + adam for the CURRENT step and record an event.
|
||||
pub fn replay_adam_and_readback(&mut self) -> Result<FusedTrainScalars, MLError> {
|
||||
// 1. Collect previous step's scalars (if any)
|
||||
let prev_scalars = if self.readback_pending {
|
||||
// Poll: is the previous step's event done?
|
||||
if let Some(ref event) = self.readback_event {
|
||||
if !event.is_complete() {
|
||||
// Still in-flight — wait on just this event (not full stream)
|
||||
event.synchronize().map_err(|e|
|
||||
MLError::ModelError(format!("readback event sync: {e}"))
|
||||
)?;
|
||||
}
|
||||
}
|
||||
// Now safe to read the scalars from the previous step
|
||||
self.readback_pending = false;
|
||||
FusedTrainScalars {
|
||||
total_loss: self.pending_loss,
|
||||
grad_norm: self.pending_grad_norm,
|
||||
}
|
||||
} else {
|
||||
FusedTrainScalars { total_loss: 0.0, grad_norm: 0.0 }
|
||||
};
|
||||
|
||||
// 2. Launch current step: grad_norm + adam
|
||||
self.compute_grad_norm_outside_graph()?;
|
||||
self.replay_adam()?;
|
||||
|
||||
// 3. Issue async DtoH for current step's scalars
|
||||
unsafe {
|
||||
cudarc::driver::sys::cuMemcpyDtoHAsync_v2(
|
||||
(&mut self.pending_loss as *mut f32).cast(),
|
||||
self.total_loss_buf.raw_ptr(),
|
||||
std::mem::size_of::<f32>(),
|
||||
*self.stream.cu_stream(),
|
||||
);
|
||||
cudarc::driver::sys::cuMemcpyDtoHAsync_v2(
|
||||
(&mut self.pending_grad_norm as *mut f32).cast(),
|
||||
self.grad_norm_f32_buf.raw_ptr(),
|
||||
std::mem::size_of::<f32>(),
|
||||
*self.stream.cu_stream(),
|
||||
);
|
||||
}
|
||||
|
||||
// 4. Record event marking when the async DtoH completes
|
||||
self.readback_event = Some(self.stream.record_event(None).map_err(|e|
|
||||
MLError::ModelError(format!("readback event record: {e}"))
|
||||
)?);
|
||||
self.readback_pending = true;
|
||||
|
||||
// Return PREVIOUS step's scalars (current step's will be collected next call)
|
||||
self.scalars_readback_host = [prev_scalars.total_loss, prev_scalars.grad_norm];
|
||||
Ok(prev_scalars)
|
||||
}
|
||||
|
||||
/// Flush: block until the last in-flight readback completes.
|
||||
/// Called once at epoch end.
|
||||
pub fn flush_readback(&mut self) -> Result<FusedTrainScalars, MLError> {
|
||||
if !self.readback_pending {
|
||||
return Ok(FusedTrainScalars { total_loss: 0.0, grad_norm: 0.0 });
|
||||
}
|
||||
if let Some(ref event) = self.readback_event {
|
||||
if !event.is_complete() {
|
||||
event.synchronize().map_err(|e|
|
||||
MLError::ModelError(format!("flush readback sync: {e}"))
|
||||
)?;
|
||||
}
|
||||
}
|
||||
self.readback_pending = false;
|
||||
Ok(FusedTrainScalars {
|
||||
total_loss: self.pending_loss,
|
||||
grad_norm: self.pending_grad_norm,
|
||||
})
|
||||
}
|
||||
```
|
||||
|
||||
- [ ] **Step 3: Update training loop to flush at epoch end**
|
||||
|
||||
In `training_loop.rs`, after the training step loop (after the `for chunk_start in ...` loop ends), add:
|
||||
|
||||
```rust
|
||||
// Flush the last in-flight readback
|
||||
if let Some(ref mut fused) = self.fused_ctx {
|
||||
let final_scalars = fused.flush_readback()?;
|
||||
// Use final_scalars for epoch-end logging if needed
|
||||
}
|
||||
```
|
||||
|
||||
- [ ] **Step 4: Verify compile**
|
||||
|
||||
Run: `SQLX_OFFLINE=true cargo check -p ml --lib`
|
||||
Expected: compiles clean
|
||||
|
||||
- [ ] **Step 5: Run local smoke test to validate loss trajectory**
|
||||
|
||||
Run: `FOXHUNT_TEST_DATA=test_data/futures-baseline FOXHUNT_FEATURE_CACHE_DIR=test_data/feature-cache SQLX_OFFLINE=true cargo test -p ml --lib -- smoke_tests::training_stability::test_50_epoch_convergence --ignored --nocapture`
|
||||
|
||||
Expected: test passes, loss values are finite, Sharpe is reasonable. Loss trajectory may differ slightly from synchronous path (1-step lag on PER priorities) but should converge similarly.
|
||||
|
||||
- [ ] **Step 6: Commit**
|
||||
|
||||
```bash
|
||||
git commit -m "perf: CUDA events replace per-step stream.synchronize in training loop"
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
### Task 2: Async Q-stats readback
|
||||
|
||||
**Files:**
|
||||
- Modify: `crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs:3640-3660`
|
||||
|
||||
The `reduce_current_q_stats()` calls `memcpy_dtoh` every 50 steps — implicit sync. Use the same event pattern.
|
||||
|
||||
- [ ] **Step 1: Add q-stats event and pending buffer**
|
||||
|
||||
Add to `GpuDqnTrainer` struct:
|
||||
|
||||
```rust
|
||||
/// Pending q-stats from previous reduction (deferred readback)
|
||||
q_stats_event: Option<Arc<CudaEvent>>,
|
||||
q_stats_pending: [f32; 5],
|
||||
q_stats_ready: bool,
|
||||
```
|
||||
|
||||
Initialize:
|
||||
```rust
|
||||
q_stats_event: None,
|
||||
q_stats_pending: [0.0; 5],
|
||||
q_stats_ready: false,
|
||||
```
|
||||
|
||||
- [ ] **Step 2: Rewrite reduce_current_q_stats for async readback**
|
||||
|
||||
Replace the synchronous `memcpy_dtoh` in `reduce_current_q_stats()`:
|
||||
|
||||
```rust
|
||||
pub fn reduce_current_q_stats(&mut self) -> Result<QValueStatsResult, MLError> {
|
||||
// Collect previous reduction if ready
|
||||
let result = if self.q_stats_ready {
|
||||
if let Some(ref event) = self.q_stats_event {
|
||||
if !event.is_complete() {
|
||||
event.synchronize().map_err(|e|
|
||||
MLError::ModelError(format!("q_stats event sync: {e}"))
|
||||
)?;
|
||||
}
|
||||
}
|
||||
self.q_stats_ready = false;
|
||||
let h = self.q_stats_pending;
|
||||
QValueStatsResult {
|
||||
avg_max_q: h[0] as f64, q_min: h[1], q_max: h[2], q_mean: h[3], q_var: h[4],
|
||||
}
|
||||
} else {
|
||||
QValueStatsResult { avg_max_q: 0.0, q_min: 0.0, q_max: 0.0, q_mean: 0.0, q_var: 0.0 }
|
||||
};
|
||||
|
||||
// Launch new reduction kernel
|
||||
// ... (existing kernel launch code stays the same) ...
|
||||
|
||||
// Async readback
|
||||
unsafe {
|
||||
cudarc::driver::sys::cuMemcpyDtoHAsync_v2(
|
||||
self.q_stats_pending.as_mut_ptr().cast(),
|
||||
self.q_stats_buf.raw_ptr(),
|
||||
5 * std::mem::size_of::<f32>(),
|
||||
*self.stream.cu_stream(),
|
||||
);
|
||||
}
|
||||
self.q_stats_event = Some(self.stream.record_event(None).map_err(|e|
|
||||
MLError::ModelError(format!("q_stats event record: {e}"))
|
||||
)?);
|
||||
self.q_stats_ready = true;
|
||||
|
||||
Ok(result)
|
||||
}
|
||||
```
|
||||
|
||||
- [ ] **Step 3: Verify compile + smoke test**
|
||||
|
||||
Run: `SQLX_OFFLINE=true cargo check -p ml --lib`
|
||||
Run: smoke test (same as Task 1 Step 5)
|
||||
|
||||
- [ ] **Step 4: Commit**
|
||||
|
||||
```bash
|
||||
git commit -m "perf: async q-stats readback via CUDA events"
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
### Task 3: Dynamic batch sizing — remove static cap
|
||||
|
||||
**Files:**
|
||||
- Modify: `config/gpu/h100.toml:3`
|
||||
- Modify: `crates/ml/src/trainers/dqn/trainer/constructor.rs` (AutoBatchSizer logic)
|
||||
|
||||
- [ ] **Step 1: Remove static batch_size from h100.toml**
|
||||
|
||||
Change `config/gpu/h100.toml` line 3:
|
||||
|
||||
```toml
|
||||
# batch_size is auto-computed by AutoBatchSizer from available VRAM.
|
||||
# Do NOT set a static value here — let the sizer maximize GPU utilization.
|
||||
# batch_size = 1024 # REMOVED: was capping H100 far below optimal
|
||||
```
|
||||
|
||||
Actually: the TOML profile must still have the field for the parser. Set it to 0 as a sentinel meaning "auto":
|
||||
|
||||
```toml
|
||||
batch_size = 0 # 0 = auto-compute from VRAM (AutoBatchSizer)
|
||||
```
|
||||
|
||||
- [ ] **Step 2: Update constructor to honor batch_size=0 as auto**
|
||||
|
||||
In `constructor.rs`, find where `hyperparams.batch_size` is set from the profile. After the profile applies, if `batch_size == 0`, let AutoBatchSizer compute it:
|
||||
|
||||
```rust
|
||||
// After profile.apply_to(&mut hyperparams):
|
||||
if hyperparams.batch_size == 0 {
|
||||
// Auto-compute: AutoBatchSizer already ran and set the VRAM ceiling.
|
||||
// Use the ceiling as the batch size, capped at a sane max (16384).
|
||||
hyperparams.batch_size = auto_batch_ceiling.min(16384);
|
||||
info!("AutoBatchSizer: batch_size auto-computed to {}", hyperparams.batch_size);
|
||||
}
|
||||
```
|
||||
|
||||
- [ ] **Step 3: Verify the auto-computed value is reasonable**
|
||||
|
||||
Run locally (RTX 3050): the sizer should still pick a small batch (64-128) for 4GB VRAM.
|
||||
Check H100 would get: look at the `AutoBatchSizer: GPU VRAM ceiling = 2085808` log — that's the theoretical max. The actual batch should be much smaller (constrained by kernel shared memory and practical gradient noise). Cap at 8192 or 16384.
|
||||
|
||||
- [ ] **Step 4: Commit**
|
||||
|
||||
```bash
|
||||
git commit -m "perf: dynamic batch sizing — AutoBatchSizer drives batch_size on H100"
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
### Task 4: Experience collection → training phase overlap
|
||||
|
||||
**Files:**
|
||||
- Modify: `crates/ml/src/trainers/dqn/trainer/training_loop.rs:1196`
|
||||
|
||||
- [ ] **Step 1: Replace stream.synchronize() with event after experience collection**
|
||||
|
||||
At line 1196, change:
|
||||
|
||||
```rust
|
||||
// BEFORE:
|
||||
stream.synchronize()
|
||||
.map_err(|e| MLError::DeviceError(format!("experience sync: {e}")))?;
|
||||
|
||||
// AFTER:
|
||||
let experience_done_event = stream.record_event(None)
|
||||
.map_err(|e| MLError::DeviceError(format!("experience event record: {e}")))?;
|
||||
```
|
||||
|
||||
Then before the first training step that samples from the replay buffer, wait on the event:
|
||||
|
||||
```rust
|
||||
// Before first replay buffer sample:
|
||||
if !experience_done_event.is_complete() {
|
||||
experience_done_event.synchronize()
|
||||
.map_err(|e| MLError::DeviceError(format!("experience event wait: {e}")))?;
|
||||
}
|
||||
```
|
||||
|
||||
This allows CPU-side training setup (lock acquisition, pre-sampling initialization) to overlap with the tail end of experience collection.
|
||||
|
||||
- [ ] **Step 2: Verify + commit**
|
||||
|
||||
```bash
|
||||
git commit -m "perf: event-based experience→training phase transition"
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
### Task 5: fxcache PVC regeneration
|
||||
|
||||
- [ ] **Step 1: Regenerate PVC cache with aligned data-dir**
|
||||
|
||||
The precompute template already uses `data-dir/SYMBOL`. Just re-run:
|
||||
|
||||
```bash
|
||||
./scripts/argo-precompute.sh ES.FUT --pool ci-compile-cpu
|
||||
```
|
||||
|
||||
- [ ] **Step 2: Verify training hits fxcache**
|
||||
|
||||
Submit a quick 2-epoch baseline and check logs for "fxcache hit":
|
||||
|
||||
```bash
|
||||
argo submit -n foxhunt --from=wftmpl/train-baseline-rl -p train-epochs=2
|
||||
```
|
||||
|
||||
Check: `argo logs <name> | grep fxcache`
|
||||
|
||||
- [ ] **Step 3: Commit (if any template changes needed)**
|
||||
|
||||
---
|
||||
|
||||
## Phase 2: Multi-Stream Parallelism
|
||||
|
||||
### Task 6: Multi-stream branch dispatch in forward pass
|
||||
|
||||
**Files:**
|
||||
- Modify: `crates/ml/src/cuda_pipeline/batched_forward.rs:265-284`
|
||||
- Modify: `crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs` (add child streams)
|
||||
|
||||
- [ ] **Step 1: Add 3 child streams to CublasForward**
|
||||
|
||||
In `batched_forward.rs`, add to the struct:
|
||||
|
||||
```rust
|
||||
/// Child streams for parallel branch dispatch (forked from main training stream)
|
||||
branch_streams: [Arc<CudaStream>; 3],
|
||||
```
|
||||
|
||||
Create them in the constructor by forking:
|
||||
|
||||
```rust
|
||||
let branch_streams = [
|
||||
stream.fork().map_err(|e| MLError::DeviceError(format!("fork branch stream 0: {e}")))?,
|
||||
stream.fork().map_err(|e| MLError::DeviceError(format!("fork branch stream 1: {e}")))?,
|
||||
stream.fork().map_err(|e| MLError::DeviceError(format!("fork branch stream 2: {e}")))?,
|
||||
];
|
||||
```
|
||||
|
||||
- [ ] **Step 2: Dispatch branches on separate streams**
|
||||
|
||||
In `forward_online_raw()`, replace the sequential `for d in 0..3` loop (lines 270-284):
|
||||
|
||||
```rust
|
||||
// Fork: dispatch each branch on its own stream
|
||||
for d in 0..3 {
|
||||
let branch_stream = &self.branch_streams[d];
|
||||
// Wait for trunk (h_s2) to be ready on this branch stream
|
||||
let trunk_event = self.stream.record_event(None)
|
||||
.map_err(|e| MLError::ModelError(format!("trunk event: {e}")))?;
|
||||
branch_stream.wait(&trunk_event)
|
||||
.map_err(|e| MLError::ModelError(format!("branch wait trunk: {e}")))?;
|
||||
|
||||
let n_d = branch_sizes[d];
|
||||
let w_fc_idx = branch_w_base[d];
|
||||
|
||||
// Hidden layer
|
||||
self.gemmex_bf16_on_stream(branch_stream, w_ptrs[w_fc_idx], h_s2_ptr, branch_h_ptrs[d], self.adv_h, b, self.shared_h2, "h_bd")?;
|
||||
self.launch_add_bias_relu_bf16_on_stream(branch_stream, branch_h_ptrs[d], w_ptrs[w_fc_idx + 1], self.adv_h, b)?;
|
||||
|
||||
// Output layer
|
||||
let adv_out_ptr = b_logits_ptr + logit_byte_offset;
|
||||
self.gemmex_bf16_to_f32_on_stream(branch_stream, w_ptrs[w_fc_idx + 2], branch_h_ptrs[d], adv_out_ptr, n_d * na, b, self.adv_h, "adv_logits")?;
|
||||
self.launch_add_bias_f32_on_stream(branch_stream, adv_out_ptr, w_ptrs[w_fc_idx + 3], n_d * na, b)?;
|
||||
|
||||
logit_byte_offset += (b * n_d * na * std::mem::size_of::<f32>()) as u64;
|
||||
}
|
||||
|
||||
// Join: main stream waits for all branches to finish
|
||||
for d in 0..3 {
|
||||
let branch_done = self.branch_streams[d].record_event(None)
|
||||
.map_err(|e| MLError::ModelError(format!("branch done event: {e}")))?;
|
||||
self.stream.wait(&branch_done)
|
||||
.map_err(|e| MLError::ModelError(format!("join branch {d}: {e}")))?;
|
||||
}
|
||||
```
|
||||
|
||||
Note: `gemmex_bf16_on_stream` and `launch_add_bias_relu_bf16_on_stream` are new variants that take an explicit stream parameter. Implement them by extracting the cuBLAS handle → set stream → call GemmEx → restore stream. cuBLAS handles are NOT stream-specific — `cublasSetStream()` switches the active stream per-call.
|
||||
|
||||
- [ ] **Step 3: Apply same pattern to f32 forward (line 325) and backward (line 461)**
|
||||
|
||||
Repeat the fork/join pattern for both the f32 experience collection forward and the backward pass branch loops.
|
||||
|
||||
- [ ] **Step 4: Verify compile + smoke test**
|
||||
|
||||
Run smoke tests — output logits must be identical (branches write to non-overlapping memory regions).
|
||||
|
||||
- [ ] **Step 5: Commit**
|
||||
|
||||
```bash
|
||||
git commit -m "perf: multi-stream branch dispatch — 3 advantage heads in parallel"
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
### Task 7: Multi-stream target/online forward parallelism
|
||||
|
||||
**Files:**
|
||||
- Modify: `crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs:4307-4323`
|
||||
|
||||
- [ ] **Step 1: Add a second stream for Double DQN pass**
|
||||
|
||||
Add to `GpuDqnTrainer`:
|
||||
|
||||
```rust
|
||||
/// Second stream for parallel target/online next-state forward passes
|
||||
double_dqn_stream: Arc<CudaStream>,
|
||||
```
|
||||
|
||||
Fork in constructor.
|
||||
|
||||
- [ ] **Step 2: Dispatch Pass 2 and Pass 3 on separate streams**
|
||||
|
||||
In `submit_forward_ops()`, after Pass 1 (online forward on current states):
|
||||
|
||||
```rust
|
||||
// Pass 2 + Pass 3 are independent: both read next_states, write to different outputs.
|
||||
// Fork: record event after stochastic depth (Pass 1 done)
|
||||
let pass1_done = self.stream.record_event(None)?;
|
||||
self.double_dqn_stream.wait(&pass1_done)?;
|
||||
|
||||
// Pass 2: target forward on next_states (main stream)
|
||||
cublas.forward_target_raw(
|
||||
&self.stream, self.ptrs.next_states_buf, &tg_w_ptrs, ...
|
||||
)?;
|
||||
|
||||
// Pass 3: online forward on next_states (double_dqn_stream)
|
||||
cublas.forward_online_raw(
|
||||
&self.double_dqn_stream, self.ptrs.next_states_buf, &on_w_ptrs, ...
|
||||
)?;
|
||||
|
||||
// Join before loss kernel
|
||||
let target_done = self.stream.record_event(None)?;
|
||||
let online_done = self.double_dqn_stream.record_event(None)?;
|
||||
self.stream.wait(&online_done)?;
|
||||
// main stream now has both pass 2 + pass 3 complete
|
||||
```
|
||||
|
||||
- [ ] **Step 3: Handle CUDA Graph capture with multi-stream**
|
||||
|
||||
CUDA Graph capture with fork/join requires `CU_STREAM_CAPTURE_MODE_THREAD_LOCAL` (already used at line 3734). The fork/join events within the graph are captured as graph dependencies. Verify this works with cudarc — may need raw `cuStreamWaitEvent` calls during capture.
|
||||
|
||||
- [ ] **Step 4: Verify + smoke test + commit**
|
||||
|
||||
```bash
|
||||
git commit -m "perf: parallel target/online forward on separate CUDA streams"
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
### Task 8: CUDA Graph for experience collection
|
||||
|
||||
**Files:**
|
||||
- Modify: `crates/ml/src/cuda_pipeline/gpu_experience_collector.rs:1506`
|
||||
|
||||
This is the most complex task. The 100-iteration host loop dispatches 400+ kernels. Capturing as a graph eliminates all host→GPU latency.
|
||||
|
||||
- [ ] **Step 1: Make per-timestep parameters GPU-resident**
|
||||
|
||||
The loop currently passes `epsilon`, step counter, and state pointers from the host. These must become GPU buffers:
|
||||
|
||||
```rust
|
||||
/// GPU-resident parameters for graph-captured timestep loop
|
||||
epsilon_buf: CudaSlice<f32>, // [1] — updated via DtoH before graph launch
|
||||
step_counter_buf: CudaSlice<u32>, // [1] — incremented by a tiny CUDA kernel per step
|
||||
```
|
||||
|
||||
- [ ] **Step 2: Create a step-increment kernel**
|
||||
|
||||
Tiny kernel that increments `step_counter_buf[0]` and computes `epsilon = eps_end + (eps_start - eps_end) * exp(-step / decay)` on GPU. No host involvement.
|
||||
|
||||
- [ ] **Step 3: Capture the 100-timestep loop as a single graph**
|
||||
|
||||
```rust
|
||||
stream.begin_capture(CU_STREAM_CAPTURE_MODE_THREAD_LOCAL)?;
|
||||
for t in 0..timesteps {
|
||||
self.launch_experience_state_gather()?;
|
||||
self.launch_cublas_forward_online()?;
|
||||
self.launch_compute_expected_q()?;
|
||||
self.launch_experience_action_select()?; // reads epsilon_buf
|
||||
self.launch_experience_env_step()?;
|
||||
self.launch_step_increment()?; // increments step_counter, updates epsilon
|
||||
}
|
||||
let graph = stream.end_capture(CUDA_GRAPH_INSTANTIATE_FLAG_AUTO_FREE_ON_LAUNCH)?;
|
||||
```
|
||||
|
||||
Before each epoch: `memcpy_htod` the initial epsilon and step counter, then `graph.launch()`.
|
||||
|
||||
- [ ] **Step 4: Handle dynamic epsilon**
|
||||
|
||||
Epsilon decays across episodes within an epoch. The step-increment kernel computes the new epsilon from the step counter using the exponential decay formula, writing to `epsilon_buf`. The action selection kernel reads `epsilon_buf` directly — no host involvement.
|
||||
|
||||
- [ ] **Step 5: Verify collected experiences match sequential path**
|
||||
|
||||
Run 1 epoch with both paths (graph and sequential), compare states/actions/rewards for the first 100 experiences. Must be numerically identical given the same Philox RNG seed.
|
||||
|
||||
- [ ] **Step 6: Commit**
|
||||
|
||||
```bash
|
||||
git commit -m "perf: CUDA Graph capture for 100-timestep experience collection loop"
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
### Task 9: Lock contention removal
|
||||
|
||||
**Files:**
|
||||
- Modify: `crates/ml/src/trainers/dqn/trainer/training_loop.rs:1336-1360`
|
||||
|
||||
- [ ] **Step 1: Pre-sample all batches for the epoch in one lock acquisition**
|
||||
|
||||
Instead of acquiring/releasing the read lock every 32 steps, sample ALL batches at epoch start:
|
||||
|
||||
```rust
|
||||
let all_batches = {
|
||||
let agent = self.agent.read().await;
|
||||
let buffer = agent.memory();
|
||||
(0..num_training_steps)
|
||||
.map(|_| buffer.sample(self.current_batch_size))
|
||||
.collect::<Vec<_>>()
|
||||
};
|
||||
// Lock released. Training loop iterates over pre-sampled batches.
|
||||
```
|
||||
|
||||
This is safe because the GPU PER buffer is immutable during training (new experiences are only added in Phase 2, which is finished before training starts).
|
||||
|
||||
- [ ] **Step 2: Remove PREFETCH_K chunking**
|
||||
|
||||
The `for chunk_start in (0..num_training_steps).step_by(PREFETCH_K)` loop becomes a simple `for batch in all_batches.iter()`.
|
||||
|
||||
- [ ] **Step 3: Verify + commit**
|
||||
|
||||
```bash
|
||||
git commit -m "perf: pre-sample all training batches in single lock acquisition"
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## Validation
|
||||
|
||||
### Task 10: H100 benchmark
|
||||
|
||||
- [ ] **Step 1: Rebuild and upload binary to PVC**
|
||||
|
||||
Compile in-cluster via the baseline template.
|
||||
|
||||
- [ ] **Step 2: Run 10-epoch baseline and measure**
|
||||
|
||||
```bash
|
||||
argo submit -n foxhunt --from=wftmpl/train-baseline-rl -p train-epochs=10
|
||||
```
|
||||
|
||||
Check:
|
||||
- Epoch time: target <10s (from 2 min)
|
||||
- `nvidia-smi`: SM utilization >85%, memory bandwidth >50%, TDP >250W
|
||||
- Loss trajectory: finite, converging
|
||||
- "fxcache hit" in logs
|
||||
|
||||
- [ ] **Step 3: Compare loss with pre-optimization baseline**
|
||||
|
||||
The 7-epoch metrics from the terminated run serve as reference:
|
||||
- Epoch 4: Sharpe 0.06, PF 0.94
|
||||
- Epoch 6: Sharpe -1.12, PF 1.17
|
||||
|
||||
Post-optimization should show similar or better convergence (larger batches may improve gradient quality).
|
||||
|
||||
- [ ] **Step 4: Run full 50-epoch baseline if metrics look good**
|
||||
|
||||
- [ ] **Step 5: Save metrics to memory**
|
||||
Reference in New Issue
Block a user