diff --git a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs index c0134548e..01481f349 100644 --- a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs +++ b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs @@ -2125,6 +2125,61 @@ impl GpuDqnTrainer { Ok(norm_f32[0].sqrt()) } + /// Per-step buffer diagnostics: check every critical buffer after forward+backward. + /// Logs L2 norms of d_value_logits, d_adv_logits, grad_buf, and h_s2 activations. + /// Costs ~4 stream syncs. Gate on step count in caller. + pub fn run_buffer_diagnostics(&mut self, step: usize) { + let b = self.config.batch_size; + let na = self.config.num_atoms; + let tba = (self.config.branch_0_size + self.config.branch_1_size + + self.config.branch_2_size + self.config.branch_3_size) * na; + + let d_val_norm = self.debug_buffer_norm_f32(self.d_value_logits_buf.raw_ptr(), b * na).unwrap_or(f32::NAN); + let d_adv_norm = self.debug_buffer_norm_f32(self.d_adv_logits_buf.raw_ptr(), b * tba).unwrap_or(f32::NAN); + let grad_norm = self.debug_buffer_norm_f32(self.grad_buf.raw_ptr(), self.total_params).unwrap_or(f32::NAN); + // h_s2 is bf16 — check via f32 conversion of first few elements would need a different kernel. + // Skip for now — the f32 buffers tell us enough. + + tracing::warn!( + step, + d_val_norm, + d_adv_norm, + grad_norm, + c51_alpha = self.c51_alpha, + "BUFFER_DIAG: per-step buffer norms after forward+backward" + ); + } + + /// Compute L2 norm of an arbitrary f32 GPU buffer (diagnostic use only). + /// Reuses the standalone grad_norm two-phase reduction. Costs 1 stream sync. + pub fn debug_buffer_norm_f32(&mut self, ptr: u64, n_elems: usize) -> Result { + let _evt_guard = EventTrackingGuard::new(self.stream.context()); + let n_i32 = n_elems as i32; + let blocks = ((n_elems + 255) / 256) as u32; + let partials_ptr = self.grad_norm_partials.raw_ptr(); + let f32_ptr = self.grad_norm_f32_buf.raw_ptr(); + let bf16_ptr = self.aux_norm_bf16_scratch.raw_ptr(); + let nb = blocks as i32; + unsafe { + self.stream.launch_builder(&self.grad_norm_standalone) + .arg(&ptr).arg(&partials_ptr).arg(&n_i32) + .launch(LaunchConfig { grid_dim: (blocks, 1, 1), block_dim: (256, 1, 1), shared_mem_bytes: 0 }) + .map_err(|e| MLError::ModelError(format!("debug_norm phase1: {e}")))?; + self.stream.launch_builder(&self.grad_norm_finalize_kernel) + .arg(&partials_ptr).arg(&f32_ptr).arg(&bf16_ptr).arg(&nb) + .launch(LaunchConfig { grid_dim: (1, 1, 1), block_dim: (256, 1, 1), shared_mem_bytes: 0 }) + .map_err(|e| MLError::ModelError(format!("debug_norm phase2: {e}")))?; + cudarc::driver::sys::cuStreamSynchronize(self.stream.cu_stream()); + } + let mut norm_sq = [0.0_f32; 1]; + unsafe { + cudarc::driver::sys::cuMemcpyDtoH_v2( + norm_sq.as_mut_ptr().cast(), self.grad_norm_f32_buf.raw_ptr(), 4, + ); + } + Ok(norm_sq[0].sqrt()) + } + /// Apply GPU multi-head feature attention to `save_h_s2` (post-graph). /// /// Runs 4-head self-attention over the trunk output `h_s2 [B, SHARED_H2]` diff --git a/crates/ml/src/trainers/dqn/fused_training.rs b/crates/ml/src/trainers/dqn/fused_training.rs index b9cd644ce..af75d3368 100644 --- a/crates/ml/src/trainers/dqn/fused_training.rs +++ b/crates/ml/src/trainers/dqn/fused_training.rs @@ -834,6 +834,11 @@ impl FusedTrainingCtx { } self.pending_vaccine_batch = None; + // ── BUFFER_DIAG: check where gradients die (first 3 steps each epoch) ── + if self.steps_since_varmap_sync < 3 { + self.trainer.run_buffer_diagnostics(self.steps_since_varmap_sync); + } + // ── Step 5: Pruning + Adam ─────────────────────────────────────── self.trainer.apply_pruning_mask() .map_err(|e| anyhow::anyhow!("Pruning mask apply: {e}"))?;