gitadmin/foxhunt
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

diag+docs(dqn): trunk-slice grad decomposition + stale IQN-trunk doc fix

Browse Source 
Extends grad_decomp_kernel to snapshot the trunk tensor slice (tensors
0..4 = w_s1, b_s1, w_s2, b_s2) in addition to the existing direction +
magnitude branch slices (8..12 / 12..16). Adds a new HEALTH_DIAG group:

  grad_trunk [iqn=<abs> ens=<abs> c51=<abs> cql=<abs> distill=<abs>
              rec=<abs> pred=<abs> cql_sx=<abs> c51_bs=<abs>]

Prior grad_split_bwd / grad_split_aux groups report mag_norm / dir_norm
ratios per loss component, computed over branch-head tensors only. That
measurement range structurally reports 0.0000 for any loss component
that writes exclusively to the trunk — IQN and Ens in particular. This
caused the persistent misdiagnosis that IQN-to-trunk was not wired; the
prior scoping in /tmp/foxhunt_research/iqn-to-trunk-wiring-scoping.md
confirmed the wiring is live (apply_iqn_trunk_gradient at
gpu_dqn_trainer.rs:4882) and that the zero reading was a blind spot in
the measurement pipeline.

Smoke confirms the diagnostic: after iqn_readiness ramps up (late
epochs), grad_trunk reports iqn=100..381 (real trunk SAXPY amplitude),
ens=0.07..3.57, c51=2.46..8.91 (value-head dueling path contributes
through trunk), while cql/cql_sx/distill/rec/pred stay near-zero — a
clean diagnostic baseline.

Also fixes stale documentation at dual-distributional-c51-iqn-design.md
that claimed "IQN trains in isolation — its gradients don't flow back
to the shared trunk": reworded to reflect current wired state with
file:function citation and explicit iqn_readiness gating note. Updated
the "What Changes" table ("IQN training") and "Implementation Order"
(Phase 1 marked DONE) with the same citation.

