fix ref-load torch_dist error

Signed-off-by: flb_ <floatlibai@gmail.com>

Author: floatlibai

tools/convert_hf_to_torch_dist.py

@@ -4,6 +4,10 @@
 
 import torch
 import torch.distributed as dist
+from vime.utils.common import is_npu
+
+if is_npu():
+    import mindspeed.megatron_adaptor  # noqa: F401
 from megatron.core.enums import ModelType
 from megatron.training.arguments import parse_args, validate_args
 from megatron.training.checkpointing import get_checkpoint_name, get_checkpoint_tracker_filename, save_checkpoint
@@ -14,7 +18,6 @@
 from vime.backends.megatron_utils.arguments import set_default_megatron_args
 from vime.backends.megatron_utils.initialize import init
 from vime.backends.megatron_utils.model_provider import get_model_provider_func
-from vime.utils.common import is_npu
 from vime.utils.logging_utils import configure_logger
 from vime.utils.memory_utils import print_memory
 
@@ -92,15 +95,19 @@ def main():
     os.environ.setdefault("LOCAL_RANK", str(local_rank))
     os.environ.setdefault("MASTER_ADDR", "localhost")
     os.environ.setdefault("MASTER_PORT", "12355")
-    backend = "nccl"
     if is_npu():
-        backend = "hccl"
-    dist.init_process_group(
-        backend=backend,
-        world_size=world_size,
-        rank=global_rank,
-        device_id=torch.device(f"cuda:{local_rank}"),
-    )
+        dist.init_process_group(
+            backend="hccl",
+            world_size=world_size,
+            rank=global_rank,
+        )
+    else:
+        dist.init_process_group(
+            backend="nccl",
+            world_size=world_size,
+            rank=global_rank,
+            device_id=torch.device(f"cuda:{local_rank}"),
+        )
     args = get_args()
     init(args)
 

vime/utils/arguments.py

@@ -10,7 +10,6 @@
 
 from vime.backends.vllm_utils.arguments import validate_args as vllm_validate_args
 from vime.backends.vllm_utils.arguments import vllm_parse_args
-from vime.utils.common import is_npu
 from vime.utils.eval_config import EvalDatasetConfig, build_eval_dataset_configs, ensure_dataset_list
 from vime.utils.logging_utils import configure_logger
 
@@ -132,14 +131,10 @@ def add_train_arguments(parser):
                 default=1024**3,
                 help="Add margin for train memory allocation. By default we will reserve 1GB as margin.",
             )
-            try:
-                default_megatron_to_hf_mode = "bridge" if is_npu() else "raw"
-            except RuntimeError:
-                default_megatron_to_hf_mode = "raw"
             parser.add_argument(
                 "--megatron-to-hf-mode",
                 choices=["raw", "bridge"],
-                default=default_megatron_to_hf_mode,
+                default="raw",
                 help="The method to convert megatron weights to hugging face weights for vLLM.",
             )
             parser.add_argument(