feat(sp4): Tasks A10+A11 — wire all 5 SP4 producer launches in training_loop

Layer A: cold-path SP4 producer launches added next to launch_h_s2_rms_ema
and the surrounding ISV producers (mag_concat h_s2_rms_ema, fold_warmup,
iqn_quantile_ema, vsn_mask_ema, aux_heads_loss_ema, moe_expert_util_ema,
moe_lambda_eff_update). All five fire once per training step before the
HEALTH_DIAG line emit per the post-cascade-fix invariant (a5f23b28f).

Producers wired:
  - launch_sp4_target_q_p99               (TARGET_Q_BOUND, slot 131)
  - launch_sp4_atom_pos_p99_all_branches  (ATOM_POS_BOUND[0..4], 132-135)
  - launch_sp4_grad_norm_p99              (GRAD_CLIP_BOUND, slot 168)
  - launch_sp4_h_s2_p99                   (H_S2_BOUND, slot 169)
  - launch_sp4_param_group_oracles_all_groups via build_sp4_aux_buffers
    (WEIGHT/ADAM_M/ADAM_V/WD_RATE × 8 groups + L1_LAMBDA[trunk]; 33 slots)

For param_group_oracles, the SP4AuxBuffers descriptor is built each step
via FusedTrainingCtx::build_sp4_aux_buffers, threading curiosity_weights
and curiosity_trainer through from gpu_experience_collector. To support
that, added GpuExperienceCollector::curiosity_trainer() accessor mirroring
the existing has_curiosity_trainer / curiosity_weight_set pair. When the
collector is absent or curiosity is disabled, the curiosity descriptor
empties and the launcher's count==0 short-circuit silently skips group 7.

All 40 SP4 ISV bound slots are now populated correctly each step via
Pearls A+D (host-side EMA + p99 application). No consumer reads them yet —
Mech 1/2/5/6/9/10 clamps still use SP3 hardcoded multipliers, so behavior
is unchanged. Layer B will wire consumers and may relocate the captured-
graph-eligible producers (target_q_p99 in particular) as part of Mech 1's
pre-clamp ordering.

Each launch uses the same `if let Some(ref fused) = self.fused_ctx` guard
+ tracing::warn-on-error pattern as launch_h_s2_rms_ema. Errors never
propagate (warn-and-continue) since Layer A is observability-only.

cargo check -p ml --lib --tests clean (11 pre-existing warnings).
state_reset_registry suite passes (3/3).

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-05-01 00:47:49 +02:00
parent 8f93c11525
commit 392fc7d698
3 changed files with 60 additions and 0 deletions

View File

@@ -4273,6 +4273,14 @@ impl GpuExperienceCollector {
&self.curiosity_weights
}
/// Access the GPU curiosity trainer (grad + Adam state). Returns `None`
/// when curiosity is disabled at collector init or after `disable_curiosity`.
/// SP4 Layer A param-group oracle wiring pairs this with `curiosity_weight_set`
/// to populate the curiosity sub-buffer descriptor in `SP4AuxBuffers`.
pub fn curiosity_trainer(&self) -> Option<&GpuCuriosityTrainer> {
self.curiosity_trainer.as_ref()
}
/// Disable curiosity training (e.g. after an async kernel crash).
pub fn disable_curiosity(&mut self) {
self.curiosity_trainer = None;

View File

@@ -3498,6 +3498,56 @@ impl DQNTrainer {
tracing::warn!("launch_moe_lambda_eff_update failed: {e}");
}
}
// SP4 Layer A Tasks A10 + A11: per-step ISV-bound producer
// launches for the 40 new SP4 slots (TARGET_Q_BOUND, ATOM_POS
// ×4 branches, GRAD_CLIP_BOUND, H_S2_BOUND, and 8 param-group
// oracles spanning WEIGHT/ADAM_M/ADAM_V/WD_RATE per group plus
// L1_LAMBDA[trunk]). Layer A is observability-only — no
// consumer reads these slots yet, so behavior is unchanged.
// All five launches go cold-path here next to the rest of
// the per-step ISV producers; Layer B may relocate the
// captured-graph eligible ones (target_q_p99 in particular)
// once the Mech 1 pre-clamp consumer lands.
//
// Producer ordering — like the surrounding ISV producers,
// these must run BEFORE the HEALTH_DIAG line emit (commit
// a5f23b28f post-cascade-fix invariant) and BEFORE any
// future consumer kernel reads the bound slots.
if let Some(ref fused) = self.fused_ctx {
if let Err(e) = fused.trainer().launch_sp4_target_q_p99() {
tracing::warn!("SP4 A10 target_q_p99 launch failed: {e}");
}
if let Err(e) = fused.trainer().launch_sp4_atom_pos_p99_all_branches() {
tracing::warn!("SP4 A10 atom_pos_p99_all_branches launch failed: {e}");
}
if let Err(e) = fused.trainer().launch_sp4_grad_norm_p99() {
tracing::warn!("SP4 A10 grad_norm_p99 launch failed: {e}");
}
if let Err(e) = fused.trainer().launch_sp4_h_s2_p99() {
tracing::warn!("SP4 A10 h_s2_p99 launch failed: {e}");
}
// SP4 A11: param_group_oracle Pearl B over 8 groups
// (DQN trunk + 4 branches groups 0-2; IQN, IQL hi/lo,
// Attn, Curiosity groups 3-7). Curiosity (group 7)
// descriptor is built from `gpu_experience_collector`'s
// `curiosity_weight_set` + `curiosity_trainer` pair —
// both `None` together (curiosity disabled at collector
// init) yields an empty group descriptor and the
// launcher silently skips it.
let curiosity_weights = self.gpu_experience_collector.as_ref()
.map(|c| c.curiosity_weight_set());
let curiosity_trainer = self.gpu_experience_collector.as_ref()
.and_then(|c| c.curiosity_trainer());
let aux_buffers = fused.build_sp4_aux_buffers(
curiosity_weights,
curiosity_trainer,
);
if let Err(e) = fused.trainer().launch_sp4_param_group_oracles_all_groups(&aux_buffers) {
tracing::warn!("SP4 A11 param_group_oracles_all_groups launch failed: {e}");
}
}
}
// Plan 4 Task 6 Commit B: refresh the aux-loss weight before the

File diff suppressed because one or more lines are too long