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

fix(producer-cadence): migrate 4 more per-step ISV producers from process_epoch_boundary to per-step hot path

Browse Source 
Continuation of SP14 B.11 cadence fix (commit 200f05fce). The same per-
epoch staleness bug affected 4 additional producers whose docstrings
explicitly claim "per-step" cadence but whose launches lived in
`process_epoch_boundary` (fires once per epoch at line ~780):

  * launch_h_s2_rms_ema           → ISV[H_S2_RMS_EMA_INDEX=96]
  * launch_fold_warmup_factor     → ISV[FOLD_WARMUP_FACTOR_INDEX=130]
  * launch_moe_expert_util_ema    → ISV[MOE_EXPERT_UTIL_EMA_BASE..+8) +
                                    ISV[MOE_GATE_ENTROPY_EMA_INDEX=126]
  * launch_moe_lambda_eff_update  → ISV[MOE_LAMBDA_EFF_INDEX=128]

Each producer's per-step consumer reads the ISV slot every training
step (thousands per epoch). With the producers stuck per-epoch, consumers
saw (steps_per_epoch − 1)-step-stale values — exactly the failure mode
SP14 B.11 documented for ALPHA_GRAD_SMOOTHED. Verified consumers:
  * h_s2_rms_ema → mag_concat_kernel (forward graph) +
    dqn_clamp_finite_f32_kernel (cuBLAS backward IQN-trunk sanitiser
    bound = 1e6 × ISV[96])
  * fold_warmup_factor → training_loop.rs:2662 per-step host read
    deriving lr_eff and clip_eff
  * moe_expert_util/entropy → consumed by launch_moe_lambda_eff_update
    immediately below (same step)
  * moe_lambda_eff → moe_load_balance_loss kernel reads ISV[128] per
    step at gpu_dqn_trainer.rs:16864

Atomic migration per `feedback_no_partial_refactor`. Producers leaving
process_epoch_boundary are replaced with one-line redirect comments
pointing to the new hook. Ordering dependencies preserved
(launch_moe_lambda_eff_update still runs AFTER launch_moe_expert_util_ema
per the kernel docstring at gpu_dqn_trainer.rs:16962). All 4 launchers
use pre-loaded CudaFunction fields (gpu_dqn_trainer.rs:12646, 30983,
16901, 16969) — graph-capture safe per
`pearl_no_host_branches_in_captured_graph`.

Cold-start ordering: h_s2_rms_ema reads `save_h_s2` populated by THIS
step's online forward (forward_child runs BEFORE aux_child per
`capture_training_graph`). MoE producer reads `moe_gate_softmax_buf`
populated by `launch_moe_forward` (in `submit_forward_ops_main`,
forward_child). Both inputs are valid by the time submit_aux_ops runs.

State-reset registry coverage already in place
(state_reset_registry.rs:304/427/380/387/398/442 — isv_h_s2_rms_ema,
isv_fold_warmup_factor, isv_moe_expert_util_ema, isv_moe_gate_entropy_ema,
isv_moe_lambda_eff, isv_grad_norm_fast_ema). HEALTH_DIAG read points
unchanged — per-epoch reads of per-step-updated slots get the latest
value (strict improvement over per-epoch reads of per-epoch-stale slots).

Producers DELIBERATELY KEPT per-epoch (audit trail):

  Collector EMAs (input data only updates per-epoch — running per-step
  computes the same EMA repeatedly off the same data):
    * launch_trade_attempt_rate_ema_inplace
    * launch_plan_threshold_update_inplace
    * launch_seed_step_counter_update_inplace
    * launch_cql_alpha_seed_update_inplace

  HEALTH_DIAG-only consumers (no per-step kernel reads these slots):
    * launch_iqn_quantile_ema  (ISV[99..103))
    * launch_vsn_mask_ema      (ISV[105..111))
    * launch_aux_heads_loss_ema (ISV[113], ISV[114])

