From 6b2cbc2965868545238d755c8e4023e1341b30f2 Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Sat, 18 Apr 2026 00:12:14 +0200 Subject: [PATCH] =?UTF-8?q?spec:=20detailed=20Phase=202=20implementation?= =?UTF-8?q?=20=E2=80=94=20workspace=20isolation,=20double=20buffer,=20dete?= =?UTF-8?q?rminism?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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) --- ...17-unified-cublas-training-graph-design.md | 180 ++++++++++++++++-- 1 file changed, 168 insertions(+), 12 deletions(-) diff --git a/docs/superpowers/specs/2026-04-17-unified-cublas-training-graph-design.md b/docs/superpowers/specs/2026-04-17-unified-cublas-training-graph-design.md index 758867363..40ce799e2 100644 --- a/docs/superpowers/specs/2026-04-17-unified-cublas-training-graph-design.md +++ b/docs/superpowers/specs/2026-04-17-unified-cublas-training-graph-design.md @@ -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; 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::(32 * 1024 * 1024)?; // 32MB +let iql_low_workspace = stream.alloc_zeros::(32 * 1024 * 1024)?; +let iqn_workspace = stream.alloc_zeros::(32 * 1024 * 1024)?; +let attn_workspace = stream.alloc_zeros::(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; 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, // [episodes * timesteps * state_dim] + actions_a: CudaSlice, + rewards_a: CudaSlice, + dones_a: CudaSlice, + + // Buffer B (same sizes) + states_b: CudaSlice, + actions_b: CudaSlice, + rewards_b: CudaSlice, + dones_b: CudaSlice, + + // 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; 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.