From 30f31563691cc690f412f5def032ab101cafcea4 Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Fri, 10 Apr 2026 08:34:01 +0200 Subject: [PATCH] =?UTF-8?q?diag:=20RECAPTURE=5FDIAG=20=E2=80=94=20per-buff?= =?UTF-8?q?er=20norms=20after=20graph=20recapture?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Fires for 3 steps after any CUDA graph recapture. Reports: - d_val/adv_norm: blended d_logits (C51 × alpha + MSE × (1-alpha)) - mse_val/adv_norm: MSE-only scratch buffers - grad_norm: backward output grad_buf Purpose: diagnose C51 grad_norm=0 on H100 at batch=16384 (epoch 2+). Co-Authored-By: Claude Opus 4.6 (1M context) --- .../ml/src/cuda_pipeline/gpu_dqn_trainer.rs | 34 +++++++++++++++++++ crates/ml/src/trainers/dqn/fused_training.rs | 4 ++- 2 files changed, 37 insertions(+), 1 deletion(-) diff --git a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs index 7b8659b26..32d98512b 100644 --- a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs +++ b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs @@ -595,6 +595,7 @@ pub struct GpuDqnTrainer { target_params_bf16: CudaSlice, // [TOTAL_PARAMS + pad] bf16 shadow for GemmEx pub(crate) grad_norm_buf: CudaSlice, // [1] bf16 L2 norm (written by finalize) grad_norm_f32_buf: CudaSlice, // [1] float sum-of-squares output + diag_recapture_remaining: u32, // diagnostic steps remaining after graph recapture grad_norm_partials: CudaSlice, // [grad_norm_blocks] per-block partial sums grad_norm_blocks: usize, // number of blocks for grad_norm kernel cql_grad_scratch: CudaSlice, // [TOTAL_PARAMS] f32 CQL gradient isolation buffer @@ -2113,6 +2114,37 @@ impl GpuDqnTrainer { ); } + /// Extended diagnostics after graph recapture: checks MSE scratch, C51 d_logits, + /// blended d_logits, and grad_buf. Fires for first 3 steps after any graph recapture. + pub fn run_recapture_diagnostics(&mut self) { + if self.diag_recapture_remaining == 0 { return; } + self.diag_recapture_remaining -= 1; + + 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; + + // After graph replay: d_value/adv_logits contain BLENDED gradients, + // MSE scratch contains MSE-only gradients, grad_buf contains backward output. + 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 mse_val_norm = self.debug_buffer_norm_f32(self.d_value_logits_mse.raw_ptr(), b * na).unwrap_or(f32::NAN); + let mse_adv_norm = self.debug_buffer_norm_f32(self.d_adv_logits_mse.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); + + tracing::warn!( + remaining = self.diag_recapture_remaining, + c51_alpha = self.c51_alpha, + d_val_norm, + d_adv_norm, + mse_val_norm, + mse_adv_norm, + grad_norm, + "RECAPTURE_DIAG: buffer norms after graph recapture replay" + ); + } + /// 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 { @@ -2997,6 +3029,7 @@ impl GpuDqnTrainer { target_params_bf16, grad_norm_buf, grad_norm_f32_buf, + diag_recapture_remaining: 3, grad_norm_partials, grad_norm_blocks, cql_grad_scratch, @@ -4511,6 +4544,7 @@ impl GpuDqnTrainer { self.graph_forward_ddqn = Some(SendSyncGraph(graph_ddqn)); self.graph_adam = Some(SendSyncGraph(graph_adam)); self.last_captured_loss_mode = Some(self.loss_mode); + self.diag_recapture_remaining = 3; // fire diagnostics for first 3 steps after recapture Ok(()) } diff --git a/crates/ml/src/trainers/dqn/fused_training.rs b/crates/ml/src/trainers/dqn/fused_training.rs index 9bbe35da7..06a3d30d0 100644 --- a/crates/ml/src/trainers/dqn/fused_training.rs +++ b/crates/ml/src/trainers/dqn/fused_training.rs @@ -834,10 +834,12 @@ impl FusedTrainingCtx { } self.pending_vaccine_batch = None; - // ── BUFFER_DIAG: check where gradients die (first 3 steps each epoch) ── + // ── BUFFER_DIAG: check where gradients die (first 3 steps each fold) ── if self.steps_since_varmap_sync < 3 { self.trainer.run_buffer_diagnostics(self.steps_since_varmap_sync); } + // ── RECAPTURE_DIAG: check C51 gradient path after graph recapture ── + self.trainer.run_recapture_diagnostics(); // ── Step 5: Pruning + Adam ─────────────────────────────────────── self.trainer.apply_pruning_mask()