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

feat(dqn): SP1 Phase B foundation — nan_flags_buf 24→48

Browse Source 
Expands the NaN flag buffer from 24 to 48 slots to make room for
backward-path NaN checks (slots 24-35 per audit doc per-slot table)
plus 12 reserved headroom slots (36-47) for SP2 framework + SP3
observer hooks.

Touches:
- gpu_dqn_trainer.rs: alloc size 24→48 (line ~11178); read_nan_flags
  signature [i32; 24]→[i32; 48] (line ~14982); field docstring updated
  (line ~2658) to reflect 48-slot layout
- fused_training.rs: pub(crate) read_nan_flags signature [i32; 24]→[i32; 48]
- training_loop.rs: BOTH name table sites (halt_nan + halt_grad_collapse
  block from commit d1808df14) updated to 48 entries

Slot names use audit allocation (docs/dqn-backward-nan-audit.md per-slot
table), which supersedes the plan's placeholder names per the audit's
plan-supersession note. Slots 24-35 cover production buffers:
d_value_logits_buf, d_adv_logits_buf, iqn_trunk_m, iqn_d_h_s2_ptr,
d_branch_logits_buf, cql_d_value_logits, aux_dh_s2_nb_buf,
ensemble_d_logits_buf, bn_d_concat_buf, bw_d_h_s2 (×3 call sites).
Slots 36-47 are rsv36-rsv47 (headroom).

