feat(dqn): SP1 Phase B instrumentation — backward NaN checks (slots 24-30, 32-35)
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.
This commit is contained in:
jgrusewski
@@ -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<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.
|
||||
///
|
||||
/// 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(())
|
||||
}
|
||||
|
||||
/// 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
|
||||
|
||||
@@ -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<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}"))?;
|
||||
|
||||
// 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
|
||||
|
||||
@@ -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<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.
|
||||
|
||||
Reference in New Issue
Block a user