diff --git a/examples/infra_features/p2p_weight_transfer/run.py b/examples/infra_features/p2p_weight_transfer/run.py index c05e93d3d2a..dd14e9b145f 100644 --- a/examples/infra_features/p2p_weight_transfer/run.py +++ b/examples/infra_features/p2p_weight_transfer/run.py @@ -18,6 +18,7 @@ import json import os +import shlex import subprocess import sys import time @@ -778,7 +779,8 @@ def cmd_run( run_cmd("pkill -9 redis || true", check=False) # --- Source model args --- - model_args_source = f'source "{MILES_ROOT}/scripts/models/{cfg.model_type}.sh"' + model_args_env = build_model_args_env(cfg) + model_args_source = build_model_args_command(cfg) # --- Worker sleep --- if not is_single_node and node_rank > 0: @@ -838,8 +840,7 @@ def cmd_run( } if not is_single_node: env_vars["MC_TRANSFER_TIMEOUT"] = str(cfg.mc_transfer_timeout) - if cfg.rotary_base is not None: - env_vars["MODEL_ARGS_ROTARY_BASE"] = str(cfg.rotary_base) + env_vars.update(model_args_env) env_vars.update(cfg.extra_env_vars) runtime_env_json = json.dumps({"env_vars": env_vars}) @@ -1067,8 +1068,6 @@ def cmd_run( # --- Submit Ray job (head node only, or single-node) --- if is_single_node or node_rank == 0: - import shlex - args_str = " ".join(shlex.quote(a) for a in args) run_cmd( f"{model_args_source} && " @@ -1099,6 +1098,19 @@ def cmd_run( print("Done.") +def build_model_args_command(cfg: RunConfig) -> str: + """A shell snippet leaving MODEL_ARGS set; the knobs must reach it, not only ray's runtime env.""" + prefix = "".join(f"{name}={shlex.quote(value)} " for name, value in build_model_args_env(cfg).items()) + return f'{prefix}source "{MILES_ROOT}/scripts/models/{cfg.model_type}.sh"' + + +def build_model_args_env(cfg: RunConfig) -> dict[str, str]: + """The MODEL_ARGS_* knobs the model definitions read, as declared by the profile.""" + if cfg.rotary_base is None: + return {} + return {"MODEL_ARGS_ROTARY_BASE": str(cfg.rotary_base)} + + # --------------------------------------------------------------------------- # CLI # --------------------------------------------------------------------------- diff --git a/tests/fast/examples/__init__.py b/tests/fast/examples/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/fast/examples/infra_features/__init__.py b/tests/fast/examples/infra_features/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/fast/examples/infra_features/p2p_weight_transfer/__init__.py b/tests/fast/examples/infra_features/p2p_weight_transfer/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/fast/examples/infra_features/p2p_weight_transfer/test_run.py b/tests/fast/examples/infra_features/p2p_weight_transfer/test_run.py new file mode 100644 index 00000000000..c602960b26f --- /dev/null +++ b/tests/fast/examples/infra_features/p2p_weight_transfer/test_run.py @@ -0,0 +1,64 @@ +import importlib.util +import subprocess +from pathlib import Path +from types import ModuleType + +import pytest + +REPO_ROOT = Path(__file__).resolve().parents[5] +RUN_SCRIPT = REPO_ROOT / "examples" / "infra_features" / "p2p_weight_transfer" / "run.py" + +_PROFILES_PINNING_A_ROTARY_BASE = ["Qwen3-235B-A22B-Instruct-2507", "Qwen3-30B-A3B", "GLM-4.5-Air"] + + +@pytest.fixture(scope="module") +def run_module() -> ModuleType: + spec = importlib.util.spec_from_file_location("p2p_weight_transfer_run", RUN_SCRIPT) + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +def expand_model_args(run_module: ModuleType, model_name: str) -> list[str]: + """Run the very snippet run.py hands to bash, so this keeps testing the real command.""" + command = run_module.build_model_args_command(run_module.RUN_CONFIGS[model_name]) + result = subprocess.run( + f'{command} && printf "%s\\n" "${{MODEL_ARGS[@]}}"', + shell=True, + executable="/bin/bash", + capture_output=True, + text=True, + check=True, + ) + return result.stdout.splitlines() + + +def test_model_args_env_is_empty_when_the_profile_pins_no_rotary_base(run_module: ModuleType) -> None: + """A profile without rotary_base must not inject any MODEL_ARGS_* override.""" + cfg = run_module.RUN_CONFIGS["Qwen3-4B"] + assert cfg.rotary_base is None + assert run_module.build_model_args_env(cfg) == {} + + +def test_model_args_env_carries_the_rotary_base_a_profile_pins(run_module: ModuleType) -> None: + """A profile pinning rotary_base must surface it as the MODEL_ARGS_* name the model definition reads.""" + cfg = run_module.RUN_CONFIGS["Qwen3-235B-A22B-Instruct-2507"] + assert cfg.rotary_base == 5000000 + assert run_module.build_model_args_env(cfg) == {"MODEL_ARGS_ROTARY_BASE": "5000000"} + + +@pytest.mark.parametrize("model_name", _PROFILES_PINNING_A_ROTARY_BASE) +def test_a_pinned_rotary_base_reaches_the_expanded_model_args(run_module: ModuleType, model_name: str) -> None: + """The knob has to survive into the shell that expands MODEL_ARGS, not only into ray's runtime env.""" + tokens = expand_model_args(run_module, model_name) + + assert str(run_module.RUN_CONFIGS[model_name].rotary_base) == tokens[tokens.index("--rotary-base") + 1] + + +@pytest.mark.parametrize("model_name", sorted({"Qwen3-4B", "GLM-4.7-Flash", *_PROFILES_PINNING_A_ROTARY_BASE})) +def test_every_profile_expands_to_a_usable_argv(run_module: ModuleType, model_name: str) -> None: + """A profile naming a model that no longer exists would submit a job with no architecture flags.""" + tokens = expand_model_args(run_module, model_name) + + assert tokens + assert tokens[0].startswith("--") diff --git a/tests/snapshots/launch_scripts/self_executing/examples/infra_features/p2p_weight_transfer/run.py/run/GLM-4.5-Air/broadcast.txt b/tests/snapshots/launch_scripts/self_executing/examples/infra_features/p2p_weight_transfer/run.py/run/GLM-4.5-Air/broadcast.txt index 0985581a519..4a24915d1fb 100644 --- a/tests/snapshots/launch_scripts/self_executing/examples/infra_features/p2p_weight_transfer/run.py/run/GLM-4.5-Air/broadcast.txt +++ b/tests/snapshots/launch_scripts/self_executing/examples/infra_features/p2p_weight_transfer/run.py/run/GLM-4.5-Air/broadcast.txt @@ -39,7 +39,7 @@ RAY_memory_monitor_refresh_ms=0 ray start python3 -c "import ray; ray.init(address='auto', ignore_reinit_error=True); print(int(ray.cluster_resources().get('GPU', 0))); ray.shutdown()" ### 11 -source "/scripts/models/glm4.5-106B-A12B.sh" && ray job submit +MODEL_ARGS_ROTARY_BASE=1000000 source "/scripts/models/glm4.5-106B-A12B.sh" && ray job submit --address='http://127.0.0.1:8265' --runtime-env-json='{"env_vars": {"PYTHONUNBUFFERED": "1", "RAY_DEBUG": "1", "PYTHONPATH": "/root/Megatron-LM/", "CUDA_DEVICE_MAX_CONNECTIONS": "1", "NCCL_NVLS_ENABLE": "0", "MILES_LOG_DIR": "", "MC_TRANSFER_TIMEOUT": "300", "MODEL_ARGS_ROTARY_BASE": "1000000"}}' -- python3 "/train.py" ${MODEL_ARGS[@]} diff --git a/tests/snapshots/launch_scripts/self_executing/examples/infra_features/p2p_weight_transfer/run.py/run/GLM-4.5-Air/p2p.txt b/tests/snapshots/launch_scripts/self_executing/examples/infra_features/p2p_weight_transfer/run.py/run/GLM-4.5-Air/p2p.txt index 83b71704059..b992f747747 100644 --- a/tests/snapshots/launch_scripts/self_executing/examples/infra_features/p2p_weight_transfer/run.py/run/GLM-4.5-Air/p2p.txt +++ b/tests/snapshots/launch_scripts/self_executing/examples/infra_features/p2p_weight_transfer/run.py/run/GLM-4.5-Air/p2p.txt @@ -39,7 +39,7 @@ RAY_memory_monitor_refresh_ms=0 ray start python3 -c "import ray; ray.init(address='auto', ignore_reinit_error=True); print(int(ray.cluster_resources().get('GPU', 0))); ray.shutdown()" ### 11 -source "/scripts/models/glm4.5-106B-A12B.sh" && ray job submit +MODEL_ARGS_ROTARY_BASE=1000000 source "/scripts/models/glm4.5-106B-A12B.sh" && ray job submit --address='http://127.0.0.1:8265' --runtime-env-json='{"env_vars": {"PYTHONUNBUFFERED": "1", "RAY_DEBUG": "1", "PYTHONPATH": "/root/Megatron-LM/", "CUDA_DEVICE_MAX_CONNECTIONS": "1", "NCCL_NVLS_ENABLE": "0", "MILES_LOG_DIR": "", "MC_TRANSFER_TIMEOUT": "300", "MODEL_ARGS_ROTARY_BASE": "1000000"}}' -- python3 "/train.py" ${MODEL_ARGS[@]} diff --git a/tests/snapshots/launch_scripts/self_executing/examples/infra_features/p2p_weight_transfer/run.py/run/Qwen3-235B-A22B-Instruct-2507/broadcast.txt b/tests/snapshots/launch_scripts/self_executing/examples/infra_features/p2p_weight_transfer/run.py/run/Qwen3-235B-A22B-Instruct-2507/broadcast.txt index 63455c46129..225afb2058d 100644 --- a/tests/snapshots/launch_scripts/self_executing/examples/infra_features/p2p_weight_transfer/run.py/run/Qwen3-235B-A22B-Instruct-2507/broadcast.txt +++ b/tests/snapshots/launch_scripts/self_executing/examples/infra_features/p2p_weight_transfer/run.py/run/Qwen3-235B-A22B-Instruct-2507/broadcast.txt @@ -39,7 +39,7 @@ RAY_memory_monitor_refresh_ms=0 ray start python3 -c "import ray; ray.init(address='auto', ignore_reinit_error=True); print(int(ray.cluster_resources().get('GPU', 0))); ray.shutdown()" ### 11 -source "/scripts/models/qwen3-235B-A22B.sh" && ray job submit +MODEL_ARGS_ROTARY_BASE=5000000 source "/scripts/models/qwen3-235B-A22B.sh" && ray job submit --address='http://127.0.0.1:8265' --runtime-env-json='{"env_vars": {"PYTHONUNBUFFERED": "1", "RAY_DEBUG": "1", "PYTHONPATH": "/root/Megatron-LM/", "CUDA_DEVICE_MAX_CONNECTIONS": "1", "NCCL_NVLS_ENABLE": "1", "MILES_LOG_DIR": "", "MC_TRANSFER_TIMEOUT": "300", "MODEL_ARGS_ROTARY_BASE": "5000000"}}' -- python3 "/train.py" ${MODEL_ARGS[@]} diff --git a/tests/snapshots/launch_scripts/self_executing/examples/infra_features/p2p_weight_transfer/run.py/run/Qwen3-235B-A22B-Instruct-2507/p2p.txt b/tests/snapshots/launch_scripts/self_executing/examples/infra_features/p2p_weight_transfer/run.py/run/Qwen3-235B-A22B-Instruct-2507/p2p.txt index 0812884fa38..0a96a138158 100644 --- a/tests/snapshots/launch_scripts/self_executing/examples/infra_features/p2p_weight_transfer/run.py/run/Qwen3-235B-A22B-Instruct-2507/p2p.txt +++ b/tests/snapshots/launch_scripts/self_executing/examples/infra_features/p2p_weight_transfer/run.py/run/Qwen3-235B-A22B-Instruct-2507/p2p.txt @@ -39,7 +39,7 @@ RAY_memory_monitor_refresh_ms=0 ray start python3 -c "import ray; ray.init(address='auto', ignore_reinit_error=True); print(int(ray.cluster_resources().get('GPU', 0))); ray.shutdown()" ### 11 -source "/scripts/models/qwen3-235B-A22B.sh" && ray job submit +MODEL_ARGS_ROTARY_BASE=5000000 source "/scripts/models/qwen3-235B-A22B.sh" && ray job submit --address='http://127.0.0.1:8265' --runtime-env-json='{"env_vars": {"PYTHONUNBUFFERED": "1", "RAY_DEBUG": "1", "PYTHONPATH": "/root/Megatron-LM/", "CUDA_DEVICE_MAX_CONNECTIONS": "1", "NCCL_NVLS_ENABLE": "1", "MILES_LOG_DIR": "", "MC_TRANSFER_TIMEOUT": "300", "MODEL_ARGS_ROTARY_BASE": "5000000"}}' -- python3 "/train.py" ${MODEL_ARGS[@]} diff --git a/tests/snapshots/launch_scripts/self_executing/examples/infra_features/p2p_weight_transfer/run.py/run/Qwen3-30B-A3B/broadcast.txt b/tests/snapshots/launch_scripts/self_executing/examples/infra_features/p2p_weight_transfer/run.py/run/Qwen3-30B-A3B/broadcast.txt index 62e0f4baf30..74ab1db0cac 100644 --- a/tests/snapshots/launch_scripts/self_executing/examples/infra_features/p2p_weight_transfer/run.py/run/Qwen3-30B-A3B/broadcast.txt +++ b/tests/snapshots/launch_scripts/self_executing/examples/infra_features/p2p_weight_transfer/run.py/run/Qwen3-30B-A3B/broadcast.txt @@ -39,7 +39,7 @@ RAY_memory_monitor_refresh_ms=0 ray start python3 -c "import ray; ray.init(address='auto', ignore_reinit_error=True); print(int(ray.cluster_resources().get('GPU', 0))); ray.shutdown()" ### 11 -source "/scripts/models/qwen3-30B-A3B.sh" && ray job submit +MODEL_ARGS_ROTARY_BASE=1000000 source "/scripts/models/qwen3-30B-A3B.sh" && ray job submit --address='http://127.0.0.1:8265' --runtime-env-json='{"env_vars": {"PYTHONUNBUFFERED": "1", "RAY_DEBUG": "1", "PYTHONPATH": "/root/Megatron-LM/", "CUDA_DEVICE_MAX_CONNECTIONS": "1", "NCCL_NVLS_ENABLE": "1", "MILES_LOG_DIR": "", "MC_TRANSFER_TIMEOUT": "300", "MODEL_ARGS_ROTARY_BASE": "1000000"}}' -- python3 "/train.py" ${MODEL_ARGS[@]} diff --git a/tests/snapshots/launch_scripts/self_executing/examples/infra_features/p2p_weight_transfer/run.py/run/Qwen3-30B-A3B/p2p.txt b/tests/snapshots/launch_scripts/self_executing/examples/infra_features/p2p_weight_transfer/run.py/run/Qwen3-30B-A3B/p2p.txt index 45a3b63863c..259940107b6 100644 --- a/tests/snapshots/launch_scripts/self_executing/examples/infra_features/p2p_weight_transfer/run.py/run/Qwen3-30B-A3B/p2p.txt +++ b/tests/snapshots/launch_scripts/self_executing/examples/infra_features/p2p_weight_transfer/run.py/run/Qwen3-30B-A3B/p2p.txt @@ -39,7 +39,7 @@ RAY_memory_monitor_refresh_ms=0 ray start python3 -c "import ray; ray.init(address='auto', ignore_reinit_error=True); print(int(ray.cluster_resources().get('GPU', 0))); ray.shutdown()" ### 11 -source "/scripts/models/qwen3-30B-A3B.sh" && ray job submit +MODEL_ARGS_ROTARY_BASE=1000000 source "/scripts/models/qwen3-30B-A3B.sh" && ray job submit --address='http://127.0.0.1:8265' --runtime-env-json='{"env_vars": {"PYTHONUNBUFFERED": "1", "RAY_DEBUG": "1", "PYTHONPATH": "/root/Megatron-LM/", "CUDA_DEVICE_MAX_CONNECTIONS": "1", "NCCL_NVLS_ENABLE": "1", "MILES_LOG_DIR": "", "MC_TRANSFER_TIMEOUT": "300", "MODEL_ARGS_ROTARY_BASE": "1000000"}}' -- python3 "/train.py" ${MODEL_ARGS[@]}