Wires per-step NaN checks on backward-path kernel outputs. Coverage
per audit per-slot accessor table (docs/dqn-backward-nan-audit.md
:530-548).
GpuDqnTrainer::run_nan_checks_post_backward (NEW) — fire-once-at-end:
- 24 d_value_logits_buf (post-c51_grad value gradient)
- 25 d_adv_logits_buf (post-c51_grad branch advantage)
- 26 iqn_trunk_m (apply_iqn_trunk_gradient cuBLAS bwd output)
- 27 iqn_d_h_s2_buf (IQN backward dh_s2; arg from FusedTrainingCtx)
- 28 d_branch_logits_buf (IQN production backward; arg from caller)
- 29 cql_d_value_logits (CQL gradient output)
- 30 aux_dh_s2_nb_buf (Aux next-bar backward dh_s2)
- 32 bn_d_concat_buf (Bottleneck Linear backward dy)
Inline checks during backward orchestration:
- 33 bw_d_h_s2 post-main (in launch_cublas_backward_to, after
backward_full + branch concat accum,
BEFORE aux_heads_backward SAXPY)
- 34 bw_d_h_s2 post-aux (in launch_cublas_backward_to, after
aux_heads_backward SAXPY)
- 35 bw_d_h_s2 post-iqn (in apply_iqn_trunk_gradient, after
graph_safe_copy_f32 DtoD overwrite,
BEFORE encoder_backward_chain consumes)
Sequential 33→34→35 fire pattern localises the NaN entry point:
- 33 alone fires → main backward chain (c51 + MSE + branch concat)
- 34 fires after 33-clean → aux SAXPY (aux_heads_backward)
- 35 fires after 34-clean → IQN DtoD or per-sample IQN backward
Slot 31 (ensemble_d_logits_buf) cleanly deferred per Task 2 commit
387335e2b — owner is FusedDqnTraining (different struct); will be
un-deferred when ensemble Phase B saxpy guards are verified.
IQN pointers (slots 27, 28) passed as Option<u64> from FusedTrainingCtx
because GpuDqnTrainer does NOT own GpuIqnHead (Task 3 deviation
finding, audit lines 535-536). None case (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 signal.
Pattern matches apply_iqn_trunk_gradient(iqn_d_h_s2_ptr: u64, ...)
at gpu_dqn_trainer.rs:6843.
Call sites (atomic — feedback_no_partial_refactor):
- fused_training.rs:1519 ungraphed step path
- fused_training.rs:2324 capture_training_graph closure (post_aux child)
- gpu_dqn_trainer.rs::launch_cublas_backward_to (slots 33, 34 — captured
in forward child via submit_forward_ops_main)
- gpu_dqn_trainer.rs::apply_iqn_trunk_gradient (slot 35 — captured in
aux child via submit_aux_ops)
Each check uses existing check_nan_f32(buf_ptr, len, flag_idx) —
single-block GPU reduce, no atomicAdd, no DtoD/HtoD/HtoH copies, no
per-step DtoH. Lengths use CudaSlice::len() where possible (auto-syncs
with allocator padding); inline arithmetic for slot 28 since
d_branch_logits_buf.len() is private to GpuIqnHead. Flags accumulate
within fold; reset at fold boundary via reset_nan_flags(). Readback
flow (commit d1808df14) consumes them in BOTH halt_nan and
halt_grad_collapse paths via name tables annotated in commit 387335e2b.
Permanent diagnostic infrastructure — stays as production-grade
regression sentinel after the surgical fix lands.
f139a63eea