Fix p2p profile's rotary_base not reaching the model script it configures

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

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("--")

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 "<REPO_ROOT>/scripts/models/glm4.5-106B-A12B.sh" && ray job submit
+MODEL_ARGS_ROTARY_BASE=1000000 source "<REPO_ROOT>/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": "<SANDBOX>", "MC_TRANSFER_TIMEOUT": "300", "MODEL_ARGS_ROTARY_BASE": "1000000"}}'
   -- python3 "<REPO_ROOT>/train.py" ${MODEL_ARGS[@]}

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 "<REPO_ROOT>/scripts/models/glm4.5-106B-A12B.sh" && ray job submit
+MODEL_ARGS_ROTARY_BASE=1000000 source "<REPO_ROOT>/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": "<SANDBOX>", "MC_TRANSFER_TIMEOUT": "300", "MODEL_ARGS_ROTARY_BASE": "1000000"}}'
   -- python3 "<REPO_ROOT>/train.py" ${MODEL_ARGS[@]}

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 "<REPO_ROOT>/scripts/models/qwen3-235B-A22B.sh" && ray job submit
+MODEL_ARGS_ROTARY_BASE=5000000 source "<REPO_ROOT>/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": "<SANDBOX>", "MC_TRANSFER_TIMEOUT": "300", "MODEL_ARGS_ROTARY_BASE": "5000000"}}'
   -- python3 "<REPO_ROOT>/train.py" ${MODEL_ARGS[@]}

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 "<REPO_ROOT>/scripts/models/qwen3-235B-A22B.sh" && ray job submit
+MODEL_ARGS_ROTARY_BASE=5000000 source "<REPO_ROOT>/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": "<SANDBOX>", "MC_TRANSFER_TIMEOUT": "300", "MODEL_ARGS_ROTARY_BASE": "5000000"}}'
   -- python3 "<REPO_ROOT>/train.py" ${MODEL_ARGS[@]}

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 "<REPO_ROOT>/scripts/models/qwen3-30B-A3B.sh" && ray job submit
+MODEL_ARGS_ROTARY_BASE=1000000 source "<REPO_ROOT>/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": "<SANDBOX>", "MC_TRANSFER_TIMEOUT": "300", "MODEL_ARGS_ROTARY_BASE": "1000000"}}'
   -- python3 "<REPO_ROOT>/train.py" ${MODEL_ARGS[@]}

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 "<REPO_ROOT>/scripts/models/qwen3-30B-A3B.sh" && ray job submit
+MODEL_ARGS_ROTARY_BASE=1000000 source "<REPO_ROOT>/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": "<SANDBOX>", "MC_TRANSFER_TIMEOUT": "300", "MODEL_ARGS_ROTARY_BASE": "1000000"}}'
   -- python3 "<REPO_ROOT>/train.py" ${MODEL_ARGS[@]}