feat: support resume_from_checkpoint=True to auto-find last checkpoint

[zhichenggeng]

PR type

• Bug Fix
• New Feature
• Document Updates
• More Models or Datasets Support

PR information

Support resume_from_checkpoint="true"((case-insensitive) to automatically detect and resume from the last checkpoint in output_dir, without needing to specify the exact checkpoint path.

Changes

• sft_args.py
  • Added _resolve_last_checkpoint() method that searches for the last checkpoint directory in output_dir.
• test_resume_from_checkpoint.py
  • Added test when resume_from_checkpoint='true'

Experiment results

N/A

[gemini-code-assist]

Summary of Changes

Hello @zhichenggeng, I'm Gemini Code Assist¹! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

This pull request enhances the training workflow by providing a convenient way to resume from the latest available checkpoint without requiring the user to specify the exact path. By setting resume_from_checkpoint to "true", the system will now automatically locate the most recent checkpoint within the designated output directory, even if it's nested within versioned subdirectories. This improvement streamlines the process of continuing interrupted training sessions and simplifies experiment management.

Highlights

• Automatic Checkpoint Resumption: Introduced support for resume_from_checkpoint="true" (case-insensitive) to automatically detect and resume training from the last checkpoint found in the output_dir.
• Checkpoint Resolution Logic: Added a new private method _resolve_last_checkpoint() in SftArguments to intelligently search for the latest checkpoint, including support for versioned subdirectories (e.g., v0-<timestamp>).
• New Test Case: Included test_resume_from_checkpoint_true() to validate the functionality of automatically finding and resuming from the last checkpoint.

🧠 New Feature in Public Preview: You can now enable Memory to help Gemini Code Assist learn from your team's feedback. This makes future code reviews more consistent and personalized to your project's style. Click here to enable Memory in your admin console.

Changelog

• swift/arguments/sft_args.py
  • Imported re module for regular expression matching and get_last_checkpoint from transformers.trainer_utils.
  • Updated the SftArguments dataclass to include a detailed description for the resume_from_checkpoint argument, explaining the new 'true' option.
  • Modified the __post_init__ method to check if resume_from_checkpoint is set to 'true' (case-insensitive) and, if so, call the new _resolve_last_checkpoint() method to determine the actual checkpoint path.
  • Added the _resolve_last_checkpoint() method, which searches for checkpoint directories in the output_dir and its versioned subdirectories, returning the path to the latest checkpoint or None if none is found.
• tests/train/test_resume_from_checkpoint.py
  • Added a new test function test_resume_from_checkpoint_true() to verify that setting resume_from_checkpoint='true' correctly identifies and resumes from the last generated checkpoint.
  • Updated the main execution block to include the call to the new test function.

Activity

• No specific activity (comments, reviews, progress updates) was provided in the context for this pull request.

Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point by creating a comment using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands on the current page.

Feature	Command	Description
Code Review	/gemini review	Performs a code review for the current pull request in its current state.
Pull Request Summary	/gemini summary	Provides a summary of the current pull request in its current state.
Comment	@gemini-code-assist	Responds in comments when explicitly tagged, both in pull request comments and review comments.
Help	/gemini help	Displays a list of available commands.

Customization

To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here.

You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.

Footnotes

1. Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

[gemini-code-assist]

Code Review

This pull request introduces a convenient feature to automatically resume training from the last checkpoint by setting resume_from_checkpoint=True. The implementation in sft_args.py correctly handles finding the latest checkpoint, including searching through versioned subdirectories. The logic is sound and the addition of a new test case in test_resume_from_checkpoint.py ensures the feature works as expected. I have one suggestion to improve the performance and robustness of directory scanning. Overall, this is a great addition.

[gemini-code-assist]

For better performance and robustness when scanning for versioned subdirectories, consider using os.scandir() instead of os.listdir(). os.scandir() is more efficient as it retrieves file type information along with the file name, avoiding extra system calls to check if an entry is a directory. I've also added error handling for the directory scanning operation, which the original code was missing.

            subdirs = []
            try:
                with os.scandir(output_dir) as it:
                    for entry in it:
                        if entry.is_dir():
                            m = re.match(r'v(\d+)', entry.name)
                            if m:
                                subdirs.append((int(m.group(1)), entry.name))
            except OSError as e:
                logger.warning(f'Could not scan directory {output_dir} for versioned subdirectories: {e}')