diag: per-step BUFFER_DIAG — d_value_logits, d_adv_logits, grad_buf norms

Logs L2 norms of critical buffers after forward+backward+aux, before Adam.
Fires on first 3 steps of each epoch only. Identifies exactly which buffer
goes zero on H100 — is it the loss gradient or the cuBLAS backward output?

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

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<f32, MLError> {
+        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]`

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}"))?;