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

fix(dqn): SP1 Phase B foundation — stale-doc cleanup + Task 4 prep

Browse Source 
Quality-review follow-ups to commit 53bc0bc50 (nan_flags_buf 24->48):

1. Update run_nan_checks_post_forward docstring (gpu_dqn_trainer.rs:14842-14873):
   replace 24-slot map with 48-slot range summary + audit-doc pointer.
   Was actively misleading after the 24->48 expansion; partial-refactor
   residue per feedback_no_partial_refactor.

2. Update '[24] system' comment in training_loop.rs (around line 2044):
   reflect the 48-slot post-expansion state (slots 0-23 fwd, 24-35 bwd,
   36-47 reserved). Also fix stale '0..11' tracing message to '0..47'.

3. Slot 31 (ensemble_d_logits_buf) annotation: flag DEFERRED + owner
   on FusedDqnTraining (different struct than 24-30, 32-35). Prevents
   Task 4 from blanket-launching check_nan_f32 on slot 31's null
   accessor.

4. Both name-table header comments now reference the future
   run_nan_checks_post_backward method (Task 4) plus the audit's
   per-slot table — pre-empts contract drift when Task 4 lands a
   3-way name-table dependency.

Audit-doc entry appended to docs/dqn-wire-up-audit.md SP1 Phase B
section. No behavioral change. Both name tables remain byte-identical.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-04-30 00:10:37 +02:00
parent 53bc0bc505
commit 387335e2b9
3 changed files with 44 additions and 48 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -14843,34 +14843,19 @@ impl GpuDqnTrainer {
/// Checks cuBLAS output logits, MSE intermediates, gradients, params, and
/// loss-component buffers (C51 target dist, MoE gate softmax, aux losses).
///
/// Flag index map (24 slots):
/// 0 states_buf (f32 padded) — input
/// 1 on_v_logits — cuBLAS forward (value stream)
/// 2 on_b_logits — cuBLAS forward (branch advantage)
/// 3 mse_loss_buf [1] — MSE loss scalar
/// 4 params_buf (f32, pre-forward) — set by run_nan_checks_pre_forward
/// 5 params_ptr (redundant pre-fwd) — set by run_nan_checks_pre_forward
/// 6 grad_buf — cuBLAS backward output
/// 7 save_current_lp — C51 online softmax probs (post-fwd)
/// 8 save_projected — C51 target distribution (most NaN-prone)
/// 9 moe_gate_softmax — MoE gate softmax (saturates → NaN)
/// 10 aux_nb_loss_scalar [1] — aux next-bar MSE loss
/// 11 aux_rg_loss_scalar [1] — aux regime CE loss
/// 12 save_h_s2 — trunk output (input to value+branch FCs)
/// 13-15 reserved (CQL / IQN / future)
/// ── P (post-bkdx5) GRN sub-stage NaN checks ──
/// 16 grn_h_s2_elu_post (state.elu_post) — post-ELU activation
/// 17 grn_h_s2_linear_b_out — Linear_b output (pre-GLU)
/// 18 grn_h_s2_glu_sigmoid (state.sigmoid_b) — GLU gate sigmoid
/// 19 grn_h_s2_ln_rstd (state.ln_rstd) — LN reciprocal-stddev
/// (variance→0 ⇒ rstd→∞ → NaN)
/// 20 grn_h_s2_ln_normed (state.ln_normed) — LN normalized (pre-affine)
/// (semantically same buffer
/// family as save_h_s2 but
/// pre-residual+affine; useful
/// to disambiguate residual
/// NaN-source vs LN-source)
/// 21-23 reserved (h_s1 stages / future)
/// SP1 Phase B 48-slot layout (high-level summary; per-slot accessor +
/// semantics live in `docs/dqn-backward-nan-audit.md` per-slot accessor
/// table — that audit is the source of truth):
/// 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.
/// 36-47 reserved headroom (SP2/SP3).
pub fn run_nan_checks_post_forward(&mut self, batch_size: usize) -> Result<(), MLError> {
let b = batch_size;
let na = self.config.num_atoms;

49
crates/ml/src/trainers/dqn/trainer/training_loop.rs
View File

@@ -1980,12 +1980,14 @@ impl DQNTrainer {
if let Some(ref mut fused) = self.fused_ctx {
if let Ok(flags) = fused.read_nan_flags() {
// Index → buffer name. Must match
// GpuDqnTrainer::run_nan_checks_post_forward and
// run_nan_checks_pre_forward. SP1 Phase B 48-slot
// layout: slots 0-23 forward-path (existing);
// slots 24-35 backward-path (per audit
// docs/dqn-backward-nan-audit.md per-slot
// accessor table — wired by Task 4); slots 36-47
// GpuDqnTrainer::run_nan_checks_pre_forward /
// run_nan_checks_post_forward /
// run_nan_checks_post_backward (Task 4 — slots
// 24-35). See docs/dqn-backward-nan-audit.md
// per-slot accessor table for slot 24-47
// semantics. SP1 Phase B 48-slot layout: slots
// 0-23 forward-path (existing); slots 24-35
// backward-path (wired by Task 4); slots 36-47
// reserved headroom for SP2/SP3.
let names = [
"states_buf", // 0 input states (post-upload)
@@ -2015,7 +2017,7 @@ impl DQNTrainer {
"d_branch_logits_buf", // 28 production IQN backward output (iqn_quantile_huber_loss)
"cql_d_value_logits", // 29 CQL gradient output
"aux_dh_s2_nb_buf", // 30 aux next-bar backward dh_s2
"ensemble_d_logits_buf", // 31 ensemble backward output (call site in fused_training.rs)
"ensemble_d_logits_buf", // 31 ensemble (DEFERRED — owner: FusedDqnTraining; instrument only after Phase B saxpy guards verified; see audit per-slot table)
"bn_d_concat_buf", // 32 bottleneck Linear backward dy
"bw_d_h_s2_post_main", // 33 bw_d_h_s2 after main backward chain
"bw_d_h_s2_post_aux", // 34 bw_d_h_s2 after aux head SAXPY
@@ -2030,7 +2032,7 @@ impl DQNTrainer {
.collect();
tracing::error!(
"NaN SOURCE at step {}: flagged=[{}] (empty = NaN entered via \
loss-component buffer not in checks 0..11; expand coverage)",
loss-component buffer not in checks 0..47; expand coverage)",
train_step_count, flagged.join(", ")
);
}
@@ -2041,10 +2043,14 @@ impl DQNTrainer {
));
}
if gr.halt_grad_collapse {
// Plan C T11 diagnostic extension (2026-04-29): the existing
// nan_flags_buf [24] system covers 18 buffers (13 base + 5
// GRN-stage P additions) and runs NaN checks every step
// inside the captured graph, but `read_nan_flags()` was
// Plan C T11 diagnostic extension (2026-04-29; SP1 Phase B
// expansion 24→48): the existing nan_flags_buf [48]
// system runs NaN checks every step inside the captured
// graph. Slots 0-23 are forward-path (13 base + 5
// GRN-stage P additions + 6 reserved), slots 24-35 are
// backward-path (Task 4 wires per-step checks via
// run_nan_checks_post_backward), slots 36-47 are
// reserved headroom for SP2/SP3. `read_nan_flags()` was
// previously only called on `halt_nan`. The training
// guard's halt_nan path checks pinned readback `grad_norm`
// — and when a kernel produces NaN that is then clamped to
@@ -2060,13 +2066,16 @@ impl DQNTrainer {
if let Some(ref mut fused) = self.fused_ctx {
if let Ok(flags) = fused.read_nan_flags() {
// Index → buffer name. Must match
// GpuDqnTrainer::run_nan_checks_post_forward
// and run_nan_checks_pre_forward. SP1 Phase B
// 48-slot layout: slots 0-23 forward-path
// (existing); slots 24-35 backward-path (per
// audit docs/dqn-backward-nan-audit.md
// per-slot accessor table — wired by Task 4);
// slots 36-47 reserved headroom for SP2/SP3.
// GpuDqnTrainer::run_nan_checks_pre_forward /
// run_nan_checks_post_forward /
// run_nan_checks_post_backward (Task 4 —
// slots 24-35). See
// docs/dqn-backward-nan-audit.md per-slot
// accessor table for slot 24-47 semantics.
// SP1 Phase B 48-slot layout: slots 0-23
// forward-path (existing); slots 24-35
// backward-path (wired by Task 4); slots
// 36-47 reserved headroom for SP2/SP3.
let names = [
"states_buf", // 0 input states (post-upload)
"on_v_logits", // 1 cuBLAS fwd: value stream
@@ -2095,7 +2104,7 @@ impl DQNTrainer {
"d_branch_logits_buf", // 28 production IQN backward output (iqn_quantile_huber_loss)
"cql_d_value_logits", // 29 CQL gradient output
"aux_dh_s2_nb_buf", // 30 aux next-bar backward dh_s2
"ensemble_d_logits_buf", // 31 ensemble backward output (call site in fused_training.rs)
"ensemble_d_logits_buf", // 31 ensemble (DEFERRED — owner: FusedDqnTraining; instrument only after Phase B saxpy guards verified; see audit per-slot table)
"bn_d_concat_buf", // 32 bottleneck Linear backward dy
"bw_d_h_s2_post_main", // 33 bw_d_h_s2 after main backward chain
"bw_d_h_s2_post_aux", // 34 bw_d_h_s2 after aux head SAXPY

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

@@ -2240,3 +2240,5 @@ P — expanded `nan_flags_buf` 16→24 with 5 new GRN h_s2 sub-stage NaN checks
SP1 Phase A audit (2026-04-29): produced `docs/dqn-backward-nan-audit.md` — read-only γ inventory of every backward-path kernel writing to `bw_d_h_s2` / `grad_buf` / `save_h_s2` accumulators, cross-referenced against `session_2026-04-05_nan_investigation.md`'s residual 8% step-2 NaN in `apply_iqn_trunk_gradient`. Per-kernel sections include: identified unsafe patterns (sqrtf-neg, 1/0, logf-≤0, expf-large, EMA variance, atomicAdd/saxpy NaN-propagation), proposed guard form, ISV bound option (existing slot leverage or new slot or Invariant 1 ε carve-out), F0 risk assessment (low/medium/high used as paper-review gate before smoke), and Phase B flag-slot allocation. Drives SP1 Phase B instrumentation (12 new slots in `nan_flags_buf` 24→48) and Phase C surgical fix decisions. Becomes durable input artifact for SP2 (framework codification) and SP3 (structural-fix scoping).
SP1 Phase B foundation (2026-04-29): expanded `nan_flags_buf` 24→48 (allocation size in `gpu_dqn_trainer.rs`; `read_nan_flags` signature `[i32; 24]` → `[i32; 48]` in both `gpu_dqn_trainer.rs` and `fused_training.rs`; name tables updated in both `training_loop.rs` consumer sites — `halt_nan` block + `halt_grad_collapse` block from commit `d1808df14`). Slot names per `docs/dqn-backward-nan-audit.md` per-slot accessor table (audit supersedes plan placeholder names): slots 24-25 are post-c51_grad `d_value_logits_buf` / `d_adv_logits_buf`; slot 26 is `iqn_trunk_m`; slot 27 is `iqn_d_h_s2_ptr`; slot 28 is `d_branch_logits_buf` (production IQN backward, `iqn_quantile_huber_loss`); slot 29 is `cql_d_value_logits`; slot 30 is `aux_dh_s2_nb_buf`; slot 31 is `ensemble_d_logits_buf` (cross-struct on `FusedDqnTraining`); slot 32 is `bn_d_concat_buf`; slots 33-35 are `bw_d_h_s2` at three different backward call sites; slots 36-47 reserved as headroom for SP2/SP3. No behavioral change in this commit (new slots stay at zero until Task 4 wires the check call sites). Buffer size reviewable by SP2 framework codification — if right-size differs (e.g., 36 with no headroom or 64 for more coverage), SP2 may resize.
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.