[megatron] fix async save

[Jintao-Huang]

No description provided.

[gemini-code-assist]

Summary of Changes

Hello @Jintao-Huang, 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 addresses an issue with asynchronous checkpoint saving within the Megatron framework by centralizing and enhancing the checkpoint saving utility. It improves the flexibility and correctness of the save_mcore_checkpoint function, allowing callers to precisely manage output directories and asynchronous operations, particularly for merged checkpoints. This ensures a more robust and consistent checkpointing process.

Highlights

• Checkpoint Saving Refactor: Refactored the checkpoint merging logic in base.py to delegate the saving process to the save_mcore_checkpoint utility, ensuring consistent handling of merged checkpoints.
• Enhanced save_mcore_checkpoint Flexibility: Modified the save_mcore_checkpoint function to accept explicit output_dir and async_save parameters, providing greater control over checkpoint saving behavior.
• Loss Reporting Cleanup: Removed an unnecessary tuple conversion for reporting_loss in trainer.py, streamlining loss handling.

🧠 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/megatron/trainers/base.py
  • Replaced manual copying of latest_checkpointed_iteration.txt and common.pt with a call to save_mcore_checkpoint for merged checkpoints.
  • Ensured args.json is copied to the merged output directory.
• swift/megatron/trainers/trainer.py
  • Removed a redundant tuple conversion for the reporting_loss variable after an all-reduce operation.
• swift/megatron/utils/megatron_lm_utils.py
  • Added output_dir and async_save as optional parameters to the save_mcore_checkpoint function, allowing them to override values from args.
  • Updated the calculation of checkpoint_dir and the async_sharded_save parameter to use the new output_dir and async_save parameters.
  • Modified the finalization function (iter_finalize_fn) to use the provided output_dir for writing the latest_checkpointed_iteration.txt and for logging messages.

Activity

• No specific activity has been recorded for this pull request yet.

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 aims to fix an issue with asynchronous saving in Megatron. The changes primarily involve refactoring the save_mcore_checkpoint function to accept output_dir and async_save parameters, allowing for more flexible checkpointing logic. This is then used in BaseMegatronTrainer to handle merged LoRA checkpoints. However, I've found a critical bug in swift/megatron/trainers/base.py where a variable src_path is used without being defined, which will cause a NameError during execution. I've provided a suggestion to fix this.

[Jintao-Huang]

/gemini review

[gemini-code-assist]

Code Review

This pull request fixes issues with asynchronous checkpoint saving in Megatron. The main change involves updating the latest_checkpointed_iteration.txt tracker file synchronously when a save is scheduled, rather than when it's finalized asynchronously. This prevents losing track of the latest checkpoint if a crash occurs. The PR also introduces a more robust async call management system with a persistent worker option, and refactors the training loop to correctly use these new utilities. Additionally, some redundant code in logging callbacks has been cleaned up. The changes are well-structured and address the async saving problem effectively. I have one minor suggestion to correct a misleading docstring in a new utility function.

[gemini-code-assist]

The docstring for the return value is incorrect and contradicts the function's implementation and name. The function returns True when there are no ongoing async calls (i.e., the queue is empty), but the docstring states the opposite. This can be misleading for developers using this function.

Suggested change

	bool: True if there is any ongoing async call.
	bool: True if there are no ongoing async calls, False otherwise.