[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')

[Jintao-Huang]

/gemini review

[gemini-code-assist]

Code Review

This pull request significantly improves the user experience for Megatron-SWIFT by updating the workflow to support direct loading and saving of safetensors weights. This removes the need for a separate, explicit weight conversion step. The changes are comprehensive, touching documentation, example scripts, and core Python logic to support this new, streamlined process. The documentation in both English and Chinese has been updated to reflect these changes, and numerous example shell scripts have been modified to use the new arguments. The core logic for argument parsing, model loading, and checkpointing has been refactored to accommodate the new workflow. Overall, this is a great enhancement. I've found one minor typo in an example script that I've commented on.

[Jintao-Huang]

/gemini review

[gemini-code-assist]

Code Review

This pull request significantly updates the Megatron integration by streamlining the workflow for loading and saving models. It transitions from a manual weight conversion process to directly using safetensors via Mcore-Bridge, which is a great improvement for user experience. The documentation and example scripts have been extensively updated to reflect this new recommended approach, while still maintaining the traditional method for backward compatibility. The core logic for handling checkpoints and arguments has been refactored accordingly. My review includes a critical bug fix for resuming training from safetensors checkpoints and a minor suggestion to improve code clarity.

[Jintao-Huang]

/gemini review

[gemini-code-assist]

Code Review

This pull request introduces a significant and welcome update to the Megatron-SWIFT workflow by better integrating Mcore-Bridge. The changes streamline the training process by enabling direct loading and saving of safetensors weights, which simplifies the user experience by removing manual weight conversion steps. The documentation and example scripts have been thoroughly updated to reflect the new recommended workflow, distinguishing it from the traditional method. The core logic changes, especially around checkpointing and model loading, appear correct and robustly support features like resuming training from safetensors checkpoints. I have one minor suggestion for improving argument parsing consistency. Overall, this is an excellent enhancement that makes Megatron-SWIFT more user-friendly.

[gemini-code-assist]

For consistency and robustness, it would be good to also handle the case where self.adapters is a string, similar to how self.ref_adapters is handled. This ensures that if a single adapter path is passed as a string (e.g., from a config file), it is correctly converted to a list.

Suggested change

	if isinstance(self.ref_adapters, str):
	self.ref_adapters = [self.ref_adapters]
	if isinstance(self.adapters, str):
	self.adapters = [self.adapters]
	if isinstance(self.ref_adapters, str):
	self.ref_adapters = [self.ref_adapters]