[NPU] Fix load torch_dist weight error - #396
Conversation
![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 refactors NPU-specific logic. In tools/convert_hf_to_torch_dist.py, it conditionally imports mindspeed.megatron_adaptor and updates the distributed process group initialization depending on whether an NPU is detected. In vime/utils/arguments.py, the default value for --megatron-to-hf-mode is simplified to always default to "raw". Feedback on these changes highlights that calling is_npu() at the module level in tools/convert_hf_to_torch_dist.py can raise a RuntimeError if no NPU device is available, and suggests wrapping the call in a try-except block to prevent import-time crashes.
| from vime.utils.common import is_npu | ||
| if is_npu(): | ||
| import mindspeed.megatron_adaptor |
There was a problem hiding this comment.
Calling is_npu() at the module level can raise a RuntimeError if torch_npu is installed but no NPU device is visible or available (for example, on a CPU-only login/compile node or during testing/building). This will cause the script to crash immediately upon import or when running basic commands like --help.
To prevent this, wrap the module-level is_npu() call in a try-except block to handle the RuntimeError gracefully.
| from vime.utils.common import is_npu | |
| if is_npu(): | |
| import mindspeed.megatron_adaptor | |
| from vime.utils.common import is_npu | |
| _is_npu = False | |
| try: | |
| _is_npu = is_npu() | |
| except RuntimeError: | |
| pass | |
| if _is_npu: | |
| import mindspeed.megatron_adaptor |
Documentation build overview
47 files changed ·
|
Signed-off-by: flb_ <floatlibai@gmail.com>
445b688 to
b93ff56
Compare
Signed-off-by: flb_ <floatlibai@gmail.com>
1 participant
This PR resolves #380 by implementing the following changes:
tools/convert_hf_to_torch_dist.pyscript.default_megatron_to_hf_modeback toraw.