fix(dqn): SP1 Phase C — preserve NaN diagnostics + correct F0 paper-review
Quality-review follow-ups to commit 97f1d25f5:
1. CRITICAL: preserve regression sentinel for slots 24, 25, 26, 32.
The post-clamps in apply_iqn_trunk_gradient + launch_cublas_backward_to
ran BEFORE run_nan_checks_post_backward, silently zeroing the buffers
the diagnostic was supposed to observe. Fix: inline check_nan_f32 calls
IMMEDIATELY BEFORE each launch_clamp_finite_f32 — same pattern as the
existing slot 35 check at line 6949. Now the diagnostic fires on NaN
entry; the clamp then sanitizes (preventing propagation but preserving
the flag — flags are sticky/OR'd in dqn_nan_check_f32, never cleared).
2. CRITICAL: correct F0 paper-review line reference + reasoning. The
commit body cited "5.0 norm-clip at line 16747" — that's a memset, not
a clip. Real norm-clip lives at lines 16827-16868. Also tightened the
L2-vs-per-element reasoning: L2 norm <= 5.0 worst-case bounds per-element
|x| <= 5.0 (single-element edge), still several orders of magnitude
below the 1e6 max_abs guard.
3. Inline comment at line 6951 mislabeled the clamp target as 'slot 27
input' — actually clamps the slot 35 buffer (bw_d_h_s2 post-DtoD). Fixed.
4. Articulated pre-clamp rationale: defense-in-depth regression protection
for input buffers (slots 24, 25, 35) against future pathologies that
could create extreme-but-finite inputs. The F1 ep2 NaN was outputs
overflowing finite inputs, but the same fix template covers both
failure modes per feedback_no_partial_refactor.
5. Renamed clamp_finite_f32_kernel -> dqn_clamp_finite_f32_kernel for
consistency with sibling utility kernels (dqn_nan_check_f32,
dqn_zero_kernel, dqn_grad_norm_kernel).
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
@@ -426,7 +426,7 @@ extern "C" __global__ void dqn_zero_kernel(
|
||||
* device memory and no DtoD/HtoH copies (per
|
||||
* `feedback_no_htod_htoh_only_mapped_pinned`).
|
||||
* ══════════════════════════════════════════════════════════════════════ */
|
||||
extern "C" __global__ void clamp_finite_f32_kernel(
|
||||
extern "C" __global__ void dqn_clamp_finite_f32_kernel(
|
||||
float* __restrict__ buf,
|
||||
int n,
|
||||
float max_abs
|
||||
|
||||
@@ -2656,7 +2656,7 @@ pub struct GpuDqnTrainer {
|
||||
/// Second NaN check kernel handle (formerly f32, now f32 — both check f32 buffers).
|
||||
pub(crate) nan_check_f32_kernel_b: CudaFunction,
|
||||
/// SP1 Phase C surgical fix: in-place finite clamp kernel
|
||||
/// (`clamp_finite_f32_kernel` in `dqn_utility_kernels.cu`). Replaces
|
||||
/// (`dqn_clamp_finite_f32_kernel` in `dqn_utility_kernels.cu`). Replaces
|
||||
/// NaN/Inf with 0 and clamps finite |v| to ±max_abs. max_abs is sourced
|
||||
/// from existing ISV slots (1e6 × ISV[H_S2_RMS_EMA_INDEX=96] for the
|
||||
/// IQN trunk path; 1e6 × ISV[Q_DIR_ABS_REF_INDEX=21] for the
|
||||
@@ -6948,9 +6948,10 @@ impl GpuDqnTrainer {
|
||||
// `docs/dqn-backward-nan-audit.md` (slot 35 row).
|
||||
self.check_nan_f32(self.ptrs.bw_d_h_s2, b * sh2, 35)?;
|
||||
|
||||
// SP1 Phase C surgical fix (slot 27 input sanitise): clamp NaN/Inf
|
||||
// to 0 + bound magnitude to 1e6 × ISV[H_S2_RMS_EMA_INDEX=96] on the
|
||||
// GEMM input gradient. Pre-call defence — bounds the cuBLAS GEMM
|
||||
// SP1 Phase C (slot 35 buffer / GEMM input sanitise): clamp NaN/Inf
|
||||
// to 0 + bound magnitude to 1e6 × ISV[H_S2_RMS_EMA_INDEX=96]. The
|
||||
// buffer here is `bw_d_h_s2` post-DtoD copy from `iqn_d_h_s2_ptr`
|
||||
// (slot 27 source). Pre-call defence — bounds the cuBLAS GEMM
|
||||
// accumulator's intermediate products so the reductions inside
|
||||
// `encoder_backward_chain` (Linear_a/Linear_b dW/dX/dB GEMMs) cannot
|
||||
// produce f32-overflow NaN from the kind of pathological inputs
|
||||
@@ -6959,7 +6960,9 @@ impl GpuDqnTrainer {
|
||||
// `feedback_isv_for_adaptive_bounds`). F0 paper-review: F0
|
||||
// ISV[96] ≈ 1.0 (LayerNorm output RMS), max_abs = 1e6 — F0-typical
|
||||
// |iqn_d_h_s2| ≤ ~10² so guard is a no-op on F0; F1 overflows
|
||||
// (≥ 1e6) trigger clamp.
|
||||
// (≥ 1e6) trigger clamp. The slot 35 NaN check at line 6949 fires
|
||||
// BEFORE this clamp, so any NaN entry is captured in the regression
|
||||
// sentinel before sanitisation — see Issue 1 fix in commit fix-up.
|
||||
let h_s2_rms_ema = self.read_isv_signal_at(H_S2_RMS_EMA_INDEX);
|
||||
let max_abs_input = (1e6_f32 * h_s2_rms_ema).max(1e3_f32);
|
||||
self.launch_clamp_finite_f32(self.ptrs.bw_d_h_s2, b * sh2, max_abs_input)?;
|
||||
@@ -7029,6 +7032,13 @@ impl GpuDqnTrainer {
|
||||
// catches that residual class. Same `max_abs_input` bound as the
|
||||
// pre-call clamp because both inputs and outputs of the trunk GEMM
|
||||
// share the H_S2_RMS_EMA-driven magnitude scale.
|
||||
//
|
||||
// Inline slot 26 NaN check IMMEDIATELY BEFORE the clamp preserves
|
||||
// the regression sentinel: `run_nan_checks_post_backward`
|
||||
// (fused_training.rs:1540) runs AFTER this clamp, so without this
|
||||
// pre-clamp check the slot 26 flag would observe sanitised zeros
|
||||
// and never fire. Mirrors the slot 35 check at line 6949 pattern.
|
||||
self.check_nan_f32(self.ptrs.iqn_trunk_m, trunk_grad_total, 26)?;
|
||||
self.launch_clamp_finite_f32(
|
||||
self.ptrs.iqn_trunk_m, trunk_grad_total, max_abs_input,
|
||||
)?;
|
||||
@@ -18596,20 +18606,34 @@ impl GpuDqnTrainer {
|
||||
// ISV[Q_DIR_ABS_REF_INDEX=21] (the direction-branch |Q| EMA — same
|
||||
// reference scale used elsewhere for q-magnitude bounds), with a
|
||||
// 1e3 ε floor for uninitialised state (Invariant 1 carve-out per
|
||||
// `feedback_isv_for_adaptive_bounds`). F0 paper-review: F0
|
||||
// ISV[21] ≈ 1.0 (Q-value scale at convergence), max_abs = 1e6 —
|
||||
// F0-typical |d_value_logits| ≤ ~10³ post the existing 5.0
|
||||
// norm-clip (line 16747) so guard is a no-op on F0; F1 overflows
|
||||
// (≥ 1e6) trigger clamp. NB: the existing line-16747 norm-clamp
|
||||
// is a per-buffer-norm operation (single L2 norm of the entire
|
||||
// buffer); this kernel is a per-element magnitude clamp — they
|
||||
// catch different failure modes (norm = aggregate; clamp_finite
|
||||
// = per-element NaN/Inf/extreme), so both stay.
|
||||
// `feedback_isv_for_adaptive_bounds`). F0 paper-review:
|
||||
// F0 |d_value_logits|, |d_adv_logits| are bounded post the existing
|
||||
// 5.0 L2 norm-clip (gpu_dqn_trainer.rs:16827-16868, max_d_logit_norm
|
||||
// = 5.0_f32 + clip_grad_kernel). L2 norm ≤ 5.0 across N elements
|
||||
// bounds per-element |x| ≤ 5.0 absolute worst-case (single-element
|
||||
// edge case) and ≤ 5.0/sqrt(N) typical-case. Both are several orders
|
||||
// of magnitude below the 1e6 max_abs guard threshold, so guards are
|
||||
// no-ops on F0 dynamics. F1 overflows enter via post-norm-clip GEMM
|
||||
// accumulator overflow; the fix targets the GEMM input/output sites
|
||||
// not the clip itself. NB: the existing 5.0 L2 norm-clip is a
|
||||
// per-buffer-norm operation (aggregate L2 over the entire buffer);
|
||||
// this kernel is a per-element magnitude clamp — they catch
|
||||
// different failure modes (norm = aggregate; clamp_finite =
|
||||
// per-element NaN/Inf/extreme), so both stay.
|
||||
//
|
||||
// Inline slot 24/25 NaN checks IMMEDIATELY BEFORE each clamp
|
||||
// preserve the regression sentinel: `run_nan_checks_post_backward`
|
||||
// (fused_training.rs:1540) runs AFTER these clamps, so without
|
||||
// these pre-clamp checks the slot 24/25 flags would observe
|
||||
// sanitised zeros and never fire. Mirrors the slot 35 check at
|
||||
// line 6949 pattern.
|
||||
let q_dir_abs_ref = self.read_isv_signal_at(Q_DIR_ABS_REF_INDEX);
|
||||
let max_abs_q = (1e6_f32 * q_dir_abs_ref).max(1e3_f32);
|
||||
let n_val_clamp = self.d_value_logits_buf.len();
|
||||
let n_adv_clamp = self.d_adv_logits_buf.len();
|
||||
self.check_nan_f32(d_value_logits_ptr, n_val_clamp, 24)?;
|
||||
self.launch_clamp_finite_f32(d_value_logits_ptr, n_val_clamp, max_abs_q)?;
|
||||
self.check_nan_f32(d_adv_logits_ptr, n_adv_clamp, 25)?;
|
||||
self.launch_clamp_finite_f32(d_adv_logits_ptr, n_adv_clamp, max_abs_q)?;
|
||||
|
||||
bw.backward_full(
|
||||
@@ -18753,9 +18777,16 @@ impl GpuDqnTrainer {
|
||||
// clamp because the same Q-magnitude reference scale bounds both
|
||||
// ends of the dueling backward chain. Skipped when bottleneck is
|
||||
// disabled (`bn_d_concat_buf` is unused / `s1_dx_output == 0`).
|
||||
//
|
||||
// Inline slot 32 NaN check IMMEDIATELY BEFORE the clamp preserves
|
||||
// the regression sentinel: `run_nan_checks_post_backward`
|
||||
// (fused_training.rs:1540) runs AFTER this clamp, so without this
|
||||
// pre-clamp check the slot 32 flag would observe sanitised zeros
|
||||
// and never fire. Mirrors the slot 35 check at line 6949 pattern.
|
||||
if self.config.bottleneck_dim > 0 {
|
||||
let bn_d_concat_len = self.bn_d_concat_buf.len();
|
||||
let bn_d_concat_ptr = self.bn_d_concat_buf.raw_ptr();
|
||||
self.check_nan_f32(bn_d_concat_ptr, bn_d_concat_len, 32)?;
|
||||
self.launch_clamp_finite_f32(bn_d_concat_ptr, bn_d_concat_len, max_abs_q)?;
|
||||
}
|
||||
|
||||
@@ -20596,8 +20627,8 @@ fn compile_training_kernels(
|
||||
// `submit_forward_ops_main` (graph-captured forward_child) and
|
||||
// `apply_iqn_trunk_gradient` (aux_child) — sharing the module is correct
|
||||
// because they belong to the same captured replay group.
|
||||
let clamp_finite_f32 = module.load_function("clamp_finite_f32_kernel")
|
||||
.map_err(|e| MLError::ModelError(format!("clamp_finite_f32_kernel load: {e}")))?;
|
||||
let clamp_finite_f32 = module.load_function("dqn_clamp_finite_f32_kernel")
|
||||
.map_err(|e| MLError::ModelError(format!("dqn_clamp_finite_f32_kernel load: {e}")))?;
|
||||
|
||||
let popart_normalize = module.load_function("popart_normalize_kernel")
|
||||
.map_err(|e| MLError::ModelError(format!("popart_normalize_kernel load: {e}")))?;
|
||||
|
||||
@@ -2253,8 +2253,10 @@ SP1 Phase C surgical fix (2026-04-29): patched two cuBLAS GEMM backward operatio
|
||||
|
||||
2. `launch_cublas_backward_to` main backward path (gpu_dqn_trainer.rs:18516+) — pre-call clamp on `d_value_logits_buf` + `d_adv_logits_buf` inputs (slot 24/25 buffers — feed the dueling/branch backward GEMMs in `bw.backward_full(...)` which produces `d_h_s2` and ultimately `bn_d_concat_buf`); post-call sanitize on `bn_d_concat_buf` output (slot 32 buffer, gated on `bottleneck_dim > 0`). ISV-driven max_abs bound = `1e6 * isv[Q_DIR_ABS_REF_INDEX=21]` with `1e3` ε floor.
|
||||
|
||||
New utility: `clamp_finite_f32_kernel` (CUDA — replaces NaN/Inf with 0 + clamps finite values to ±max_abs); `launch_clamp_finite_f32` (Rust wrapper). Pattern matches the existing `isfinite`-guard precedent in `iqn_backward_per_sample` (line 612) + `iqn_adam_kernel` (line 873). Kernel lives in `dqn_utility_kernels.cu` (already in `build.rs`); loaded into the same `module` as the NaN check kernels in `compile_training_kernels`. Both call sites land inside captured children (forward_child for slot 32 inputs/output via `launch_cublas_backward_to` → `submit_forward_ops_main`; aux_child for slot 26 input/output via `apply_iqn_trunk_gradient` → `submit_aux_ops`); the kernel uses only stream-bound `launch_builder` like `check_nan_f32`, so it is graph-safe.
|
||||
New utility: `dqn_clamp_finite_f32_kernel` (CUDA — replaces NaN/Inf with 0 + clamps finite values to ±max_abs); `launch_clamp_finite_f32` (Rust wrapper). Kernel name follows the `dqn_*` prefix convention used by sibling utility kernels (`dqn_nan_check_f32`, `dqn_zero_kernel`, `dqn_grad_norm_kernel`). Pattern matches the existing `isfinite`-guard precedent in `iqn_backward_per_sample` (line 612) + `iqn_adam_kernel` (line 873). Kernel lives in `dqn_utility_kernels.cu` (already in `build.rs`); loaded into the same `module` as the NaN check kernels in `compile_training_kernels`. Both call sites land inside captured children (forward_child for slot 32 inputs/output via `launch_cublas_backward_to` → `submit_forward_ops_main`; aux_child for slot 26 input/output via `apply_iqn_trunk_gradient` → `submit_aux_ops`); the kernel uses only stream-bound `launch_builder` like `check_nan_f32`, so it is graph-safe.
|
||||
|
||||
F0 risk: low — guards activate at 1e6× ISV magnitude (several orders of magnitude wider than F0-typical inputs); F0 paper-review confirmed guards are no-ops on F0 dynamics (F0 ISV[96] ≈ 1.0 LayerNorm RMS; F0 ISV[21] ≈ 1.0 Q-value scale; F0 |iqn_d_h_s2| ≤ ~10²; F0 |d_value_logits| ≤ ~10³ post the existing 5.0 norm-clip at line 16747).
|
||||
F0 risk: low — guards activate at 1e6× ISV magnitude (several orders of magnitude wider than F0-typical inputs); F0 paper-review confirmed guards are no-ops on F0 dynamics (F0 ISV[96] ≈ 1.0 LayerNorm RMS; F0 ISV[21] ≈ 1.0 Q-value scale; F0 |iqn_d_h_s2| ≤ ~10²). For the dueling-branch path, F0 |d_value_logits| / |d_adv_logits| are bounded post the existing 5.0 L2 norm-clip (gpu_dqn_trainer.rs:16827-16868, `max_d_logit_norm = 5.0_f32` + `clip_grad_kernel`). L2 norm ≤ 5.0 across N elements bounds per-element |x| ≤ 5.0 absolute worst-case (single-element edge case) and ≤ 5.0/sqrt(N) typical-case — both several orders of magnitude below the 1e6 max_abs guard threshold, so guards are no-ops on F0. F1 overflows enter via post-norm-clip GEMM accumulator overflow; the fix targets the GEMM input/output sites not the clip itself.
|
||||
|
||||
Permanent diagnostic infrastructure (Phase B slots 24-35) stays as regression sentinel — will detect any future regression of the cuBLAS-overflow NaN class. Combined RELATED commit per `feedback_no_partial_refactor` — both kernels patched in one commit since they share the same unsafe-pattern (cuBLAS GEMM extreme-intermediate-product overflow) + the same fix template.
|
||||
|
||||
SP1 Phase C quality fix-up (2026-04-29): commit-quality follow-up to `97f1d25f5`. (1) The post-clamps in `apply_iqn_trunk_gradient` (slot 26 site) and `launch_cublas_backward_to` (slots 24/25/32 sites) ran BEFORE `run_nan_checks_post_backward` (fused_training.rs:1540), silently zeroing the buffers the diagnostic was supposed to observe — the regression sentinel for these 4 slots was effectively dead. Inline `check_nan_f32` calls now fire IMMEDIATELY BEFORE each `launch_clamp_finite_f32`, mirroring the existing slot 35 check at line 6949: the diagnostic captures NaN entry, the clamp then sanitises so propagation stops, and the flag remains for the regression sentinel readback path. (2) Corrected the F0 paper-review reference: the original commit cited `5.0 norm-clip at line 16747` which is actually a `cuMemsetD32Async` zero-init; the real L2 norm-clip lives at lines 16827-16868 (`max_d_logit_norm = 5.0_f32` + `clip_grad_kernel`). Reasoning tightened: L2 norm ≤ 5.0 worst-case bounds per-element |x| ≤ 5.0 (single-element edge), still several orders of magnitude below the 1e6 max_abs guard. (3) Inline comment at `gpu_dqn_trainer.rs:6951` mislabeled the clamp target as `slot 27 input` — the kernel actually clamps `bw_d_h_s2` post-DtoD, which is the slot 35 buffer; slot 27 is the source-side `iqn_d_h_s2_ptr`. Comment fixed. (4) Pre-clamp rationale articulated: the pre-clamps (slots 24/25/35) target buffers Phase B smoke confirmed CLEAN at F1 first-fire, so they are NOT the smoking gun. They are defense-in-depth regression protection — sanitise input buffers to forestall any FUTURE pathology that creates extreme-but-finite cuBLAS inputs (different from the F1 ep2 NaN where outputs overflow on finite-but-large inputs). They cost 4 additional graph nodes per step (~µs) and are no-ops on F0/F1 typical inputs (≪ 1e6 max_abs guard). The pattern covers both failure modes in one fix per `feedback_no_partial_refactor`. (5) Renamed `clamp_finite_f32_kernel` → `dqn_clamp_finite_f32_kernel` for consistency with sibling utility kernels (`dqn_nan_check_f32`, `dqn_zero_kernel`, `dqn_grad_norm_kernel`); the Rust wrapper retains its `launch_clamp_finite_f32` name (matches `launch_check_nan_f32` precedent — wrapper names don't carry the `dqn_*` prefix).
|
||||
|
||||
Reference in New Issue
Block a user