diff --git a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs index a4c605981..752f30e5e 100644 --- a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs +++ b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs @@ -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, // [branch01_elems] — snapshot before IQN trunk grad 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 @@ -1073,25 +1083,31 @@ pub struct GpuDqnTrainer { 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 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, // [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) @@ -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, ctx: &str) -> Result<(), MLError> { let elem_size = std::mem::size_of::() 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, @@ -2432,6 +2477,8 @@ impl GpuDqnTrainer { self.stream.launch_builder(&self.grad_decomp_kernel) .arg(¤t_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(¶m_sizes, 8) as usize; - let mag_start_byte = padded_byte_offset(¶m_sizes, 12) as usize; + let trunk_start_byte = padded_byte_offset(¶m_sizes, 0) as usize; + let trunk_end_byte = padded_byte_offset(¶m_sizes, 4) as usize; + let dir_start_byte = padded_byte_offset(¶m_sizes, 8) as usize; + let mag_start_byte = padded_byte_offset(¶m_sizes, 12) as usize; let branch2_start_byte = padded_byte_offset(¶m_sizes, 16) as usize; let elem_size = std::mem::size_of::(); + 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::(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::(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::(), flags) + cudarc::driver::result::malloc_host(27 * 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..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, diff --git a/crates/ml/src/cuda_pipeline/grad_decomp_kernel.cu b/crates/ml/src/cuda_pipeline/grad_decomp_kernel.cu index 466872b75..43d2eef8c 100644 --- a/crates/ml/src/cuda_pipeline/grad_decomp_kernel.cu +++ b/crates/ml/src/cuda_pipeline/grad_decomp_kernel.cu @@ -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]); } } diff --git a/crates/ml/src/trainers/dqn/fused_training.rs b/crates/ml/src/trainers/dqn/fused_training.rs index fe150b25f..2cf9254d1 100644 --- a/crates/ml/src/trainers/dqn/fused_training.rs +++ b/crates/ml/src/trainers/dqn/fused_training.rs @@ -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). diff --git a/crates/ml/src/trainers/dqn/trainer/training_loop.rs b/crates/ml/src/trainers/dqn/trainer/training_loop.rs index 4ceb98825..487edc75e 100644 --- a/crates/ml/src/trainers/dqn/trainer/training_loop.rs +++ b/crates/ml/src/trainers/dqn/trainer/training_loop.rs @@ -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). diff --git a/docs/superpowers/specs/2026-03-23-dual-distributional-c51-iqn-design.md b/docs/superpowers/specs/2026-03-23-dual-distributional-c51-iqn-design.md index 33c0e5d5f..ba921da97 100644 --- a/docs/superpowers/specs/2026-03-23-dual-distributional-c51-iqn-design.md +++ b/docs/superpowers/specs/2026-03-23-dual-distributional-c51-iqn-design.md @@ -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)