From dde8adbd9c48ac829757bbd1e9196ae98b392d7b Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Mon, 6 Apr 2026 00:16:52 +0200 Subject: [PATCH] diag: per-line logging inside replay_adam_and_readback --- crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs index fc72672d8..3b52fdb41 100644 --- a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs +++ b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs @@ -2990,6 +2990,10 @@ 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 { tracing::info!(adam_step, "adam_readback: enter"); } + // 1. Collect previous step's scalars (if pending) let prev_scalars = if self.readback_pending { if let Some(ref event) = self.readback_event { @@ -3012,8 +3016,11 @@ impl GpuDqnTrainer { }; // 2. Launch current step: grad_norm + adam + if adam_step < 3 { tracing::info!(adam_step, "adam_readback: compute_grad_norm"); } self.compute_grad_norm_outside_graph()?; + if adam_step < 3 { tracing::info!(adam_step, "adam_readback: replay_adam"); } self.replay_adam()?; + if adam_step < 3 { tracing::info!(adam_step, "adam_readback: async readback"); } // 3. Async DtoH for current step's scalars (non-blocking) unsafe {