DEFERRED follow-up (NOT changed in this commit per the user-specified
strict heuristic "docstring per-step + per-step consumer ⇒ migrate";
docstrings of the SP4/SP5 producers below don't explicitly claim
per-step cadence — they're silent — so they fall outside the heuristic
even though their consumers ARE per-step):
  * launch_sp4_target_q_p99 + atom_pos_p99 + grad_norm_p99 + h_s2_p99
  * launch_sp4_param_group_oracles_all_groups
    (consumed by `read_group_adam_bounds` per-step in submit_aux_ops)
  * launch_sp5_pearl_1_atom + pearl_3_sigma + pearl_2_budget +
    pearl_4_adam_hparams + pearl_5_iqn_tau + pearl_8_trail +
    pearl_1_ext_num_atoms
  * launch_max_budget_compute + launch_loss_balance_controller (SP7/SP8)
This staleness will be addressed in a follow-up audit pass; for now the
SP4/SP5 epoch-block producers stay where they are.

Verification:
  * SQLX_OFFLINE=true CUDA_COMPUTE_CAP=86 cargo check -p ml --tests
    --all-targets clean (only pre-existing unused-var warnings).
  * sp14_oracle_tests 6/6 pass (alpha_grad_adaptive_beta,
    alpha_grad_schmitt_hysteresis, dir_concat_qaux_correct,
    gradient_hack_circuit_breaker_fires,
    q_disagreement_all_hold_no_contribution, q_disagreement_k4_k2_mapping).
  * HEALTH_DIAG validation pending L40S re-dispatch.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-05-07 18:55:29 +02:00
