feat(dqn): SP1 Phase B accessors — backward-buffer NaN check pointers
Adds 5 new pub(crate) accessor methods exposing backward-path buffer
device pointers for Task 4's per-step NaN checks (slots 24, 25, 28,
29, 30 per audit per-slot table at docs/dqn-backward-nan-audit.md
:530-548):
GpuDqnTrainer (4 new):
- d_value_logits_buf_ptr (slot 24 — post-c51_grad value gradient)
- d_adv_logits_buf_ptr (slot 25 — post-c51_grad branch advantage)
- cql_d_value_logits_ptr (slot 29 — CQL gradient output)
- aux_dh_s2_nb_buf_ptr (slot 30 — aux next-bar backward dh_s2)
GpuIqnHead (1 new):
- d_branch_logits_buf_ptr (slot 28 — production IQN backward output,
iqn_quantile_huber_loss)
Slot 27 (iqn_d_h_s2_buf) reuses existing GpuIqnHead::d_h_s2_raw_ptr()
at gpu_iqn_head.rs:1660 — no new method per feedback_no_legacy_aliases:
the existing accessor is already public and sufficient; renaming +
chasing the single call site adds churn without value.
Slots 26, 32, 33-35 reuse pre-existing handles:
- 26: self.ptrs.iqn_trunk_m
- 32: self.bn_d_concat_buf() (existing, returns &CudaSlice<f32>)
- 33-35: self.ptrs.bw_d_h_s2 (3 different Task 4 call sites)
Slot 31 (ensemble_d_logits_buf) deferred per Task 2 commit 387335e2b
(cross-struct on FusedDqnTraining).
DEVIATION FROM PLAN: the plan called for "delegate accessors on
GpuDqnTrainer for slots 27/28" — structurally invalid because
GpuDqnTrainer does NOT own GpuIqnHead. The IQN head is owned by
FusedTrainingCtx (fused_training.rs:289) alongside the trainer at
line 234. The audit's per-slot accessor table (lines 535-536) is
correct: accessors land on GpuIqnHead. Task 4's
run_nan_checks_post_backward will receive 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, ...)).
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 (one buffer, three Task 4 invocations).
Pattern follows commit e9096c7be's GRN-block accessors (concise
pub(crate) fn name_ptr(&self) -> u64 with doc-comment referencing
slot number + audit doc + buffer semantics). Additive — no behavioral
change; new accessors consumed by Task 4's NaN check call sites.
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
This commit is contained in:
@@ -6764,6 +6764,48 @@ impl GpuDqnTrainer {
|
||||
self.config.bottleneck_dim + (ml_core::state_layout::STATE_DIM.saturating_sub(self.config.market_dim))
|
||||
}
|
||||
|
||||
/// SP1 Phase B (slot 24): post-`c51_grad_kernel` value-stream gradient buffer
|
||||
/// `[B, NA]` (f32, atomicAdd target). Wired into `run_nan_checks_post_backward`
|
||||
/// (Task 4) as a symmetric orchestrator-side check; the launch-site local at
|
||||
/// `gpu_dqn_trainer.rs:17707` reads this same field. Per-slot semantics in
|
||||
/// `docs/dqn-backward-nan-audit.md` per-slot accessor table (slot 24 row).
|
||||
/// Distinct from the `d_value_logits` *staging* buffer at line 3131 — that one
|
||||
/// is unrelated to `c51_grad_kernel`.
|
||||
pub(crate) fn d_value_logits_buf_ptr(&self) -> u64 {
|
||||
self.d_value_logits_buf.raw_ptr()
|
||||
}
|
||||
|
||||
/// SP1 Phase B (slot 25): post-`c51_grad_kernel` branch-advantage gradient
|
||||
/// buffer `[B, (b0+b1+b2+b3) × NA]` (f32, atomicAdd target). Companion to
|
||||
/// slot 24 — fires independently iff a per-branch path corrupted (e.g. a
|
||||
/// magnitude-only `inv_a_std` divide). Per-slot semantics in
|
||||
/// `docs/dqn-backward-nan-audit.md`. Distinct from the `d_adv_logits` staging
|
||||
/// buffer at line 3133 (unrelated to `c51_grad_kernel`).
|
||||
pub(crate) fn d_adv_logits_buf_ptr(&self) -> u64 {
|
||||
self.d_adv_logits_buf.raw_ptr()
|
||||
}
|
||||
|
||||
/// SP1 Phase B (slot 29): CQL gradient buffer `[B, NA]` (f32) — written by
|
||||
/// `cql_logit_grad_kernel`. Wired into `run_nan_checks_post_backward` only if
|
||||
/// the slot is un-deferred from the audit's priority list; until then this
|
||||
/// accessor exists for symmetric coverage but the call site stays optional.
|
||||
/// Per-slot semantics in `docs/dqn-backward-nan-audit.md` (slot 29 row).
|
||||
pub(crate) fn cql_d_value_logits_ptr(&self) -> u64 {
|
||||
self.cql_d_value_logits.raw_ptr()
|
||||
}
|
||||
|
||||
/// SP1 Phase B (slot 30): aux next-bar backward dh_s2 contribution `[B, SH2]`
|
||||
/// (f32) — produced by `aux_next_bar_backward`, SAXPY-accumulated into
|
||||
/// `bw_d_h_s2` before `encoder_backward_chain` runs. Wired into
|
||||
/// `run_nan_checks_post_backward` (Task 4). Companion buffer
|
||||
/// `aux_dh_s2_rg_buf` (regime variant) is intentionally NOT instrumented —
|
||||
/// the audit limits slot 30 to the next-bar variant; both share the same
|
||||
/// caller and shape so a single check provides sufficient diagnostic
|
||||
/// resolution. Per-slot semantics in `docs/dqn-backward-nan-audit.md`.
|
||||
pub(crate) fn aux_dh_s2_nb_buf_ptr(&self) -> u64 {
|
||||
self.aux_dh_s2_nb_buf.raw_ptr()
|
||||
}
|
||||
|
||||
/// Reference to the next_states buffer on GPU.
|
||||
///
|
||||
/// Shape: `[B, STATE_DIM]` — contains the batch's next states after `train_step()`.
|
||||
|
||||
@@ -1661,6 +1661,23 @@ impl GpuIqnHead {
|
||||
self.d_h_s2_buf.raw_ptr()
|
||||
}
|
||||
|
||||
/// SP1 Phase B (slot 28): production IQN backward output buffer
|
||||
/// `d_branch_logits_buf [TBA, B*Q]` (f32) — written by
|
||||
/// `iqn_quantile_huber_loss` (`iqn_dual_head_kernel.cu:1346-1413`,
|
||||
/// `d_q_online[idx] = qw * d_huber / Q` with no inline `isfinite` guard).
|
||||
/// Wired into the NaN-check infrastructure by `run_nan_checks_post_backward`
|
||||
/// (Task 4) via the `FusedTrainingCtx` call site (the IQN head is owned by
|
||||
/// `FusedTrainingCtx`, not `GpuDqnTrainer`, so the pointer is forwarded as a
|
||||
/// `u64` argument to the trainer-side check method — same pattern as
|
||||
/// `apply_iqn_trunk_gradient(iqn_d_h_s2_ptr: u64, ...)` at
|
||||
/// `gpu_dqn_trainer.rs:6843`). Per-slot semantics in
|
||||
/// `docs/dqn-backward-nan-audit.md` per-slot accessor table (slot 28 row).
|
||||
/// NOTE: distinct from the kernel-internal `save_dL_dq` name used inside
|
||||
/// the dormant `iqn_backward_per_sample` path (not loaded in production).
|
||||
pub(crate) fn d_branch_logits_buf_ptr(&self) -> u64 {
|
||||
self.d_branch_logits_buf.raw_ptr()
|
||||
}
|
||||
|
||||
/// Read IQN loss from GPU (call sparingly — synchronizes the stream).
|
||||
/// Use for epoch-end logging, NOT per-step monitoring.
|
||||
pub fn read_loss(&self) -> Result<f32, MLError> {
|
||||
|
||||
@@ -2242,3 +2242,5 @@ SP1 Phase A audit (2026-04-29): produced `docs/dqn-backward-nan-audit.md` — re
|
||||
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.
|
||||
|
||||
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.
|
||||
|
||||
Reference in New Issue
Block a user