From 29459ab1f7b51c12c26b507067e009165e081005 Mon Sep 17 00:00:00 2001
From: Lars Bogner <upyxo@student.kit.edu>
Date: Mon, 27 Jul 2026 16:20:05 +0200
Subject: [PATCH] Log batch-level metrics to W&B, not just per-epoch summaries

giant train --wandb now also logs loss/grad_norm/lr every N optimizer
steps (--wandb-log-every, default 50) so W&B shows within-epoch trends,
not just one point per epoch. Both share global_step as a single
monotonic step axis (wandb.Run.log requires step to never decrease
across calls), which also fixes global_step previously only advancing
in wgan mode.

Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
---
 giant/cli.py      |  9 +++++++++
 giant/config.py   |  6 ++++++
 giant/pipeline.py |  1 +
 giant/train.py    | 25 +++++++++++++++++++++++--
 4 files changed, 39 insertions(+), 2 deletions(-)

diff --git a/giant/cli.py b/giant/cli.py
index 310a342..73993fe 100644
--- a/giant/cli.py
+++ b/giant/cli.py
@@ -415,6 +415,14 @@ def train(
         Optional[str],
         typer.Option("--wandb-run-name", help="W&B run name (default: out_dir name)"),
     ] = None,
+    wandb_log_every: Annotated[
+        Optional[int],
+        typer.Option(
+            "--wandb-log-every",
+            help="Log batch-level loss/grad_norm/lr to W&B every N optimizer "
+            "steps (default: 50); per-epoch metrics always log in full",
+        ),
+    ] = None,
 ) -> None:
     """Train the GIANT surrogate model."""
     batch_size_auto = False
@@ -455,6 +463,7 @@ def train(
             "wandb": wandb,
             "wandb_project": wandb_project,
             "wandb_run_name": wandb_run_name,
+            "wandb_log_every": wandb_log_every,
         }.items()
         if v is not None
     }
diff --git a/giant/config.py b/giant/config.py
index 5d3dd9f..54cbcd6 100644
--- a/giant/config.py
+++ b/giant/config.py
@@ -42,6 +42,12 @@ DEFAULT_CONFIG: dict = {
         "wandb": False,
         "wandb_project": "giant",
         "wandb_run_name": "",
+        # Batch-granularity metrics (loss/grad_norm/lr) are logged every N
+        # optimizer steps, not every batch — a single epoch can be tens of
+        # thousands of steps (see steps_per_epoch above), and logging every
+        # one of them would flood the run with points the UI has to downsample
+        # anyway. Per-epoch metrics (the metrics.csv row) always log in full.
+        "wandb_log_every": 50,
     },
     "model": {
         "hidden_dim": 256,
diff --git a/giant/pipeline.py b/giant/pipeline.py
index 2a7ba31..de49898 100644
--- a/giant/pipeline.py
+++ b/giant/pipeline.py
@@ -262,4 +262,5 @@ def run_train_job(
         use_wandb=t.get("wandb", False),
         wandb_project=t.get("wandb_project", "giant"),
         wandb_run_name=t.get("wandb_run_name", ""),
+        wandb_log_every=t.get("wandb_log_every", 50),
     )
diff --git a/giant/train.py b/giant/train.py
index 8a160a3..a520368 100644
--- a/giant/train.py
+++ b/giant/train.py
@@ -323,6 +323,7 @@ def train(
     use_wandb: bool = False,
     wandb_project: str = "giant",
     wandb_run_name: str = "",
+    wandb_log_every: int = 50,
 ) -> None:
     out_dir = Path(out_dir)
     out_dir.mkdir(parents=True, exist_ok=True)
@@ -539,7 +540,6 @@ def train(
                         lambda_nsec,
                         lambda_s2,
                     )
-                    global_step += 1
                     if stats["did_g_step"]:
                         lr_sched.step()
                         if ema_decay > 0:
@@ -609,6 +609,23 @@ def train(
                     f"loss={ema_loss:.4f} gnorm={ema_grad_norm:.3f}", refresh=False
                 )
 
+                global_step += 1
+                if (
+                    wandb_run is not None
+                    and wandb_log_every > 0
+                    and global_step % wandb_log_every == 0
+                ):
+                    wandb_run.log(
+                        {
+                            "batch/epoch": epoch,
+                            "batch/loss": batch_loss,
+                            "batch/loss_ema": ema_loss,
+                            "batch/grad_norm": batch_grad_norm,
+                            "batch/lr": optimizer.param_groups[0]["lr"],
+                        },
+                        step=global_step,
+                    )
+
                 if shutdown.requested:
                     break
             bar.close()
@@ -747,7 +764,11 @@ def train(
             metrics_writer.writerow(metrics_row)
             metrics_file.flush()
             if wandb_run is not None:
-                wandb_run.log(metrics_row, step=epoch)
+                # Shares the same monotonic step axis as the per-batch
+                # `batch/*` logs above (global_step) rather than `epoch`,
+                # since a wandb run's `step` argument across `log()` calls
+                # must never decrease.
+                wandb_run.log(metrics_row, step=global_step)
 
             ckpt: dict = {
                 "model": stage1_model.state_dict(),