parent 200f05fcef
commit 5608b866b6
3 changed files with 105 additions and 46 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -2039,6 +2039,82 @@ impl FusedTrainingCtx {
self.trainer.launch_sp14_alpha_grad_compute(alpha_var)
.map_err(|e| anyhow::anyhow!("SP14 B.11 alpha_grad_compute: {e}"))?;
// 2026-05-07 cadence-fix continuation (after SP14 B.11 commit
// 200f05fce): four MORE producers whose docstrings explicitly
// claim "per-step" cadence but whose launches lived in
// `process_epoch_boundary` (fires once per epoch). Each has at
// least one captured-graph or per-step-host consumer that read
// (steps_per_epoch − 1)-step-stale ISV values per the same
// failure mode SP14 B.11 documented for ALPHA_GRAD_SMOOTHED.
//
// Atomic migration per `feedback_no_partial_refactor`. Producers
// are ordered to preserve the existing epoch-block dependency
// chain (moe_lambda_eff_update reads MOE_GATE_ENTROPY_EMA which
// moe_expert_util_ema produces — same-step ordering preserved).
//
// * launch_h_s2_rms_ema → ISV[H_S2_RMS_EMA_INDEX=96]
// Consumer: mag_concat_kernel (per-step in forward graph,
// reads ISV[96] for adaptive scale of save_h_s2 → magnitude
// branch concat). Also `dqn_clamp_finite_f32_kernel` cuBLAS
// backward sanitiser uses 1e6 × ISV[96] for the IQN trunk
// path. Both per-step. Producer kernel reads `save_h_s2`
// populated by THIS step's online forward.
//
// * launch_fold_warmup_factor → ISV[FOLD_WARMUP_FACTOR_INDEX=130]
// Consumer: training_loop.rs:2662 `read_isv_signal_at(130)`
// in per-step `update_adaptive_clip` block, derives
// lr_eff and clip_eff. Producer kernel reads fast/slow
// grad-norm EMAs filled by `update_grad_norm_emas_kernel`
// per step. Note: `WARMUP_FACTOR_MIN_STEPS` and
// `steps_observed_i32` kernel args are FROZEN at capture
// time — the gate is monotone (once passed, stays passed),
// and the host-side step counter increments outside graph
// capture, so freezing the captured value at >200 keeps
// the gate permanently open as designed.
//
// * launch_moe_expert_util_ema → ISV[MOE_EXPERT_UTIL_EMA_BASE
// ..MOE_EXPERT_UTIL_EMA_BASE+8) = [118..126) +
// ISV[MOE_GATE_ENTROPY_EMA_INDEX=126]
// Consumer: launch_moe_lambda_eff_update reads ISV[126]
// (next entry below — same-step). Also HEALTH_DIAG. Producer
// reads `moe_gate_softmax_buf` populated by
// `submit_forward_ops_main`'s `launch_moe_forward` THIS step.
//
// * launch_moe_lambda_eff_update → ISV[MOE_LAMBDA_EFF_INDEX=128]
// Consumer: `moe_load_balance_loss` kernel (in MoE backward
// forward — captured per step) reads ISV[128] at runtime
// to scale the load-balance loss. Producer reads ISV[126]
// written by `launch_moe_expert_util_ema` immediately
// above on the same stream.
//
// All four launchers use pre-loaded `CudaFunction` fields
// (no per-call `load_cubin` / `load_function`) — graph-capture
// safe per `pearl_no_host_branches_in_captured_graph` (verified
// gpu_dqn_trainer.rs:12646, 30983, 16901, 16969).
//
// Cold-start ordering: h_s2_rms_ema reads `save_h_s2` from the
// forward this step (forward_child runs BEFORE aux_child per
// `capture_training_graph`). MoE producer reads
// `moe_gate_softmax_buf` populated by `launch_moe_forward`
// inside the same forward_child. Both inputs are valid by
// the time submit_aux_ops runs.
//
// The per-epoch slots written here are STILL read by HEALTH_DIAG
// emits in `process_epoch_boundary` further down — those reads
// continue to see the latest per-step value (per-epoch read of
// per-step-updated slot is a strict improvement).
self.trainer.launch_h_s2_rms_ema()
.map_err(|e| anyhow::anyhow!("Plan 4 2c.3c.5 h_s2_rms_ema: {e}"))?;
self.trainer.launch_fold_warmup_factor()
.map_err(|e| anyhow::anyhow!("Plan C K fold_warmup_factor: {e}"))?;
// MoE producer + λ-controller chain — order load-bearing per
// gpu_dqn_trainer.rs:16962 ("MUST be called AFTER
// launch_moe_expert_util_ema").
self.trainer.launch_moe_expert_util_ema()
.map_err(|e| anyhow::anyhow!("Phase 3 T3.5 moe_expert_util_ema: {e}"))?;
self.trainer.launch_moe_lambda_eff_update()
.map_err(|e| anyhow::anyhow!("moe_lambda_eff_update: {e}"))?;
// 2. Gather Q(s, a_taken) for both IQL trainers.
let q_out = self.trainer.q_out_buf();
let actions = self.trainer.actions_buf();

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

@@ -4137,32 +4137,23 @@ impl DQNTrainer {
}
}
// Plan 4 Task 2c.3c.5: per-step H_S2 RMS EMA into ISV[96].
// Single-block 256-thread shmem-reduction kernel reads the
// online trunk's `save_h_s2 [B, SH2]` (populated by the most
// recent forward) and EMAs sqrt(sum_sq / N) into ISV[96] at
// α=0.05. Producer-only — 2c.3c.6 wires the consumer in
// `mag_concat_qdir`'s adaptive-scale path. Same launch
// cadence as the Plan 3 producers above.
if let Some(ref fused) = self.fused_ctx {
if let Err(e) = fused.trainer().launch_h_s2_rms_ema() {
tracing::warn!("Plan 4 2c.3c.5 h_s2_rms_ema launch failed: {e}");
}
}
// 2026-05-07 cadence-fix continuation (after SP14 B.11
// commit 200f05fce): launch_h_s2_rms_ema MOVED to
// submit_aux_ops (per-step, captured into aux_child graph).
// Consumer (mag_concat_kernel + dqn_clamp_finite_f32_kernel)
// reads ISV[96] every training step — per-epoch producer
// here meant (steps_per_epoch − 1)-step-stale values for
// every captured-graph replay. See submit_aux_ops in
// fused_training.rs for the new hook.
// Plan C Phase 2 follow-up K (2026-04-29): per-step adaptive
// fold-boundary warmup factor producer. Reads the fast/slow
// grad-norm EMA mapped-pinned scalars (filled by
// `update_adaptive_clip` in the per-step guard block above)
// and writes ISV[FOLD_WARMUP_FACTOR_INDEX=130] = clamp(fast/slow, 0, 1).
// Two consumers below derive lr_eff (via set_lr) and
// clip_eff (via update_adaptive_clip's pinned slot). Same
// launch cadence as h_s2_rms_ema — once per step.
if let Some(ref fused) = self.fused_ctx {
if let Err(e) = fused.trainer().launch_fold_warmup_factor() {
tracing::warn!("Plan C K fold_warmup_factor launch failed: {e}");
}
}
// 2026-05-07 cadence-fix continuation (after SP14 B.11
// commit 200f05fce): launch_fold_warmup_factor MOVED to
// submit_aux_ops (per-step, captured into aux_child graph).
// Consumer at training_loop.rs:2662 reads
// ISV[FOLD_WARMUP_FACTOR_INDEX=130] every training step to
// derive lr_eff and clip_eff — per-epoch producer here
// meant the warmup factor lagged by an entire epoch. See
// submit_aux_ops in fused_training.rs for the new hook.
// Plan 4 Task 3 (E.3): per-step IQN multi-quantile diagnostic
// EMAs into ISV[99..103). 4-block kernel reads the IQN online
@@ -4239,27 +4230,17 @@ impl DQNTrainer {
// breaker stays in `process_epoch_boundary` below (its
// lockout decrement IS one-per-epoch by design).
// Phase 3 T3.5: MoE expert-utilisation EMA + gate entropy EMA.
// Reads `moe_gate_softmax_buf [B, K]` (valid after the captured
// forward graph ran) and EMA-updates ISV[118..127).
if let Some(ref fused) = self.fused_ctx {
if let Err(e) = fused.trainer().launch_moe_expert_util_ema() {
tracing::warn!("Phase 3 T3.5 launch_moe_expert_util_ema failed: {e}");
}
}
// Adaptive MoE load-balance λ controller (2026-04-27).
// MUST run AFTER `launch_moe_expert_util_ema` so the gate-
// entropy EMA fed into the controller is fresh for the NEXT
// step's `moe_load_balance_loss` consumer. Reads ISV[126]
// (entropy EMA), writes ISV[MOE_LAMBDA_EFF_INDEX=128].
// Per `pearl_blend_formulas_must_have_permanent_floor.md`,
// λ_floor is a permanent minimum.
if let Some(ref fused) = self.fused_ctx {
if let Err(e) = fused.trainer().launch_moe_lambda_eff_update() {
tracing::warn!("launch_moe_lambda_eff_update failed: {e}");
}
}
// 2026-05-07 cadence-fix continuation (after SP14 B.11
// commit 200f05fce): launch_moe_expert_util_ema +
// launch_moe_lambda_eff_update MOVED to submit_aux_ops
// (per-step, captured into aux_child graph). Consumer
// `moe_load_balance_loss` reads ISV[MOE_LAMBDA_EFF_INDEX=128]
// every training step (and the controller reads
// ISV[MOE_GATE_ENTROPY_EMA_INDEX=126] from the producer);
// per-epoch producers here meant the load-balance λ lagged
// by an entire epoch. Order preserved (expert_util BEFORE
// lambda_eff_update). See submit_aux_ops in
// fused_training.rs for the new hook.
// SP4 Layer A Tasks A10 + A11: per-step ISV-bound producer
// launches for the 40 new SP4 slots (TARGET_Q_BOUND, ATOM_POS

2
docs/dqn-wire-up-audit.md
View File

File diff suppressed because one or more lines are too long