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

feat(sp6): Pearl 2 — SAXPY per-branch loss budgets (correction-factor pattern)

Browse Source 
SP6 sub-project 1 (Pearl 2): converts the 4 SAXPY launchers from a single
scalar budget (mean of 4 ISV branch slots) to per-branch differentiated
scaling via the correction-factor pattern.

Problem: SP5 Layer B read ISV[190..210) per-branch budget slots but
collapsed them to a scalar via sum/4.0, then passed the single scalar to
apply_c51_budget_scale / apply_cql_saxpy / apply_iqn_trunk_gradient.
Branch HEAD parameters received the same budget as trunk, defeating
per-branch differentiation.

Fix: compute_adaptive_budgets() now returns ([f32;4], [f32;4], [f32;4],
[f32;4], f32, f32, f32, f32) — four per-branch arrays + four trunk-mean
scalars. The trunk mean (D3 decision) is used for the full-buffer trunk/value
SAXPY call (preserving SP5 Layer B behavior for shared params). Branch HEAD
parameter slices receive a correction sub-launch:

  correction = branch_budget[b] / trunk_mean    (skip if |correction-1| <= 1e-6)

After both launches, branch HEAD slice is effectively scaled by
branch_budget[b], trunk/value is scaled by trunk_mean. No double-scaling.

New helpers added to GpuDqnTrainer:
  apply_c51_budget_scale_branch(branch_idx, correction): scale_f32_ungraphed
    on branch-slice [f32 elements], offset via padded_byte_offset.
  apply_cql_saxpy_branch(branch_idx, correction): saxpy_f32_aux on
    branch-slice of both grad_buf and cql_grad_scratch.

IQN trunk gradient: uses iqn_trunk (mean of 4 branch IQN budgets) — IQN
backward flows through trunk only; per-branch IQN routing is beyond SP6 scope.

HEALTH_DIAG: three new per-epoch info! lines emit per-branch c51/iqn/cql
budget arrays (dir/mag/ord/urg) after the intent_dist line.

State: per-branch budget arrays cached on GpuDqnTrainer
(last_*_budget_per_branch: [f32;4]) and on DqnTrainer
(last_*_budget_per_branch: Option<[f32;4]>) for diagnostics.

docs/isv-slots.md: updated Pearl 2 slot rows to reflect SP6 consumer wiring.

Verification: cargo check + release build clean (13 warnings, pre-existing).
13 sp5+sp4+state_reset_registry lib tests pass. sp5_producer_unit_tests
--no-run clean. Sanity grep for old sum/4.0 averaging pattern: empty.

Files changed (6): gpu_dqn_trainer.rs, fused_training.rs, constructor.rs,
mod.rs, training_loop.rs, docs/isv-slots.md. No Pearl 3 or Pearl 5 files touched.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-05-02 02:08:54 +02:00
parent 6b01b70c3e
commit f62860ea82
6 changed files with 190 additions and 52 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -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::<f32>();
let param_sizes = compute_param_sizes(&self.config);
let first_tensor = 8 + branch_idx * 4;
let start_bytes = padded_byte_offset(&param_sizes, first_tensor) as usize;
let end_bytes = padded_byte_offset(&param_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::<f32>();
let param_sizes = compute_param_sizes(&self.config);
let first_tensor = 8 + branch_idx * 4;
let start_bytes = padded_byte_offset(&param_sizes, first_tensor) as usize;
let end_bytes = padded_byte_offset(&param_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,

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

@@ -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::<f32>()
/ 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::<f32>()
/ 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::<f32>()
/ 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::<f32>()
/ 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::<f32>() / 4.0_f32;
let iqn_trunk = iqn.iter().sum::<f32>() / 4.0_f32;
let cql_trunk = cql.iter().sum::<f32>() / 4.0_f32;
let ens_trunk = ens.iter().sum::<f32>() / 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).

3
crates/ml/src/trainers/dqn/trainer/constructor.rs
View File

@@ -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,

4
crates/ml/src/trainers/dqn/trainer/mod.rs
View File

@@ -755,6 +755,10 @@ pub struct DQNTrainer {
pub(crate) last_iqn_budget_eff: Option<f32>,
pub(crate) last_cql_budget_eff: Option<f32>,
pub(crate) last_c51_budget_eff: Option<f32>,
/// 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<f32>,
/// B3/G4 — Expected SARSA temperature scaling factor

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

@@ -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.

8
docs/isv-slots.md
View File

@@ -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 |