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

diag(policy-quality): Task 2.0 extension — instrument 5 more grad writers

Browse Source 
First Task 2.0 pass (commit d60e5375a) covered IQN/CQL/C51/Ens. Result:
IQN+Ens are architecturally zero (don't touch branches); C51+CQL send
30-400× more gradient to magnitude than direction. But Task 0.4's
grad_ratio_mag_dir=0.0000 at epoch end — so magnitude gradient gets
cancelled somewhere in the aux-graph gap between CQL/C51 and epoch end.

This commit instruments the 5 unmeasured writers/scalings:
  - apply_distillation_gradient
  - launch_recursive_confidence_backward
  - compute_predictive_coding_loss
  - apply_c51_budget_scale (multiplicative; signed delta reveals shrinkage)
  - apply_cql_saxpy (multiplicative)

Schema extended 4 → 9 components. Pinned result slot grew 8 → 18 floats.
9 device scratch buffers (~90 MB total on RTX 3050 Ti — 2.2% of 4 GB,
acceptable for diagnostic).

HEALTH_DIAG now emits:
  grad_split_bwd [iqn=… cql=… c51=… ens=…]
  grad_split_aux [distill=… rec=… pred=… cql_sx=… c51_bs=…]

No atomicAdd (block-internal tree reduction). memcpy_dtod_async is not
captureable → used graph_safe_copy_f32 (same as first pass).

Last 5 epochs (HEALTH_DIAG[15..19]), RTX 3050 Ti, magnitude_distribution
smoke test:

  grad_split_bwd [iqn=0.0000 cql=18.4618  c51=102.2001 ens=0.0000]
  grad_split_aux [distill=1.2281 rec=0.0000 pred=0.0000 cql_sx=18.4618  c51_bs=102.2001]

  grad_split_bwd [iqn=0.0000 cql=55.0089  c51=125.1374 ens=0.0000]
  grad_split_aux [distill=1.2381 rec=0.0000 pred=0.0000 cql_sx=55.0089  c51_bs=125.1374]

  grad_split_bwd [iqn=0.0000 cql=30.1931  c51=262.0410 ens=0.0000]
  grad_split_aux [distill=0.7916 rec=0.0000 pred=0.0000 cql_sx=30.1931  c51_bs=262.0410]

  grad_split_bwd [iqn=0.0000 cql=70.2364  c51=126.6637 ens=0.0000]
  grad_split_aux [distill=0.7397 rec=0.0000 pred=0.0000 cql_sx=70.2364  c51_bs=126.6637]

  grad_split_bwd [iqn=0.0000 cql=103.3925 c51=104.4165 ens=0.0000]
  grad_split_aux [distill=0.8085 rec=0.0000 pred=0.0000 cql_sx=103.3925 c51_bs=104.4165]

Keystone findings for Task 2.1:

* rec (recursive confidence) + pred (predictive coding) report 0.0000
  every epoch — architecturally expected: MSE/predictive gradients flow
  through bw_d_h_s2 / trunk tensors 0..8 only, they never touch branch
  tensors 8..16. Same architectural-zero class as IQN+Ens.

* distill is the only aux writer that actually touches branches, but at
  ratios ~0.5–1.2 (nearly balanced mag/dir), so it cannot be the source
  of Task 0.4's grad_ratio_mag_dir=0 cancellation.

* cql_sx ≡ cql and c51_bs ≡ c51 exactly (ratio, not norm). This is
  expected — multiplicative scalings (saxpy budget / c51 budget scale)
  uniformly scale BOTH direction and magnitude slices, so the mag/dir
  RATIO is invariant. These ops change amplitude, not balance.

Conclusion: NONE of the 9 instrumented stages zero the magnitude signal.
CQL+C51 reach grad_buf with huge mag/dir ratios (30-400×); distill adds
a balanced micro-contribution; rec/pred/iqn/ens/scalings all architectur-
ally neutral for the mag/dir ratio. Yet Task 0.4's post-step
grad_ratio_mag_dir=0 (meaning dir < 1e-9).

Re-reading Task 0.4's impl: `per_branch_grad_norms` returns [dir, mag, …]
and the ratio is `mag / dir if dir > 1e-9 else 0.0`. So 0.0000 means
dir_norm is zero, NOT mag cancelled. Task 0.4's measurement runs at
EPOCH boundary after Adam consumes grad_buf — which means the "cancel"
is either (a) Adam consuming/zeroing grad_buf pre-readback, (b) the
next step's `submit_forward_ops_main` memset firing before Task 0.4's
DtoH reads, or (c) something Adam-related that drops dir to zero.

Five unmeasured writers verified present (all 5 exist via grep +
read — none renamed or deleted since the first pass's report). Writers
that DON'T touch branches (rec/pred) are now explicit 0.0000 evidence,
not speculation.

Feeds Task 2.1's decision tree: H4 is NOT caused by cancellation in
the aux-graph gap. Next hypothesis: Adam grad_buf lifecycle / Task 0.4
timing relative to memset.

Build clean; magnitude_distribution smoke green (21.90 s local, RTX 3050 Ti).
Test result: 1 passed; 0 failed.

Push status: deferred — first Task 2.0 pass reported network unreachable
(d60e5375a local-only). Will attempt in same commit cycle; if it fails
again, that's consistent with the prior pass's finding.

Per plan Task 2.0 extension (from Task 2.0 first-pass escalation report).

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-04-22 10:35:53 +02:00
parent d60e5375a9
commit 980f3b07f3
3 changed files with 301 additions and 74 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -971,31 +971,57 @@ 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.
/// Task 2.0 (+ Task 2.0 extension) — per-component magnitude-branch grad
/// decomposition.
///
/// Four device-side scratch buffers sized to the combined length of
/// Nine 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.
/// 12..16). BEFORE each component's op 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 op,
/// `grad_component_delta_norm` computes `‖current−snapshot‖` on direction
/// and magnitude slices and writes two f32 values into the pinned result
/// slot at the component's 2-float offset.
///
/// 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.
/// The C51 (post-backward) 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.
///
/// Task 2.0 extension adds five measurement points in the aux-graph gap
/// between CQL and Ens where Task 0.4's `grad_ratio_mag_dir=0` indicates
/// the magnitude signal gets cancelled:
/// - `grad_snapshot_cql_sx` — before `apply_cql_saxpy`
/// - `grad_snapshot_c51_bs` — before `apply_c51_budget_scale`
/// - `grad_snapshot_distill` — before `apply_distillation_gradient`
/// - `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).
grad_snapshot_iqn: CudaSlice<f32>, // [branch01_elems] — snapshot before IQN trunk grad
grad_snapshot_cql: CudaSlice<f32>, // [branch01_elems] — snapshot before CQL SAXPY
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
grad_zero_ref_buf: CudaSlice<f32>, // [branch01_elems] — permanent zeros; C51 snapshot substitute
grad_snapshot_c51_bs: CudaSlice<f32>, // [branch01_elems] — snapshot before apply_c51_budget_scale
grad_snapshot_distill: CudaSlice<f32>, // [branch01_elems] — snapshot before apply_distillation_gradient
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 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
/// 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
grad_decomp_result_pinned: *mut f32,
grad_decomp_result_dev_ptr: u64,
/// Cached grad-decomp slice metadata (element indices into `grad_buf`).
@@ -1011,9 +1037,9 @@ pub struct GpuDqnTrainer {
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],
/// accessors: `[iqn, cql, cql_sx, c51, c51_bs, distill, rec, pred, ens]`.
grad_component_norms_mag: [f32; 9],
grad_component_norms_dir: [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)
@@ -2280,33 +2306,80 @@ impl GpuDqnTrainer {
/// 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.
/// the component's backward / scaling op 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")
}
/// Task 2.0 extension — snapshot before `apply_cql_saxpy` (isolates the
/// SAXPY step from the CQL-logit-gradient + cuBLAS-backward portion of
/// `apply_cql_gradient`, which writes into `cql_grad_scratch` not
/// `grad_buf`).
pub(crate) fn grad_decomp_snapshot_cql_sx(&self) -> Result<(), MLError> {
self.grad_decomp_snapshot(&self.grad_snapshot_cql_sx, "grad_snapshot_cql_sx")
}
/// Task 2.0 extension — snapshot before `apply_c51_budget_scale`. The
/// scale op multiplies `grad_buf` in place by `c51_budget`, so delta =
/// current − snapshot exposes the amplitude change the scale applied
/// (negative-signed on both dir and mag components when c51_budget < 1).
pub(crate) fn grad_decomp_snapshot_c51_bs(&self) -> Result<(), MLError> {
self.grad_decomp_snapshot(&self.grad_snapshot_c51_bs, "grad_snapshot_c51_bs")
}
/// Task 2.0 extension — snapshot before `apply_distillation_gradient`.
pub(crate) fn grad_decomp_snapshot_distill(&self) -> Result<(), MLError> {
self.grad_decomp_snapshot(&self.grad_snapshot_distill, "grad_snapshot_distill")
}
/// Task 2.0 extension — snapshot before
/// `launch_recursive_confidence_backward`.
pub(crate) fn grad_decomp_snapshot_rec(&self) -> Result<(), MLError> {
self.grad_decomp_snapshot(&self.grad_snapshot_rec, "grad_snapshot_rec")
}
/// Task 2.0 extension — snapshot before `compute_predictive_coding_loss`.
pub(crate) fn grad_decomp_snapshot_pred(&self) -> Result<(), MLError> {
self.grad_decomp_snapshot(&self.grad_snapshot_pred, "grad_snapshot_pred")
}
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.
///
/// 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.
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")
}
pub(crate) fn grad_decomp_launch_cql_sx(&self) -> Result<(), MLError> {
self.launch_grad_decomp(&self.grad_snapshot_cql_sx, 4, "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.
pub(crate) fn grad_decomp_launch_c51(&self) -> Result<(), MLError> {
self.launch_grad_decomp(&self.grad_zero_ref_buf, 4, "c51")
self.launch_grad_decomp(&self.grad_zero_ref_buf, 6, "c51")
}
pub(crate) fn grad_decomp_launch_c51_bs(&self) -> Result<(), MLError> {
self.launch_grad_decomp(&self.grad_snapshot_c51_bs, 8, "c51_bs")
}
pub(crate) fn grad_decomp_launch_distill(&self) -> Result<(), MLError> {
self.launch_grad_decomp(&self.grad_snapshot_distill, 10, "distill")
}
pub(crate) fn grad_decomp_launch_rec(&self) -> Result<(), MLError> {
self.launch_grad_decomp(&self.grad_snapshot_rec, 12, "rec")
}
pub(crate) fn grad_decomp_launch_pred(&self) -> Result<(), MLError> {
self.launch_grad_decomp(&self.grad_snapshot_pred, 14, "pred")
}
pub(crate) fn grad_decomp_launch_ens(&self) -> Result<(), MLError> {
self.launch_grad_decomp(&self.grad_snapshot_ens, 6, "ens")
self.launch_grad_decomp(&self.grad_snapshot_ens, 16, "ens")
}
/// Task 2.0 — populate the host-side `grad_component_norms_mag/_dir`
@@ -2319,12 +2392,21 @@ impl GpuDqnTrainer {
/// 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
/// 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
///
/// Host caches use `[iqn, cql, c51, ens]` index order matching the
/// `grad_mag_{iqn,cql,c51,ens}_ratio` accessor families.
/// 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,
/// distill,rec,pred,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(
@@ -2335,9 +2417,9 @@ 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, 8)
std::slice::from_raw_parts(self.grad_decomp_result_pinned, 18)
};
for comp in 0..4_usize {
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];
}
@@ -2346,11 +2428,11 @@ impl GpuDqnTrainer {
/// 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 }
/// `[iqn, cql, cql_sx, c51, c51_bs, distill, rec, pred, ens]`.
pub fn grad_component_norms_mag(&self) -> [f32; 9] { 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 }
pub fn grad_component_norms_dir(&self) -> [f32; 9] { self.grad_component_norms_dir }
/// Task 0.6 — per-branch target/online parameter drift (H8 detection signal).
///
@@ -5996,6 +6078,10 @@ impl GpuDqnTrainer {
.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}")))?;
// Task 2.0 extension — isolate `apply_cql_saxpy` (the point where the
// CQL scratch is actually added into grad_buf with budget scaling).
let grad_snapshot_cql_sx = stream.alloc_zeros::<f32>(grad_decomp_snapshot_len)
.map_err(|e| MLError::ModelError(format!("alloc grad_snapshot_cql_sx: {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
@@ -6003,20 +6089,30 @@ impl GpuDqnTrainer {
// 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}")))?;
// Task 2.0 extension — multiplicative scale + additive-writer snapshots.
let grad_snapshot_c51_bs = stream.alloc_zeros::<f32>(grad_decomp_snapshot_len)
.map_err(|e| MLError::ModelError(format!("alloc grad_snapshot_c51_bs: {e}")))?;
let grad_snapshot_distill = stream.alloc_zeros::<f32>(grad_decomp_snapshot_len)
.map_err(|e| MLError::ModelError(format!("alloc grad_snapshot_distill: {e}")))?;
let grad_snapshot_rec = stream.alloc_zeros::<f32>(grad_decomp_snapshot_len)
.map_err(|e| MLError::ModelError(format!("alloc grad_snapshot_rec: {e}")))?;
let grad_snapshot_pred = stream.alloc_zeros::<f32>(grad_decomp_snapshot_len)
.map_err(|e| MLError::ModelError(format!("alloc grad_snapshot_pred: {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
// 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).
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)
cudarc::driver::result::malloc_host(18 * 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; }
for i in 0..18 { *grad_decomp_result_pinned.add(i) = 0.0; }
}
let grad_decomp_result_dev_ptr = unsafe {
let mut dp = 0u64;
@@ -7899,7 +7995,12 @@ impl GpuDqnTrainer {
grad_readback_pinned_capacity,
grad_snapshot_iqn,
grad_snapshot_cql,
grad_snapshot_cql_sx,
grad_zero_ref_buf,
grad_snapshot_c51_bs,
grad_snapshot_distill,
grad_snapshot_rec,
grad_snapshot_pred,
grad_snapshot_ens,
grad_decomp_result_pinned,
grad_decomp_result_dev_ptr,
@@ -7909,8 +8010,8 @@ impl GpuDqnTrainer {
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],
grad_component_norms_mag: [0.0_f32; 9],
grad_component_norms_dir: [0.0_f32; 9],
params_buf,
target_params_buf,
m_buf,

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

@@ -903,30 +903,72 @@ impl FusedTrainingCtx {
/// Task 2.0 — IQN magnitude/direction grad-norm ratio. Index 0 in the
/// `grad_component_norms_*` caches (populated by
/// `refresh_grad_component_norms`).
/// `refresh_grad_component_norms`). Index order (9 slots, Task 2.0
/// extension):
/// `[iqn, cql, cql_sx, c51, c51_bs, distill, rec, pred, ens]`.
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).
/// Task 2.0 — CQL magnitude/direction grad-norm ratio (index 1; spans
/// the snapshot→reduction pair that wraps `apply_cql_gradient` +
/// `apply_cql_saxpy` together).
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 {
/// Task 2.0 extension — CQL-SAXPY magnitude/direction grad-norm ratio
/// (index 2; isolates only `apply_cql_saxpy` — the budget-scaled add of
/// `cql_grad_scratch` into `grad_buf`).
pub(crate) fn grad_mag_cql_sx_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 {
/// Task 2.0 — C51 magnitude/direction grad-norm ratio (index 3; pre-
/// budget-scale — same semantics as first-pass Task 2.0).
pub(crate) fn grad_mag_c51_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 2.0 extension — C51-budget-scale magnitude/direction delta ratio
/// (index 4; isolates the multiplicative amplitude change applied by
/// `apply_c51_budget_scale`).
pub(crate) fn grad_mag_c51_bs_ratio(&self) -> f32 {
let m = self.trainer.grad_component_norms_mag()[4];
let d = self.trainer.grad_component_norms_dir()[4];
if d > 1e-9 { m / d } else { 0.0 }
}
/// Task 2.0 extension — Distillation magnitude/direction grad-norm ratio
/// (index 5; isolates `apply_distillation_gradient`).
pub(crate) fn grad_mag_distill_ratio(&self) -> f32 {
let m = self.trainer.grad_component_norms_mag()[5];
let d = self.trainer.grad_component_norms_dir()[5];
if d > 1e-9 { m / d } else { 0.0 }
}
/// Task 2.0 extension — Recursive-confidence magnitude/direction grad-
/// norm ratio (index 6; isolates `launch_recursive_confidence_backward`).
pub(crate) fn grad_mag_rec_ratio(&self) -> f32 {
let m = self.trainer.grad_component_norms_mag()[6];
let d = self.trainer.grad_component_norms_dir()[6];
if d > 1e-9 { m / d } else { 0.0 }
}
/// Task 2.0 extension — Predictive-coding magnitude/direction grad-norm
/// ratio (index 7; isolates `compute_predictive_coding_loss`).
pub(crate) fn grad_mag_pred_ratio(&self) -> f32 {
let m = self.trainer.grad_component_norms_mag()[7];
let d = self.trainer.grad_component_norms_dir()[7];
if d > 1e-9 { m / d } else { 0.0 }
}
/// Task 2.0 — Ensemble magnitude/direction grad-norm ratio (index 8).
pub(crate) fn grad_mag_ens_ratio(&self) -> f32 {
let m = self.trainer.grad_component_norms_mag()[8];
let d = self.trainer.grad_component_norms_dir()[8];
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‖.
@@ -1231,8 +1273,16 @@ impl FusedTrainingCtx {
// C51 backward already wrote into grad_buf during submit_forward_ops_main (Phase 2).
// This scale MUST run before CQL/IQN SAXPYs below so only C51's portion is scaled.
// Order: c51_budget × C51_grad (here), then cql_budget × CQL_grad, iqn_budget × IQN_grad.
//
// Task 2.0 extension — snapshot BEFORE the multiplicative scale so the
// post-scale reduction measures the amplitude change on branch 0+1.
// When `c51_budget ≈ 1.0` the scale is a no-op and delta = 0.0.
self.trainer.grad_decomp_snapshot_c51_bs()
.map_err(|e| anyhow::anyhow!("Task 2.0 grad_decomp_snapshot_c51_bs: {e}"))?;
self.trainer.apply_c51_budget_scale(c51_budget)
.map_err(|e| anyhow::anyhow!("c51_budget_scale: {e}"))?;
self.trainer.grad_decomp_launch_c51_bs()
.map_err(|e| anyhow::anyhow!("Task 2.0 grad_decomp c51_bs: {e}"))?;
// DIAGNOSTIC: sync between each aux op to find hanging kernel
@@ -1568,14 +1618,40 @@ impl FusedTrainingCtx {
if self.trainer.has_cql() {
match self.trainer.apply_cql_gradient(f5_barrier_weight) {
Ok(true) => {
// Task 2.0 extension — isolate `apply_cql_saxpy`. The
// preceding `apply_cql_gradient` only writes into
// `cql_grad_scratch`, not `grad_buf`, so the cql_sx
// delta captures exactly the budget-scaled SAXPY into
// `grad_buf`.
self.trainer.grad_decomp_snapshot_cql_sx()
.map_err(|e| anyhow::anyhow!("Task 2.0 grad_decomp_snapshot_cql_sx: {e}"))?;
self.trainer.apply_cql_saxpy(cql_budget)
.map_err(|e| anyhow::anyhow!("CQL SAXPY: {e}"))?;
self.trainer.grad_decomp_launch_cql_sx()
.map_err(|e| anyhow::anyhow!("Task 2.0 grad_decomp cql_sx: {e}"))?;
}
Ok(false) => {
// Keep cql_sx slot populated on skip path — delta = 0.
self.trainer.grad_decomp_snapshot_cql_sx()
.map_err(|e| anyhow::anyhow!("Task 2.0 grad_decomp_snapshot_cql_sx (skip): {e}"))?;
self.trainer.grad_decomp_launch_cql_sx()
.map_err(|e| anyhow::anyhow!("Task 2.0 grad_decomp cql_sx (skip): {e}"))?;
}
Ok(false) => {}
Err(e) => {
tracing::warn!("CQL gradient failed (non-fatal): {e}");
// Keep cql_sx slot populated on error path — delta = 0.
self.trainer.grad_decomp_snapshot_cql_sx()
.map_err(|e| anyhow::anyhow!("Task 2.0 grad_decomp_snapshot_cql_sx (err): {e}"))?;
self.trainer.grad_decomp_launch_cql_sx()
.map_err(|e| anyhow::anyhow!("Task 2.0 grad_decomp cql_sx (err): {e}"))?;
}
}
} else {
// Keep cql_sx slot populated on CQL-disabled path — delta = 0.
self.trainer.grad_decomp_snapshot_cql_sx()
.map_err(|e| anyhow::anyhow!("Task 2.0 grad_decomp_snapshot_cql_sx (no-cql): {e}"))?;
self.trainer.grad_decomp_launch_cql_sx()
.map_err(|e| anyhow::anyhow!("Task 2.0 grad_decomp cql_sx (no-cql): {e}"))?;
}
// Task 2.0 — CQL reduction. Delta = grad_buf − snapshot over branch 0+1.
@@ -1593,14 +1669,27 @@ impl FusedTrainingCtx {
// epoch-boundary placement wrote to grad_buf AFTER the last step,
// and the next step's graph_forward memset zeroed it before Adam
// could consume it.
//
// Task 2.0 extension — isolate `apply_distillation_gradient` delta.
self.trainer.grad_decomp_snapshot_distill()
.map_err(|e| anyhow::anyhow!("Task 2.0 grad_decomp_snapshot_distill: {e}"))?;
self.trainer.apply_distillation_gradient()
.map_err(|e| anyhow::anyhow!("distill SAXPY: {e}"))?;
self.trainer.grad_decomp_launch_distill()
.map_err(|e| anyhow::anyhow!("Task 2.0 grad_decomp distill: {e}"))?;
// Recursive confidence backward: MSE grad into trunk + conf weight gradients.
// Must run before Adam (which reads grad_buf for the parameter update).
//
// Task 2.0 extension — isolate `launch_recursive_confidence_backward`
// delta.
self.trainer.grad_decomp_snapshot_rec()
.map_err(|e| anyhow::anyhow!("Task 2.0 grad_decomp_snapshot_rec: {e}"))?;
self.trainer.launch_recursive_confidence_backward(self.batch_size)
.map_err(|e| anyhow::anyhow!("Recursive confidence backward: {e}"))?;
self.trainer.grad_decomp_launch_rec()
.map_err(|e| anyhow::anyhow!("Task 2.0 grad_decomp rec: {e}"))?;
// G12: Predictive coding auxiliary loss + backward.
// Self-supervised temporal smoothness on the enriched trunk h_s2 —
@@ -1608,8 +1697,15 @@ impl FusedTrainingCtx {
// bw_d_h_s2. Runs HERE so the gradient lands in bw_d_h_s2 before
// the trunk backward GEMMs consume it (and before recursive
// confidence's own trunk gradient — both accumulate via plain +=).
//
// Task 2.0 extension — isolate `compute_predictive_coding_loss`
// delta.
self.trainer.grad_decomp_snapshot_pred()
.map_err(|e| anyhow::anyhow!("Task 2.0 grad_decomp_snapshot_pred: {e}"))?;
self.trainer.compute_predictive_coding_loss(self.batch_size)
.map_err(|e| anyhow::anyhow!("Predictive coding (G12): {e}"))?;
self.trainer.grad_decomp_launch_pred()
.map_err(|e| anyhow::anyhow!("Task 2.0 grad_decomp pred: {e}"))?;
// G6 + G10 measurement (commit 61ab27ff3) found both penalties are
// dominated by existing mechanisms:

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

@@ -2228,27 +2228,40 @@ 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 2.0 (+ extension) — per-component grad decomposition across
// 9 measurement points: IQN / CQL / CQL-SAXPY / C51 / C51-budget-
// scale / Distillation / Recursive-confidence / Predictive-coding /
// Ensemble. 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_cql_sx,
grad_mag_c51, grad_mag_c51_bs,
grad_mag_distill, grad_mag_rec, grad_mag_pred,
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_cql_sx_ratio(),
fused.grad_mag_c51_ratio(),
fused.grad_mag_c51_bs_ratio(),
fused.grad_mag_distill_ratio(),
fused.grad_mag_rec_ratio(),
fused.grad_mag_pred_ratio(),
fused.grad_mag_ens_ratio(),
)
} 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)
};
// Task 0.6 — per-branch NoisyNets σ mean. H7 detection signal.
// 0 = direction head, 1 = magnitude head.
@@ -2592,7 +2605,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 [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}]",
"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}] 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,
@@ -2635,12 +2648,29 @@ 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.
// Task 2.0 (+ extension) — per-component grad decomposition
// (9 f32 split across two HEALTH_DIAG groups to keep each
// under ~500 chars and separate backward writers from aux-
// graph writers + multiplicative scalings):
//
// grad_split_bwd [iqn cql c51 ens]
// The four original Task 2.0 measurement points — full
// backward-gradient contributions from each loss family.
//
// grad_split_aux [distill rec pred cql_sx c51_bs]
// The five Task 2.0 extension points inside the aux-graph
// gap between CQL and Ens: three auxiliary-loss writers
// (distill / recursive-confidence / predictive-coding)
// and two multiplicative scalings (cql_sx = apply_cql_
// saxpy; c51_bs = apply_c51_budget_scale). Localizes
// exactly which stage zeroes magnitude gradient.
//
// 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,
// 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],