diff --git a/configs/nvidia-master.yaml b/configs/nvidia-master.yaml index 74eb03c53..9af864b6d 100644 --- a/configs/nvidia-master.yaml +++ b/configs/nvidia-master.yaml @@ -9064,6 +9064,97 @@ dsv4-fp4-gb200-dynamo-vllm: # MTP2 variant of dsv4-fp4-gb200-dynamo-vllm. Uses the vLLM 0.20.1 image # and hand-picked 8k/1k Pareto points mirrored from NVIDIA/srt-slurm. dsv4-fp4-gb200-dynamo-vllm-mtp2: + image: vllm/vllm-openai:v0.20.1-ubuntu2404 + model: deepseek-ai/DeepSeek-V4-Pro + model-prefix: dsv4 + runner: gb200 + precision: fp4 + framework: dynamo-vllm + multinode: true + disagg: true + scenarios: + fixed-seq-len: + - isl: 8192 + osl: 1024 + search-space: + # Aggregate low latency: TP=8, max-num-seqs=4. + - conc-list: [1] + spec-decoding: mtp + prefill: + num-worker: 1 + tp: 8 + ep: 1 + dp-attn: false + additional-settings: + - "SYNTHETIC_ACCEPTANCE=true" + - "SYNTHETIC_ACCEPTANCE_LENGTH=2.27" + - "CONFIG_FILE=recipes/vllm/deepseek-v4/8k1k/agg-gb200-low-latency-mtp2.yaml" + decode: + num-worker: 0 + tp: 8 + ep: 1 + dp-attn: false + + # Low-latency bridge: 1 prefill (DEP=8) + 4 decode (TP=8), no offload. + - conc-list: [16, 32, 64] + spec-decoding: mtp + prefill: + num-worker: 1 + tp: 8 + ep: 8 + dp-attn: true + additional-settings: + - "SYNTHETIC_ACCEPTANCE=true" + - "SYNTHETIC_ACCEPTANCE_LENGTH=2.27" + - "CONFIG_FILE=recipes/vllm/deepseek-v4/8k1k/disagg-gb200-low-latency-mtp2.yaml" + decode: + num-worker: 4 + tp: 8 + ep: 1 + dp-attn: false + + # MegaMOE mid curve: 1 prefill (DEP=8) + 1 decode (DEP=8). + # 5 nodes total with a dedicated NATS/etcd infra node. + - conc-list: [128, 256, 512, 1024] + spec-decoding: mtp + prefill: + num-worker: 1 + tp: 8 + ep: 8 + dp-attn: true + additional-settings: + - "SYNTHETIC_ACCEPTANCE=true" + - "SYNTHETIC_ACCEPTANCE_LENGTH=2.27" + - "CONFIG_FILE=recipes/vllm/deepseek-v4/8k1k/disagg-gb200-mid-curve-megamoe-mtp2.yaml" + decode: + num-worker: 1 + tp: 8 + ep: 8 + dp-attn: true + + # MegaMOE high throughput: 2 prefill (DEP=8 each) + 1 decode (DEP=8). + # 7 nodes total with a dedicated NATS/etcd infra node. + - conc-list: [1024] + spec-decoding: mtp + prefill: + num-worker: 2 + tp: 8 + ep: 8 + dp-attn: true + additional-settings: + - "SYNTHETIC_ACCEPTANCE=true" + - "SYNTHETIC_ACCEPTANCE_LENGTH=2.27" + - "CONFIG_FILE=recipes/vllm/deepseek-v4/8k1k/disagg-gb200-high-tpt-megamoe-mtp2.yaml" + decode: + num-worker: 1 + tp: 8 + ep: 8 + dp-attn: true + +# Baseline (no synthetic acceptance) variant of dsv4-fp4-gb200-dynamo-vllm-mtp2 +# for before/after Pareto comparison. Identical topology and recipes, only the +# SYNTHETIC_ACCEPTANCE envs are removed. +dsv4-fp4-gb200-dynamo-vllm-mtp2-nosynthetic: image: vllm/vllm-openai:v0.20.1-ubuntu2404 model: deepseek-ai/DeepSeek-V4-Pro model-prefix: dsv4 diff --git a/runners/inject_synthetic_acceptance.py b/runners/inject_synthetic_acceptance.py new file mode 100644 index 000000000..4bda75a9a --- /dev/null +++ b/runners/inject_synthetic_acceptance.py @@ -0,0 +1,164 @@ +#!/usr/bin/env python3 +"""Inject synthetic acceptance parameters into an srt-slurm recipe (generic driver). + +This is the framework-agnostic half of the synthetic-acceptance mechanism. It +decides *whether* to inject (the ``SYNTHETIC_ACCEPTANCE`` flag) and *what* mean +acceptance length to inject, then delegates the actual recipe rewrite to a +framework-specific backend (see ``runners/synthetic_injectors/``). + +The script is a no-op (exit 0, file untouched) when: + - SYNTHETIC_ACCEPTANCE is unset/false, +so existing callers that do not opt in get exactly the previous behavior. When +enabled it requires a backend registered for the given framework; the vLLM +backend is added in a follow-up framework-support change. + +Environment variables: + SYNTHETIC_ACCEPTANCE "true" to enable (default: "false") + SYNTHETIC_ACCEPTANCE_LENGTH target mean acceptance length; if unset, it is + auto-resolved from the reference AL YAML using + MODEL_PREFIX (+ NUM_SPEC_TOKENS / THINKING_MODE) + NUM_SPEC_TOKENS number of speculative tokens (for auto-lookup; + falls back to the value parsed from the recipe) + MODEL_PREFIX model prefix key in the reference YAML (e.g. "dsv4") + THINKING_MODE "thinking_on" / "thinking_off" — only used when the + reference YAML is in the thinking matrix form + (default: "thinking_on") + FRAMEWORK framework key selecting the backend (e.g. + "dynamo-vllm"); may also be passed as argv[2]. + +Usage (from a runner; use an absolute path since runners cd into the srt-slurm +clone before invoking this): + python3 "$GITHUB_WORKSPACE/runners/inject_synthetic_acceptance.py" "${CONFIG_FILE%%:*}" "$FRAMEWORK" +""" + +import os +import sys + +from synthetic_injectors import get_injector + +# MODEL_PREFIX -> top-level key in speedbench-reference-al.yaml. +MODEL_PREFIX_TO_YAML_KEY = { + "dsv4": "deepseek-v4-pro", + "dsr1": "deepseek-r1", +} + + +def _log(msg): + print(f"[Synthetic AR] {msg}") + + +def _yaml_key(model_prefix): + return MODEL_PREFIX_TO_YAML_KEY.get(model_prefix, model_prefix) + + +def _lookup_al(model_block, num_spec_tokens): + """Resolve AL for num_spec_tokens from either reference-YAML shape. + + Flat list form: [ {1: 1.90}, {2: 2.60}, ... ] + Thinking matrix: { thinking_on: {1: ...}, thinking_off: {1: ...} } + """ + # Flat list form (each item is a single-key {level: al} mapping). + if isinstance(model_block, list): + for item in model_block: + if num_spec_tokens in item: + return item[num_spec_tokens] + return None + + if isinstance(model_block, dict): + # Thinking matrix form: pick the requested mode, then index by level. + if any(str(k).startswith("thinking") for k in model_block): + mode = os.environ.get("THINKING_MODE", "thinking_on").strip() or "thinking_on" + mode_block = model_block.get(mode) + if mode_block is None: + sys.exit( + f"ERROR: THINKING_MODE='{mode}' not found in reference YAML " + f"(available: {sorted(model_block)})" + ) + return mode_block.get(num_spec_tokens) + # Plain {level: al} mapping. + return model_block.get(num_spec_tokens) + + return None + + +def _resolve_al(config_text, injector, ref_yaml): + explicit = os.environ.get("SYNTHETIC_ACCEPTANCE_LENGTH", "").strip() + if explicit: + return float(explicit) + + if not os.path.isfile(ref_yaml): + sys.exit( + "ERROR: SYNTHETIC_ACCEPTANCE_LENGTH not set and reference YAML not " + f"found: {ref_yaml}" + ) + + import yaml # local import: only needed on the auto-lookup path + + with open(ref_yaml) as f: + data = yaml.safe_load(f) + + key = _yaml_key(os.environ.get("MODEL_PREFIX", "")) + model_block = data.get(key) + if model_block is None: + sys.exit(f'ERROR: model key "{key}" not found in {ref_yaml}') + + nst_env = os.environ.get("NUM_SPEC_TOKENS", "").strip() + if nst_env: + num_spec_tokens = int(nst_env) + else: + num_spec_tokens = injector.spec_tokens_from_recipe(config_text) or 2 + + al = _lookup_al(model_block, num_spec_tokens) + if al is None: + sys.exit(f"ERROR: num_spec_tokens={num_spec_tokens} not found for {key} in {ref_yaml}") + + _log( + f"Auto-resolved AL={al} from {ref_yaml} " + f"(model={key}, num_spec_tokens={num_spec_tokens})" + ) + return float(al) + + +def inject(config_file, framework): + if os.environ.get("SYNTHETIC_ACCEPTANCE", "false") != "true": + return 0 + + injector = get_injector(framework) + if injector is None: + sys.exit( + "ERROR: SYNTHETIC_ACCEPTANCE=true but no synthetic-acceptance " + f"injector is registered for FRAMEWORK='{framework}'" + ) + + with open(config_file) as f: + content = f.read() + + al = _resolve_al( + content, + injector, + os.path.join(os.path.dirname(__file__), "..", "benchmarks", "speedbench-reference-al.yaml"), + ) + + _log(f"Injecting synthetic acceptance (length={al}) into {config_file}") + + new_content, count = injector.rewrite(content, al, _log) + + if count == 0: + _log("WARNING: No speculative-config entries found to modify; leaving recipe unchanged") + return 0 + + with open(config_file, "w") as f: + f.write(new_content) + _log(f"Modified {count} speculative-config entries") + return 0 + + +def main(argv): + if len(argv) not in (2, 3): + sys.exit("Usage: inject_synthetic_acceptance.py CONFIG_FILE [FRAMEWORK]") + framework = argv[2] if len(argv) == 3 else os.environ.get("FRAMEWORK", "") + return inject(argv[1], framework) + + +if __name__ == "__main__": + sys.exit(main(sys.argv)) diff --git a/runners/launch_gb200-nv.sh b/runners/launch_gb200-nv.sh index ba5fffa83..d125971f3 100755 --- a/runners/launch_gb200-nv.sh +++ b/runners/launch_gb200-nv.sh @@ -52,11 +52,15 @@ elif [[ $FRAMEWORK == "dynamo-vllm" ]]; then export MODEL_PATH="/mnt/lustre01/models/kimi-k2.5-nvfp4" export SRT_SLURM_MODEL_PREFIX="kimi-k2.5-nvfp4" elif [[ $MODEL_PREFIX == "dsv4" && $PRECISION == "fp4" ]]; then - # The FP4 checkpoint is staged on compute-visible Lustre. The former - # /mnt/numa1 path is no longer present on watchtower compute nodes; - # the lowercase Lustre sibling is the FP8 checkpoint, so keep the - # NVFP4 path explicit here. - export MODEL_PATH="/mnt/lustre01/models/DeepSeek-V4-Pro-NVFP4/" + # FP4 checkpoint on compute-visible Lustre (the /mnt/numa1 path is gone + # on watchtower compute nodes). Use the base DeepSeek-V4-Pro checkpoint, + # NOT the -NVFP4 re-quant: the recipe's served identity is plain + # deepseek-ai/DeepSeek-V4-Pro and the pinned v0.20.1 container's + # deepseek_v4 loader doesn't define the NVFP4 export's extra quant + # params (e.g. ffn.experts.w13_input_scale), which KeyErrors at load. + # The lowercase Lustre sibling is the FP8 checkpoint, so name the + # CamelCase FP4 path explicitly (Linux is case-sensitive). + export MODEL_PATH="/mnt/lustre01/models/DeepSeek-V4-Pro" export SRT_SLURM_MODEL_PREFIX="deepseek-v4-pro" elif [[ $MODEL_PREFIX == "minimaxm2.5" && $PRECISION == "fp4" ]]; then export MODEL_PATH="/mnt/lustre01/models/MiniMax-M2.5-NVFP4" @@ -84,8 +88,12 @@ NGINX_IMAGE="nginx:1.27.4" uses_watchtower_shared_fs() { case "$MODEL_PREFIX" in minimaxm2.5|minimaxm3|kimik2.5) return 0 ;; - *) return 1 ;; esac + # dsv4 multinode runs only under dynamo-vllm on watchtower, which likewise + # needs the srt-slurm workspace/outputs on a compute-visible shared FS + # (the runner home is not cross-mounted to compute nodes). + [[ "$FRAMEWORK" == "dynamo-vllm" && "$MODEL_PREFIX" == "dsv4" ]] && return 0 + return 1 } # === Cluster diagnostic probe for watchtower-hosted sweeps === @@ -462,6 +470,11 @@ fi # Keep the Slurm job name aligned with the GitHub runner name. sed -i "s/^name:.*/name: \"${RUNNER_NAME}\"/" "$CONFIG_PATH" +# Optionally inject synthetic acceptance into the recipe's speculative-config +# when SYNTHETIC_ACCEPTANCE=true (no-op otherwise). Must run after the name +# override and before srtctl apply so the rendered job picks it up. +python3 "$GITHUB_WORKSPACE/runners/inject_synthetic_acceptance.py" "$CONFIG_PATH" "$FRAMEWORK" + # Don't leak the login-node venv to the compute-node orchestrator. sbatch's # default --export=ALL propagates VIRTUAL_ENV (set by `source # .venv/bin/activate` above) into job_script_minimal.j2, whose diff --git a/runners/synthetic_injectors/__init__.py b/runners/synthetic_injectors/__init__.py new file mode 100644 index 000000000..5bc4d8a36 --- /dev/null +++ b/runners/synthetic_injectors/__init__.py @@ -0,0 +1,38 @@ +"""Framework-specific synthetic-acceptance injectors. + +The generic driver (``inject_synthetic_acceptance.py``) resolves *whether* to +inject and *what* acceptance length to use, then hands the recipe text to a +framework-specific injector registered here. This keeps the driver +framework-agnostic: adding a new backend (e.g. sglang, trtllm) means dropping a +module in this package and registering it, without touching the driver. + +A backend is any object exposing: + + rewrite(content: str, al: float, log) -> tuple[str, int] + Return the recipe text with the synthetic acceptance parameters applied + and the number of entries modified. ``log`` is a ``callable(str)`` used + for human-readable progress output. + + spec_tokens_from_recipe(content: str) -> int | None + Best-effort extraction of the speculative-token count from the recipe, + used only for reference-AL auto-lookup. May return ``None``. +""" + +# framework value passed by the runner (e.g. "dynamo-vllm") -> backend module. +# Populated by backend modules registering themselves at import time (see the +# `from . import` block below). +_INJECTORS = {} + + +def register(framework, backend): + _INJECTORS[framework] = backend + + +def get_injector(framework): + """Return the backend for ``framework`` (as passed by the runner), or None.""" + return _INJECTORS.get(framework) + + +# Import backends after register/get_injector are defined so each module can +# call register() at import time. Add new frameworks (sglang, trtllm, ...) here. +from . import vllm # noqa: E402,F401 diff --git a/runners/synthetic_injectors/vllm.py b/runners/synthetic_injectors/vllm.py new file mode 100644 index 000000000..e71a5df2b --- /dev/null +++ b/runners/synthetic_injectors/vllm.py @@ -0,0 +1,69 @@ +"""vLLM synthetic-acceptance backend (FRAMEWORK=dynamo-vllm). + +Rewrites every ``speculative-config: ''`` entry in an srt-slurm recipe to +use synthetic rejection sampling: it adds ``rejection_sample_method=synthetic`` +and ``synthetic_acceptance_length=`` to the JSON so the engine emits a +controlled mean acceptance length instead of running the real draft model. + +Registered under the "dynamo-vllm" framework key at import time, so importing +the ``synthetic_injectors`` package is enough for the generic driver to resolve +this backend. +""" + +import json +import re +import sys + +from . import register + +# Matches `speculative-config: ''` (single-quoted JSON, as written in the +# recipe YAML). Capturing the JSON lets us edit it with the json module instead +# of string-munging, so we never produce malformed quoting. +_SPEC_CONFIG_RE = re.compile(r"speculative-config:\s*'([^']+)'") + + +def spec_tokens_from_recipe(text): + """Best-effort: read num_speculative_tokens from the recipe itself.""" + for m in _SPEC_CONFIG_RE.finditer(text): + try: + spec = json.loads(m.group(1)) + except json.JSONDecodeError: + continue + n = spec.get("num_speculative_tokens") + if n: + return int(n) + return 2 + + +def rewrite(content, al, log): + """Rewrite every speculative-config entry to synthetic acceptance. + + Returns ``(new_content, count)`` where count is the number of entries + modified (0 => nothing matched, recipe left unchanged by the driver). + """ + before = [ln.strip() for ln in content.splitlines() if _SPEC_CONFIG_RE.search(ln)] + if before: + log("Before:") + for ln in before: + print(f" {ln}") + + def _replace(match): + spec = json.loads(match.group(1)) + spec["rejection_sample_method"] = "synthetic" + spec["synthetic_acceptance_length"] = al + # Compact separators keep the same style as the hand-written recipes. + return "speculative-config: '" + json.dumps(spec, separators=(",", ":")) + "'" + + new_content, count = _SPEC_CONFIG_RE.subn(_replace, content) + + if count: + after = [ln.strip() for ln in new_content.splitlines() if _SPEC_CONFIG_RE.search(ln)] + if after: + log("After:") + for ln in after: + print(f" {ln}") + + return new_content, count + + +register("dynamo-vllm", sys.modules[__name__])