diag: add per-step STEP_DIAG logging (zero-cost, uses existing readback)

Logs grad_norm, loss, c51_alpha every 50 steps via existing async readback
scalars. No stream sync or extra GPU work. Works with mega-graph path.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-04-09 16:35:26 +02:00
parent 7036ba3589
commit c26f0ae451

View File

@@ -853,6 +853,20 @@ impl FusedTrainingCtx {
.map_err(|e| anyhow::anyhow!("Fused GPU bookkeeping: {e}"))?;
self.steps_since_varmap_sync += 1;
// Per-step gradient diagnostic — zero cost (uses existing readback, no sync).
// Log every 50 steps to see gradient trajectory during collapse.
if self.steps_since_varmap_sync % 50 == 1 {
let alpha = self.trainer.c51_alpha();
tracing::warn!(
step = self.steps_since_varmap_sync,
grad_norm = fused_result.grad_norm,
loss = fused_result.total_loss,
c51_alpha = alpha,
batch_size = self.batch_size,
"STEP_DIAG: per-step gradient readback"
);
}
self.last_combined_norm = fused_result.grad_norm;
// Capture mega-graph at step 2 (all sub-trainers initialized).