diff --git a/baseline/nanogpt_one_head/src/rg_nanogpt_one_head/train_loop.py b/baseline/nanogpt_one_head/src/rg_nanogpt_one_head/train_loop.py index 77de987..dad7310 100644 --- a/baseline/nanogpt_one_head/src/rg_nanogpt_one_head/train_loop.py +++ b/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( diff --git a/baseline/nanogpt_one_head/tests/test_muonclip_large_20260830.py b/baseline/nanogpt_one_head/tests/test_muonclip_large_20260830.py index aa11fb9..eff6d50 100644 --- a/baseline/nanogpt_one_head/tests/test_muonclip_large_20260830.py +++ b/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(