Expand the model args in python before building the command

miles/utils/external_utils/command_utils.py

@@ -9,6 +9,7 @@
 import random
 import shlex
 import socket
+import subprocess
 from dataclasses import dataclass, field
 from functools import partial
 from pathlib import Path
@@ -31,6 +32,18 @@ def _pythonpath_with_sources(megatron_path: str, *additional_pythonpaths: str |
     return os.pathsep.join(dict.fromkeys(entries))
 
 
+def load_model_args(megatron_model_type: str) -> list[str]:
+    """Expand the MODEL_ARGS array that scripts/models/<megatron_model_type>.sh declares."""
+    script = f"{repo_base_dir}/scripts/models/{megatron_model_type}.sh"
+    assert os.path.exists(script), f"no model args script at {script}"
+    expansion = f'source {shlex.quote(script)} && printf "%s\\0" "${{MODEL_ARGS[@]}}"'
+    result = subprocess.run(["bash", "-c", expansion], capture_output=True, text=True, check=True)
+    tokens = result.stdout.split("\0")[:-1]
+    for token in tokens:
+        assert token.split() == [token], f"model args token must be one whitespace-free word: {token!r}"
+    return tokens
+
+
 def convert_checkpoint(
     model_name,
     megatron_model_type,
@@ -63,13 +76,12 @@ def convert_checkpoint(
         fn = exec_command_gpu
     pythonpath = shlex.quote(_pythonpath_with_sources(megatron_path))
     fn(
-        f"source {repo_base_dir}/scripts/models/{megatron_model_type}.sh && "
         f"PYTHONPATH={pythonpath} "
         f"torchrun "
         f"--nproc-per-node {num_gpus_per_node} "
         f"{multinode_args}"
         f"{repo_base_dir}/tools/convert_hf_to_torch_dist.py "
-        "${MODEL_ARGS[@]} "
+        f"{' '.join(load_model_args(megatron_model_type))} "
         f"--hf-checkpoint {hf_checkpoint} "
         f"--save {path_dst} "
         f"{extra_args}"
@@ -195,18 +207,13 @@ def execute_train(
     runtime_env_json = json.dumps({"env_vars": runtime_env_vars})
 
     if get_bool_env_var("MILES_SCRIPT_ENABLE_RAY_SUBMIT", "1"):
-        cmd_megatron_model_source = (
-            f'source "{repo_base_dir}/scripts/models/{megatron_model_type}.sh" && '
-            if megatron_model_type is not None
-            else ""
-        )
+        model_args = " ".join(load_model_args(megatron_model_type)) if megatron_model_type is not None else ""
         exec_command_cpu(
             f"export no_proxy=127.0.0.1 && export PYTHONUNBUFFERED=1 && "
-            f"{cmd_megatron_model_source}"
             f"""ray job submit {'' if 'RAY_ADDRESS' in os.environ else '--address="http://127.0.0.1:8265" '}"""
             f"--runtime-env-json={shlex.quote(runtime_env_json)} "
             f"-- python3 {train_script} "
-            f"{'${MODEL_ARGS[@]}' if megatron_model_type is not None else ''} "
+            f"{model_args} "
             f"{train_args}"
         )
 

tests/fast/launch_scripts/py_harness.py

@@ -30,19 +30,23 @@
     "WANDB_API_KEY": "frozen-wandb-api-key",
 }
 
-_CLEARED_ENV = (
+CLEARED_ENV = (
     "CUDA_VISIBLE_DEVICES",
     "GITHUB_COMMIT_NAME",
     "GLOO_SOCKET_IFNAME",
     "KEEP_MOE_LORA",
     "MILES_SCRIPT_EXTERNAL_RAY",
+    "MODEL_ARGS_FIRST_K_DENSE_REPLACE",
+    "MODEL_ARGS_NUM_LAYERS",
+    "MODEL_ARGS_ROTARY_BASE",
     "NCCL_DEBUG",
     "NCCL_DEBUG_FILE",
     "NCCL_NVLS_ENABLE",
     "NCCL_SOCKET_IFNAME",
     "NO_PROXY",
     "OPTIMIZER_CPU_OFFLOAD",
     "RAY_ADDRESS",
+    "ROTARY_SCALING_FACTOR",
     "SLURM_JOB_NUM_NODES",
 )
 
@@ -103,7 +107,7 @@ def fake_run(command, *args, **kwargs):
 def freeze_environment(monkeypatch) -> None:
     for key, value in _FROZEN_ENV.items():
         monkeypatch.setenv(key, value)
-    for key in _CLEARED_ENV:
+    for key in CLEARED_ENV:
         monkeypatch.delenv(key, raising=False)
 
 

tests/fast/launch_scripts/test_py_launch_scripts.py

@@ -1,10 +1,12 @@
 import json
+import re
 from collections.abc import Callable
 from pathlib import Path
 
 import pytest
 
 from tests.fast.launch_scripts.py_harness import (
+    CLEARED_ENV,
     call_entrypoint,
     format_recording,
     freeze_environment,
@@ -147,6 +149,19 @@ def test_the_uncovered_launcher_really_is_uncoverable_here(self, rel):
         with pytest.raises(ImportError, match="execute_train_npu"):
             import_launch_script(REPO_ROOT / rel)
 
+    def test_every_environment_knob_a_model_script_reads_is_frozen(self):
+        """The snapshots now pin expanded model args, so a developer's exported override would fail them."""
+        knobs = set()
+        for script in sorted((REPO_ROOT / "scripts" / "models").iterdir()):
+            if not script.is_file():
+                continue
+            text = script.read_text()
+            knobs |= set(re.findall(r"\$\{([A-Z][A-Z0-9_]*):-", text))
+            knobs |= set(re.findall(r"environ\.get\(\s*\"([A-Z][A-Z0-9_]*)\"", text))
+
+        assert knobs
+        assert knobs <= set(CLEARED_ENV)
+
     def test_execute_train_config_defaults_are_not_taken_from_a_slurm_allocation(self, monkeypatch):
         """SLURM_JOB_NUM_NODES is read at import time, so a stale allocation would skew every snapshot."""
         import miles.utils.external_utils.command_utils as command_utils

tests/fast/utils/test_command_utils.py

@@ -44,7 +44,7 @@ def test_preserves_source_paths_on_the_pythonpath(self, monkeypatch, tmp_path):
 
         command_utils.convert_checkpoint(
             model_name="model",
-            megatron_model_type="model_type",
+            megatron_model_type="qwen3-4B",
             num_gpus_per_node=1,
             dir_dst=str(tmp_path),
             megatron_path="/megatron",
@@ -234,7 +234,7 @@ def test_exports_unbuffered_python_to_ray(self, monkeypatch):
         command_utils.execute_train(
             train_args="",
             num_gpus_per_node=1,
-            megatron_model_type="model_type",
+            megatron_model_type="qwen3-4B",
         )
 
         exports = [command for command in commands if "export PYTHONUNBUFFERED" in command]
@@ -250,7 +250,7 @@ def test_unbuffers_the_ray_workers_too(self, monkeypatch):
         monkeypatch.setattr(command_utils, "exec_command_cpu", commands.append)
         monkeypatch.setattr(command_utils, "check_has_nvlink", lambda: False)
 
-        command_utils.execute_train(train_args="", num_gpus_per_node=1, megatron_model_type="model_type")
+        command_utils.execute_train(train_args="", num_gpus_per_node=1, megatron_model_type="qwen3-4B")
 
         runtime_env_arg = next(arg for arg in shlex.split(commands[-1]) if arg.startswith("--runtime-env-json="))
         assert json.loads(runtime_env_arg.split("=", 1)[1])["env_vars"]["PYTHONUNBUFFERED"] == "1"
@@ -267,7 +267,7 @@ def test_preserves_source_paths_in_the_ray_runtime(self, monkeypatch):
         command_utils.execute_train(
             train_args="",
             num_gpus_per_node=1,
-            megatron_model_type="model_type",
+            megatron_model_type="qwen3-4B",
             megatron_path="/megatron",
             extra_env_vars={"PYTHONPATH": "/custom:/sglang", "QUOTED_VALUE": "it's preserved"},
         )
@@ -328,21 +328,20 @@ def test_can_skip_the_ray_job_submit(self, commands, monkeypatch):
 
         assert not any("ray job submit" in command for command in commands)
 
-    def test_sources_the_model_config_and_expands_model_args(self, commands):
-        """The megatron model type is turned into a `source` plus a ${MODEL_ARGS[@]} expansion."""
+    def test_expands_the_model_config_into_the_submitted_command(self, commands):
+        """The megatron model type is expanded into the argv its model script declares."""
         command_utils.execute_train(train_args="--x 1", num_gpus_per_node=8, megatron_model_type="qwen3-4B")
 
         submit = commands[-1]
-        assert f'source "{command_utils.repo_base_dir}/scripts/models/qwen3-4B.sh" && ' in submit
-        assert "${MODEL_ARGS[@]}" in submit
+        assert "--num-layers 36 " in submit
+        assert "source" not in submit
         assert submit.endswith("--x 1")
 
-    def test_omits_the_model_source_for_fsdp(self, commands):
-        """FSDP has no megatron model config to source."""
+    def test_omits_the_model_args_for_fsdp(self, commands):
+        """FSDP has no megatron model config to expand."""
         command_utils.execute_train(train_args="--train-backend fsdp", num_gpus_per_node=8, megatron_model_type=None)
 
-        assert "scripts/models/" not in commands[-1]
-        assert "${MODEL_ARGS[@]}" not in commands[-1]
+        assert "--num-layers" not in commands[-1]
 
     def test_drops_cuda_device_max_connections_for_fsdp(self, commands):
         """Pinning it to 1 breaks computation/communication overlap on FSDP."""

tests/snapshots/launch_scripts/py/scripts/amd/run_deepseek_v4.py/full_train.txt

@@ -31,8 +31,64 @@ python <REPO_ROOT>/tools/fp8_cast_bf16.py
   --output-bf16-hf-path /root/models/DeepSeek-V4-Flash-FP8-bf16/ 
 
 ### 7
-source <REPO_ROOT>/scripts/models/deepseek-v4-flash.sh && PYTHONPATH=<REPO_ROOT>:/root/Megatron-LM:/frozen/pythonpath torchrun
-  --nproc-per-node 8 <REPO_ROOT>/tools/convert_hf_to_torch_dist.py ${MODEL_ARGS[@]}
+PYTHONPATH=<REPO_ROOT>:/root/Megatron-LM:/frozen/pythonpath torchrun
+  --nproc-per-node 8 <REPO_ROOT>/tools/convert_hf_to_torch_dist.py
+  --disable-bias-linear
+  --num-layers 43
+  --hidden-size 4096
+  --ffn-hidden-size 2048
+  --num-attention-heads 64
+  --normalization RMSNorm
+  --position-embedding-type rope
+  --norm-epsilon 1e-6
+  --swiglu
+  --untie-embeddings-and-output-weights
+  --vocab-size 129280
+  --hidden-dropout 0.0
+  --attention-dropout 0.0
+  --multi-latent-attention
+  --q-lora-rank 1024
+  --kv-lora-rank 512
+  --qk-head-dim 512
+  --qk-pos-emb-head-dim 64
+  --v-head-dim 512
+  --qk-layernorm
+  --rotary-scaling-factor 16
+  --rotary-base 10000
+  --original-max-position-embeddings 65536
+  --beta-fast 32
+  --beta-slow 1
+  --attention-softmax-in-fp32
+  --no-rope-fusion
+  --num-experts 256
+  --moe-layer-freq [1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1]
+  --moe-ffn-hidden-size 2048
+  --moe-router-topk 6
+  --moe-shared-expert-intermediate-size 2048
+  --moe-router-pre-softmax
+  --moe-router-score-function sqrtsoftplus
+  --moe-router-enable-expert-bias
+  --moe-router-load-balancing-type seq_aux_loss
+  --moe-token-dispatcher-type alltoall
+  --moe-aux-loss-coeff 0
+  --moe-grouped-gemm
+  --moe-router-topk-scaling-factor 1.5
+  --experimental-attention-variant dsv4
+  --dsv4-hc-mult 4
+  --dsv4-hc-sinkhorn-iters 20
+  --dsv4-compress-ratios 0 0 4 128 4 128 4 128 4 128 4 128 4 128 4 128 4 128 4 128 4 128 4 128 4 128 4 128 4 128 4 128 4 128 4 128 4 128 4 128 4 128 4 0
+  --dsv4-compress-rope-theta 160000
+  --dsv4-o-groups 8
+  --dsv4-o-lora-rank 1024
+  --dsv4-n-hash-layers 3
+  --dsv4-window-size 128
+  --dsa-indexer-n-heads 64
+  --dsa-indexer-head-dim 128
+  --dsa-indexer-topk 512
+  --spec miles_plugins.models.deepseek_v4.deepseek_v4 get_dsv4_spec
+  --activation-func-clamp-value 10
+  --no-bias-swiglu-fusion
+  --no-activation-func-clamp-shared-expert
   --hf-checkpoint /root/models/DeepSeek-V4-Flash-FP8-bf16
   --save /root/models/DeepSeek-V4-Flash-FP8_torch_dist
   --expert-tensor-parallel-size 1
@@ -56,10 +112,66 @@ export PYTHONUNBUFFERED=1 && ray start
 nvidia-smi topo -m 2>/dev/null | grep -o 'NV[0-9][0-9]*' | wc -l
 
 ### 11
-export no_proxy=127.0.0.1 && export PYTHONUNBUFFERED=1 && source "<REPO_ROOT>/scripts/models/deepseek-v4-flash.sh" && ray job submit
+export no_proxy=127.0.0.1 && export PYTHONUNBUFFERED=1 && ray job submit
   --address="http://127.0.0.1:8265"
   --runtime-env-json='{"env_vars": {"PYTHONUNBUFFERED": "1", "CUDA_DEVICE_MAX_CONNECTIONS": "1", "NCCL_NVLS_ENABLE": "0", "no_proxy": "127.0.0.1,127.0.0.1", "MASTER_ADDR": "127.0.0.1", "SGLANG_SKIP_CHECKPOINT_LOAD_CHECK": "1", "SGLANG_DSV4_FP4_EXPERTS": "0", "SGLANG_HACK_FLASHMLA_BACKEND": "unified_kv_triton", "SGLANG_OPT_USE_COMPRESSOR_V2": "true", "SGLANG_OPT_USE_TILELANG_INDEXER": "true", "SGLANG_OPT_USE_JIT_NORM": "true", "SGLANG_OPT_USE_FUSED_COMPRESS": "true", "SGLANG_HEALTH_CHECK_TIMEOUT": "120", "AITER_BF16_FP8_MOE_BOUND": "0", "NCCL_ALGO": "Ring", "NVTE_ALLOW_NONDETERMINISTIC_ALGO": "0", "CUBLAS_WORKSPACE_CONFIG": ":4096:8", "PYTHONPATH": "<REPO_ROOT>:/root/Megatron-LM:/frozen/pythonpath"}}'
-  -- python3 <REPO_ROOT>/train.py ${MODEL_ARGS[@]}
+  -- python3 <REPO_ROOT>/train.py
+  --disable-bias-linear
+  --num-layers 43
+  --hidden-size 4096
+  --ffn-hidden-size 2048
+  --num-attention-heads 64
+  --normalization RMSNorm
+  --position-embedding-type rope
+  --norm-epsilon 1e-6
+  --swiglu
+  --untie-embeddings-and-output-weights
+  --vocab-size 129280
+  --hidden-dropout 0.0
+  --attention-dropout 0.0
+  --multi-latent-attention
+  --q-lora-rank 1024
+  --kv-lora-rank 512
+  --qk-head-dim 512
+  --qk-pos-emb-head-dim 64
+  --v-head-dim 512
+  --qk-layernorm
+  --rotary-scaling-factor 16
+  --rotary-base 10000
+  --original-max-position-embeddings 65536
+  --beta-fast 32
+  --beta-slow 1
+  --attention-softmax-in-fp32
+  --no-rope-fusion
+  --num-experts 256
+  --moe-layer-freq [1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1]
+  --moe-ffn-hidden-size 2048
+  --moe-router-topk 6
+  --moe-shared-expert-intermediate-size 2048
+  --moe-router-pre-softmax
+  --moe-router-score-function sqrtsoftplus
+  --moe-router-enable-expert-bias
+  --moe-router-load-balancing-type seq_aux_loss
+  --moe-token-dispatcher-type alltoall
+  --moe-aux-loss-coeff 0
+  --moe-grouped-gemm
+  --moe-router-topk-scaling-factor 1.5
+  --experimental-attention-variant dsv4
+  --dsv4-hc-mult 4
+  --dsv4-hc-sinkhorn-iters 20
+  --dsv4-compress-ratios 0 0 4 128 4 128 4 128 4 128 4 128 4 128 4 128 4 128 4 128 4 128 4 128 4 128 4 128 4 128 4 128 4 128 4 128 4 128 4 128 4 128 4 0
+  --dsv4-compress-rope-theta 160000
+  --dsv4-o-groups 8
+  --dsv4-o-lora-rank 1024
+  --dsv4-n-hash-layers 3
+  --dsv4-window-size 128
+  --dsa-indexer-n-heads 64
+  --dsa-indexer-head-dim 128
+  --dsa-indexer-topk 512
+  --spec miles_plugins.models.deepseek_v4.deepseek_v4 get_dsv4_spec
+  --activation-func-clamp-value 10
+  --no-bias-swiglu-fusion
+  --no-activation-func-clamp-shared-expert
   --hf-checkpoint /root/models/DeepSeek-V4-Flash-FP8
   --ref-load /root/models/DeepSeek-V4-Flash-FP8_torch_dist
   --load /root/models/260101-000000-000/checkpoints

tests/snapshots/launch_scripts/py/scripts/amd/run_deepseek_v4.py/prepare_spmd.txt

@@ -1,6 +1,62 @@
 ### 0
-source <REPO_ROOT>/scripts/models/deepseek-v4-flash.sh && PYTHONPATH=<REPO_ROOT>:/root/Megatron-LM:/frozen/pythonpath torchrun
-  --nproc-per-node 8 <REPO_ROOT>/tools/convert_hf_to_torch_dist.py ${MODEL_ARGS[@]}
+PYTHONPATH=<REPO_ROOT>:/root/Megatron-LM:/frozen/pythonpath torchrun
+  --nproc-per-node 8 <REPO_ROOT>/tools/convert_hf_to_torch_dist.py
+  --disable-bias-linear
+  --num-layers 43
+  --hidden-size 4096
+  --ffn-hidden-size 2048
+  --num-attention-heads 64
+  --normalization RMSNorm
+  --position-embedding-type rope
+  --norm-epsilon 1e-6
+  --swiglu
+  --untie-embeddings-and-output-weights
+  --vocab-size 129280
+  --hidden-dropout 0.0
+  --attention-dropout 0.0
+  --multi-latent-attention
+  --q-lora-rank 1024
+  --kv-lora-rank 512
+  --qk-head-dim 512
+  --qk-pos-emb-head-dim 64
+  --v-head-dim 512
+  --qk-layernorm
+  --rotary-scaling-factor 16
+  --rotary-base 10000
+  --original-max-position-embeddings 65536
+  --beta-fast 32
+  --beta-slow 1
+  --attention-softmax-in-fp32
+  --no-rope-fusion
+  --num-experts 256
+  --moe-layer-freq [1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1]
+  --moe-ffn-hidden-size 2048
+  --moe-router-topk 6
+  --moe-shared-expert-intermediate-size 2048
+  --moe-router-pre-softmax
+  --moe-router-score-function sqrtsoftplus
+  --moe-router-enable-expert-bias
+  --moe-router-load-balancing-type seq_aux_loss
+  --moe-token-dispatcher-type alltoall
+  --moe-aux-loss-coeff 0
+  --moe-grouped-gemm
+  --moe-router-topk-scaling-factor 1.5
+  --experimental-attention-variant dsv4
+  --dsv4-hc-mult 4
+  --dsv4-hc-sinkhorn-iters 20
+  --dsv4-compress-ratios 0 0 4 128 4 128 4 128 4 128 4 128 4 128 4 128 4 128 4 128 4 128 4 128 4 128 4 128 4 128 4 128 4 128 4 128 4 128 4 128 4 128 4 0
+  --dsv4-compress-rope-theta 160000
+  --dsv4-o-groups 8
+  --dsv4-o-lora-rank 1024
+  --dsv4-n-hash-layers 3
+  --dsv4-window-size 128
+  --dsa-indexer-n-heads 64
+  --dsa-indexer-head-dim 128
+  --dsa-indexer-topk 512
+  --spec miles_plugins.models.deepseek_v4.deepseek_v4 get_dsv4_spec
+  --activation-func-clamp-value 10
+  --no-bias-swiglu-fusion
+  --no-activation-func-clamp-shared-expert
   --hf-checkpoint /root/models/DeepSeek-V4-Flash-FP8-bf16
   --save /root/models/DeepSeek-V4-Flash-FP8_torch_dist
   --expert-tensor-parallel-size 1