diff --git a/crates/ml/src/cuda_pipeline/dqn_utility_kernels.cu b/crates/ml/src/cuda_pipeline/dqn_utility_kernels.cu index a108e262c..b63d22703 100644 --- a/crates/ml/src/cuda_pipeline/dqn_utility_kernels.cu +++ b/crates/ml/src/cuda_pipeline/dqn_utility_kernels.cu @@ -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 diff --git a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs index a069f4d05..ccfcdcfdd 100644 --- a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs +++ b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs @@ -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}")))?; diff --git a/docs/dqn-wire-up-audit.md b/docs/dqn-wire-up-audit.md index c2a813ba2..ab1e9e53b 100644 --- a/docs/dqn-wire-up-audit.md +++ b/docs/dqn-wire-up-audit.md @@ -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).