plan: Phase 2 multi-stream parallelism — 7 tasks
Task 1: Allocate aux streams + workspaces + events Task 2: IQN/Attention accept workspace+stream override Task 3: Fork-join in submit_aux_ops (IQN+Attention parallel) Task 4: Switch capture mode GLOBAL → RELAXED Task 5: Forward backward branch parallelism Task 6: Fix adam child event timing (verify) Task 7: Deploy + verify child graph breakdown Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
349
docs/superpowers/plans/2026-04-18-multi-stream-phase2.md
Normal file
349
docs/superpowers/plans/2026-04-18-multi-stream-phase2.md
Normal file
@@ -0,0 +1,349 @@
|
||||
# Phase 2: Multi-Stream Graph Parallelism Implementation 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 per-step GPU time from 4.0s to ~2.0s by parallelizing aux_child (295ms, 69%) and forward_child backward branches (part of 130ms).
|
||||
|
||||
**Architecture:** Aux trainers (IQN, Attention) run on 2 dedicated streams forked from main. IQL high+low stay sequential (grouped GEMM future optimization). Forward backward branches use existing branch_streams infrastructure. Capture mode switches to RELAXED for multi-stream graph capture.
|
||||
|
||||
**Tech Stack:** CUDA 13, cudarc (vendored), cublasLtMatmul, CU_STREAM_CAPTURE_MODE_RELAXED, cuEventRecord fork-join
|
||||
|
||||
---
|
||||
|
||||
### Task 1: Aux Stream Infrastructure — Allocate 2 streams + workspaces + events
|
||||
|
||||
**Files:**
|
||||
- Modify: `crates/ml/src/trainers/dqn/fused_training.rs` (FusedTrainingCtx struct + constructor)
|
||||
- Modify: `crates/ml/src/cuda_pipeline/shared_cublas_handle.rs` (workspace allocation helper)
|
||||
|
||||
- [ ] **Step 1: Add aux stream fields to FusedTrainingCtx**
|
||||
|
||||
In `fused_training.rs`, add to the struct after line ~272 (`gpu_attention`):
|
||||
|
||||
```rust
|
||||
// ── Phase 2: Aux multi-stream ──
|
||||
/// Dedicated stream for IQN training (parallel with Attention).
|
||||
iqn_stream: Option<Arc<CudaStream>>,
|
||||
/// Dedicated stream for Attention training (parallel with IQN).
|
||||
attn_stream: Option<Arc<CudaStream>>,
|
||||
/// cublasLt workspace for IQN stream (32MB, separate from shared handle).
|
||||
iqn_workspace: Option<CudaSlice<u8>>,
|
||||
/// cublasLt workspace for Attention stream (32MB, separate from shared handle).
|
||||
attn_workspace: Option<CudaSlice<u8>>,
|
||||
/// Fork event: IQL done on main stream, IQN/Attention can start.
|
||||
iql_done_event: Option<CudaEvent>,
|
||||
/// Join event: IQN done on iqn_stream.
|
||||
iqn_done_event: Option<CudaEvent>,
|
||||
/// Join event: Attention done on attn_stream.
|
||||
attn_done_event: Option<CudaEvent>,
|
||||
```
|
||||
|
||||
- [ ] **Step 2: Allocate in constructor**
|
||||
|
||||
After IQL/IQN/Attention initialization, allocate the streams and workspaces:
|
||||
|
||||
```rust
|
||||
let iqn_stream = Arc::new(stream.context().new_stream()
|
||||
.map_err(|e| anyhow::anyhow!("iqn_stream: {e}"))?);
|
||||
let attn_stream = Arc::new(stream.context().new_stream()
|
||||
.map_err(|e| anyhow::anyhow!("attn_stream: {e}"))?);
|
||||
let iqn_workspace = stream.alloc_zeros::<u8>(32 * 1024 * 1024)
|
||||
.map_err(|e| anyhow::anyhow!("iqn workspace: {e}"))?;
|
||||
let attn_workspace = stream.alloc_zeros::<u8>(32 * 1024 * 1024)
|
||||
.map_err(|e| anyhow::anyhow!("attn workspace: {e}"))?;
|
||||
let iql_done_event = stream.context().new_event(None)
|
||||
.map_err(|e| anyhow::anyhow!("iql_done_event: {e}"))?;
|
||||
let iqn_done_event = stream.context().new_event(None)
|
||||
.map_err(|e| anyhow::anyhow!("iqn_done_event: {e}"))?;
|
||||
let attn_done_event = stream.context().new_event(None)
|
||||
.map_err(|e| anyhow::anyhow!("attn_done_event: {e}"))?;
|
||||
```
|
||||
|
||||
- [ ] **Step 3: Initialize fields in struct initializer and add to `reset_for_fold`**
|
||||
|
||||
- [ ] **Step 4: Build and test**
|
||||
|
||||
Run: `SQLX_OFFLINE=true cargo check -p ml && SQLX_OFFLINE=true cargo test -p ml --lib`
|
||||
|
||||
- [ ] **Step 5: Commit**
|
||||
|
||||
```bash
|
||||
git commit -m "feat: allocate aux streams, workspaces, events for Phase 2 multi-stream"
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
### Task 2: IQN/Attention Accept Per-Call Workspace + Stream Override
|
||||
|
||||
**Files:**
|
||||
- Modify: `crates/ml/src/cuda_pipeline/gpu_iqn_head.rs`
|
||||
- Modify: `crates/ml/src/cuda_pipeline/gpu_attention.rs`
|
||||
|
||||
Currently IQN and Attention use the shared handle's workspace and stream for all cublasLtMatmul calls. They need to accept an optional override workspace pointer and stream reference so they can run on dedicated streams with dedicated workspaces.
|
||||
|
||||
- [ ] **Step 1: Add workspace/stream override to IQN's execute_training_pipeline**
|
||||
|
||||
Add optional parameters to `execute_training_pipeline`:
|
||||
|
||||
```rust
|
||||
pub fn execute_training_pipeline(
|
||||
&mut self,
|
||||
// ... existing params ...
|
||||
override_stream: Option<&Arc<CudaStream>>,
|
||||
override_workspace: Option<(u64, usize)>, // (ptr, size)
|
||||
) -> Result<(), MLError> {
|
||||
let stream = override_stream.unwrap_or(&self.stream);
|
||||
let (ws_ptr, ws_size) = override_workspace
|
||||
.unwrap_or((self.shared_handle.lt_workspace_ptr, self.shared_handle.lt_workspace_size));
|
||||
// Use `stream` and `ws_ptr/ws_size` for all cublasLtMatmul calls in this function
|
||||
```
|
||||
|
||||
- [ ] **Step 2: Same for Attention forward/backward/adam_step**
|
||||
|
||||
Add `override_stream` and `override_workspace` parameters to `forward()`, `backward()`, `adam_step()`.
|
||||
|
||||
- [ ] **Step 3: Update all call sites to pass `None, None` (no behavior change)**
|
||||
|
||||
All existing callers in `submit_aux_ops` and elsewhere pass `None, None`.
|
||||
|
||||
- [ ] **Step 4: Build and test — behavior unchanged**
|
||||
|
||||
- [ ] **Step 5: Commit**
|
||||
|
||||
```bash
|
||||
git commit -m "feat: IQN/Attention accept workspace+stream override for multi-stream"
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
### Task 3: Fork-Join in submit_aux_ops
|
||||
|
||||
**Files:**
|
||||
- Modify: `crates/ml/src/trainers/dqn/fused_training.rs` (submit_aux_ops)
|
||||
|
||||
- [ ] **Step 1: Fork IQN and Attention to dedicated streams after IQL**
|
||||
|
||||
In `submit_aux_ops`, after the IQL pipeline (gather → train → advantages → support → branch_scales → gap → epsilon), fork IQN and Attention to their dedicated streams:
|
||||
|
||||
```rust
|
||||
// IQL pipeline runs on main stream (sequential) ...
|
||||
// ... existing IQL code ...
|
||||
|
||||
// Fork: IQN and Attention run in parallel on dedicated streams.
|
||||
// Record IQL completion on main stream.
|
||||
if let Some(ref event) = self.iql_done_event {
|
||||
event.record(&self.stream)
|
||||
.map_err(|e| anyhow::anyhow!("iql_done record: {e}"))?;
|
||||
}
|
||||
|
||||
// IQN on iqn_stream
|
||||
if let (Some(ref mut iqn), Some(ref iqn_stream), Some(ref iqn_ws), Some(ref iql_event)) =
|
||||
(&mut self.gpu_iqn, &self.iqn_stream, &self.iqn_workspace, &self.iql_done_event)
|
||||
{
|
||||
iqn_stream.wait(iql_event)
|
||||
.map_err(|e| anyhow::anyhow!("iqn wait iql: {e}"))?;
|
||||
// Execute IQN training pipeline on iqn_stream with iqn workspace
|
||||
iqn.execute_training_pipeline(
|
||||
/* ... existing args ... */
|
||||
Some(iqn_stream),
|
||||
Some((iqn_ws.raw_ptr(), 32 * 1024 * 1024)),
|
||||
).map_err(|e| anyhow::anyhow!("IQN parallel: {e}"))?;
|
||||
// Record IQN completion
|
||||
if let Some(ref event) = self.iqn_done_event {
|
||||
event.record(iqn_stream)
|
||||
.map_err(|e| anyhow::anyhow!("iqn_done record: {e}"))?;
|
||||
}
|
||||
}
|
||||
|
||||
// Attention on attn_stream (same pattern)
|
||||
if let (Some(ref mut attn), Some(ref attn_stream), Some(ref attn_ws), Some(ref iql_event)) =
|
||||
(&mut self.gpu_attention, &self.attn_stream, &self.attn_workspace, &self.iql_done_event)
|
||||
{
|
||||
attn_stream.wait(iql_event)
|
||||
.map_err(|e| anyhow::anyhow!("attn wait iql: {e}"))?;
|
||||
self.trainer.apply_attention_forward_on_stream(attn, self.batch_size, attn_stream, attn_ws)
|
||||
.map_err(|e| anyhow::anyhow!("Attention parallel fwd: {e}"))?;
|
||||
// ... backward + adam on attn_stream ...
|
||||
if let Some(ref event) = self.attn_done_event {
|
||||
event.record(attn_stream)
|
||||
.map_err(|e| anyhow::anyhow!("attn_done record: {e}"))?;
|
||||
}
|
||||
}
|
||||
|
||||
// Join: main stream waits for IQN and Attention to complete
|
||||
if let Some(ref event) = self.iqn_done_event {
|
||||
self.stream.wait(event)
|
||||
.map_err(|e| anyhow::anyhow!("main wait iqn: {e}"))?;
|
||||
}
|
||||
if let Some(ref event) = self.attn_done_event {
|
||||
self.stream.wait(event)
|
||||
.map_err(|e| anyhow::anyhow!("main wait attn: {e}"))?;
|
||||
}
|
||||
|
||||
// CQL + ensemble + recursive_confidence_bwd run on main stream after join
|
||||
```
|
||||
|
||||
- [ ] **Step 2: Move IQN code out of the sequential section**
|
||||
|
||||
Remove the existing IQN training from the sequential section (currently after IQL).
|
||||
|
||||
- [ ] **Step 3: Move Attention code out of the sequential section**
|
||||
|
||||
Remove the existing Attention forward+backward+adam from the sequential section.
|
||||
|
||||
- [ ] **Step 4: Build and test**
|
||||
|
||||
- [ ] **Step 5: Commit**
|
||||
|
||||
```bash
|
||||
git commit -m "feat: fork IQN+Attention to parallel streams in submit_aux_ops"
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
### Task 4: Switch Capture Mode to RELAXED
|
||||
|
||||
**Files:**
|
||||
- Modify: `crates/ml/src/trainers/dqn/fused_training.rs` (capture_child_graph)
|
||||
|
||||
- [ ] **Step 1: Change capture mode from GLOBAL to RELAXED**
|
||||
|
||||
In `capture_child_graph`, change:
|
||||
|
||||
```rust
|
||||
// Before:
|
||||
let begin_result = self.stream.begin_capture(
|
||||
cuda_sys::CUstreamCaptureMode::CU_STREAM_CAPTURE_MODE_GLOBAL,
|
||||
);
|
||||
|
||||
// After:
|
||||
let begin_result = self.stream.begin_capture(
|
||||
cuda_sys::CUstreamCaptureMode::CU_STREAM_CAPTURE_MODE_RELAXED,
|
||||
);
|
||||
```
|
||||
|
||||
RELAXED mode is required for multi-stream capture — it allows other streams to join the capture graph via `cuStreamWaitEvent`. GLOBAL mode would reject non-captured streams.
|
||||
|
||||
- [ ] **Step 2: Build and test**
|
||||
|
||||
- [ ] **Step 3: Deploy to H100 and verify child graph breakdown shows aux < 200ms**
|
||||
|
||||
Run: `./scripts/argo-train.sh dqn --baseline --epochs 2`
|
||||
|
||||
Check logs for: `GPU child graph breakdown: spectral=X forward=Y ddqn=Z aux=W adam=V`
|
||||
|
||||
Expected: aux drops from ~295ms to ~180ms.
|
||||
|
||||
- [ ] **Step 4: Commit**
|
||||
|
||||
```bash
|
||||
git commit -m "perf: switch to RELAXED capture mode — enables multi-stream graph capture"
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
### Task 5: Forward Backward Branch Parallelism
|
||||
|
||||
**Files:**
|
||||
- Modify: `crates/ml/src/cuda_pipeline/batched_backward.rs` (backward_full)
|
||||
|
||||
The forward pass already uses `branch_streams[4]` with fork-join. The backward needs the same pattern. However, branches share scratch buffers (`scratch_d_glu_value`, `scratch_d_glu_gate`) that prevent parallel execution.
|
||||
|
||||
- [ ] **Step 1: Add per-branch GLU scratch buffers to CublasBackwardSet**
|
||||
|
||||
```rust
|
||||
// In CublasBackwardSet struct:
|
||||
/// Per-branch GLU backward scratch (Phase 2: parallel branch backward).
|
||||
/// 4 × [B, AH] for d_value and 4 × [B, AH] for d_gate_pre.
|
||||
branch_d_glu_value: [CudaSlice<f32>; 4],
|
||||
branch_d_glu_gate: [CudaSlice<f32>; 4],
|
||||
```
|
||||
|
||||
Allocate in constructor: `4 × 2 × B × AH × 4 bytes = 4 × 2 × 16384 × 32 × 4 = 16MB`.
|
||||
|
||||
- [ ] **Step 2: Add branch_streams + events to CublasBackwardSet**
|
||||
|
||||
Reuse the same `branch_streams` and events from `CublasForward` (passed at construction or shared via Arc).
|
||||
|
||||
- [ ] **Step 3: Fork-join in backward_full branch loop**
|
||||
|
||||
Replace the sequential `for d in 0..4` loop with fork-join:
|
||||
|
||||
```rust
|
||||
// Record pre-branch state on main stream
|
||||
self.trunk_done_event.record(stream)?;
|
||||
|
||||
for d in 0..4 {
|
||||
let bs = &self.branch_streams[d];
|
||||
bs.wait(&self.trunk_done_event)?;
|
||||
|
||||
// Run branch d backward on branch_stream[d]
|
||||
// Use per-branch scratch: branch_d_glu_value[d], branch_d_glu_gate[d]
|
||||
self.backward_branch_on_stream(bs, d, /* ... args with per-branch scratch ... */)?;
|
||||
|
||||
self.branch_done_events[d].record(bs)?;
|
||||
stream.wait(&self.branch_done_events[d])?;
|
||||
}
|
||||
|
||||
// Trunk backward runs on main stream after all branches complete
|
||||
```
|
||||
|
||||
- [ ] **Step 4: Extract single-branch backward into `backward_branch_on_stream`**
|
||||
|
||||
Move the body of the `for d in 0..4` loop into a method that takes a stream parameter.
|
||||
|
||||
- [ ] **Step 5: Build and test**
|
||||
|
||||
- [ ] **Step 6: Commit**
|
||||
|
||||
```bash
|
||||
git commit -m "perf: parallel backward branches using fork-join on branch_streams"
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
### Task 6: Fix Adam Child Event Timing
|
||||
|
||||
**Files:**
|
||||
- Modify: `crates/ml/src/trainers/dqn/fused_training.rs`
|
||||
|
||||
This was already done in commit `8ebb65a94` — added `adam_child_end` event. Verify the output shows a real value instead of `-1.0ms`.
|
||||
|
||||
- [ ] **Step 1: Deploy and check logs**
|
||||
|
||||
Expected: `GPU child graph breakdown: spectral=X forward=Y ddqn=Z aux=W adam=V` with V > 0.
|
||||
|
||||
- [ ] **Step 2: If adam is unexpectedly large (>100ms), investigate**
|
||||
|
||||
The adam_child contains: mamba2_backward (12K params), step_mamba2_adam (saxpy), apply_pruning_mask, compute_grad_norm (2 kernels), launch_adam_update (660K params), unflatten (24 DtoD or zero-copy skip), update_isv_signals. Total should be <50ms.
|
||||
|
||||
---
|
||||
|
||||
### Task 7: Verify and Profile
|
||||
|
||||
- [ ] **Step 1: Deploy with all changes**
|
||||
|
||||
```bash
|
||||
./scripts/argo-train.sh dqn --baseline --epochs 5 --watch
|
||||
```
|
||||
|
||||
- [ ] **Step 2: Check child graph breakdown**
|
||||
|
||||
Expected:
|
||||
```
|
||||
GPU child graph breakdown:
|
||||
spectral = ~0.2ms
|
||||
forward = ~100ms (backward branches parallel)
|
||||
ddqn = ~1ms
|
||||
aux = ~180ms (IQN + Attention parallel)
|
||||
adam = ~30ms
|
||||
total = ~310ms (down from ~430ms)
|
||||
```
|
||||
|
||||
- [ ] **Step 3: Check epoch time**
|
||||
|
||||
Expected: ~310ms/step × 178 steps = ~55s training + 3.5s step0 + 2s experience + 1s validation = **~62s per epoch** (down from ~718s).
|
||||
|
||||
- [ ] **Step 4: Commit final timing results to spec**
|
||||
|
||||
Update `docs/superpowers/specs/2026-04-18-multi-stream-phase2-design.md` with measured results.
|
||||
Reference in New Issue
Block a user