Skip to content

fix: sort checkpoints numerically by step suffix #25

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

Merged
Neonkraft merged 2 commits into from
May 19, 2026
Merged

fix: sort checkpoints numerically by step suffix #25

Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
dfa4ccb
fix: sort checkpoints numerically by step suffix
Neonkraft May 18, 2026
aa1239e
refactor: use transformers.get_last_checkpoint for auto-resume
Neonkraft May 18, 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
5 changes: 3 additions & 2 deletions scripts/train.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -124,9 +124,10 @@ def main() -> None:
return

# Auto-resume from the latest checkpoint if one exists.
from transformers.trainer_utils import get_last_checkpoint

@KonstiNik KonstiNik May 19, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

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

What is the reason for the in-line import?

@Neonkraft Neonkraft May 19, 2026

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

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

This is the latest change. I decided to just use the get_last_checkpoint function from transformers to get the latest checkpoint. It sorts the checkpoints and ignores checkpoints without a numerical suffix. Source code:

PREFIX_CHECKPOINT_DIR = "checkpoint"
_re_checkpoint = re.compile(r"^" + PREFIX_CHECKPOINT_DIR + r"\-(\d+)$")


def get_last_checkpoint(folder):
    content = os.listdir(folder)
    checkpoints = [
        path
        for path in content
        if _re_checkpoint.search(path) is not None and os.path.isdir(os.path.join(folder, path))
    ]
    if len(checkpoints) == 0:
        return
    return os.path.join(folder, max(checkpoints, key=lambda x: int(_re_checkpoint.search(x).groups()[0])))

@KonstiNik KonstiNik May 19, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

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

I should have been a bit more precise – Why import there and not at the beginning of the file?

@Neonkraft Neonkraft May 19, 2026

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

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

Because we won't need this import at all if we're using --tokenize_only. This shouldn't make much of a difference. I can move it to the top if you feel it's cleaner that way.

@KonstiNik KonstiNik May 19, 2026 •
edited
Loading

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

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

I see. I go with how you prefer it.


checkpoints_dir = run_dir / "checkpoints"
existing = sorted(checkpoints_dir.glob("checkpoint-*"))
resume_from = str(existing[-1]) if existing else None
resume_from = get_last_checkpoint(str(checkpoints_dir)) if checkpoints_dir.exists() else None
if resume_from:
logger.info("Resuming from checkpoint: %s", resume_from)

Expand Down
Loading