[NPU] Fix load torch_dist weight error

tests/test_qwen3_30B_A3B_npu.py

@@ -1,36 +1,52 @@
 import os
 import shlex
+from pathlib import Path
 
 import vime.utils.external_utils.command_utils as U
 
 
 TEST_ROOT = os.environ.get("HF_HOME") or "/root"
 MODEL_DIR = f"{TEST_ROOT}/models/Qwen3-30B-A3B"
+CHECKPOINT_DIR = f"{TEST_ROOT}/models/Qwen3-30B-A3B_torch_dist"
 DATASET_DIR = f"{TEST_ROOT}/datasets/dapo-math-17k"
 
 
 def prepare():
     models_dir = shlex.quote(f"{TEST_ROOT}/models")
     datasets_dir = shlex.quote(f"{TEST_ROOT}/datasets")
     model_dir = shlex.quote(MODEL_DIR)
+    checkpoint_dir = shlex.quote(CHECKPOINT_DIR)
     dataset_dir = shlex.quote(DATASET_DIR)
 
     U.exec_command(f"mkdir -p {models_dir} {datasets_dir}")
     U.exec_command(f"hf download Qwen/Qwen3-30B-A3B --local-dir {model_dir}")
     U.exec_command("hf download --repo-type dataset zhuzilin/dapo-math-17k " f"--local-dir {dataset_dir}")
+    U.exec_command(f"rm -rf {checkpoint_dir}")
+    U.exec_command(
+        "source scripts/models/qwen3-30B-A3B.sh && "
+        "PYTHONPATH=/root/Megatron-LM "
+        f"torchrun --nproc-per-node 8 tools/convert_hf_to_torch_dist.py "
+        "${MODEL_ARGS[@]} "
+        f"--hf-checkpoint {model_dir} --save {checkpoint_dir}"
+    )
+
+    checkpoint_path = Path(CHECKPOINT_DIR)
+    tracker = checkpoint_path / "latest_checkpointed_iteration.txt"
+    assert tracker.read_text().strip() == "release"
+    weight_files = [
+        path
+        for path in checkpoint_path.rglob("*")
+        if path.is_file() and path.name != "latest_checkpointed_iteration.txt"
+    ]
+    assert weight_files, f"No checkpoint weights found under {checkpoint_path}"
 
 
 def execute():
     model_dir = shlex.quote(MODEL_DIR)
+    checkpoint_dir = shlex.quote(CHECKPOINT_DIR)
     prompt_data = shlex.quote(f"{DATASET_DIR}/dapo-math-17k.jsonl")
 
-    checkpoint_args = (
-        f"--hf-checkpoint {model_dir} "
-        f"--load {model_dir} "
-        f"--ref-load {model_dir} "
-        "--megatron-to-hf-mode bridge "
-        "--no-load-optim "
-    )
+    checkpoint_args = f"--hf-checkpoint {model_dir} " f"--ref-load {checkpoint_dir} " "--no-load-optim "
 
     rollout_args = (
         f"--prompt-data {prompt_data} "

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 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

@@ -132,14 +132,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(