Fix MuonClip rolling checkpoint diagnostics

baseline/nanogpt_one_head/src/rg_nanogpt_one_head/train_loop.py

@@ -97,6 +97,43 @@ def _evaluation_due(
     )
 
 
+def _checkpoint_due(
+    step: int,
+    *,
+    cfg: dict,
+    epoch_steps: dict[int, float],
+    total_steps: int,
+) -> bool:
+    return (
+        step % int(cfg["training"]["checkpoint_interval_steps"]) == 0
+        or step in epoch_steps
+        or step == total_steps
+    )
+
+
+def _resume_diagnostics_due(
+    step: int,
+    *,
+    cfg: dict,
+    epoch_steps: dict[int, float],
+    total_steps: int,
+) -> bool:
+    # A rolling checkpoint must carry finite gradient diagnostics even when it
+    # falls between evaluation rows. Otherwise a checkpoint interval shorter
+    # than the evaluation interval persists the initialization NaN sentinels.
+    return _evaluation_due(
+        step,
+        cfg=cfg,
+        epoch_steps=epoch_steps,
+        total_steps=total_steps,
+    ) or _checkpoint_due(
+        step,
+        cfg=cfg,
+        epoch_steps=epoch_steps,
+        total_steps=total_steps,
+    )
+
+
 def _resume_diagnostics(
     previous_eval_snapshot: list[torch.Tensor],
     *,
@@ -486,7 +523,7 @@ def execute_training_loop(
         last_update_lrs = dict(next_update_lrs)
 
         new_step = completed_steps + 1
-        if _evaluation_due(
+        if _resume_diagnostics_due(
             new_step,
             cfg=cfg,
             epoch_steps=epoch_steps,
@@ -506,14 +543,11 @@ def execute_training_loop(
                 else False
             )
 
-        checkpoint_due = (
-            new_step
-            % int(
-                cfg["training"]["checkpoint_interval_steps"]
-            )
-            == 0
-            or new_step in epoch_steps
-            or new_step == total_steps
+        checkpoint_due = _checkpoint_due(
+            new_step,
+            cfg=cfg,
+            epoch_steps=epoch_steps,
+            total_steps=total_steps,
         )
         if checkpoint_due:
             _require_finite_model(

baseline/nanogpt_one_head/tests/test_muonclip_large_20260830.py

@@ -26,6 +26,10 @@
 from rg_nanogpt_one_head.model import GPT, GPTConfig, transformer_matrix_items
 from rg_nanogpt_one_head.muonclip import install_muonclip_extension
 from rg_nanogpt_one_head.spectral import _validate_weightwatcher_frame
+from rg_nanogpt_one_head.train_loop import (
+    _evaluation_due,
+    _resume_diagnostics_due,
+)
 
 
 def _load_config() -> dict:
@@ -58,6 +62,27 @@ def test_large_muonclip_protocol_has_the_declared_scale_and_schedule() -> None:
     assert profile["warmup_fraction"] == pytest.approx(0.016)
 
 
+def test_checkpoint_before_first_evaluation_materializes_gradients() -> None:
+    cfg = _load_config()
+    epoch_steps = epoch_step_map(cfg)
+    total_steps = max_steps(cfg)
+
+    assert cfg["training"]["checkpoint_interval_steps"] == 250
+    assert cfg["training"]["eval_interval_steps"] == 500
+    assert not _evaluation_due(
+        250,
+        cfg=cfg,
+        epoch_steps=epoch_steps,
+        total_steps=total_steps,
+    )
+    assert _resume_diagnostics_due(
+        250,
+        cfg=cfg,
+        epoch_steps=epoch_steps,
+        total_steps=total_steps,
+    )
+
+
 def test_gpt_and_matrix_inventory_support_multiple_blocks() -> None:
     model = GPT(
         GPTConfig(