From 32e5375ac83b2087f78abcd09a95c70f776fcc93 Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Tue, 28 Apr 2026 13:02:19 +0200 Subject: [PATCH] =?UTF-8?q?diag(nan):=20add=20flag=2012=20=3D=20save=5Fh?= =?UTF-8?q?=5Fs2=20=E2=80=94=20differentiate=20trunk=20vs=20branch=20GEMM?= =?UTF-8?q?=20as=20NaN=20source?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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). --- .../ml/src/cuda_pipeline/gpu_dqn_trainer.rs | 9 +++++- .../src/trainers/dqn/trainer/training_loop.rs | 3 +- docs/dqn-wire-up-audit.md | 30 +++++++++++++++++++ 3 files changed, 40 insertions(+), 2 deletions(-) diff --git a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs index 04c128f5d..6e25d9acf 100644 --- a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs +++ b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs @@ -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(()) } diff --git a/crates/ml/src/trainers/dqn/trainer/training_loop.rs b/crates/ml/src/trainers/dqn/trainer/training_loop.rs index d85f315f2..7a3d8bbaa 100644 --- a/crates/ml/src/trainers/dqn/trainer/training_loop.rs +++ b/crates/ml/src/trainers/dqn/trainer/training_loop.rs @@ -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 = flags.iter().enumerate() .filter(|(_, &f)| f != 0) diff --git a/docs/dqn-wire-up-audit.md b/docs/dqn-wire-up-audit.md index 5f9c77149..e0843cb3a 100644 --- a/docs/dqn-wire-up-audit.md +++ b/docs/dqn-wire-up-audit.md @@ -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