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

fix(aux-head): stop-gradient on aux's h_s2 input — fixes aux-loss-rises-during-training pathology

Browse Source 
Root cause from train-v8ztm 9-epoch HEALTH_DIAG aux next_bar_mse trajectory:
- Ep 0: 0.352 (learnable signal — below random baseline ln(2)≈0.693)
- Ep 9: 0.717 (above random baseline — aux is now WORSE than random)
- aux_dir_acc_long stuck at 0.19 (anti-correlated with truth)

Aux head's backward gradient was flowing back to shared trunk activation
h_s2 via dh_s2_out write at aux_heads_kernel.cu:599-613. Q-loss
gradient on h_s2 dominates (larger magnitude, structurally different
objective: cumulative discounted reward vs next-bar direction). h_s2
evolves to support Q's task; aux's CE loss climbs as h_s2 features
become anti-aligned with direction prediction.

Fix: stop-gradient. Aux reads h_s2 via forward, trains its own w1/b1/w2/b2
from CE loss, but does NOT propagate to h_s2. Q-loss is the sole shaping
force on h_s2. Aux must adapt to whatever h_s2 happens to be — if the
representation has direction signal, aux's params will extract it; if
not, aux can't learn (separate-trunk Option 2 deferred for that case).

This was the SEVENTH fix in today's chain (after 6 SP14 EGF cadence/
gate/saturation fixes). The EGF was a scaffold over a broken aux head;
fixing aux first is the architectural prerequisite for EGF to route
useful signal.

Verification: cargo check clean; sp14_oracle_tests 7/7 pass.
Validation: aux next_bar_mse should now DECREASE during training in
the next L40S smoke (vs the rising-from-0.35-to-0.72 pattern in v8ztm).

Deferred follow-up: aux_regime_backward has the same architecture
(propagates dh_s2 to trunk). Same fix is a candidate once next_bar
result validates the approach.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-05-07 22:49:37 +02:00
parent c260dca8bd
commit 872bd73927
2 changed files with 36 additions and 10 deletions
Download Patch File Download Diff File Expand all files Collapse all files

44
crates/ml/src/cuda_pipeline/aux_heads_kernel.cu
View File

@@ -596,19 +596,43 @@ extern "C" __global__ void aux_next_bar_backward(
}
}
/* Step 3: dh_s2[b, j] = sum_k sh_dh_pre[k] * w1[k, j].
* Per-sample, per-feature serial reduction over H lanes — H is small
* (32) and SH2 is moderate; stride loop over j. Masked rows propagate
* zero d_h_pre through this sum, so dh_s2 row is zero for skipped
* samples (the trunk SAXPY adds zero — no spurious gradient). */
/* Step 3: SP14 stop-gradient (2026-05-07).
*
* Diagnostic from L40S `train-v8ztm` 9-epoch HEALTH_DIAG aux next_bar_mse:
* Ep 0: 0.352 (learnable signal — below random baseline ln(2)≈0.693)
* Ep 9: 0.717 (above random baseline — aux is now WORSE than random)
* aux_dir_acc_long stuck at 0.19 (anti-correlated with truth)
*
* Root cause: the aux head's backward gradient was flowing back into the
* shared trunk activation `h_s2` via the SAXPY of `dh_s2_out` into the
* trunk's main-path d_h_s2. Q-loss gradient on h_s2 dominates (larger
* magnitude, structurally different objective: cumulative discounted
* reward vs next-bar direction). h_s2 evolves to support Q's task; aux's
* CE loss climbs as h_s2 features become anti-aligned with direction
* prediction.
*
* Fix: stop-gradient on aux's input. We zero `dh_s2_out` rather than
* computing the sum-of-w1·dh_pre. The aux head still trains its own
* params (w1/b1/w2/b2) from the CE loss via the dW1/db1/dW2/db2
* partials written above — those are unaffected. But h_s2 is no longer
* pulled by aux's gradient. Q-loss is the sole shaping force on h_s2;
* aux must adapt to whatever h_s2 happens to be.
*
* If aux still cannot learn after this fix, the deeper diagnosis is that
* h_s2 carries no direction signal at all — separate-trunk Option 2
* (deferred). The trunk SAXPY adds zero either way (same arithmetic the
* masked-row branch already produced for skipped samples), so no other
* consumer is affected — `dh_s2_out` is read only by the trunk
* accumulation step.
*
* NOTE: `aux_regime_backward` (the K=5 regime CE head) has the same
* architecture and is a candidate for the same fix once next-bar
* validation lands; deliberately scoped tight here because the v8ztm
* diagnostic evidence is direct only for the next-bar head. */
{
float* dh_row = dh_s2_out + (size_t)b * SH2;
for (int j = tid; j < SH2; j += blockDim.x) {
float acc = 0.0f;
for (int k = 0; k < H; ++k) {
acc += sh_dh_pre[k] * w1[(size_t)k * SH2 + j];
}
dh_row[j] = acc;
dh_row[j] = 0.0f;
}
}
}

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

