diff --git a/crates/ml/src/trainers/dqn/trainer/training_loop.rs b/crates/ml/src/trainers/dqn/trainer/training_loop.rs index d00c2ab0e..648b1c083 100644 --- a/crates/ml/src/trainers/dqn/trainer/training_loop.rs +++ b/crates/ml/src/trainers/dqn/trainer/training_loop.rs @@ -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 = 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) diff --git a/docs/dqn-wire-up-audit.md b/docs/dqn-wire-up-audit.md index 0d95acfa8..f78172fdf 100644 --- a/docs/dqn-wire-up-audit.md +++ b/docs/dqn-wire-up-audit.md @@ -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.