diff --git a/src/hyperloom/inference_optimizer/cli/parser.py b/src/hyperloom/inference_optimizer/cli/parser.py index e96982fffe..46e44da14f 100644 --- a/src/hyperloom/inference_optimizer/cli/parser.py +++ b/src/hyperloom/inference_optimizer/cli/parser.py @@ -387,10 +387,11 @@ def _build_parser() -> argparse.ArgumentParser: default=None, help="Expert-parallel size for MoE inference. 1 (default) keeps " "experts sharded by TP (legacy behaviour). >=2 enables true " - "expert parallelism: sglang adds `--expert-parallel-size N`, " - "vllm adds `--enable-expert-parallel`. Typical: EP=TP for " - "DSr1/DSv3 on multi-node. Default: 1. " - "EP > TP is rejected at server-restart time.", + "expert parallelism: sglang adds `--ep-size N`, vllm and atom " + "add `--enable-expert-parallel` (EP then spans all TP ranks). " + "Typical: EP=TP. Default: 1. On a single node an EP that does " + "not divide TP is dropped with a warning; on multi-node EP > TP " + "is rejected at server-restart time.", ) opt.add_argument( "--pd-mode", diff --git a/src/hyperloom/inference_optimizer/framework_registry.py b/src/hyperloom/inference_optimizer/framework_registry.py index d77727171d..e6d450fc05 100644 --- a/src/hyperloom/inference_optimizer/framework_registry.py +++ b/src/hyperloom/inference_optimizer/framework_registry.py @@ -82,6 +82,15 @@ class FrameworkSpec: DEFAULT_FRAMEWORK = "sglang" +# Server option that turns on expert parallelism: (flag to write, every +# argparse alias of it, whether it takes the EP size). vLLM and ATOM take a +# switch and derive the size from TP x DP. Frameworks absent here have none. +EXPERT_PARALLEL_FLAGS: dict[str, tuple[str, tuple[str, ...], bool]] = { + "sglang": ("--ep-size", ("--ep-size", "--expert-parallel-size", "--ep"), True), + "vllm": ("--enable-expert-parallel", ("--enable-expert-parallel", "-ep"), False), + "atom": ("--enable-expert-parallel", ("--enable-expert-parallel",), False), +} + def names() -> tuple[str, ...]: """Return the canonical tuple of supported framework names.""" @@ -150,6 +159,15 @@ def server_args_env_name(framework: str | None) -> str: return extra_args_env(DEFAULT_FRAMEWORK) +def expert_parallel_flag(framework: str | None) -> tuple[str, tuple[str, ...], bool] | None: + """Return the :data:`EXPERT_PARALLEL_FLAGS` row for the server reading ``framework``'s server args.""" + env_name = server_args_env_name(framework) + return next( + (row for name, row in EXPERT_PARALLEL_FLAGS.items() if FRAMEWORKS[name].extra_args_env == env_name), + None, + ) + + def throughput_unit(framework: str | None) -> str: """Return the throughput unit string for ``framework``.""" return _spec_or_default(framework).throughput_unit diff --git a/src/hyperloom/inference_optimizer/tests/test_server_argv_seal.py b/src/hyperloom/inference_optimizer/tests/test_server_argv_seal.py index 4ae3aa20a2..a680c648a7 100644 --- a/src/hyperloom/inference_optimizer/tests/test_server_argv_seal.py +++ b/src/hyperloom/inference_optimizer/tests/test_server_argv_seal.py @@ -13,7 +13,9 @@ from hyperloom.inference_optimizer.framework_registry import names, server_args_env_name from hyperloom.orchestrator.actions.executors import _grid_runner, _workload_envs +from hyperloom.orchestrator.actions.executors._grid_base import GridVariant from hyperloom.orchestrator.actions.executors._server_argv import ( + ServerArgv, add_server_arg_unless_pinned, config_launch_env, config_server_argv, @@ -269,3 +271,143 @@ def test_a_repaired_argv_reaches_the_launch_through_the_rendered_config(tmp_path assert resealed.argv == ("--tp", "8") # Read back off disk, because the file is what the launch opens. assert config_server_argv(config).argv == ("--tp", "8") + + +def _seal_with_ep(framework: str, args: str, ep: object) -> ServerArgv: + """Seal ``args`` for ``framework`` with ``ep`` in the envs; ``None`` leaves EP unset.""" + envs: dict[str, object] = {server_args_env_name(framework): args} + if ep is not None: + envs["EP"] = ep + return seal_server_argv(envs, framework) + + +@pytest.mark.parametrize( + ("framework", "args", "ep", "expected"), + [ + ("sglang", "--tp 8", 4, ("--tp", "8", "--ep-size", "4")), + ("sglang", "", "8", ("--ep-size", "8")), + ("sglang", "--ep-size 4", 4, ("--ep-size", "4")), + ("sglang", "--expert-parallel-size=4", 4, ("--expert-parallel-size=4",)), + ("sglang", "--ep 4", "4", ("--ep", "4")), + # A flag that merely starts with an alias is a different option. + ("sglang", "--ep-num-redundant-experts 32", 4, ("--ep-num-redundant-experts", "32", "--ep-size", "4")), + ("vllm", "--foo", 8, ("--foo", "--enable-expert-parallel")), + ("vllm", "", 8, ("--enable-expert-parallel",)), + ("vllm", "--enable-expert-parallel", 8, ("--enable-expert-parallel",)), + ("vllm", "-ep", 8, ("-ep",)), + ("atom", "--level 2", 8, ("--level", "2", "--enable-expert-parallel")), + ("vllm", "--foo", 1, ("--foo",)), + ("vllm", "--foo", None, ("--foo",)), + # EP 1 is a no-op, so a recipe that pins its own EP size is left standing. + ("sglang", "--ep-size 4", 1, ("--ep-size", "4")), + ("xdit", "--foo", 8, ("--foo",)), + ], +) +def test_the_seal_hands_the_server_its_expert_parallel_flag(framework, args, ep, expected): + """EP > 1 reaches the server's argv exactly once, in the framework's own spelling.""" + assert _seal_with_ep(framework, args, ep).argv == expected + + +@pytest.mark.parametrize("pinned", ["--ep-size 2", "--expert-parallel-size=8", "--ep 1", "--ep-size"]) +def test_an_expert_parallel_pin_that_contradicts_ep_fails_the_seal(pinned): + """Neither the operator's pin nor the CLI's EP may win silently.""" + with pytest.raises(ValueError, match="EP=4"): + _seal_with_ep("sglang", pinned, 4) + + +def test_a_contradicting_pin_fails_the_seal_even_when_ep_does_not_divide_tp(): + """The skip for an EP that TP cannot hold must not let a different pinned size launch unreported.""" + with pytest.raises(ValueError, match="EP=3"): + seal_server_argv({"EXTRA_SGLANG_ARGS": "--ep-size 2", "EP": 3, "TP": 8}, "sglang") + + +def test_sealing_twice_adds_the_expert_parallel_flag_once(): + """The seal reads its own output back as a pin.""" + envs = {"EXTRA_SGLANG_ARGS": "--tp 8", "EP": 8} + seal_server_argv(envs, "sglang") + + assert seal_server_argv(envs, "sglang").argv == ("--tp", "8", "--ep-size", "8") + + +def test_a_resealed_argv_keeps_the_expert_parallel_flag(tmp_path): + """A repair rewrites the argument string, not the run's EP.""" + config = tmp_path / "bench.yaml" + config.write_text( + yaml.safe_dump( + {"benchmark": {"framework": "sglang", "model": "/m", "envs": {"EP": 4, "EXTRA_SGLANG_ARGS": "--bogus x"}}} + ), + encoding="utf-8", + ) + + assert reseal_config_argv(config, "--tp 4").argv == ("--tp", "4", "--ep-size", "4") + assert config_server_argv(config).argv == ("--tp", "4", "--ep-size", "4") + + +@pytest.mark.parametrize(("variant_mode", "base_mode"), [("replace", "append"), ("append", "replace")]) +def test_a_replacing_grid_variant_keeps_the_expert_parallel_flag(tmp_path, variant_mode, base_mode): + """Replacing the argument string does not replace the run's EP.""" + base = tmp_path / "base.yaml" + base.write_text( + yaml.safe_dump( + { + "benchmark": { + "framework": "vllm", + "model": "/m", + "envs": {"TP": 8, "EP": 8, "EXTRA_VLLM_ARGS": "--enable-expert-parallel --block-size 64"}, + } + } + ), + encoding="utf-8", + ) + variant = GridVariant(name="v", extra_server_args="--max-num-seqs 64", args_mode=variant_mode) + + config = _grid_runner._build_variant_yaml( + base, "--max-num-seqs 32", variant, output_subdir=tmp_path / "v", base_args_mode=base_mode + ) + + argv = config_server_argv(config).argv + assert "--block-size" not in argv + assert argv.count("--enable-expert-parallel") == 1 + + +@pytest.mark.parametrize( + ("nodes", "tp", "ep"), + [("1", 8, 3), ("2", 16, 16)], + ids=["ep-does-not-divide-tp", "multi-node"], +) +def test_a_resealed_argv_gets_no_expert_parallel_flag_the_launch_cannot_take(tmp_path, monkeypatch, nodes, tp, ep): + """A repair goes through the same EP rule as the first seal.""" + monkeypatch.setenv("INFERENCE_OPTIMIZER_NODES", nodes) + config = tmp_path / "bench.yaml" + config.write_text( + yaml.safe_dump( + {"benchmark": {"framework": "sglang", "model": "/m", "envs": {"TP": tp, "EP": ep, "EXTRA_SGLANG_ARGS": ""}}} + ), + encoding="utf-8", + ) + + assert reseal_config_argv(config, "--mem-fraction-static 0.9").argv == ("--mem-fraction-static", "0.9") + + +def test_a_grid_variant_ep_that_does_not_divide_tp_adds_no_flag(tmp_path, caplog): + """A variant's extra_envs reach the seal unfiltered, so the seal is where EP meets TP.""" + base = tmp_path / "base.yaml" + base.write_text( + yaml.safe_dump( + { + "benchmark": { + "framework": "sglang", + "model": "/m", + "envs": {"TP": 8, "EXTRA_SGLANG_ARGS": "--mem-fraction-static 0.9"}, + } + } + ), + encoding="utf-8", + ) + variant = GridVariant(name="v", extra_envs={"EP": "3"}) + + with caplog.at_level("WARNING"): + config = _grid_runner._build_variant_yaml(base, "", variant, output_subdir=tmp_path / "v") + + assert "--ep-size" not in config_server_argv(config).argv + assert "EP=3 does not divide the launched TP=8" in caplog.text diff --git a/src/hyperloom/inference_optimizer/tests/test_workload_envs.py b/src/hyperloom/inference_optimizer/tests/test_workload_envs.py index ada5452881..5e3b7ac0a0 100644 --- a/src/hyperloom/inference_optimizer/tests/test_workload_envs.py +++ b/src/hyperloom/inference_optimizer/tests/test_workload_envs.py @@ -5,12 +5,15 @@ from __future__ import annotations +import json import sys from types import SimpleNamespace import pytest import yaml +from hyperloom.common.env_safety import OPERATOR_EXTRA_ENV_VAR +from hyperloom.inference_optimizer.framework_registry import server_args_env_name from hyperloom.orchestrator.actions.executors import _workload_envs as we @@ -272,79 +275,118 @@ def test_server_args_merge_existing(monkeypatch, tmp_path): assert "chunked-prefill-size" in merged -def test_vllm_ep_injects_expert_parallel(monkeypatch, tmp_path): +def _ep_materialize(monkeypatch, tmp_path, *, framework, args, tp, ep): + """Materialize ``framework`` at ``tp``/``ep`` and return its benchmark envs.""" _clear_env(monkeypatch) monkeypatch.setenv("INFERENCE_OPTIMIZER_DISABLE_TP_CLAMP", "1") - monkeypatch.setenv("EP", "8") - src = _write(tmp_path / "cfg.yaml", framework="vllm", envs={"EXTRA_VLLM_ARGS": "--trust-remote-code"}) + monkeypatch.setenv("TP", str(tp)) + monkeypatch.setenv("EP", str(ep)) + src = _write(tmp_path / "cfg.yaml", framework=framework, envs={server_args_env_name(framework): args}) + return _materialize(src, tmp_path / "out")["envs"] - bench = _materialize(src, tmp_path / "out") - args = bench["envs"]["EXTRA_VLLM_ARGS"] - assert "--trust-remote-code" in args - assert "--enable-expert-parallel" in args +def test_materialize_hands_sglang_its_expert_parallel_size(monkeypatch, tmp_path): + envs = _ep_materialize(monkeypatch, tmp_path, framework="sglang", args="--mem-fraction-static 0.9", tp=8, ep=4) + argv = envs["EXTRA_SGLANG_ARGS"].split() + assert envs["EP"] == 4 + assert "--mem-fraction-static" in argv + assert argv.count("--ep-size") == 1 + assert argv[argv.index("--ep-size") + 1] == "4" -def test_vllm_ep_one_does_not_inject_expert_parallel(monkeypatch, tmp_path): - _clear_env(monkeypatch) - monkeypatch.setenv("INFERENCE_OPTIMIZER_DISABLE_TP_CLAMP", "1") - monkeypatch.setenv("EP", "1") - src = _write(tmp_path / "cfg.yaml", framework="vllm", envs={"EXTRA_VLLM_ARGS": "--trust-remote-code"}) - bench = _materialize(src, tmp_path / "out") +@pytest.mark.parametrize("pinned", ["--ep-size 4", "--expert-parallel-size=4"]) +def test_materialize_keeps_an_operator_pinned_expert_parallel_size(monkeypatch, tmp_path, pinned): + monkeypatch.setenv("INFERENCE_OPTIMIZER_SERVER_ARGS", pinned) - args = bench["envs"]["EXTRA_VLLM_ARGS"] - assert "--trust-remote-code" in args - assert "--enable-expert-parallel" not in args + envs = _ep_materialize(monkeypatch, tmp_path, framework="sglang", args="", tp=8, ep=4) + flags = [tok.split("=", 1)[0] for tok in envs["EXTRA_SGLANG_ARGS"].split()] + assert sum(flag in {"--ep-size", "--expert-parallel-size", "--ep"} for flag in flags) == 1 -def test_vllm_ep_preserves_profile_args(monkeypatch, tmp_path): - _clear_env(monkeypatch) - monkeypatch.setenv("INFERENCE_OPTIMIZER_DISABLE_TP_CLAMP", "1") - monkeypatch.setenv("EP", "8") - src = _write( - tmp_path / "cfg.yaml", - framework="vllm", - envs={"EXTRA_VLLM_ARGS": "--profiler-config.ignore_frontend True"}, - ) - bench = _materialize(src, tmp_path / "out") +@pytest.mark.parametrize( + "args", + ["--trust-remote-code", "--profiler-config.ignore_frontend True", "--enable-expert-parallel"], +) +def test_materialize_enables_vllm_expert_parallel_once(monkeypatch, tmp_path, args): + envs = _ep_materialize(monkeypatch, tmp_path, framework="vllm", args=args, tp=8, ep=8) + + assert args in envs["EXTRA_VLLM_ARGS"] + assert envs["EXTRA_VLLM_ARGS"].split().count("--enable-expert-parallel") == 1 + + +def test_materialize_ep_one_adds_no_expert_parallel_flag(monkeypatch, tmp_path): + envs = _ep_materialize(monkeypatch, tmp_path, framework="vllm", args="--trust-remote-code", tp=8, ep=1) - args = bench["envs"]["EXTRA_VLLM_ARGS"] - assert "--profiler-config.ignore_frontend True" in args - assert "--enable-expert-parallel" in args + assert envs["EXTRA_VLLM_ARGS"] == "--trust-remote-code" -def test_vllm_ep_does_not_duplicate_existing_flag(monkeypatch, tmp_path): +def test_materialize_drops_an_ep_the_clamped_tp_cannot_hold(monkeypatch, tmp_path, caplog): _clear_env(monkeypatch) - monkeypatch.setenv("INFERENCE_OPTIMIZER_DISABLE_TP_CLAMP", "1") + monkeypatch.setenv("INFERENCE_OPTIMIZER_VISIBLE_GPU_COUNT", "4") + monkeypatch.setenv("TP", "8") monkeypatch.setenv("EP", "8") - src = _write( - tmp_path / "cfg.yaml", - framework="vllm", - envs={"EXTRA_VLLM_ARGS": "--enable-expert-parallel"}, - ) + src = _write(tmp_path / "cfg.yaml", envs={"EXTRA_SGLANG_ARGS": "--mem-fraction-static 0.9"}) - bench = _materialize(src, tmp_path / "out") + with caplog.at_level("WARNING"): + envs = _materialize(src, tmp_path / "out")["envs"] + + assert envs["TP"] == 4 + assert envs["EP"] == 8 + assert "--ep-size" not in envs["EXTRA_SGLANG_ARGS"] + assert "EP=8 does not divide the launched TP=4" in caplog.text - args = bench["envs"]["EXTRA_VLLM_ARGS"] - assert args.count("--enable-expert-parallel") == 1 + +def test_materialize_drops_an_ep_that_does_not_divide_tp(monkeypatch, tmp_path, caplog): + with caplog.at_level("WARNING"): + envs = _ep_materialize(monkeypatch, tmp_path, framework="sglang", args="--mem-fraction-static 0.9", tp=8, ep=3) + + assert envs["EP"] == 3 + assert "--ep-size" not in envs["EXTRA_SGLANG_ARGS"] + assert "EP=3 does not divide the launched TP=8" in caplog.text @pytest.mark.parametrize( - "server_args, framework, ep, expected", + ("environ", "extra_envs"), [ - ("--foo", "vllm", 8, "--foo --enable-expert-parallel"), - ("--foo", "vllm", 1, "--foo"), - ("--foo", "sglang", 8, "--foo"), - ("--foo", "vllm", "bad", "--foo"), - ("--foo", "vllm", None, "--foo"), - ("--enable-expert-parallel", "vllm", 8, "--enable-expert-parallel"), - ("", "vllm", 8, "--enable-expert-parallel"), + # ``--extra-env EP=3`` serializes the pin and also exports it under its own name. + ({OPERATOR_EXTRA_ENV_VAR: json.dumps({"EP": "3"}), "EP": "3"}, None), + ({}, {"EP": "3"}), ], + ids=["operator-pin", "extra-envs"], ) -def test_inject_vllm_expert_parallel_unit(server_args, framework, ep, expected): - assert we.inject_vllm_expert_parallel(server_args, framework, ep) == expected +def test_an_extra_env_ep_that_does_not_divide_tp_adds_no_flag(monkeypatch, tmp_path, caplog, environ, extra_envs): + """An EP written by the extra-env merge meets the same TP rule as the run's own EP.""" + _clear_env(monkeypatch) + monkeypatch.setenv("INFERENCE_OPTIMIZER_DISABLE_TP_CLAMP", "1") + monkeypatch.setenv("TP", "8") + for name, value in environ.items(): + monkeypatch.setenv(name, value) + src = _write(tmp_path / "cfg.yaml", envs={"EXTRA_SGLANG_ARGS": "--mem-fraction-static 0.9"}) + + with caplog.at_level("WARNING"): + envs = _materialize(src, tmp_path / "out", extra_envs=extra_envs)["envs"] + + assert "--ep-size" not in envs["EXTRA_SGLANG_ARGS"].split() + assert "EP=3 does not divide the launched TP=8" in caplog.text + + +@pytest.mark.parametrize("operator_pin", [False, True]) +def test_materialize_leaves_multi_node_ep_to_the_launcher(monkeypatch, tmp_path, operator_pin): + """The multi-node launcher passes EP itself; this pod's GPU count is not the cluster's TP.""" + _clear_env(monkeypatch) + monkeypatch.setenv("INFERENCE_OPTIMIZER_NODES", "2") + monkeypatch.setenv("INFERENCE_OPTIMIZER_VISIBLE_GPU_COUNT", "8") + monkeypatch.setenv("TP", "16") + monkeypatch.setenv("EP", "16") + if operator_pin: + monkeypatch.setenv(OPERATOR_EXTRA_ENV_VAR, json.dumps({"EP": "16"})) + src = _write(tmp_path / "cfg.yaml", envs={"EXTRA_SGLANG_ARGS": "--mem-fraction-static 0.9"}) + + envs = _materialize(src, tmp_path / "out")["envs"] + + assert "--ep-size" not in envs["EXTRA_SGLANG_ARGS"].split() def test_mimo_v2_injects_triton_attention(monkeypatch, tmp_path): diff --git a/src/hyperloom/inference_optimizer/tests/test_workload_envs_branches_unit.py b/src/hyperloom/inference_optimizer/tests/test_workload_envs_branches_unit.py index ff6dee9fd7..dc6cd50947 100644 --- a/src/hyperloom/inference_optimizer/tests/test_workload_envs_branches_unit.py +++ b/src/hyperloom/inference_optimizer/tests/test_workload_envs_branches_unit.py @@ -292,81 +292,6 @@ def test_server_args_merge_existing(monkeypatch, tmp_path): assert "chunked-prefill-size" in merged -def test_vllm_ep_injects_expert_parallel(monkeypatch, tmp_path): - _clear_env(monkeypatch) - monkeypatch.setenv("INFERENCE_OPTIMIZER_DISABLE_TP_CLAMP", "1") - monkeypatch.setenv("EP", "8") - src = _write(tmp_path / "cfg.yaml", framework="vllm", envs={"EXTRA_VLLM_ARGS": "--trust-remote-code"}) - - bench = _materialize(src, tmp_path / "out") - - args = bench["envs"]["EXTRA_VLLM_ARGS"] - assert "--trust-remote-code" in args - assert "--enable-expert-parallel" in args - - -def test_vllm_ep_one_does_not_inject_expert_parallel(monkeypatch, tmp_path): - _clear_env(monkeypatch) - monkeypatch.setenv("INFERENCE_OPTIMIZER_DISABLE_TP_CLAMP", "1") - monkeypatch.setenv("EP", "1") - src = _write(tmp_path / "cfg.yaml", framework="vllm", envs={"EXTRA_VLLM_ARGS": "--trust-remote-code"}) - - bench = _materialize(src, tmp_path / "out") - - args = bench["envs"]["EXTRA_VLLM_ARGS"] - assert "--trust-remote-code" in args - assert "--enable-expert-parallel" not in args - - -def test_vllm_ep_preserves_profile_args(monkeypatch, tmp_path): - _clear_env(monkeypatch) - monkeypatch.setenv("INFERENCE_OPTIMIZER_DISABLE_TP_CLAMP", "1") - monkeypatch.setenv("EP", "8") - src = _write( - tmp_path / "cfg.yaml", - framework="vllm", - envs={"EXTRA_VLLM_ARGS": "--profiler-config.ignore_frontend True"}, - ) - - bench = _materialize(src, tmp_path / "out") - - args = bench["envs"]["EXTRA_VLLM_ARGS"] - assert "--profiler-config.ignore_frontend True" in args - assert "--enable-expert-parallel" in args - - -def test_vllm_ep_does_not_duplicate_existing_flag(monkeypatch, tmp_path): - _clear_env(monkeypatch) - monkeypatch.setenv("INFERENCE_OPTIMIZER_DISABLE_TP_CLAMP", "1") - monkeypatch.setenv("EP", "8") - src = _write( - tmp_path / "cfg.yaml", - framework="vllm", - envs={"EXTRA_VLLM_ARGS": "--enable-expert-parallel"}, - ) - - bench = _materialize(src, tmp_path / "out") - - args = bench["envs"]["EXTRA_VLLM_ARGS"] - assert args.count("--enable-expert-parallel") == 1 - - -@pytest.mark.parametrize( - "server_args, framework, ep, expected", - [ - ("--foo", "vllm", 8, "--foo --enable-expert-parallel"), - ("--foo", "vllm", 1, "--foo"), - ("--foo", "sglang", 8, "--foo"), - ("--foo", "vllm", "bad", "--foo"), - ("--foo", "vllm", None, "--foo"), - ("--enable-expert-parallel", "vllm", 8, "--enable-expert-parallel"), - ("", "vllm", 8, "--enable-expert-parallel"), - ], -) -def test_inject_vllm_expert_parallel_unit(server_args, framework, ep, expected): - assert we.inject_vllm_expert_parallel(server_args, framework, ep) == expected - - def test_mimo_v2_injects_triton_attention(monkeypatch, tmp_path): _clear_env(monkeypatch) monkeypatch.setenv("INFERENCE_OPTIMIZER_DISABLE_TP_CLAMP", "1") diff --git a/src/hyperloom/orchestrator/actions/executors/_server_argv.py b/src/hyperloom/orchestrator/actions/executors/_server_argv.py index bfd5aed24c..1bfba9750b 100644 --- a/src/hyperloom/orchestrator/actions/executors/_server_argv.py +++ b/src/hyperloom/orchestrator/actions/executors/_server_argv.py @@ -6,6 +6,7 @@ from __future__ import annotations import hashlib +import logging from collections.abc import Mapping, MutableMapping, Sequence from dataclasses import dataclass from pathlib import Path @@ -13,12 +14,16 @@ import yaml -from hyperloom.inference_optimizer.framework_registry import server_args_env_name +from hyperloom.inference_optimizer.framework_registry import expert_parallel_flag, server_args_env_name from hyperloom.inference_optimizer.grid_server_args import merge_server_args from hyperloom.inference_optimizer.grid_server_args import tokenize_server_args_preserving_json from hyperloom.inference_optimizer.grid_server_args import validate_server_args_shell_safe +from ._multi_node_env import is_multi_node + +log = logging.getLogger(__name__) + @dataclass(frozen=True) class ServerArgv: @@ -85,12 +90,66 @@ def add_server_arg_unless_pinned( return True +def _pinned_values(tokens: Sequence[str], aliases: Sequence[str]) -> list[str | None]: + """Return the value of every whole-token occurrence of ``aliases``; ``None`` where one has none.""" + values: list[str | None] = [] + for index, token in enumerate(tokens): + name, assigned, value = token.partition("=") + if name not in aliases: + continue + if assigned: + values.append(value) + elif index + 1 < len(tokens) and not tokens[index + 1].startswith("-"): + values.append(tokens[index + 1]) + else: + values.append(None) + return values + + +def _expert_parallel_arg(text: str, framework: str | None, envs: Mapping[str, Any]) -> str: + """Return the flag ``text`` still owes the server to run at ``envs["EP"]``; empty when it owes none. + + The flag shards experts across the ``envs["TP"]`` ranks this config launches with, so an EP that does not divide + them gets none. A multi-node server is launched by the multi-node launcher, which passes EP itself. + + Raises: + ValueError: When ``text`` pins an expert-parallel size other than ``envs["EP"]``. + """ + size = int(envs.get("EP") or 1) + row = expert_parallel_flag(framework) + if row is None or size <= 1 or is_multi_node(): + return "" + flag, aliases, sized = row + split = tokenize_server_args_preserving_json(text) + pinned = _pinned_values(split[1] if split else [], aliases) + if sized and any(value is None or not value.isdigit() or int(value) != size for value in pinned): + raise ValueError( + f"EP={size} but the server args already pin an expert-parallel size of {pinned} " + f"({' / '.join(aliases)}); drop the pin or pass --ep to match it" + ) + if pinned: + return "" + tp = int(envs.get("TP") or size) + if tp % size: + log.warning( + "EP=%d does not divide the launched TP=%d (after any clamp to the visible GPUs); " + "launching without expert parallelism.", + size, + tp, + ) + return "" + return f"{flag} {size}" if sized else flag + + def seal_server_argv( envs: MutableMapping[str, Any], framework: str | None, ) -> ServerArgv: """Write the final server argument string into ``envs`` and return its argv. + The run's ``EP`` env is honoured here, below every composer, so no + replacement of the argument string can drop the expert-parallel flag. + Args: envs: The benchmark env mapping being materialised. framework: The framework the config serves. @@ -99,10 +158,12 @@ def seal_server_argv( ServerArgv: The sealed argv. Raises: - ValueError: When the composed string carries shell control syntax. + ValueError: When the composed string carries shell control syntax, or + pins an expert-parallel size that contradicts ``envs["EP"]``. """ env_name = server_args_env_name(framework) text = validate_server_args_shell_safe(str(envs.get(env_name) or "")) + text = merge_server_args(text, _expert_parallel_arg(text, framework, envs)) sealed = _sealed(framework, env_name, text) if sealed.text or env_name in envs: envs[env_name] = sealed.text diff --git a/src/hyperloom/orchestrator/actions/executors/_workload_envs.py b/src/hyperloom/orchestrator/actions/executors/_workload_envs.py index 5de77f5a20..b498fe3d5a 100644 --- a/src/hyperloom/orchestrator/actions/executors/_workload_envs.py +++ b/src/hyperloom/orchestrator/actions/executors/_workload_envs.py @@ -977,26 +977,6 @@ def _client_tokenizer_mode(model_path: str | None) -> str: return model_type -def inject_vllm_expert_parallel( - server_args: str | None, - framework: Any, - ep: Any, -) -> str: - """Append vLLM expert-parallel flag when EP is enabled.""" - args = str(server_args or "").strip() - if "vllm" not in str(framework or "").lower(): - return args - try: - ep_int = int(ep if ep not in (None, "") else 1) - except (TypeError, ValueError): - return args - if ep_int <= 1: - return args - if re.search(r"(?:^|\s)--enable-expert-parallel(?:\s|$)", args): - return args - return f"{args} --enable-expert-parallel".strip() - - class FrameworkScriptMismatchError(ValueError): """Raised when benchmark_script targets a different framework than the run. @@ -1166,7 +1146,7 @@ def _finalize_framework_server_args( drop_moe_runner_backend: bool = False, ) -> None: """Apply the final framework server-arg guard pipeline in place - (context-length/watchdog/attention/MoE/EP/dedup/compact/shell-safe); order is fixed. + (context-length/watchdog/attention/MoE/dedup/compact/shell-safe); order is fixed. 1. --context-length cap: sglang sizes max_total_tokens off the model's max_position_embeddings, so a huge native window balloons the aiter @@ -1211,11 +1191,6 @@ def _finalize_framework_server_args( # crashed the server. if drop_moe_runner_backend and framework_env == "EXTRA_SGLANG_ARGS": resolved_server_args = _remove_moe_runner_backend_arg(resolved_server_args) - resolved_server_args = inject_vllm_expert_parallel( - resolved_server_args, - bench.get("framework"), - os.environ.get("EP", "").strip() or envs.get("EP"), - ) # 5. vLLM/atom argparse dedup: collapse repeated single-value flags to # last-wins (vLLM crashes EngineCoreProc on a duplicate); no-op for # sglang. @@ -1475,6 +1450,7 @@ def materialize_config_with_envs( "OSL", "MAX_MODEL_LEN", "TP", + "EP", "PORT", ): val = os.environ.get(env_key, "").strip()