diff --git a/tests/test_qwen3_30B_A3B_npu.py b/tests/test_qwen3_30B_A3B_npu.py index a88197043..cfe0a932f 100644 --- a/tests/test_qwen3_30B_A3B_npu.py +++ b/tests/test_qwen3_30B_A3B_npu.py @@ -1,11 +1,13 @@ 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" @@ -13,24 +15,38 @@ 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} " diff --git a/tools/convert_hf_to_torch_dist.py b/tools/convert_hf_to_torch_dist.py index 55334bd1b..117969fa4 100644 --- a/tools/convert_hf_to_torch_dist.py +++ b/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) diff --git a/vime/utils/arguments.py b/vime/utils/arguments.py index 968bec729..0dc1145ff 100644 --- a/vime/utils/arguments.py +++ b/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(