From 387335e2b936bcddfc830fceabcd9a419aacfcb4 Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Thu, 30 Apr 2026 00:10:37 +0200 Subject: [PATCH] =?UTF-8?q?fix(dqn):=20SP1=20Phase=20B=20foundation=20?= =?UTF-8?q?=E2=80=94=20stale-doc=20cleanup=20+=20Task=204=20prep?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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) --- .../ml/src/cuda_pipeline/gpu_dqn_trainer.rs | 41 +++++----------- .../src/trainers/dqn/trainer/training_loop.rs | 49 +++++++++++-------- docs/dqn-wire-up-audit.md | 2 + 3 files changed, 44 insertions(+), 48 deletions(-) diff --git a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs index e18f1a9ee..894dc54a8 100644 --- a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs +++ b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs @@ -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; diff --git a/crates/ml/src/trainers/dqn/trainer/training_loop.rs b/crates/ml/src/trainers/dqn/trainer/training_loop.rs index 1e4042fb4..492d72faa 100644 --- a/crates/ml/src/trainers/dqn/trainer/training_loop.rs +++ b/crates/ml/src/trainers/dqn/trainer/training_loop.rs @@ -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 diff --git a/docs/dqn-wire-up-audit.md b/docs/dqn-wire-up-audit.md index 800700159..d6b8b2a32 100644 --- a/docs/dqn-wire-up-audit.md +++ b/docs/dqn-wire-up-audit.md @@ -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.