From a63753816c4e0a8108f9c6d02b33b05ddb83acf2 Mon Sep 17 00:00:00 2001 From: Zhehao Hu Date: Tue, 20 Oct 2026 17:06:28 +0000 Subject: [PATCH 1/2] specdec-release: add training and serving readiness gates for drafters Add the specdec-release skill with two stdlib-only gates that follow the day0-release contract (JSON verdict; exit 0 pass, 1 fail, 2 USER_CONFIG_ERROR): - gate_train.py: the drafter trained cleanly (check_regression.py's verdict, with the recipe's thresholds set and their metrics present) and its export fits the target (architecture, projector_type, hidden/vocab size, mask token, target layer ids). Checks use only fields both configs carry, so any target + drafter family works. - gate_smoke.py: classifies vllm_smoke_test.sh's SMOKE_JSON as not servable, wrong drafter/target, speculation inactive, or below the recipe's acceptance-length floor. No per-model performance bar is applied; release criteria come later. Co-Authored-By: Claude Opus 5.5 Signed-off-by: Zhehao Hu --- .claude/skills/specdec-release | 1 + .../modelopt/skills/specdec-release/SKILL.md | 52 +++ .../specdec-release/scripts/gate_smoke.py | 137 ++++++++ .../specdec-release/scripts/gate_train.py | 326 ++++++++++++++++++ .../tests/test_specdec_gates.py | 309 +++++++++++++++++ 5 files changed, 825 insertions(+) create mode 120000 .claude/skills/specdec-release create mode 100644 plugins/modelopt/skills/specdec-release/SKILL.md create mode 100644 plugins/modelopt/skills/specdec-release/scripts/gate_smoke.py create mode 100644 plugins/modelopt/skills/specdec-release/scripts/gate_train.py create mode 100644 plugins/modelopt/skills/specdec-release/tests/test_specdec_gates.py diff --git a/.claude/skills/specdec-release b/.claude/skills/specdec-release new file mode 120000 index 00000000000..36bf6f699ec --- /dev/null +++ b/.claude/skills/specdec-release @@ -0,0 +1 @@ +../../.agents/skills/specdec-release \ No newline at end of file diff --git a/plugins/modelopt/skills/specdec-release/SKILL.md b/plugins/modelopt/skills/specdec-release/SKILL.md new file mode 100644 index 00000000000..d594c308476 --- /dev/null +++ b/plugins/modelopt/skills/specdec-release/SKILL.md @@ -0,0 +1,52 @@ +--- +name: specdec-release +description: Deterministic readiness gates for a speculative-decoding drafter — training/export (gate_train), and serving canary (gate_smoke). Use when the user asks whether a trained drafter (EAGLE3, DFlash, DFlash2, DSpark, Domino, LiLiCorr) is ready for release review, or to gate the artifacts of a draft-training pipeline run. Do NOT use to configure, debug or triage a draft-training run (use speculative-decoding), to quantize (use ptq), or to serve a model (use deployment). +license: Apache-2.0 +--- + +# Spec-Dec Release + +Gates that decide whether a drafter is **ready for release review**: trained cleanly, +exported to fit its target, and servable with speculation active. They set no per-model +performance bar; AL is returned as evidence. The +checkpoint release criteria (Addendum A: AL and throughput bars) are a later stage and are +not implemented here. + +## Gate contract + +Same as `day0-release`. Each script prints JSON `{pass, failure_class, detail, notes, +inputs, ...}` and exits `0` pass, `1` gate failed, `2` unreadable or invalid input +(`USER_CONFIG_ERROR`). Run the gates on local copies of the artifacts or on the cluster +login node; they are stdlib-only. **Never read a verdict from a log or a Slurm state.** + +| Stage | Command | Reads | +| --- | --- | --- | +| Train + export | `python "$SKILL_DIR/scripts/gate_train.py" --regression /regression_check.json --export /exported-checkpoint-final --target-config /config.json --algorithm ` | `check_regression.py --json` output, export dir, target config | +| Smoke | `python "$SKILL_DIR/scripts/gate_smoke.py" --smoke --draft --target ` | `vllm_smoke_test.sh` `SMOKE_JSON` | + +A gate never passes a run its producing script failed, classifies the failure itself only +when the script cannot (so `gate_smoke` reads raw values; `gate_train` keeps +`check_regression.py`'s verdict), and never re-implements a script's comparison. + +Artifact files are read +from host paths (`/cicd//...`), but `--draft`/`--target` are compared to the paths the +job recorded, so pass them as the job saw them (e.g. `/scratchspace/...`). `gate_smoke` +classifies failures from the raw values and its `stage`; a SMOKE_JSON without `stage` comes +from an older script, so re-run the smoke test. + +## Failure classes + +| `failure_class` | Gate | Action | +| --- | --- | --- | +| `USER_CONFIG_ERROR` | any | Fix the invocation (paths, `--algorithm`, unknown family → `--expect-architecture`); re-run the gate only | +| `METRICS_MISSING` | train | No regression verdict, or the recipe set no thresholds. Check the training log; set `MAX_FINAL_LOSS`/`MIN_FINAL_ACC` in the launcher YAML | +| `TRAIN_REGRESSION` | train | The recipe's own floors failed. Triage with the **speculative-decoding** skill; do not serve | +| `EXPORT_MISSING` | train | Re-run the training task to re-export (an export counts only once `config.json` and weights exist) | +| `DRAFT_CONFIG_MISMATCH` | train | Export does not fit the target (architecture, hidden/vocab size, mask token, layer ids). Fix the recipe overrides per the algorithm sheet; retrain | +| `CHECKPOINT_NOT_SERVABLE` | smoke | Same as `day0-release`: fix the export or serving flags; do not benchmark | +| `SPEC_NOT_ACTIVE` | smoke | AL ≤ 1.05: drafter not wired up or useless. Check `SPEC_METHOD`, token count, drafter path | +| `ACCEPTANCE_REGRESSION` | smoke | Serves and speculates, but AL < the recipe's `MIN_ACCEPTANCE_LENGTH`. Triage training with **speculative-decoding**; do not benchmark | +| `SPEC_CONFIG_MISMATCH` | smoke | Served the wrong drafter or target. Fix the submission; re-run | + +Verdict: both gates passed → a servable drafter, ready for a matched benchmark; +training, export or serving failed for good → `INFEASIBLE`. diff --git a/plugins/modelopt/skills/specdec-release/scripts/gate_smoke.py b/plugins/modelopt/skills/specdec-release/scripts/gate_smoke.py new file mode 100644 index 00000000000..78a84f42692 --- /dev/null +++ b/plugins/modelopt/skills/specdec-release/scripts/gate_smoke.py @@ -0,0 +1,137 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Spec-dec serving canary gate: target + drafter load, generate, and speculate. + +Reads the SMOKE_JSON written by tools/launcher/common/specdec/vllm_smoke_test.sh and +classifies the failure from its raw values: the script's own ``pass`` cannot tell a serving +failure from an acceptance-length regression. A failing ``pass`` the values do not explain +still fails the gate. Pure decision logic in ``evaluate_smoke``. +""" + +from __future__ import annotations + +import argparse +import json +import math +import sys + +# Mean acceptance length counts the bonus token, so exactly 1.0 means no draft token was +# ever accepted. Anything at or below this is a drafter that is wired up but useless, or +# not wired up at all; generic, not a performance bar. +SPEC_ON_MIN_AL = 1.05 + +_RANKING = ( + "CHECKPOINT_NOT_SERVABLE", + "SPEC_CONFIG_MISMATCH", + "SPEC_NOT_ACTIVE", + "ACCEPTANCE_REGRESSION", +) + + +def evaluate_smoke(smoke, draft=None, target=None): + """Validate a smoke-test result. + + Args: + smoke: dict from SMOKE_JSON. + draft: expected drafter path or HF id, as the smoke job saw it; optional. + target: expected target model path or HF id, as the smoke job saw it; optional. + + Returns: + dict ``{pass, failure_class, detail, notes, mean_acceptance_length}``. + """ + if "stage" not in smoke: + # Written by a vllm_smoke_test.sh that predates "stage": it cannot say how far the + # run got, so re-run the smoke test rather than guess. + return { + "pass": False, + "failure_class": "USER_CONFIG_ERROR", + "detail": "SMOKE_JSON has no 'stage'; re-run the smoke test with the current script", + "notes": [], + "mean_acceptance_length": smoke.get("mean_acceptance_length"), + } + problems = [] + passed, failed = smoke.get("prompts_passed"), smoke.get("prompts_failed") + stage = smoke["stage"] + if stage != "completed": + problems.append(("CHECKPOINT_NOT_SERVABLE", f"stopped at {stage}: {smoke.get('reason')}")) + elif failed or not passed: + problems.append(("CHECKPOINT_NOT_SERVABLE", f"{passed} prompts passed, {failed} failed")) + for label, expected in (("draft", draft), ("target", target)): + if expected is not None and smoke.get(label) != expected: + problems.append( + ("SPEC_CONFIG_MISMATCH", f"{label} {smoke.get(label)!r} != expected {expected!r}") + ) + al = smoke.get("mean_acceptance_length") + valid = isinstance(al, (int, float)) and not isinstance(al, bool) and math.isfinite(al) + if not valid: + problems.append(("SPEC_NOT_ACTIVE", f"no acceptance length reported ({al!r})")) + elif al <= SPEC_ON_MIN_AL: + problems.append(("SPEC_NOT_ACTIVE", f"acceptance length {al} <= {SPEC_ON_MIN_AL}")) + min_al = smoke.get("min_acceptance_length") + if valid and min_al is not None and al < min_al: + problems.append( + ("ACCEPTANCE_REGRESSION", f"acceptance length {al} < recipe floor {min_al}") + ) + if not problems and smoke.get("pass") is not True: + # Every failure the script knows is re-derived above; one this gate does not + # recognize must still not pass a run whose launcher task failed. + problems.append( + ("CHECKPOINT_NOT_SERVABLE", f"smoke test failed: {smoke.get('reason') or 'no reason'}") + ) + + notes = [] + if min_al is None: + notes.append("recipe set no MIN_ACCEPTANCE_LENGTH; only the generic floor applied") + result = {"mean_acceptance_length": al, "notes": notes} + if not problems: + return {"pass": True, "failure_class": None, "detail": "servable, speculating", **result} + fc = min((c for c, _ in problems), key=_RANKING.index) + return { + "pass": False, + "failure_class": fc, + "detail": "; ".join(f"{c}: {m}" for c, m in problems), + **result, + } + + +def main(argv=None): + """CLI entry point: read SMOKE_JSON and print the verdict.""" + p = argparse.ArgumentParser(description="Spec-dec serving canary gate") + p.add_argument("--smoke", required=True, help="SMOKE_JSON from vllm_smoke_test.sh") + p.add_argument("--draft", help="expected drafter path or HF id") + p.add_argument("--target", help="expected target model path or HF id") + args = p.parse_args(argv) + + try: + with open(args.smoke) as f: + smoke = json.load(f) + if not isinstance(smoke, dict): + raise ValueError(f"{args.smoke} does not hold a JSON object") + except (OSError, ValueError) as e: + error = {"pass": False, "failure_class": "USER_CONFIG_ERROR", "detail": str(e)} + print(json.dumps({**error, "notes": [], "inputs": {"smoke": args.smoke}}, indent=2)) + return 2 + + result = evaluate_smoke(smoke, args.draft, args.target) + result["inputs"] = {"smoke": args.smoke} + print(json.dumps(result, indent=2)) + if result["failure_class"] == "USER_CONFIG_ERROR": + return 2 + return 0 if result["pass"] else 1 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/plugins/modelopt/skills/specdec-release/scripts/gate_train.py b/plugins/modelopt/skills/specdec-release/scripts/gate_train.py new file mode 100644 index 00000000000..dedc1976fa0 --- /dev/null +++ b/plugins/modelopt/skills/specdec-release/scripts/gate_train.py @@ -0,0 +1,326 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Spec-dec training gate: the drafter trained cleanly and its export fits the target. + +Pure decision logic in ``evaluate_training`` (unit-tested without a cluster); ``main`` +reads the files and prints the verdict. Inputs: + + regression_check.json written by tools/launcher/common/check_regression.py --json. Its + thresholds belong to the recipe (launcher YAML); this gate requires + that they were set and passed, and does not apply numbers of its own. + exported drafter: config.json plus *.safetensors. + target config.json the target model's HF config (text_config/llm_config is used when + the target is multimodal). + +Draft-to-target checks use only fields both configs carry, so a drafter family this gate +does not know degrades to notes ("not checked") instead of failing. +""" + +from __future__ import annotations + +import argparse +import json +import math +import os +import sys +from glob import glob + +# Exported "architectures" per algorithm, from modelopt/torch/export/plugins/hf_spec_export.py +# and hf_spec_configs.py. +_ARCHITECTURES = { + "eagle3": {"LlamaForCausalLMEagle3", "Eagle3DeepseekV2ForCausalLM"}, + "dflash": {"DFlashDraftModel"}, + "dflash2": {"DFlash2DraftModel"}, + "dspark": {"DFlashDraftModel"}, + "domino": {"DFlashDraftModel"}, + "lilicorr": {"LiLiCorrDraftModel"}, +} +# Metric each check_regression.py threshold is compared against. +_THRESHOLD_METRICS = { + "MAX_FINAL_LOSS": "loss", + "MIN_FINAL_ACC": "accuracy", + "MAX_FINAL_PERPLEXITY": "perplexity", +} +# dflash_config.projector_type each algorithm's export may carry. DFlash, DSpark and Domino +# all export as DFlashDraftModel, so this is what tells them apart; plain DFlash writes none. +_PROJECTOR_TYPES: dict[str, set[str | None]] = { + "dflash": {None, "dflash"}, + "dflash2": {"dflash2"}, + "dspark": {"dspark"}, + "domino": {"domino"}, + "lilicorr": {"lilicorr"}, +} + +# Highest priority first: report the most fundamental problem. +_RANKING = ( + "USER_CONFIG_ERROR", + "EXPORT_MISSING", + "DRAFT_CONFIG_MISMATCH", + "METRICS_MISSING", + "TRAIN_REGRESSION", +) + + +def _target_text_config(target): + return target.get("text_config") or target.get("llm_config") or target + + +def _token_ids(value): + if value is None: + return set() + return set(value) if isinstance(value, list) else {value} + + +def _check_compatibility(draft, target, problems, notes): + """Draft <-> target checks on fields both configs carry.""" + tgt = _target_text_config(target) + for key in ("hidden_size", "vocab_size"): + if draft.get(key) is None or tgt.get(key) is None: + notes.append(f"{key} not checked (absent from draft or target config)") + elif draft[key] != tgt[key]: + problems.append( + ("DRAFT_CONFIG_MISMATCH", f"{key}: draft {draft[key]} != target {tgt[key]}") + ) + + vocab = tgt.get("vocab_size") + special = {} + for key in ("bos_token_id", "eos_token_id", "pad_token_id"): + for tid in _token_ids(tgt.get(key)): + special[tid] = key + sub = draft.get("dflash_config") or {} + mask = sub.get("mask_token_id") + if mask is None: + notes.append("mask_token_id not checked (absent from draft config)") + elif not isinstance(mask, int) or isinstance(mask, bool) or mask < 0: + problems.append( + ("DRAFT_CONFIG_MISMATCH", f"mask_token_id {mask!r} is not a token id (int >= 0)") + ) + elif vocab is not None and mask >= vocab: + problems.append( + ("DRAFT_CONFIG_MISMATCH", f"mask_token_id {mask} >= target vocab_size {vocab}") + ) + elif mask in special: + problems.append( + ("DRAFT_CONFIG_MISMATCH", f"mask_token_id {mask} is the target's {special[mask]}") + ) + + n_layers = tgt.get("num_hidden_layers") + layer_ids = sub.get("target_layer_ids") + if layer_ids is None: + layer_ids = (draft.get("eagle_config") or {}).get("eagle_aux_hidden_state_layer_ids") + if layer_ids is None or n_layers is None: + notes.append("target layer ids not checked (absent from draft or target config)") + elif not isinstance(layer_ids, list): + problems.append(("DRAFT_CONFIG_MISMATCH", f"target layer ids {layer_ids!r} is not a list")) + else: + bad = [ + i + for i in layer_ids + if not isinstance(i, int) or isinstance(i, bool) or not 0 <= i < n_layers + ] + if bad: + problems.append( + ("DRAFT_CONFIG_MISMATCH", f"target layer ids {bad} outside [0, {n_layers})") + ) + if draft.get("num_target_layers") not in (None, n_layers) and n_layers is not None: + problems.append( + ( + "DRAFT_CONFIG_MISMATCH", + f"num_target_layers {draft['num_target_layers']} != target {n_layers}", + ) + ) + + +def evaluate_training( + regression, export_config, export_weights, target_config, algorithm, expect_architecture=None +): + """Decide whether training ran cleanly and the export fits the target. + + Args: + regression: dict from regression_check.json, or None if the file is absent. + export_config: dict from the export's config.json, or None if absent or unreadable. + export_weights: list of (filename, size_bytes) for the export's *.safetensors. + target_config: dict from the target's config.json. + algorithm: drafter algorithm (key of ``_ARCHITECTURES``), case-insensitive. + expect_architecture: required when ``algorithm`` is not in ``_ARCHITECTURES``. + + Returns: + dict ``{pass, failure_class, detail, notes, recipe_thresholds, metrics}``. + """ + algorithm = (algorithm or "").lower() + problems, notes = [], [] + expected = {expect_architecture} if expect_architecture else _ARCHITECTURES.get(algorithm) + if not expected: + problems.append( + ( + "USER_CONFIG_ERROR", + f"unknown algorithm {algorithm!r}: pass --expect-architecture " + f"(known: {sorted(_ARCHITECTURES)})", + ) + ) + + if regression is None: + problems.append(("METRICS_MISSING", "no regression_check.json")) + regression = {} + elif not isinstance(regression, dict) or not regression: + problems.append(("METRICS_MISSING", "regression_check.json holds no verdict")) + regression = {} + elif not regression.get("thresholds"): + problems.append( + ("METRICS_MISSING", "regression check ran without thresholds (recipe sets none)") + ) + elif regression.get("pass") is not True: + cls = regression.get("failure_class") or "TRAIN_REGRESSION" + if cls not in ("METRICS_MISSING", "TRAIN_REGRESSION"): + cls = "TRAIN_REGRESSION" + problems.append((cls, "; ".join(regression.get("messages") or []) or "check failed")) + else: + # A pass must carry the metric behind every threshold. The comparison itself stays + # check_regression.py's; this only refuses a verdict with nothing behind it. + metrics = regression.get("metrics") + metrics = metrics if isinstance(metrics, dict) else {} + for name in regression["thresholds"]: + key = _THRESHOLD_METRICS.get(name) + value = metrics.get(key) if key else None + if not ( + isinstance(value, (int, float)) + and not isinstance(value, bool) + and math.isfinite(value) + ): + problems.append( + ("METRICS_MISSING", f"passed with no finite '{key or name}' for {name}") + ) + + if export_config is None: + problems.append(("EXPORT_MISSING", "export has no readable config.json")) + elif not any(size > 0 for _, size in export_weights): + problems.append(("EXPORT_MISSING", "export has no non-empty *.safetensors")) + else: + archs = export_config.get("architectures") or [] + if expected and (len(archs) != 1 or archs[0] not in expected): + problems.append( + ("DRAFT_CONFIG_MISMATCH", f"architectures {archs} not one of {sorted(expected)}") + ) + sub = export_config.get("dflash_config") or {} + projectors = _PROJECTOR_TYPES.get(algorithm) + if projectors and sub.get("projector_type") not in projectors: + problems.append( + ( + "DRAFT_CONFIG_MISMATCH", + f"projector_type {sub.get('projector_type')!r} not one of " + f"{sorted(projectors, key=str)}", + ) + ) + top, inner = export_config.get("block_size"), sub.get("block_size") + if top is not None and inner is not None and top != inner: + problems.append(("DRAFT_CONFIG_MISMATCH", f"block_size {top} != dflash_config {inner}")) + _check_compatibility(export_config, target_config, problems, notes) + + result = { + "recipe_thresholds": regression.get("thresholds"), + "metrics": regression.get("metrics"), + "notes": notes, + } + if not problems: + return { + "pass": True, + "failure_class": None, + "detail": "training and export valid", + **result, + } + fc = min((c for c, _ in problems), key=_RANKING.index) + return { + "pass": False, + "failure_class": fc, + "detail": "; ".join(f"{c}: {m}" for c, m in problems), + **result, + } + + +def _load_json(path, missing_ok=False): + if missing_ok and not os.path.exists(path): + return None + with open(path) as f: + return json.load(f) + + +def _load_object(path): + """A JSON object from path; anything else is a bad input (ValueError).""" + obj = _load_json(path) + if not isinstance(obj, dict): + raise ValueError(f"{path} does not hold a JSON object") + return obj + + +def main(argv=None): + """CLI entry point: read the training artifacts and print the verdict.""" + p = argparse.ArgumentParser(description="Spec-dec training gate") + p.add_argument("--regression", required=True, help="regression_check.json") + p.add_argument("--export", required=True, help="exported drafter directory") + p.add_argument("--target-config", required=True, help="target model config.json") + p.add_argument("--algorithm", required=True, help=f"one of {sorted(_ARCHITECTURES)}") + p.add_argument("--expect-architecture", help="required for an algorithm not listed above") + args = p.parse_args(argv) + inputs = { + "regression": args.regression, + "export": args.export, + "target_config": args.target_config, + } + + # A missing export or verdict inside an existing training output dir is what a failed + # run leaves behind; a missing output dir is a mistyped path, and "re-run training" is + # the wrong advice for a typo. + for flag, path in (("--export", args.export), ("--regression", args.regression)): + parent = os.path.dirname(os.path.abspath(path)) + if not os.path.isdir(parent): + error = { + "pass": False, + "failure_class": "USER_CONFIG_ERROR", + "detail": f"directory of {flag} {parent} does not exist", + } + print(json.dumps({**error, "notes": [], "inputs": inputs}, indent=2)) + return 2 + try: + target = _load_object(args.target_config) + regression = _load_json(args.regression, missing_ok=True) + except (OSError, ValueError) as e: + error = {"pass": False, "failure_class": "USER_CONFIG_ERROR", "detail": str(e)} + print(json.dumps({**error, "notes": [], "inputs": inputs}, indent=2)) + return 2 + try: + # A missing export dir reads as a missing config.json: EXPORT_MISSING. + export_config = _load_json(os.path.join(args.export, "config.json"), missing_ok=True) + except (OSError, ValueError): + export_config = None # a truncated config.json is a broken export, not a bad input + if not isinstance(export_config, dict): + export_config = None + weights = [ + (os.path.basename(w), os.path.getsize(w)) + for w in glob(os.path.join(args.export, "*.safetensors")) + ] + + result = evaluate_training( + regression, export_config, weights, target, args.algorithm, args.expect_architecture + ) + result["inputs"] = inputs + print(json.dumps(result, indent=2)) + if result["failure_class"] == "USER_CONFIG_ERROR": + return 2 + return 0 if result["pass"] else 1 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/plugins/modelopt/skills/specdec-release/tests/test_specdec_gates.py b/plugins/modelopt/skills/specdec-release/tests/test_specdec_gates.py new file mode 100644 index 00000000000..ffd8e84520f --- /dev/null +++ b/plugins/modelopt/skills/specdec-release/tests/test_specdec_gates.py @@ -0,0 +1,309 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Unit tests for the spec-dec readiness gates. + +These are deterministic — no GPU, cluster, or network. Inputs carry the values of a real +Qwen3-8B DFlash2 run (2000 steps on GB300), trimmed to the fields the gates read. Run with: + + python -m pytest "$SKILL_DIR/tests/test_specdec_gates.py" +""" + +import copy +import json +import sys +from pathlib import Path + +import pytest + +sys.path.insert(0, str(Path(__file__).parent.parent / "scripts")) + +from gate_smoke import evaluate_smoke +from gate_smoke import main as smoke_main +from gate_train import evaluate_training +from gate_train import main as train_main + +TARGET = "/hf-local/Qwen/Qwen3-8B" +DRAFT = "/cicd/m0-drafts/qwen3-8b-dflash2/exported-checkpoint-2000" +SERVED_DRAFT = "/scratchspace/dflash2_bs16/exported-checkpoint-final" +WEIGHTS = [("model.safetensors", 1 << 20)] + + +def _regression(): + """check_regression.py --json output of the training run.""" + return { + "pass": True, + "failure_class": None, + "metrics": {"loss": 6.5775, "accuracy": 0.1767}, + "thresholds": {"MAX_FINAL_LOSS": 7.0, "MIN_FINAL_ACC": 0.15}, + "messages": ["PASS: loss=6.578", "PASS: accuracy=0.177"], + } + + +def _export_config(): + """The exported DFlash2 drafter's config.json, fields the gate reads.""" + return { + "architectures": ["DFlash2DraftModel"], + "block_size": 16, + "hidden_size": 4096, + "vocab_size": 151936, + "num_target_layers": 36, + "dflash_config": { + "mask_token_id": 151669, + "target_layer_ids": [1, 9, 17, 25, 33], + "projector_type": "dflash2", + "block_size": 16, + }, + } + + +def _target_config(): + """Qwen/Qwen3-8B config.json, fields the gate reads.""" + return { + "hidden_size": 4096, + "vocab_size": 151936, + "num_hidden_layers": 36, + "bos_token_id": 151643, + "eos_token_id": 151645, + } + + +def _smoke(): + """vllm_smoke_test.sh SMOKE_JSON of the same drafter.""" + return { + "stage": "completed", + "pass": True, + "reason": None, + "target": TARGET, + "draft": SERVED_DRAFT, + "num_speculative_tokens": 7, + "prompts_passed": 8, + "prompts_failed": 0, + "mean_acceptance_length": 1.79, + "min_acceptance_length": 1.5, + } + + +def _with(obj, path, value): + """A copy of obj with one nested key replaced (value None deletes it).""" + obj = copy.deepcopy(obj) + *parents, leaf = path + node = obj + for key in parents: + node = node[key] + if value is None: + del node[leaf] + else: + node[leaf] = value + return obj + + +# ── gate_train ──────────────────────────────────────────────────────── + + +def _train(algorithm="dflash2", weights=WEIGHTS, **overrides): + inputs = { + "regression": _regression(), + "export_config": _export_config(), + "target_config": _target_config(), + } + inputs.update(overrides) + return evaluate_training( + inputs["regression"], inputs["export_config"], weights, inputs["target_config"], algorithm + ) + + +def test_train_real_run_passes(): + r = _train() + assert r["pass"], r["detail"] + assert r["recipe_thresholds"] == {"MAX_FINAL_LOSS": 7.0, "MIN_FINAL_ACC": 0.15} + + +@pytest.mark.parametrize( + ("overrides", "failure_class"), + [ + ({"regression": None}, "METRICS_MISSING"), + ({"regression": _with(_regression(), ["thresholds"], {})}, "METRICS_MISSING"), + ({"regression": ["not", "an", "object"]}, "METRICS_MISSING"), + ({"regression": _with(_regression(), ["pass"], False)}, "TRAIN_REGRESSION"), + ({"regression": _with(_regression(), ["pass"], "false")}, "TRAIN_REGRESSION"), + ({"regression": _with(_regression(), ["metrics", "accuracy"], None)}, "METRICS_MISSING"), + ({"export_config": None}, "EXPORT_MISSING"), + ({"weights": [("model.safetensors", 0)]}, "EXPORT_MISSING"), + ({"algorithm": "dflash"}, "DRAFT_CONFIG_MISMATCH"), # DFlash2 export, DFlash requested + ( + { + "algorithm": "dflash", # a DSpark export shares DFlash's architecture + "export_config": _with( + _with(_export_config(), ["architectures"], ["DFlashDraftModel"]), + ["dflash_config", "projector_type"], + "dspark", + ), + }, + "DRAFT_CONFIG_MISMATCH", + ), + ({"algorithm": "medusa"}, "USER_CONFIG_ERROR"), + ( + {"export_config": _with(_export_config(), ["dflash_config", "mask_token_id"], 151645)}, + "DRAFT_CONFIG_MISMATCH", # the target's eos + ), + ( + { + "export_config": _with( + _export_config(), ["dflash_config", "target_layer_ids"], [1, 36] + ) + }, + "DRAFT_CONFIG_MISMATCH", + ), + ( + {"export_config": _with(_export_config(), ["dflash_config", "target_layer_ids"], 9)}, + "DRAFT_CONFIG_MISMATCH", + ), + ( + {"export_config": _with(_export_config(), ["dflash_config", "projector_type"], None)}, + "DRAFT_CONFIG_MISMATCH", # a DFlash-architecture export can't prove it is DFlash2 + ), + ( + {"export_config": _with(_export_config(), ["hidden_size"], 2048)}, + "DRAFT_CONFIG_MISMATCH", + ), + ( + {"export_config": _with(_export_config(), ["dflash_config", "block_size"], 8)}, + "DRAFT_CONFIG_MISMATCH", + ), + ], +) +def test_train_failures(overrides, failure_class): + r = _train(**overrides) + assert not r["pass"] + assert r["failure_class"] == failure_class, r["detail"] + + +def test_train_generic_eagle3_llama_pair_passes(): + """A family with different config fields passes on what both configs carry.""" + target = { + "hidden_size": 4096, + "vocab_size": 128256, + "num_hidden_layers": 32, + "eos_token_id": [128001, 128009], + } + export = { + "architectures": ["LlamaForCausalLMEagle3"], + "hidden_size": 4096, + "vocab_size": 128256, + "draft_vocab_size": 32000, + "eagle_config": {"eagle_aux_hidden_state_layer_ids": [1, 15, 28]}, + } + r = evaluate_training(_regression(), export, WEIGHTS, target, "eagle3") + assert r["pass"], r["detail"] + assert any("mask_token_id not checked" in n for n in r["notes"]) + + +def test_train_inherited_qwen_mask_token_on_smaller_vocab_fails(): + """A DSpark export keeping the Qwen3 mask id on a 128K-vocab target trains on garbage.""" + target = {"hidden_size": 4096, "vocab_size": 128256, "num_hidden_layers": 32} + export = _with(_export_config(), ["architectures"], ["DFlashDraftModel"]) + export.update(vocab_size=128256, num_target_layers=32) + export["dflash_config"].update(projector_type="dspark", target_layer_ids=[1, 8, 15, 22, 29]) + r = evaluate_training(_regression(), export, WEIGHTS, target, "dspark") + assert r["failure_class"] == "DRAFT_CONFIG_MISMATCH" + assert "mask_token_id 151669 >= target vocab_size 128256" in r["detail"] + + +# ── gate_smoke ──────────────────────────────────────────────────────── + + +@pytest.mark.parametrize( + ("change", "kwargs", "failure_class"), + [ + ({}, {}, None), + ({"pass": False, "prompts_passed": 7, "prompts_failed": 1}, {}, "CHECKPOINT_NOT_SERVABLE"), + ({"prompts_passed": 0}, {}, "CHECKPOINT_NOT_SERVABLE"), + ( + {"stage": "server_died", "pass": False, "mean_acceptance_length": None}, + {}, + "CHECKPOINT_NOT_SERVABLE", + ), + ( + {"pass": False, "reason": "a failure the gate does not know"}, + {}, + "CHECKPOINT_NOT_SERVABLE", + ), + ({"stage": None}, {}, "USER_CONFIG_ERROR"), # deleted below: predates "stage" + ({"mean_acceptance_length": 1.0}, {}, "SPEC_NOT_ACTIVE"), + ({"mean_acceptance_length": None}, {}, "SPEC_NOT_ACTIVE"), + ({"pass": False, "mean_acceptance_length": 1.4}, {}, "ACCEPTANCE_REGRESSION"), + ({"mean_acceptance_length": 1.5}, {}, None), # the script's floor is inclusive too + ({}, {"draft": SERVED_DRAFT, "target": TARGET}, None), + ({}, {"draft": DRAFT}, "SPEC_CONFIG_MISMATCH"), # a host path is not what was served + ], +) +def test_smoke(change, kwargs, failure_class): + smoke = {**_smoke(), **change} + if smoke["stage"] is None: + del smoke["stage"] + r = evaluate_smoke(smoke, **kwargs) + assert r["failure_class"] == failure_class, r["detail"] + assert r["pass"] == (failure_class is None) + + +# ── CLIs ────────────────────────────────────────────────────────────── + + +def _write(path, obj): + path.write_text(json.dumps(obj)) + return str(path) + + +def test_cli_exit_codes(tmp_path, capsys): + export = tmp_path / "export" + export.mkdir() + _write(export / "config.json", _export_config()) + (export / "model.safetensors").write_bytes(b"\0" * 8) + common = [ + "--regression", _write(tmp_path / "regression_check.json", _regression()), + "--export", str(export), + "--target-config", _write(tmp_path / "target.json", _target_config()), + ] # fmt: skip + assert train_main([*common, "--algorithm", "dflash2"]) == 0 + assert train_main([*common, "--algorithm", "dspark"]) == 1 # projector_type mismatch + assert train_main([*common, "--algorithm", "medusa"]) == 2 + bad_target = tmp_path / "bad_target.json" + bad_target.write_text("{") + capsys.readouterr() + assert train_main([*common, "--target-config", str(bad_target), "--algorithm", "dflash2"]) == 2 + assert json.loads(capsys.readouterr().out)["inputs"]["target_config"] == str(bad_target) + (export / "config.json").write_text("{") # truncated mid-write: a broken export + assert train_main([*common, "--algorithm", "dflash2"]) == 1 + _write(export / "config.json", []) # not an object: also a broken export + assert train_main([*common, "--algorithm", "dflash2"]) == 1 + assert ( + train_main([*common, "--target-config", _write(bad_target, []), "--algorithm", "dflash2"]) + == 2 + ) + # A missing directory is a mistyped path; a missing file in an existing one is a failed run. + _write(export / "config.json", _export_config()) + nope = str(tmp_path / "nope" / "exported-checkpoint-final") + assert train_main([*common, "--export", nope, "--algorithm", "dflash2"]) == 2 + crashed = str(tmp_path / "exported-checkpoint-final") # training died before the export + assert train_main([*common, "--export", crashed, "--algorithm", "dflash2"]) == 1 + missing = str(tmp_path / "nope" / "regression_check.json") + assert train_main([*common, "--regression", missing, "--algorithm", "dflash2"]) == 2 + missing = str(tmp_path / "regression_missing.json") + assert train_main([*common, "--regression", missing, "--algorithm", "dflash2"]) == 1 + + for not_an_object in ([], "stage", None): + bad = _write(tmp_path / "bad.json", not_an_object) + assert smoke_main(["--smoke", bad]) == 2 From 17cfcfe8045342acb192306edfda1ab00692bf5e Mon Sep 17 00:00:00 2001 From: Zhehao Hu Date: Fri, 23 Oct 2026 00:08:21 +0000 Subject: [PATCH 2/2] specdec-release: gate_smoke reports a rejected smoke config as USER_CONFIG_ERROR vllm_smoke_test.sh ([2/8]) now stops with stage invalid_config when MIN_ACCEPTANCE_LENGTH or NUM_SPEC_TOKENS is not a valid number. gate_smoke filed every stage other than completed as CHECKPOINT_NOT_SERVABLE, which blames the drafter for a typo in the recipe. It now returns USER_CONFIG_ERROR with the script's reason and says to fix the recipe and re-run the smoke task. Signed-off-by: Zhehao Hu --- .../skills/specdec-release/scripts/gate_smoke.py | 13 ++++++++++++- .../specdec-release/tests/test_specdec_gates.py | 5 +++++ 2 files changed, 17 insertions(+), 1 deletion(-) diff --git a/plugins/modelopt/skills/specdec-release/scripts/gate_smoke.py b/plugins/modelopt/skills/specdec-release/scripts/gate_smoke.py index 78a84f42692..6c2f2a71f81 100644 --- a/plugins/modelopt/skills/specdec-release/scripts/gate_smoke.py +++ b/plugins/modelopt/skills/specdec-release/scripts/gate_smoke.py @@ -62,9 +62,20 @@ def evaluate_smoke(smoke, draft=None, target=None): "notes": [], "mean_acceptance_length": smoke.get("mean_acceptance_length"), } + stage = smoke["stage"] + if stage == "invalid_config": + # The script rejected a recipe value before serving anything: nothing about the + # drafter is known, and re-running unchanged fails the same way. + return { + "pass": False, + "failure_class": "USER_CONFIG_ERROR", + "detail": f"smoke task config: {smoke.get('reason')}; fix it in the recipe and " + "re-run the smoke task", + "notes": [], + "mean_acceptance_length": smoke.get("mean_acceptance_length"), + } problems = [] passed, failed = smoke.get("prompts_passed"), smoke.get("prompts_failed") - stage = smoke["stage"] if stage != "completed": problems.append(("CHECKPOINT_NOT_SERVABLE", f"stopped at {stage}: {smoke.get('reason')}")) elif failed or not passed: diff --git a/plugins/modelopt/skills/specdec-release/tests/test_specdec_gates.py b/plugins/modelopt/skills/specdec-release/tests/test_specdec_gates.py index ffd8e84520f..53145fbf7a1 100644 --- a/plugins/modelopt/skills/specdec-release/tests/test_specdec_gates.py +++ b/plugins/modelopt/skills/specdec-release/tests/test_specdec_gates.py @@ -242,6 +242,11 @@ def test_train_inherited_qwen_mask_token_on_smaller_vocab_fails(): "CHECKPOINT_NOT_SERVABLE", ), ({"stage": None}, {}, "USER_CONFIG_ERROR"), # deleted below: predates "stage" + ( + {"stage": "invalid_config", "pass": False, "mean_acceptance_length": None}, + {}, + "USER_CONFIG_ERROR", + ), ({"mean_acceptance_length": 1.0}, {}, "SPEC_NOT_ACTIVE"), ({"mean_acceptance_length": None}, {}, "SPEC_NOT_ACTIVE"), ({"pass": False, "mean_acceptance_length": 1.4}, {}, "ACCEPTANCE_REGRESSION"),