plan(policy-quality): Task 2.0 revised approach — in-graph pinned snapshots
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).
This commit is contained in:
jgrusewski
@@ -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<f32>, // [TOTAL_PARAMS]
|
||||
grad_snapshot_cql: CudaSlice<f32>,
|
||||
grad_snapshot_c51: CudaSlice<f32>,
|
||||
grad_snapshot_ens: CudaSlice<f32>,
|
||||
|
||||
/// 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::<f32>(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::<f32>(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:
|
||||
|
||||
Reference in New Issue
Block a user