Changes:
  - grad_decomp_kernel.cu: per-component result slot 2 → 3 floats
    (mag_norm, dir_norm, trunk_norm); extra __shared__ sum_trunk +
    tree reduction; new grad_trunk_start/trunk_len kernel args.
  - gpu_dqn_trainer.rs: pinned result buffer 18 → 27 floats; snapshot
    now does two copy_f32 passes (trunk → dst[0..trunk_len), branch →
    dst[trunk_len..]); per-component slot offsets 0/2/… → 0/3/…;
    grad_component_norms_trunk cached field + accessor; compute trunk
    range from padded_byte_offset(&param_sizes, 0..4).
  - fused_training.rs: grad_trunk_norms_by_component() + per-component
    grad_trunk_*_abs() accessors.
  - training_loop.rs: HEALTH_DIAG emits new grad_trunk group ordered
    [iqn ens c51 cql distill rec pred cql_sx c51_bs]; extended doc
    comment explaining the three groups' roles.
  - design spec (Problem #1 + What Changes row + Implementation Order):
    stale "IQN trains in isolation" replaced by current wired-state
    description, cites gpu_dqn_trainer.rs:4882 and readiness ramp at
    gpu_dqn_trainer.rs:4228-4243.

Pure diagnostic — no training dynamics change, no atomicAdd, no tuning
knobs, no TF32 changes. Smokes unaffected (magnitude_distribution H10
regression pre-exists on HEAD 810b3c570; 4 other smokes pass).

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-04-22 22:03:20 +02:00
parent 810b3c5703
commit b8cd4e1d94
5 changed files with 300 additions and 113 deletions
Download Patch File Download Diff File Expand all files Collapse all files

254
crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs
View File

@@ -1061,9 +1061,19 @@ pub struct GpuDqnTrainer {
/// - `grad_snapshot_rec` — before `launch_recursive_confidence_backward`
/// - `grad_snapshot_pred` — before `compute_predictive_coding_loss`
///
/// Total device memory: 9 × branch01_elems × 4 bytes (~90 MB on a model
/// with ~TOTAL_PARAMS ≈ 2.5 M branch01 elems — 2.2 % of the 4 GB
/// RTX 3050 Ti frame buffer, acceptable for diagnostic).
/// Task 2.Z diagnostic extension: each snapshot now covers the TRUNK
/// slice (tensors 0..4 — w_s1, b_s1, w_s2, b_s2) in addition to the
/// branch-head slice. Layout per snapshot:
/// [0 .. trunk_len) = trunk
/// [trunk_len .. trunk_len+dir_len) = direction
/// [trunk_len+dir_len .. trunk_len+dir_len+mag_len) = magnitude
/// The trunk slot exists to make IQN (`apply_iqn_trunk_gradient`) and
/// Ens (`apply_ensemble_diversity_backward`) contributions measurable
/// — both SAXPY exclusively into the trunk, which the prior branch-
/// only pipeline reported as `0.0000` regardless of the real delta.
///
/// Total device memory: 9 × (trunk+branch01)_elems × 4 bytes — a
/// small bump on top of the pre-extension snapshot size.
grad_snapshot_iqn: CudaSlice<f32>, // [branch01_elems] — snapshot before IQN trunk grad
grad_snapshot_cql: CudaSlice<f32>, // [branch01_elems] — snapshot before CQL block (gradient+saxpy)
grad_snapshot_cql_sx: CudaSlice<f32>, // [branch01_elems] — snapshot before apply_cql_saxpy only
@@ -1073,25 +1083,31 @@ pub struct GpuDqnTrainer {
grad_snapshot_rec: CudaSlice<f32>, // [branch01_elems] — snapshot before launch_recursive_confidence_backward
grad_snapshot_pred: CudaSlice<f32>, // [branch01_elems] — snapshot before compute_predictive_coding_loss
grad_snapshot_ens: CudaSlice<f32>, // [branch01_elems] — snapshot before ensemble diversity
/// Pinned device-mapped 18-float result slot written by the reduction
/// kernel. Layout (matches plan Task 2.0 extension Step 4, index order
/// chosen to match host-cache `[iqn, cql, cql_sx, c51, c51_bs, distill,
/// rec, pred, ens]`):
/// [ 0]=iqn_mag [ 1]=iqn_dir
/// [ 2]=cql_mag [ 3]=cql_dir
/// [ 4]=cql_sx_mag [ 5]=cql_sx_dir
/// [ 6]=c51_mag [ 7]=c51_dir
/// [ 8]=c51_bs_mag [ 9]=c51_bs_dir
/// [10]=distill_mag [11]=distill_dir
/// [12]=rec_mag [13]=rec_dir
/// [14]=pred_mag [15]=pred_dir
/// [16]=ens_mag [17]=ens_dir
/// Pinned device-mapped 27-float result slot written by the reduction
/// kernel. Layout (Task 2.Z diagnostic — 3 floats per component: mag,
/// dir, trunk — preserving the original mag/dir order at [0]/[1] of
/// each 3-float slot so pre-existing consumers continue to read
/// unchanged offsets). Index order matches host-cache `[iqn, cql,
/// cql_sx, c51, c51_bs, distill, rec, pred, ens]`:
/// [ 0]=iqn_mag [ 1]=iqn_dir [ 2]=iqn_trunk
/// [ 3]=cql_mag [ 4]=cql_dir [ 5]=cql_trunk
/// [ 6]=cql_sx_mag [ 7]=cql_sx_dir [ 8]=cql_sx_trunk
/// [ 9]=c51_mag [10]=c51_dir [11]=c51_trunk
/// [12]=c51_bs_mag [13]=c51_bs_dir [14]=c51_bs_trunk
/// [15]=distill_mag [16]=distill_dir [17]=distill_trunk
/// [18]=rec_mag [19]=rec_dir [20]=rec_trunk
/// [21]=pred_mag [22]=pred_dir [23]=pred_trunk
/// [24]=ens_mag [25]=ens_dir [26]=ens_trunk
grad_decomp_result_pinned: *mut f32,
grad_decomp_result_dev_ptr: u64,
/// Cached grad-decomp slice metadata (element indices into `grad_buf`).
/// `grad_decomp_trunk_start`/`trunk_len` cover trunk tensors 0..4
/// (w_s1, b_s1, w_s2, b_s2) — Task 2.Z diagnostic addition.
/// `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_snapshot_len` = trunk_len + dir_len + mag_len.
grad_decomp_trunk_start: i32,
grad_decomp_trunk_len: i32,
grad_decomp_dir_start: i32,
grad_decomp_dir_len: i32,
grad_decomp_mag_start: i32,
@@ -1104,6 +1120,14 @@ pub struct GpuDqnTrainer {
/// accessors: `[iqn, cql, cql_sx, c51, c51_bs, distill, rec, pred, ens]`.
grad_component_norms_mag: [f32; 9],
grad_component_norms_dir: [f32; 9],
/// Task 2.Z diagnostic — cached per-component TRUNK slice L2 norms
/// populated by `refresh_grad_component_norms`. Index order matches
/// mag/dir: `[iqn, cql, cql_sx, c51, c51_bs, distill, rec, pred, ens]`.
/// Non-zero for components that actually write into trunk tensors 0..4
/// (IQN via `apply_iqn_trunk_gradient`, Ens via
/// `apply_ensemble_diversity_backward`); near-zero for components that
/// write only to branch heads (C51, CQL, etc.).
grad_component_norms_trunk: [f32; 9],
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)
@@ -2387,32 +2411,53 @@ impl GpuDqnTrainer {
Ok((n[0], n[1])) // dir, mag
}
/// 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).
/// Task 2.0 — in-graph DtoD snapshot of `grad_buf`'s trunk + branch 0+1
/// slices 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`.
/// Two non-contiguous ranges in `grad_buf` are copied into one dst
/// buffer laid out as `[trunk | dir | mag]` (matching the kernel's
/// snapshot layout convention documented on
/// `grad_component_delta_norm`):
/// * Trunk slice (tensors 0..4 — w_s1, b_s1, w_s2, b_s2) →
/// `dst[0 .. trunk_len)` (Task 2.Z diagnostic — captures IQN +
/// Ens contributions that SAXPY exclusively into the trunk).
/// * Branch slice (tensors 8..16 — direction + magnitude heads) →
/// `dst[trunk_len .. trunk_len + dir_len + mag_len)`. The branch
/// slice is a single contiguous element range in `grad_buf` —
/// tensors 12..16 (magnitude) directly follow 8..12 (direction)
/// in the flat parameter layout per `compute_param_sizes`.
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)
// ── Trunk copy (tensors 0..4) → dst[0..trunk_len) ─────────────────
{
let src_ptr = self.ptrs.grad_buf + (self.grad_decomp_trunk_start as u64) * elem_size;
let dst_ptr = dst.raw_ptr();
let n_bytes = (self.grad_decomp_trunk_len as usize) * (elem_size as usize);
self.graph_safe_copy_f32(dst_ptr, src_ptr, n_bytes, ctx)?;
}
// ── Branch copy (tensors 8..16, contiguous dir+mag) → dst[trunk..] ─
{
let src_ptr = self.ptrs.grad_buf + (self.grad_decomp_dir_start as u64) * elem_size;
let dst_ptr = dst.raw_ptr() + (self.grad_decomp_trunk_len as u64) * elem_size;
let branch_elems = (self.grad_decomp_dir_len + self.grad_decomp_mag_len) as usize;
let n_bytes = branch_elems * (elem_size as usize);
self.graph_safe_copy_f32(dst_ptr, src_ptr, n_bytes, ctx)?;
}
Ok(())
}
/// 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.
/// snapshot, writing (mag_norm, dir_norm, trunk_norm) into the pinned
/// result slot at the component's 3-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`).
/// `slot_offset_elems` is the element offset into the 27-float pinned
/// slot — 0=iqn, 3=cql, 6=cql_sx, 9=c51, 12=c51_bs, 15=distill, 18=rec,
/// 21=pred, 24=ens (Task 2.Z extension: 3 floats per component; matches
/// pinned-layout convention documented on `grad_decomp_result_pinned`).
pub(crate) fn launch_grad_decomp(
&self,
snapshot: &CudaSlice<f32>,
@@ -2432,6 +2477,8 @@ impl GpuDqnTrainer {
self.stream.launch_builder(&self.grad_decomp_kernel)
.arg(&current_ptr)
.arg(&snapshot_ptr)
.arg(&self.grad_decomp_trunk_start)
.arg(&self.grad_decomp_trunk_len)
.arg(&self.grad_decomp_dir_start)
.arg(&self.grad_decomp_dir_len)
.arg(&self.grad_decomp_mag_start)
@@ -2486,62 +2533,64 @@ impl GpuDqnTrainer {
/// Task 2.0 — convenience: launch the reduction kernel for each component
/// with its pre-configured snapshot buffer and pinned-slot offset.
///
/// Pinned-slot offsets (2 floats per component — mag then dir):
/// iqn = 0, cql = 2, cql_sx = 4,
/// c51 = 6, c51_bs = 8, distill = 10,
/// rec = 12, pred = 14, ens = 16.
/// Pinned-slot offsets (3 floats per component — mag, dir, trunk —
/// Task 2.Z extension):
/// iqn = 0, cql = 3, cql_sx = 6,
/// c51 = 9, c51_bs = 12, distill = 15,
/// rec = 18, pred = 21, ens = 24.
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")
self.launch_grad_decomp(&self.grad_snapshot_cql, 3, "cql")
}
pub(crate) fn grad_decomp_launch_cql_sx(&self) -> Result<(), MLError> {
self.launch_grad_decomp(&self.grad_snapshot_cql_sx, 4, "cql_sx")
self.launch_grad_decomp(&self.grad_snapshot_cql_sx, 6, "cql_sx")
}
/// 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.
/// Kernel computes `‖grad_buf − 0‖` on trunk + branch 0+1 slices.
pub(crate) fn grad_decomp_launch_c51(&self) -> Result<(), MLError> {
self.launch_grad_decomp(&self.grad_zero_ref_buf, 6, "c51")
self.launch_grad_decomp(&self.grad_zero_ref_buf, 9, "c51")
}
pub(crate) fn grad_decomp_launch_c51_bs(&self) -> Result<(), MLError> {
self.launch_grad_decomp(&self.grad_snapshot_c51_bs, 8, "c51_bs")
self.launch_grad_decomp(&self.grad_snapshot_c51_bs, 12, "c51_bs")
}
pub(crate) fn grad_decomp_launch_distill(&self) -> Result<(), MLError> {
self.launch_grad_decomp(&self.grad_snapshot_distill, 10, "distill")
self.launch_grad_decomp(&self.grad_snapshot_distill, 15, "distill")
}
pub(crate) fn grad_decomp_launch_rec(&self) -> Result<(), MLError> {
self.launch_grad_decomp(&self.grad_snapshot_rec, 12, "rec")
self.launch_grad_decomp(&self.grad_snapshot_rec, 18, "rec")
}
pub(crate) fn grad_decomp_launch_pred(&self) -> Result<(), MLError> {
self.launch_grad_decomp(&self.grad_snapshot_pred, 14, "pred")
self.launch_grad_decomp(&self.grad_snapshot_pred, 21, "pred")
}
pub(crate) fn grad_decomp_launch_ens(&self) -> Result<(), MLError> {
self.launch_grad_decomp(&self.grad_snapshot_ens, 16, "ens")
self.launch_grad_decomp(&self.grad_snapshot_ens, 24, "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.
/// Task 2.0 — populate the host-side `grad_component_norms_mag/_dir/
/// _trunk` 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). Task 2.0 extension grows this from 8 → 18
/// floats to cover 9 measurement points:
/// [ 0]=iqn_mag [ 1]=iqn_dir
/// [ 2]=cql_mag [ 3]=cql_dir
/// [ 4]=cql_sx_mag [ 5]=cql_sx_dir
/// [ 6]=c51_mag [ 7]=c51_dir
/// [ 8]=c51_bs_mag [ 9]=c51_bs_dir
/// [10]=distill_mag [11]=distill_dir
/// [12]=rec_mag [13]=rec_dir
/// [14]=pred_mag [15]=pred_dir
/// [16]=ens_mag [17]=ens_dir
/// Pinned layout (written by `grad_component_delta_norm`, 3 floats per
/// component — mag, dir, trunk — Task 2.Z extension grows this from
/// 18 → 27 floats to cover the trunk slice for all 9 measurement
/// points):
/// [ 0]=iqn_mag [ 1]=iqn_dir [ 2]=iqn_trunk
/// [ 3]=cql_mag [ 4]=cql_dir [ 5]=cql_trunk
/// [ 6]=cql_sx_mag [ 7]=cql_sx_dir [ 8]=cql_sx_trunk
/// [ 9]=c51_mag [10]=c51_dir [11]=c51_trunk
/// [12]=c51_bs_mag [13]=c51_bs_dir [14]=c51_bs_trunk
/// [15]=distill_mag [16]=distill_dir [17]=distill_trunk
/// [18]=rec_mag [19]=rec_dir [20]=rec_trunk
/// [21]=pred_mag [22]=pred_dir [23]=pred_trunk
/// [24]=ens_mag [25]=ens_dir [26]=ens_trunk
///
/// Host caches use `[iqn, cql, cql_sx, c51, c51_bs, distill, rec, pred,
/// ens]` index order matching the `grad_mag_{iqn,cql,cql_sx,c51,c51_bs,
@@ -2556,11 +2605,12 @@ impl GpuDqnTrainer {
MLError::ModelError(format!("grad_component_norms sync: {e}"))
})?;
let host_slice: &[f32] = unsafe {
std::slice::from_raw_parts(self.grad_decomp_result_pinned, 18)
std::slice::from_raw_parts(self.grad_decomp_result_pinned, 27)
};
for comp in 0..9_usize {
self.grad_component_norms_mag[comp] = host_slice[2 * comp];
self.grad_component_norms_dir[comp] = host_slice[2 * comp + 1];
self.grad_component_norms_mag[comp] = host_slice[3 * comp];
self.grad_component_norms_dir[comp] = host_slice[3 * comp + 1];
self.grad_component_norms_trunk[comp] = host_slice[3 * comp + 2];
}
Ok(())
}
@@ -2572,6 +2622,14 @@ impl GpuDqnTrainer {
/// Task 2.0 — accessor for cached per-component direction norms
/// populated by `refresh_grad_component_norms`.
pub fn grad_component_norms_dir(&self) -> [f32; 9] { self.grad_component_norms_dir }
/// Task 2.Z diagnostic — accessor for cached per-component TRUNK
/// (tensors 0..4 = w_s1, b_s1, w_s2, b_s2) L2 norms populated by
/// `refresh_grad_component_norms`. Index order:
/// `[iqn, cql, cql_sx, c51, c51_bs, distill, rec, pred, ens]`.
/// Expected non-zero for IQN + Ens (both SAXPY into trunk); expected
/// near-zero for C51/CQL/CQL-SX/C51-BS/distill/rec/pred (all write to
/// branch heads only).
pub fn grad_component_norms_trunk(&self) -> [f32; 9] { self.grad_component_norms_trunk }
/// Task 0.6 — per-branch target/online parameter drift (H8 detection signal).
///
@@ -6325,25 +6383,44 @@ impl GpuDqnTrainer {
// 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) = {
// Task 2.Z diagnostic: each snapshot now covers three slices of
// `grad_buf`:
// * Trunk (tensors 0..4 — w_s1, b_s1, w_s2, b_s2). Captures
// `apply_iqn_trunk_gradient` + `apply_ensemble_diversity_backward`
// SAXPY contributions that were structurally invisible to the
// prior branch-only measurement pipeline.
// * Direction (branch 0, tensors 8..12).
// * Magnitude (branch 1, tensors 12..16).
//
// Trunk tensors are at the start of the flat parameter layout
// (offset 0). Direction and magnitude are 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 bounds stay in lock-step with the flat layout.
let (grad_decomp_trunk_start, grad_decomp_trunk_len,
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(&param_sizes, 8) as usize;
let mag_start_byte = padded_byte_offset(&param_sizes, 12) as usize;
let trunk_start_byte = padded_byte_offset(&param_sizes, 0) as usize;
let trunk_end_byte = padded_byte_offset(&param_sizes, 4) as usize;
let dir_start_byte = padded_byte_offset(&param_sizes, 8) as usize;
let mag_start_byte = padded_byte_offset(&param_sizes, 12) as usize;
let branch2_start_byte = padded_byte_offset(&param_sizes, 16) as usize;
let elem_size = std::mem::size_of::<f32>();
let trunk_start = trunk_start_byte / elem_size;
let trunk_end = trunk_end_byte / elem_size;
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 mag_end = branch2_start_byte / elem_size;
let trunk_len = trunk_end - trunk_start;
let dir_len = mag_start - dir_start;
let mag_len = mag_end - mag_start;
let snap_len = trunk_len + dir_len + mag_len;
(trunk_start as i32, trunk_len as i32,
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}")))?;
@@ -6372,18 +6449,19 @@ impl GpuDqnTrainer {
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 18-float result slot (Task 2.0 extension grows
// from 8 → 18 floats for 9 measurement components). Device writes via
// `grad_decomp_result_dev_ptr` (captured in graph); host reads via
// `grad_decomp_result_pinned` at epoch boundary (zero-copy).
// Pinned device-mapped 27-float result slot (Task 2.Z extension
// grows from 18 → 27 floats — 3 floats per component × 9
// components; the third slot is the trunk L2 norm). 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(18 * std::mem::size_of::<f32>(), flags)
cudarc::driver::result::malloc_host(27 * 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..18 { *grad_decomp_result_pinned.add(i) = 0.0; }
for i in 0..27 { *grad_decomp_result_pinned.add(i) = 0.0; }
}
let grad_decomp_result_dev_ptr = unsafe {
let mut dp = 0u64;
@@ -6404,8 +6482,9 @@ impl GpuDqnTrainer {
.map_err(|e| MLError::ModelError(format!("grad_component_delta_norm load: {e}")))?
};
info!(
"GpuDqnTrainer: grad_decomp kernel loaded — snapshot_len={} dir=[{}..+{}] mag=[{}..+{}]",
"GpuDqnTrainer: grad_decomp kernel loaded — snapshot_len={} trunk=[{}..+{}] dir=[{}..+{}] mag=[{}..+{}]",
grad_decomp_snapshot_len,
grad_decomp_trunk_start, grad_decomp_trunk_len,
grad_decomp_dir_start, grad_decomp_dir_len,
grad_decomp_mag_start, grad_decomp_mag_len,
);
@@ -8373,6 +8452,8 @@ impl GpuDqnTrainer {
grad_snapshot_ens,
grad_decomp_result_pinned,
grad_decomp_result_dev_ptr,
grad_decomp_trunk_start,
grad_decomp_trunk_len,
grad_decomp_dir_start,
grad_decomp_dir_len,
grad_decomp_mag_start,
@@ -8381,6 +8462,7 @@ impl GpuDqnTrainer {
grad_decomp_kernel,
grad_component_norms_mag: [0.0_f32; 9],
grad_component_norms_dir: [0.0_f32; 9],
grad_component_norms_trunk: [0.0_f32; 9],
params_buf,
target_params_buf,
m_buf,

70
crates/ml/src/cuda_pipeline/grad_decomp_kernel.cu
View File

@@ -1,15 +1,29 @@
/**
* 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.
* Computes L2 norms of (current − snapshot) restricted to three slices of
* `grad_buf` and writes three f32 results to a pinned device-mapped output
* slot:
* result_out[0] = ‖Δ‖ on the magnitude branch slice (tensors 12..16)
* result_out[1] = ‖Δ‖ on the direction branch slice (tensors 8..12)
* result_out[2] = ‖Δ‖ on the trunk slice (tensors 0..4 — w_s1, b_s1,
* w_s2, b_s2)
*
* 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.
* Launched once per loss component (9 measurement points: IQN / CQL /
* CQL-SAXPY / C51 / C51-budget-scale / Distillation / Recursive-confidence
* / Predictive-coding / Ensemble) — 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 three byte ranges
* so that `current − snapshot` isolates exactly the component's additive
* delta.
*
* The trunk slice was added (Task 2.Z diagnostic) so the measurement
* pipeline can report IQN and Ens contributions — both `apply_iqn_trunk_
* gradient` and `apply_ensemble_diversity_backward` SAXPY exclusively into
* the trunk tensors 0..4, which sit entirely OUTSIDE the branch-head range
* measured by the prior two-slot kernel. Without this third slot, IQN
* and Ens always report 0.0000 on HEALTH_DIAG regardless of the actual
* contribution magnitude.
*
* All reductions are block-internal via shared-memory tree reduction —
* NO atomicAdd (Task 0.5 / 0.8 no-atomics rule on the captured path).
@@ -17,42 +31,54 @@
* 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
* trunk_len + dir_len + mag_len), laid out as:
* [0 .. trunk_len) = trunk slice
* [trunk_len .. trunk_len+dir_len) = direction slice
* [trunk_len+dir_len .. trunk_len+dir_len+mag_len) = magnitude slice
* indexed from zero; the caller is responsible for
* having copied the corresponding grad_buf range into
* having copied the corresponding grad_buf ranges into
* this buffer BEFORE the component's backward fired.
* grad_dir_start/len, grad_mag_start/len
* grad_trunk_start/len, 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.
* slice (so the kernel can read both sides of the
* subtraction from the correct places).
* result_out — [3] output: (mag_norm, dir_norm, trunk_norm) — mag/dir
* order chosen to match the grad_ratio_mag_dir
* convention; trunk appended last so existing pinned-
* slot consumers continue to read mag at [0] and dir
* at [1] unchanged.
*
* Launch config: one block, 256 threads.
*/
extern "C" __global__ void grad_component_delta_norm(
const float* __restrict__ current,
const float* __restrict__ snapshot,
int grad_trunk_start, int trunk_len,
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];
__shared__ float sum_trunk[256];
int tid = threadIdx.x;
sum_mag[tid] = 0.0f;
sum_dir[tid] = 0.0f;
sum_trunk[tid] = 0.0f;
// Direction slice — snapshot offset starts at 0.
// Trunk slice — snapshot offset starts at 0.
for (int i = tid; i < trunk_len; i += blockDim.x) {
float d = current[grad_trunk_start + i] - snapshot[i];
sum_trunk[tid] += d * d;
}
// Direction slice — snapshot offset starts at trunk_len.
for (int i = tid; i < dir_len; i += blockDim.x) {
float d = current[grad_dir_start + i] - snapshot[i];
float d = current[grad_dir_start + i] - snapshot[trunk_len + i];
sum_dir[tid] += d * d;
}
// Magnitude slice — snapshot offset starts at dir_len.
// Magnitude slice — snapshot offset starts at trunk_len + dir_len.
for (int i = tid; i < mag_len; i += blockDim.x) {
float d = current[grad_mag_start + i] - snapshot[dir_len + i];
float d = current[grad_mag_start + i] - snapshot[trunk_len + dir_len + i];
sum_mag[tid] += d * d;
}
__syncthreads();
@@ -62,11 +88,13 @@ extern "C" __global__ void grad_component_delta_norm(
if (tid < s) {
sum_mag[tid] += sum_mag[tid + s];
sum_dir[tid] += sum_dir[tid + s];
sum_trunk[tid] += sum_trunk[tid + s];
}
__syncthreads();
}
if (tid == 0) {
result_out[0] = sqrtf(sum_mag[0]);
result_out[1] = sqrtf(sum_dir[0]);
result_out[2] = sqrtf(sum_trunk[0]);
}
}

26
crates/ml/src/trainers/dqn/fused_training.rs
View File

@@ -996,6 +996,32 @@ impl FusedTrainingCtx {
if d > 1e-9 { m / d } else { 0.0 }
}
/// Task 2.Z diagnostic — absolute trunk-slice L2 norms per component
/// populated by `refresh_grad_component_norms`. Index order matches
/// the mag/dir accessors:
/// `[iqn, cql, cql_sx, c51, c51_bs, distill, rec, pred, ens]`.
///
/// Values report the L2 norm of the additive delta each component's
/// backward / scaling op applied to trunk tensors 0..4 (w_s1, b_s1,
/// w_s2, b_s2). Expected non-zero on `[iqn]` (via
/// `apply_iqn_trunk_gradient`) and `[ens]` (via
/// `apply_ensemble_diversity_backward`) whenever their readiness /
/// scale factors are non-zero; expected near-zero on all others
/// (they write exclusively to branch heads 8..16).
pub(crate) fn grad_trunk_norms_by_component(&self) -> [f32; 9] {
self.trainer.grad_component_norms_trunk()
}
/// Task 2.Z diagnostic — per-component absolute trunk L2 norm accessors.
pub(crate) fn grad_trunk_iqn_abs(&self) -> f32 { self.trainer.grad_component_norms_trunk()[0] }
pub(crate) fn grad_trunk_cql_abs(&self) -> f32 { self.trainer.grad_component_norms_trunk()[1] }
pub(crate) fn grad_trunk_cql_sx_abs(&self) -> f32 { self.trainer.grad_component_norms_trunk()[2] }
pub(crate) fn grad_trunk_c51_abs(&self) -> f32 { self.trainer.grad_component_norms_trunk()[3] }
pub(crate) fn grad_trunk_c51_bs_abs(&self) -> f32 { self.trainer.grad_component_norms_trunk()[4] }
pub(crate) fn grad_trunk_distill_abs(&self) -> f32 { self.trainer.grad_component_norms_trunk()[5] }
pub(crate) fn grad_trunk_rec_abs(&self) -> f32 { self.trainer.grad_component_norms_trunk()[6] }
pub(crate) fn grad_trunk_pred_abs(&self) -> f32 { self.trainer.grad_component_norms_trunk()[7] }
pub(crate) fn grad_trunk_ens_abs(&self) -> f32 { self.trainer.grad_component_norms_trunk()[8] }
/// 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).

43
crates/ml/src/trainers/dqn/trainer/training_loop.rs
View File

@@ -2274,6 +2274,9 @@ impl DQNTrainer {
grad_mag_c51, grad_mag_c51_bs,
grad_mag_distill, grad_mag_rec, grad_mag_pred,
grad_mag_ens,
// Task 2.Z — absolute trunk-slice L2 norms per component
// (iqn, cql, cql_sx, c51, c51_bs, distill, rec, pred, ens).
grad_trunk_norms,
) = if let Some(fused) = self.fused_ctx.as_mut() {
if let Err(e) = fused.refresh_grad_component_norms() {
tracing::warn!(
@@ -2290,10 +2293,12 @@ impl DQNTrainer {
fused.grad_mag_rec_ratio(),
fused.grad_mag_pred_ratio(),
fused.grad_mag_ens_ratio(),
fused.grad_trunk_norms_by_component(),
)
} else {
(0.0_f32, 0.0_f32, 0.0_f32, 0.0_f32, 0.0_f32,
0.0_f32, 0.0_f32, 0.0_f32, 0.0_f32)
0.0_f32, 0.0_f32, 0.0_f32, 0.0_f32,
[0.0_f32; 9])
};
// Task 0.6 — per-branch NoisyNets σ mean. H7 detection signal.
@@ -2671,7 +2676,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}] grad_split_bwd [iqn={:.4} cql={:.4} c51={:.4} ens={:.4}] grad_split_aux [distill={:.4} rec={:.4} pred={:.4} cql_sx={:.4} c51_bs={:.4}] grad_abs [dir={:.6e} mag={:.6e}] trail [fire_q={:.3} fire_h={:.3} fire_f={:.3} hold_q={:.2} hold_h={:.2} hold_f={:.2}] mag_stats [wr_q={:.3} wr_h={:.3} wr_f={:.3} var_q={:.4e} var_h={:.4e} var_f={:.4e}] 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}] 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_bwd [iqn={:.4} cql={:.4} c51={:.4} ens={:.4}] grad_split_aux [distill={:.4} rec={:.4} pred={:.4} cql_sx={:.4} c51_bs={:.4}] grad_trunk [iqn={:.4} ens={:.4} c51={:.4} cql={:.4} distill={:.4} rec={:.4} pred={:.4} cql_sx={:.4} c51_bs={:.4}] grad_abs [dir={:.6e} mag={:.6e}] trail [fire_q={:.3} fire_h={:.3} fire_f={:.3} hold_q={:.2} hold_h={:.2} hold_f={:.2}] mag_stats [wr_q={:.3} wr_h={:.3} wr_f={:.3} var_q={:.4e} var_h={:.4e} var_f={:.4e}] 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}] 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,
@@ -2721,7 +2726,9 @@ impl DQNTrainer {
//
// grad_split_bwd [iqn cql c51 ens]
// The four original Task 2.0 measurement points — full
// backward-gradient contributions from each loss family.
// backward-gradient contributions from each loss family,
// reported as mag_norm / dir_norm ratios over branch
// tensors 8..16 only.
//
// grad_split_aux [distill rec pred cql_sx c51_bs]
// The five Task 2.0 extension points inside the aux-graph
@@ -2731,12 +2738,42 @@ impl DQNTrainer {
// saxpy; c51_bs = apply_c51_budget_scale). Localizes
// exactly which stage zeroes magnitude gradient.
//
// grad_trunk [iqn ens c51 cql distill rec pred cql_sx c51_bs]
// Task 2.Z diagnostic extension — absolute L2 norms of
// each component's additive delta on TRUNK tensors 0..4
// (w_s1, b_s1, w_s2, b_s2). IQN and Ens are the two
// components wired to SAXPY into the trunk (via
// `apply_iqn_trunk_gradient` + `apply_ensemble_
// diversity_backward`); all other components write
// exclusively to branch heads and should report
// near-zero here. Required because grad_split_bwd /
// grad_split_aux structurally measure branch heads
// only, so IQN + Ens reported `0.0000` in those groups
// regardless of their actual trunk contribution.
//
// 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,
grad_mag_distill, grad_mag_rec, grad_mag_pred,
grad_mag_cql_sx, grad_mag_c51_bs,
// Task 2.Z diagnostic — absolute TRUNK-slice L2 norms per
// loss component (tensors 0..4 = w_s1, b_s1, w_s2, b_s2).
// Ordered as [iqn, ens, c51, cql, distill, rec, pred, cql_sx,
// c51_bs] — IQN + Ens first because they are the two
// components that actually SAXPY into trunk (via
// `apply_iqn_trunk_gradient` at gpu_dqn_trainer.rs:4882 and
// `apply_ensemble_diversity_backward` at 5051); the remaining
// seven slots are near-zero baselines (C51/CQL/etc. write
// exclusively to branch heads 8..16) and are included so the
// measurement is symmetric across the same 9 components
// tracked by grad_split_bwd + grad_split_aux.
//
// grad_trunk_norms index order matches the mag/dir caches:
// [iqn, cql, cql_sx, c51, c51_bs, distill, rec, pred, ens].
grad_trunk_norms[0], grad_trunk_norms[8], grad_trunk_norms[3],
grad_trunk_norms[1], grad_trunk_norms[5], grad_trunk_norms[6],
grad_trunk_norms[7], grad_trunk_norms[2], grad_trunk_norms[4],
// Task 2.0 confirmation — absolute per-branch grad L2 norms
// (dir, mag) in scientific notation. Disambiguates the
// `grad_ratio_mag_dir` threshold collapse (0.0 when dir < 1e-9).

20
docs/superpowers/specs/2026-03-23-dual-distributional-c51-iqn-design.md
View File

@@ -3,7 +3,17 @@
## Problem
The system has both C51 (101 atoms) and IQN (32 quantiles) implemented, but:
1. IQN trains in isolation — its gradients don't flow back to the shared trunk
1. **[Phase 1 DONE]** IQN's backward gradient is now wired into the shared
trunk via `apply_iqn_trunk_gradient` (`crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs:4882-5034`),
captured unconditionally inside the aux child graph from
`crates/ml/src/trainers/dqn/fused_training.rs:1582-1587`. The effective
scale is `iqn_lambda * iqn_readiness * iqn_budget`, where
`iqn_readiness` is a loss-improvement-driven ramp (0→1) at
`gpu_dqn_trainer.rs:4228-4243`. Trunk contribution is measurable via
the Task 2.Z `grad_trunk [iqn=...]` HEALTH_DIAG group — the previous
`grad_split_bwd [iqn=0.0000]` reading was a measurement blind spot in
the branch-only grad_decomp snapshot range, not evidence of missing
wiring.
2. IQN's quantile output is never used for action selection or risk management
3. `use_cvar_action_selection` exists in config but is hardcoded to `false`
4. The trunk learns only from C51's cross-entropy loss, missing the risk signal from IQN
@@ -56,7 +66,7 @@ Wire the existing IQN dual-head for two purposes:
| Component | Today | After | Effort |
|-----------|-------|-------|--------|
| C51 training | Working | No change | 0 |
| IQN training | Working, disconnected | Gradient flows to trunk | LOW |
| IQN training | Working, trunk gradient wired (Phase 1 landed — `apply_iqn_trunk_gradient`) | No further change | DONE |
| IQN inference | Kernel compiled, never called | Wired for CVaR | MEDIUM |
| CVaR computation | Not implemented | New trivial CUDA kernel | LOW |
| C51/IQN disagreement | Not implemented | |E_c51 - E_iqn| | LOW |
@@ -67,7 +77,11 @@ Wire the existing IQN dual-head for two purposes:
## Implementation Order
1. Wire IQN gradient to shared trunk (Phase 1)
1. ~~Wire IQN gradient to shared trunk (Phase 1)~~ — **DONE**
(`apply_iqn_trunk_gradient` at `gpu_dqn_trainer.rs:4882-5034`,
scheduled from `fused_training.rs:1582-1587`, scale
`iqn_lambda * iqn_readiness * iqn_budget`). Measurable via the
Task 2.Z `grad_trunk [iqn=...]` HEALTH_DIAG group.
2. Wire IQN inference + CVaR computation (Phase 2)
3. Add disagreement uncertainty signal (Phase 3)
4. Hyperopt integration for cvar_alpha (Phase 4)