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

diag(nan): add flag 12 = save_h_s2 — differentiate trunk vs branch GEMM as NaN source

Browse Source 
Previous L40S 15-epoch repro fired the wired NaN diagnostic with
flagged=[2=on_b_logits, 3=mse_loss_scalar, 6=grad_buf, 7=save_current_lp,
8=save_projected] while value stream (flag 1) and pre-forward params
(flags 4-5) stayed clean. That isolates the corruption to the branch
advantage GEMM/activation but leaves one open question: is the GEMM
input (h_s2, shared between value and branch heads) finite or already
NaN?

Flag 12 = save_h_s2 covers the post-trunk activation. Outcomes:
  - flag 12 clean + flag 2 NaN → branch GEMM overflowed (advantage-side
    fp32 overflow under post-S&P + adversarial reward stress).
  - flag 12 NaN → trunk encoder corrupted upstream of both heads.

Mechanically: one new check_nan_f32 call against save_h_s2 in
run_nan_checks_post_forward, names array slot 12 updated rsv12 →
save_h_s2 in training_loop.rs error message. Same ~5µs cost as the
other 12 checks. Slots 13-15 still reserved.

Per `feedback_no_partial_refactor.md` (extending the NaN diagnostic
contract atomically) and `feedback_no_stubs.md` (rsv12 was a real
reservation; now consumed).
This commit is contained in:
jgrusewski
2026-04-28 13:02:19 +02:00
parent 1c50a4e8a4
commit 32e5375ac8
3 changed files with 40 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -14277,7 +14277,8 @@ impl GpuDqnTrainer {
/// 9 moe_gate_softmax — MoE gate softmax (saturates → NaN)
/// 10 aux_nb_loss_scalar [1] — aux next-bar MSE loss
/// 11 aux_rg_loss_scalar [1] — aux regime CE loss
/// 12-15 reserved (CQL / IQN / future)
/// 12 save_h_s2 — trunk output (input to value+branch FCs)
/// 13-15 reserved (CQL / IQN / future)
pub fn run_nan_checks_post_forward(&mut self, batch_size: usize) -> Result<(), MLError> {
let b = batch_size;
let na = self.config.num_atoms;
@@ -14314,6 +14315,12 @@ impl GpuDqnTrainer {
self.check_nan_f32(self.aux_nb_loss_scalar_buf.raw_ptr(), 1, 10)?;
// Flag 11: aux_rg_loss_scalar [1] (aux regime CE loss)
self.check_nan_f32(self.aux_rg_loss_scalar_buf.raw_ptr(), 1, 11)?;
// Flag 12: save_h_s2 (post-trunk activation, input to value FC + branch FC).
// Diagnostic differentiator: if `on_b_logits` (flag 2) is NaN but
// `save_h_s2` is finite, the branch GEMM itself overflowed; if both are
// NaN, the corruption is upstream in the trunk encoder.
let sh2 = self.config.shared_h2;
self.check_nan_f32(self.save_h_s2.raw_ptr(), b * sh2, 12)?;
Ok(())
}

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

@@ -1860,7 +1860,8 @@ impl DQNTrainer {
"moe_gate_softmax", // 9 MoE gate over 8 experts
"aux_nb_loss_scalar", // 10 aux next-bar MSE [1]
"aux_rg_loss_scalar", // 11 aux regime CE [1]
"rsv12", "rsv13", "rsv14", "rsv15", // reserved
"save_h_s2", // 12 trunk output (input to FCs)
"rsv13", "rsv14", "rsv15", // reserved
];
let flagged: Vec<String> = flags.iter().enumerate()
.filter(|(_, &f)| f != 0)

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

@@ -2,6 +2,36 @@
**Status:** Populated during Plan 1 Task 6 (A.5 orphan audit). Updated on every commit per Invariant 7.
NaN diagnostic flag 12 = `save_h_s2` (2026-04-28): the previous wire-up
commit (`756b1ef31`) extended NaN coverage to flags 0-11 and proved that
the L40S fold-1 NaN entered via `flagged=[2=on_b_logits, 3=mse_loss_scalar,
6=grad_buf, 7=save_current_lp, 8=save_projected]` while `params_buf_pre_fwd`
(flag 4) and `on_v_logits` (flag 1) stayed clean. That isolates the
corruption to the branch advantage GEMM/activation (the value stream is
fine despite sharing the same trunk input). To complete the differential
diagnosis we add flag 12 = `save_h_s2` — the post-trunk activation that
is the SHARED input to both the value FC and the branch FC. Two
outcomes:
- flag 12 clean + flag 2 NaN → branch GEMM itself overflowed
(likely fp32-overflow under post-S&P + adversarial reward stress
in the advantage logits accumulation). Fix is at the branch
forward call site / advantage scale clamp.
- flag 12 NaN → trunk encoder corrupted upstream of both heads
(GRN block produced NaN). Trace further into the GRN backward
chain or the OFI/state input.
`run_nan_checks_post_forward` in `gpu_dqn_trainer.rs` adds one
`check_nan_f32` call against `self.save_h_s2.raw_ptr()` with size
`batch_size * config.shared_h2`, flag index 12. Same ~5µs cost as the
other 12 checks already wired. `training_loop.rs` names array slot 12
updated `"rsv12"` → `"save_h_s2"`. Slots 13-15 still reserved (CQL,
IQN, per-branch backward). Touched files: `gpu_dqn_trainer.rs`
(+1 line in flag map docstring, +5 lines in
`run_nan_checks_post_forward`); `training_loop.rs` (1 line in names
array). cargo check clean at 13 warnings (workspace baseline). No
fingerprint change.
P5T5 Phase H Site 3 — fuse VSN feature-selection + GLU value/gate
GEMM+bias (2026-04-28): migrates the remaining 4 forward sites:
1. **VSN Linear_1** (×6 feature groups) — adds 6 per-group RELU_BIAS