From bb699d41b21a272115283052bddd033253836fb4 Mon Sep 17 00:00:00 2001 From: Lars Bogner Date: Wed, 29 Jul 2026 10:42:35 +0200 Subject: [PATCH] Persist global_step across --resume so W&B step stays monotonic Previously global_step always reset to 0, even on --resume. Since the W&B run reattaches to the same run id on resume, logging with step=global_step after a restart passed step values below what was already recorded, silently dropping the resumed portion's metrics. Co-Authored-By: Claude Sonnet 5 --- giant/train.py | 9 ++++++++- 1 file changed, 8 insertions(+), 1 deletion(-) diff --git a/giant/train.py b/giant/train.py index cd80217..473652c 100644 --- a/giant/train.py +++ b/giant/train.py @@ -483,6 +483,7 @@ def train( start_epoch = 1 best_val_loss = float("inf") + resumed_global_step = 0 if resume_path is not None: ckpt = torch.load(resume_path, map_location=device, weights_only=False) stage1_model.load_state_dict(ckpt["model"]) @@ -506,6 +507,7 @@ def train( lr_sched.load_state_dict(ckpt["lr_sched"]) start_epoch = ckpt.get("epoch", 0) + 1 best_val_loss = ckpt.get("best_val_loss", float("inf")) + resumed_global_step = ckpt.get("global_step", 0) # optimizer/lr_sched.load_state_dict() above restore the checkpoint's # own base LR, which would otherwise silently override an explicit @@ -535,7 +537,11 @@ def train( epoch_w = len(str(epochs)) last_completed_epoch = start_epoch - 1 - global_step = 0 + # Restored from the checkpoint on --resume so wandb_run.log(..., step=...) + # keeps advancing monotonically instead of restarting at 0 mid-run (a + # reattached wandb run — see wandb.init(id=..., resume="allow") below — + # would otherwise silently drop every post-resume point). + global_step = resumed_global_step with _GracefulShutdown() as shutdown: for epoch in range(start_epoch, epochs + 1): epoch_start = time.monotonic() @@ -950,6 +956,7 @@ def train( "lr_sched": lr_sched.state_dict(), "epoch": epoch, "best_val_loss": best_val_loss, + "global_step": global_step, } if mode == "wgan": assert (