fix(dqn): R1 — eliminate K Adam shrink + P — GRN-stage NaN checks
R1: K's hardcoded shrink-and-perturb (m×0.1, v×0.01) at fold boundary
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).
K was introduced (commit 4ef1d8ebb) BEFORE fold_warmup_factor existed
in the same commit's "K + adaptive warmup" pair. With warmup_factor
in place — ISV-driven, dampens lr+clip via lr_eff = lr_base ×
max(MIN_WARMUP_LR_FRAC, fold_warmup_factor) — K is redundant. Single
mechanism, ISV-driven, no hardcoded constants. Eliminating K leaves
m,v reset to 0 at fold boundary; warmup_factor handles cold-start.
P: expanded nan_flags_buf 16→24 with 5 GRN-stage checks
(linear_a_out, elu_out, linear_b_out, glu_sigmoid_out,
layernorm_var/out) for finer-grained source identification if R1
alone doesn't fix F1.
Predicted outcomes:
- If K's tiny-v_hat was the cause: F1 trains successfully (R1 alone)
- If different mechanism: new GRN-stage flags pinpoint which sub-
stage produces NaN, enabling layer-level fix
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
@@ -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(
|
||||
|
||||
@@ -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::<i32>(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::<i32>(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)
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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()
|
||||
}
|
||||
|
||||
|
||||
@@ -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<String> = 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<String> = flags.iter().enumerate()
|
||||
.filter(|(_, &f)| f != 0)
|
||||
|
||||
@@ -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).
|
||||
|
||||
Reference in New Issue
Block a user