From 3e726648bc1aebeebb41838fbcf2f5a1003a4b92 Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Mon, 6 Apr 2026 14:44:47 +0200 Subject: [PATCH] cleanup: remove all H100 diagnostic eprintln, delete debug cuStreamSynchronize Remove diagnostic eprintln calls from the H100 hang investigation: - H100_STEP, H100_LOOP, H100_HANG4 prefixes in fused_training.rs - adam_readback prefix + ADAM_CTR static counter in gpu_dqn_trainer.rs - H100_LOOP + FUSED STEP ERROR in training_loop.rs - Debug cuStreamSynchronize block after PER update in fused_training.rs Integration tests and smoke tests already match current signatures (ext_stream: Option<&Arc> with None). seg_tree_kernel.cu already deleted in prior commit. Co-Authored-By: Claude Opus 4.6 (1M context) --- crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs | 11 ----------- crates/ml/src/trainers/dqn/fused_training.rs | 18 ------------------ .../src/trainers/dqn/trainer/training_loop.rs | 4 ---- 3 files changed, 33 deletions(-) diff --git a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs index a93d57c87..c8d854c3c 100644 --- a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs +++ b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs @@ -2934,10 +2934,6 @@ impl GpuDqnTrainer { } pub fn replay_adam_and_readback(&mut self) -> Result { - static ADAM_CTR: std::sync::atomic::AtomicUsize = std::sync::atomic::AtomicUsize::new(0); - let adam_step = ADAM_CTR.fetch_add(1, std::sync::atomic::Ordering::Relaxed); - if adam_step < 3 { eprintln!("adam_readback: step {adam_step} — enter"); } - // 1. Collect previous step's scalars (if pending) // Non-blocking: read from pinned buffer without synchronizing. // These values are monitoring-only (the training guard reads GPU buffers @@ -2962,11 +2958,8 @@ impl GpuDqnTrainer { }; // 2. Launch current step: grad_norm + adam - if adam_step < 3 { eprintln!("adam_readback: step {adam_step} — compute_grad_norm"); } self.compute_grad_norm_outside_graph()?; - if adam_step < 3 { eprintln!("adam_readback: step {adam_step} — replay_adam"); } self.replay_adam()?; - if adam_step < 3 { eprintln!("adam_readback: step {adam_step} — async readback start"); } // 3. Async DtoH into pinned host buffer (truly async — no CPU blocking) unsafe { @@ -2993,14 +2986,10 @@ impl GpuDqnTrainer { ); } - if adam_step < 3 { eprintln!("adam_readback: step {adam_step} — async readback done"); } - // No event recording needed — flush_readback reads pinned buffer // directly without synchronizing (async copy completes in <1 µs). self.readback_pending = true; - if adam_step < 3 { eprintln!("adam_readback: step {adam_step} — returning"); } - // Return previous step's values self.scalars_readback_host = [prev_scalars.total_loss, prev_scalars.grad_norm]; Ok(prev_scalars) diff --git a/crates/ml/src/trainers/dqn/fused_training.rs b/crates/ml/src/trainers/dqn/fused_training.rs index 1850959c8..071ff8291 100644 --- a/crates/ml/src/trainers/dqn/fused_training.rs +++ b/crates/ml/src/trainers/dqn/fused_training.rs @@ -679,8 +679,6 @@ impl FusedTrainingCtx { .ok_or_else(|| anyhow::anyhow!("Fused training requires gpu_batch (GPU PER)"))?; // ── Step 1: Spectral normalization BEFORE forward pass ───────── - let step = self.steps_since_varmap_sync; - if step < 3 { eprintln!("H100_STEP: step {step} — spectral norm (graphed={})", self.graph_spectral.is_some()); } if let Some(ref graph) = self.graph_spectral { // Steps 2+: zero-overhead graph replay (batched kernel + bf16→f32 sync). graph.launch(self.stream.cu_stream())?; @@ -696,7 +694,6 @@ impl FusedTrainingCtx { } // ── Step 2: Upload batch + replay graph_forward ────────────────── - if step < 3 { eprintln!("H100_STEP: step {step} — train_step_gpu (forward)"); } let _fused_placeholder = self.trainer.train_step_gpu( gpu_batch, &self.online_dueling, &self.online_branching, @@ -730,7 +727,6 @@ impl FusedTrainingCtx { } // ── Step 3: Auxiliary ops (graph_aux captured or replayed) ──────── - if step < 3 { eprintln!("H100_STEP: step {step} — aux_ops (graphed={})", self.graph_aux.is_some()); } if self.graph_aux.is_some() { self.pre_replay_state_update(agent); let graph_aux = self.graph_aux.as_ref().unwrap(); @@ -772,11 +768,9 @@ impl FusedTrainingCtx { self.trainer.apply_pruning_mask() .map_err(|e| anyhow::anyhow!("Pruning mask apply: {e}"))?; - if step < 3 { eprintln!("H100_STEP: step {step} — replay_adam"); } let fused_result = self.trainer.replay_adam_and_readback() .map_err(|e| anyhow::anyhow!("graph_adam replay: {e}"))?; - if step < 3 { eprintln!("H100_STEP: step {step} — PER priority update"); } // ── Step 6: PER priority update via seg_tree_update ───────────── // Uses the replay buffer's segment tree kernel: writes priorities, // updates tree leaves, propagates sums to root. @@ -785,33 +779,21 @@ impl FusedTrainingCtx { self.batch_size, &self.stream, ).map_err(|e| anyhow::anyhow!("PER seg_tree_update: {e}"))?; - // DEBUG: force sync to surface any CUDA errors from PER kernels - if step < 3 { - eprintln!("H100_HANG4: step {step} — PER returned, syncing stream to check errors"); - unsafe { cudarc::driver::sys::cuStreamSynchronize(self.stream.cu_stream()); } - let err = self.stream.context().check_err(); - eprintln!("H100_HANG4: step {step} — stream sync done, err={err:?}"); - } // Bookkeeping. agent.fused_post_step_gpu_bookkeeping() .map_err(|e| anyhow::anyhow!("Fused GPU bookkeeping: {e}"))?; - if step < 3 { eprintln!("H100_HANG4: step {step} — bookkeeping done"); } self.steps_since_varmap_sync += 1; self.last_combined_norm = fused_result.grad_norm; - if step < 3 { eprintln!("H100_HANG4: step {step} — varmap_sync={}, graph_aux={}", self.steps_since_varmap_sync, self.graph_aux.is_some()); } // Capture graph_aux after the second step (all sub-trainers initialized). if self.graph_aux.is_none() && self.steps_since_varmap_sync == 2 { - if step < 3 { eprintln!("H100_HANG4: step {step} — capturing graph_aux"); } if let Err(e) = self.capture_graph_aux(agent, gpu_batch) { tracing::warn!("graph_aux capture failed (non-fatal, continuing ungraphed): {e}"); } - if step < 3 { eprintln!("H100_HANG4: step {step} — graph_aux capture done"); } } - if step < 3 { eprintln!("H100_HANG4: step {step} — building result"); } tracing::debug!( step = self.steps_since_varmap_sync, graphed_aux = self.graph_aux.is_some(), diff --git a/crates/ml/src/trainers/dqn/trainer/training_loop.rs b/crates/ml/src/trainers/dqn/trainer/training_loop.rs index 1c5daa541..fbecdab13 100644 --- a/crates/ml/src/trainers/dqn/trainer/training_loop.rs +++ b/crates/ml/src/trainers/dqn/trainer/training_loop.rs @@ -1175,7 +1175,6 @@ impl DQNTrainer { } for _step in 0..num_training_steps { - if _step < 3 { eprintln!("H100_LOOP: step {_step}/{num_training_steps} — sampling PER"); } let sample_start = std::time::Instant::now(); let (batch, vaccine_batch) = { let agent = self.agent.read().await; @@ -1200,11 +1199,9 @@ impl DQNTrainer { fused.pending_vaccine_batch = vb.gpu_batch; } let _fused_result = fused.run_full_step(&batch, &mut *agent, &self.device) - .map_err(|e| { eprintln!("!!! FUSED STEP ERROR: {:#}", e); e }) .context("Fused CUDA training step failed")?; fused_total_us += fused_start.elapsed().as_micros() as u64; - if _step < 3 { eprintln!("H100_LOOP: step {_step} — fused done, running guard"); } let guard_start = std::time::Instant::now(); if let Some(ref mut guard) = self.training_guard { @@ -1257,7 +1254,6 @@ impl DQNTrainer { } train_step_count += 1; self.gradient_logging_step += 1; - if _step < 3 { eprintln!("H100_LOOP: step {_step} — complete, looping"); } } }