From bf4e92bdb729acec0cc17b50129caad6fb3b9ef0 Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Tue, 14 Apr 2026 08:10:46 +0200 Subject: [PATCH] cleanup: remove BUFFER_DIAG, RECAPTURE_DIAG, GRAD_DIAG diagnostic infrastructure Root causes are fixed (cuBLASLt heuristic, entropy inv_batch, overflow, IQL OOB). Removes: run_buffer_diagnostics, run_recapture_diagnostics, debug_buffer_norm_f32, diag_recapture_remaining field, gradient_stage_diagnostics field + FOXHUNT_GRAD_DIAG env var. Eliminates ~4 stream syncs per step for first 3 steps each fold. Co-Authored-By: Claude Opus 4.6 (1M context) --- .../ml/src/cuda_pipeline/gpu_dqn_trainer.rs | 86 ------------------- crates/ml/src/trainers/dqn/fused_training.rs | 38 +------- 2 files changed, 1 insertion(+), 123 deletions(-) diff --git a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs index 6f12f038a..2695ba0e7 100644 --- a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs +++ b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs @@ -613,7 +613,6 @@ pub struct GpuDqnTrainer { m_buf: CudaSlice, // [TOTAL_PARAMS] Adam first moment (f32 for precision) v_buf: CudaSlice, // [TOTAL_PARAMS] Adam second moment (f32 for precision) pub(crate) grad_norm_buf: CudaSlice, // [1] f32 L2 norm (written by finalize) - 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 @@ -2124,89 +2123,6 @@ impl GpuDqnTrainer { Ok(norm_f32[0]) } - /// 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.ptrs.grad_buf, self.total_params).unwrap_or(f32::NAN); - - 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" - ); - } - - /// 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.ptrs.grad_buf, 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 { - 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 buf_ptr = self.aux_norm_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(&buf_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_val = [0.0_f32; 1]; - unsafe { - cudarc::driver::sys::cuMemcpyDtoH_v2( - norm_val.as_mut_ptr().cast(), self.aux_norm_scratch.raw_ptr(), 4, - ); - } - Ok(norm_val[0]) - } - /// 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]` @@ -3074,7 +2990,6 @@ impl GpuDqnTrainer { m_buf, v_buf, grad_norm_buf, - diag_recapture_remaining: 3, grad_norm_partials, grad_norm_blocks, cql_grad_scratch, @@ -4635,7 +4550,6 @@ impl GpuDqnTrainer { self.graph_forward = Some(SendSyncGraph(graph_fwd)); self.graph_forward_ddqn = Some(SendSyncGraph(graph_ddqn)); self.graph_adam = Some(SendSyncGraph(graph_adam)); - 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 7e96854a0..5dd10ce5c 100644 --- a/crates/ml/src/trainers/dqn/fused_training.rs +++ b/crates/ml/src/trainers/dqn/fused_training.rs @@ -286,10 +286,6 @@ pub(crate) struct FusedTrainingCtx { pub(crate) hindsight_fraction: f64, /// v8: Curriculum learning enabled (false by default). pub(crate) curriculum_enabled: bool, - /// When true, disables mega-graph/graph_aux capture and logs per-stage - /// gradient norms (raw → post-primary-clip → post-IQN → post-CQL). - /// Costs ~4 stream syncs per step. For debugging H100 gradient collapse. - pub(crate) gradient_stage_diagnostics: bool, } impl Drop for FusedTrainingCtx { @@ -743,7 +739,6 @@ impl FusedTrainingCtx { prev_popart_var: 0.0, hindsight_fraction: hyperparams.hindsight_fraction, curriculum_enabled: hyperparams.curriculum_enabled, - gradient_stage_diagnostics: std::env::var("FOXHUNT_GRAD_DIAG").is_ok(), }) } @@ -1018,13 +1013,6 @@ impl FusedTrainingCtx { .map_err(|e| anyhow::anyhow!("pre-Adam grad_buf clip: {e}"))?; } - // ── 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() .map_err(|e| anyhow::anyhow!("Pruning mask apply: {e}"))?; @@ -1073,12 +1061,11 @@ impl FusedTrainingCtx { // Capture mega-graph at step 2 (all sub-trainers initialized). // Mega-graph fuses spectral + forward + backward + aux → 1 launch. - // Skip when gradient_stage_diagnostics is on — need ungraphed path for per-stage reads. // Capture graph_mega on step 0 — ONE capture session for determinism. // Individual graphs at steps 0-1 caused cross-process non-determinism: // each separate capture runs cublasLtMatmul ungraphed, contaminating // the shared handle's internal state for subsequent captures. - if self.graph_mega.is_none() && !self.gradient_stage_diagnostics { + if self.graph_mega.is_none() { if let Err(e) = self.capture_graph_mega(agent, gpu_batch) { tracing::warn!( "graph_mega capture failed (non-fatal, falling back to individual graphs): {e}" @@ -1306,17 +1293,6 @@ impl FusedTrainingCtx { iqn.compute_iqr() .map_err(|e| anyhow::anyhow!("IQN IQR: {e}"))?; - // ── DIAG: post-IQN trunk SAXPY gradient norm ────────── - if self.gradient_stage_diagnostics { - let post_iqn_norm = self.trainer.read_grad_norm_standalone_sync() - .unwrap_or(f32::NAN); - tracing::warn!( - batch_size = self.batch_size, - post_iqn_norm, - stage = "post_iqn_saxpy", - "GRAD_DIAG: gradient norm AFTER IQN trunk SAXPY" - ); - } } Err(e) => { tracing::warn!("IQN step failed (non-fatal): {e}"); @@ -1346,18 +1322,6 @@ impl FusedTrainingCtx { } } - // ── DIAG: final gradient norm (after all component injections) ──── - if self.gradient_stage_diagnostics { - let final_norm = self.trainer.read_grad_norm_standalone_sync() - .unwrap_or(f32::NAN); - tracing::warn!( - batch_size = self.batch_size, - final_norm, - stage = "final_pre_adam", - "GRAD_DIAG: final gradient norm before Adam (all components injected)" - ); - } - // Regime-adaptive PER scaling. self.trainer.regime_scale_td_errors() .map_err(|e| anyhow::anyhow!("Regime PER scaling: {e}"))?;