diff --git a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs index 878b67fee..4771a1c9d 100644 --- a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs +++ b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs @@ -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, // [branch01_elems] — snapshot before IQN trunk grad - grad_snapshot_cql: CudaSlice, // [branch01_elems] — snapshot before CQL SAXPY + grad_snapshot_cql: CudaSlice, // [branch01_elems] — snapshot before CQL block (gradient+saxpy) + grad_snapshot_cql_sx: CudaSlice, // [branch01_elems] — snapshot before apply_cql_saxpy only grad_zero_ref_buf: CudaSlice, // [branch01_elems] — permanent zeros; C51 snapshot substitute + grad_snapshot_c51_bs: CudaSlice, // [branch01_elems] — snapshot before apply_c51_budget_scale + grad_snapshot_distill: CudaSlice, // [branch01_elems] — snapshot before apply_distillation_gradient + grad_snapshot_rec: CudaSlice, // [branch01_elems] — snapshot before launch_recursive_confidence_backward + grad_snapshot_pred: CudaSlice, // [branch01_elems] — snapshot before compute_predictive_coding_loss grad_snapshot_ens: CudaSlice, // [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, // [TOTAL_PARAMS] f32 master online parameters (Adam operates here) target_params_buf: CudaSlice, // [TOTAL_PARAMS] f32 master target parameters (EMA operates here) m_buf: CudaSlice, // [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::(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::(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::(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::(grad_decomp_snapshot_len) + .map_err(|e| MLError::ModelError(format!("alloc grad_snapshot_c51_bs: {e}")))?; + let grad_snapshot_distill = stream.alloc_zeros::(grad_decomp_snapshot_len) + .map_err(|e| MLError::ModelError(format!("alloc grad_snapshot_distill: {e}")))?; + let grad_snapshot_rec = stream.alloc_zeros::(grad_decomp_snapshot_len) + .map_err(|e| MLError::ModelError(format!("alloc grad_snapshot_rec: {e}")))?; + let grad_snapshot_pred = stream.alloc_zeros::(grad_decomp_snapshot_len) + .map_err(|e| MLError::ModelError(format!("alloc grad_snapshot_pred: {e}")))?; let grad_snapshot_ens = stream.alloc_zeros::(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::(), flags) + cudarc::driver::result::malloc_host(18 * std::mem::size_of::(), 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, diff --git a/crates/ml/src/trainers/dqn/fused_training.rs b/crates/ml/src/trainers/dqn/fused_training.rs index 2680d1774..eb1ceee3d 100644 --- a/crates/ml/src/trainers/dqn/fused_training.rs +++ b/crates/ml/src/trainers/dqn/fused_training.rs @@ -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: diff --git a/crates/ml/src/trainers/dqn/trainer/training_loop.rs b/crates/ml/src/trainers/dqn/trainer/training_loop.rs index 63d9cd158..e3ce0a3d3 100644 --- a/crates/ml/src/trainers/dqn/trainer/training_loop.rs +++ b/crates/ml/src/trainers/dqn/trainer/training_loop.rs @@ -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],