@@ -2,6 +2,8 @@
**Status:** Populated during Plan 1 Task 6 (A.5 orphan audit). Updated on every commit per Invariant 7.
SP14 aux-head stop-gradient — fix aux-loss-rises-during-training pathology (2026-05-07): root-caused the `train-v8ztm` 9-epoch HEALTH_DIAG `aux next_bar_mse` trajectory (Ep 0: 0.352 — learnable signal below random baseline `ln(2)≈0.693`; Ep 9: 0.717 — above random baseline, aux is now WORSE than random; `aux_dir_acc_long` stuck at 0.19 — anti-correlated with truth) to the aux head's backward gradient flowing back into the shared trunk activation `h_s2` via `dh_s2_out` at `crates/ml/src/cuda_pipeline/aux_heads_kernel.cu:599-613` and the trunk SAXPY that accumulates it into the main-path `d_h_s2`. Q-loss gradient on h_s2 dominates (larger magnitude, structurally different objective: cumulative discounted reward vs next-bar direction). h_s2 evolves to support Q's task; aux's CE loss climbs as h_s2 features become anti-aligned with direction prediction — aux's own params can only extract whatever signal is left in h_s2, which becomes inverted by the dominant Q-shaping force. Fix (atomic, single kernel): replace the `dh_s2[b, j] = sum_k sh_dh_pre[k] * w1[k, j]` reduction in `aux_next_bar_backward`'s Step 3 with an unconditional zero write `dh_row[j] = 0.0f`. The aux head still trains its own params (`w1`, `b1`, `w2`, `b2`) from the CE loss via the `dW1_partial` / `db1_partial` / `dW2_partial` / `db2_partial` writes above (those are unaffected — they read `sh_dh_pre`, `sh_dlogits`, `sh_h_post`, `h_row` which are computed pre-write); only the trunk-bound gradient is severed. Q-loss is now the sole shaping force on h_s2; aux must adapt to whatever h_s2 happens to be — if the representation has direction signal, aux's params will extract it; if not, aux can't learn (separate-trunk Option 2 deferred for that case). The trunk SAXPY adds zero, identical arithmetic to the masked-row branch's existing zero-row behaviour (existing code already writes a zero row when `labels[b] == -1` because all `sh_dlogits` entries are zero, propagating to `sh_dh_pre = 0` and downstream `acc = 0`). No other consumer reads `dh_s2_out` — only the trunk accumulation step. This was the SEVENTH fix in today's chain (after 6 SP14 EGF cadence/gate/saturation fixes — β-migration steps 1-5 plus the empty-batch decay gate). The EGF was a scaffold over a broken aux head; fixing aux first is the architectural prerequisite for EGF to route useful signal. **Deferred follow-up**: `aux_regime_backward` (the K=5 regime CE head, lines 634-743) has the SAME architecture — it propagates `dh_s2` to the trunk via the same SAXPY pattern. Same fix is a candidate once next-bar validation lands; deliberately NOT changed in this commit because the v8ztm diagnostic evidence is direct only for the next-bar head (the regime head's loss trajectory was not flagged as rising in the same diagnostic; applying the fix without evidence would risk losing useful regime-direction trunk shaping if regime's gradient happens to be Q-aligned). **Verification**: `SQLX_OFFLINE=true CUDA_COMPUTE_CAP=86 cargo check -p ml --tests --all-targets` clean (dev profile, only pre-existing unused-var warnings unrelated to this change); `cargo test -p ml --test sp14_oracle_tests --release -- --ignored --nocapture` 7/7 GPU tests pass on RTX 3050 Ti (`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_empty_batch_preserves_ema`, `q_disagreement_k4_k2_mapping`) — these don't directly cover `aux_next_bar_backward`'s Step 3 but they exercise the consumer EGF chain that depends on aux signal quality, so a regression in CE loss math (which would also affect dW1/dW2 partials) would surface as ISV-state divergence. The change is "set output to 0" so no math regression in oracle tests is expected. **Validation pending L40S re-dispatch**: `aux next_bar_mse` should now DECREASE during training (vs the rising-from-0.35-to-0.72 pattern in v8ztm); `aux_dir_acc_long` should approach 0.5 (random baseline) at minimum and ideally above it if h_s2 carries direction signal at all. If the metric stays flat at the bootstrap value (sentinel = 0.5 from `pearl_first_observation_bootstrap`), the deeper diagnosis is that h_s2 has no direction signal — separate-trunk follow-up. Files modified: `crates/ml/src/cuda_pipeline/aux_heads_kernel.cu` (−14 / +37 lines: replace Step 3 reduction with zero-write + 30-line audit-comment block documenting diagnostic, mechanism, fix, and deferred regime follow-up).
SP14 β-migration step 5 — remove training-time EGF launches from `submit_aux_ops` (2026-05-07): Final atomic step of the 5-commit β migration. Deletes the 3 trainer-time launches at `crates/ml/src/trainers/dqn/fused_training.rs:2031-2040` (`launch_sp13_aux_dir_metrics` + `launch_sp14_q_disagreement_update` + `launch_sp14_alpha_grad_compute`) — these were the producers introduced by SP14 B.11 commit `200f05fce` (per-step migration) and stabilised by commit `9d0c124ce` (empty-batch decay gate). Even with both fixes the trainer-time path could only fire against replay-batch Q-direction picks dominated by Hold/Flat (the natural argmax distribution of the converged 4-way head), so `total_cnt=0` every step and the gate kept EMAs at zero. The collector-native β path (commit c691bd381 / step 4) reads rollout-time q_values where Thompson sampling forces Short/Long picks every step — that's the genuine source of aux↔Q disagreement signal. The collector-native chain is now the SINGLE production caller. **Trainer launcher methods preserved**: `launch_sp14_q_disagreement_update` (line 8671), `launch_sp14_alpha_grad_compute` (line 8733), `launch_sp13_aux_dir_metrics` (line 15270) and the 3 sub-launchers (`launch_aux_dir_acc_reduce`, `launch_apply_fixed_alpha_ema`, `launch_aux_pred_to_isv_tanh`) stay defined as `pub(crate)` (no Rust compiler warnings — `pub(crate)` methods don't trigger dead-code warnings). Oracle tests in `crates/ml/tests/sp14_oracle_tests.rs` exercise the SAME kernels directly via `load_cubin` / `load_function`, not through these launchers — verified by `grep -n "launch_sp14|launch_sp13_aux" tests/sp14_oracle_tests.rs` returning 0 hits — so deletion of the launchers would not affect oracle coverage. They are retained for atomic-rollback potential and the possibility that the trainer-time path becomes useful again if a future curriculum stage reverses the rollout-only signal-quality assumption. **`launch_sp14_gradient_hack_detect` (per-epoch circuit breaker) unchanged** — the lockout-counter decrement is one-per-epoch by design, stays at `process_epoch_boundary`. The replacement comment at the deletion site in `fused_training.rs` documents the migration trail (β-migration step 5 → c691bd381 → 9d0c124ce) so a future archeologist can reconstruct the chain. **Verification**: `SQLX_OFFLINE=true cargo check -p ml --lib` clean (dev profile, only pre-existing unused-var warnings unrelated to this change); `cargo test -p ml --test sp14_oracle_tests` 2/2 non-GPU pass (7 GPU tests ignored on RTX 3050 Ti host — verify kernel correctness, unchanged); `grep -rn "trainer.launch_sp14|trainer.launch_sp13_aux" crates/ml/src/` returns 0 production callers (the 3 calls in fused_training.rs are deleted). L40S smoke validation pending — expect EGF chain to drive non-zero ISV[383/384/389/393] across the full curriculum, unblocking Gate 1 from its closed-forever state. Files modified: `crates/ml/src/trainers/dqn/fused_training.rs` (−10 / +14 lines: deleted 3 launches, replaced 60-line audit-comment block with 12-line redirect comment).
SP14 β-migration step 4 — wire collector-native SP14 EGF producer chain after aux forward (2026-05-07): Step 4 of 5. Inserts the 3 SP14/SP13-EGF producer launches into the collector's per-rollout-step body, immediately after the `expected_q_kernel` (which produces `q_values`) and BEFORE the IQR/ensemble/noise SAXPY blocks (so the SP14 q_disagreement consumer sees noise-free `q_values` matching the trainer's `q_out_buf` semantic). Producer order matches the trainer's `submit_aux_ops` chain (preserved per `feedback_no_partial_refactor.md`): (1) SP13 dir-acc reduce + 2 fixed-α EMAs (α=0.3 short → ISV[373], α=0.05 long → ISV[374]) + aux_pred → ISV[375] tanh — same 4 launches the trainer's `launch_sp13_aux_dir_metrics` orchestrates inline (gpu_dqn_trainer.rs:15270), reading the collector's `exp_aux_nb_softmax_buf [n_episodes, 2]` (from step 3) + `exp_aux_nb_label_buf [n_episodes]` (from step 3's per-step label producer) + writing to the shared ISV via `isv_signals_dev_ptr`. Sentinel 0.5 = random-guessing baseline per `pearl_first_observation_bootstrap`. (2) SP14 `q_disagreement_update_kernel` — reads `exp_aux_nb_softmax_buf [n_episodes, K=2]` + collector's `q_values [n_episodes, q_stride=13]` (post-`expected_q_kernel` output); writes ISV[383/384/389] (q_dis_short, q_dis_long, var_q EMAs). q_stride=13 matches the trainer's q_out_buf stride; the kernel's argmax-over-K_DIR=4 finds the direction-head pick from the first 4 columns of each row. Single block, 256 threads, smem = 2 × 256 × f32. (3) SP14 `alpha_grad_compute_kernel` — pure ISV-state state machine reading the q_dis EMAs just written above (same-stream serial ordering enforces the dependency) plus ISV[372..395) drivers; writes ISV[392/393] (α_grad raw + smoothed) plus the Schmitt-trigger / k_q / k_aux state machine bookkeeping at ISV[385..391]. Single thread (we launch with 32 to mirror the B.4 oracle test launch shape). The kernel's empty-batch decay gate (commit 9d0c124ce) preserves prior EMA values when `total_cnt == 0` for a given step — same protection extends to the rollout path. **Gating**: `if self.isv_signals_dev_ptr != 0 && self.trainer_params_ptr != 0` — without ISV the EMAs have nowhere to write; without trainer params the aux forward in step 3 would have been skipped and the softmax tile would still be `alloc_zeros`. **No `seed_phase_active_cache` gate** here (unlike SP13 hold-rate which prices Hold-cost only post-seed): the EGF Schmitt-trigger Gate 1 needs to observe BOTH seed-phase scripted-policy actions and post-seed Q-policy actions to discriminate aux signal quality across the curriculum. **q_logits placement note**: the SP14 chain runs AFTER `expected_q_kernel` (line 4581) and BEFORE the IQR (3b) / ensemble (3c) / noise (3d) SAXPY bonuses — so `q_values` here is bit-equivalent to the trainer's `q_out_buf` semantic at the trainer's `launch_sp14_q_disagreement_update` call site. The IQR/ensemble bonuses below are exploration noise added for action selection, not gradient flow. Per `pearl_no_host_branches_in_captured_graph` the launches are OUTSIDE any captured graph (the captured `forward` graph has already ended). The trainer's existing private `launch_sp14_q_disagreement_update` / `launch_sp14_alpha_grad_compute` / `launch_sp13_aux_dir_metrics` methods at gpu_dqn_trainer.rs:8671 / 8733 / 15270 are still defined and still pass kernel correctness via the 7 sp14_oracle GPU tests — production has migrated to the collector path; step 5 will delete the now-unused `submit_aux_ops` invocations (commit `200f05fce` lines 2031-2040). Files modified: `crates/ml/src/cuda_pipeline/gpu_experience_collector.rs` (+200 lines: 3 producer launch blocks at line ~4583 with comprehensive doc-comments). Verification: `SQLX_OFFLINE=true cargo check -p ml --lib` clean (dev profile); `cargo test -p ml --test sp14_oracle_tests` 2/2 non-GPU pass (7 GPU tests ignored on RTX 3050 Ti host — they verify kernel correctness, which is unchanged by this commit).