patch for megatron core part 2

[koutianzhengwen]

patch for megatron core part 2

[gemini-code-assist]

Summary of Changes

Hello @koutianzhengwen, 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 significantly extends the Megatron-LM framework by integrating support for Huawei Atlas800TA3 NPUs and introducing robust mechanisms for heterogeneous training. It updates model configurations for Qwen2.5-VL, enhances checkpointing capabilities, and refines argument and feature management, paving the way for more diverse and efficient large language model training environments.

Highlights

• Huawei Atlas800TA3 (NPU) Integration: Extensive adaptations have been introduced to support Huawei Atlas800TA3 NPUs, including NPU-specific rotary position embedding, matmul/add operations, and a comprehensive Megatron-LM adaptation layer (megatron_adaptor.py) to bridge with NPU functionalities.
• Heterogeneous Training Support: New configurations and logic have been added to enable heterogeneous training, allowing for different device types and pipeline layer splits. This includes modifications to parallel state management, optimizer parameter grouping (e.g., is_vision_model_param), and communication handling for CPU-based groups.
• Qwen2.5-VL Model Configuration & Checkpointing: Configuration files for Qwen2.5-VL 3B and 7B models have been added/updated, specifying system, model, optimizer, and data parameters suitable for training. A dedicated script for converting Qwen2.5-VL HuggingFace checkpoints to Megatron-Core format (and vice-versa) is also included, supporting vision models and tensor/pipeline parallelism.
• DualPipeV Pipeline Scheduling: The use_dualpipev feature has been integrated into pipeline scheduling, allowing for specialized handling of model chunks and mixed-precision operations within the pipeline, particularly for MoE models.
• Enhanced Argument and Feature Management: A MindSpeedFeaturesManager and associated patching utilities have been introduced to dynamically manage and apply various features and arguments, providing a more flexible and extensible way to configure Megatron-LM.
• Tokenizer and Dataset Expansions: New tokenizer types (AquilaTokenizerFS, HFTokenizerFS, Llama3TokenizerFS, QwenTokenizerFS, Qwen2TokenizerFS, Qwen2VLTokenizer, RWKVTokenizer) and dataset utilities (e.g., for Aquila models and LLaVA conversion) have been added, broadening the range of supported models and data formats.

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 large set of patches for Megatron-LM, primarily aimed at enabling and optimizing training on Huawei Atlas800TA3 hardware. The changes span configuration files, model adaptation layers, custom optimizers, and hardware-specific communication logic. While the patch adds significant new functionality, there are several critical issues that need to be addressed, including hardcoded device types that will cause crashes on non-CUDA platforms, potential NameError and IndexError exceptions, and a security vulnerability related to torch.load. Additionally, there are maintainability concerns like hardcoded paths in configuration files and the inclusion of backup files in the repository. A number of unit tests have also been disabled, which is a significant concern for code quality and should be rectified.

[gemini-code-assist]

The function _rotate_half is called here but it is not defined or imported within this file. This will result in a NameError at runtime when args.use_fused_rotary_pos_emb is False. Please ensure this function is available in the scope, for example by importing it from its definition location.

[gemini-code-assist]

The device for the optimizer's step tensor is hardcoded to cuda. This will cause a crash on non-CUDA hardware like the target NPU platform. The device should be dynamically inferred from the model parameters to ensure portability.

                if group['step'].device != group['params'][0].device:
                    group['step'] = group['step'].to(group['params'][0].device)
            else:
                group['step'] = torch.tensor(1, dtype=torch.int64, device=group['params'][0].device)

[gemini-code-assist]

Accessing actual_seq_len[0] without checking if the list is empty will raise an IndexError if the input samples list is empty. This can crash the data loading process. It's important to add a guard clause to handle this edge case gracefully.

        if not actual_seq_len:
            return batch
        seq_len = actual_seq_len[0]
        actual_seq_len = [elem + i * seq_len for i, elem in enumerate(actual_seq_len)]
        batch['actual_seq_len'] = torch.tensor(actual_seq_len, dtype=torch.long)

[gemini-code-assist]

The data_path is hardcoded to /share/project/lizhiyu/data/pile_wikipedia_demo. Hardcoding paths makes the configuration less portable and harder to use in different environments. Consider using a placeholder or an environment variable, like ${data_path:??} which is used in other configuration files in this repository.

[gemini-code-assist]

This patch adds a .bak file to the repository. Backup files should not be committed to version control. The history of the file should be managed by Git itself. Please remove this backup file to keep the repository clean.

[gemini-code-assist]

The data_path is hardcoded to /home/dataset/pile_wikipedia_demo. Hardcoding paths makes the configuration less portable and harder to use in different environments. It is recommended to use a placeholder like ${data_path:??} to allow for easier configuration in different setups.

[CLAassistant]

Thank you for your submission! We really appreciate it. Like many open source projects, we ask that you sign our Contributor License Agreement before we can accept your contribution.

---

zhaoxuncheng seems not to be a GitHub user. You need a GitHub account to be able to sign the CLA. If you have already a GitHub account, please add the email address used for this commit to your account.
_{You have signed the CLA already but the status is still pending? Let us recheck it.}