[megatron] Update megatron shells

[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 focuses on refining the Megatron integration within the ms-swift framework by updating documentation and example configurations. The primary goal is to streamline model handling, particularly by transitioning towards Safetensors for loading and saving model weights, which offers improved security and efficiency. The changes also include necessary adjustments to argument parsing and checkpoint saving mechanisms to fully support this new approach, enhancing the overall robustness of large model training workflows.

Highlights

• Documentation Update: Simplified environment setup instructions in both Chinese and English best practices documents by removing specific mentions of HF and MCore model weight conversion, streamlining the guidance for users.
• Safetensors Adoption: Updated example training scripts in the documentation to explicitly use --model for specifying models and added --load_safetensors true --save_safetensors true, promoting the use of the Safetensors format for model loading and saving.
• Megatron Argument Refinement: Adjusted the default behavior of load_safetensors and save_safetensors in ExtraMegatronArguments, making save_safetensors default to True and introducing logic to automatically set load_safetensors if not explicitly provided.
• Checkpoint Saving Enhancement: Modified the save_checkpoint method in the base trainer to explicitly accept the model and introduced conditional logic to handle the saving process, including temporary adjustment of args.save and options to skip saving optimizer or RNG states.
• Robust State Dictionary Generation: Added a check in generate_state_dict to ensure robustness by gracefully handling cases where the state dictionary might not contain a 'model' key, preventing potential errors during state dictionary processing.

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 updates Megatron-related shell scripts and underlying Python code. The changes primarily focus on improving the model loading and saving mechanisms, particularly by introducing and defaulting to safetensors.

Key changes include:

• Updating documentation and example scripts to use --model instead of --load and enabling safetensors by default.
• Modifying MegatronArguments to support the new safetensors logic with smart defaults.
• Refactoring the save_checkpoint logic in MegatronTrainer to handle safetensors for model weights while still using Megatron's checkpointing for optimizer and RNG states.

My main feedback is on the save_checkpoint implementation, where there appears to be a bug that could lead to redundant saving of model weights. I've left a detailed comment on this. Additionally, a defensive check has been added in utils.py which is a good improvement.

[gemini-code-assist]

This logic appears to have a bug. When save_safetensors is true, model weights are already saved via self.bridge.save_weights. This block is intended to save only the optimizer and/or RNG states. However, if args.no_save_optim is False, the original model object is passed to _origin_save_checkpoint, which causes the model weights to be saved again in the Megatron checkpoint format. This is redundant and wastes storage.

To prevent this, the model variable should be set to an empty list ([]) unconditionally before calling _origin_save_checkpoint to signal that it should skip saving the model's state dictionary. The optimizer state can still be saved correctly as the optimizer object holds its own reference to the model if needed (e.g., for distributed optimizers).

[Jintao-Huang]

/gemini review

[gemini-code-assist]

Code Review

This pull request introduces significant updates to the Megatron integration, primarily focused on standardizing model loading via safetensors and improving checkpointing for resuming training. The changes are extensive, touching many example shell scripts to use --model instead of --load, and updating documentation to reflect the new practices. The core logic for saving and loading checkpoints has been refactored to support these new workflows, including better handling of LoRA adapters and optimizer states. While the changes are generally positive and well-aligned with the goal, I've identified a critical issue in the checkpoint loading logic that could prevent training resumption, along with a few minor inconsistencies and potential typos in the shell scripts and documentation. Addressing these points will ensure the new functionality is robust and reliable.

[gemini-code-assist]

The iteration number is hardcoded as 5 here. This will cause the function to always look for the checkpoint of iteration 5, instead of the latest one specified in latest_checkpointed_iteration.txt. This will prevent training from resuming correctly. The iteration variable should be used instead.

Suggested change

	common_path = os.path.join(ckpt_dir, f'iter_{5:07d}', 'common.pt')
	common_path = os.path.join(ckpt_dir, f'iter_{int(iteration):07d}', 'common.pt')