diff --git a/docs/superpowers/plans/2026-04-17-unified-cublas-training-graph.md b/docs/superpowers/plans/2026-04-17-unified-cublas-training-graph.md index 205b7cd0c..c1fd674d1 100644 --- a/docs/superpowers/plans/2026-04-17-unified-cublas-training-graph.md +++ b/docs/superpowers/plans/2026-04-17-unified-cublas-training-graph.md @@ -1951,6 +1951,109 @@ Task 6 (Attention cubins)──┘ --- +## AMENDMENTS (from plan review) + +### Amendment A: train_step_gpu must not capture its own graphs + +**Affects:** Task 9 (Step 9.3) and Task 10. + +The ungraphed path at step 0 calls `self.trainer.train_step_gpu()` which internally calls `capture_training_graphs()` and creates `graph_forward`, `graph_forward_ddqn`, `graph_adam`. These MUST NOT be created during the unified path — the unified graph captures everything. + +**Fix:** In Task 10, when removing `capture_training_graphs()` from `gpu_dqn_trainer.rs`, also modify `train_step_gpu()` to ONLY submit ops (call `submit_forward_ops_main()` + `submit_forward_ops_ddqn()`) WITHOUT capturing any graphs. Rename to `submit_training_ops_ungraphed()` for clarity. The graph capture happens ONLY in `capture_graph_unified()`. + +```rust +/// Submit all training ops ungraphed (step 0 only, before unified graph capture). +pub(crate) fn submit_training_ops_ungraphed( + &mut self, + gpu_batch: &GpuBatch, + online_d: &DuelingWeightSet, + online_b: &BranchingWeightSet, + target_d: &DuelingWeightSet, + target_b: &BranchingWeightSet, +) -> Result<(), MLError> { + self.upload_batch_gpu(gpu_batch)?; + self.update_stochastic_depth_mask()?; + self.adam_step += 1; + unsafe { *self.t_pinned = self.adam_step; } + self.submit_forward_ops_main()?; + self.submit_forward_ops_ddqn()?; + Ok(()) +} +``` + +### Amendment B: Mamba2 backward must be inside unified graph + +**Affects:** Task 7 and Task 9. + +Mamba2 forward is inside the unified graph (via `submit_forward_ops_main` from Task 7). Mamba2 backward is currently outside the graph (line 1693 in Step 9.3). This means during replay, the backward operates on stale activations from the PREVIOUS capture. + +**Fix:** Move `mamba2_backward()` + `step_mamba2_adam()` inside the unified graph capture in Task 9 (after Phase 6 Adam, before Phase 7 ISV). Update `capture_graph_unified()`: + +```rust + // ── Phase 6: Pruning + grad norm + Adam + unflatten ── + let adam_result = if conf_result.is_ok() { + self.trainer.apply_pruning_mask() + .and_then(|_| self.trainer.submit_adam_ops()) + .map_err(|e| anyhow::anyhow!("unified adam: {e}")) + } else { Ok(()) }; + + // ── Phase 6b: Mamba2 BPTT + SGD (must be inside graph with forward) ── + let mamba_result = if adam_result.is_ok() { + self.mamba2_backward(self.batch_size) + .and_then(|_| self.step_mamba2_adam()) + .map_err(|e| anyhow::anyhow!("unified mamba2 bwd: {e}")) + } else { Ok(()) }; + + // ── Phase 7: ISV signal update ── + let isv_result = if mamba_result.is_ok() { + // ... +``` + +Remove `self.mamba2_backward()` and `self.step_mamba2_adam()` from the post-graph section of `run_full_step()`. + +### Amendment C: submit_adam_ops must exist as public method + +**Affects:** Task 9 (Step 9.2). + +The capture references `self.trainer.submit_adam_ops()` but the current code's public method is `replay_adam_and_readback()` which replays a pre-captured graph_adam. Task 10 removes graph_adam. The underlying op submission function `submit_adam_ops()` already exists at `gpu_dqn_trainer.rs:8125` — it just needs to be `pub(crate)` (currently it is). + +**Verify:** `submit_adam_ops()` signature at gpu_dqn_trainer.rs:8125: +```rust +pub(crate) fn submit_adam_ops( + &mut self, + online_d: &DuelingWeightSet, + online_b: &BranchingWeightSet, +) -> Result<(), MLError> +``` + +Note: it takes weight set refs. In the unified capture, these are available as `self.online_dueling` / `self.online_branching`. Update the capture call: +```rust +self.trainer.submit_adam_ops(&self.online_dueling, &self.online_branching) +``` + +### Amendment D: cuBLAS workspace sharing verification + +**Affects:** Tasks 1, 3, 5. + +All IQL (8 GEMMs), IQN (~10 GEMMs), and Attention (~12 GEMMs) share the single `lt_workspace_ptr` (32 MB) from `SharedCublasHandle`. During unified graph capture, all GEMMs are stream-ordered (one at a time) so workspace sharing is safe — no concurrent access. + +**Verification step:** Add to each GEMM descriptor creation: log the `heuristic.workspaceSize` returned by the heuristic. If ANY GEMM needs >32MB, increase `SharedCublasHandle::lt_workspace_size`. + +```rust +let heuristic = cublaslt_result::get_matmul_algo_heuristic(...)?; +tracing::info!(%label, ws_needed = heuristic.workspaceSize, "GEMM desc created"); +if heuristic.workspaceSize > lt_ws_size { + return Err(MLError::ModelError(format!( + "{label}: needs {}MB workspace but only {}MB available", + heuristic.workspaceSize / (1024*1024), lt_ws_size / (1024*1024) + ))); +} +``` + +Add this to Task 1 Step 1.2, Task 3, and Task 5 constructors. + +--- + ## Task 13: Hardcoded Dimension Cleanup **Files:**