From da71b46af80b8ba2dea450b3e8edd03e21657a0a Mon Sep 17 00:00:00 2001 From: camilobrownpinilla Date: Fri, 26 Jun 2026 16:40:15 -0400 Subject: [PATCH 1/5] Fix clean runs not always saving on last step --- CHANGELOG.md | 3 +++ docs/training/training-loop.md | 10 ++++++++++ scripts/train.py | 20 ++++++++++++++++++++ tests/e2e/test_training_e2e.py | 31 +++++++++++++++++++++++++++++++ 4 files changed, 64 insertions(+) diff --git a/CHANGELOG.md b/CHANGELOG.md index 5a2c8c9..ea0067d 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -81,6 +81,9 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 - `docs/checkpointing/dcp-model.md`: updated the save/load snippets and the "shape to fill" explanation to the DCP-aware helpers. - Tests (fail on the pre-fix code, pass after): `tests/integration/test_checkpoint_roundtrip.py::test_manager_restores_optimizer_moments_single_gpu` and `tests/distributed/test_checkpoint.py::test_resume_restores_optimizer_moments` assert `exp_avg` / `exp_avg_sq` are restored bit-exactly into a *fresh* optimizer (single-GPU + distributed); `tests/e2e/test_training_e2e.py::test_resume_determinism_single_gpu` / `test_resume_determinism_2gpu_fsdp` assert end-to-end bit-exact loss across an interrupt-and-resume on a learnable dataset. - **On-disk format note:** optimizer state is now keyed by parameter fully-qualified name rather than positional index. Checkpoints written before this fix will not restore optimizer state on resume (training continues with a fresh optimizer); model state is unaffected. +- **Always save a final checkpoint on clean completion.** A run that finished cleanly previously wrote a final checkpoint only when `max_steps` happened to land on the save schedule; otherwise the fully-trained model (including the entire WSD decay tail) was never persisted and `latest` resolved to the last scheduled step. The training-loop teardown now writes one unconditional checkpoint at `max_steps` on normal completion when that step is off-schedule, matching the emergency-checkpoint guarantee on the preemption path. Purely additive: no existing checkpoint is changed, moved, or pruned, and the schedule/retention knobs (`interval`, `keep_last_n`, `dyn_ckpt_window`) are untouched. + - `scripts/train.py`: a `completed_normally` marker set on the final iteration, plus an off-schedule final `ckpt_mgr.save(...)` (and `on_checkpoint_save` hook) before `ckpt_mgr.wait()`. + - Tests: `tests/e2e/test_training_e2e.py` — a clean off-schedule completion writes `step_{max_steps}` and points `latest` at it. ## [0.1.0] — 2026-04-16 diff --git a/docs/training/training-loop.md b/docs/training/training-loop.md index 641880d..4112d64 100644 --- a/docs/training/training-loop.md +++ b/docs/training/training-loop.md @@ -185,12 +185,22 @@ After the loop: ```python prof.stop() +# Clean off-schedule finish: persist the fully-trained final step +if completed_normally and not config.checkpoint.should_save(step): + ckpt_mgr.save(step, ...) # final checkpoint; updates `latest` ckpt_mgr.wait() # drain last async save hook_runner.on_train_end(step, tokens_seen) tracker.close() destroy_distributed() ``` +On a clean finish, an unconditional checkpoint is written at `max_steps` +when that step is not already on the save schedule — so a completed run's +fully-trained model (including the WSD decay tail) is always recoverable +and `latest` points at it. This mirrors the emergency checkpoint the +preemption path writes on shutdown. The `should_save` guard avoids a +duplicate when `max_steps` already coincided with the schedule. + `ckpt_mgr.wait()` is load-bearing — without it, a rank can exit before its async DCP write completes, corrupting the checkpoint for everyone else on the same save. See diff --git a/scripts/train.py b/scripts/train.py index d9b8d29..531f14d 100644 --- a/scripts/train.py +++ b/scripts/train.py @@ -649,6 +649,8 @@ def __getitem__(self, idx: int) -> dict[str, torch.Tensor]: ) hook_runner.on_checkpoint_save(0, config.checkpoint.dir) + completed_normally = False + while step < tc.max_steps: # Refresh data iterator at start / epoch boundary if dataloader is not None and data_iter is None: @@ -1016,11 +1018,29 @@ def __getitem__(self, idx: int) -> dict[str, torch.Tensor]: shutdown_handler.finish() break + # Clean-completion marker for the unconditional final save after the + # loop. Reached only on the iteration that hits max_steps with no + # break, so it is never set after a NaN/NCCL/shutdown break and never + # on a zero-iteration resume — guaranteeing ckpt_extra is bound when + # the final save runs. + if step >= tc.max_steps: + completed_normally = True + if prof is not None: prof.stop() if rank == 0: print_profiler_summary(prof, trace_dir=config.profiling.trace_dir) + if completed_normally and not config.checkpoint.should_save(step): + ckpt_mgr.save( + step=step, + tokens_seen=tokens_seen, + scheduler=scheduler, + dataloader=dataloader, + extra=ckpt_extra, + ) + hook_runner.on_checkpoint_save(step, config.checkpoint.dir) + # Flush any pending async checkpoint before tearing down process group ckpt_mgr.wait() diff --git a/tests/e2e/test_training_e2e.py b/tests/e2e/test_training_e2e.py index 86983bc..875108f 100644 --- a/tests/e2e/test_training_e2e.py +++ b/tests/e2e/test_training_e2e.py @@ -606,6 +606,37 @@ def test_sigterm_triggers_emergency_checkpoint(tmp_path): assert "Checkpoint saved" in output, f"Emergency checkpoint was not saved:\n{output[-2000:]}" +@pytest.mark.e2e +def test_final_checkpoint_on_clean_completion(tmp_path): + """A clean finish at an off-schedule step must still persist a final + checkpoint, with `latest` pointing at it.""" + ckpt_dir = tmp_path / "final_ckpt" + + # max_steps=10 with interval=1000 => no scheduled save lands (debug.toml + # configures no dyn_ckpt_window), so the only checkpoint must be the + # unconditional final save at step 10. + result = _run_training( + [ + DEBUG_CONFIG, + "--train.max_steps=10", + "--metrics.log_interval=5", + f"--checkpoint.dir={ckpt_dir}", + "--checkpoint.interval=1000", # no scheduled checkpoint + ], + nproc=1, + ) + _assert_training_complete(result, expected_steps=10) + + output = result.stdout + result.stderr + final = ckpt_dir / "step_10" + latest = ckpt_dir / "latest" + assert final.is_dir(), f"final checkpoint step_10 was not written:\n{output[-2000:]}" + assert latest.exists(), f"`latest` pointer is missing:\n{output[-2000:]}" + assert latest.resolve().name == "step_10", ( + f"`latest` should resolve to step_10, got {latest.resolve().name!r}" + ) + + # ============================================================================ # MoE Training # ============================================================================ From 2a9256521f5d2c098eaa576b06bd8ff96fcc4ff4 Mon Sep 17 00:00:00 2001 From: Camilo Brown-Pinilla Date: Fri, 26 Jun 2026 17:06:51 -0400 Subject: [PATCH 2/5] Fix inaccurate comment Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com> --- docs/training/training-loop.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs/training/training-loop.md b/docs/training/training-loop.md index 4112d64..e7eea1f 100644 --- a/docs/training/training-loop.md +++ b/docs/training/training-loop.md @@ -187,7 +187,7 @@ After the loop: prof.stop() # Clean off-schedule finish: persist the fully-trained final step if completed_normally and not config.checkpoint.should_save(step): - ckpt_mgr.save(step, ...) # final checkpoint; updates `latest` + ckpt_mgr.save(step, ...) # final checkpoint; `latest` committed after wait ckpt_mgr.wait() # drain last async save hook_runner.on_train_end(step, tokens_seen) tracker.close() From 7d51dd9ecbdbf5b8c6c773b76c8aa49179c4a95b Mon Sep 17 00:00:00 2001 From: camilobrownpinilla Date: Mon, 29 Jun 2026 09:38:15 -0400 Subject: [PATCH 3/5] Improve comments --- scripts/train.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/scripts/train.py b/scripts/train.py index 1fff8bc..bb57634 100644 --- a/scripts/train.py +++ b/scripts/train.py @@ -1042,10 +1042,10 @@ def __getitem__(self, idx: int) -> dict[str, torch.Tensor]: break # Clean-completion marker for the unconditional final save after the - # loop. Reached only on the iteration that hits max_steps with no - # break, so it is never set after a NaN/NCCL/shutdown break and never - # on a zero-iteration resume — guaranteeing ckpt_extra is bound when - # the final save runs. + # loop. Only reached when training completes without any errors, e.g., + # no NaN/NCCl/shutdown breaks. If a run encounters a NaN, the last step + # is intentionally *not* saved because the actual model state would be + # `max_steps` - 1, not `max_steps`. if step >= tc.max_steps: completed_normally = True From 43e1af29f6279f66d3948dd88f0682bdc8f22bdd Mon Sep 17 00:00:00 2001 From: Camilo Brown-Pinilla Date: Tue, 30 Jun 2026 10:55:32 -0400 Subject: [PATCH 4/5] Fix comment spelling & formatting Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com> --- scripts/train.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/scripts/train.py b/scripts/train.py index bb57634..40e8358 100644 --- a/scripts/train.py +++ b/scripts/train.py @@ -1043,9 +1043,9 @@ def __getitem__(self, idx: int) -> dict[str, torch.Tensor]: # Clean-completion marker for the unconditional final save after the # loop. Only reached when training completes without any errors, e.g., - # no NaN/NCCl/shutdown breaks. If a run encounters a NaN, the last step + # no NaN/NCCL/shutdown breaks. If a run encounters a NaN, the last step # is intentionally *not* saved because the actual model state would be - # `max_steps` - 1, not `max_steps`. + # `max_steps - 1`, not `max_steps`. if step >= tc.max_steps: completed_normally = True From a88270eee2cf3dc465f51c377703ac9d1539cd5c Mon Sep 17 00:00:00 2001 From: Camilo Brown-Pinilla Date: Tue, 30 Jun 2026 16:02:07 -0400 Subject: [PATCH 5/5] Fix inaccurate changelog entry --- CHANGELOG.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index eb2203b..d0fff79 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -87,7 +87,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 - `docs/checkpointing/dcp-model.md`: updated the save/load snippets and the "shape to fill" explanation to the DCP-aware helpers. - Tests (fail on the pre-fix code, pass after): `tests/integration/test_checkpoint_roundtrip.py::test_manager_restores_optimizer_moments_single_gpu` and `tests/distributed/test_checkpoint.py::test_resume_restores_optimizer_moments` assert `exp_avg` / `exp_avg_sq` are restored bit-exactly into a *fresh* optimizer (single-GPU + distributed); `tests/e2e/test_training_e2e.py::test_resume_determinism_single_gpu` / `test_resume_determinism_2gpu_fsdp` assert end-to-end bit-exact loss across an interrupt-and-resume on a learnable dataset. - **On-disk format note:** optimizer state is now keyed by parameter fully-qualified name rather than positional index. Checkpoints written before this fix will not restore optimizer state on resume (training continues with a fresh optimizer); model state is unaffected. -- **Always save a final checkpoint on clean completion.** A run that finished cleanly previously wrote a final checkpoint only when `max_steps` happened to land on the save schedule; otherwise the fully-trained model (including the entire WSD decay tail) was never persisted and `latest` resolved to the last scheduled step. The training-loop teardown now writes one unconditional checkpoint at `max_steps` on normal completion when that step is off-schedule, matching the emergency-checkpoint guarantee on the preemption path. Purely additive: no existing checkpoint is changed, moved, or pruned, and the schedule/retention knobs (`interval`, `keep_last_n`, `dyn_ckpt_window`) are untouched. +- **Always save a final checkpoint on clean completion.** A run that finished cleanly previously wrote a final checkpoint only when `max_steps` happened to land on the save schedule; otherwise the fully-trained model (including the entire WSD decay tail) was never persisted and `latest` resolved to the last scheduled step. The training-loop teardown now writes one unconditional checkpoint at `max_steps` on normal completion when that step is off-schedule, matching the emergency-checkpoint guarantee on the preemption path. Purely additive: the schedule/retention knobs (`interval`, `keep_last_n`, `dyn_ckpt_window`) are untouched. - `scripts/train.py`: a `completed_normally` marker set on the final iteration, plus an off-schedule final `ckpt_mgr.save(...)` (and `on_checkpoint_save` hook) before `ckpt_mgr.wait()`. - Tests: `tests/e2e/test_training_e2e.py` — a clean off-schedule completion writes `step_{max_steps}` and points `latest` at it.