From 37fe7a5adb91f7fba2cd3e5b2a96da0f6d28c3c9 Mon Sep 17 00:00:00 2001 From: flb_ Date: Sat, 22 Aug 2026 16:40:36 +0800 Subject: [PATCH 1/9] fix ref-load torch_dist error Signed-off-by: flb_ --- tools/convert_hf_to_torch_dist.py | 25 ++++++++++++++++--------- vime/utils/arguments.py | 7 +------ 2 files changed, 17 insertions(+), 15 deletions(-) diff --git a/tools/convert_hf_to_torch_dist.py b/tools/convert_hf_to_torch_dist.py index 55334bd1b..afed0b3f1 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 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) diff --git a/vime/utils/arguments.py b/vime/utils/arguments.py index 968bec729..12aa12be2 100644 --- a/vime/utils/arguments.py +++ b/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( From 8f44e6edffc2138da3b954ecc23a4268bbf80542 Mon Sep 17 00:00:00 2001 From: flb_ Date: Thu, 27 Aug 2026 20:41:49 +0800 Subject: [PATCH 2/9] change to megatron adaptor Signed-off-by: flb_ --- tools/convert_hf_to_torch_dist.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tools/convert_hf_to_torch_dist.py b/tools/convert_hf_to_torch_dist.py index afed0b3f1..117969fa4 100644 --- a/tools/convert_hf_to_torch_dist.py +++ b/tools/convert_hf_to_torch_dist.py @@ -7,7 +7,7 @@ from vime.utils.common import is_npu if is_npu(): - import mindspeed.megatron_adaptor # noqa: F401 + 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 From 33ec930456a5f8c63575e04e36a233d17be99449 Mon Sep 17 00:00:00 2001 From: flb_ Date: Wed, 2 Sep 2026 15:22:06 +0800 Subject: [PATCH 3/9] add testcases to verify args and convertion Signed-off-by: flb_ --- .buildkite/npu_suites.py | 1 + .buildkite/pipeline.yml | 1 + tests/test_convert_hf_to_torch_dist_args.py | 53 ++++++++++++++++++++ tests/test_convert_hf_to_torch_dist_npu.py | 55 +++++++++++++++++++++ 4 files changed, 110 insertions(+) create mode 100644 tests/test_convert_hf_to_torch_dist_args.py create mode 100644 tests/test_convert_hf_to_torch_dist_npu.py diff --git a/.buildkite/npu_suites.py b/.buildkite/npu_suites.py index 647a7a1f7..6341c721c 100644 --- a/.buildkite/npu_suites.py +++ b/.buildkite/npu_suites.py @@ -28,6 +28,7 @@ # (test_name, resource_class, extra_args, env_overrides) SUITES = { "smk": [ + ("test_convert_hf_to_torch_dist_npu.py", "npu-8", "", {}), ("test_qwen3_4B_npu.py", "npu-8", "", {}), ("test_qwen3_30B_A3B_npu.py", "npu-16", "", {}), ("test_qwen3_vl_8B_npu.py", "npu-8", "", {}), diff --git a/.buildkite/pipeline.yml b/.buildkite/pipeline.yml index 5ae2e5900..73ab81da9 100644 --- a/.buildkite/pipeline.yml +++ b/.buildkite/pipeline.yml @@ -61,6 +61,7 @@ steps: pip install -q pytest numpy packaging pyyaml omegaconf tqdm httpx requests ray pybase64 pylatexenc sympy aiohttp pillow safetensors transformers cloudpickle pip install -q -e . --no-deps python tests/test_megatron_argument_validation.py + python tests/test_convert_hf_to_torch_dist_args.py python tests/test_value_temperature.py python tests/test_rollout_validation.py python tests/plugin_contracts/test_plugin_rollout_contracts.py diff --git a/tests/test_convert_hf_to_torch_dist_args.py b/tests/test_convert_hf_to_torch_dist_args.py new file mode 100644 index 000000000..3b553e261 --- /dev/null +++ b/tests/test_convert_hf_to_torch_dist_args.py @@ -0,0 +1,53 @@ +"""Unit tests for arguments specific to ``convert_hf_to_torch_dist.py``.""" + +from __future__ import annotations + +import argparse +import ast +from pathlib import Path + +import pytest + + +NUM_GPUS = 0 + + +@pytest.mark.unit +def test_megatron_to_hf_mode_defaults_to_raw(): + """The conversion mode must remain explicit and deterministic by default.""" + root = Path(__file__).resolve().parents[1] + source = root / "tools" / "convert_hf_to_torch_dist.py" + tree = ast.parse(source.read_text()) + for node in ast.walk(tree): + if not isinstance(node, ast.FunctionDef) or node.name != "add_convertion_args": + continue + namespace: dict[str, object] = {} + # The argument-registration function is self-contained. Executing only + # this function keeps the test CPU-only while exercising argparse itself. + exec(compile(ast.Module(body=[node], type_ignores=[]), str(source), "exec"), namespace) + parser = argparse.ArgumentParser() + namespace["add_convertion_args"](parser) + args = parser.parse_args(["--hf-checkpoint", "/tmp/model"]) + assert args.megatron_to_hf_mode == "raw" + + # The same option is exposed by the training argument provider. Keep + # that entry point aligned with the standalone converter as well. + training_source = root / "vime" / "utils" / "arguments.py" + training_tree = ast.parse(training_source.read_text()) + training_defaults = [] + for call in ast.walk(training_tree): + if not isinstance(call, ast.Call) or not isinstance(call.func, ast.Attribute): + continue + if call.func.attr != "add_argument" or not call.args: + continue + if isinstance(call.args[0], ast.Constant) and call.args[0].value == "--megatron-to-hf-mode": + training_defaults.extend(kw.value for kw in call.keywords if kw.arg == "default") + assert len(training_defaults) == 1 + assert isinstance(training_defaults[0], ast.Constant) + assert training_defaults[0].value == "raw" + return + pytest.fail("--megatron-to-hf-mode argument was not found") + + +if __name__ == "__main__": + raise SystemExit(pytest.main([__file__])) diff --git a/tests/test_convert_hf_to_torch_dist_npu.py b/tests/test_convert_hf_to_torch_dist_npu.py new file mode 100644 index 000000000..3c04c6183 --- /dev/null +++ b/tests/test_convert_hf_to_torch_dist_npu.py @@ -0,0 +1,55 @@ +"""NPU smoke test for HuggingFace-to-Megatron checkpoint conversion.""" + +from __future__ import annotations + +import os +import shlex +import shutil +from pathlib import Path + +import vime.utils.external_utils.command_utils as U + + +TEST_ROOT = Path(os.environ.get("HF_HOME") or "/root") +MODEL_NAME = "Qwen3-0.6B" +NUM_GPUS = 8 +MODEL_DIR = TEST_ROOT / "models" / MODEL_NAME +CHECKPOINT_DIR = TEST_ROOT / "models" / f"{MODEL_NAME}_torch_dist_npu_test" + + +def prepare(): + """Download the small smoke-test model and reset its output directory.""" + U.exec_command(f"mkdir -p {shlex.quote(str(MODEL_DIR.parent))}") + U.exec_command(f"hf download Qwen/{MODEL_NAME} --local-dir {shlex.quote(str(MODEL_DIR))}") + shutil.rmtree(CHECKPOINT_DIR, ignore_errors=True) + + +def execute(): + """Run conversion on all NPUs and verify a release checkpoint was written.""" + model_dir = shlex.quote(str(MODEL_DIR)) + checkpoint_dir = shlex.quote(str(CHECKPOINT_DIR)) + U.exec_command( + "source scripts/models/qwen3-0.6B.sh && " + "PYTHONPATH=/root/Megatron-LM " + f"torchrun --nproc-per-node {NUM_GPUS} tools/convert_hf_to_torch_dist.py " + "${MODEL_ARGS[@]} " + f"--hf-checkpoint {model_dir} --save {checkpoint_dir}" + ) + + tracker = CHECKPOINT_DIR / "latest_checkpointed_iteration.txt" + assert tracker.read_text().strip() == "release" + weight_files = [ + path + for path in CHECKPOINT_DIR.rglob("*") + if path.is_file() and path.name != "latest_checkpointed_iteration.txt" + ] + assert weight_files, f"No checkpoint weights found under {CHECKPOINT_DIR}" + + +def main(): + prepare() + execute() + + +if __name__ == "__main__": + main() From dfb2c3c126f8993e1af720f944861a05e834cae8 Mon Sep 17 00:00:00 2001 From: flb_ Date: Wed, 2 Sep 2026 15:36:54 +0800 Subject: [PATCH 4/9] fix pre-commit Signed-off-by: flb_ --- vime/utils/arguments.py | 1 + 1 file changed, 1 insertion(+) diff --git a/vime/utils/arguments.py b/vime/utils/arguments.py index 12aa12be2..0dc1145ff 100644 --- a/vime/utils/arguments.py +++ b/vime/utils/arguments.py @@ -10,6 +10,7 @@ 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 From 60c24e91a15fbefcb05ded89d7519bbc9b5b59a9 Mon Sep 17 00:00:00 2001 From: flb_ Date: Wed, 2 Sep 2026 17:03:53 +0800 Subject: [PATCH 5/9] combine to one testcase Signed-off-by: flb_ --- .buildkite/pipeline.yml | 1 - tests/test_convert_hf_to_torch_dist_args.py | 53 ------------------- tests/test_convert_hf_to_torch_dist_npu.py | 58 +++++++++++++++++++++ 3 files changed, 58 insertions(+), 54 deletions(-) delete mode 100644 tests/test_convert_hf_to_torch_dist_args.py diff --git a/.buildkite/pipeline.yml b/.buildkite/pipeline.yml index 73ab81da9..5ae2e5900 100644 --- a/.buildkite/pipeline.yml +++ b/.buildkite/pipeline.yml @@ -61,7 +61,6 @@ steps: pip install -q pytest numpy packaging pyyaml omegaconf tqdm httpx requests ray pybase64 pylatexenc sympy aiohttp pillow safetensors transformers cloudpickle pip install -q -e . --no-deps python tests/test_megatron_argument_validation.py - python tests/test_convert_hf_to_torch_dist_args.py python tests/test_value_temperature.py python tests/test_rollout_validation.py python tests/plugin_contracts/test_plugin_rollout_contracts.py diff --git a/tests/test_convert_hf_to_torch_dist_args.py b/tests/test_convert_hf_to_torch_dist_args.py deleted file mode 100644 index 3b553e261..000000000 --- a/tests/test_convert_hf_to_torch_dist_args.py +++ /dev/null @@ -1,53 +0,0 @@ -"""Unit tests for arguments specific to ``convert_hf_to_torch_dist.py``.""" - -from __future__ import annotations - -import argparse -import ast -from pathlib import Path - -import pytest - - -NUM_GPUS = 0 - - -@pytest.mark.unit -def test_megatron_to_hf_mode_defaults_to_raw(): - """The conversion mode must remain explicit and deterministic by default.""" - root = Path(__file__).resolve().parents[1] - source = root / "tools" / "convert_hf_to_torch_dist.py" - tree = ast.parse(source.read_text()) - for node in ast.walk(tree): - if not isinstance(node, ast.FunctionDef) or node.name != "add_convertion_args": - continue - namespace: dict[str, object] = {} - # The argument-registration function is self-contained. Executing only - # this function keeps the test CPU-only while exercising argparse itself. - exec(compile(ast.Module(body=[node], type_ignores=[]), str(source), "exec"), namespace) - parser = argparse.ArgumentParser() - namespace["add_convertion_args"](parser) - args = parser.parse_args(["--hf-checkpoint", "/tmp/model"]) - assert args.megatron_to_hf_mode == "raw" - - # The same option is exposed by the training argument provider. Keep - # that entry point aligned with the standalone converter as well. - training_source = root / "vime" / "utils" / "arguments.py" - training_tree = ast.parse(training_source.read_text()) - training_defaults = [] - for call in ast.walk(training_tree): - if not isinstance(call, ast.Call) or not isinstance(call.func, ast.Attribute): - continue - if call.func.attr != "add_argument" or not call.args: - continue - if isinstance(call.args[0], ast.Constant) and call.args[0].value == "--megatron-to-hf-mode": - training_defaults.extend(kw.value for kw in call.keywords if kw.arg == "default") - assert len(training_defaults) == 1 - assert isinstance(training_defaults[0], ast.Constant) - assert training_defaults[0].value == "raw" - return - pytest.fail("--megatron-to-hf-mode argument was not found") - - -if __name__ == "__main__": - raise SystemExit(pytest.main([__file__])) diff --git a/tests/test_convert_hf_to_torch_dist_npu.py b/tests/test_convert_hf_to_torch_dist_npu.py index 3c04c6183..ef7143b41 100644 --- a/tests/test_convert_hf_to_torch_dist_npu.py +++ b/tests/test_convert_hf_to_torch_dist_npu.py @@ -2,6 +2,8 @@ from __future__ import annotations +import argparse +import ast import os import shlex import shutil @@ -16,9 +18,65 @@ MODEL_DIR = TEST_ROOT / "models" / MODEL_NAME CHECKPOINT_DIR = TEST_ROOT / "models" / f"{MODEL_NAME}_torch_dist_npu_test" +# Root of the repo, used to locate the converter/training arg sources for the +# pre-flight static check below. +REPO_ROOT = Path(__file__).resolve().parents[1] + + +def _assert_megatron_to_hf_mode_defaults_to_raw() -> None: + """Fail fast, before spending NPU time, if the mode default drifted. + + Mirrors the standalone unit test: parses ``add_convertion_args`` out of + the converter script via AST (kept CPU-only / import-free) and confirms + the training argument provider defines the same default. This is a + pre-flight guard for :func:`execute` — it must pass before we burn + 8-NPU time on a distributed conversion that would be misconfigured + anyway. + """ + source = REPO_ROOT / "tools" / "convert_hf_to_torch_dist.py" + tree = ast.parse(source.read_text()) + for node in ast.walk(tree): + if not isinstance(node, ast.FunctionDef) or node.name != "add_convertion_args": + continue + namespace: dict[str, object] = {} + exec(compile(ast.Module(body=[node], type_ignores=[]), str(source), "exec"), namespace) + parser = argparse.ArgumentParser() + namespace["add_convertion_args"](parser) + args = parser.parse_args(["--hf-checkpoint", "/tmp/model"]) + assert args.megatron_to_hf_mode == "raw", ( + "convert_hf_to_torch_dist.py: --megatron-to-hf-mode default drifted " + f"from 'raw' to {args.megatron_to_hf_mode!r}" + ) + + training_source = REPO_ROOT / "vime" / "utils" / "arguments.py" + training_tree = ast.parse(training_source.read_text()) + training_defaults = [] + for call in ast.walk(training_tree): + if not isinstance(call, ast.Call) or not isinstance(call.func, ast.Attribute): + continue + if call.func.attr != "add_argument" or not call.args: + continue + if isinstance(call.args[0], ast.Constant) and call.args[0].value == "--megatron-to-hf-mode": + training_defaults.extend(kw.value for kw in call.keywords if kw.arg == "default") + assert len(training_defaults) == 1, ( + "vime/utils/arguments.py: expected exactly one --megatron-to-hf-mode " + f"default declaration, found {len(training_defaults)}" + ) + assert isinstance(training_defaults[0], ast.Constant) + assert training_defaults[0].value == "raw", ( + "vime/utils/arguments.py: --megatron-to-hf-mode default drifted " + f"from 'raw' to {training_defaults[0].value!r}" + ) + return + raise AssertionError("--megatron-to-hf-mode argument was not found in add_convertion_args") + def prepare(): """Download the small smoke-test model and reset its output directory.""" + # Pre-flight: confirm the default conversion mode hasn't drifted before + # spending 8-NPU time on a run that would be silently misconfigured. + _assert_megatron_to_hf_mode_defaults_to_raw() + U.exec_command(f"mkdir -p {shlex.quote(str(MODEL_DIR.parent))}") U.exec_command(f"hf download Qwen/{MODEL_NAME} --local-dir {shlex.quote(str(MODEL_DIR))}") shutil.rmtree(CHECKPOINT_DIR, ignore_errors=True) From ce4dbd661f9ce16c4c719ab0f32823486770b86e Mon Sep 17 00:00:00 2001 From: flb_ Date: Wed, 2 Sep 2026 18:04:07 +0800 Subject: [PATCH 6/9] add ref-load test Signed-off-by: flb_ --- tests/test_convert_hf_to_torch_dist_npu.py | 31 ++++++++++++++++++++-- 1 file changed, 29 insertions(+), 2 deletions(-) diff --git a/tests/test_convert_hf_to_torch_dist_npu.py b/tests/test_convert_hf_to_torch_dist_npu.py index ef7143b41..7e8b8ee44 100644 --- a/tests/test_convert_hf_to_torch_dist_npu.py +++ b/tests/test_convert_hf_to_torch_dist_npu.py @@ -13,7 +13,7 @@ TEST_ROOT = Path(os.environ.get("HF_HOME") or "/root") -MODEL_NAME = "Qwen3-0.6B" +MODEL_NAME = "Qwen3-4B" NUM_GPUS = 8 MODEL_DIR = TEST_ROOT / "models" / MODEL_NAME CHECKPOINT_DIR = TEST_ROOT / "models" / f"{MODEL_NAME}_torch_dist_npu_test" @@ -87,7 +87,7 @@ def execute(): model_dir = shlex.quote(str(MODEL_DIR)) checkpoint_dir = shlex.quote(str(CHECKPOINT_DIR)) U.exec_command( - "source scripts/models/qwen3-0.6B.sh && " + "source scripts/models/qwen3-4B.sh && " "PYTHONPATH=/root/Megatron-LM " f"torchrun --nproc-per-node {NUM_GPUS} tools/convert_hf_to_torch_dist.py " "${MODEL_ARGS[@]} " @@ -103,6 +103,33 @@ def execute(): ] assert weight_files, f"No checkpoint weights found under {CHECKPOINT_DIR}" + U.execute_train( + train_args=( + f"--hf-checkpoint {model_dir} " + f"--ref-load {checkpoint_dir} " + "--debug-train-only " + "--num-rollout 1 " + "--start-rollout-id 1 " + "--no-load-optim " + "--no-load-rng " + "--rollout-batch-size 1 " + "--global-batch-size 1 " + "--kl-coef 0.1 " + "--lr-decay-iters 1 " + "--optimizer adam " + "--lr 1e-6 " + "--lr-decay-style constant " + "--weight-decay 0.0 " + "--adam-beta1 0.9 " + "--adam-beta2 0.98 " + "--actor-num-nodes 1 " + f"--actor-num-gpus-per-node {NUM_GPUS} " + "--ci-test " + ), + num_gpus_per_node=NUM_GPUS, + megatron_model_type="qwen3-4B", + ) + def main(): prepare() From 8211f82b13b01fdc7aec47dcd166a5806cd1197e Mon Sep 17 00:00:00 2001 From: flb_ Date: Thu, 3 Sep 2026 11:53:37 +0800 Subject: [PATCH 7/9] use exiting 30b case Signed-off-by: flb_ --- .buildkite/npu_suites.py | 1 - tests/test_convert_hf_to_torch_dist_npu.py | 140 --------------------- tests/test_qwen3_30B_A3B_npu.py | 26 ++-- 3 files changed, 19 insertions(+), 148 deletions(-) delete mode 100644 tests/test_convert_hf_to_torch_dist_npu.py diff --git a/.buildkite/npu_suites.py b/.buildkite/npu_suites.py index 6341c721c..647a7a1f7 100644 --- a/.buildkite/npu_suites.py +++ b/.buildkite/npu_suites.py @@ -28,7 +28,6 @@ # (test_name, resource_class, extra_args, env_overrides) SUITES = { "smk": [ - ("test_convert_hf_to_torch_dist_npu.py", "npu-8", "", {}), ("test_qwen3_4B_npu.py", "npu-8", "", {}), ("test_qwen3_30B_A3B_npu.py", "npu-16", "", {}), ("test_qwen3_vl_8B_npu.py", "npu-8", "", {}), diff --git a/tests/test_convert_hf_to_torch_dist_npu.py b/tests/test_convert_hf_to_torch_dist_npu.py deleted file mode 100644 index 7e8b8ee44..000000000 --- a/tests/test_convert_hf_to_torch_dist_npu.py +++ /dev/null @@ -1,140 +0,0 @@ -"""NPU smoke test for HuggingFace-to-Megatron checkpoint conversion.""" - -from __future__ import annotations - -import argparse -import ast -import os -import shlex -import shutil -from pathlib import Path - -import vime.utils.external_utils.command_utils as U - - -TEST_ROOT = Path(os.environ.get("HF_HOME") or "/root") -MODEL_NAME = "Qwen3-4B" -NUM_GPUS = 8 -MODEL_DIR = TEST_ROOT / "models" / MODEL_NAME -CHECKPOINT_DIR = TEST_ROOT / "models" / f"{MODEL_NAME}_torch_dist_npu_test" - -# Root of the repo, used to locate the converter/training arg sources for the -# pre-flight static check below. -REPO_ROOT = Path(__file__).resolve().parents[1] - - -def _assert_megatron_to_hf_mode_defaults_to_raw() -> None: - """Fail fast, before spending NPU time, if the mode default drifted. - - Mirrors the standalone unit test: parses ``add_convertion_args`` out of - the converter script via AST (kept CPU-only / import-free) and confirms - the training argument provider defines the same default. This is a - pre-flight guard for :func:`execute` — it must pass before we burn - 8-NPU time on a distributed conversion that would be misconfigured - anyway. - """ - source = REPO_ROOT / "tools" / "convert_hf_to_torch_dist.py" - tree = ast.parse(source.read_text()) - for node in ast.walk(tree): - if not isinstance(node, ast.FunctionDef) or node.name != "add_convertion_args": - continue - namespace: dict[str, object] = {} - exec(compile(ast.Module(body=[node], type_ignores=[]), str(source), "exec"), namespace) - parser = argparse.ArgumentParser() - namespace["add_convertion_args"](parser) - args = parser.parse_args(["--hf-checkpoint", "/tmp/model"]) - assert args.megatron_to_hf_mode == "raw", ( - "convert_hf_to_torch_dist.py: --megatron-to-hf-mode default drifted " - f"from 'raw' to {args.megatron_to_hf_mode!r}" - ) - - training_source = REPO_ROOT / "vime" / "utils" / "arguments.py" - training_tree = ast.parse(training_source.read_text()) - training_defaults = [] - for call in ast.walk(training_tree): - if not isinstance(call, ast.Call) or not isinstance(call.func, ast.Attribute): - continue - if call.func.attr != "add_argument" or not call.args: - continue - if isinstance(call.args[0], ast.Constant) and call.args[0].value == "--megatron-to-hf-mode": - training_defaults.extend(kw.value for kw in call.keywords if kw.arg == "default") - assert len(training_defaults) == 1, ( - "vime/utils/arguments.py: expected exactly one --megatron-to-hf-mode " - f"default declaration, found {len(training_defaults)}" - ) - assert isinstance(training_defaults[0], ast.Constant) - assert training_defaults[0].value == "raw", ( - "vime/utils/arguments.py: --megatron-to-hf-mode default drifted " - f"from 'raw' to {training_defaults[0].value!r}" - ) - return - raise AssertionError("--megatron-to-hf-mode argument was not found in add_convertion_args") - - -def prepare(): - """Download the small smoke-test model and reset its output directory.""" - # Pre-flight: confirm the default conversion mode hasn't drifted before - # spending 8-NPU time on a run that would be silently misconfigured. - _assert_megatron_to_hf_mode_defaults_to_raw() - - U.exec_command(f"mkdir -p {shlex.quote(str(MODEL_DIR.parent))}") - U.exec_command(f"hf download Qwen/{MODEL_NAME} --local-dir {shlex.quote(str(MODEL_DIR))}") - shutil.rmtree(CHECKPOINT_DIR, ignore_errors=True) - - -def execute(): - """Run conversion on all NPUs and verify a release checkpoint was written.""" - model_dir = shlex.quote(str(MODEL_DIR)) - checkpoint_dir = shlex.quote(str(CHECKPOINT_DIR)) - U.exec_command( - "source scripts/models/qwen3-4B.sh && " - "PYTHONPATH=/root/Megatron-LM " - f"torchrun --nproc-per-node {NUM_GPUS} tools/convert_hf_to_torch_dist.py " - "${MODEL_ARGS[@]} " - f"--hf-checkpoint {model_dir} --save {checkpoint_dir}" - ) - - tracker = CHECKPOINT_DIR / "latest_checkpointed_iteration.txt" - assert tracker.read_text().strip() == "release" - weight_files = [ - path - for path in CHECKPOINT_DIR.rglob("*") - if path.is_file() and path.name != "latest_checkpointed_iteration.txt" - ] - assert weight_files, f"No checkpoint weights found under {CHECKPOINT_DIR}" - - U.execute_train( - train_args=( - f"--hf-checkpoint {model_dir} " - f"--ref-load {checkpoint_dir} " - "--debug-train-only " - "--num-rollout 1 " - "--start-rollout-id 1 " - "--no-load-optim " - "--no-load-rng " - "--rollout-batch-size 1 " - "--global-batch-size 1 " - "--kl-coef 0.1 " - "--lr-decay-iters 1 " - "--optimizer adam " - "--lr 1e-6 " - "--lr-decay-style constant " - "--weight-decay 0.0 " - "--adam-beta1 0.9 " - "--adam-beta2 0.98 " - "--actor-num-nodes 1 " - f"--actor-num-gpus-per-node {NUM_GPUS} " - "--ci-test " - ), - num_gpus_per_node=NUM_GPUS, - megatron_model_type="qwen3-4B", - ) - - -def main(): - prepare() - execute() - - -if __name__ == "__main__": - main() diff --git a/tests/test_qwen3_30B_A3B_npu.py b/tests/test_qwen3_30B_A3B_npu.py index a88197043..cfafab53a 100644 --- a/tests/test_qwen3_30B_A3B_npu.py +++ b/tests/test_qwen3_30B_A3B_npu.py @@ -6,6 +6,7 @@ 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 +14,35 @@ 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( + "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}" + ) + tracker = CHECKPOINT_DIR / "latest_checkpointed_iteration.txt" + assert tracker.read_text().strip() == "release" + weight_files = [ + path + for path in CHECKPOINT_DIR.rglob("*") + if path.is_file() and path.name != "latest_checkpointed_iteration.txt" + ] + assert weight_files, f"No checkpoint weights found under {CHECKPOINT_DIR}" 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} " From 0970086f7f2e6c52dd383b27df571467bcbc8986 Mon Sep 17 00:00:00 2001 From: flb_ Date: Thu, 3 Sep 2026 12:31:35 +0800 Subject: [PATCH 8/9] fix path Signed-off-by: flb_ --- tests/test_qwen3_30B_A3B_npu.py | 9 ++++++--- 1 file changed, 6 insertions(+), 3 deletions(-) diff --git a/tests/test_qwen3_30B_A3B_npu.py b/tests/test_qwen3_30B_A3B_npu.py index cfafab53a..d661f6124 100644 --- a/tests/test_qwen3_30B_A3B_npu.py +++ b/tests/test_qwen3_30B_A3B_npu.py @@ -1,5 +1,6 @@ import os import shlex +from pathlib import Path import vime.utils.external_utils.command_utils as U @@ -27,14 +28,16 @@ def prepare(): "${MODEL_ARGS[@]} " f"--hf-checkpoint {model_dir} --save {checkpoint_dir}" ) - tracker = CHECKPOINT_DIR / "latest_checkpointed_iteration.txt" + + 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_DIR.rglob("*") + 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_DIR}" + assert weight_files, f"No checkpoint weights found under {checkpoint_path}" def execute(): From 81eff1fd9cd9abdd7c881b2b11e5fb70865afcde Mon Sep 17 00:00:00 2001 From: flb_ Date: Thu, 3 Sep 2026 14:51:03 +0800 Subject: [PATCH 9/9] remove torch_dist first Signed-off-by: flb_ --- tests/test_qwen3_30B_A3B_npu.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/test_qwen3_30B_A3B_npu.py b/tests/test_qwen3_30B_A3B_npu.py index d661f6124..cfe0a932f 100644 --- a/tests/test_qwen3_30B_A3B_npu.py +++ b/tests/test_qwen3_30B_A3B_npu.py @@ -21,6 +21,7 @@ def prepare(): 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 "