spec: detailed Phase 2 implementation — workspace isolation, double buffer, determinism
Phase 2a: branch parallel — cuBLAS workspace isolation per branch stream, buffer isolation verified, backward join before trunk reduction. Phase 2b: aux parallel — per-trainer workspace (128MB additional), 4 aux streams, dependency graph within aux_child detailed. Phase 2c: double buffer — ping-pong experience buffers (3GB additional), collection/training stream separation, replay buffer staleness analysis, child graph re-capture for buffer pointer swap. Determinism: CUDA Graph replay is bit-identical via fixed dependency edges. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
@@ -381,32 +381,188 @@ stream.end_capture()?; // Graph captures the fork-join topology
|
||||
The captured graph replays all 4 branches in parallel automatically — the
|
||||
graph scheduler sees the dependency edges and dispatches to multiple SMs.
|
||||
|
||||
### Expected Performance (Phase 2)
|
||||
### Implementation Details
|
||||
|
||||
| Metric | Phase 1 | Phase 2 |
|
||||
|--------|---------|---------|
|
||||
| Forward pass | ~3ms (sequential branches) | ~1.2ms (parallel branches) |
|
||||
| Aux ops | ~5ms (sequential IQL+IQN+attn) | ~2ms (parallel trainers) |
|
||||
| SM utilization | ~40% (sequential) | ~85% (parallel) |
|
||||
| **Epoch time** | **<80s** | **<40s** |
|
||||
#### 2a. Parallel Branch Forward (forward_child)
|
||||
|
||||
**cuBLAS workspace isolation:** Each branch stream needs its own cuBLAS
|
||||
workspace to avoid contention. `batched_forward.rs` already allocates:
|
||||
- `branch_workspace_ptrs: [u64; 4]` — 4 × 32MB = 128MB total
|
||||
- `branch_streams: [Arc<CudaStream>; 4]` — forked from main stream
|
||||
- `trunk_done_event` + `branch_done_events[4]` — fork/join synchronization
|
||||
|
||||
**cublasLt handle sharing:** `cublasLtMatmul` is thread-safe when using
|
||||
different streams + different workspaces. The single `lt_handle` is shared
|
||||
across all 5 streams. Each branch calls `cublasLtMatmul(..., branch_ws_ptr,
|
||||
branch_ws_size, branch_stream)` — no internal state conflict because cublasLt
|
||||
is stateless per-call (unlike legacy cublas which has handle-bound stream).
|
||||
|
||||
**Buffer isolation:** Each branch writes to its own output buffer:
|
||||
- Branch 0 (value): `on_v_logits_buf[B, NA]`
|
||||
- Branch 1-3 (advantage): `on_b_logits_buf` at offsets `[B*b0*NA, B*b1*NA, ...]`
|
||||
- No write conflicts — each branch has exclusive output region.
|
||||
|
||||
**Read sharing:** All branches READ from `save_h_s2[B, SH2]` (trunk output).
|
||||
Read-read sharing is safe — no synchronization needed beyond the fork event.
|
||||
|
||||
**Backward parallelism:** Within each branch's backward:
|
||||
- `dW_bi = d_logits_bi @ h_bi^T` (weight gradient) — writes to `grad_buf[offset_bi]`
|
||||
- `dh_bi = W_bi^T @ d_logits_bi` (activation gradient) — writes to `d_h_bi_buf`
|
||||
- These are independent for different branches. The trunk backward aggregates
|
||||
`d_h_s2 = sum(W_bi^T @ d_h_bi)` — this requires a join before the reduction.
|
||||
|
||||
#### 2b. Parallel Aux Trainers (aux_child)
|
||||
|
||||
**Workspace requirement:** IQL, IQN, and Attention each need their own cuBLAS
|
||||
workspace. Current: all share `SharedCublasHandle::lt_workspace_ptr` (32MB).
|
||||
Phase 2 adds per-trainer workspace buffers:
|
||||
|
||||
```rust
|
||||
// In FusedTrainingCtx constructor:
|
||||
let iql_workspace = stream.alloc_zeros::<u8>(32 * 1024 * 1024)?; // 32MB
|
||||
let iql_low_workspace = stream.alloc_zeros::<u8>(32 * 1024 * 1024)?;
|
||||
let iqn_workspace = stream.alloc_zeros::<u8>(32 * 1024 * 1024)?;
|
||||
let attn_workspace = stream.alloc_zeros::<u8>(32 * 1024 * 1024)?;
|
||||
// Total: 128MB additional VRAM
|
||||
```
|
||||
|
||||
Each trainer receives its own `(workspace_ptr, workspace_size)` at construction.
|
||||
|
||||
**Stream allocation:** 4 aux streams forked from main:
|
||||
```rust
|
||||
let aux_streams: [Arc<CudaStream>; 4] = [
|
||||
stream.fork()?, // IQL high
|
||||
stream.fork()?, // IQL low
|
||||
stream.fork()?, // IQN
|
||||
stream.fork()?, // Attention
|
||||
];
|
||||
```
|
||||
|
||||
**Dependency graph within aux_child:**
|
||||
```
|
||||
[fork from main after backward]
|
||||
│
|
||||
├── stream_0: HER relabel → EMA target update → IQL high (fwd+bwd+adam)
|
||||
├── stream_1: IQL low (fwd+bwd+adam)
|
||||
├── stream_2: IQN (fwd+bwd+adam+EMA)
|
||||
└── stream_3: Attention (fwd+bwd+adam)
|
||||
│
|
||||
[join all 4 → main stream]
|
||||
│
|
||||
└── CQL gradient (reads from all trainers) → recursive_confidence_bwd
|
||||
```
|
||||
|
||||
HER + EMA run on stream_0 before IQL high because IQL reads from the updated
|
||||
target weights. IQL low can start immediately (reads from online weights only).
|
||||
CQL gradient reads from all trainers' outputs — must wait for all 4 to complete.
|
||||
|
||||
**Buffer isolation:** Each trainer has its own:
|
||||
- `params_buf`, `m_buf`, `v_buf`, `grad_buf` — exclusive
|
||||
- `states_buf` (read-only, shared) — safe
|
||||
- `q_out_buf` (read by IQL gather) — written by forward_child, read by aux_child.
|
||||
Dependency enforced by child graph ordering (forward_child before aux_child).
|
||||
|
||||
#### 2c. Double-buffered Experience Collection
|
||||
|
||||
**Problem:** Experience collection (4096 episodes × 1000 timesteps = ~1.2s)
|
||||
runs sequentially before training (~56s on single stream). The GPU idles
|
||||
during collection because training graphs can't run simultaneously.
|
||||
|
||||
**Architecture:** Two sets of experience output buffers (ping-pong):
|
||||
|
||||
```rust
|
||||
struct DoubleBufferedExperience {
|
||||
// Buffer A
|
||||
states_a: CudaSlice<f32>, // [episodes * timesteps * state_dim]
|
||||
actions_a: CudaSlice<i32>,
|
||||
rewards_a: CudaSlice<f32>,
|
||||
dones_a: CudaSlice<f32>,
|
||||
|
||||
// Buffer B (same sizes)
|
||||
states_b: CudaSlice<f32>,
|
||||
actions_b: CudaSlice<i32>,
|
||||
rewards_b: CudaSlice<f32>,
|
||||
dones_b: CudaSlice<f32>,
|
||||
|
||||
// Which buffer is "training" (being read by graph) vs "collecting" (being written)
|
||||
active: usize, // 0 = A training, B collecting. 1 = swap.
|
||||
}
|
||||
```
|
||||
|
||||
**Stream separation:**
|
||||
- `training_stream`: runs the parent graph (reads from active buffer)
|
||||
- `collection_stream`: runs experience collection (writes to inactive buffer)
|
||||
- Synchronization via events:
|
||||
```
|
||||
Epoch N:
|
||||
collection_stream: collect into buffer B
|
||||
training_stream: train from buffer A (parent_graph.launch())
|
||||
[event: collection B done]
|
||||
[event: training A done]
|
||||
swap active buffer
|
||||
Epoch N+1:
|
||||
collection_stream: collect into buffer A
|
||||
training_stream: train from buffer B
|
||||
...
|
||||
```
|
||||
|
||||
**VRAM cost:** Doubles the experience output buffer: ~1.5GB × 2 = 3GB additional.
|
||||
On H100 (80GB) with ~47GB used, this fits within headroom.
|
||||
|
||||
**Replay buffer interaction:** The replay buffer is written by collection and
|
||||
read by training. With double buffering, the PER segment tree updates happen
|
||||
on the collection stream AFTER experience is written. Training reads from
|
||||
the PREVIOUS epoch's segment tree state. One-epoch staleness in priorities
|
||||
is acceptable — PER priorities are approximate by nature.
|
||||
|
||||
**Graph capture implication:** The parent graph reads from fixed buffer
|
||||
addresses (captured at step 0). With double buffering, the parent graph must
|
||||
be re-instantiated each epoch to point at the swapped buffer — or use
|
||||
`cudaGraphExecKernelNodeSetParams` to update the buffer pointers. The child
|
||||
graph architecture helps: only `forward_child` reads from the experience
|
||||
buffer (states/rewards/dones). Re-capture just that one child, re-compose
|
||||
parent. Other children unchanged.
|
||||
|
||||
### Expected Performance
|
||||
|
||||
| Metric | Phase 1 | Phase 2a+b | Phase 2c |
|
||||
|--------|---------|------------|----------|
|
||||
| Forward pass | ~3ms | ~1.2ms | ~1.2ms |
|
||||
| Aux ops | ~5ms | ~2ms | ~2ms |
|
||||
| Experience collection | 1.2s (blocking) | 1.2s (blocking) | 0s (overlapped) |
|
||||
| SM utilization | ~40% | ~85% | ~90% |
|
||||
| **Epoch time** | **<80s** | **<40s** | **<38s** |
|
||||
|
||||
### What Already Exists
|
||||
|
||||
`batched_forward.rs` already has the infrastructure:
|
||||
`batched_forward.rs` infrastructure ready for Phase 2a:
|
||||
- `branch_streams: [Arc<CudaStream>; 4]` — 4 dedicated branch streams
|
||||
- `branch_workspace_ptrs: [u64; 4]` — per-branch cuBLAS workspace (no contention)
|
||||
- `trunk_done_event: CudaEvent` — fork event after trunk
|
||||
- `branch_done_events: [CudaEvent; 4]` — join events per branch
|
||||
- `sgemm_f32_branch()` — already routes GEMMs to per-branch stream+workspace
|
||||
|
||||
These are currently unused during graph capture (single-stream mode). Phase 2
|
||||
These are currently unused during graph capture (single-stream mode). Phase 2a
|
||||
activates them inside the `forward_child` capture with RELAXED mode.
|
||||
|
||||
### Determinism Guarantee
|
||||
|
||||
Multi-stream execution is non-deterministic in general (kernel scheduling
|
||||
order varies). But CUDA Graph replay is deterministic — the graph captures
|
||||
the exact dependency topology and replays it identically every time. The
|
||||
fork-join events become fixed graph edges. As long as:
|
||||
1. No atomicAdd across streams (we have zero atomicAdd policy)
|
||||
2. No floating-point reduction across streams (each stream reduces independently)
|
||||
3. All cross-stream joins happen via captured events (not ad-hoc synchronization)
|
||||
|
||||
...the multi-stream graph replay is bit-identical across runs.
|
||||
|
||||
### Implementation Order
|
||||
|
||||
1. Phase 1: single-stream child graphs (this spec) — validate correctness
|
||||
2. Phase 2: multi-stream forward_child — parallelize 4 branch heads
|
||||
3. Phase 2b: multi-stream aux_child — parallelize IQL/IQN/attention
|
||||
4. Phase 2c: double-buffered experience collection — overlap with training
|
||||
2. Phase 2a: multi-stream forward_child — parallelize 4 branch heads
|
||||
3. Phase 2b: multi-stream aux_child — parallelize IQL/IQN/attention (128MB additional workspace)
|
||||
4. Phase 2c: double-buffered experience collection (3GB additional VRAM)
|
||||
|
||||
Each phase is additive — Phase 1 child graphs don't change, they just gain
|
||||
internal parallelism. The parent graph composition is unchanged.
|
||||
|
||||
Reference in New Issue
Block a user