gitadmin/foxhunt
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

feat(dqn): SP2 — replace 8 check_nan_f32 calls with fused launch

Browse Source 
run_nan_checks_post_backward now fires a single fused kernel launch
covering slots 24-35 in 12 blocks. Reduces graph-capture overhead from
8 launches to 1 per step. Slots 33-35 inline checks at backward
orchestration sites (gpu_dqn_trainer.rs:18580/:18601/:6939) unchanged
— multi-point localization preserved.

Method signature simplified — IQN pointers no longer per-step args.
populate_nan_check_meta (called once at construction in fused_training.rs)
baked Option<u64> nullity into the metadata buffer entries; the fused
kernel's null-pointer guard handles slots 27/28 when IQN inactive
without per-step branching at the Rust caller.

Both call sites in fused_training.rs (ungraphed + graph-captured paths)
updated together per feedback_no_partial_refactor.

This is the F0 regression fix — Phase B's per-step kernel-launch
overhead was the F0 cause across 3 SP1 smokes (F0 ~ 35 vs baseline
55.87). Gate 1 smoke (Task A6) validates F0 >= 53.08.
This commit is contained in:
jgrusewski
2026-04-30 09:22:52 +02:00
parent fbf48df9de
commit b5a064f6ca
3 changed files with 37 additions and 126 deletions
Download Patch File Download Diff File Expand all files Collapse all files

122
crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs
View File

@@ -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<u64>` 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<u64> 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<u64>,
iqn_d_branch_logits_buf_ptr: Option<u64>,
) -> 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::<f32>(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<f32>`;
// 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.

39
crates/ml/src/trainers/dqn/fused_training.rs
View File

@@ -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<u64>` 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<u64> 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<u64> 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

2
docs/dqn-wire-up-audit.md
View File

@@ -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<u64>` (12 entries, one per slot 24-35) and `nan_check_buf_lens: CudaSlice<i32>` (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<u64>` / `nan_check_buf_lens: CudaSlice<i32>` 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<u64>` 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<u64>` 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.