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:
@@ -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;
|
||||
|
||||
@@ -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
Reference in New Issue
Block a user