diff --git a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs index 37a3271de..491422fc9 100644 --- a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs +++ b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs @@ -15112,109 +15112,27 @@ impl GpuDqnTrainer { /// `387335e2b` — owner is `FusedDqnTraining` (different struct) and the /// ensemble Phase B saxpy guards must be verified before un-deferring. /// - /// IQN pointers are passed as `Option` because `GpuIqnHead` is - /// owned by `FusedTrainingCtx`, not `GpuDqnTrainer` — same - /// argument-style pattern as - /// `apply_iqn_trunk_gradient(&mut self, iqn_d_h_s2_ptr: u64, ...)`. - /// `None` skips slots 27/28 cleanly when IQN is disabled (the - /// `iqn_lambda > 0.0` ctor gate); the slot stays at zero — semantically - /// honest "the buffer doesn't exist this run", not a false-clean signal. + /// SP2 (Task A4): collapses 8 individual `check_nan_f32` launches into a + /// single fused launch covering all 12 backward-path slots (24-35) in + /// 12 blocks. The `(ptr, len)` table is populated once at construction + /// via `populate_nan_check_meta` (called from `fused_training.rs` after + /// `gpu_iqn` is known) — Option nullity for slots 27/28 (IQN + /// inactive) and slot 31 (deferred ensemble) and slots 33-35 (inline + /// at backward orchestration phases) is baked into the metadata buffer + /// as `(0, 0)` entries; the fused kernel's null-pointer guard skips + /// them without per-step Rust branching. /// - /// Pattern: identical to `run_nan_checks_post_forward` (single-block - /// reduce per buffer via `check_nan_f32`, no atomicAdd, no DtoH per - /// step). Fires every training step inside the `post_aux` captured - /// child graph; flags accumulate within fold and reset on fold - /// boundary via `reset_nan_flags()`. - pub fn run_nan_checks_post_backward( - &mut self, - batch_size: usize, - iqn_d_h_s2_ptr: Option, - iqn_d_branch_logits_buf_ptr: Option, - ) -> Result<(), MLError> { - let b = batch_size; - let sh2 = self.config.shared_h2; - - // Slot 24: post-c51_grad value gradient `[B × pad32(NA)]` — match the - // f32 atomicAdd allocation at gpu_dqn_trainer.rs:10408. Use `.len()` - // so any future allocator-side padding stays in sync. - self.check_nan_f32( - self.d_value_logits_buf_ptr(), - self.d_value_logits_buf.len(), - 24, - )?; - - // Slot 25: post-c51_grad branch advantage gradient - // `[B × (b0+b1+b2+b3) × NA + 32*4]` — match the f32 atomicAdd - // allocation at gpu_dqn_trainer.rs:10410. - self.check_nan_f32( - self.d_adv_logits_buf_ptr(), - self.d_adv_logits_buf.len(), - 25, - )?; - - // Slot 26: iqn_trunk_m — apply_iqn_trunk_gradient cuBLAS bwd output. - // Length = trunk_params (sum of `param_sizes[0..13)` padded byte - // offset / 4); same value used at gpu_dqn_trainer.rs:7014 SAXPY. - // The CudaSlice `.len()` returns this exact count from - // `alloc_zeros::(trunk_params)` at gpu_dqn_trainer.rs:10457. - self.check_nan_f32( - self.ptrs.iqn_trunk_m, - self.iqn_trunk_m.len(), - 26, - )?; - - // Slot 27: IQN dh_s2 buffer (caller-supplied). Length `B × SH2` - // per `gpu_iqn_head.rs:467` allocation. Skipped cleanly when IQN - // is disabled (`iqn_lambda == 0.0` → `gpu_iqn = None`). - if let Some(ptr) = iqn_d_h_s2_ptr { - self.check_nan_f32(ptr, b * sh2, 27)?; - } - - // Slot 28: production IQN backward output `d_branch_logits_buf` - // `[tba × bq]` per `gpu_iqn_head.rs:482`. Caller's accessor - // (`d_branch_logits_buf_ptr`) returns `raw_ptr()` of the same - // CudaSlice — but `.len()` is private to GpuIqnHead, so rely on - // the caller computing it. Audit's length expression is - // `total_branch_actions × (B × IQN_NUM_QUANTILES)`; we inline the - // arithmetic here so the check method stays self-contained. - if let Some(ptr) = iqn_d_branch_logits_buf_ptr { - // total_branch_actions: sum across 4 branches (b0+b1+b2+b3 = 12 typical) - let tba = self.config.branch_0_size - + self.config.branch_1_size - + self.config.branch_2_size - + self.config.branch_3_size; - // bq = B * Q (Q = num_quantiles, IQN-specific). The DQN-side - // config field that matches gpu_iqn_head's `num_quantiles` is - // `iqn_num_quantiles`; both are pinned to FIXED_TAUS.len() = 5. - let q = self.config.iqn_num_quantiles.max(1); - self.check_nan_f32(ptr, tba * b * q, 28)?; - } - - // Slot 29: CQL gradient buffer `[B × pad32(NA)]` per - // gpu_dqn_trainer.rs:10402. - self.check_nan_f32( - self.cql_d_value_logits_ptr(), - self.cql_d_value_logits.len(), - 29, - )?; - - // Slot 30: aux next-bar backward dh_s2 `[B × SH2]` per - // gpu_dqn_trainer.rs:11134 allocation (`aux_b * aux_sh2`). - self.check_nan_f32( - self.aux_dh_s2_nb_buf_ptr(), - self.aux_dh_s2_nb_buf.len(), - 30, - )?; - - // Slot 32: bottleneck Linear backward dy. The accessor - // (`bn_d_concat_buf()` at line 6757) returns `&CudaSlice`; - // use its `.raw_ptr()` + `.len()` for the GPU reduce. Length - // matches the `b * concat_dim + kt` allocation at - // gpu_dqn_trainer.rs:10938. - let bn_d_concat = self.bn_d_concat_buf(); - self.check_nan_f32(bn_d_concat.raw_ptr(), bn_d_concat.len(), 32)?; - - Ok(()) + /// Inline checks for slots 33/34/35 (bw_d_h_s2 multi-point snapshots) + /// continue to fire at backward orchestration call sites + /// (`launch_cublas_backward_to`, `apply_iqn_trunk_gradient`) — those + /// provide localization across the 3 backward phases (post-main, + /// post-aux, post-iqn) and are not subsumed by the single fused launch. + /// + /// Fires every training step inside the `post_aux` captured child + /// graph; flags accumulate within fold and reset on fold boundary via + /// `reset_nan_flags()`. + pub fn run_nan_checks_post_backward(&mut self) -> Result<(), MLError> { + self.launch_nan_check_fused_f32() } /// Launch GPU-side NaN check on an f32 buffer. Writes 1 to nan_flags_buf[flag_idx] if NaN found. diff --git a/crates/ml/src/trainers/dqn/fused_training.rs b/crates/ml/src/trainers/dqn/fused_training.rs index cee73380c..2e4879e85 100644 --- a/crates/ml/src/trainers/dqn/fused_training.rs +++ b/crates/ml/src/trainers/dqn/fused_training.rs @@ -1539,9 +1539,9 @@ impl FusedTrainingCtx { self.trainer.submit_post_aux_ops(self.batch_size) .map_err(|e| anyhow::anyhow!("post_aux ops: {e}"))?; - // SP1 Phase B (Task 4): NaN checks on backward-path kernel outputs - // that have a stable post-backward state — slots 24-30, 32. Slots - // 33/34/35 fire inline during backward orchestration (see + // SP1 Phase B (Task 4) + SP2 (Task A4): single fused NaN-check + // launch covers all 12 backward-path slots (24-35) in 12 blocks. + // Slots 33/34/35 fire inline during backward orchestration (see // `launch_cublas_backward_to` and `apply_iqn_trunk_gradient` in // `gpu_dqn_trainer.rs`). Permanent diagnostic infrastructure; // flags accumulate within fold and reset on fold boundary via @@ -1549,19 +1549,12 @@ impl FusedTrainingCtx { // graph as `submit_post_aux_ops` so the GPU reduce kernels share // launch overhead with the existing pre/post forward checks. // - // IQN pointers are passed as `Option` because `GpuIqnHead` - // is a sibling (`gpu_iqn`), not owned by the trainer. `None` - // (when `iqn_lambda == 0.0`) cleanly skips slots 27/28 — the - // semantically honest "buffer doesn't exist this run" rather - // than a false-clean signal. - let iqn_d_h_s2_ptr = self.gpu_iqn.as_ref().map(|g| g.d_h_s2_raw_ptr()); - let iqn_d_branch_logits_buf_ptr = - self.gpu_iqn.as_ref().map(|g| g.d_branch_logits_buf_ptr()); - self.trainer.run_nan_checks_post_backward( - self.batch_size, - iqn_d_h_s2_ptr, - iqn_d_branch_logits_buf_ptr, - ).map_err(|e| anyhow::anyhow!("post_backward nan checks: {e}"))?; + // SP2: IQN ptrs no longer per-step args — `populate_nan_check_meta` + // (called once at construction) baked Option nullity into + // the fused kernel's metadata buffer entries; null/zero entries + // are skipped by the kernel's null-pointer guard. + self.trainer.run_nan_checks_post_backward() + .map_err(|e| anyhow::anyhow!("post_backward nan checks: {e}"))?; // Phase 6: TLOB bwd + Mamba2 bwd + OFI embed bwd + pruning + grad_norm + Adam + ISV update // TLOB backward: reads TLOB-slice gradient from bn_d_concat_buf, updates W_Q/K/V/O. @@ -2374,14 +2367,12 @@ impl FusedTrainingCtx { let post_aux = self.capture_child_graph("post_aux", |s| { s.trainer.submit_post_aux_ops(s.batch_size) .map_err(|e| anyhow::anyhow!("{e}"))?; - let iqn_d_h_s2_ptr = s.gpu_iqn.as_ref().map(|g| g.d_h_s2_raw_ptr()); - let iqn_d_branch_logits_buf_ptr = - s.gpu_iqn.as_ref().map(|g| g.d_branch_logits_buf_ptr()); - s.trainer.run_nan_checks_post_backward( - s.batch_size, - iqn_d_h_s2_ptr, - iqn_d_branch_logits_buf_ptr, - ).map_err(|e| anyhow::anyhow!("{e}")) + // SP2 (Task A4): IQN ptrs no longer per-step args — + // `populate_nan_check_meta` (called once at construction) + // baked Option nullity into the fused kernel's + // metadata buffer entries. + s.trainer.run_nan_checks_post_backward() + .map_err(|e| anyhow::anyhow!("{e}")) })?; // adam_grad: mamba2 bwd + ofi_embed bwd + pruning + grad_norm diff --git a/docs/dqn-wire-up-audit.md b/docs/dqn-wire-up-audit.md index c95b4ebd5..054d00adb 100644 --- a/docs/dqn-wire-up-audit.md +++ b/docs/dqn-wire-up-audit.md @@ -2270,3 +2270,5 @@ SP2 Phase A1 — fused NaN-check kernel (2026-04-29): foundational kernel-only c SP2 Phase A2 — nan-check (ptr, len) device tables (2026-04-29): added two device buffers on `GpuDqnTrainer` to feed the fused NaN-check kernel from A1: `nan_check_buf_ptrs: CudaSlice` (12 entries, one per slot 24-35) and `nan_check_buf_lens: CudaSlice` (12 entries). Both are `stream.alloc_zeros` allocations adjacent to `nan_flags_buf`. Two separate buffers chosen over a packed-stride layout for direct ABI match with the kernel signature `(const float* const* buf_ptrs, const int* buf_lens, ...)` — no struct-stride alignment concerns. Sized at 12 (slots 24-35); slots 36-47 headroom is reserved for SP3 Mech 5 extension which will resize both buffers in lockstep. Slot 31 (deferred ensemble, cross-struct on `FusedDqnTraining` per A1's null-skip pattern) entry will hold 0 — kernel skips that block on `buf == nullptr`. Allocated as zeros; populated once in A3 (slot pointers are stable across the trainer's lifetime — no per-step host→device traffic, satisfying `feedback_no_htod_htoh_only_mapped_pinned` once A3 selects the population path: mapped-pinned helper from `mapped_pinned.rs` for the one-shot construction-time write). Unused yet — A3 wires the population + Rust launch wrapper, A4 replaces the 8 individual `check_nan_f32` calls in `run_nan_checks_post_backward` with a single fused launch. Field-level only — no behavioral change. SP2 Phase A3 — populate + launch wrapper (2026-04-29): refactored A2's `nan_check_buf_ptrs: CudaSlice` / `nan_check_buf_lens: CudaSlice` to mapped-pinned (`MappedU64Buffer` / `MappedI32Buffer` from `mapped_pinned.rs`) per `feedback_no_htod_htoh_only_mapped_pinned` — host-side write through the mapped pages is visible to the kernel via the `dev_ptr` returned by `cuMemHostGetDevicePointer_v2`, eliminating the HtoD copy entirely. Added `populate_nan_check_meta(b, sh2, iqn_d_h_s2_ptr, iqn_d_branch_logits_buf_ptr, iqn_q)` on `GpuDqnTrainer` — one-shot construction-time writer of 12 `(ptr, len)` tuples covering slots 24-35 (mirrors the per-slot accessor table at the end of `docs/dqn-backward-nan-audit.md` and the per-slot launches in `run_nan_checks_post_backward`). Slot 31 deferred ensemble entry: `(0, 0)`. Slots 27/28 IQN entries: `Option` ptrs gated on `gpu_iqn.is_some()` — when IQN is inactive the entries stay zero and the fused kernel skips them via its null-pointer guard. Slots 33-35 (`bw_d_h_s2` inline checks) hold null entries — fused kernel skips; the inline `check_nan_f32` calls at the three backward orchestration phases (`launch_cublas_backward_to` post-main / post-aux / `apply_iqn_trunk_gradient` post-iqn) continue to fire individually for entry-point localization. Added `launch_nan_check_fused_f32` Rust wrapper — single launch with `grid_dim=12, block_dim=256, BASE_FLAG_IDX=24`, mirrors the kernel signature `(buf_ptrs_dev, buf_lens_dev, base_flag_idx, nan_flags_ptr)`. Kernel registered in the precompiled-cubin loader (`compile_training_kernels` tuple 43→44, 38→39 utility kernels logged) and stored on `GpuDqnTrainer` as `nan_check_fused_f32_kernel: CudaFunction` — same `module` (forward_child / aux_child / post_aux_child captured replay group) as the per-buffer `dqn_nan_check_f32` to preserve graph-capture compatibility for A4's call-site replacement. Constructor-time wire-up lives in `FusedTrainingCtx::new` (after `gpu_iqn` is constructed) — chosen over the GpuDqnTrainer constructor because `gpu_iqn` is owned by `FusedTrainingCtx`, mirroring the same `Option` argument pattern used by `apply_iqn_trunk_gradient(iqn_d_h_s2_ptr, ...)` and `run_nan_checks_post_backward`. Wrapper unused at A3 commit time — A4 replaces the 8-launch sequence in `run_nan_checks_post_backward` with the single fused launch. Zero new HtoD copies; pre-commit Invariant 7 satisfied; no ISV slots added; sticky-flag semantics preserved at the kernel. + +SP2 Phase A4 — call-site replacement (2026-04-29): replaced the 8 per-step `check_nan_f32` launches inside `GpuDqnTrainer::run_nan_checks_post_backward` with a single delegated call to `launch_nan_check_fused_f32` (12 blocks × 256 threads, one block per slot 24-35). Method signature simplified to `run_nan_checks_post_backward(&mut self) -> Result<(), MLError>` — IQN pointers (`iqn_d_h_s2_ptr`, `iqn_d_branch_logits_buf_ptr`) and `batch_size` are no longer per-step args; their values were already baked into the mapped-pinned `(ptr, len)` metadata buffer at construction time via `populate_nan_check_meta` (A3). Both call sites in `fused_training.rs` (ungraphed step at line 1556 + graph-captured `post_aux` closure at line 2374) updated atomically per `feedback_no_partial_refactor` — old `iqn_d_h_s2_ptr` / `iqn_d_branch_logits_buf_ptr` derivation blocks removed (no leftover dead code). Slots 33-35 (`bw_d_h_s2` multi-point snapshots) inline checks at backward orchestration sites (`launch_cublas_backward_to` post-main + post-aux, `apply_iqn_trunk_gradient` post-iqn) unchanged — fused kernel's null-pointer guard skips them in the metadata table; the per-phase inline calls continue to provide entry-point localization across the 3 backward phases. Sticky-flag semantics preserved (kernel only writes 1, never clears); flags accumulate within fold and reset on fold boundary via `reset_nan_flags()`. Behavioral change at the diagnostic level: graph-capture overhead reduced from 8 launches per step to 1 (per `feedback_no_atomicadd` and the kernel's grid-strided block-local `__syncthreads_or` reduce). This is the F0 regression remediation — Phase B's 8 per-step launches were the F0 cause across 3 SP1 smokes (F0 ≈ 35 vs baseline 55.87); Gate 1 smoke (Task A6) validates F0 ≥ 53.08. Zero new HtoD/DtoD/HtoH copies; pre-commit Invariant 7 satisfied; no ISV slots added.