gitadmin/foxhunt
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

diag: per-line logging inside replay_adam_and_readback

Browse Source
This commit is contained in:
jgrusewski
2026-04-06 00:16:52 +02:00
parent e2da8255e1
commit dde8adbd9c
1 changed files with 7 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

7
crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs
View File

@@ -2990,6 +2990,10 @@ impl GpuDqnTrainer {
}
pub fn replay_adam_and_readback(&mut self) -> Result<FusedTrainScalars, MLError> {
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 {