diff --git a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs index 5b2ef82d3..e18f1a9ee 100644 --- a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs +++ b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs @@ -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, /// #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::(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::(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) diff --git a/crates/ml/src/trainers/dqn/fused_training.rs b/crates/ml/src/trainers/dqn/fused_training.rs index 1bf77f23f..b2836ffdd 100644 --- a/crates/ml/src/trainers/dqn/fused_training.rs +++ b/crates/ml/src/trainers/dqn/fused_training.rs @@ -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() } diff --git a/crates/ml/src/trainers/dqn/trainer/training_loop.rs b/crates/ml/src/trainers/dqn/trainer/training_loop.rs index ed627e998..1e4042fb4 100644 --- a/crates/ml/src/trainers/dqn/trainer/training_loop.rs +++ b/crates/ml/src/trainers/dqn/trainer/training_loop.rs @@ -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 = 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 = flags.iter().enumerate() .filter(|(_, &f)| f != 0) diff --git a/docs/dqn-wire-up-audit.md b/docs/dqn-wire-up-audit.md index f306013fd..800700159 100644 --- a/docs/dqn-wire-up-audit.md +++ b/docs/dqn-wire-up-audit.md @@ -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.