fix(dqn): SP1 Phase C — F1 NaN cuBLAS-overflow surgical fix (slots 26 + 32)
Patches two cuBLAS GEMM backward operations identified by SP1 Phase B
smoke smoke-test-xvzgk (commit f139a63ee) as the F1 ep2 NaN sources:
1. apply_iqn_trunk_gradient (gpu_dqn_trainer.rs:6885+):
- Pre-call: clamp bw_d_h_s2 (the IQN DtoD-overwrite target that feeds
the GRN trunk encoder's cuBLAS Linear_a/Linear_b dW/dX/dB GEMMs;
symmetric with slot 27 source-side check) to ±1e6×ISV[96].
- Post-call: sanitize iqn_trunk_m (slot 26 output) — isfinite-or-zero
+ magnitude clamp; defence-in-depth.
2. launch_cublas_backward_to main backward path:
- Pre-call: clamp d_value_logits_buf + d_adv_logits_buf (slots 24/25
inputs that feed bw.backward_full's dueling/branch backward GEMMs)
to ±1e6×ISV[21].
- Post-call: sanitize bn_d_concat_buf (slot 32 output, gated on
bottleneck_dim > 0).
Both fixes use the new clamp_finite_f32_kernel utility (CUDA — replaces
NaN/Inf with 0 + clamps finite values to ±max_abs) + launch_clamp_finite_f32
(Rust wrapper). 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, aux_child for slot 26) and use only stream-
bound launch_builder, so the kernel is graph-safe.
ISV-driven max_abs bound = 1e6 × isv[<slot>] with 1e3 ε floor for
uninitialized state (Invariant 1 carve-out per
feedback_isv_for_adaptive_bounds). Wide guard band — F0 inputs sit
several orders of magnitude below the threshold (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), so guards are no-ops on F0; F1 overflows trigger clamp.
F0 paper-review pre-smoke confirmed.
Combined RELATED commit per feedback_no_partial_refactor — both kernels
share the same unsafe-pattern (cuBLAS GEMM extreme-intermediate-product
overflow) + the same fix template, so they ship together.
Phase B slots 24-35 remain as permanent diagnostic infrastructure;
will catch any future regression of this NaN class.
SP1 Phase D validation smoke pending.
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
@@ -404,6 +404,40 @@ extern "C" __global__ void dqn_zero_kernel(
|
||||
if (i < n) buf[i] = 0.0f;
|
||||
}
|
||||
|
||||
/* ══════════════════════════════════════════════════════════════════════
|
||||
* IN-PLACE FINITE CLAMP KERNEL — SP1 Phase C surgical fix infrastructure.
|
||||
*
|
||||
* For each element of `buf`:
|
||||
* - If !isfinite(v) (NaN or ±Inf): set to 0.0f.
|
||||
* - Else clamp finite |v| to ±max_abs (ISV-driven bound supplied by host).
|
||||
*
|
||||
* Used as pre-call sanitisation on cuBLAS GEMM backward inputs (slots 27 +
|
||||
* 24/25 per `docs/dqn-backward-nan-audit.md`) and as post-call defence-in-
|
||||
* depth on outputs (slots 26 + 32) where extreme intermediate-product
|
||||
* accumulator overflow has been observed (smoke-test-xvzgk @ commit
|
||||
* f139a63ee). 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 Bottleneck Linear path
|
||||
* with a 1e3 floor applied host-side when the ISV slot is uninitialised
|
||||
* (Invariant 1 ε carve-out — prevents zero-clamping a hot buffer when the
|
||||
* ISV producer hasn't run yet).
|
||||
*
|
||||
* Launch config: grid=(ceil(n/256)), block=(256). In-place — no extra
|
||||
* device memory and no DtoD/HtoH copies (per
|
||||
* `feedback_no_htod_htoh_only_mapped_pinned`).
|
||||
* ══════════════════════════════════════════════════════════════════════ */
|
||||
extern "C" __global__ void clamp_finite_f32_kernel(
|
||||
float* __restrict__ buf,
|
||||
int n,
|
||||
float max_abs
|
||||
) {
|
||||
int idx = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
if (idx >= n) return;
|
||||
float v = buf[idx];
|
||||
/* NaN/Inf → 0; finite values clamped to ±max_abs */
|
||||
buf[idx] = isfinite(v) ? fmaxf(fminf(v, max_abs), -max_abs) : 0.0f;
|
||||
}
|
||||
|
||||
/* ══════════════════════════════════════════════════════════════════════
|
||||
* REGIME-ADAPTIVE PER SCALING KERNEL
|
||||
*
|
||||
|
||||
@@ -2655,6 +2655,16 @@ pub struct GpuDqnTrainer {
|
||||
pub(crate) nan_check_f32_kernel: CudaFunction,
|
||||
/// 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
|
||||
/// 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
|
||||
/// Bottleneck Linear path) with a 1e3 floor for uninitialised state.
|
||||
/// Used as pre-call sanitisation on cuBLAS GEMM backward inputs and
|
||||
/// post-call defence-in-depth on outputs at slots 26/32 per
|
||||
/// `docs/dqn-backward-nan-audit.md`.
|
||||
pub(crate) clamp_finite_f32_kernel: CudaFunction,
|
||||
/// SP1 Phase B: NaN flag buffer expanded from 24 → 48 slots. Slots 0-23
|
||||
/// are forward-path checks (existing); slots 24-35 are backward-path
|
||||
/// checks (new — wired by Task 4); slots 36-47 are reserved headroom
|
||||
@@ -6938,6 +6948,22 @@ 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
|
||||
// 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
|
||||
// observed at smoke-test-xvzgk (commit f139a63ee). 1e3 ε floor
|
||||
// covers the uninitialised-ISV case (Invariant 1 carve-out per
|
||||
// `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.
|
||||
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)?;
|
||||
|
||||
// ── 3. GRN trunk backward (encoder_backward_chain) ──
|
||||
let w_ptrs = f32_weight_ptrs_from_base(self.ptrs.params_ptr, ¶m_sizes);
|
||||
let states_ptr = if self.config.bottleneck_dim > 0 {
|
||||
@@ -6992,6 +7018,21 @@ impl GpuDqnTrainer {
|
||||
h_s1_save_ptr,
|
||||
)?;
|
||||
|
||||
// SP1 Phase C surgical fix (slot 26 output sanitise): post-call
|
||||
// defence-in-depth — `encoder_backward_chain`'s cuBLAS Linear_a/
|
||||
// Linear_b/Linear_residual dW/dX/dB GEMMs accumulate onto
|
||||
// `iqn_trunk_m` (= `trunk_scratch_base`) with beta=1; an extreme
|
||||
// intermediate-product overflow can produce NaN/Inf even when the
|
||||
// input gradient was finite (residual 8% from
|
||||
// `session_2026-04-05_nan_investigation.md` cross-checked against
|
||||
// smoke-test-xvzgk slot 26). isfinite-or-zero + magnitude clamp
|
||||
// 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.
|
||||
self.launch_clamp_finite_f32(
|
||||
self.ptrs.iqn_trunk_m, trunk_grad_total, max_abs_input,
|
||||
)?;
|
||||
|
||||
// ── 3b. Plan 4 Task 1B-iv-ext: VSN backward chain ─────────────────
|
||||
// Consumes the bn_d_concat just produced by encoder_backward_chain,
|
||||
// assembles `d_vsn_gated_state`, and writes 24 VSN dW/dB into
|
||||
@@ -9374,7 +9415,7 @@ impl GpuDqnTrainer {
|
||||
// per array. Stack is set once in DQNTrainer::new() (64KB for all kernels).
|
||||
|
||||
// ── Compile utility kernels (grad_norm, adam_update, etc.) ─
|
||||
let (grad_norm_kernel, grad_norm_finalize_kernel, grad_norm_finalize_adam, grad_norm_finalize_post_aux, adam_update_kernel, adam_update_post_aux, scale_f32_post_aux, saxpy_kernel, zero_kernel, regime_scale_kernel, shrink_perturb, _relu_mask_in_module, spectral_norm_kernel, clip_grad_kernel, pad_states_kernel, saxpy_f32_kernel, scale_f32_kernel, stochastic_depth_kernel, stochastic_depth_rng_kernel, bn_tanh_backward_kernel, bn_bias_grad_kernel, bn_tanh_concat_kernel_fn, vaccine_dot_kernel, vaccine_project_kernel, vaccine_dot_finalize, causal_intervene_kernel_fn, causal_reduce_kernel_fn, causal_mean_reduce_kernel_fn, pruning_mask_kernel, pruning_compute_kernel, her_inplace_kernel, nan_check_f32_kernel, nan_check_f32_kernel_b, popart_normalize_kernel, saxpy_f32_adam_grad, saxpy_f32_aux, scale_f32_aux, distill_saxpy_aux, grad_norm_standalone_post_aux, shrink_perturb_ungraphed, scale_f32_ungraphed, popart_robust_kernel) =
|
||||
let (grad_norm_kernel, grad_norm_finalize_kernel, grad_norm_finalize_adam, grad_norm_finalize_post_aux, adam_update_kernel, adam_update_post_aux, scale_f32_post_aux, saxpy_kernel, zero_kernel, regime_scale_kernel, shrink_perturb, _relu_mask_in_module, spectral_norm_kernel, clip_grad_kernel, pad_states_kernel, saxpy_f32_kernel, scale_f32_kernel, stochastic_depth_kernel, stochastic_depth_rng_kernel, bn_tanh_backward_kernel, bn_bias_grad_kernel, bn_tanh_concat_kernel_fn, vaccine_dot_kernel, vaccine_project_kernel, vaccine_dot_finalize, causal_intervene_kernel_fn, causal_reduce_kernel_fn, causal_mean_reduce_kernel_fn, pruning_mask_kernel, pruning_compute_kernel, her_inplace_kernel, nan_check_f32_kernel, nan_check_f32_kernel_b, clamp_finite_f32_kernel, popart_normalize_kernel, saxpy_f32_adam_grad, saxpy_f32_aux, scale_f32_aux, distill_saxpy_aux, grad_norm_standalone_post_aux, shrink_perturb_ungraphed, scale_f32_ungraphed, popart_robust_kernel) =
|
||||
compile_training_kernels(&stream, &config)?;
|
||||
|
||||
// Separate grad_norm instance for non-graph launches (outside CUDA graph).
|
||||
@@ -12618,6 +12659,7 @@ impl GpuDqnTrainer {
|
||||
her_inplace_kernel,
|
||||
nan_check_f32_kernel,
|
||||
nan_check_f32_kernel_b,
|
||||
clamp_finite_f32_kernel,
|
||||
nan_flags_buf,
|
||||
pruning_epoch: prune_ep,
|
||||
pruning_fraction: prune_frac,
|
||||
@@ -15147,6 +15189,44 @@ impl GpuDqnTrainer {
|
||||
.map_err(|e| MLError::ModelError(format!("nan_flags reset: {e}")))
|
||||
}
|
||||
|
||||
/// SP1 Phase C surgical fix: launch in-place finite clamp on `buf_ptr`.
|
||||
///
|
||||
/// Replaces NaN/Inf with 0 and clamps finite |v| to ±max_abs (host-side
|
||||
/// `f32`). 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 Bottleneck Linear path) with a
|
||||
/// `1e3` ε floor applied at the call site for uninitialised ISV state
|
||||
/// (Invariant 1 carve-out per `feedback_isv_for_adaptive_bounds`).
|
||||
///
|
||||
/// Pattern matches the existing `isfinite`-guard precedent in
|
||||
/// `iqn_backward_per_sample` (line 612) and `iqn_adam_kernel` (line 873).
|
||||
/// The kernel is in-place — no DtoD/HtoD/HtoH copies (per
|
||||
/// `feedback_no_htod_htoh_only_mapped_pinned`). Graph-safe — both call
|
||||
/// sites land inside captured children (forward_child for slot 32,
|
||||
/// aux_child for slot 26) and the kernel uses only stream-bound
|
||||
/// `launch_builder` like `check_nan_f32`.
|
||||
pub fn launch_clamp_finite_f32(
|
||||
&self,
|
||||
buf_ptr: u64,
|
||||
n: usize,
|
||||
max_abs: f32,
|
||||
) -> Result<(), MLError> {
|
||||
let n_i32 = n as i32;
|
||||
let blocks = ((n + 255) / 256) as u32;
|
||||
unsafe {
|
||||
self.stream
|
||||
.launch_builder(&self.clamp_finite_f32_kernel)
|
||||
.arg(&buf_ptr).arg(&n_i32).arg(&max_abs)
|
||||
.launch(LaunchConfig {
|
||||
grid_dim: (blocks.max(1), 1, 1),
|
||||
block_dim: (256, 1, 1),
|
||||
shared_mem_bytes: 0,
|
||||
})
|
||||
.map_err(|e| MLError::ModelError(format!("clamp_finite_f32 launch: {e}")))?;
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Read PopArt running variance from GPU (single scalar DtoH).
|
||||
pub fn read_popart_var(&self, out: &mut [f32; 1]) -> Result<(), MLError> {
|
||||
self.stream.memcpy_dtoh(&self.popart_var, out)
|
||||
@@ -18506,6 +18586,32 @@ impl GpuDqnTrainer {
|
||||
} else {
|
||||
0u64
|
||||
};
|
||||
|
||||
// SP1 Phase C surgical fix (slots 24/25 input sanitise): pre-call
|
||||
// clamp on the cuBLAS GEMM inputs that feed `backward_full`. NaN
|
||||
// or extreme finite values in `d_value_logits_buf` /
|
||||
// `d_adv_logits_buf` propagate through the dueling backward GEMMs
|
||||
// into `d_h_s2`/`bn_d_concat_buf` and produce the slot 32 NaN
|
||||
// observed at smoke-test-xvzgk (commit f139a63ee). Bound by 1e6 ×
|
||||
// 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.
|
||||
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.launch_clamp_finite_f32(d_value_logits_ptr, n_val_clamp, max_abs_q)?;
|
||||
self.launch_clamp_finite_f32(d_adv_logits_ptr, n_adv_clamp, max_abs_q)?;
|
||||
|
||||
bw.backward_full(
|
||||
&self.stream,
|
||||
d_value_logits_ptr,
|
||||
@@ -18637,6 +18743,22 @@ impl GpuDqnTrainer {
|
||||
h_s1_ptr,
|
||||
)?;
|
||||
|
||||
// SP1 Phase C surgical fix (slot 32 output sanitise): post-call
|
||||
// defence-in-depth on `bn_d_concat_buf` (= `s1_dx_output` when
|
||||
// bottleneck active). The trunk-encoder backward writes here via
|
||||
// Linear_a (beta=0) + Linear_residual (beta=1) cuBLAS GEMMs whose
|
||||
// intermediate-product accumulator can overflow even on finite
|
||||
// inputs at extreme F1-ep2 magnitudes (smoke-test-xvzgk slot 32
|
||||
// signature). Reuses `max_abs_q` from the slot 24/25 pre-call
|
||||
// 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`).
|
||||
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.launch_clamp_finite_f32(bn_d_concat_ptr, bn_d_concat_len, max_abs_q)?;
|
||||
}
|
||||
|
||||
// ── #31 Bottleneck backward: chain rule through tanh + GEMM ──
|
||||
// After backward_full, d_bn_concat_buf has dL/d(bn_concat) [B, concat_dim] f32.
|
||||
// Chain rule (post-1B-iii — bottleneck Linear's forward consumes
|
||||
@@ -20335,7 +20457,7 @@ fn deflate_rank_one(mat: &[f32], n: usize, lambda_1: f32) -> Vec<f32> {
|
||||
fn compile_training_kernels(
|
||||
stream: &Arc<CudaStream>,
|
||||
config: &GpuDqnTrainConfig,
|
||||
) -> Result<(CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction), MLError> {
|
||||
) -> Result<(CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction), MLError> {
|
||||
info!(
|
||||
state_dim = ml_core::state_layout::STATE_DIM,
|
||||
total_params = compute_total_params(config),
|
||||
@@ -20468,14 +20590,22 @@ fn compile_training_kernels(
|
||||
.map_err(|e| MLError::ModelError(format!("dqn_nan_check_f32 load: {e}")))?;
|
||||
let nan_check_f32_b = module.load_function("dqn_nan_check_f32_b")
|
||||
.map_err(|e| MLError::ModelError(format!("dqn_nan_check_f32_b load: {e}")))?;
|
||||
// SP1 Phase C surgical fix: in-place finite clamp kernel (replaces NaN/Inf
|
||||
// with 0; clamps finite |v| to ±max_abs from caller). Loaded from the same
|
||||
// `module` as the NaN check kernels because both are launched inside
|
||||
// `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 popart_normalize = module.load_function("popart_normalize_kernel")
|
||||
.map_err(|e| MLError::ModelError(format!("popart_normalize_kernel load: {e}")))?;
|
||||
let popart_robust = ungraphed_module.load_function("popart_normalize_robust")
|
||||
.map_err(|e| MLError::ModelError(format!("popart_normalize_robust load: {e}")))?;
|
||||
|
||||
info!("GpuDqnTrainer: 37 utility kernels loaded from precompiled cubin (5 CUmodules)");
|
||||
Ok((grad_norm, grad_norm_finalize, grad_norm_finalize_adam, grad_norm_finalize_post_aux, adam_update, adam_update_post_aux, scale_f32_post_aux, saxpy, zero, regime_scale, shrink_perturb, _relu_mask_from_module, spectral_norm, clip_grad, pad_states, saxpy_f32, scale_f32, stochastic_depth, stochastic_depth_rng, bn_tanh_bw, bn_bias_grad, bn_tanh_concat, vaccine_dot, vaccine_project, vaccine_dot_finalize, causal_intervene, causal_reduce, causal_mean_reduce, pruning_mask_fn, pruning_compute_fn, her_inplace, nan_check_f32, nan_check_f32_b, popart_normalize, saxpy_f32_adam_grad, saxpy_f32_aux, scale_f32_aux, distill_saxpy_aux, grad_norm_standalone_post_aux, shrink_perturb_ungraphed, scale_f32_ungraphed, popart_robust))
|
||||
info!("GpuDqnTrainer: 38 utility kernels loaded from precompiled cubin (5 CUmodules)");
|
||||
Ok((grad_norm, grad_norm_finalize, grad_norm_finalize_adam, grad_norm_finalize_post_aux, adam_update, adam_update_post_aux, scale_f32_post_aux, saxpy, zero, regime_scale, shrink_perturb, _relu_mask_from_module, spectral_norm, clip_grad, pad_states, saxpy_f32, scale_f32, stochastic_depth, stochastic_depth_rng, bn_tanh_bw, bn_bias_grad, bn_tanh_concat, vaccine_dot, vaccine_project, vaccine_dot_finalize, causal_intervene, causal_reduce, causal_mean_reduce, pruning_mask_fn, pruning_compute_fn, her_inplace, nan_check_f32, nan_check_f32_b, clamp_finite_f32, popart_normalize, saxpy_f32_adam_grad, saxpy_f32_aux, scale_f32_aux, distill_saxpy_aux, grad_norm_standalone_post_aux, shrink_perturb_ungraphed, scale_f32_ungraphed, popart_robust))
|
||||
}
|
||||
|
||||
/// Load the standalone Polyak EMA kernel from precompiled cubin.
|
||||
|
||||
@@ -2246,3 +2246,15 @@ SP1 Phase B foundation — stale-doc cleanup + Task 4 prep (2026-04-29): doc-onl
|
||||
SP1 Phase B accessors (2026-04-29): added 5 new accessor methods exposing backward-path buffer device pointers for Task 4's per-step NaN checks. **Scope diverged from plan** — the plan's "delegate accessors on `GpuDqnTrainer` for slots 27 + 28" is structurally invalid: `GpuDqnTrainer` does NOT own `GpuIqnHead` (the IQN head is owned by `FusedTrainingCtx` at `fused_training.rs:289`, alongside the trainer at line 234). The audit's per-slot accessor table (lines 535-536) is correct: slot 27/28 accessors land on `GpuIqnHead`, and Task 4's `run_nan_checks_post_backward(batch_size, iqn_d_h_s2_ptr, d_branch_logits_ptr, ...)` will receive the IQN pointers as `u64` arguments from the `FusedTrainingCtx` call site — same pattern already in use at `gpu_dqn_trainer.rs:6843` (`apply_iqn_trunk_gradient(&mut self, iqn_d_h_s2_ptr: u64, ...)`). Per `feedback_no_legacy_aliases.md`: no delegate wrappers — Task 4 either calls the IQN accessor directly via the `gpu_iqn` field on `FusedTrainingCtx` or receives the pointer via method args; no shadow accessor on `GpuDqnTrainer`. New accessors landed: `GpuDqnTrainer::{d_value_logits_buf_ptr, d_adv_logits_buf_ptr, cql_d_value_logits_ptr, aux_dh_s2_nb_buf_ptr}` (slots 24, 25, 29, 30 — 4 methods) + `GpuIqnHead::d_branch_logits_buf_ptr` (slot 28). Slot 27 (`iqn_d_h_s2_buf`) reuses the existing `GpuIqnHead::d_h_s2_raw_ptr()` accessor at `gpu_iqn_head.rs:1660` — no new method (per `feedback_no_legacy_aliases.md`, the existing accessor is sufficient; renaming + chasing one call site adds churn without value). Slots 26, 32, 33-35 reuse pre-existing handles (`self.ptrs.iqn_trunk_m`, `self.bn_d_concat_buf()` accessor, `self.ptrs.bw_d_h_s2`). Slot 31 (`ensemble_d_logits_buf`) deferred per Task 2 commit `387335e2b` (cross-struct on `FusedDqnTraining`). NO new scratch buffer added — the plan's `bw_d_h_s2_pre_saxpy` scratch + DtoD-copy approach was superseded by the audit's 3-call-site reformulation (slots 33/34/35 are post-main / post-aux / post-iqn snapshots of the same `bw_d_h_s2`, taken via inline call sites in Task 4). Pattern follows commit `e9096c7be`'s GRN-block accessor style (concise `pub(crate) fn name_ptr(&self) -> u64`, doc-comment referencing slot number + audit doc + buffer semantics). Additive — no behavioral change; new accessors consumed by Task 4's NaN check call sites.
|
||||
|
||||
SP1 Phase B instrumentation — Task 4 (2026-04-29): wired per-step backward-path NaN checks. New `GpuDqnTrainer::run_nan_checks_post_backward(batch_size, iqn_d_h_s2_ptr: Option<u64>, iqn_d_branch_logits_buf_ptr: Option<u64>)` mirrors `run_nan_checks_post_forward`'s pattern (single-block GPU reduce per buffer via `check_nan_f32`, no atomicAdd, no DtoH per step) and covers slots 24-30 + 32 — the eight backward-path buffers with stable post-backward state. Three INLINE `check_nan_f32` calls at orchestration call sites cover slots 33/34/35: **slot 33** in `launch_cublas_backward_to` AFTER `backward_full` + branch-concat accumulations and BEFORE `aux_heads_backward` (main-path entry-point check); **slot 34** in `launch_cublas_backward_to` AFTER `aux_heads_backward` SAXPY (compared against 33 to localise aux-head poisoning); **slot 35** in `apply_iqn_trunk_gradient` AFTER `graph_safe_copy_f32(bw_d_h_s2 ← iqn_d_h_s2_ptr, ...)` and BEFORE `encoder_backward_chain` (covers the IQN trunk DtoD overwrite — symmetric with slot 27 which checks the SOURCE side). Sequential 33→34→35 fire pattern localises the entry point: 33 alone = main backward chain; 34 fires after clean 33 = aux SAXPY; 35 fires after clean 34 = IQN DtoD or per-sample backward. **IQN pointers passed as `Option<u64>`** because `GpuIqnHead` is a sibling on `FusedTrainingCtx` (per Task 3 deviation finding, audit lines 535-536); `None` (when `iqn_lambda == 0.0` and `gpu_iqn = None`) cleanly skips slots 27/28 — semantically honest "buffer doesn't exist this run" rather than false-clean. **Call sites:** `fused_training.rs` ungraphed step at line 1519 + capture closure at line 2324 (both updated atomically — `feedback_no_partial_refactor`). The post-backward checks join the same `post_aux` captured child as the existing pre/post forward checks; the inline 33/34 checks land in the `forward` child (because `launch_cublas_backward_to` is called inside `submit_forward_ops_main`); slot 35 lands in the `aux` child (because `apply_iqn_trunk_gradient` is called inside `submit_aux_ops`). Slot 31 (`ensemble_d_logits_buf`) cleanly deferred per Task 2 commit `387335e2b` — its owner is `FusedDqnTraining`, and the ensemble Phase B saxpy guards must be verified before un-deferring; not partial — slot 31's deferral is documented in the new method's docstring + the audit doc. Lengths use `CudaSlice::len()` where the slice is owned by `GpuDqnTrainer` (slots 24, 25, 29, 30, 32 — auto-syncs with allocator padding), `iqn_trunk_m.len()` for slot 26 (matches the `trunk_params` allocation), `b * sh2` for slots 27/30/33/34/35, and inline `tba * b * q` arithmetic for slot 28 (since `GpuIqnHead::d_branch_logits_buf.len()` is private and computing it on the trainer side avoids adding a length accessor for one call site). Each check uses existing `check_nan_f32(buf_ptr, len, flag_idx)` infrastructure — no new kernel, no new scratch buffer, no DtoD/HtoD/HtoH copies. Flags accumulate within fold; reset on fold boundary via existing `reset_nan_flags()`. Readback flow (commit `d1808df14`) consumes them in BOTH `halt_nan` and `halt_grad_collapse` paths via the name tables annotated in commit `387335e2b`. Permanent diagnostic infrastructure — stays as production-grade regression sentinel after the surgical fix lands.
|
||||
|
||||
SP1 Phase C surgical fix (2026-04-29): patched two cuBLAS GEMM backward operations identified by Phase B smoke `smoke-test-xvzgk` at commit `f139a63ee`:
|
||||
|
||||
1. `apply_iqn_trunk_gradient` (gpu_dqn_trainer.rs:6885+) — pre-call clamp on `bw_d_h_s2` (the IQN DtoD-overwrite target that feeds the GRN trunk encoder's cuBLAS Linear_a/Linear_b dW/dX/dB GEMMs; symmetric with slot 27 source-side check); post-call sanitize on `iqn_trunk_m` (slot 26 buffer). ISV-driven max_abs bound = `1e6 * isv[H_S2_RMS_EMA_INDEX=96]` with `1e3` ε floor (Invariant 1 carve-out for ISV-uninitialized state).
|
||||
|
||||
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.
|
||||
|
||||
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).
|
||||
|
||||
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.
|
||||
|
||||
Reference in New Issue
Block a user