diag(dqn): read NaN flags on grad-collapse path too (find F1 ep2 explosion source)
Existing nan_flags_buf [16] covers 13 buffers with per-step NaN checks inside the captured training graph. But read_nan_flags() was only called when training guard's halt_nan fired — which checks pinned- readback grad_norm. NaN-clamped-to-zero gradients reach the pinned scalar as 0, not NaN, so halt_nan never fires for the explosion case. The collapse path (halt_grad_collapse, grad < 0.01) was firing instead without reading the flags. Add flag readback in that path when gr.raw_grad_norm < 1e-6 (suspicious zero) — logs NaN-CLAMPED-TO-ZERO with flagged buffer names. Genuine near-zero gradients get a separate "no NaN flags set" log so we can distinguish the two cases. This should pinpoint which kernel produces the F1 ep2 NaN first (C51 KL projection? IQN aux? CQL? aux heads?). Diagnostic only — keep after fix lands; correct gate for future regressions.
This commit is contained in:
jgrusewski
@@ -2015,6 +2015,64 @@ 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.
|
||||
// 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
|
||||
// flags set" log distinguishing the two cases.
|
||||
if gr.raw_grad_norm < 1e-6 {
|
||||
if let Some(ref mut fused) = self.fused_ctx {
|
||||
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.
|
||||
let names = [
|
||||
"states_buf", // 0
|
||||
"on_v_logits", // 1
|
||||
"on_b_logits", // 2
|
||||
"mse_loss_scalar", // 3
|
||||
"params_buf_pre_fwd", // 4
|
||||
"params_ptr_pre_fwd", // 5
|
||||
"grad_buf", // 6
|
||||
"save_current_lp", // 7
|
||||
"save_projected", // 8
|
||||
"moe_gate_softmax", // 9
|
||||
"aux_nb_loss_scalar", // 10
|
||||
"aux_rg_loss_scalar", // 11
|
||||
"save_h_s2", // 12
|
||||
"rsv13", "rsv14", "rsv15",
|
||||
];
|
||||
let flagged: Vec<String> = flags.iter().enumerate()
|
||||
.filter(|(_, &f)| f != 0)
|
||||
.map(|(i, _)| format!("{}={}", i, names[i]))
|
||||
.collect();
|
||||
if !flagged.is_empty() {
|
||||
tracing::error!(
|
||||
"NaN-CLAMPED-TO-ZERO at step {} (grad collapse path): \
|
||||
flagged=[{}] — the collapse is NaN-derived, not \
|
||||
legitimate; this is the F1 ep2 explosion source",
|
||||
train_step_count, flagged.join(", ")
|
||||
);
|
||||
} else {
|
||||
tracing::warn!(
|
||||
"Genuine grad collapse at step {} (no NaN flags set; \
|
||||
legitimate near-zero gradients)",
|
||||
train_step_count
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
agent.check_gradient_collapse(gr.raw_grad_norm).map_err(|e| {
|
||||
tracing::info!("Early stopping (gradient collapse): {}", e);
|
||||
anyhow::anyhow!("Early stopping: {}", e)
|
||||
|
||||
@@ -2227,4 +2227,6 @@ Per `pearl_adaptive_moe_lambda.md` — canonical "EMA-tracked diagnostic drives
|
||||
|
||||
Plan C T11 follow-up N (2026-04-29): Winsorize the adaptive_clip EMA input — clamp single-sample raw_grad_norm to `GRAD_CLIP_OUTLIER_K × previous_adaptive_clip` before EMA update. K = 100 is a numerical-stability bound. Discovered when smoke-test-xw4c6 showed F1 ep2 grad_norm=8.4B polluting adaptive_clip's EMA → subsequent extreme grads pass unclipped → NaN propagation. After N, single outlier → EMA absorbs at most 100× growth → next adaptive_clip caps at ~1000 (vs 30 steady-state) → subsequent F1 outliers get clipped → bounded grads pass to Adam → no NaN. Companion observability log emits GRAD_CLIP_OUTLIER warning per clamp event. Fast/slow grad_norm EMAs (warmup factor inputs) intentionally NOT winsorized — they're a stability signal that should respond to outliers (further dampens warmup_factor → extra defense). Composes with K+warmup (4ef1d8ebb), A.2 adaptive tau, A.1/A.3 fold-boundary state resets, F+H Q-drift kill robustness.
|
||||
|
||||
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.
|
||||
|
||||
Reference in New Issue
Block a user