diff --git a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs index 4a380c7ba..ea4221799 100644 --- a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs +++ b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs @@ -6926,6 +6926,18 @@ impl GpuDqnTrainer { self.graph_safe_copy_f32(self.ptrs.bw_d_h_s2, iqn_d_h_s2_ptr, n_bytes, "iqn_d_h_s2")?; } + // SP1 Phase B (slot 35): bw_d_h_s2 snapshot AFTER the IQN DtoD + // overwrite, BEFORE the IQN aux `encoder_backward_chain` consumes + // it. Symmetric with slot 27 (which checks the SOURCE side + // `iqn_d_h_s2_ptr`); both are needed because the DtoD might be a + // graph-capture artefact (cudarc cuMemcpy in graph capture has + // historic edge cases — see `bf16_session_2026-03-28.md` and + // `session_2026-04-05_h100_hang.md`). If slot 27 is clean and + // slot 35 fires, the DtoD itself is suspect; if slot 27 fires, + // the IQN backward is the seed. Per-slot semantics in + // `docs/dqn-backward-nan-audit.md` (slot 35 row). + self.check_nan_f32(self.ptrs.bw_d_h_s2, b * sh2, 35)?; + // ── 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 { @@ -14891,12 +14903,15 @@ impl GpuDqnTrainer { /// 0-23 forward-path (this method populates 0-3, 6-12, 16-20; slots /// 4-5 are populated by `run_nan_checks_pre_forward`; 13-15 and /// 21-23 are reserved headroom). - /// 24-35 backward-path (wired by Task 4 in `run_nan_checks_post_backward`; - /// covers post-c51_grad, IQN/CQL/aux/ensemble backward outputs and - /// bw_d_h_s2 post-stage snapshots — see audit per-slot table). - /// Slot 31 (ensemble_d_logits_buf) is DEFERRED — owner is on - /// `FusedDqnTraining` (different struct) and Task 4 must skip it - /// until the ensemble Phase B saxpy guards are verified. + /// 24-35 backward-path (wired by Task 4 — see + /// `run_nan_checks_post_backward` for slots 24-30 + 32 with + /// stable post-backward state, plus three inline `check_nan_f32` + /// calls during backward orchestration for slots 33/34/35 + /// (`bw_d_h_s2` post-main / post-aux / post-IQN snapshots + /// localising the entry point of NaN poisoning into the + /// backward chain). Slot 31 (ensemble_d_logits_buf) DEFERRED — + /// owner is on `FusedDqnTraining` (different struct); will be + /// un-deferred when ensemble Phase B saxpy guards are verified. /// 36-47 reserved headroom (SP2/SP3). pub fn run_nan_checks_post_forward(&mut self, batch_size: usize) -> Result<(), MLError> { let b = batch_size; @@ -14964,6 +14979,137 @@ impl GpuDqnTrainer { Ok(()) } + /// SP1 Phase B (Task 4): NaN checks on backward-path kernel outputs that + /// have a stable post-backward state. Runs after the entire backward + /// orchestration has completed (main backward chain + IQN aux + CQL + + /// aux heads + bottleneck) but before Adam consumes `grad_buf`. + /// + /// Index → buffer mapping (mirrors training_loop.rs name table; per-slot + /// semantics + length expressions in `docs/dqn-backward-nan-audit.md` + /// per-slot accessor table at the end of that file): + /// 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 (caller-supplied) + /// 28 d_branch_logits_buf — production IQN backward (iqn_quantile_huber_loss) + /// 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 + /// + /// Slots 33/34/35 (`bw_d_h_s2` at three orchestration phases) are NOT + /// checked here — they fire inline during backward orchestration to + /// localize the NaN entry point. See `launch_cublas_backward_to` + /// (slots 33, 34) and `apply_iqn_trunk_gradient` (slot 35). + /// + /// Slot 31 (`ensemble_d_logits_buf`) is DEFERRED per Task 2 commit + /// `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. + /// + /// 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(()) + } + /// Launch GPU-side NaN check on an f32 buffer. Writes 1 to nan_flags_buf[flag_idx] if NaN found. /// No CPU sync — stays entirely on GPU. Call reset_nan_flags() before a batch of checks. pub fn check_nan_f32(&self, buf_ptr: u64, n: usize, flag_idx: i32) -> Result<(), MLError> { @@ -18420,6 +18566,20 @@ impl GpuDqnTrainer { )?; } + // SP1 Phase B (slot 33): bw_d_h_s2 snapshot AFTER backward_full + + // branch concat accumulations, BEFORE aux_heads_backward SAXPY. The + // buffer holds the "main path" gradient — branches 0/1/2/3 accumulated + // via `accumulate_d_h_s2_from_concat` (beta=0 for branch 0, beta=1 + // thereafter) plus the value-FC contribution via `backward_full`. + // Localises whether the main path corrupted the buffer before any + // aux-head SAXPY runs. Per-slot semantics in + // `docs/dqn-backward-nan-audit.md` (slot 33 row). + { + let b_post = self.config.batch_size; + let sh2_post = self.config.shared_h2; + self.check_nan_f32(self.ptrs.bw_d_h_s2, b_post * sh2_post, 33)?; + } + // ── Plan 4 Task 6 Commit B: aux-heads backward ── // Runs the per-head backward kernels, reduces 8 per-tensor partials // into final dW/dB, SAXPYs them into `grad_base[119..127)` with @@ -18429,6 +18589,18 @@ impl GpuDqnTrainer { // augmented dh_s2 (= main + aux contributions). self.aux_heads_backward(grad_base)?; + // SP1 Phase B (slot 34): bw_d_h_s2 snapshot AFTER aux_heads_backward + // SAXPYed `aux_dh_s2_nb_buf` + `aux_dh_s2_rg_buf` into the trunk + // accumulator. Compared against slot 33: any NEW NaN here came from + // the aux-head SAXPY path (the saxpy kernel currently has no + // isfinite guard — Phase C proposal). Per-slot semantics in + // `docs/dqn-backward-nan-audit.md` (slot 34 row). + { + let b_post = self.config.batch_size; + let sh2_post = self.config.shared_h2; + self.check_nan_f32(self.ptrs.bw_d_h_s2, b_post * sh2_post, 34)?; + } + // ── Phase 3 T3.6: MoE backward ───────────────────────────────── // Runs AFTER aux_heads_backward (which SAXPYed into bw_d_h_s2) and // BEFORE encoder_backward_chain. Computes dW/dB for MoE params diff --git a/crates/ml/src/trainers/dqn/fused_training.rs b/crates/ml/src/trainers/dqn/fused_training.rs index b2836ffdd..6d6d58a0c 100644 --- a/crates/ml/src/trainers/dqn/fused_training.rs +++ b/crates/ml/src/trainers/dqn/fused_training.rs @@ -1519,6 +1519,30 @@ 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 + // `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 + // `reset_nan_flags()`. Captures into the same `post_aux` child + // 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}"))?; + // 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. { @@ -2320,10 +2344,24 @@ impl FusedTrainingCtx { s.submit_aux_ops(agent, gpu_batch) })?; - // Capture post_aux child (selectivity + denoise + Q-stats + risk_sgd). + // Capture post_aux child (selectivity + denoise + Q-stats + risk_sgd + // + SP1 Phase B post-backward NaN checks). The post-backward NaN + // checks join the same captured child as the existing pre/post + // forward checks since both are GPU-side reduce kernels writing + // into `nan_flags_buf` — keeping them adjacent in the graph + // minimises launch overhead and makes the slot-name table + // (in `training_loop.rs`) align with one capture point per group. 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}")) + .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}")) })?; // 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 08e3615a2..88729ad7f 100644 --- a/docs/dqn-wire-up-audit.md +++ b/docs/dqn-wire-up-audit.md @@ -2244,3 +2244,5 @@ SP1 Phase B foundation (2026-04-29): expanded `nan_flags_buf` 24→48 (allocatio SP1 Phase B foundation — stale-doc cleanup + Task 4 prep (2026-04-29): doc-only follow-up to commit `53bc0bc50`. Replaced the 24-slot index map docstring on `GpuDqnTrainer::run_nan_checks_post_forward` with a 48-slot range summary that defers per-slot semantics to `docs/dqn-backward-nan-audit.md` per-slot accessor table (DRY — audit is the source of truth). Updated the `halt_grad_collapse` diagnostic comment in `training_loop.rs` from `[24] system` (13 base + 5 GRN-stage) to the post-expansion 48-slot layout (slots 0-23 forward / 24-35 backward / 36-47 reserved). Fixed the stale `0..11` tracing message (kept from before the 16→24 GRN expansion) to `0..47`. Annotated slot 31 (`ensemble_d_logits_buf`) in BOTH `training_loop.rs` name tables as DEFERRED with cross-struct ownership note (`FusedDqnTraining`) — this prevents Task 4 from blanket-launching `check_nan_f32` on slot 31's null `GpuDqnTrainer` accessor before the ensemble Phase B saxpy guards are verified. Pre-emptively updated both name-table header comments to reference the future `run_nan_checks_post_backward` method (Task 4) plus the audit's per-slot table — drops the 3-way name-table contract drift risk when Task 4 lands. Both name tables remain byte-identical (modulo indentation). Per `feedback_no_partial_refactor.md`: the 24→48 expansion shipped without doc-coverage on the consumer side; this cleanup commit closes that residue before Task 4 starts. 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, iqn_d_branch_logits_buf_ptr: Option)` 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`** 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.