From 5a4d56145141cef8807266495b94cdb956286d3d Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Wed, 22 Apr 2026 09:46:17 +0200 Subject: [PATCH] =?UTF-8?q?plan(policy-quality):=20Task=202.0=20revised=20?= =?UTF-8?q?approach=20=E2=80=94=20in-graph=20pinned=20snapshots?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit First Task 2.0 dispatch escalated BLOCKED: the four loss-component backward kernels are captured inside the fused training graph, so host-side snapshot-between-components isn't possible mid-graph without a force-ungraphed diagnostic step (~210 LOC + cross-stream sync risk). Revised approach (chosen after cost analysis): cudaMemcpyAsync(device → pinned host) IS captureable in a CUDA graph. Even better: DtoD into per-component scratch buffers, then an in-graph reduction kernel computes per-component (mag_norm, dir_norm) and writes 8 floats to a pinned result slot. Only the 8-float result crosses PCIe (at epoch boundary), keeping per-step PCIe traffic to zero. Changes to the plan's Step 2 + Step 3 + Step 4: - Step 2: added 4 device-side scratch buffers (one per component, ~10 MB each = 40 MB device) + 8-float pinned result slot + new reduction kernel grad_decomp_kernel.cu spec'd out. - Step 3: clarified that DtoD snapshot + backward + reduction kernel are ALL captured in the graph; graph replays them every step; no force-ungraphed dance needed. - Step 4: added refresh_grad_component_norms() accessor that reads the 8-float pinned slot at epoch boundary (zero-copy) and populates the host-side cache. Approach matches Task 0.4 pattern (commit bb42c9963) extended four-fold. No atomicAdd (reduction uses shared-mem tree), no non-captured replay, no cross-stream sync risk. LOC estimate: ~90 (was ~210 for the rejected option). --- .../plans/2026-04-21-policy-quality-phase2.md | 188 +++++++++++++----- 1 file changed, 135 insertions(+), 53 deletions(-) diff --git a/docs/superpowers/plans/2026-04-21-policy-quality-phase2.md b/docs/superpowers/plans/2026-04-21-policy-quality-phase2.md index 4572763dc..6243e2179 100644 --- a/docs/superpowers/plans/2026-04-21-policy-quality-phase2.md +++ b/docs/superpowers/plans/2026-04-21-policy-quality-phase2.md @@ -92,87 +92,169 @@ grep -n "iqn_grad\|cql_grad\|c51_grad\|ens_grad\|component_grad\|iqn_backward\|c ``` Expected: each component adds into the shared `adam_grad` buffer via its own kernel. The current Task 0.4 accessor sums the final `adam_grad` post-all-components; this task captures the per-component deltas. -- [ ] **Step 2: Add a `grad_snapshot_pinned_ptr` + snapshot helper** +- [ ] **Step 2: Add pinned snapshot slots + in-graph reduction kernel** -In `crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs`, alongside the existing Task 0.4 `grad_readback_pinned_ptr`: +**Approach decision (from Task 2.0 first-dispatch escalation):** the four loss-component backward kernels are all captured inside the fused training graph (`fused_training.rs:1667-1776`), so host-side snapshots with a synchronize + memcpy-out pattern are not possible mid-graph. Three implementation options were evaluated: + +| Option | Cost | 4-way split preserved? | Decision | +|---|---|---|---| +| Force-ungraphed diagnostic step (non-captured replay) | ~210 LOC + cross-stream sync risk + ~200ms recapture | ✅ | rejected (complexity) | +| 2-way split only (C51 vs bundle {IQN+CQL+Ens}) | ~40 LOC | ❌ loses information | rejected (insufficient for Task 2.1) | +| **In-graph pinned snapshots + in-graph reduction kernel** | ~90 LOC | ✅ | **accepted** | + +The accepted approach: `cudaMemcpyAsync(device → pinned host, …)` IS captureable in a CUDA graph. The memcpy fires every step (implicit in graph replay), but the PCIe cost is benign — 4 × TOTAL_PARAMS × 4 bytes per step ≈ 40 MB/step on PCIe at ~25 GB/s sustained is <2% of bandwidth, fully async with kernel execution, wall-clock impact negligible. Pattern matches Task 0.4 exactly (commit `bb42c9963`), just extended four-fold. + +**Important**: to keep HEALTH_DIAG traffic tiny, snapshot TOTAL_PARAMS to device-side scratch buffers (one per component), then launch a small reduction kernel (also captured in the graph) that computes `‖Δmag‖` + `‖Δdir‖` over the per-component delta and stores 2 floats per component in a result slot. DtoH only the final 8-float result at epoch boundary. This keeps per-step PCIe traffic to zero and matches the "host-reduce-at-epoch-boundary" rule established by Task 0.5/0.8. + +Schema to add in `crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs`: ```rust /// Task 2.0 — per-component magnitude-branch grad decomposition. -/// Clone the adam_grad buffer into a pinned host slot BETWEEN each -/// component's backward kernel, then compute magnitude-head L2 norm -/// on the delta (post − pre). Produces [iqn, cql, c51, ens] norms. -grad_snapshot_pinned_ptr: *mut f32, +/// Four device-side scratch buffers, one per loss component. Each slot +/// holds a snapshot of adam_grad taken BEFORE that component's backward. +/// After the component's backward, a reduction kernel computes +/// ‖Δmag‖ + ‖Δdir‖ on the (current_adam_grad − snapshot) slices and +/// writes them into grad_component_norms_device. Only the final +/// 8-float result crosses PCIe (DtoH at epoch boundary). +grad_snapshot_iqn: CudaSlice, // [TOTAL_PARAMS] +grad_snapshot_cql: CudaSlice, +grad_snapshot_c51: CudaSlice, +grad_snapshot_ens: CudaSlice, + +/// Pinned 8-float result slot: [iqn_mag, iqn_dir, cql_mag, cql_dir, +/// c51_mag, c51_dir, ens_mag, ens_dir]. +/// Written by the reduction kernel in-graph; read host-side at epoch +/// boundary (zero-copy via pinned mapping). +grad_component_norms_pinned_ptr: *mut f32, // length 8 + +/// Cached host-side ratios for HEALTH_DIAG readout. grad_component_norms_mag: [f32; 4], // [iqn, cql, c51, ens] -grad_component_norms_dir: [f32; 4], // [iqn, cql, c51, ens] for ratio denominator +grad_component_norms_dir: [f32; 4], ``` -Allocate the pinned slot at construction (`stream.alloc_pinned::(TOTAL_PARAMS)`); free in the trainer's `Drop` alongside other pinned slots. +Allocate the four device scratch buffers at construction via the standard `stream.alloc_zeros::(TOTAL_PARAMS)?` (each ~10 MB on RTX 3050 Ti = 40 MB total device; <1% of a 4 GB budget). Allocate the 8-float pinned slot via `cuMemAllocHost` (same pattern as Task 0.4's `grad_readback_pinned_ptr`). Free all 5 in the trainer's `Drop` alongside the existing pinned slots. -Add the snapshot helper (non-graph, epoch-boundary or pre-/post-component-kernel — **not** inside the captured CUDA graph; this is diagnostic-only): +**Reduction kernel** (create `crates/ml/src/cuda_pipeline/grad_decomp_kernel.cu`): -```rust -/// Copy current adam_grad to pinned host. Synchronous (sync stream first). -/// Not inside graph capture. -fn snapshot_adam_grad(&mut self) -> Result<(), MLError> { - self.stream.synchronize()?; - unsafe { - let slice = std::slice::from_raw_parts_mut( - self.grad_snapshot_pinned_ptr, - TOTAL_PARAMS, - ); - self.stream.memcpy_dtoh(&self.adam_grad, slice)?; +```cuda +// Computes L2 norms of (current − snapshot) on magnitude and direction +// branch slices, writes 2 floats to result_out. +// Called four times per step, once after each component's backward. +__global__ void grad_component_delta_norm( + const float* __restrict__ current, // adam_grad now [TOTAL_PARAMS] + const float* __restrict__ snapshot, // adam_grad before this component [TOTAL_PARAMS] + int mag_start, int mag_len, // byte-offset slice for branch 1 (magnitude) + int dir_start, int dir_len, // byte-offset slice for branch 0 (direction) + float* __restrict__ result_out // [2] — (mag_norm, dir_norm) +) { + __shared__ float sum_mag[256]; + __shared__ float sum_dir[256]; + int tid = threadIdx.x; + sum_mag[tid] = 0.0f; + sum_dir[tid] = 0.0f; + + // Stride-loop over magnitude slice + for (int i = tid; i < mag_len; i += blockDim.x) { + float d = current[mag_start + i] - snapshot[mag_start + i]; + sum_mag[tid] += d * d; } - Ok(()) -} + // Stride-loop over direction slice + for (int i = tid; i < dir_len; i += blockDim.x) { + float d = current[dir_start + i] - snapshot[dir_start + i]; + sum_dir[tid] += d * d; + } + __syncthreads(); -/// Compute magnitude-branch and direction-branch L2 norms on -/// (adam_grad_now − snapshot). Uses Task 0.4's compute_param_sizes -/// + padded_byte_offset to slice branch tensors 8–11 (dir) and -/// 12–15 (mag) out of the 42-tensor layout. -fn component_grad_norms(&mut self, prev_snapshot: &[f32]) -> Result<(f32, f32), MLError> { - /* sync, dtoh current adam_grad, subtract prev_snapshot, L2 on mag/dir slices */ + // Block-level tree reduction (256 → 1) + for (int s = blockDim.x / 2; s > 0; s >>= 1) { + if (tid < s) { + sum_mag[tid] += sum_mag[tid + s]; + sum_dir[tid] += sum_dir[tid + s]; + } + __syncthreads(); + } + if (tid == 0) { + result_out[0] = sqrtf(sum_mag[0]); + result_out[1] = sqrtf(sum_dir[0]); + } } ``` -- [ ] **Step 3: Wire per-component snapshots in the training step** +One-block launch with 256 threads. No atomicAdd (block-internal `__syncthreads` + shared-mem tree reduction — the established deterministic pattern). -In `crates/ml/src/cuda_pipeline/fused_training.rs`, find the loss-component backward sequence (grep for `iqn_backward`, `cql_backward`, `c51_grad_kernel`, `ensemble_backward` or equivalent). Insert snapshot + post-compute-norm calls: +- [ ] **Step 3: Wire in-graph snapshot + reduction around each backward** + +In `crates/ml/src/cuda_pipeline/fused_training.rs` (aux-graph orchestration around lines 1146-1559 per the escalation report), insert captured DtoD snapshot + captured reduction-kernel launch around each component: ```rust -// Before IQN backward -trainer.snapshot_adam_grad()?; -iqn_backward(/* args */)?; -let (mag_iqn, dir_iqn) = trainer.component_grad_norms(pinned_snapshot_prev)?; -trainer.grad_component_norms_mag[0] = mag_iqn; -trainer.grad_component_norms_dir[0] = dir_iqn; - -// Before CQL backward -trainer.snapshot_adam_grad()?; -cql_backward(/* args */)?; -let (mag_cql, dir_cql) = trainer.component_grad_norms(pinned_snapshot_prev)?; -trainer.grad_component_norms_mag[1] = mag_cql; -trainer.grad_component_norms_dir[1] = dir_cql; - -// ... repeat for c51 (index 2) and ens (index 3) ... +// Before C51 backward (example; repeat for IQN/CQL/Ens): +unsafe { + self.stream.memcpy_dtod( + &trainer.adam_grad.slice(..total_params), + &mut trainer.grad_snapshot_c51.slice_mut(..total_params), + )?; +} +// Existing backward: +launch_cublas_backward_c51(/* args */)?; +// After: +launch_grad_component_delta_norm( + &trainer.adam_grad, + &trainer.grad_snapshot_c51, + mag_slice_start, mag_slice_len, + dir_slice_start, dir_slice_len, + &mut trainer.grad_component_norms_device.slice_mut(4..6), // offsets for c51 = 2 +)?; ``` -**Important:** this breaks graph capture for the diagnostic steps. Guard the snapshots inside the non-captured epoch-boundary path — per-step capture is preserved. If the loss backwards are all inside one captured graph, fall back to running the diagnostic every Nth *epoch* (not step) by replaying without capture. Document the chosen approach in the commit message. +All three operations (DtoD memcpy, component backward, reduction kernel launch) are fully captureable. The graph replays them every step — intended. Each step's 8-float result is OVERWRITTEN into the same pinned slot; host reads the current value at epoch-boundary HEALTH_DIAG, so what we see is the last step of each epoch. For Task 2.1's decision tree, last-step-of-epoch is a sufficient proxy for epoch-averaged (the signal we care about is whether the ratio is ~0 or ~1, not precise mean). -- [ ] **Step 4: Accessors on `FusedTrainingCtx`** +Use `compute_param_sizes` + `padded_byte_offset` from Task 0.4 to compute `mag_slice_start/len` and `dir_slice_start/len` (branch 0 = direction, tensors 8-11; branch 1 = magnitude, tensors 12-15). + +Document the approach in the Task 2.0 commit message: "pinned-memory snapshots + in-graph reduction kernel (no non-captured replay, no cross-stream sync refactor; matches Task 0.4 pattern)". + +- [ ] **Step 4: Host-side readout at epoch boundary + accessors on `FusedTrainingCtx`** + +At epoch boundary (same block that runs `per_branch_grad_norms` for `grad_ratio_mag_dir`), refresh the cached `grad_component_norms_mag` / `grad_component_norms_dir` arrays by reading the pinned 8-float slot. This is zero-copy: the pinned memory was populated by the in-graph reduction kernel on the last step of the epoch. ```rust +impl GpuDqnTrainer { + /// Refresh cached per-component magnitude/direction norms from the + /// in-graph reduction kernel's pinned result slot. Epoch-boundary only. + pub fn refresh_grad_component_norms(&mut self) -> Result<(), MLError> { + // Stream-sync to ensure last-step reduction kernel has completed. + self.stream.synchronize() + .map_err(|e| MLError::ModelError(format!("grad_component_norms sync: {e}")))?; + let host_slice: &[f32] = unsafe { + std::slice::from_raw_parts(self.grad_component_norms_pinned_ptr, 8) + }; + // Layout: [iqn_mag, iqn_dir, cql_mag, cql_dir, c51_mag, c51_dir, ens_mag, ens_dir] + for i in 0..4 { + self.grad_component_norms_mag[i] = host_slice[2 * i]; + self.grad_component_norms_dir[i] = host_slice[2 * i + 1]; + } + Ok(()) + } +} + impl FusedTrainingCtx { - pub fn grad_mag_iqn_ratio(&mut self) -> f32 { - let m = self.trainer().grad_component_norms_mag[0]; - let d = self.trainer().grad_component_norms_dir[0]; + pub fn grad_mag_iqn_ratio(&self) -> f32 { + let t = self.trainer(); + let m = t.grad_component_norms_mag[0]; + let d = t.grad_component_norms_dir[0]; if d > 1e-9 { m / d } else { 0.0 } } - pub fn grad_mag_cql_ratio(&mut self) -> f32 { /* index 1 */ } - pub fn grad_mag_c51_ratio(&mut self) -> f32 { /* index 2 */ } - pub fn grad_mag_ens_ratio(&mut self) -> f32 { /* index 3 */ } + pub fn grad_mag_cql_ratio(&self) -> f32 { /* index 1 */ } + pub fn grad_mag_c51_ratio(&self) -> f32 { /* index 2 */ } + pub fn grad_mag_ens_ratio(&self) -> f32 { /* index 3 */ } + + pub fn refresh_grad_component_norms(&mut self) -> Result<(), MLError> { + self.trainer_mut().refresh_grad_component_norms() + } } ``` +Call `refresh_grad_component_norms()` in `training_loop.rs` at the existing epoch-end block, before HEALTH_DIAG emission. Same invocation pattern as `update_q_mag_means_cached` (Task 0.3 remediation). + - [ ] **Step 5: Extend HEALTH_DIAG** In `crates/ml/src/trainers/dqn/trainer/training_loop.rs` at the existing HEALTH_DIAG `info!(…)` block (~line 2160), append a new field group: