Skip to content

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

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
jxiaof wants to merge 4 commits into
base: main
Choose a base branch
from
Open

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

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
cd19b6e
feat: add --save-total-limit to auto-cleanup old checkpoints
jxiaof Jun 23, 2026
8a69d5a
fix: address PR review on checkpoint cleanup
jxiaof Jul 7, 2026
e889fb5
test: cover print_on_rank0 logging behavior
jxiaof Jul 8, 2026
62b7393
fix: make checkpoint cleanup best-effort
jxiaof Jul 8, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
11 changes: 11 additions & 0 deletions scripts/train_dflash.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -37,6 +37,7 @@
from specforge.optimizer import BF16Optimizer
from specforge.tracker import create_tracker
from specforge.utils import (
cleanup_checkpoints,
get_last_checkpoint,
get_local_device,
load_tokenizer,
Expand Down Expand Up @@ -165,6 +166,13 @@ def parse_args():
output_group.add_argument("--log-interval", type=int, default=50)
output_group.add_argument("--eval-interval", type=int, default=1000)
output_group.add_argument("--save-interval", type=int, default=1000)
output_group.add_argument(
"--save-total-limit",
type=int,
default=None,
help="Maximum number of checkpoints to keep. Older checkpoints are "
"removed when this limit is exceeded. None (default) keeps all.",
)

optimization_group = parser.add_argument_group("optimization")
optimization_group.add_argument(
Expand Down Expand Up @@ -637,6 +645,9 @@ def main():
save_checkpoint(
args, epoch, global_step, dflash_model, draft_model, optimizer
)
if args.save_total_limit is not None and args.save_total_limit > 0:
cleanup_checkpoints(args.output_dir, args.save_total_limit)
dist.barrier()

save_checkpoint(
args, args.num_epochs, global_step, dflash_model, draft_model, optimizer
Expand Down
11 changes: 11 additions & 0 deletions scripts/train_domino.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -34,6 +34,7 @@
from specforge.optimizer import BF16Optimizer
from specforge.tracker import create_tracker
from specforge.utils import (
cleanup_checkpoints,
get_last_checkpoint,
get_local_device,
load_tokenizer,
Expand Down Expand Up @@ -141,6 +142,13 @@ def parse_args():
output_group.add_argument("--log-interval", type=int, default=50)
output_group.add_argument("--eval-interval", type=int, default=1000)
output_group.add_argument("--save-interval", type=int, default=1000)
output_group.add_argument(
"--save-total-limit",
type=int,
default=None,
help="Maximum number of checkpoints to keep. Older checkpoints are "
"removed when this limit is exceeded. None (default) keeps all.",
)

optimization_group = parser.add_argument_group("optimization")
optimization_group.add_argument(
Expand Down Expand Up @@ -665,6 +673,9 @@ def main():
save_checkpoint(
args, epoch, global_step, domino_model, draft_model, optimizer
)
if args.save_total_limit is not None and args.save_total_limit > 0:
cleanup_checkpoints(args.output_dir, args.save_total_limit)
dist.barrier()

save_checkpoint(
args, args.num_epochs, global_step, domino_model, draft_model, optimizer
Expand Down
11 changes: 11 additions & 0 deletions scripts/train_eagle3.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -51,6 +51,7 @@
from specforge.optimizer import BF16Optimizer
from specforge.tracker import Tracker, create_tracker, get_tracker_class
from specforge.utils import (
cleanup_checkpoints,
create_draft_config_from_target,
get_last_checkpoint,
load_tokenizer,
Expand Down Expand Up @@ -194,6 +195,13 @@ def build_parser() -> ArgumentParser:
)
training_group.add_argument("--eval-interval", type=int, default=5000)
training_group.add_argument("--save-interval", type=int, default=5000)
training_group.add_argument(
"--save-total-limit",
type=int,
default=None,
help="Maximum number of checkpoints to keep. Older checkpoints are "
"removed when this limit is exceeded. None (default) keeps all.",
)
training_group.add_argument(
"--log-interval",
type=int,
Expand Down Expand Up @@ -1357,6 +1365,9 @@ def main():
if global_step % (args.save_interval * args.draft_accumulation_steps) == 0:
# Save the model
save_checkpoints(args, epoch, global_step, eagle3_model, optimizer)
if args.save_total_limit is not None and args.save_total_limit > 0:
cleanup_checkpoints(args.output_dir, args.save_total_limit)
dist.barrier()

if args.max_num_steps is not None and global_step >= args.max_num_steps:
break
Expand Down
11 changes: 11 additions & 0 deletions scripts/train_peagle.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -43,6 +43,7 @@
from specforge.optimizer import BF16Optimizer
from specforge.tracker import Tracker, create_tracker, get_tracker_class
from specforge.utils import (
cleanup_checkpoints,
get_last_checkpoint,
print_args_with_dots,
print_on_rank0,
Expand Down Expand Up @@ -137,6 +138,13 @@ def parse_args() -> Tuple[ArgumentParser, Namespace]:
training_group.add_argument("--ckpt-dir", type=str, default=None)
training_group.add_argument("--eval-interval", type=int, default=5000)
training_group.add_argument("--save-interval", type=int, default=5000)
training_group.add_argument(
"--save-total-limit",
type=int,
default=None,
help="Maximum number of checkpoints to keep. Older checkpoints are "
"removed when this limit is exceeded. None (default) keeps all.",
)
training_group.add_argument("--log-interval", type=int, default=50)
training_group.add_argument("--seed", type=int, default=0)
training_group.add_argument("--draft-accumulation-steps", type=int, default=1)
Expand Down Expand Up @@ -800,6 +808,9 @@ def main():
# Save Checkpoints
if global_step % args.save_interval == 0:
save_checkpoints(args, epoch, global_step, peagle_model, optimizer)
if args.save_total_limit is not None and args.save_total_limit > 0:
cleanup_checkpoints(args.output_dir, args.save_total_limit)
dist.barrier()

if args.max_num_steps is not None and global_step >= args.max_num_steps:
break
Expand Down
72 changes: 70 additions & 2 deletions specforge/utils.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,7 @@
import logging
import os
import re
import shutil
from contextlib import contextmanager

import torch
Expand Down Expand Up @@ -191,8 +192,10 @@ def print_args_with_dots(args):


def print_on_rank0(message):
if dist.get_rank() == 0:
logger.info(message)
if dist.is_available() and dist.is_initialized() and dist.get_rank() != 0:
return

logger.info(message)


def get_last_checkpoint(folder, prefix="epoch"):
Expand Down Expand Up @@ -237,6 +240,71 @@ def sort_key(x):
return os.path.join(folder, last_checkpoint), (epoch, step)


def cleanup_checkpoints(output_dir, save_total_limit, prefix="epoch"):
"""
Remove oldest checkpoint directories exceeding save_total_limit.

Only rank 0 performs deletion. Callers should invoke ``dist.barrier()``
after this function in distributed settings.

Args:
output_dir: Directory containing checkpoint folders.
save_total_limit: Maximum number of checkpoints to keep.
``None`` or non-positive means no cleanup (backward compatible).
prefix: Checkpoint directory prefix, default is "epoch".

Returns:
Number of checkpoints removed.
"""
if save_total_limit is None or save_total_limit <= 0:
return 0

if dist.is_available() and dist.is_initialized() and dist.get_rank() != 0:
return 0

if not os.path.isdir(output_dir):
return 0
Comment on lines +259 to +266

@gemini-code-assist gemini-code-assist Bot Jun 23, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

high

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


content = os.listdir(output_dir)
_re_checkpoint = re.compile(rf"^{re.escape(prefix)}_(\d+)(?:_step_(\d+))?$")

checkpoints = [
path
for path in content
if _re_checkpoint.search(path) is not None
and os.path.isdir(os.path.join(output_dir, path))
]

if len(checkpoints) <= save_total_limit:
return 0

def sort_key(x):
match = _re_checkpoint.search(x)
epoch = int(match.group(1))
step = int(match.group(2)) if match.group(2) else 0
return (epoch, step)

checkpoints.sort(key=sort_key)
num_to_remove = len(checkpoints) - save_total_limit
removed = 0
for ckpt in checkpoints[:num_to_remove]:
ckpt_path = os.path.join(output_dir, ckpt)
try:
shutil.rmtree(ckpt_path)
except FileNotFoundError:
Comment thread
jxiaof marked this conversation as resolved.
continue
except OSError as error:
print_on_rank0(
f"Failed to remove old checkpoint {ckpt_path}: {error}; will retry on next save"
)
continue

print_on_rank0(f"Removed old checkpoint: {ckpt_path}")
removed += 1

return removed


def generate_draft_model_config(
target_model_path: str, template_config_path: str = None, cache_dir: str = None
):
Expand Down
Loading