feat: add --save-total-limit to auto-cleanup old checkpoints

[jxiaof]

Training scripts save a checkpoint every --save-interval steps with no limit on how many are kept. For long runs (e.g. 500K steps with interval=2000), this produces 250+ checkpoints consuming 1.1TB+ of disk.

This PR adds a --save-total-limit argument to all three training scripts (train_domino, train_dflash, train_eagle3). When set, the oldest checkpoints beyond the limit are automatically removed after each save. Default is None (no cleanup) to preserve backward compatibility.

Changes:

• Add cleanup_checkpoints() utility in specforge/utils.py
• Add --save-total-limit arg to all three training scripts
• Call cleanup_checkpoints() after each periodic save (not after final)
• Add comprehensive unit tests in tests/test_utils/test_checkpoint_cleanup.py

Motivation

Modifications

Related Issues

Accuracy Test

Benchmark & Profiling

Checklist

• Format your code according to the Code Formatting with Pre-Commit.
• Add unit tests as outlined in the Running Unit Tests.
• Update documentation / docstrings / example tutorials as needed, according to Writing Documentation.
• Provide throughput / latency benchmark results and accuracy evaluation results as needed, according to Benchmark and Profiling and Accuracy Results.
• For reviewers: If you haven't made any contributions to this PR and are only assisting with merging the main branch, please remove yourself as a co-author when merging the PR.
• Please feel free to join our Slack channel at https://sgl-fru7574.slack.com/archives/C09784E3EN6 to discuss your PR.

[gemini-code-assist]

Code Review

This pull request introduces a checkpoint cleanup utility to limit the total number of saved checkpoints during training, adding a --save-total-limit argument to the DFlash, Domino, and Eagle3 training scripts. The utility parses, sorts, and removes older checkpoint directories exceeding the specified limit. The review feedback correctly identifies a critical issue in distributed training settings: the cleanup_checkpoints function does not enforce that only rank 0 performs the deletion, which can lead to race conditions and crashes on shared filesystems. A code suggestion is provided to restrict the cleanup logic to rank 0.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.

[gemini-code-assist]

In distributed training settings with a shared filesystem (e.g., NFS, Lustre), having all ranks attempt to delete the same checkpoint directories simultaneously will cause race conditions, leading to FileNotFoundError or OSError on non-zero ranks and crashing the training run.

Although the docstring mentions that only rank 0 performs the deletion, the implementation does not actually enforce this check. Adding a check for dist.is_initialized() and dist.get_rank() != 0 ensures that only rank 0 executes the cleanup, while other ranks return early and wait at the subsequent dist.barrier() in the training scripts.

Suggested change

	if save_total_limit is None or save_total_limit <= 0:
	return 0
	if not os.path.isdir(output_dir):
	return 0
	if save_total_limit is None or save_total_limit <= 0:
	return 0
	if dist.is_initialized() and dist.get_rank() != 0:
	return 0
	if not os.path.isdir(output_dir):
	return 0

[Copilot]

Pull request overview

Adds an optional checkpoint retention limit to SpecForge training scripts to prevent unbounded disk growth by automatically deleting the oldest checkpoints after periodic saves.

Changes:

• Introduces cleanup_checkpoints() in specforge/utils.py to delete oldest epoch_*_step_* checkpoint directories beyond a configured limit.
• Adds --save-total-limit CLI flag to train_domino.py, train_dflash.py, and train_eagle3.py, and calls cleanup after periodic checkpoint saves.
• Adds unit tests covering limit behavior and ignoring non-checkpoint files/dirs.

Reviewed changes

Copilot reviewed 5 out of 5 changed files in this pull request and generated 5 comments.

Show a summary per file

File	Description
specforge/utils.py	Adds cleanup_checkpoints() utility for pruning oldest checkpoint directories.
scripts/train_domino.py	Adds --save-total-limit argument and calls checkpoint cleanup after periodic saves.
scripts/train_dflash.py	Adds --save-total-limit argument and calls checkpoint cleanup after periodic saves.
scripts/train_eagle3.py	Adds --save-total-limit argument and calls checkpoint cleanup after periodic saves.
tests/test_utils/test_checkpoint_cleanup.py	Adds unit tests for checkpoint cleanup behavior.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[jxiaof]

Addressed all bot review comments and rebased the branch onto the latest main with a clean linear 2-commit history.

What changed:

• restricted checkpoint cleanup to rank 0 in distributed runs
• made print_on_rank0 safe in non-distributed contexts
• ignored FileNotFoundError during checkpoint deletion races
• gated cleanup + barrier on enabled save_total_limit only
• added focused regression tests for non-zero-rank skip and delete-race handling

Validation:

• ran python -m pre_commit run --all-files
• verified the touched files for syntax/editor errors
• ran targeted checkpoint cleanup behavior validation

Ready for human review.
@FlamingoPg @shuaills could you please take a look when you have time?

[jxiaof]

@sleepcoo @shuaills

[Copilot]

Pull request overview

Copilot reviewed 5 out of 5 changed files in this pull request and generated 1 comment.

[jessiewei7]

This wires --save-total-limit into train_dflash.py / train_domino.py / train_eagle3.py, but scripts/train_peagle.py has the same save path and is left out. Worth either adding it here (looks like the same three-line insertion) or noting in the description why it's excluded, so it doesn't get lost.

[jxiaof]

Updates:

made checkpoint cleanup best-effort for generic OSError cases, so failures to remove old checkpoints no longer abort training
added a warning log and retry-on-next-save behavior for failed deletions
wired --save-total-limit into train_peagle.py as well, since it uses the same checkpoint save path
added focused regression coverage for the OSError cleanup path
Validation:

ran pre-commit on the touched files
ran focused syntax/error checks
re-validated the targeted cleanup behavior
If you have time, I’d appreciate another look.
@jessiewei7 @FlamingoPg @shuaills @sleepcoo @FrankLeeeee

[jxiaof]

Follow-up update after the latest review comments:

• extended checkpoint cleanup to be best-effort for generic OSError cases as well, so failures to remove old checkpoints no longer abort training
• added a warning + retry-on-next-save behavior for failed deletions
• wired --save-total-limit into train_peagle.py too, since it uses the same checkpoint save path
• added focused regression coverage for the OSError cleanup path

Validation:

• ran pre-commit on the touched files
• ran focused syntax/editor checks on the touched files
• re-validated targeted checkpoint cleanup behavior, including the OSError fallback path

At this point all training scripts in this PR that save periodic checkpoints (train_domino.py, train_dflash.py, train_eagle3.py, and train_peagle.py) are covered consistently.