diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 40696d8344c..88efc552be5 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -38,6 +38,7 @@ Changelog *Misc* - The ``evaluation`` agent skill can stop looping ("doom loop") reasoning early: ask it to avoid doom loops and it adds vLLM's ``repetition_detection`` request field (vLLM 0.17.0 or later) to every task, ending a request once its output repeats. Use the same setting for baseline and candidate runs. +- Add the ``specdec-release`` agent skill: it trains, smoke-tests, freezes and A|B-benchmarks a speculative-decoding drafter with a gate after each stage, resumes from a state file, and returns a readiness verdict (not a publish decision). **Backward Breaking Changes** diff --git a/plugins/modelopt/skills/specdec-release/SKILL.md b/plugins/modelopt/skills/specdec-release/SKILL.md index 8501f9a900e..e562cd87e8b 100644 --- a/plugins/modelopt/skills/specdec-release/SKILL.md +++ b/plugins/modelopt/skills/specdec-release/SKILL.md @@ -1,60 +1,324 @@ --- name: specdec-release -description: Deterministic readiness gates for a speculative-decoding drafter — training/export (gate_train), serving canary (gate_smoke), matched A|B benchmark validity (gate_bench), and a non-blocking greedy-match report (gate_correctness). 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). +description: Deterministic end-to-end driver that takes a target model and a speculative-decoding algorithm (DFlash, DFlash2; not yet EAGLE3, DSpark, Domino or LiLiCorr) to a drafter readiness verdict — chains draft training → export → serving smoke test → freeze → matched A|B benchmark with enforced gates between stages, resumes from a state file, and returns DRAFTER_READY / INFEASIBLE / ANOMALOUS. Use when the user asks to "train a drafter for model X and tell me if it's ready", "run the spec-dec release flow", "resume the spec-dec run", or to gate the artifacts of a draft-training pipeline run. Do NOT use for single-stage requests — configuring, debugging or triaging a draft-training run (use speculative-decoding), serving only (use deployment), or quantizing (use ptq). 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, servable with speculation active, and measured validly. They -set no per-model performance bar. AL, speedup and TPS are returned as evidence. The -checkpoint release criteria (Addendum A: AL and throughput bars) are a later stage and are -not implemented here. +Drive a target model from no drafter to a **readiness verdict** for a trained drafter, in a +fixed sequence with a gate after every stage. This skill is a **conductor**: it submits the +existing launcher pipelines and enforces the gates. It does not re-implement training, +serving or benchmarking. -## Gate contract +**Goal:** a drafter that trained cleanly, exported to fit its target, serves with speculation +active, and was measured validly against the target alone. `DRAFTER_READY` means ready for +**release review**, not publishable: the gates set no per-model performance bar, and AL, +speedup and TPS are evidence. The release criteria, drafter PTQ and publishing are later +stages and are not implemented here. -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.** +## When to use -| 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\|B bench | `python "$SKILL_DIR/scripts/gate_bench.py" --summary /ab_summary.json --concurrencies 1,8,32 --repeats N --algorithm DFLASH --num-spec-tokens 7 --draft --target ` | `compare_ab.py` summary | -| Correctness | `python "$SKILL_DIR/scripts/gate_correctness.py" --summary /ab_summary.json` | same summary; **report-only**, always passes | +Use only for the full goal-driven flow, or to gate a pipeline run's artifacts. For a single +stage, route to the domain skill: configure a new model, review or triage a training run → +**speculative-decoding**; serve → **deployment**; quantize → **ptq**. + +## Inputs + +Resolve these before starting (ask the user for anything missing): + +- **Target** — HF handle or `/hf-local/...` path; Step 1 stages a handle the job host lacks. +- **Algorithm** — one with launcher examples: `tools/launcher/examples///hf_online_.yaml` + (training, export, smoke) and `specdec_ab_.yaml` (A|B bench). No training example → + configure one with **speculative-decoding** first; a missing bench example is added in Step 4. + Training must run through `common/specdec/dflash_online_training.sh`, the only script that + writes `regression_check.json` (EAGLE3 and DSpark train with `train_eagle*.sh`), and the + algorithm must serve on vLLM (Domino has no path; LiLiCorr serves only on SGLang). Refuse + any other algorithm before submitting anything. +- **Cluster / launcher** — from `clusters.yaml` (see the common skill's + `environment-setup.md`). It fixes the **job host**, which holds the launcher's ``: + the login node with `$SLURM_JOB_DIR`, or this machine with `tools/launcher/local_experiments` + for local Docker (submit with `hf_local=`). +- **Budget** — `dev` (the example as committed) or `full` (the overrides its header lists). + Never invent training knobs. +- **`existing_drafter`** (optional) — an exported drafter to freeze and benchmark instead of + training one. +- **`approve_submissions`** — `each` (default): before every cluster job (launcher, + `sbatch`/`srun`, weight staging), retries included and whoever submits it, show the + resolved command, job count, node-hours and reason, and wait for an OK. An OK covers that + exact job once: a retry, or a changed partition, image or command, needs a new one (an + attempt that created no job, e.g. an SSH error, may be repeated once unchanged). A + sub-agent's brief says: submit at most once; on failure, report back, never retry. `none`: + submit within the retry limits below. Either way, log every attempt, even one that created + no job. +- **Run dir** — the working directory, unless the user names another. Don't create a + subdirectory for the run. + +## The chain + +```text +setup ─▶ train pipeline (data → train+export → smoke) ─▶ freeze ─▶ A|B bench ─▶ closeout + │ │ │ + gate_train + gate_smoke │ gate_bench + gate_correctness (report) + freeze_drafter +``` + +Training, export and the smoke test are one launcher experiment, so the smoke test serves the +unfrozen export; the bench serves the frozen copy, which freeze proves byte-identical to the +export at freeze time. Bench in a second experiment, only after both training gates pass. + +Run each stage, then its gate. **Do not advance past a failed gate.** Copy this checklist and +track progress: + +```text +- [ ] Step 0: Resolve inputs; create the run dir, specdec_state.json, work_log.md +- [ ] Step 1: Setup gate — creds present, cluster reachable, job host python3 ≥ 3.8 +- [ ] Step 2: Train pipeline → gate_train.py → gate_smoke.py +- [ ] Step 3: Freeze → freeze_drafter.py +- [ ] Step 4: A|B bench → gate_bench.py → gate_correctness.py (report) +- [ ] Step 5: Closeout — SPECDEC_REPORT.md + readiness verdict +``` + +### Running a gate + +Gates and `freeze_drafter.py` print JSON `{pass, failure_class, detail, ...}` and exit `0` +pass, `1` failed, `2` bad input (`USER_CONFIG_ERROR`: fix the invocation or the recipe value it +names, not the drafter). Run +them on the job host so artifacts are read in place: + +```bash +# local Docker +python "$SKILL_DIR/scripts/.py" > "$RUN_DIR/gates/.json" +# remote (scripts synced in Step 1). Every Bash call is a new shell, so load the cluster +# in each one, on resume too: +source "/remote_exec.sh" +remote_load_cluster && remote_check_ssh +remote_run "python3 specdec-gates/.py " > "$RUN_DIR/gates/.json" +``` + +If `remote_load_cluster` or `remote_check_ssh` fails, stop with `SYSTEMIC`. Don't fall back to +raw `ssh`. + +Trust the exit code first: any exit other than `0`/`1`/`2`, or output that is not one JSON +object (`remote_run` also captures the remote stderr), means the script crashed — re-run it; +never parse a verdict out of it. + +**Never read a verdict from a log or a job state** — without `REGRESSION_STRICT=1`, a +failed regression check still ends `COMPLETED`. + +**Paths.** Arguments that name a file the gate reads take job-host paths: the job's +`/scratchspace/` is `/cicd//`, and `/cicd/` is `/cicd/`. +`--draft`/`--target` are matched as text against what the job recorded, so pass the exact +string it was submitted with (e.g. `/hf-local/Qwen/Qwen3-8B`, not `Qwen/Qwen3-8B`). + +### Step 0 — Run dir and state + +Create `specdec_state.json` (schema under **State and resume**), `work_log.md` and `gates/` +in the run dir. If the state file exists, this is a **resume**: follow its rules. + +### Step 1 — Setup gate + +Use the common skill's `credentials.md` and `remote-execution.md` to confirm credentials and +cluster reachability. On a remote job host, `remote_sync_to "$SKILL_DIR/scripts" specdec-gates` +and check its `python3 -V` (≥ 3.8, for the gates). If any fails, stop with `SYSTEMIC` — do not +submit training. + +**Stage the target** (here until it moves to the common skill, for **ptq** and +**day0-release** too). Jobs see `hf_local` (`$SLURM_HF_LOCAL` on Slurm) at `/hf-local`. If +`///config.json` reads, use it; a copy that lists but can't be read is +missing, not a reason to hunt for another. Otherwise submit one download job: a CPU partition +from `clusters.yaml`, `HF_TOKEN` from the environment and never echoed, `huggingface_hub` +installed if the image lacks it, and `snapshot_download(, revision=, +local_dir=//)`. Check `config.json`, the tokenizer and every shard the +index lists, and record the revision in `inputs`. A 401/403 on a gated model is +`USER_CONFIG_ERROR`: ask the user to accept its license. + +### Step 2 — Train pipeline + +If the example has no smoke task (`vllm_smoke_test.sh`), add one modeled on +`examples/Qwen/Qwen3-8B/hf_online_dflash2.yaml` task_2. + +**Overrides.** Use `pipeline.task_.slurm_config.=`, and append to `args` or +`environment` with `'pipeline.task_.args+=["key=value"]'` or +`'pipeline.task_.environment+=[{"KEY":"value"}]'` (a later entry wins). Dropping an entry +(e.g. the `full` budget's `training.max_steps`) needs an edited copy of the example in the +run dir, passed to `--yaml`. + +Fit every submission, smoke and bench included, to the cluster per **speculative-decoding**'s +`references/stages/configure.md` (Step 3). Record those as overrides, not budget changes: +they keep the global batch. + +From `tools/launcher`, preview, then rerun with `--yes detach=true` in place of `--to-yaml` +(asking first if `approve_submissions` is `each`): + +```bash +uv run launch.py --yaml examples///hf_online_.yaml \ + pipeline.global_vars.output_dir=/scratchspace/ [hf_local=] \ + --to-yaml /train.yaml +``` -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. +`launch.py` exits `0` even when submission fails, so check its output for `Unexpected error`, +not the exit code. Tasks queued before the failure stay queued: on Slurm, find them with +`squeue`/`sacct` and cancel them before resubmitting. Record the job ids as `submitted` only +once all are confirmed. Never resubmit into a run's `output_dir`: a retrain is a new run, with +a new `run_id`, run dir and state file. Gate with +`` = `/cicd//`: -Pass `gate_bench` what the run was **submitted** with, not targets. 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/...`). An empty -summary is `USER_CONFIG_ERROR`: check `--summary`; if the path is right, the bench job died -before writing any arm, so resubmit as for `BENCH_INCOMPLETE`. `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. +```bash +# gate_train.py +--regression /regression_check.json --export /exported-checkpoint-final \ + --target-config /config.json --algorithm +# gate_smoke.py (after gate_train passes) +--smoke /smoke_check.json --draft /scratchspace//exported-checkpoint-final \ + --target +``` -## Failure classes +**Smoke-only submission.** `/smoke.yaml` is the example's smoke task alone, with +`DRAFT_MODEL: ` replacing `DRAFT_CKPT_DIR` and +`output_dir=/scratchspace/`; gate it with `--draft `, the string +the smoke JSON records. Use it: + +- with `existing_drafter` (skips training, not its gates), after the `gate_train.py` above + with `--export-only` in place of `--regression` and `--export `; +- to resubmit a smoke test (Triage), with + `DRAFT_MODEL=/cicd///exported-checkpoint-final`. + +### Step 3 — Freeze + +```bash +# freeze_drafter.py (mkdir -p the destination's parent first) +--src /exported-checkpoint-final --dst /cicd/drafts/// +# with existing_drafter: --src +``` + +It copies to a `.partial` sibling, re-hashes it, writes `SHA256SUMS`, makes it world-readable +and its files read-only, then renames it into place. Record the frozen path, checksums and +`manifest_sha256` in the state. + +### Step 4 — A|B bench + +If the target has no `specdec_ab_.yaml`, add one modeled on +`examples/Qwen/Qwen3-8B/specdec_ab_dflash2.yaml`, taking the target, `NUM_SPEC_TOKENS`, +container and GPU count from the smoke task that passed. Preview and submit as in Step 2: + +```bash +uv run launch.py --yaml examples///specdec_ab_.yaml \ + pipeline.global_vars.draft_model=/cicd/drafts/// \ + pipeline.global_vars.output_dir=/scratchspace/ab [hf_local=] --to-yaml /ab.yaml +``` + +`REPEATS`, `CONCURRENCIES` and `NUM_SPEC_TOKENS` are `environment` entries: change them as in +**Overrides** (Step 2). + +Both arms run in one job on one node. Gate with what the run was **submitted** with, taking +each `` from the `environment` in `/ab.yaml` (the last entry, if repeated), +with `` = `/cicd//ab`: + +```bash +# gate_bench.py +--summary /ab_summary.json --concurrencies --repeats \ + --algorithm --num-spec-tokens \ + --draft /cicd/drafts/// --target +# gate_correctness.py — report only, never blocks; A-vs-B match next to target-alone drift +--summary /ab_summary.json +``` + +`SPEC_ALGO` is `run.py`'s serving name, not the training name (`DFLASH` also serves +DFlash2). An empty summary is +`USER_CONFIG_ERROR`: check `--summary`; if it is right, the job died before writing any arm, +so resubmit as for `BENCH_INCOMPLETE`. + +### Step 5 — Closeout + +Set `verdict` in the state, then write `SPECDEC_REPORT.md` in the run dir from the state, +`gates/` and `work_log.md`: + +- **Setup:** inputs and pins (ModelOpt commit, images, launcher example and overrides), + experiment and job ids. +- **Training:** metrics against the recipe's floors (not checked with `existing_drafter`), + smoke-test AL. +- **Drafter:** frozen path and checksums. +- **Benchmark:** AL, TPS A/B and speedup per concurrency from `gate_bench`'s `evidence`, its + `measurement_setup`, and `gate_correctness`'s greedy exact match next to its target-alone + control (a low match at a matching control is the target's numerics, not the drafter). +- **Verdicts:** every gate's, the `HUMAN_INTERVENTION` count with reasons, and the readiness + verdict. + +State it is **not** a publish decision, and end with "Next: release review" (SPEED-Bench +measurement on GB200, its gates, model card, publish). + +## Watching jobs + +**Start watching in the same turn as every submission**, and re-arm the watcher whenever it +expires, until the result file appears or the jobs end. Never end a turn while a stage is +`submitted`, unless the user asks you to stop. Register Slurm jobs with the **monitor** skill, +but track progress by the result files the gates read (`regression_check.json`, +`smoke_check.json`, `ab_summary.json`), not job states. +Read a task's log (`_/` in the experiment dir) within 10 min of it +starting. A job that ends without its result file is still gated. Poll and run every remote +command through `remote_run`, which runs bash: the login shell may be csh, where a bare +`ssh … '… 2>/dev/null'` breaks. Never `set -x` with secrets in scope. + +## State and resume + +`specdec_state.json`, one per run: + +```json +{"schema": 1, "run_id": "...", "inputs": {...}, + "stages": {"train": {"status": "passed", "experiment_id": "cicd_...", "job_ids": [...], + "attempt": 1, "overrides": [...], + "gates": {"train": {"pass": true, "failure_class": null, "detail": "...", + "file": "gates/train.json"}, "smoke": {...}}, + "artifacts": {"export": "/cicd/cicd_...//exported-checkpoint-final"}}, + "freeze": {"status": "passed", "path": "...", "sha256": {...}, "manifest_sha256": "..."}, + "bench": {"status": "submitted", "experiment_id": "...", "job_ids": [...], + "attempt": 1, "gates": {"bench": {...}, "correctness": {...}}}}, + "verdict": null} +``` + +`status` is `pending`, `submitted`, `passed`, `failed` or `not run`. Each gate entry keeps +`pass`, `failure_class`, `detail` and the path of its full output under `gates/`. `attempt` +counts the stage's submissions. Paths are absolute job-host paths, `` expanded. On +resume: + +- never re-run a `passed` stage; +- **re-watch a `submitted` stage, never resubmit it**, until its result file appears or its + jobs are gone; then gate it; +- use the host paths recorded in the state, never `/scratchspace`; +- `freeze_drafter.py --dst --verify --manifest-sha256 ` before + benchmarking. + +`work_log.md` lines: ` | | ACTION or HUMAN_INTERVENTION | `. +Submissions and approvals are `ACTION`. Any step a human had to take is +`HUMAN_INTERVENTION`, including a nudge to check progress or to start watching. + +## Triage (gate failure → decision) | `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 | +| `USER_CONFIG_ERROR` | any | Fix the invocation (paths, `--algorithm`, unknown family → `--expect-architecture`), or permissions or disk space (freeze); re-run the gate only. From `gate_smoke` with `stage: invalid_config`: fix `MIN_ACCEPTANCE_LENGTH`/`NUM_SPEC_TOKENS` in the YAML and resubmit as a smoke-only submission; training needn't rerun | +| `METRICS_MISSING` | train | No regression verdict (training crashed, or the recipe set no thresholds). Review logs with **speculative-decoding**; set `MAX_FINAL_LOSS`/`MIN_FINAL_ACC` in the YAML | +| `TRAIN_REGRESSION` | train | The recipe's own floors failed. Triage with **speculative-decoding**; never waive the floor | +| `EXPORT_MISSING` | train, freeze | Training died before exporting, or the export is incomplete. Triage the training task; retrain as a new run | +| `DRAFT_CONFIG_MISMATCH` | train | Export does not fit the target (architecture, projector type, hidden/vocab size, mask token, layer ids). Fix the recipe overrides per the algorithm sheet; retrain | +| `CHECKPOINT_NOT_SERVABLE` | smoke | Act on the smoke JSON's `stage`, not the log's last traceback (the script's exit tears the container down, so later errors are noise). `server_timeout`: the server was still starting; resubmit as a smoke-only submission with a larger `SERVER_TIMEOUT` (900 for a large or hybrid target). `server_died`: fix the export or serving flags with **deployment**, or resubmit once unchanged if the log shows a transient fault. 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, bench | Wrong or unrecorded (pass `--block_size`/`--draft_length`) drafter/target/algorithm/token count, or speculation on in arm A / off in arm B. Fix the submission; re-run | -| `BENCH_INCOMPLETE` | bench | Missing rows or arms, appended results, arms with different response counts. Resubmit into a **fresh** `SAVE_ROOT` (≤ 2 automatic retries) | -| `ARMS_UNMATCHED` | bench | Arms or rows differ in target, engine, GPU, TP, dataset, length or temperature, or an arm ran at another concurrency than its directory says. The comparison is invalid; re-run in one job | +| `SPEC_CONFIG_MISMATCH` | smoke, bench | Wrong or unrecorded drafter, target, algorithm or token count, or speculation on in arm A / off in arm B. Fix the submission; re-run | +| `FREEZE_MISMATCH` | freeze | The copy differs from the source (re-run the freeze), or the destination already holds another drafter (freeze to a new `run_id`; never overwrite) | +| `BENCH_INCOMPLETE` | bench | Missing rows or arms, appended results, arms with different response counts. Resubmit into a **fresh** `output_dir` (≤ 2 automatic retries) | +| `ARMS_UNMATCHED` | bench | Arms or rows differ in target, engine, GPU, TP, dataset, length, temperature or concurrency. The comparison is invalid; re-run in one job | | `ANOMALOUS` | bench | Implausible numbers (AL or speedup above tokens + 1, non-positive TPS, unpaired turns). Do not read the numbers; investigate | -Verdict: every blocking gate passed → `DRAFTER_READY` (not a publish decision); -training, export or serving failed for good → `INFEASIBLE`; the measurement is invalid → -`ANOMALOUS`. +The conductor does not patch and retrain on its own: on a training or serving failure it +stops with the triage finding. `SYSTEMIC` (cluster or credentials down) aborts the run. + +## Output + +Return a verdict, not raw artifacts: + +- `DRAFTER_READY` — every blocking gate passed (`gate_train`, `gate_smoke`, `freeze_drafter`, + `gate_bench`); ready for release review, not a publish decision. +- `INFEASIBLE` — training, export or serving failed and triage found no fix within the run. +- `ANOMALOUS` — the measurement is invalid after the allowed retries, or freeze failed. +- Always: run dir, experiment ids, frozen path, the `HUMAN_INTERVENTION` count. diff --git a/plugins/modelopt/skills/specdec-release/scripts/freeze_drafter.py b/plugins/modelopt/skills/specdec-release/scripts/freeze_drafter.py new file mode 100644 index 00000000000..64034792a17 --- /dev/null +++ b/plugins/modelopt/skills/specdec-release/scripts/freeze_drafter.py @@ -0,0 +1,266 @@ +# 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. + +"""Freeze a drafter export: a read-only, checksummed copy the benchmark serves. + +Copies to a ``.partial`` sibling of , re-hashes it against the source, writes +``SHA256SUMS`` in ``sha256sum -c`` format, makes the tree world-readable and its files +read-only, and renames it to . A ``dst`` already holding the same files is accepted as +is, so a resumed run can re-freeze safely. ``--verify`` re-hashes ``dst`` against its +``SHA256SUMS``, after checking that file against ``--manifest-sha256`` when given (before a +benchmark that follows a resume). Both return ``manifest_sha256``, the hash of ``SHA256SUMS`` +itself, for the caller to record outside the copy. + +Same contract as the gates: JSON ``{pass, failure_class, detail, notes, inputs, ...}``; +exit 0 pass, 1 failed, 2 ``USER_CONFIG_ERROR`` (including unreadable or unwritable paths). +Stdlib only. Pure comparison in ``evaluate_copy``. +""" + +from __future__ import annotations + +import argparse +import hashlib +import json +import os +import re +import shutil +import sys + +SUMS = "SHA256SUMS" + + +def hash_file(path): + """sha256 hex digest of one file, read in 1 MiB chunks.""" + digest = hashlib.sha256() + with open(path, "rb") as f: + for chunk in iter(lambda: f.read(1 << 20), b""): + digest.update(chunk) + return digest.hexdigest() + + +def hash_tree(root): + """{relative path: sha256} of every file under root except SHA256SUMS.""" + hashes = {} + for dirpath, _, filenames in os.walk(root): + for name in filenames: + path = os.path.join(dirpath, name) + rel = os.path.relpath(path, root) + if rel != SUMS: + hashes[rel] = hash_file(path) + return hashes + + +def read_sums(root): + """{relative path: sha256} from root/SHA256SUMS, or None if absent.""" + path = os.path.join(root, SUMS) + if not os.path.exists(path): + return None + sums = {} + with open(path) as f: + for line in f: + if line.strip(): + digest, sep, rel = line.rstrip("\n").partition(" ") + if not sep: + raise ValueError(f"malformed {SUMS} line: {line.rstrip()!r}") + sums[rel] = digest + return sums + + +def evaluate_copy(expected, actual): + """Problems between two {relative path: sha256} maps; empty if identical.""" + problems = [f"missing {rel}" for rel in sorted(set(expected) - set(actual))] + problems += [f"unexpected {rel}" for rel in sorted(set(actual) - set(expected))] + problems += [ + f"checksum differs: {rel}" + for rel in sorted(set(expected) & set(actual)) + if expected[rel] != actual[rel] + ] + return problems + + +def _is_export(hashes, root): + weights = [ + rel + for rel in hashes + if rel.endswith(".safetensors") and os.path.getsize(os.path.join(root, rel)) > 0 + ] + return "config.json" in hashes and bool(weights) + + +def _make_read_only(root): + # Directories stay writable so a failed or rejected copy can still be removed. + for dirpath, _, filenames in os.walk(root): + os.chmod(dirpath, os.stat(dirpath).st_mode | 0o555) + for name in filenames: + path = os.path.join(dirpath, name) + os.chmod(path, (os.stat(path).st_mode | 0o444) & ~0o222) + + +def _write_sums(root, hashes): + with open(os.path.join(root, SUMS), "w") as f: + f.writelines(f"{hashes[rel]} {rel}\n" for rel in sorted(hashes)) + + +def _result(problems, failure_class, detail, **extra): + if problems: + return { + "pass": False, + "failure_class": failure_class, + "detail": "; ".join(problems), + **extra, + } + return {"pass": True, "failure_class": None, "detail": detail, **extra} + + +def freeze(src, dst): + """Copy src to dst with checksums. Returns the verdict dict (without inputs).""" + notes: list[str] = [] + src_hashes = hash_tree(src) + if not _is_export(src_hashes, src): + return _result( + ["source has no config.json and non-empty *.safetensors"], + "EXPORT_MISSING", + "", + notes=notes, + ) + if os.path.exists(dst): + recorded = read_sums(dst) + problems = evaluate_copy(src_hashes, hash_tree(dst)) + if recorded is not None: + problems += [f"{SUMS}: {p}" for p in evaluate_copy(src_hashes, recorded)] + else: + problems.append(f"{dst} exists without {SUMS}") + if problems: + return _result(problems, "FREEZE_MISMATCH", "", notes=notes, path=dst) + notes.append("destination already frozen; not copied again") + return _result( + [], + None, + "already frozen", + notes=notes, + path=dst, + sha256=src_hashes, + manifest_sha256=hash_file(os.path.join(dst, SUMS)), + ) + + partial = dst.rstrip("/") + ".partial" + shutil.rmtree(partial, ignore_errors=True) + shutil.copytree(src, partial) + problems = evaluate_copy(src_hashes, hash_tree(partial)) + if problems: + notes.append(f"bad copy left at {partial}, not published; re-run to retry") + return _result(problems, "FREEZE_MISMATCH", "", notes=notes, path=dst) + _write_sums(partial, src_hashes) + _make_read_only(partial) + os.rename(partial, dst) + return _result( + [], + None, + "frozen", + notes=notes, + path=dst, + sha256=src_hashes, + manifest_sha256=hash_file(os.path.join(dst, SUMS)), + ) + + +def verify(dst, manifest_sha256=None): + """Re-hash dst against its SHA256SUMS, first checking SHA256SUMS against manifest_sha256.""" + sums_path = os.path.join(dst, SUMS) + if not os.path.exists(sums_path): + return _result([f"no {SUMS} in {dst}"], "FREEZE_MISMATCH", "", notes=[], path=dst) + actual_manifest = hash_file(sums_path) + if manifest_sha256 and actual_manifest != manifest_sha256.lower(): + return _result( + [f"{SUMS} differs from the recorded manifest_sha256"], + "FREEZE_MISMATCH", + "", + notes=[], + path=dst, + ) + recorded = read_sums(dst) + actual = hash_tree(dst) + problems = evaluate_copy(recorded, actual) + if not _is_export(actual, dst): + problems.append("no config.json and non-empty *.safetensors") + return _result( + problems, + "FREEZE_MISMATCH", + "verified", + notes=[] if manifest_sha256 else [f"{SUMS} not checked against a recorded manifest_sha256"], + path=dst, + sha256=recorded, + manifest_sha256=actual_manifest, + ) + + +def main(argv=None): + """CLI entry point: freeze (or verify) and print the verdict.""" + p = argparse.ArgumentParser(description="Freeze a drafter export with SHA256 checksums") + p.add_argument("--src", help="exported drafter directory (exported-checkpoint-*)") + p.add_argument("--dst", required=True, help="frozen destination; must not be inside --src") + p.add_argument("--verify", action="store_true", help="only re-hash --dst against SHA256SUMS") + p.add_argument( + "--manifest-sha256", + help="with --verify: the manifest_sha256 freeze returned; SHA256SUMS must hash to it", + ) + args = p.parse_args(argv) + inputs = {"src": args.src, "dst": args.dst, "manifest_sha256": args.manifest_sha256} + + error = None + if args.verify: + if args.src: + error = "--src is not used with --verify" + elif not os.path.isdir(args.dst): + error = f"{args.dst} does not exist" + elif args.manifest_sha256 and not re.fullmatch(r"[0-9a-fA-F]{64}", args.manifest_sha256): + error = "--manifest-sha256 is not a 64-character hex digest" + elif args.manifest_sha256: + error = "--manifest-sha256 is only used with --verify" + elif not args.src or not os.path.isdir(args.src): + error = f"source directory {args.src} does not exist" + else: + real_src, real_dst = os.path.realpath(args.src), os.path.realpath(args.dst) + if real_dst == real_src or real_dst.startswith(real_src + os.sep): + error = "--dst is --src or inside it" + elif not os.path.isdir(os.path.dirname(os.path.abspath(args.dst))): + error = f"parent of --dst {args.dst} does not exist" + elif os.path.isdir(args.dst) and not os.listdir(args.dst): + error = ( + f"destination {args.dst} is an empty directory; " + "remove it or pass the path without creating it" + ) + + if not error: + try: + if args.verify: + result = verify(args.dst, args.manifest_sha256) + else: + result = freeze(args.src, args.dst) + except OSError as e: + error = str(e) + except ValueError as e: + result = _result([str(e)], "FREEZE_MISMATCH", "", notes=[], path=args.dst) + if error: + result = {"pass": False, "failure_class": "USER_CONFIG_ERROR", "detail": error, "notes": []} + result["inputs"] = inputs + print(json.dumps(result, indent=2)) + if 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 index dedc1976fa0..9452c84802a 100644 --- a/plugins/modelopt/skills/specdec-release/scripts/gate_train.py +++ b/plugins/modelopt/skills/specdec-release/scripts/gate_train.py @@ -145,17 +145,25 @@ def _check_compatibility(draft, target, problems, notes): def evaluate_training( - regression, export_config, export_weights, target_config, algorithm, expect_architecture=None + regression, + export_config, + export_weights, + target_config, + algorithm, + expect_architecture=None, + export_only=False, ): """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. + Ignored when ``export_only``. 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``. + export_only: check only the export, for a drafter trained outside this run. Returns: dict ``{pass, failure_class, detail, notes, recipe_thresholds, metrics}``. @@ -172,7 +180,10 @@ def evaluate_training( ) ) - if regression is None: + if export_only: + notes.append("training not checked (--export-only)") + regression = {} + elif regression is None: problems.append(("METRICS_MISSING", "no regression_check.json")) regression = {} elif not isinstance(regression, dict) or not regression: @@ -238,7 +249,7 @@ def evaluate_training( return { "pass": True, "failure_class": None, - "detail": "training and export valid", + "detail": "export valid" if export_only else "training and export valid", **result, } fc = min((c for c, _ in problems), key=_RANKING.index) @@ -268,7 +279,11 @@ def _load_object(path): 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") + source = p.add_mutually_exclusive_group(required=True) + source.add_argument("--regression", help="regression_check.json") + source.add_argument( + "--export-only", action="store_true", help="drafter trained elsewhere: check the export" + ) 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)}") @@ -284,6 +299,8 @@ def main(argv=None): # 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)): + if path is None: + continue parent = os.path.dirname(os.path.abspath(path)) if not os.path.isdir(parent): error = { @@ -295,7 +312,7 @@ def main(argv=None): return 2 try: target = _load_object(args.target_config) - regression = _load_json(args.regression, missing_ok=True) + regression = None if args.export_only else _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)) @@ -313,7 +330,13 @@ def main(argv=None): ] result = evaluate_training( - regression, export_config, weights, target, args.algorithm, args.expect_architecture + regression, + export_config, + weights, + target, + args.algorithm, + args.expect_architecture, + export_only=args.export_only, ) result["inputs"] = inputs print(json.dumps(result, indent=2)) diff --git a/plugins/modelopt/skills/specdec-release/tests/evals.json b/plugins/modelopt/skills/specdec-release/tests/evals.json new file mode 100644 index 00000000000..1efa93bd7f4 --- /dev/null +++ b/plugins/modelopt/skills/specdec-release/tests/evals.json @@ -0,0 +1,86 @@ +[ + { + "name": "full-specdec-flow-triggers", + "skills": [ + "specdec-release" + ], + "query": "Train a DFlash2 drafter for Qwen/Qwen3-8B on my cluster with the dev budget, and tell me whether it's ready for release review.", + "files": [], + "expected_behavior": [ + "Selects the specdec-release skill for the full goal-driven flow", + "Resolves target, algorithm, launcher examples, cluster and budget before submitting anything", + "Asks for approval before each launcher submission (approve_submissions defaults to each)", + "Runs stages in fixed order: setup, train pipeline, gate_train, gate_smoke, freeze, A|B bench, gate_bench, gate_correctness, closeout", + "Reads every verdict from a gate script's JSON, never from a log or a job state", + "Does not benchmark before gate_train and gate_smoke pass", + "Returns DRAFTER_READY / INFEASIBLE / ANOMALOUS and says DRAFTER_READY is not a publish decision" + ] + }, + { + "name": "resume-skips-passed-stages", + "skills": [ + "specdec-release" + ], + "query": "Resume the spec-dec run in this directory.", + "files": [], + "expected_behavior": [ + "Reads specdec_state.json from the run dir before doing anything else", + "Does not re-run a stage whose status is passed", + "Re-watches a submitted stage instead of resubmitting it", + "Re-verifies the frozen drafter with freeze_drafter.py --verify --manifest-sha256 before benchmarking" + ] + }, + { + "name": "existing-drafter-skips-training", + "skills": [ + "specdec-release" + ], + "query": "I already have an exported DFlash2 drafter for Qwen3-8B at /cicd/m0-drafts/qwen3-8b-dflash2/exported-checkpoint-2000. Benchmark it against the target alone and tell me if it's ready for release review.", + "files": [], + "expected_behavior": [ + "Selects specdec-release with existing_drafter set", + "Does not submit a training pipeline", + "Checks the export with gate_train.py --export-only and smoke-tests the drafter with a smoke-only submission, gating it with --draft set to the DRAFT_MODEL value", + "Freezes the drafter with freeze_drafter.py, then benchmarks the frozen copy in one matched A|B job", + "Gates the benchmark with gate_bench using the submitted concurrencies, repeats, algorithm and token count" + ] + }, + { + "name": "train-gate-failure-stops-the-flow", + "skills": [ + "specdec-release", + "speculative-decoding" + ], + "query": "Run the spec-dec release flow for Qwen3-8B DFlash2. The training finished but gate_train returned TRAIN_REGRESSION.", + "files": [], + "expected_behavior": [ + "Stops at the failed gate and does not serve, freeze or benchmark the drafter", + "Triages the regression with the speculative-decoding skill", + "Does not waive or raise the recipe's training thresholds to make the gate pass" + ] + }, + { + "name": "triage-only-does-not-trigger-specdec-release", + "skills": [ + "speculative-decoding" + ], + "query": "Why did my EAGLE3 training job fail? The experiment is cicd_1234.", + "files": [], + "expected_behavior": [ + "Selects the speculative-decoding skill, not specdec-release", + "Does not submit any new job or run the readiness gates" + ] + }, + { + "name": "quantize-only-does-not-trigger-specdec-release", + "skills": [ + "ptq" + ], + "query": "Quantize Qwen3-8B to NVFP4 and save the checkpoint.", + "files": [], + "expected_behavior": [ + "Selects the ptq skill, not specdec-release", + "Does not train a drafter or run a speculative-decoding benchmark" + ] + } +] 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 a683a0b3e02..27197fdc97d 100644 --- a/plugins/modelopt/skills/specdec-release/tests/test_specdec_gates.py +++ b/plugins/modelopt/skills/specdec-release/tests/test_specdec_gates.py @@ -23,7 +23,9 @@ """ import copy +import hashlib import json +import shutil import sys from pathlib import Path @@ -31,6 +33,7 @@ sys.path.insert(0, str(Path(__file__).parent.parent / "scripts")) +from freeze_drafter import main as freeze_main from gate_bench import evaluate_bench from gate_bench import main as bench_main from gate_correctness import evaluate_correctness @@ -430,6 +433,79 @@ def test_correctness_reports_greedy_match_without_gating(): assert r["pass"] and not r["measurable"] and r["notes"] +# ── freeze_drafter ──────────────────────────────────────────────────── + + +def test_freeze_copies_checksums_and_detects_tampering(tmp_path, capsys, monkeypatch): + src = tmp_path / "exported-checkpoint-final" + src.mkdir() + _write(src / "config.json", _export_config()) + (src / "model.safetensors").write_bytes(b"\1" * 64) + dst = tmp_path / "drafts" / "run-1" + dst.parent.mkdir() + args = ["--src", str(src), "--dst", str(dst)] + + capsys.readouterr() + assert freeze_main(args) == 0 + manifest_sha256 = json.loads(capsys.readouterr().out)["manifest_sha256"] + sums = (dst / "SHA256SUMS").read_text().splitlines() + assert sorted(line.split(" ")[1] for line in sums) == ["config.json", "model.safetensors"] + mode = (dst / "model.safetensors").stat().st_mode + assert mode & 0o004 and not mode & 0o222 # world-readable, read-only + assert not (tmp_path / "drafts" / "run-1.partial").exists() + assert freeze_main(args) == 0 # a resumed run re-freezes into the same copy + assert json.loads(capsys.readouterr().out)["manifest_sha256"] == manifest_sha256 + verify = ["--dst", str(dst), "--verify"] + assert freeze_main([*verify, "--manifest-sha256", manifest_sha256]) == 0 + + # Tamper with a weight and rewrite SHA256SUMS to match: only the recorded hash catches it. + (dst / "model.safetensors").chmod(0o644) + (dst / "model.safetensors").write_bytes(b"\2" * 64) + (dst / "SHA256SUMS").chmod(0o644) + (dst / "SHA256SUMS").write_text( + "".join( + f"{hashlib.sha256((dst / rel).read_bytes()).hexdigest()} {rel}\n" + for rel in ("config.json", "model.safetensors") + ) + ) + assert freeze_main(verify) == 0 + assert freeze_main([*verify, "--manifest-sha256", manifest_sha256]) == 1 + capsys.readouterr() + assert freeze_main(args) == 1 # an existing, different copy is never overwritten + assert not json.loads(capsys.readouterr().out)["notes"] + + (dst / "SHA256SUMS").write_text("not a manifest\n") + assert freeze_main(verify) == 1 + assert json.loads(capsys.readouterr().out)["failure_class"] == "FREEZE_MISMATCH" + + empty = tmp_path / "drafts" / "empty" + empty.mkdir() + assert freeze_main(["--src", str(src), "--dst", str(empty)]) == 2 # a pre-created dst + (empty / "SHA256SUMS").write_text("") + assert freeze_main(["--dst", str(empty), "--verify"]) == 1 # an empty manifest is no export + + real_copytree = shutil.copytree + + def corrupting_copytree(from_dir, to_dir): + real_copytree(from_dir, to_dir) + (Path(to_dir) / "model.safetensors").write_bytes(b"\3" * 64) + + monkeypatch.setattr(shutil, "copytree", corrupting_copytree) + run_3 = tmp_path / "drafts" / "run-3" + assert freeze_main(["--src", str(src), "--dst", str(run_3)]) == 1 + assert not run_3.exists() # a bad copy is never published + monkeypatch.undo() + + (src / "model.safetensors").write_bytes(b"") + assert freeze_main(["--src", str(src), "--dst", str(tmp_path / "drafts" / "run-2")]) == 1 + assert freeze_main(["--src", str(tmp_path / "nope"), "--dst", str(dst)]) == 2 + assert freeze_main(["--src", str(src), "--dst", str(src / "inside")]) == 2 + assert freeze_main(["--src", str(src), "--dst", str(src)]) == 2 + assert freeze_main(["--src", str(src), "--dst", str(dst), "--verify"]) == 2 + assert freeze_main([*args, "--manifest-sha256", manifest_sha256]) == 2 + assert freeze_main([*verify, "--manifest-sha256", "abc"]) == 2 + + # ── CLIs ────────────────────────────────────────────────────────────── @@ -474,6 +550,14 @@ def test_cli_exit_codes(tmp_path, capsys): 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 + # An existing drafter has no training run: --export-only checks the export alone. + export_only = [*common[2:], "--export-only"] + assert train_main([*export_only, "--algorithm", "dflash2"]) == 0 + assert train_main([*export_only, "--algorithm", "dspark"]) == 1 + for both_or_neither in ([*common, "--export-only"], common[2:]): + with pytest.raises(SystemExit) as exc: + train_main([*both_or_neither, "--algorithm", "dflash2"]) + assert exc.value.code == 2 for not_an_object in ([], "stage", None): bad = _write(tmp_path / "bad.json", not_an_object) diff --git a/plugins/modelopt/skills/speculative-decoding/SKILL.md b/plugins/modelopt/skills/speculative-decoding/SKILL.md index a6ed38a07ec..0101f15722c 100644 --- a/plugins/modelopt/skills/speculative-decoding/SKILL.md +++ b/plugins/modelopt/skills/speculative-decoding/SKILL.md @@ -7,7 +7,8 @@ description: >- wants experiment logs reviewed, or wants to check whether a run's acceptance rate meets threshold. Triggers on "EAGLE3", "DFlash", "DSpark", "Domino", "draft model", "acceptance rate", "speculative decoding pipeline". Do NOT use for quantizing a - model (use ptq) or serving a checkpoint (use deployment). + model (use ptq), serving a checkpoint (use deployment), or the end-to-end flow that + trains, benchmarks and returns a drafter readiness verdict (use specdec-release). user_invocable: true --- diff --git a/plugins/modelopt/skills/speculative-decoding/references/stages/configure.md b/plugins/modelopt/skills/speculative-decoding/references/stages/configure.md index 59b541bc6e1..a417a6c9b18 100644 --- a/plugins/modelopt/skills/speculative-decoding/references/stages/configure.md +++ b/plugins/modelopt/skills/speculative-decoding/references/stages/configure.md @@ -41,7 +41,15 @@ backend (vLLM / HF / TRT-LLM). The sheet's *Pipeline tasks* section says how to Copy node/GPU counts from the reference example, then sanity-check against the target: the base model's BF16 weights must fit in the allocated GPU memory for the serving and dump tasks, so scale `tensor-parallel-size`, `gpus_per_node`, or `nodes` -if the target is larger than the reference. +if the target is larger than the reference. Then fit every task to the cluster: + +- Clusters that only accept whole nodes reject smaller jobs at `sbatch` (`QOSMinGRES` or + `Requested node configuration is not available`; the launcher's `--mem=0` also trips + clusters that scale RAM per GPU). Set every task's `gpus_per_node` to the node size + (`sinfo -o '%P %G'`); serving tasks keep their TP. +- Training starts one process per visible GPU, so when a training task's GPU count + changes, set `training.gradient_accumulation_steps` to keep nodes × GPUs × per-device + batch × accumulation unchanged. If no integer does, ask the user. ## Step 4 — Preview