diff --git a/crates/ml/src/cuda_pipeline/batched_forward.rs b/crates/ml/src/cuda_pipeline/batched_forward.rs index ed938646d..d94d5d759 100644 --- a/crates/ml/src/cuda_pipeline/batched_forward.rs +++ b/crates/ml/src/cuda_pipeline/batched_forward.rs @@ -838,6 +838,25 @@ impl CublasGemmSet { pub(crate) fn batch_size(&self) -> usize { self.batch_size } pub(crate) fn shared_h2(&self) -> usize { self.shared_h2 } + /// Online h_s2 GRN block accessor (read-only) — exposes the wrapper's + /// save-for-backward state ptrs (elu_post / linear_b_out / sigmoid_b / + /// ln_rstd / ln_normed) for the trainer's GRN-stage NaN diagnostic + /// instrumentation (P, post-bkdx5 — slots 16-20 of nan_flags_buf). + /// h_s2 is the bottleneck stage where save_h_s2 NaN was observed at + /// F1 ~step 1745 in smoke-test-bkdx5; instrumenting only h_s2 covers + /// the failing stage without doubling the per-step kernel-launch cost. + pub(crate) fn grn_h_s2_online(&self) -> &super::gpu_grn::GrnBlock { + &self.grn_h_s2_online + } + + /// Online h_s2 Linear_b output buffer pointer (Linear_b output, [B, 2*SH2]). + /// Same buffer that `state.linear_b_out` is DtoD-copied from inside + /// `forward_raw_phase2` — exposed directly so the trainer can NaN-check + /// the live cuBLAS output without depending on the save-DtoD landing. + pub(crate) fn grn_h_s2_linear_b_ptr(&self) -> u64 { + self.grn_h_s2_linear_b_ptr + } + /// Diagnostic: test cublasLtMatmul with raw cudaMalloc buffers. /// Uses the SAME lt_handle and workspace as the real forward. pub(crate) fn test_lt_matmul_raw( diff --git a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs index bf27ac4c5..5b2ef82d3 100644 --- a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs +++ b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs @@ -3717,64 +3717,30 @@ impl GpuDqnTrainer { /// Zeroes momentum (m), variance (v), and step counter (t). /// Weights are preserved (warm-start for the new fold). pub fn reset_adam_state(&mut self) -> Result<(), MLError> { - // Plan C Phase 2 follow-up K (2026-04-29): Adam shrink-and-perturb at - // fold boundary. Replaces the prior full-zero reset of `m_buf` / - // `v_buf` with a multiplicative shrink — preserves directional - // information while damping magnitude. Composes with the existing - // param shrink-and-perturb (`alpha=0.8` in - // `FusedTrainingCtx::reset_for_fold`) but is more aggressive - // because Adam moments are more cold-start-fragile than parameters - // (zeroed `v` makes the Adam denominator `sqrt(v_hat) + ε ≈ ε`, - // so the FIRST post-reset step is `lr × m̂ / ε` — a catastrophic - // overshoot when the new fold's data distribution differs from - // the previous fold's). + // R1 (post-bkdx5 + ISV-design correction): eliminated K's hardcoded + // shrink ratios. K was a stepping-stone fix introduced before + // fold_warmup_factor (commit 4ef1d8ebb) existed. fold_warmup_factor + // already dampens lr+clip via ISV-driven ramp during the first ~50 + // post-reset steps, which IS the principled first-step protection. // - // Numerical-stability bounds (Invariant 1 carve-out per - // `feedback_isv_for_adaptive_bounds.md`): the shrink factors 0.1 - // (m) and 0.01 (v) are STRUCTURAL constants ("preserve the - // optimisation direction, lose the magnitude history") not tuned - // values. Shrinking `m` by 0.1 keeps the gradient direction's - // sign+ratio across tensors; shrinking `v` by 0.01 keeps the - // per-tensor scale ratios while letting the new fold's gradients - // re-grow the absolute magnitude. Smoke-test-s9h4h F0 successful - // (Best Sharpe = 36.03), F1 ep1 grad_norm = 355,009 (5 OoM - // larger than F0 steady-state ~10) → root-caused to full-zero - // Adam reset's denominator pathology. + // K's hardcoded shrink ratios (m×0.1, v×0.01) violated + // feedback_adaptive_not_tuned (untracked tunable knobs) AND created + // a downstream pathology: tiny v_hat denominator → oversized Adam + // updates 50+ steps post-reset → trunk param overshoot → save_h_s2 + // NaN at F1 ~step 1745 (smoke-test-bkdx5 diagnostic). // - // `t_pinned` (Adam step counter) IS still zeroed: the bias - // correction `1 - β^t` should restart so the new fold's first - // few steps are properly bias-corrected. Only `m`/`v` magnitude - // history is shrunk (not direction). - const ADAM_M_SHRINK: f32 = 0.1; - const ADAM_V_SHRINK: f32 = 0.01; - let n = self.total_params as i32; - let blocks = ((self.total_params as u32 + 255) / 256) as u32; - let cfg = LaunchConfig { - grid_dim: (blocks, 1, 1), - block_dim: (256, 1, 1), - shared_mem_bytes: 0, - }; - let m_ptr = self.m_buf.raw_ptr(); - let v_ptr = self.v_buf.raw_ptr(); - unsafe { - // dqn_scale_f32_kernel signature: (y, alpha, n) → y[i] *= alpha. - // Use scale_f32_ungraphed (separate CUmodule) — `reset_adam_state` - // is called outside the captured forward_child / adam_child - // graphs (fold boundary), so the ungraphed kernel module - // applies (mirrors `scale_adam_momentum`'s rationale). - self.stream.launch_builder(&self.scale_f32_ungraphed) - .arg(&m_ptr) - .arg(&ADAM_M_SHRINK) - .arg(&n) - .launch(cfg) - .map_err(|e| MLError::ModelError(format!("reset_adam_state shrink m: {e}")))?; - self.stream.launch_builder(&self.scale_f32_ungraphed) - .arg(&v_ptr) - .arg(&ADAM_V_SHRINK) - .arg(&n) - .launch(cfg) - .map_err(|e| MLError::ModelError(format!("reset_adam_state shrink v: {e}")))?; - } + // With K removed: Adam restarts at m=0, v=0 (standard fold reset); + // fold_warmup_factor caps effective lr to MIN_WARMUP_LR_FRAC ≈ 0.05 + // of base for the first cold-start window → bounded effective updates + // → no overshoot. The two mechanisms attacking the same failure mode + // were over-engineered; warmup_factor subsumes K cleanly. + // + // `t_pinned` (Adam step counter) is still zeroed so bias correction + // `1 - β^t` restarts cleanly for the new fold's first steps. + self.stream.memset_zeros(&mut self.m_buf) + .map_err(|e| MLError::ModelError(format!("m_buf reset: {e}")))?; + self.stream.memset_zeros(&mut self.v_buf) + .map_err(|e| MLError::ModelError(format!("v_buf reset: {e}")))?; unsafe { *self.t_pinned = 0; } self.adam_step = 0; @@ -11205,7 +11171,11 @@ impl GpuDqnTrainer { aux_knb, aux_kr, aux_h, max_aux_tensor_len, ); - let nan_flags_buf = stream.alloc_zeros::(16) + // P (post-bkdx5): expanded 16→24 to add 5 GRN-stage NaN-check slots + // (16-20). The new slots pinpoint which GRN sub-stage produces NaN + // first (elu_post → linear_b_out → glu_sigmoid → ln_rstd → ln_normed) + // when R1's K-removal alone doesn't fix the F1 explosion. + let nan_flags_buf = stream.alloc_zeros::(24) .map_err(|e| MLError::ModelError(format!("nan_flags alloc: {e}")))?; // v8: PopArt running statistics buffers @@ -14869,7 +14839,7 @@ impl GpuDqnTrainer { /// Checks cuBLAS output logits, MSE intermediates, gradients, params, and /// loss-component buffers (C51 target dist, MoE gate softmax, aux losses). /// - /// Flag index map (16 slots): + /// Flag index map (24 slots): /// 0 states_buf (f32 padded) — input /// 1 on_v_logits — cuBLAS forward (value stream) /// 2 on_b_logits — cuBLAS forward (branch advantage) @@ -14884,6 +14854,19 @@ impl GpuDqnTrainer { /// 11 aux_rg_loss_scalar [1] — aux regime CE loss /// 12 save_h_s2 — trunk output (input to value+branch FCs) /// 13-15 reserved (CQL / IQN / future) + /// ── P (post-bkdx5) GRN sub-stage NaN checks ── + /// 16 grn_h_s2_elu_post (state.elu_post) — post-ELU activation + /// 17 grn_h_s2_linear_b_out — Linear_b output (pre-GLU) + /// 18 grn_h_s2_glu_sigmoid (state.sigmoid_b) — GLU gate sigmoid + /// 19 grn_h_s2_ln_rstd (state.ln_rstd) — LN reciprocal-stddev + /// (variance→0 ⇒ rstd→∞ → NaN) + /// 20 grn_h_s2_ln_normed (state.ln_normed) — LN normalized (pre-affine) + /// (semantically same buffer + /// family as save_h_s2 but + /// pre-residual+affine; useful + /// to disambiguate residual + /// NaN-source vs LN-source) + /// 21-23 reserved (h_s1 stages / 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; @@ -14926,6 +14909,27 @@ impl GpuDqnTrainer { // 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)?; + + // P (post-bkdx5): GRN h_s2 sub-stage NaN checks (slots 16-20). Drills + // into the trunk-encoder pathway to locate which sub-step of the GRN + // forward composition (Linear_a → ELU → Linear_b → GLU → residual+LN) + // produces NaN first when save_h_s2 (slot 12) goes NaN. + // + // Pragmatic instrumentation scope: only h_s2 (the stage that actually + // failed at F1 ~step 1745 in smoke-test-bkdx5). Slot 16/17/18/19/20 + // map to elu_post / linear_b_out / sigmoid_b / ln_rstd / ln_normed of + // the h_s2 GRN block — see the index map in this method's docstring. + let grn_h_s2 = self.cublas_forward.grn_h_s2_online(); + let elu_post_p = grn_h_s2.elu_post_ptr(); + let lin_b_out_p = self.cublas_forward.grn_h_s2_linear_b_ptr(); + let sigmoid_b_p = grn_h_s2.sigmoid_b_ptr(); + let ln_rstd_p = grn_h_s2.ln_rstd_ptr(); + let ln_normed_p = grn_h_s2.ln_normed_ptr(); + self.check_nan_f32(elu_post_p, b * sh2, 16)?; + self.check_nan_f32(lin_b_out_p, b * 2 * sh2, 17)?; + self.check_nan_f32(sigmoid_b_p, b * sh2, 18)?; + self.check_nan_f32(ln_rstd_p, b, 19)?; + self.check_nan_f32(ln_normed_p, b * sh2, 20)?; Ok(()) } @@ -14973,10 +14977,10 @@ impl GpuDqnTrainer { Ok(()) } - /// Read NaN flags back to CPU (synchronizes stream). Returns [16] flags. + /// Read NaN flags back to CPU (synchronizes stream). Returns [24] flags. /// Index → buffer mapping is documented in `run_nan_checks_post_forward`. - pub fn read_nan_flags(&self) -> Result<[i32; 16], MLError> { - let mut host = [0_i32; 16]; + pub fn read_nan_flags(&self) -> Result<[i32; 24], MLError> { + let mut host = [0_i32; 24]; self.stream.memcpy_dtoh(&self.nan_flags_buf, &mut host) .map_err(|e| MLError::ModelError(format!("nan_flags read: {e}")))?; Ok(host) diff --git a/crates/ml/src/cuda_pipeline/gpu_grn.rs b/crates/ml/src/cuda_pipeline/gpu_grn.rs index f3e6f1706..d0b34cd92 100644 --- a/crates/ml/src/cuda_pipeline/gpu_grn.rs +++ b/crates/ml/src/cuda_pipeline/gpu_grn.rs @@ -258,6 +258,50 @@ impl GrnBlock { self.batch_size } + /// Raw device pointer to `state.elu_post` (post-ELU activation, [B, H]). + /// Populated by `forward_raw` / `forward_raw_phase1` when + /// `save_for_backward = true`. Used by the trainer's GRN-stage NaN + /// diagnostic instrumentation (P, post-bkdx5 — slot 16/17 of nan_flags). + #[inline] + pub(crate) fn elu_post_ptr(&self) -> u64 { + self.elu_post.raw_ptr() + } + + /// Raw device pointer to `state.linear_b_out` (Linear_b output pre-GLU, + /// [B, 2*H]). Populated by `forward_raw` / `forward_raw_phase2` when + /// `save_for_backward = true`. Used by the trainer's GRN-stage NaN + /// diagnostic instrumentation (P, post-bkdx5). + #[inline] + pub(crate) fn linear_b_out_ptr(&self) -> u64 { + self.linear_b_out.raw_ptr() + } + + /// Raw device pointer to `state.sigmoid_b` (sigmoid of GLU gate half, + /// [B, H]). Populated inside `grn_glu_forward`. Used by the trainer's + /// GRN-stage NaN diagnostic instrumentation (P, post-bkdx5). + #[inline] + pub(crate) fn sigmoid_b_ptr(&self) -> u64 { + self.sigmoid_b.raw_ptr() + } + + /// Raw device pointer to `state.ln_rstd` (LN reciprocal-stddev, [B]). + /// Populated inside `grn_residual_layernorm_forward`. Used by the + /// trainer's GRN-stage NaN diagnostic instrumentation (P, post-bkdx5): + /// if forward variance hits zero, rstd → ∞ → downstream NaN. + #[inline] + pub(crate) fn ln_rstd_ptr(&self) -> u64 { + self.ln_rstd.raw_ptr() + } + + /// Raw device pointer to `state.ln_normed` (LN normalized values pre- + /// affine, [B, H]). Populated inside `grn_residual_layernorm_forward`. + /// Used by the trainer's GRN-stage NaN diagnostic instrumentation + /// (P, post-bkdx5). + #[inline] + pub(crate) fn ln_normed_ptr(&self) -> u64 { + self.ln_normed.raw_ptr() + } + /// Forward pass — sequences the 3 GRN forward kernels (CudaSlice variant). /// /// Prefer `forward_raw` for graph-captured callers that operate on raw diff --git a/crates/ml/src/trainers/dqn/fused_training.rs b/crates/ml/src/trainers/dqn/fused_training.rs index 8162bc49c..1bf77f23f 100644 --- a/crates/ml/src/trainers/dqn/fused_training.rs +++ b/crates/ml/src/trainers/dqn/fused_training.rs @@ -3621,7 +3621,7 @@ impl FusedTrainingCtx { /// Read NaN detection flags (synchronizes stream). Used by training guard on NaN halt. /// Index → buffer mapping is documented in /// `GpuDqnTrainer::run_nan_checks_post_forward`. - pub(crate) fn read_nan_flags(&self) -> Result<[i32; 16], crate::MLError> { + pub(crate) fn read_nan_flags(&self) -> Result<[i32; 24], crate::MLError> { self.trainer.read_nan_flags() } diff --git a/crates/ml/src/trainers/dqn/trainer/training_loop.rs b/crates/ml/src/trainers/dqn/trainer/training_loop.rs index 648b1c083..ed627e998 100644 --- a/crates/ml/src/trainers/dqn/trainer/training_loop.rs +++ b/crates/ml/src/trainers/dqn/trainer/training_loop.rs @@ -1981,7 +1981,8 @@ impl DQNTrainer { if let Ok(flags) = fused.read_nan_flags() { // Index → buffer name. Must match // GpuDqnTrainer::run_nan_checks_post_forward and - // run_nan_checks_pre_forward. 16-slot layout. + // run_nan_checks_pre_forward. 24-slot layout + // (P, post-bkdx5: added GRN-stage slots 16-20). let names = [ "states_buf", // 0 input states (post-upload) "on_v_logits", // 1 cuBLAS fwd: value stream @@ -1996,7 +1997,13 @@ impl DQNTrainer { "aux_nb_loss_scalar", // 10 aux next-bar MSE [1] "aux_rg_loss_scalar", // 11 aux regime CE [1] "save_h_s2", // 12 trunk output (input to FCs) - "rsv13", "rsv14", "rsv15", // reserved + "rsv13", "rsv14", "rsv15", // 13-15 reserved + "grn_h_s2_elu_post", // 16 GRN h_s2 post-ELU activation + "grn_h_s2_linear_b_out", // 17 GRN h_s2 Linear_b output + "grn_h_s2_glu_sigmoid", // 18 GRN h_s2 GLU gate sigmoid + "grn_h_s2_ln_rstd", // 19 GRN h_s2 LN reciprocal-stddev + "grn_h_s2_ln_normed", // 20 GRN h_s2 LN normed (pre-affine) + "rsv21", "rsv22", "rsv23", // reserved (h_s1 stages / future) ]; let flagged: Vec = flags.iter().enumerate() .filter(|(_, &f)| f != 0) @@ -2016,15 +2023,16 @@ impl DQNTrainer { } if gr.halt_grad_collapse { // Plan C T11 diagnostic extension (2026-04-29): the existing - // nan_flags_buf [16] system covers 13 buffers and runs NaN - // checks every step inside the captured graph, but - // `read_nan_flags()` was previously only called on - // `halt_nan`. The training guard's halt_nan path checks - // pinned readback `grad_norm` — and when a kernel produces - // NaN that is then clamped to 0 by downstream gradient - // sanitization, the pinned scalar reads 0 (finite), so - // halt_nan never fires. The collapse path (raw_grad_norm - // < 0.01) fires instead, masking the underlying NaN. + // nan_flags_buf [24] system covers 18 buffers (13 base + 5 + // GRN-stage P additions) and runs NaN checks every step + // inside the captured graph, but `read_nan_flags()` was + // previously only called on `halt_nan`. The training + // guard's halt_nan path checks pinned readback `grad_norm` + // — and when a kernel produces NaN that is then clamped to + // 0 by downstream gradient sanitization, the pinned scalar + // reads 0 (finite), so halt_nan never fires. The collapse + // path (raw_grad_norm < 0.01) fires instead, masking the + // underlying NaN. // Read the flags here too — gated to suspicious-zero // (`raw_grad_norm < 1e-6`) so legitimate near-zero // gradients (genuine collapse) get a separate "no NaN @@ -2034,7 +2042,8 @@ impl DQNTrainer { if let Ok(flags) = fused.read_nan_flags() { // Index → buffer name. Must match // GpuDqnTrainer::run_nan_checks_post_forward - // and run_nan_checks_pre_forward. 16-slot layout. + // and run_nan_checks_pre_forward. 24-slot + // layout (P, post-bkdx5: GRN-stage 16-20). let names = [ "states_buf", // 0 "on_v_logits", // 1 @@ -2050,6 +2059,12 @@ impl DQNTrainer { "aux_rg_loss_scalar", // 11 "save_h_s2", // 12 "rsv13", "rsv14", "rsv15", + "grn_h_s2_elu_post", // 16 + "grn_h_s2_linear_b_out", // 17 + "grn_h_s2_glu_sigmoid", // 18 + "grn_h_s2_ln_rstd", // 19 + "grn_h_s2_ln_normed", // 20 + "rsv21", "rsv22", "rsv23", ]; 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 f78172fdf..013427daa 100644 --- a/docs/dqn-wire-up-audit.md +++ b/docs/dqn-wire-up-audit.md @@ -2230,3 +2230,9 @@ Plan C T11 follow-up N (2026-04-29): Winsorize the adaptive_clip EMA input — c Plan C T11 diagnostic extension (2026-04-29): the existing `nan_flags_buf [16]` system covers 13 buffers and runs NaN checks every step in the captured graph, but its `read_nan_flags()` readout was previously only triggered on `halt_nan` (which checks pinned readback `grad_norm` — always `0` after NaN-clamp because downstream gradient sanitization clamps NaN to zero before the pinned scalar copy, so the clamp-to-zero case never fires `halt_nan`). Extend the readout to also fire when `halt_grad_collapse` triggers AND `gr.raw_grad_norm < 1e-6` — the collapse-with-zero-grad case is the NaN-clamped path. Logs `NaN-CLAMPED-TO-ZERO at step N (grad collapse path): flagged=[buf_idx=name, ...]` identifying the kernel responsible. Genuine near-zero gradients (legitimate collapse with no NaN flags set) emit a separate `Genuine grad collapse at step N (no NaN flags set; legitimate near-zero gradients)` warn log distinguishing the two cases. Diagnostic-only — keep the path even after the F1 ep2 root cause fix lands; it's the right gate for future regressions of the same NaN-clamped-to-zero class. Companion infrastructure (`fused_training.rs::read_nan_flags` pub(crate) accessor and `reset_nan_flags` at fold boundary in `FusedTrainingCtx::reset_for_fold`) was already in place from the earlier `halt_nan`-only readout — only the new `halt_grad_collapse` consumer site landed in this commit. Per `feedback_kill_runs_on_anomaly_quickly.md`: the diagnostic enables fast root-cause identification of the F1 ep2 explosion source (which kernel produced the first NaN — C51 KL projection? IQN aux? CQL? aux heads?). Per `feedback_no_partial_refactor.md`: error propagation preserved (early-stop still fires after diagnostic emits); the diagnostic is strictly additive to the existing collapse-halt contract. Plan C T11 diagnostic (2026-04-29): per-step `FOLD_EXPLOSION_DIAG` log emits when `fold >= 1` and `grad_norm > 1000` or jumps `5x` from the prior guard step. Captures loss decomposition (`c51`, `mse`, `iqn` from already-pinned device-mapped scalars + `total` from `gr.raw_loss`), Q range (`q_min`/`q_mean`/`q_max` via `reduce_current_q_stats` cold-path reduce) and atom positions (`per_sample_support[sample=0, dir=0]` and `[sample=0, dir=2]` triples `(v_min, v_max, delta_z)` via 12-element cold-path DtoH). After the trigger fires, continues emitting unconditionally for `STEPS_AFTER_EXPLOSION = 5` guard steps so the explosion trajectory is captured. Plus an unconditional `FOLD_EXPLOSION_DIAG[F1_END_EP1]` baseline log emitted once at the end of fold 1 epoch 0 — the healthy state immediately preceding the F1 ep2 explosion. Removed after root cause identified — not for production. Companion fields `current_fold`, `last_logged_grad_norm`, `explosion_diag_steps_remaining` on `DQNTrainer` (init `(0, None, 0)` in constructor; `last_logged_grad_norm`/`explosion_diag_steps_remaining` reset in `reset_for_fold`; `current_fold` overwritten unconditionally at fold-loop entry). Companion accessors `FusedTrainingCtx::explosion_diag_{loss_components,atom_range,q_range}` and `GpuDqnTrainer::{c51_loss_pinned_value, mse_loss_pinned_value, stream_for_diag}`. CQL and ensemble losses are not pinned-readback (they live in transient device buffers consumed inside the training graph) and are explicitly absent from the decomposition: if none of `{c51, iqn, mse}` is the runaway driver but `total_loss` still explodes, that implicates the unexposed CQL/ens path — the absence is itself diagnostic information. + +Plan C T11 follow-up R1+P (2026-04-29): + +R1 — eliminated K's Adam shrink-and-perturb at fold boundary (m_buf and v_buf back to memset_zeros, K's hardcoded ADAM_M_SHRINK=0.1 / ADAM_V_SHRINK=0.01 constants removed). K was a stepping-stone fix introduced before fold_warmup_factor existed (commit 4ef1d8ebb). fold_warmup_factor — ISV-driven, drives lr+clip dampening — already provides principled first-step protection via the cold-start `lr_eff = lr_base × max(MIN_WARMUP_LR_FRAC, fold_warmup_factor)` formula. K's hardcoded shrink ratios violated `feedback_adaptive_not_tuned` AND created a downstream pathology: tiny v_hat denominator → oversized Adam updates ~50+ steps post-reset → trunk param overshoot → save_h_s2 NaN at F1 ~step 1745 (smoke-test-bkdx5 diagnostic). With K removed, Adam restarts cleanly at m=0, v=0 and warmup_factor handles the cold-start window via lr/clip damping — single mechanism, ISV-driven, no hardcoded constants. R1.B verification (no edits): `state_reset_registry.rs` already lists both `isv_fold_warmup_factor` and `isv_grad_norm_fast_ema` as `FoldReset` entries with matching dispatch arms in `training_loop.rs::reset_named_state` — the K+warmup commit landed both halves of the warmup-factor input contract atomically per `feedback_no_partial_refactor.md`. + +P — expanded `nan_flags_buf` 16→24 with 5 new GRN h_s2 sub-stage NaN checks (`grn_h_s2_elu_post`=16, `grn_h_s2_linear_b_out`=17, `grn_h_s2_glu_sigmoid`=18, `grn_h_s2_ln_rstd`=19, `grn_h_s2_ln_normed`=20) for finer-grained source identification if R1 alone doesn't fix F1. Slot 19 (`ln_rstd`) is the divide-by-zero diagnostic (variance→0 ⇒ rstd→∞ ⇒ NaN downstream). New accessors: `GrnBlock::{elu_post_ptr, linear_b_out_ptr, sigmoid_b_ptr, ln_rstd_ptr, ln_normed_ptr}` (additive, reads-only) and `CublasGemmSet::{grn_h_s2_online, grn_h_s2_linear_b_ptr}`. h_s2 only — pragmatic instrumentation scope: that's the GRN block where save_h_s2 went NaN at F1 step 1745, and instrumenting only h_s2 covers the failing stage without doubling per-step kernel-launch cost. h_s1 stages can be added by name table slots 21-23 if the F1 explosion turns out to originate above h_s2. The K-removal + GRN-stage checks combined: if the F1 explosion was due to K's tiny-v_hat pathology, R1 alone fixes it. If a different mechanism, the new flags pinpoint which GRN sub-stage produces NaN first. Per `feedback_no_partial_refactor.md`: R1 and P land together because both touch the fold-boundary contract (K-removal changes the post-reset Adam state; GRN-stage NaN checks observe the consequences of that state on the trunk encoder's first forward). `read_nan_flags()` signature update [16]→[24] propagates through `GpuDqnTrainer`, `FusedTrainingCtx`, and both name-table sites in `training_loop.rs` (halt_nan + halt_grad_collapse).