No behavioral change — new slots stay at zero until Task 4 wires the
kernel-output NaN checks. Buffer size reviewable by SP2.
This commit is contained in:
jgrusewski
2026-04-30 00:00:43 +02:00
parent 2406ebf2f7
commit 53bc0bc505
4 changed files with 82 additions and 34 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -2655,7 +2655,11 @@ pub struct GpuDqnTrainer {
pub(crate) nan_check_f32_kernel: CudaFunction,
/// Second NaN check kernel handle (formerly f32, now f32 — both check f32 buffers).
pub(crate) nan_check_f32_kernel_b: CudaFunction,
/// NaN flags buffer [8] — one flag per checkpoint. Reset per step, read at epoch boundary.
/// SP1 Phase B: NaN flag buffer expanded from 24 → 48 slots. Slots 0-23
/// are forward-path checks (existing); slots 24-35 are backward-path
/// checks (new — wired by Task 4); slots 36-47 are reserved headroom
/// for SP2 framework + SP3 observers. Reset per step, read at epoch
/// boundary or on training-guard halts.
pub(crate) nan_flags_buf: CudaSlice<i32>,
/// #20 Pruning epoch (epoch at which to compute the mask).
pruning_epoch: usize,
@@ -11171,11 +11175,11 @@ impl GpuDqnTrainer {
aux_knb, aux_kr, aux_h, max_aux_tensor_len,
);
// P (post-bkdx5): expanded 16→24 to add 5 GRN-stage NaN-check slots
// (16-20). The new slots pinpoint which GRN sub-stage produces NaN
// first (elu_post → linear_b_out → glu_sigmoid → ln_rstd → ln_normed)
// when R1's K-removal alone doesn't fix the F1 explosion.
let nan_flags_buf = stream.alloc_zeros::<i32>(24)
// SP1 Phase B: expanded 24 → 48 to make room for backward-path NaN
// checks (slots 24-35) plus 12 reserved headroom slots (36-47) for
// SP2 framework checker hooks + SP3 structural observers. Buffer
// size reviewable by SP2 framework codification.
let nan_flags_buf = stream.alloc_zeros::<i32>(48)
.map_err(|e| MLError::ModelError(format!("nan_flags alloc: {e}")))?;
// v8: PopArt running statistics buffers
@@ -14977,10 +14981,12 @@ impl GpuDqnTrainer {
Ok(())
}
/// Read NaN flags back to CPU (synchronizes stream). Returns [24] flags.
/// Read NaN flags back to CPU (synchronizes stream). Returns [48] flags.
/// Index → buffer mapping is documented in `run_nan_checks_post_forward`.
pub fn read_nan_flags(&self) -> Result<[i32; 24], MLError> {
let mut host = [0_i32; 24];
/// SP1 Phase B: slots 24-35 reserved for backward-path checks (wired by
/// Task 4); slots 36-47 reserved for SP2/SP3 headroom.
pub fn read_nan_flags(&self) -> Result<[i32; 48], MLError> {
let mut host = [0_i32; 48];
self.stream.memcpy_dtoh(&self.nan_flags_buf, &mut host)
.map_err(|e| MLError::ModelError(format!("nan_flags read: {e}")))?;
Ok(host)

4
crates/ml/src/trainers/dqn/fused_training.rs
View File

@@ -3621,7 +3621,9 @@ impl FusedTrainingCtx {
/// Read NaN detection flags (synchronizes stream). Used by training guard on NaN halt.
/// Index → buffer mapping is documented in
/// `GpuDqnTrainer::run_nan_checks_post_forward`.
pub(crate) fn read_nan_flags(&self) -> Result<[i32; 24], crate::MLError> {
/// SP1 Phase B: returns 48 flags (was 24). Slots 24-35 cover backward-path
/// checks; slots 36-47 reserved for SP2/SP3 headroom.
pub(crate) fn read_nan_flags(&self) -> Result<[i32; 48], crate::MLError> {
self.trainer.read_nan_flags()
}

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

@@ -1981,8 +1981,12 @@ impl DQNTrainer {
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. 24-slot layout
// (P, post-bkdx5: added GRN-stage slots 16-20).
// 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.
let names = [
"states_buf", // 0 input states (post-upload)
"on_v_logits", // 1 cuBLAS fwd: value stream
@@ -2004,6 +2008,21 @@ impl DQNTrainer {
"grn_h_s2_ln_rstd", // 19 GRN h_s2 LN reciprocal-stddev
"grn_h_s2_ln_normed", // 20 GRN h_s2 LN normed (pre-affine)
"rsv21", "rsv22", "rsv23", // reserved (h_s1 stages / future)
"d_value_logits_buf", // 24 post-c51_grad value gradient
"d_adv_logits_buf", // 25 post-c51_grad branch advantage gradient
"iqn_trunk_m", // 26 apply_iqn_trunk_gradient cuBLAS bwd output
"iqn_d_h_s2_ptr", // 27 IQN backward dh_s2 input (= gpu_iqn_head::d_h_s2_buf)
"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)
"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
"bw_d_h_s2_post_iqn", // 35 bw_d_h_s2 after IQN DtoD
"rsv36", "rsv37", "rsv38", "rsv39", // 36-39 SP2/SP3 headroom
"rsv40", "rsv41", "rsv42", "rsv43", // 40-43 SP2/SP3 headroom
"rsv44", "rsv45", "rsv46", "rsv47", // 44-47 SP2/SP3 headroom
];
let flagged: Vec<String> = flags.iter().enumerate()
.filter(|(_, &f)| f != 0)
@@ -2042,29 +2061,48 @@ impl DQNTrainer {
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. 24-slot
// layout (P, post-bkdx5: GRN-stage 16-20).
// 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.
let names = [
"states_buf", // 0
"on_v_logits", // 1
"on_b_logits", // 2
"mse_loss_scalar", // 3
"params_buf_pre_fwd", // 4
"params_ptr_pre_fwd", // 5
"grad_buf", // 6
"save_current_lp", // 7
"save_projected", // 8
"moe_gate_softmax", // 9
"aux_nb_loss_scalar", // 10
"aux_rg_loss_scalar", // 11
"save_h_s2", // 12
"rsv13", "rsv14", "rsv15",
"grn_h_s2_elu_post", // 16
"grn_h_s2_linear_b_out", // 17
"grn_h_s2_glu_sigmoid", // 18
"grn_h_s2_ln_rstd", // 19
"grn_h_s2_ln_normed", // 20
"rsv21", "rsv22", "rsv23",
"states_buf", // 0 input states (post-upload)
"on_v_logits", // 1 cuBLAS fwd: value stream
"on_b_logits", // 2 cuBLAS fwd: branch advantage
"mse_loss_scalar", // 3 MSE loss [1]
"params_buf_pre_fwd", // 4 f32 params (pre-forward)
"params_ptr_pre_fwd", // 5 same buffer (redundant pre-fwd)
"grad_buf", // 6 cuBLAS bwd output
"save_current_lp", // 7 C51 online softmax probs
"save_projected", // 8 C51 target distribution
"moe_gate_softmax", // 9 MoE gate over 8 experts
"aux_nb_loss_scalar", // 10 aux next-bar MSE [1]
"aux_rg_loss_scalar", // 11 aux regime CE [1]
"save_h_s2", // 12 trunk output (input to FCs)
"rsv13", "rsv14", "rsv15", // 13-15 reserved
"grn_h_s2_elu_post", // 16 GRN h_s2 post-ELU activation
"grn_h_s2_linear_b_out", // 17 GRN h_s2 Linear_b output
"grn_h_s2_glu_sigmoid", // 18 GRN h_s2 GLU gate sigmoid
"grn_h_s2_ln_rstd", // 19 GRN h_s2 LN reciprocal-stddev
"grn_h_s2_ln_normed", // 20 GRN h_s2 LN normed (pre-affine)
"rsv21", "rsv22", "rsv23", // reserved (h_s1 stages / future)
"d_value_logits_buf", // 24 post-c51_grad value gradient
"d_adv_logits_buf", // 25 post-c51_grad branch advantage gradient
"iqn_trunk_m", // 26 apply_iqn_trunk_gradient cuBLAS bwd output
"iqn_d_h_s2_ptr", // 27 IQN backward dh_s2 input (= gpu_iqn_head::d_h_s2_buf)
"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)
"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
"bw_d_h_s2_post_iqn", // 35 bw_d_h_s2 after IQN DtoD
"rsv36", "rsv37", "rsv38", "rsv39", // 36-39 SP2/SP3 headroom
"rsv40", "rsv41", "rsv42", "rsv43", // 40-43 SP2/SP3 headroom
"rsv44", "rsv45", "rsv46", "rsv47", // 44-47 SP2/SP3 headroom
];
let flagged: Vec<String> = flags.iter().enumerate()
.filter(|(_, &f)| f != 0)

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

@@ -2238,3 +2238,5 @@ R1 — eliminated K's Adam shrink-and-perturb at fold boundary (m_buf and v_buf
P — expanded `nan_flags_buf` 16→24 with 5 new GRN h_s2 sub-stage NaN checks (`grn_h_s2_elu_post`=16, `grn_h_s2_linear_b_out`=17, `grn_h_s2_glu_sigmoid`=18, `grn_h_s2_ln_rstd`=19, `grn_h_s2_ln_normed`=20) for finer-grained source identification if R1 alone doesn't fix F1. Slot 19 (`ln_rstd`) is the divide-by-zero diagnostic (variance→0 ⇒ rstd→∞ ⇒ NaN downstream). New accessors: `GrnBlock::{elu_post_ptr, linear_b_out_ptr, sigmoid_b_ptr, ln_rstd_ptr, ln_normed_ptr}` (additive, reads-only) and `CublasGemmSet::{grn_h_s2_online, grn_h_s2_linear_b_ptr}`. h_s2 only — pragmatic instrumentation scope: that's the GRN block where save_h_s2 went NaN at F1 step 1745, and instrumenting only h_s2 covers the failing stage without doubling per-step kernel-launch cost. h_s1 stages can be added by name table slots 21-23 if the F1 explosion turns out to originate above h_s2. The K-removal + GRN-stage checks combined: if the F1 explosion was due to K's tiny-v_hat pathology, R1 alone fixes it. If a different mechanism, the new flags pinpoint which GRN sub-stage produces NaN first. Per `feedback_no_partial_refactor.md`: R1 and P land together because both touch the fold-boundary contract (K-removal changes the post-reset Adam state; GRN-stage NaN checks observe the consequences of that state on the trunk encoder's first forward). `read_nan_flags()` signature update [16]→[24] propagates through `GpuDqnTrainer`, `FusedTrainingCtx`, and both name-table sites in `training_loop.rs` (halt_nan + halt_grad_collapse).
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.