Implement NPU_ENV handling in init.py to support megatron in NPU. #6661
Conversation
Add NPU_ENV check and repatch configuration for Megatron.
Summary of ChangesHello @vx120, I'm Gemini Code Assist1! 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 introduces support for NPU (Neural Processing Unit) environments within the Megatron framework. It achieves this by adding NPU environment detection and conditionally applying a specialized Highlights
Using Gemini Code AssistThe 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
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 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
|
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces support for NPU environments by conditionally importing mindspeed and applying necessary patches. The logic is sound, but I have a couple of suggestions to enhance code quality and maintainability. Specifically, I recommend using the logger for warning messages instead of print, and defining a static configuration dictionary as a module-level constant.
| override_transformer_config = { | ||
| "recompute_granularity": None, | ||
| "recompute_modules": ["core_attn"], | ||
| "recompute_method": None, | ||
| "recompute_num_layers": None | ||
| } |
There was a problem hiding this comment.
The override_transformer_config dictionary is static. It's more efficient and readable to define it as a constant at the module level, outside of the _patch_megatron function.
For example:
# At the top of the file, e.g., after the logger is defined.
_NPU_OVERRIDE_TRANSFORMER_CONFIG = {
"recompute_granularity": None,
"recompute_modules": ["core_attn"],
"recompute_method": None,
"recompute_num_layers": None,
}
# ...
def _patch_megatron():
if NPU_ENV:
if repatch is not None:
repatch(_NPU_OVERRIDE_TRANSFORMER_CONFIG)
# ...| else: | ||
| print('repatch没有正确导入') |
There was a problem hiding this comment.
It's better to use the logger for warnings or errors instead of print(). This ensures consistent logging behavior. Also, using an English message would improve maintainability and consistency with the rest of the codebase.
| else: | |
| print('repatch没有正确导入') | |
| else: | |
| logger.warning('`repatch` from `mindspeed.megatron_adaptor` was not imported correctly, skipping NPU patch.') |
| if repatch is not None: | ||
| repatch(override_transformer_config) | ||
| else: | ||
| print('repatch没有正确导入') |
There was a problem hiding this comment.
Hello. Please use English comments. 😊
| from swift.llm import git_clone_github | ||
| from swift.utils import (get_logger, is_flash_attn_3_available, is_megatron_available, safe_ddp_context, split_list, | ||
| subprocess_run) | ||
| try: |
There was a problem hiding this comment.
这里加一个判断吧，安全一些
from transformers.utils import is_torch_npu_available
|
好的
发自我的iPhone
…------------------ Original ------------------
From: Jintao ***@***.***>
Date: Wed,Nov 19,2025 9:57 PM
To: modelscope/ms-swift ***@***.***>
Cc: vx120 ***@***.***>, Mention ***@***.***>
Subject: Re: [modelscope/ms-swift] Implement NPU_ENV handling in init.py tosupport megatron in NPU. (PR #6661)
@Jintao-Huang commented on this pull request.
In swift/megatron/init.py:
> @@ -18,6 +18,11 @@ from swift.llm import git_clone_github from swift.utils import (get_logger, is_flash_attn_3_available, is_megatron_available, safe_ddp_context, split_list, subprocess_run) +try:
这里加一个判断吧，安全一些
from transformers.utils import is_torch_npu_available
—
Reply to this email directly, view it on GitHub, or unsubscribe.
You are receiving this because you were mentioned.Message ID: ***@***.***>
|
|
请运行以下代码（<py312），过一下lint测试 然后建议在文档中简单加一下NPU的内容。Megatron-SWIFT的快速开始文档。 |
2 participants
PR type
PR information
Add NPU_ENV check and repatch configuration for Megatron in NPU.