diff --git a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs index 86bfbf8b9..c10850a38 100644 --- a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs +++ b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs @@ -4829,6 +4829,12 @@ pub struct GpuDqnTrainer { /// B4/G5: Adaptive gradient budget for ensemble (constant 0.05). /// Populated by `FusedTrainingCtx::compute_adaptive_budgets`. pub(crate) last_ens_budget_eff: f32, + /// SP6 Pearl 2: per-branch budget arrays — set by compute_adaptive_budgets. + /// Index order: dir=0, mag=1, ord=2, urg=3. + pub(crate) last_c51_budget_per_branch: [f32; 4], + pub(crate) last_iqn_budget_per_branch: [f32; 4], + pub(crate) last_cql_budget_per_branch: [f32; 4], + pub(crate) last_ens_budget_per_branch: [f32; 4], /// D1/N1: Ring buffer of best-health weight snapshots for temporal self-distillation. pub(crate) q_snapshots: crate::cuda_pipeline::q_snapshot::SnapshotRing, @@ -9461,6 +9467,46 @@ impl GpuDqnTrainer { Ok(()) } + /// SP6 Pearl 2: scale only branch `branch_idx` parameter slice of grad_buf by `budget`. + /// + /// Called after `apply_c51_budget_scale(trunk_mean)` has scaled the entire grad_buf. + /// The correction factor is `branch_budget / trunk_mean`, so the branch slice ends up + /// scaled by `branch_budget` rather than `trunk_mean`. + /// + /// Branch b owns tensors `(8+4b)..(12+4b)` in the padded flat layout. + /// Padding bytes in `grad_buf` are always zero so scaling them is a no-op. + pub fn apply_c51_budget_scale_branch(&mut self, branch_idx: usize, budget: f32) -> Result<(), MLError> { + if (budget - 1.0).abs() < 1e-6 { + return Ok(()); + } + let f32_size = std::mem::size_of::(); + let param_sizes = compute_param_sizes(&self.config); + let first_tensor = 8 + branch_idx * 4; + let start_bytes = padded_byte_offset(¶m_sizes, first_tensor) as usize; + let end_bytes = padded_byte_offset(¶m_sizes, first_tensor + 4) as usize; + let start_elems = (start_bytes / f32_size) as i32; + let len_elems = ((end_bytes - start_bytes) / f32_size) as i32; + if len_elems == 0 { return Ok(()); } + let grad_slice_ptr = self.ptrs.grad_buf + (start_elems as u64) * (f32_size as u64); + let blocks = ((len_elems as u32 + 255) / 256) as u32; + unsafe { + self.stream + .launch_builder(&self.scale_f32_ungraphed) + .arg(&grad_slice_ptr) + .arg(&budget) + .arg(&len_elems) + .launch(LaunchConfig { + grid_dim: (blocks, 1, 1), + block_dim: (256, 1, 1), + shared_mem_bytes: 0, + }) + .map_err(|e| MLError::ModelError(format!( + "c51_budget_scale_branch[{branch_idx}]: {e}" + )))?; + } + Ok(()) + } + /// Add CQL gradient scratch to grad_buf via plain SAXPY. /// /// Called after `apply_cql_gradient` populated `cql_grad_scratch`. @@ -9493,6 +9539,45 @@ impl GpuDqnTrainer { Ok(()) } + /// SP6 Pearl 2: SAXPY only branch `branch_idx` parameter slice: + /// grad_buf[branch_slice] += budget * cql_grad_scratch[branch_slice] + /// + /// Called after `apply_cql_saxpy(trunk_mean)` has performed the full-buffer SAXPY. + /// The correction factor is `branch_budget / trunk_mean`, so the branch slice ends up + /// accumulating `branch_budget * cql_scratch` rather than `trunk_mean * cql_scratch`. + /// + /// Branch b owns tensors `(8+4b)..(12+4b)` in the padded flat layout. + pub fn apply_cql_saxpy_branch(&mut self, branch_idx: usize, budget: f32) -> Result<(), MLError> { + let f32_size = std::mem::size_of::(); + let param_sizes = compute_param_sizes(&self.config); + let first_tensor = 8 + branch_idx * 4; + let start_bytes = padded_byte_offset(¶m_sizes, first_tensor) as usize; + let end_bytes = padded_byte_offset(¶m_sizes, first_tensor + 4) as usize; + let start_elems = (start_bytes / f32_size) as i32; + let len_elems = ((end_bytes - start_bytes) / f32_size) as i32; + if len_elems == 0 { return Ok(()); } + let grad_ptr = self.ptrs.grad_buf + (start_elems as u64) * (f32_size as u64); + let scratch_ptr = self.ptrs.cql_grad_scratch + (start_elems as u64) * (f32_size as u64); + let blocks = ((len_elems as u32 + 255) / 256) as u32; + unsafe { + self.stream + .launch_builder(&self.saxpy_f32_aux) + .arg(&grad_ptr) + .arg(&scratch_ptr) + .arg(&budget) + .arg(&len_elems) + .launch(LaunchConfig { + grid_dim: (blocks, 1, 1), + block_dim: (256, 1, 1), + shared_mem_bytes: 0, + }) + .map_err(|e| MLError::ModelError(format!( + "cql_saxpy_branch[{branch_idx}]: {e}" + )))?; + } + Ok(()) + } + /// F5/D2: Inject the Q-gap barrier gradient into cql_d_adv_logits and cql_d_value_logits /// (direction branch b0 only) via atomicAdd. /// @@ -16943,6 +17028,10 @@ impl GpuDqnTrainer { last_cql_budget_eff: 0.00, last_c51_budget_eff: 0.55, last_ens_budget_eff: 0.05, + last_c51_budget_per_branch: [0.55; 4], + last_iqn_budget_per_branch: [0.40; 4], + last_cql_budget_per_branch: [0.00; 4], + last_ens_budget_per_branch: [0.05; 4], q_snapshots: crate::cuda_pipeline::q_snapshot::SnapshotRing::new( stream_for_snapshots, total_params, diff --git a/crates/ml/src/trainers/dqn/fused_training.rs b/crates/ml/src/trainers/dqn/fused_training.rs index aebff917a..40826dbf0 100644 --- a/crates/ml/src/trainers/dqn/fused_training.rs +++ b/crates/ml/src/trainers/dqn/fused_training.rs @@ -1878,8 +1878,9 @@ impl FusedTrainingCtx { // not as norm clips. This preserves gradient magnitude variation while steering // component balance. Safety is still handled by the single global clip in Adam. - // B4/G5: Compute adaptive budgets once per step from ISV health/regime signals. - let (c51_budget, iqn_budget, cql_budget, _ens_budget) = self.compute_adaptive_budgets(); + // SP6 Pearl 2: Compute adaptive budgets — returns per-branch arrays + trunk means. + let (c51_branch, iqn_branch, cql_branch, _ens_branch, + c51_trunk, iqn_trunk, cql_trunk, _ens_trunk) = self.compute_adaptive_budgets(); // F8/G5: Scale C51 contribution in grad_buf by c51_budget. // C51 backward already wrote into grad_buf during submit_forward_ops_main (Phase 2). @@ -1891,8 +1892,18 @@ impl FusedTrainingCtx { // 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}"))?; + // SP6 Pearl 2: trunk/value first (mean budget), then per-branch correction sub-launches. + // Branch slice has already been scaled by c51_trunk; correction = branch[b]/trunk_mean + // so branch slice ends up scaled by c51_branch[b]. + self.trainer.apply_c51_budget_scale(c51_trunk) + .map_err(|e| anyhow::anyhow!("c51_budget_scale_trunk: {e}"))?; + for b in 0..4_usize { + let correction = if c51_trunk > 1e-6 { c51_branch[b] / c51_trunk } else { 1.0_f32 }; + if (correction - 1.0).abs() > 1e-6 { + self.trainer.apply_c51_budget_scale_branch(b, correction) + .map_err(|e| anyhow::anyhow!("c51_budget_scale_branch[{b}]: {e}"))?; + } + } self.trainer.grad_decomp_launch_c51_bs() .map_err(|e| anyhow::anyhow!("Task 2.0 grad_decomp c51_bs: {e}"))?; @@ -2265,8 +2276,9 @@ impl FusedTrainingCtx { if iqn_ok { if let Some(ref mut iqn) = self.gpu_iqn { let d_h_s2_ptr = iqn.d_h_s2_raw_ptr(); + // SP6 Pearl 2: IQN operates on trunk only — use trunk mean (no per-branch sub-launch). self.trainer.apply_iqn_trunk_gradient( - d_h_s2_ptr, &mut self.online_dueling, iqn_budget, + d_h_s2_ptr, &mut self.online_dueling, iqn_trunk, ).map_err(|e| anyhow::anyhow!("IQN trunk gradient: {e}"))?; let tau = self.trainer.read_isv_signal_at(TAU_EFF_INDEX); @@ -2321,8 +2333,16 @@ impl FusedTrainingCtx { // `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}"))?; + // SP6 Pearl 2: trunk/value SAXPY first (mean budget), then per-branch corrections. + self.trainer.apply_cql_saxpy(cql_trunk) + .map_err(|e| anyhow::anyhow!("CQL SAXPY trunk: {e}"))?; + for b in 0..4_usize { + let correction = if cql_trunk > 1e-6 { cql_branch[b] / cql_trunk } else { 1.0_f32 }; + if (correction - 1.0).abs() > 1e-6 { + self.trainer.apply_cql_saxpy_branch(b, correction) + .map_err(|e| anyhow::anyhow!("CQL SAXPY branch[{b}]: {e}"))?; + } + } self.trainer.grad_decomp_launch_cql_sx() .map_err(|e| anyhow::anyhow!("Task 2.0 grad_decomp cql_sx: {e}"))?; } @@ -3278,49 +3298,46 @@ impl FusedTrainingCtx { self.trainer.last_sarsa_tau_factor } - /// B4/G5: Compute adaptive gradient budgets from learning_health and regime_stability. - /// Returns (c51_budget, iqn_budget, cql_budget, ens_budget). Caches to trainer fields - /// for HEALTH_DIAG logging. + /// SP6 Pearl 2: Compute adaptive gradient budgets from ISV[190..210). /// - /// SP5 Layer B (Pearl 2): reads per-branch budget signals from ISV[190..210) - /// and collapses them to a single effective scale per loss component by - /// averaging across the 4 branches. The Layer A producer (SP5 Task A3) runs - /// the Wiener-optimal EMA (Pearls A+D) so the values already encode - /// health/regime dynamics — no second-hand formula is needed here. + /// Returns `(c51_branch, iqn_branch, cql_branch, ens_branch, c51_trunk, iqn_trunk, + /// cql_trunk, ens_trunk)`. The per-branch arrays feed per-branch SAXPY correction + /// sub-launches on branch HEAD parameter slices. The trunk scalars (mean of 4 branch + /// budgets) feed the full-buf trunk/value scaling call — preserving SP5 Layer B + /// behaviour for shared parameters (D3 decision in SP6 spec). /// - /// Cold-start floor: ISV reads 0 before the first observation; the - /// normalisation `/ 4.0` and clamp ensure we never pass 0 to the loss - /// scalers. The floor values (0.05 C51, 0.05 IQN, 0.02 CQL, 0.02 Ens) - /// are Invariant 1 carve-outs (numerical-stability), not tuned constants. - pub(crate) fn compute_adaptive_budgets(&mut self) -> (f32, f32, f32, f32) { - // Mean of 4 per-branch slots → single effective scale. - let c51_budget = ((0..4_usize) - .map(|b| self.read_isv_signal_at(BUDGET_C51_BASE + b)) - .sum::() - / 4.0_f32) - .max(0.05_f32); - let iqn_budget = ((0..4_usize) - .map(|b| self.read_isv_signal_at(BUDGET_IQN_BASE + b)) - .sum::() - / 4.0_f32) - .max(0.05_f32); - let cql_budget = ((0..4_usize) - .map(|b| self.read_isv_signal_at(BUDGET_CQL_BASE + b)) - .sum::() - / 4.0_f32) - .max(0.02_f32); - let ens_budget = ((0..4_usize) - .map(|b| self.read_isv_signal_at(BUDGET_ENS_BASE + b)) - .sum::() - / 4.0_f32) - .max(0.02_f32); - - self.trainer.last_iqn_budget_eff = iqn_budget; - self.trainer.last_cql_budget_eff = cql_budget; - self.trainer.last_c51_budget_eff = c51_budget; - self.trainer.last_ens_budget_eff = ens_budget; - - (c51_budget, iqn_budget, cql_budget, ens_budget) + /// Cold-start floor: ISV slots read 0 before first observation. Floors (0.05 C51, + /// 0.05 IQN, 0.02 CQL, 0.02 Ens) are Invariant 1 carve-outs (numerical-stability, + /// not tuned constants). + pub(crate) fn compute_adaptive_budgets( + &mut self, + ) -> ([f32; 4], [f32; 4], [f32; 4], [f32; 4], f32, f32, f32, f32) { + let mut c51 = [0.0_f32; 4]; + let mut iqn = [0.0_f32; 4]; + let mut cql = [0.0_f32; 4]; + let mut ens = [0.0_f32; 4]; + for b in 0..4_usize { + c51[b] = self.read_isv_signal_at(BUDGET_C51_BASE + b).max(0.05_f32); + iqn[b] = self.read_isv_signal_at(BUDGET_IQN_BASE + b).max(0.05_f32); + cql[b] = self.read_isv_signal_at(BUDGET_CQL_BASE + b).max(0.02_f32); + ens[b] = self.read_isv_signal_at(BUDGET_ENS_BASE + b).max(0.02_f32); + } + // Trunk/value use mean of 4 branch budgets (D3 decision in SP6 spec). + let c51_trunk = c51.iter().sum::() / 4.0_f32; + let iqn_trunk = iqn.iter().sum::() / 4.0_f32; + let cql_trunk = cql.iter().sum::() / 4.0_f32; + let ens_trunk = ens.iter().sum::() / 4.0_f32; + // Cache per-branch arrays for HEALTH_DIAG logging. + self.trainer.last_c51_budget_per_branch = c51; + self.trainer.last_iqn_budget_per_branch = iqn; + self.trainer.last_cql_budget_per_branch = cql; + self.trainer.last_ens_budget_per_branch = ens; + // Cache trunk scalars for backward-compat accessors. + self.trainer.last_c51_budget_eff = c51_trunk; + self.trainer.last_iqn_budget_eff = iqn_trunk; + self.trainer.last_cql_budget_eff = cql_trunk; + self.trainer.last_ens_budget_eff = ens_trunk; + (c51, iqn, cql, ens, c51_trunk, iqn_trunk, cql_trunk, ens_trunk) } /// B4/G5: Last adaptive IQN gradient budget (0.10 + 0.30×health). diff --git a/crates/ml/src/trainers/dqn/trainer/constructor.rs b/crates/ml/src/trainers/dqn/trainer/constructor.rs index d972ced64..520238bb4 100644 --- a/crates/ml/src/trainers/dqn/trainer/constructor.rs +++ b/crates/ml/src/trainers/dqn/trainer/constructor.rs @@ -786,6 +786,9 @@ impl DQNTrainer { last_iqn_budget_eff: None, last_cql_budget_eff: None, last_c51_budget_eff: None, + last_c51_budget_per_branch: None, + last_iqn_budget_per_branch: None, + last_cql_budget_per_branch: None, last_tau_eff: None, last_sarsa_tau_factor: None, last_gamma_eff: None, diff --git a/crates/ml/src/trainers/dqn/trainer/mod.rs b/crates/ml/src/trainers/dqn/trainer/mod.rs index f899b18ff..11e6d5ae9 100644 --- a/crates/ml/src/trainers/dqn/trainer/mod.rs +++ b/crates/ml/src/trainers/dqn/trainer/mod.rs @@ -755,6 +755,10 @@ pub struct DQNTrainer { pub(crate) last_iqn_budget_eff: Option, pub(crate) last_cql_budget_eff: Option, pub(crate) last_c51_budget_eff: Option, + /// SP6 Pearl 2 — per-branch budget arrays (dir=0, mag=1, ord=2, urg=3) + pub(crate) last_c51_budget_per_branch: Option<[f32; 4]>, + pub(crate) last_iqn_budget_per_branch: Option<[f32; 4]>, + pub(crate) last_cql_budget_per_branch: Option<[f32; 4]>, /// B2/G3 — tau effective pub(crate) last_tau_eff: Option, /// B3/G4 — Expected SARSA temperature scaling factor diff --git a/crates/ml/src/trainers/dqn/trainer/training_loop.rs b/crates/ml/src/trainers/dqn/trainer/training_loop.rs index f19ed77ce..6eac61be6 100644 --- a/crates/ml/src/trainers/dqn/trainer/training_loop.rs +++ b/crates/ml/src/trainers/dqn/trainer/training_loop.rs @@ -2749,6 +2749,10 @@ impl DQNTrainer { self.last_iqn_budget_eff = Some(fused.last_iqn_budget_eff()); self.last_cql_budget_eff = Some(fused.last_cql_budget_eff()); self.last_c51_budget_eff = Some(fused.last_c51_budget_eff()); + // SP6 Pearl 2: cache per-branch budget arrays. + self.last_c51_budget_per_branch = Some(fused.trainer().last_c51_budget_per_branch); + self.last_iqn_budget_per_branch = Some(fused.trainer().last_iqn_budget_per_branch); + self.last_cql_budget_per_branch = Some(fused.trainer().last_cql_budget_per_branch); } // Plan 2 D.2: propagate direction-branch gamma_eff from ISV for logging. @@ -3975,6 +3979,27 @@ impl DQNTrainer { self.last_eval_intent_magnitude_dist[2], ); + // SP6 Pearl 2: per-branch budget HEALTH_DIAG lines — enables Layer C debugging + // by showing how much each branch's loss budget differs from the trunk mean. + if let Some(arr) = self.last_c51_budget_per_branch { + tracing::info!( + "HEALTH_DIAG[{}]: c51_budget_per_branch [dir={:.4} mag={:.4} ord={:.4} urg={:.4}]", + epoch, arr[0], arr[1], arr[2], arr[3], + ); + } + if let Some(arr) = self.last_iqn_budget_per_branch { + tracing::info!( + "HEALTH_DIAG[{}]: iqn_budget_per_branch [dir={:.4} mag={:.4} ord={:.4} urg={:.4}]", + epoch, arr[0], arr[1], arr[2], arr[3], + ); + } + if let Some(arr) = self.last_cql_budget_per_branch { + tracing::info!( + "HEALTH_DIAG[{}]: cql_budget_per_branch [dir={:.4} mag={:.4} ord={:.4} urg={:.4}]", + epoch, arr[0], arr[1], arr[2], arr[3], + ); + } + // C.2 Plan 3 Task 1 (spec §4.C.2): reward_split HEALTH_DIAG line. // Reads 6 ISV EMA slots updated by the GPU reward_component_ema kernel // launched just above. CPU-side code only reads; GPU wrote the values. diff --git a/docs/isv-slots.md b/docs/isv-slots.md index af78e52dd..46ad869d2 100644 --- a/docs/isv-slots.md +++ b/docs/isv-slots.md @@ -133,10 +133,10 @@ Constants live in `crates/ml/src/cuda_pipeline/sp5_isv_slots.rs`. | [178..182) | `ATOM_V_HALF_BASE` | per-branch [4] | Pearl 1 | C51 atom half-width | | [182..186) | `ATOM_HEADROOM_BASE` | per-branch [4] | Pearl 1 | C51 atom headroom | | [186..190) | `ATOM_CLIP_RATE_BASE` | per-branch [4] | Pearl 1 | C51 atom clip rate | -| [190..194) | `BUDGET_C51_BASE` | per-branch [4] | Pearl 2 | C51 loss budget weight | -| [194..198) | `BUDGET_IQN_BASE` | per-branch [4] | Pearl 2 | IQN loss budget weight | -| [198..202) | `BUDGET_CQL_BASE` | per-branch [4] | Pearl 2 | CQL loss budget weight | -| [202..206) | `BUDGET_ENS_BASE` | per-branch [4] | Pearl 2 | Ensemble loss budget weight | +| [190..194) | `BUDGET_C51_BASE` | per-branch [4] | Pearl 2 | C51 loss budget weight. SP6 Pearl 2: `compute_adaptive_budgets()` reads individually, applies correction-factor sub-launches via `apply_c51_budget_scale_branch`. | +| [194..198) | `BUDGET_IQN_BASE` | per-branch [4] | Pearl 2 | IQN loss budget weight. SP6 Pearl 2: used as trunk-mean only (`iqn_trunk`) — IQN backward targets trunk params exclusively. | +| [198..202) | `BUDGET_CQL_BASE` | per-branch [4] | Pearl 2 | CQL loss budget weight. SP6 Pearl 2: `compute_adaptive_budgets()` reads individually, applies correction-factor sub-launches via `apply_cql_saxpy_branch`. | +| [202..206) | `BUDGET_ENS_BASE` | per-branch [4] | Pearl 2 | Ensemble loss budget weight. SP6 Pearl 2: used as trunk-mean only. | | [206..210) | `FLATNESS_BASE` | per-branch [4] | Pearl 2 | Loss flatness diagnostic | | [210..214) | `NOISY_SIGMA_BASE` | per-branch [4] | Pearl 3 | NoisyNet σ level | | [214..218) | `SIGMA_FRACTION_BASE` | per-branch [4] | Pearl 3 | NoisyNet σ fraction |