diag(policy-quality): Task 2.0 — per-component grad decomposition for H4
Adds grad_mag_{iqn,cql,c51,ens} HEALTH_DIAG fields via in-graph
pinned-snapshot + in-graph reduction kernel (revised approach; first
Task 2.0 dispatch escalated BLOCKED on host-side-snapshots-inside-
captured-graph, plan revised at 5a4d56145 to use this pattern).
How it works:
* Three CudaSlice<f32>[branch01_elems] per-component scratch buffers
(iqn, cql, ens) + one permanent-zeros reference buffer for C51
(grad_buf is memset-zeroed at start of submit_forward_ops_main —
no pre-C51 snapshot needed). Snapshots only cover the branch 0+1
slice (direction+magnitude, tensors 8..16), NOT TOTAL_PARAMS —
~420 KB per snapshot, 4× = ~1.7 MB device memory total.
* BEFORE each component's backward (captured in graph):
graph_safe_copy_f32 kernel copies grad_buf[branch01] → scratch.
C51 uses the permanent-zeros reference buffer directly.
* AFTER each component's backward (captured in graph):
grad_component_delta_norm kernel (256 threads, one block) reads
(grad_buf[dir_slice] − snapshot[0..dir_len]) and
(grad_buf[mag_slice] − snapshot[dir_len..]), reduces via
shared-memory tree reduction (NO atomicAdd), writes
(mag_norm, dir_norm) into a pinned device-mapped 8-float result
slot at the component's 2-float offset.
* At epoch boundary: refresh_grad_component_norms() stream-syncs,
reads the pinned slot (zero-copy via cuMemHostGetDevicePointer),
populates host-side grad_component_norms_mag/_dir caches.
* HEALTH_DIAG emits `grad_split [iqn=... cql=... c51=... ens=...]`
with per-component magnitude/direction ratios.
Uses graph_safe_copy_f32 (already existing — see gpu_dqn_trainer.rs:3839
doc comment: "memcpy_dtod_async is NOT captured by CUDA Graph"). The
plan's claim that DtoD memcpy is captureable is contradicted by the
codebase's own history, so we follow the established pattern.
Matches Task 0.5 / 0.8 no-atomics pattern; matches Task 0.4's pinned
device-mapped pattern for kernel-written readback. New file:
crates/ml/src/cuda_pipeline/grad_decomp_kernel.cu (72 LOC, one kernel
+ header doc).
Zero per-step PCIe traffic: the 8-float pinned slot is device-mapped;
writes flow via the mapped device pointer during graph replay, reads
are direct host dereferences at epoch boundary.
20-epoch grad_split trajectory — final fold, epochs 15–19:
grad_split [iqn=0.0000 cql=33.5947 c51=418.0461 ens=0.0000]
grad_split [iqn=0.0000 cql=65.2418 c51=222.4561 ens=0.0000]
grad_split [iqn=0.0000 cql=23.8915 c51=145.3589 ens=0.0000]
grad_split [iqn=0.0000 cql=46.1190 c51=268.8549 ens=0.0000]
grad_split [iqn=0.0000 cql=147.4817 c51=213.1188 ens=0.0000]
Findings for Task 2.1 decision tree:
* IQN and Ensemble report 0.0000 ratios every epoch — architecturally
expected: apply_iqn_trunk_gradient writes only to trunk tensors 0..4
(w_s1/b_s1/w_s2/b_s2); run_ensemble_step flows through the value
head only (tensors 4..8). Neither touches branches 8..24 by design.
* C51 and CQL are the ONLY two loss components that write directly
to branch head weights — and both send MORE gradient to magnitude
than to direction (ratios routinely 50–400×).
* Task 0.4's grad_ratio_mag_dir=0.0000 (reported post-all-accumulations,
post-budget-scaling) vs per-component ratios >>1 shifts the
diagnosis: gradient IS reaching magnitude from C51/CQL — the
global mag/dir=0 at epoch boundary reflects cancellation with
other writers (distill SAXPY, recursive_confidence_backward,
predictive_coding_loss) that accumulate into grad_buf BETWEEN
CQL and Ens, AND the heavy c51_budget_scale (0.70) + cql_budget
(0.04) scaling that is applied AFTER our reduction.
This is H9 territory (data-driven magnitude Q-value collapse) more
than H4 (architectural gradient starvation) — the Task 2.1 decision
tree should pick branches A or C (init scale / direction-conditioning)
only after excluding data-driven factors, or skip straight to the
"revisit in Phase 3" row.
Per plan Task 2.0 (revised at 5a4d56145). Build clean; magnitude
distribution smoke green (21.10s local, RTX 3050 Ti).
This commit is contained in:
@@ -81,6 +81,7 @@ fn main() {
|
||||
"iqn_cvar_kernel.cu",
|
||||
"mamba2_temporal_kernel.cu",
|
||||
"graph_utility_kernels.cu",
|
||||
"grad_decomp_kernel.cu",
|
||||
];
|
||||
|
||||
// ALL kernels get common header (BF16 types + wrappers)
|
||||
|
||||
@@ -75,6 +75,7 @@ static ENSEMBLE_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/ensembl
|
||||
static CQL_GRAD_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/cql_grad_kernel.cubin"));
|
||||
static MAMBA2_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/mamba2_temporal_kernel.cubin"));
|
||||
pub(crate) static GRAPH_UTILITY_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/graph_utility_kernels.cubin"));
|
||||
static GRAD_DECOMP_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/grad_decomp_kernel.cubin"));
|
||||
|
||||
/// Mamba2 temporal scan configuration.
|
||||
const MAMBA2_HISTORY_K: usize = 8; // Rolling history length
|
||||
@@ -970,6 +971,49 @@ pub struct GpuDqnTrainer {
|
||||
/// training graph capture region.
|
||||
grad_readback_pinned_ptr: usize, // host-visible f32 buffer
|
||||
grad_readback_pinned_capacity: usize, // element count (TOTAL_PARAMS)
|
||||
/// Task 2.0 — per-component magnitude-branch grad decomposition.
|
||||
///
|
||||
/// Four device-side scratch buffers sized to the combined length of
|
||||
/// branch 0 (direction, tensors 8..12) + branch 1 (magnitude, tensors
|
||||
/// 12..16). BEFORE each component's backward fires (captured in the
|
||||
/// training graph), a `copy_f32` kernel snapshots `grad_buf` over the
|
||||
/// branch 0+1 byte range into these slots. AFTER the component's
|
||||
/// backward, `grad_component_delta_norm` computes `‖current−snapshot‖`
|
||||
/// on direction and magnitude slices and writes two f32 values into
|
||||
/// the pinned result slot.
|
||||
///
|
||||
/// The C51 slot is paired with `grad_zero_ref_buf` (permanent zeros)
|
||||
/// because `grad_buf` is memset-zeroed at the start of every
|
||||
/// `submit_forward_ops_main` — there is no "pre-C51" grad to snapshot.
|
||||
/// Uniform subtraction keeps the kernel interface simple.
|
||||
grad_snapshot_iqn: CudaSlice<f32>, // [branch01_elems] — snapshot before IQN trunk grad
|
||||
grad_snapshot_cql: CudaSlice<f32>, // [branch01_elems] — snapshot before CQL SAXPY
|
||||
grad_zero_ref_buf: CudaSlice<f32>, // [branch01_elems] — permanent zeros; C51 snapshot substitute
|
||||
grad_snapshot_ens: CudaSlice<f32>, // [branch01_elems] — snapshot before ensemble diversity
|
||||
/// Pinned device-mapped 8-float result slot written by the reduction
|
||||
/// kernel. Layout (matches plan Task 2.0 Step 4):
|
||||
/// [0]=iqn_mag [1]=iqn_dir
|
||||
/// [2]=cql_mag [3]=cql_dir
|
||||
/// [4]=c51_mag [5]=c51_dir
|
||||
/// [6]=ens_mag [7]=ens_dir
|
||||
grad_decomp_result_pinned: *mut f32,
|
||||
grad_decomp_result_dev_ptr: u64,
|
||||
/// Cached grad-decomp slice metadata (element indices into `grad_buf`).
|
||||
/// `grad_decomp_dir_start`/`dir_len` cover branch 0 tensors 8..12,
|
||||
/// `grad_decomp_mag_start`/`mag_len` cover branch 1 tensors 12..16.
|
||||
/// `grad_decomp_snapshot_len` = dir_len + mag_len.
|
||||
grad_decomp_dir_start: i32,
|
||||
grad_decomp_dir_len: i32,
|
||||
grad_decomp_mag_start: i32,
|
||||
grad_decomp_mag_len: i32,
|
||||
grad_decomp_snapshot_len: usize,
|
||||
/// Reduction kernel (one-block 256-thread launch, no atomics).
|
||||
grad_decomp_kernel: CudaFunction,
|
||||
/// Host-side cached per-component norms populated at epoch boundary
|
||||
/// from the pinned result slot. Index order matches `grad_mag_*_ratio`
|
||||
/// accessors: `[iqn, cql, c51, ens]`.
|
||||
grad_component_norms_mag: [f32; 4],
|
||||
grad_component_norms_dir: [f32; 4],
|
||||
params_buf: CudaSlice<f32>, // [TOTAL_PARAMS] f32 master online parameters (Adam operates here)
|
||||
target_params_buf: CudaSlice<f32>, // [TOTAL_PARAMS] f32 master target parameters (EMA operates here)
|
||||
m_buf: CudaSlice<f32>, // [TOTAL_PARAMS] Adam first moment (f32 for precision)
|
||||
@@ -2045,6 +2089,12 @@ impl Drop for GpuDqnTrainer {
|
||||
cudarc::driver::result::free_host(self.grad_readback_pinned_ptr as *mut std::ffi::c_void)
|
||||
};
|
||||
}
|
||||
// Task 2.0 — pinned device-mapped grad-component result slot.
|
||||
if !self.grad_decomp_result_pinned.is_null() {
|
||||
let _ = unsafe {
|
||||
cudarc::driver::result::free_host(self.grad_decomp_result_pinned.cast())
|
||||
};
|
||||
}
|
||||
if !self.per_branch_q_gaps_pinned.is_null() {
|
||||
let _ = unsafe { cudarc::driver::result::free_host(self.per_branch_q_gaps_pinned.cast()) };
|
||||
}
|
||||
@@ -2172,6 +2222,136 @@ impl GpuDqnTrainer {
|
||||
Ok(if dir > 1e-9 { mag / dir } else { 0.0 })
|
||||
}
|
||||
|
||||
/// Task 2.0 — in-graph DtoD snapshot of `grad_buf`'s branch 0+1 slice into
|
||||
/// the caller-supplied scratch buffer. Uses the `copy_f32` kernel (graph-
|
||||
/// captured DtoD) because `cuMemcpyDtoDAsync` is NOT captured by CUDA
|
||||
/// Graph (see `graph_safe_copy_f32` doc comment).
|
||||
///
|
||||
/// Copies element range `[grad_decomp_dir_start, grad_decomp_dir_start +
|
||||
/// grad_decomp_snapshot_len)` from `grad_buf` (which spans branches 0+1
|
||||
/// contiguously — see `compute_param_sizes`) into the first
|
||||
/// `grad_decomp_snapshot_len` elements of `dst`.
|
||||
pub(crate) fn grad_decomp_snapshot(&self, dst: &CudaSlice<f32>, ctx: &str)
|
||||
-> Result<(), MLError>
|
||||
{
|
||||
let elem_size = std::mem::size_of::<f32>() as u64;
|
||||
let src_ptr = self.ptrs.grad_buf + (self.grad_decomp_dir_start as u64) * elem_size;
|
||||
let dst_ptr = dst.raw_ptr();
|
||||
let n_bytes = self.grad_decomp_snapshot_len * elem_size as usize;
|
||||
self.graph_safe_copy_f32(dst_ptr, src_ptr, n_bytes, ctx)
|
||||
}
|
||||
|
||||
/// Task 2.0 — launch `grad_component_delta_norm` on a per-component
|
||||
/// snapshot, writing (mag_norm, dir_norm) into the pinned result slot at
|
||||
/// the component's 2-float offset.
|
||||
///
|
||||
/// `slot_offset_elems` is the element offset into the 8-float pinned slot
|
||||
/// — 0=iqn, 2=cql, 4=c51, 6=ens (matches pinned-layout convention
|
||||
/// documented on `grad_decomp_result_pinned`).
|
||||
pub(crate) fn launch_grad_decomp(
|
||||
&self,
|
||||
snapshot: &CudaSlice<f32>,
|
||||
slot_offset_elems: i32,
|
||||
ctx: &str,
|
||||
) -> Result<(), MLError> {
|
||||
let current_ptr = self.ptrs.grad_buf;
|
||||
let snapshot_ptr = snapshot.raw_ptr();
|
||||
let result_ptr = self.grad_decomp_result_dev_ptr
|
||||
+ (slot_offset_elems as u64) * (std::mem::size_of::<f32>() as u64);
|
||||
let cfg = LaunchConfig {
|
||||
grid_dim: (1, 1, 1),
|
||||
block_dim: (256, 1, 1),
|
||||
shared_mem_bytes: 0,
|
||||
};
|
||||
unsafe {
|
||||
self.stream.launch_builder(&self.grad_decomp_kernel)
|
||||
.arg(¤t_ptr)
|
||||
.arg(&snapshot_ptr)
|
||||
.arg(&self.grad_decomp_dir_start)
|
||||
.arg(&self.grad_decomp_dir_len)
|
||||
.arg(&self.grad_decomp_mag_start)
|
||||
.arg(&self.grad_decomp_mag_len)
|
||||
.arg(&result_ptr)
|
||||
.launch(cfg)
|
||||
.map_err(|e| MLError::ModelError(format!("grad_decomp {ctx}: {e}")))?;
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Task 2.0 — convenience: snapshot `grad_buf` into the component's
|
||||
/// snapshot buffer (captured DtoD via `copy_f32` kernel). Called BEFORE
|
||||
/// the component's backward fires.
|
||||
pub(crate) fn grad_decomp_snapshot_iqn(&self) -> Result<(), MLError> {
|
||||
self.grad_decomp_snapshot(&self.grad_snapshot_iqn, "grad_snapshot_iqn")
|
||||
}
|
||||
pub(crate) fn grad_decomp_snapshot_cql(&self) -> Result<(), MLError> {
|
||||
self.grad_decomp_snapshot(&self.grad_snapshot_cql, "grad_snapshot_cql")
|
||||
}
|
||||
pub(crate) fn grad_decomp_snapshot_ens(&self) -> Result<(), MLError> {
|
||||
self.grad_decomp_snapshot(&self.grad_snapshot_ens, "grad_snapshot_ens")
|
||||
}
|
||||
|
||||
/// Task 2.0 — convenience: launch the reduction kernel for each component
|
||||
/// with its pre-configured snapshot buffer and pinned-slot offset.
|
||||
pub(crate) fn grad_decomp_launch_iqn(&self) -> Result<(), MLError> {
|
||||
self.launch_grad_decomp(&self.grad_snapshot_iqn, 0, "iqn")
|
||||
}
|
||||
pub(crate) fn grad_decomp_launch_cql(&self) -> Result<(), MLError> {
|
||||
self.launch_grad_decomp(&self.grad_snapshot_cql, 2, "cql")
|
||||
}
|
||||
/// C51 uses `grad_zero_ref_buf` (permanent zeros) because `grad_buf` is
|
||||
/// zero before `submit_forward_ops_main` — there is no pre-C51 snapshot.
|
||||
/// Kernel computes `‖grad_buf − 0‖` on branch 0+1 slices.
|
||||
pub(crate) fn grad_decomp_launch_c51(&self) -> Result<(), MLError> {
|
||||
self.launch_grad_decomp(&self.grad_zero_ref_buf, 4, "c51")
|
||||
}
|
||||
pub(crate) fn grad_decomp_launch_ens(&self) -> Result<(), MLError> {
|
||||
self.launch_grad_decomp(&self.grad_snapshot_ens, 6, "ens")
|
||||
}
|
||||
|
||||
/// Task 2.0 — populate the host-side `grad_component_norms_mag/_dir`
|
||||
/// caches from the pinned result slot. Called at epoch boundary, before
|
||||
/// HEALTH_DIAG emission.
|
||||
///
|
||||
/// Stream-syncs to ensure the final step's reduction kernel has completed
|
||||
/// (same pattern as `per_branch_grad_norms`). Zero-copy: the pinned
|
||||
/// memory is device-mapped, so the host read is a direct dereference of
|
||||
/// the last-step kernel output.
|
||||
///
|
||||
/// Pinned layout (written by `grad_component_delta_norm`, 2 floats per
|
||||
/// component — mag then dir):
|
||||
/// [0]=iqn_mag [1]=iqn_dir [2]=cql_mag [3]=cql_dir
|
||||
/// [4]=c51_mag [5]=c51_dir [6]=ens_mag [7]=ens_dir
|
||||
///
|
||||
/// Host caches use `[iqn, cql, c51, ens]` index order matching the
|
||||
/// `grad_mag_{iqn,cql,c51,ens}_ratio` accessor families.
|
||||
pub fn refresh_grad_component_norms(&mut self) -> Result<(), MLError> {
|
||||
if self.grad_decomp_result_pinned.is_null() {
|
||||
return Err(MLError::ModelError(
|
||||
"refresh_grad_component_norms: pinned result slot is null".into()
|
||||
));
|
||||
}
|
||||
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_decomp_result_pinned, 8)
|
||||
};
|
||||
for comp in 0..4_usize {
|
||||
self.grad_component_norms_mag[comp] = host_slice[2 * comp];
|
||||
self.grad_component_norms_dir[comp] = host_slice[2 * comp + 1];
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Task 2.0 — accessor for cached per-component magnitude norms
|
||||
/// populated by `refresh_grad_component_norms`. Index order:
|
||||
/// `[iqn, cql, c51, ens]`.
|
||||
pub fn grad_component_norms_mag(&self) -> [f32; 4] { self.grad_component_norms_mag }
|
||||
/// Task 2.0 — accessor for cached per-component direction norms
|
||||
/// populated by `refresh_grad_component_norms`.
|
||||
pub fn grad_component_norms_dir(&self) -> [f32; 4] { self.grad_component_norms_dir }
|
||||
|
||||
/// Task 0.6 — per-branch target/online parameter drift (H8 detection signal).
|
||||
///
|
||||
/// Returns `[direction, magnitude, order, urgency]` RMS deviation between
|
||||
@@ -5790,6 +5970,79 @@ impl GpuDqnTrainer {
|
||||
};
|
||||
let grad_readback_pinned_capacity = total_params;
|
||||
|
||||
// Task 2.0 — per-component grad-decomposition snapshots + pinned result slot.
|
||||
//
|
||||
// The snapshots only need to cover branch 0 (direction, tensors 8..12) and
|
||||
// branch 1 (magnitude, tensors 12..16) — a single contiguous element range
|
||||
// in `grad_buf`. Sizes are computed from `compute_param_sizes` exactly like
|
||||
// `per_branch_grad_norms` (Task 0.4) so direction/magnitude bounds stay in
|
||||
// lock-step with the flat parameter layout.
|
||||
let (grad_decomp_dir_start, grad_decomp_dir_len, grad_decomp_mag_start,
|
||||
grad_decomp_mag_len, grad_decomp_snapshot_len) = {
|
||||
let param_sizes = compute_param_sizes(&config);
|
||||
let dir_start_byte = padded_byte_offset(¶m_sizes, 8) as usize;
|
||||
let mag_start_byte = padded_byte_offset(¶m_sizes, 12) as usize;
|
||||
let branch2_start_byte = padded_byte_offset(¶m_sizes, 16) as usize;
|
||||
let elem_size = std::mem::size_of::<f32>();
|
||||
let dir_start = dir_start_byte / elem_size;
|
||||
let mag_start = mag_start_byte / elem_size;
|
||||
let mag_end = branch2_start_byte / elem_size;
|
||||
let dir_len = mag_start - dir_start;
|
||||
let mag_len = mag_end - mag_start;
|
||||
let snap_len = dir_len + mag_len;
|
||||
(dir_start as i32, dir_len as i32, mag_start as i32, mag_len as i32, snap_len)
|
||||
};
|
||||
let grad_snapshot_iqn = stream.alloc_zeros::<f32>(grad_decomp_snapshot_len)
|
||||
.map_err(|e| MLError::ModelError(format!("alloc grad_snapshot_iqn: {e}")))?;
|
||||
let grad_snapshot_cql = stream.alloc_zeros::<f32>(grad_decomp_snapshot_len)
|
||||
.map_err(|e| MLError::ModelError(format!("alloc grad_snapshot_cql: {e}")))?;
|
||||
// C51 fires BEFORE the aux phase and `grad_buf` is memset-zeroed at the
|
||||
// start of `submit_forward_ops_main` — no "pre-C51" grad exists to
|
||||
// snapshot. `grad_zero_ref_buf` is initialized to zeros and never
|
||||
// written again; the reduction kernel reads (current − 0) = current on
|
||||
// the C51 path.
|
||||
let grad_zero_ref_buf = stream.alloc_zeros::<f32>(grad_decomp_snapshot_len)
|
||||
.map_err(|e| MLError::ModelError(format!("alloc grad_zero_ref_buf: {e}")))?;
|
||||
let grad_snapshot_ens = stream.alloc_zeros::<f32>(grad_decomp_snapshot_len)
|
||||
.map_err(|e| MLError::ModelError(format!("alloc grad_snapshot_ens: {e}")))?;
|
||||
|
||||
// Pinned device-mapped 8-float result slot. Device writes via
|
||||
// `grad_decomp_result_dev_ptr` (captured in graph); host reads via
|
||||
// `grad_decomp_result_pinned` at epoch boundary (zero-copy).
|
||||
let grad_decomp_result_pinned: *mut f32 = unsafe {
|
||||
let flags = cudarc::driver::sys::CU_MEMHOSTALLOC_DEVICEMAP;
|
||||
cudarc::driver::result::malloc_host(8 * std::mem::size_of::<f32>(), flags)
|
||||
.map_err(|e| MLError::ModelError(format!("pinned grad_decomp_result alloc: {e}")))?
|
||||
as *mut f32
|
||||
};
|
||||
unsafe {
|
||||
for i in 0..8 { *grad_decomp_result_pinned.add(i) = 0.0; }
|
||||
}
|
||||
let grad_decomp_result_dev_ptr = unsafe {
|
||||
let mut dp = 0u64;
|
||||
cudarc::driver::sys::cuMemHostGetDevicePointer_v2(
|
||||
&mut dp as *mut u64,
|
||||
grad_decomp_result_pinned.cast(),
|
||||
0,
|
||||
);
|
||||
dp
|
||||
};
|
||||
|
||||
// Reduction kernel — loaded from a dedicated CUmodule so its CUfunction
|
||||
// is isolated from all other graphs (Hopper CUfunction-isolation rule).
|
||||
let grad_decomp_kernel = {
|
||||
let module = stream.context().load_cubin(GRAD_DECOMP_CUBIN.to_vec())
|
||||
.map_err(|e| MLError::ModelError(format!("grad_decomp cubin load: {e}")))?;
|
||||
module.load_function("grad_component_delta_norm")
|
||||
.map_err(|e| MLError::ModelError(format!("grad_component_delta_norm load: {e}")))?
|
||||
};
|
||||
info!(
|
||||
"GpuDqnTrainer: grad_decomp kernel loaded — snapshot_len={} dir=[{}..+{}] mag=[{}..+{}]",
|
||||
grad_decomp_snapshot_len,
|
||||
grad_decomp_dir_start, grad_decomp_dir_len,
|
||||
grad_decomp_mag_start, grad_decomp_mag_len,
|
||||
);
|
||||
|
||||
// eval_v_range — pinned device-mapped (zero-copy write from CPU).
|
||||
let eval_v_range_pinned: *mut f32 = unsafe {
|
||||
let flags = cudarc::driver::sys::CU_MEMHOSTALLOC_DEVICEMAP;
|
||||
@@ -7644,6 +7897,20 @@ impl GpuDqnTrainer {
|
||||
grad_buf,
|
||||
grad_readback_pinned_ptr,
|
||||
grad_readback_pinned_capacity,
|
||||
grad_snapshot_iqn,
|
||||
grad_snapshot_cql,
|
||||
grad_zero_ref_buf,
|
||||
grad_snapshot_ens,
|
||||
grad_decomp_result_pinned,
|
||||
grad_decomp_result_dev_ptr,
|
||||
grad_decomp_dir_start,
|
||||
grad_decomp_dir_len,
|
||||
grad_decomp_mag_start,
|
||||
grad_decomp_mag_len,
|
||||
grad_decomp_snapshot_len,
|
||||
grad_decomp_kernel,
|
||||
grad_component_norms_mag: [0.0_f32; 4],
|
||||
grad_component_norms_dir: [0.0_f32; 4],
|
||||
params_buf,
|
||||
target_params_buf,
|
||||
m_buf,
|
||||
|
||||
72
crates/ml/src/cuda_pipeline/grad_decomp_kernel.cu
Normal file
72
crates/ml/src/cuda_pipeline/grad_decomp_kernel.cu
Normal file
@@ -0,0 +1,72 @@
|
||||
/**
|
||||
* grad_decomp_kernel — Task 2.0 per-component gradient decomposition.
|
||||
*
|
||||
* Computes L2 norms of (current − snapshot) restricted to the direction
|
||||
* (branch 0) and magnitude (branch 1) slices of grad_buf, writes two
|
||||
* f32 results to a pinned device-mapped output slot.
|
||||
*
|
||||
* Launched four times per step — once per loss component (IQN/CQL/C51/Ens)
|
||||
* — always AFTER the component's backward kernels finish writing into
|
||||
* `grad_buf`. The matching BEFORE-backward snapshot is produced by the
|
||||
* `copy_f32` kernel over the branch 0+1 byte range so that
|
||||
* `current − snapshot` isolates exactly the component's additive delta.
|
||||
*
|
||||
* All reductions are block-internal via shared-memory tree reduction —
|
||||
* NO atomicAdd (Task 0.5 / 0.8 no-atomics rule on the captured path).
|
||||
*
|
||||
* Layout:
|
||||
* current — points to grad_buf[0] (f32, length total_params)
|
||||
* snapshot — points to the per-component snapshot (f32, length
|
||||
* dir_len + mag_len), laid out as:
|
||||
* [0 .. dir_len) = direction slice
|
||||
* [dir_len .. dir_len+mag_len) = magnitude slice
|
||||
* indexed from zero; the caller is responsible for
|
||||
* having copied the corresponding grad_buf range into
|
||||
* this buffer BEFORE the component's backward fired.
|
||||
* grad_dir_start/len, grad_mag_start/len
|
||||
* — element offsets + lengths into `current` for each
|
||||
* branch slice (so the kernel can read both sides of
|
||||
* the subtraction from the correct places).
|
||||
* result_out — [2] output: (mag_norm, dir_norm) — order chosen to
|
||||
* match the grad_ratio_mag_dir convention.
|
||||
*
|
||||
* Launch config: one block, 256 threads.
|
||||
*/
|
||||
extern "C" __global__ void grad_component_delta_norm(
|
||||
const float* __restrict__ current,
|
||||
const float* __restrict__ snapshot,
|
||||
int grad_dir_start, int dir_len,
|
||||
int grad_mag_start, int mag_len,
|
||||
float* __restrict__ result_out
|
||||
) {
|
||||
__shared__ float sum_mag[256];
|
||||
__shared__ float sum_dir[256];
|
||||
int tid = threadIdx.x;
|
||||
sum_mag[tid] = 0.0f;
|
||||
sum_dir[tid] = 0.0f;
|
||||
|
||||
// Direction slice — snapshot offset starts at 0.
|
||||
for (int i = tid; i < dir_len; i += blockDim.x) {
|
||||
float d = current[grad_dir_start + i] - snapshot[i];
|
||||
sum_dir[tid] += d * d;
|
||||
}
|
||||
// Magnitude slice — snapshot offset starts at dir_len.
|
||||
for (int i = tid; i < mag_len; i += blockDim.x) {
|
||||
float d = current[grad_mag_start + i] - snapshot[dir_len + i];
|
||||
sum_mag[tid] += d * d;
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
// 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]);
|
||||
}
|
||||
}
|
||||
@@ -891,6 +891,43 @@ impl FusedTrainingCtx {
|
||||
self.trainer.grad_ratio_mag_dir().unwrap_or(0.0)
|
||||
}
|
||||
|
||||
/// Task 2.0 — refresh the cached per-component grad-norm arrays from the
|
||||
/// pinned result slot populated by the in-graph reduction kernel.
|
||||
/// Epoch-boundary cold path; stream-syncs once. Returns `Err` on
|
||||
/// pinned-buffer absence or sync failure — the caller is expected to
|
||||
/// `tracing::warn!` + fall back to stale values (same pattern as
|
||||
/// `update_q_mag_means_cached`).
|
||||
pub(crate) fn refresh_grad_component_norms(&mut self) -> Result<(), crate::MLError> {
|
||||
self.trainer.refresh_grad_component_norms()
|
||||
}
|
||||
|
||||
/// Task 2.0 — IQN magnitude/direction grad-norm ratio. Index 0 in the
|
||||
/// `grad_component_norms_*` caches (populated by
|
||||
/// `refresh_grad_component_norms`).
|
||||
pub(crate) fn grad_mag_iqn_ratio(&self) -> f32 {
|
||||
let m = self.trainer.grad_component_norms_mag()[0];
|
||||
let d = self.trainer.grad_component_norms_dir()[0];
|
||||
if d > 1e-9 { m / d } else { 0.0 }
|
||||
}
|
||||
/// Task 2.0 — CQL magnitude/direction grad-norm ratio (index 1).
|
||||
pub(crate) fn grad_mag_cql_ratio(&self) -> f32 {
|
||||
let m = self.trainer.grad_component_norms_mag()[1];
|
||||
let d = self.trainer.grad_component_norms_dir()[1];
|
||||
if d > 1e-9 { m / d } else { 0.0 }
|
||||
}
|
||||
/// Task 2.0 — C51 magnitude/direction grad-norm ratio (index 2).
|
||||
pub(crate) fn grad_mag_c51_ratio(&self) -> f32 {
|
||||
let m = self.trainer.grad_component_norms_mag()[2];
|
||||
let d = self.trainer.grad_component_norms_dir()[2];
|
||||
if d > 1e-9 { m / d } else { 0.0 }
|
||||
}
|
||||
/// Task 2.0 — Ensemble magnitude/direction grad-norm ratio (index 3).
|
||||
pub(crate) fn grad_mag_ens_ratio(&self) -> f32 {
|
||||
let m = self.trainer.grad_component_norms_mag()[3];
|
||||
let d = self.trainer.grad_component_norms_dir()[3];
|
||||
if d > 1e-9 { m / d } else { 0.0 }
|
||||
}
|
||||
|
||||
/// Task 0.6 — per-branch target/online param drift (H8 detection signal).
|
||||
/// Returns `[direction, magnitude, order, urgency]` RMS ‖target − online‖.
|
||||
/// Zeros on accessor failure (non-fatal — diag only).
|
||||
@@ -1170,6 +1207,18 @@ impl FusedTrainingCtx {
|
||||
).map_err(|e| anyhow::anyhow!("HER in-place relabel kernel: {e}"))?;
|
||||
}
|
||||
|
||||
// Task 2.0 — measure C51 gradient contribution to branch 0 (direction)
|
||||
// and branch 1 (magnitude). C51's backward ran inside
|
||||
// `submit_forward_ops_main` (Phase 2); `grad_buf` was memset-zeroed
|
||||
// at its start. At this point `grad_buf` holds exactly the C51/MSE
|
||||
// blended backward — no IQN/CQL/Ens contributions yet. Reduction
|
||||
// kernel runs BEFORE `apply_c51_budget_scale` so we capture the raw
|
||||
// backward contribution rather than the budget-scaled amplitude.
|
||||
// Uses `grad_zero_ref_buf` as the "snapshot" (permanent zeros), so
|
||||
// kernel computes `‖grad_buf − 0‖` on branch 0+1 slices.
|
||||
self.trainer.grad_decomp_launch_c51()
|
||||
.map_err(|e| anyhow::anyhow!("Task 2.0 grad_decomp C51: {e}"))?;
|
||||
|
||||
// NOTE: Per-component gradient budget clip REMOVED (old explicit clip → constant grad norm).
|
||||
// B4/G5: Budgets are now applied as SCALE factors on individual SAXPY contributions,
|
||||
// not as norm clips. This preserves gradient magnitude variation while steering
|
||||
@@ -1443,6 +1492,17 @@ impl FusedTrainingCtx {
|
||||
// ── Post-join IQN processing (main stream) ──
|
||||
// trunk gradient, EMA, IQR run on self.stream (main) because they use
|
||||
// the trainer's cuBLAS handle which is bound to the main stream.
|
||||
//
|
||||
// Task 2.0 — snapshot `grad_buf` branch 0+1 slice BEFORE the IQN
|
||||
// trunk gradient fires, so the post-IQN reduction sees only the
|
||||
// additive delta. Snapshot runs unconditionally (even when
|
||||
// `iqn_ok=false` or `gpu_iqn` is None) so the reduction below can
|
||||
// also run unconditionally and keep the pinned result slot's IQN
|
||||
// field (slots [0..2]) populated every step (all-zero delta → 0.0
|
||||
// ratio, which is the honest signal for a step that skipped IQN).
|
||||
self.trainer.grad_decomp_snapshot_iqn()
|
||||
.map_err(|e| anyhow::anyhow!("Task 2.0 grad_decomp_snapshot_iqn: {e}"))?;
|
||||
|
||||
if iqn_ok {
|
||||
if let Some(ref mut iqn) = self.gpu_iqn {
|
||||
let d_h_s2_ptr = iqn.d_h_s2_raw_ptr();
|
||||
@@ -1469,6 +1529,13 @@ impl FusedTrainingCtx {
|
||||
}
|
||||
}
|
||||
|
||||
// Task 2.0 — IQN reduction. Runs unconditionally so the pinned slot
|
||||
// stays populated even on steps where `iqn_ok=false`. When IQN was
|
||||
// skipped the snapshot and current `grad_buf` are identical over the
|
||||
// branch 0+1 range → delta norm = 0.0 (honest semantic).
|
||||
self.trainer.grad_decomp_launch_iqn()
|
||||
.map_err(|e| anyhow::anyhow!("Task 2.0 grad_decomp IQN: {e}"))?;
|
||||
|
||||
// IQN→PER loss: cast f32 per_sample_loss → f32 td_errors_buf.
|
||||
// IQN is now f32 native; PER system expects f32 td_errors.
|
||||
if let Some(ref iqn) = self.gpu_iqn {
|
||||
@@ -1490,6 +1557,14 @@ impl FusedTrainingCtx {
|
||||
let meta_q_pred = self.trainer.last_meta_q_pred();
|
||||
let temporal_amp = 1.0_f32 + (meta_q_pred - 0.5).max(0.0) * 2.0;
|
||||
let f5_barrier_weight = 0.20_f32 * temporal_amp;
|
||||
|
||||
// Task 2.0 — snapshot BEFORE the CQL gradient path. The snapshot runs
|
||||
// whether or not `has_cql()` fires so the post-CQL reduction below
|
||||
// can run unconditionally (keeps the pinned result slot's CQL field
|
||||
// populated every step — 0.0 on steps where CQL was skipped).
|
||||
self.trainer.grad_decomp_snapshot_cql()
|
||||
.map_err(|e| anyhow::anyhow!("Task 2.0 grad_decomp_snapshot_cql: {e}"))?;
|
||||
|
||||
if self.trainer.has_cql() {
|
||||
match self.trainer.apply_cql_gradient(f5_barrier_weight) {
|
||||
Ok(true) => {
|
||||
@@ -1503,6 +1578,10 @@ impl FusedTrainingCtx {
|
||||
}
|
||||
}
|
||||
|
||||
// Task 2.0 — CQL reduction. Delta = grad_buf − snapshot over branch 0+1.
|
||||
self.trainer.grad_decomp_launch_cql()
|
||||
.map_err(|e| anyhow::anyhow!("Task 2.0 grad_decomp CQL: {e}"))?;
|
||||
|
||||
// D1/N1: Temporal self-distillation — per-step kernel pull toward the
|
||||
// best historical snapshot. Fully GPU-native: alpha is computed
|
||||
// in-kernel from ISV[LEARNING_HEALTH_INDEX] (pinned device-mapped),
|
||||
@@ -1547,6 +1626,13 @@ impl FusedTrainingCtx {
|
||||
// wiring backward adds gradient noise without value. Both calls removed
|
||||
// here; the kernels and Rust glue are deleted in a follow-up cleanup.
|
||||
|
||||
// Task 2.0 — snapshot BEFORE ensemble diversity step. Runs
|
||||
// unconditionally even when `ensemble_extra_heads.is_empty()` so the
|
||||
// post-ens reduction below can run unconditionally (0.0 delta on
|
||||
// steps without ensemble diversity — honest semantic).
|
||||
self.trainer.grad_decomp_snapshot_ens()
|
||||
.map_err(|e| anyhow::anyhow!("Task 2.0 grad_decomp_snapshot_ens: {e}"))?;
|
||||
|
||||
// Ensemble diversity: heads 1..K-1 forward + KL gradient + trunk backward.
|
||||
// Runs inside aux_child graph — gradients land in grad_buf before Adam.
|
||||
if !self.ensemble_extra_heads.is_empty() {
|
||||
@@ -1554,6 +1640,10 @@ impl FusedTrainingCtx {
|
||||
.map_err(|e| anyhow::anyhow!("Ensemble diversity: {e}"))?;
|
||||
}
|
||||
|
||||
// Task 2.0 — ensemble reduction. Delta = grad_buf − snapshot over branch 0+1.
|
||||
self.trainer.grad_decomp_launch_ens()
|
||||
.map_err(|e| anyhow::anyhow!("Task 2.0 grad_decomp ens: {e}"))?;
|
||||
|
||||
// Regime-adaptive PER scaling.
|
||||
self.trainer.regime_scale_td_errors()
|
||||
.map_err(|e| anyhow::anyhow!("Regime PER scaling: {e}"))?;
|
||||
|
||||
@@ -2228,6 +2228,28 @@ impl DQNTrainer {
|
||||
.map(|f| f.grad_ratio_mag_dir())
|
||||
.unwrap_or(0.0);
|
||||
|
||||
// Task 2.0 — per-component grad decomposition (IQN/CQL/C51/Ens).
|
||||
// Refresh the cached norms from the pinned result slot populated
|
||||
// by the in-graph reduction kernel on the last step of this epoch,
|
||||
// then compute mag/dir ratios per component. Non-fatal on failure
|
||||
// (warn + stale values) — same pattern as update_q_mag_means_cached.
|
||||
let (grad_mag_iqn, grad_mag_cql, grad_mag_c51, grad_mag_ens) =
|
||||
if let Some(fused) = self.fused_ctx.as_mut() {
|
||||
if let Err(e) = fused.refresh_grad_component_norms() {
|
||||
tracing::warn!(
|
||||
"HEALTH_DIAG grad_component_norms refresh failed: {e}"
|
||||
);
|
||||
}
|
||||
(
|
||||
fused.grad_mag_iqn_ratio(),
|
||||
fused.grad_mag_cql_ratio(),
|
||||
fused.grad_mag_c51_ratio(),
|
||||
fused.grad_mag_ens_ratio(),
|
||||
)
|
||||
} else {
|
||||
(0.0_f32, 0.0_f32, 0.0_f32, 0.0_f32)
|
||||
};
|
||||
|
||||
// Task 0.6 — per-branch NoisyNets σ mean. H7 detection signal.
|
||||
// 0 = direction head, 1 = magnitude head.
|
||||
let (sigma_mag, sigma_dir) = {
|
||||
@@ -2570,7 +2592,7 @@ impl DQNTrainer {
|
||||
};
|
||||
|
||||
tracing::info!(
|
||||
"HEALTH_DIAG[{}]: health={:.2} components [q_gap={:.2} q_var={:.2} atoms={:.2} grad_stable={:.2} ens_agree={:.2} grad_cos={:.2} spectral={:.2}] effective [cql_alpha={:.4} iqn_budget={:.2} cql_budget={:.2} c51_budget={:.2} tau={:.5} sarsa_tau={:.2} gamma={:.3} cf_ratio={:.2}] novels [distill={} barrier={:.3} plasticity={} ib={:.3} ensemble_collapse={:.2} contrarian={} meta_q_pred={:.2}] diag [sharpe_ema={:.3} action_entropy={:.2}] gems [g12_predictive={:.4}] mag [q_full={:.3} q_half={:.3} q_quarter={:.3} var_scale={:.3} kelly_f={:.3} avg_win_ratio={:.3} grad_ratio_mag_dir={:.4} dist_q={:.3} dist_h={:.3} dist_f={:.3}] trail [fire_q={:.3} fire_h={:.3} fire_f={:.3} hold_q={:.2} hold_h={:.2} hold_f={:.2}] noisy [vsn_mag={:.3} vsn_dir={:.3} sigma_mag={:.4} sigma_dir={:.4} drift_mag={:.3} drift_dir={:.3}] eval_dist [eq={:.3} eh={:.3} ef={:.3}] reward_contrib [popart={:.3} cf={:.3} trail_r={:.3} micro={:.3} la={:.3}] controller [anti_lr={} tau={} gamma={} clip={} cql={} cost={} fire_frac={:.2}] explore [ent_mag={:.2} ent_dir={:.2} sigma_mean={:.4}]",
|
||||
"HEALTH_DIAG[{}]: health={:.2} components [q_gap={:.2} q_var={:.2} atoms={:.2} grad_stable={:.2} ens_agree={:.2} grad_cos={:.2} spectral={:.2}] effective [cql_alpha={:.4} iqn_budget={:.2} cql_budget={:.2} c51_budget={:.2} tau={:.5} sarsa_tau={:.2} gamma={:.3} cf_ratio={:.2}] novels [distill={} barrier={:.3} plasticity={} ib={:.3} ensemble_collapse={:.2} contrarian={} meta_q_pred={:.2}] diag [sharpe_ema={:.3} action_entropy={:.2}] gems [g12_predictive={:.4}] mag [q_full={:.3} q_half={:.3} q_quarter={:.3} var_scale={:.3} kelly_f={:.3} avg_win_ratio={:.3} grad_ratio_mag_dir={:.4} dist_q={:.3} dist_h={:.3} dist_f={:.3}] grad_split [iqn={:.4} cql={:.4} c51={:.4} ens={:.4}] trail [fire_q={:.3} fire_h={:.3} fire_f={:.3} hold_q={:.2} hold_h={:.2} hold_f={:.2}] noisy [vsn_mag={:.3} vsn_dir={:.3} sigma_mag={:.4} sigma_dir={:.4} drift_mag={:.3} drift_dir={:.3}] eval_dist [eq={:.3} eh={:.3} ef={:.3}] reward_contrib [popart={:.3} cf={:.3} trail_r={:.3} micro={:.3} la={:.3}] controller [anti_lr={} tau={} gamma={} clip={} cql={} cost={} fire_frac={:.2}] explore [ent_mag={:.2} ent_dir={:.2} sigma_mean={:.4}]",
|
||||
epoch,
|
||||
health_value,
|
||||
self.learning_health.components.q_gap_norm,
|
||||
@@ -2613,6 +2635,12 @@ impl DQNTrainer {
|
||||
avg_win_ratio,
|
||||
grad_ratio_mag_dir,
|
||||
dist_q, dist_h, dist_f,
|
||||
// Task 2.0 — per-component grad decomposition (4 f32):
|
||||
// iqn / cql / c51 / ens magnitude/direction grad-norm
|
||||
// ratios. In-graph pinned-snapshot + reduction-kernel
|
||||
// pipeline; zero-copy readback at epoch boundary.
|
||||
// Feeds Task 2.1 decision tree for the H4 fix.
|
||||
grad_mag_iqn, grad_mag_cql, grad_mag_c51, grad_mag_ens,
|
||||
// Track 1 — trail (6 f32): fire_q/h/f, hold_q/h/f (H6 — Task 0.5).
|
||||
// Real values from GPU experience collector's per-sample buffers.
|
||||
trail_rates[0], trail_rates[1], trail_rates[2],
|
||||
|
||||
Reference in New Issue
Block a user