Repository navigation
[1/8] launcher: make the regression gate machine-readable and opt-in strict - #2839
Conversation
|
Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually. Contributors can view more details about this message here. |
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info
📝 Walkthrough
Merge Risk: ⚪ Minimal · up to
|
| Check name | Status | Explanation |
|---|---|---|
| Description Check |  |
Check skipped - CodeRabbit’s high-level summary is enabled. |
| Title check | |
The title names the main change: the DFlash regression gate is made machine-readable through a JSON verdict and made opt-in strict via REGRESSION_STRICT. The raw summary and objectives both center on … |
| Docstring Coverage | |
Docstring coverage is 95.00% which is sufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Analyzed 20 functions across 4 files. (3 skipped: 3 … |
| Linked Issues check | |
Check skipped because no linked issues were found for this pull request. |
| Out of Scope Changes check | |
Check skipped because no linked issues were found for this pull request. |
| Security Anti-Patterns | |
No specified security anti-pattern was introduced. The PR changes only launcher Python, shell, YAML, Markdown, and tests; no modelopt package or examples Python files changed. Added lines contain no u… |
Full details: Title check
-
Explanation
The title names the main change: the DFlash regression gate is made machine-readable through a JSON verdict and made opt-in strict via REGRESSION_STRICT. The raw summary and objectives both center on this: check_regression.py gains a --json verdict file, dflash_online_training.sh writes regression_check.json, and REGRESSION_STRICT=1 makes a regression fail the task. The '[1/8]' stack prefix is a series marker and does not affect the assessment. The title is short, single-sentence, and specific enough for a reviewer scanning history.
Full details: Docstring Coverage
-
Explanation
Docstring coverage is 95.00% which is sufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Analyzed 20 functions across 4 files. (3 skipped: 3 unsupported.)
Full details: Security Anti-Patterns
-
Explanation
No specified security anti-pattern was introduced. The PR changes only launcher Python, shell, YAML, Markdown, and tests; no modelopt package or examples Python files changed. Added lines contain no unsafe torch.load, numpy.load(..., allow_pickle=True), hardcoded trust_remote_code=True, eval/exec on external input, or # nosec. No pyproject.toml or requirements.txt dependency changes were made.
✨ Finishing Touches
-
📝 Generate docstrings
-
- Commit to this branch
- Create a new PR
🧪 Generate unit tests (beta)
-
- Commit to this branch
- Create a new PR
-
- Autofix · Keep fixing CodeRabbit findings and required CI, and resolving merge conflicts
Comment @coderabbitai help to get the list of available commands.
Codecov Report✅ All modified and coverable lines are covered by tests. Additional details and impacted files@@ Coverage Diff @@
## main #2839 +/- ##
==========================================
- Coverage 79.59% 79.18% -0.42%
==========================================
Files 659 659
Lines 74077 76516 +2439
==========================================
+ Hits 58964 60591 +1627
- Misses 15113 15925 +812
Flags with carried forward coverage won't be shown. Click here to find out more. ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
| checkpoint_states = sorted( | ||
| glob(os.path.join(output_dir, "checkpoint-*", "trainer_state.json")), | ||
| key=_checkpoint_step, | ||
| ) |
There was a problem hiding this comment.
Check checkpoint subdirs first, by numeric step: a lexicographic sort would rank checkpoint-8000 after checkpoint-12000.
| - dflash.dflash_swa_window_size=2048 | ||
| environment: | ||
| - MAX_FINAL_LOSS: "5.0" | ||
| - MAX_FINAL_LOSS: "7.0" |
There was a problem hiding this comment.
Calibrated on this config (global batch 8): a GB300 run finished at step 2000 with loss 6.51 and acc 0.181. The 5.0 inherited from the DFlash example is unreachable here: DFlash2 adds a selector loss term and stops at 2000 steps.
check_regression.py could not stop a bad DFlash training run:
- find_trainer_state() sorted checkpoints as strings, so checkpoint-8000
ranked after checkpoint-12000; it now sorts by step number.
- With thresholds set, a missing or empty trainer_state.json only warned;
it now fails (METRICS_MISSING).
- An empty threshold variable (e.g. MAX_FINAL_LOSS="") crashed on float('');
it now counts as unset.
- dflash_online_training.sh ignored the result (|| true) and the verdict
existed only as log text.
check_regression.py gains --json and prints "Regression check FAILED" on a
failure; its other log lines are unchanged. The training script writes the
verdict to ${OUTPUT_DIR}/regression_check.json and, with REGRESSION_STRICT=1,
fails the task on a regression (default unchanged). hf_online_dflash2.yaml
opts in and recalibrates MAX_FINAL_LOSS from 5.0 (inherited from the DFlash
example) to 7.0: DFlash2 starts from a higher loss, so 5.0 failed healthy
2000-step runs.
Signed-off-by: Zhehao Hu <zhehaoh@nvidia.com>
915d296 to
60fce3f
Compare
Zhehao-Hu
changed the title
|
/claude review |
cjluo-nv
left a comment
cjluo-nv
left a comment
There was a problem hiding this comment.
Bot review (gpt-6.1-sol) — DM the bot to share feedback.
Changes requested: non-finite training metrics can produce an invalid JSON verdict precisely when downstream gates need to read a failed run.
Needs action:
- Fix non-finite metric handling in
tools/launcher/common/check_regression.pyand extend the CLI verdict test to cover NaN and infinity.
No action needed:
- Existing test edits are justified by passing thresholds explicitly; their assertions remain intact.
- Numeric checkpoint ordering and opt-in strict-mode trap wiring look correct. Tests were inspected, not executed.
|
|
||
| if args.json_path: | ||
| with open(args.json_path, "w") as f: | ||
| json.dump(verdict, f, indent=2) |
There was a problem hiding this comment.
Bot comment.
json.dump defaults to allow_nan=True. If a failed training run logs NaN or infinity (Python's JSON reader accepts these, and float() also accepts their string forms), the new artifact contains bare NaN/Infinity, which are not valid JSON and are rejected by strict downstream parsers. Non-finite metrics should produce a failed verdict with JSON-safe metric values (for example, omit them or use null), rather than an unreadable artifact. Please validate finiteness before comparison/serialization, use strict serialization as a safeguard, and extend test_cli_exit_code_and_json_verdict with non-finite cases that assert failure and standards-compliant JSON.
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
| # Regression check (uses env vars MAX_FINAL_LOSS, MIN_FINAL_ACC, etc.). The verdict | ||
| # lands next to the exports for downstream gates. A regression only warns unless | ||
| # REGRESSION_STRICT=1, which fails the task (via the ERR trap) so Slurm reports it. | ||
| if [ -n "$OUTPUT_DIR" ]; then | ||
| python3 common/check_regression.py "${OUTPUT_DIR}" || true | ||
| REGRESSION_ARGS=("${OUTPUT_DIR}" --json "${OUTPUT_DIR}/regression_check.json") | ||
| if [ "${REGRESSION_STRICT:-0}" = "1" ]; then | ||
| python3 common/check_regression.py "${REGRESSION_ARGS[@]}" | ||
| else | ||
| python3 common/check_regression.py "${REGRESSION_ARGS[@]}" || true |
There was a problem hiding this comment.
[IMPORTANT Correctness] A crashed training run can still write pass: true to regression_check.json.
error_handler in common/service_utils.sh:27 only sets FAIL=1/FAIL_EXIT=1 — it does not exit, and this script has no set -e. So when accelerate launch (line 120) dies (CUDA OOM, NCCL timeout, a crash during the final save), execution falls straight through to the export loop and then to this gate. The gate reads whatever trainer_state.json survived and, if that checkpoint's metrics happen to clear the thresholds, writes {"pass": true, "failure_class": null}.
Why it matters: the verdict file is the machine-readable contract this PR introduces, and per the description the specdec-release skill ([8/8]) gates training on it. A consumer that reads pass sees a green verdict for a job that failed. The window is narrow but not hypothetical — a crash in the last few hundred steps of the 92K-step recipe leaves a checkpoint well under MAX_FINAL_LOSS. It is exactly the failure class this PR is closing (an unchecked/bad run stamped PASS), just arriving through the training exit status rather than through a missing trainer_state.json.
Suggested fix — record the training status in the verdict so the gate cannot be read out of context, e.g. skip the check and let the ERR trap stand:
if [ -n "$OUTPUT_DIR" ]; then
if [ "$FAIL_EXIT" = "1" ]; then
# Training (or export) already failed; a surviving checkpoint must not
# produce a passing verdict.
python3 -c 'import json,sys; json.dump({"pass": False, "failure_class": "TRAIN_FAILED",
"trainer_state": None, "metrics": {}, "thresholds": {}, "messages":
["REGRESSION: training task failed before the regression check"]},
open(sys.argv[1], "w"), indent=2)' "${OUTPUT_DIR}/regression_check.json"
else
REGRESSION_ARGS=("${OUTPUT_DIR}" --json "${OUTPUT_DIR}/regression_check.json")
...
fi
fiA --failure-class TRAIN_FAILED flag on check_regression.py (writing the verdict and exiting 1) would be cleaner than the inline python3 -c, and keeps the JSON schema in one place.
There was a problem hiding this comment.
Agreed that it's real at this commit. It's fixed in [3/8] #2931, which rewrites this block. That commit checks the accelerate launch status, skips export and the regression check when training failed, and deletes any regression_check.json left from an earlier run before training starts. A crashed run then has no verdict, and the [6/8] gate (#2934) reports that as METRICS_MISSING, never as a pass. The series merges in order, so the [8/8] skill never reads a verdict without that fix. I kept it in [3/8] rather than duplicating it here.
| checkpoint_states = sorted( | ||
| glob(os.path.join(output_dir, "checkpoint-*", "trainer_state.json")), | ||
| key=_checkpoint_step, | ||
| ) | ||
| if checkpoint_states: | ||
| return checkpoint_states[-1] |
There was a problem hiding this comment.
[SUGGESTION] The step-number sort fixes ordering among checkpoints, but the final state in output_dir is still only a fallback — so "latest" can still be wrong.
Any numeric checkpoint-*/trainer_state.json wins over <output_dir>/trainer_state.json, even when the latter is the newer, end-of-training state. That happens whenever max_steps is not a multiple of save_steps: HF Trainer's final save_state() lands in output_dir while the last periodic checkpoint is some steps behind, and the gate then grades the older one. Note hf_online_dflash2.yaml has save_steps=5000 with max_steps=2000, so that example relies entirely on the fallback path today — the preference only bites once save_steps < max_steps (the 92K-step recipe this PR targets).
Since the point of the change is "grade the newest state", consider making the choice uniform by comparing global_step from each candidate file (the HF trainer_state.json records it), or at minimum preferring <output_dir>/trainer_state.json when its global_step is higher than the best checkpoint's. A test with checkpoint-1000/ plus a higher-step output_dir/trainer_state.json would pin the behaviour either way.
There was a problem hiding this comment.
The HF Trainer saves a checkpoint at the final step whatever save_steps is, so the newest numbered checkpoint never lags a completed run's final state. I have tested this example (save_steps=5000, max_steps=2000). It produced checkpoint-2000, whose trainer_state.json has the same global_step (2000) as <output_dir>/trainer_state.json, and the gate graded checkpoint-2000. If training stops early, save_state() never runs, so the output_dir copy isn't newer in that case either.
| verdict["trainer_state"] = trainer_state | ||
| where = trainer_state or f"{output_dir} (no trainer_state.json)" | ||
| try: | ||
| metrics = get_final_metrics(trainer_state) if trainer_state else {} | ||
| missing = f"no training logs to check in {where}" | ||
| except (OSError, ValueError) as e: | ||
| metrics, missing = {}, f"cannot read {where}: {e}" | ||
| if not metrics: |
There was a problem hiding this comment.
[SUGGESTION] Two unexpected-exception paths produce no verdict file, which the non-strict branch then swallows.
get_final_metrics()can raise outside(OSError, ValueError).state.get("log_history", ...)raisesAttributeErrorif the JSON parses to a non-object ([],null), andfloat(last["loss"])raisesTypeErroron"loss": null. These are less likely than the truncation case you test, but they land in the same bucket as "a job killed mid-save".thresholds = {v: float(os.environ[v]) ...}on line 156 raisesValueErrorbefore the verdict dict exists, so a typo'dMAX_FINAL_LOSS: "7,0"in a recipe is a bare traceback.
In both cases the process exits non-zero with no regression_check.json written, and in the default (non-strict) branch || true discards the traceback's significance — so a downstream gate cannot tell "the checker was misconfigured / crashed" from "the training script never got far enough to run it". Given that distinction is the whole point of --json, consider widening to except Exception for the metrics read, and validating/reporting a malformed threshold as its own failure_class (e.g. THRESHOLDS_INVALID) that still writes the file.
There was a problem hiding this comment.
Fixed the metrics read in 9a5904c. TypeError/AttributeError from a well-formed but wrong-shaped trainer_state.json now give METRICS_MISSING, with tests for [] and "loss": null.
| if [ -n "$OUTPUT_DIR" ]; then | ||
| python3 common/check_regression.py "${OUTPUT_DIR}" || true | ||
| REGRESSION_ARGS=("${OUTPUT_DIR}" --json "${OUTPUT_DIR}/regression_check.json") | ||
| if [ "${REGRESSION_STRICT:-0}" = "1" ]; then |
There was a problem hiding this comment.
[SUGGESTION] REGRESSION_STRICT silently falls back to warn-only for any value other than the literal 1 — "true", "yes", "1 " all disable the gate the operator thought they enabled. Per the repo's "fail loudly rather than silently ignore an option" guidance, either accept the usual truthy spellings or reject unknown values:
case "${REGRESSION_STRICT:-0}" in
1|true|TRUE|yes) REGRESSION_STRICT_ON=1 ;;
0|false|FALSE|no|"") REGRESSION_STRICT_ON=0 ;;
*) echo "[ERROR] REGRESSION_STRICT must be 0/1, got '${REGRESSION_STRICT}'" >&2; false ;;
esacThere was a problem hiding this comment.
Fixed in 9a5904c. Only 0/unset and 1 are accepted. Any other value, true included, fails the task with REGRESSION_STRICT must be 0 or 1. I kept one documented spelling instead of accepting several truthy ones.
There was a problem hiding this comment.
Claude review — [1/8] launcher: make the regression gate machine-readable and opt-in strict
Scope: the trigger comment was a bare /claude review, so this is a full review. All 4 changed files (153 additions) were read: tools/launcher/common/check_regression.py, tools/launcher/common/specdec/dflash_online_training.sh, tools/launcher/examples/Qwen/Qwen3-8B/hf_online_dflash2.yaml, tools/launcher/tests/test_check_regression.py. Nothing under modelopt/ is touched, so there are no mode-registration, modelopt_state, or export-compat surfaces in play. I also read common/service_utils.sh for the trap semantics the shell change relies on, and grepped for other callers of check_regression() — there are none outside the test, so the signature change is safe.
Findings: CRITICAL: 0 · IMPORTANT: 1 · SUGGESTION: 3
The three bugs the PR claims to fix do get fixed
- The string-sort bug is real and the
key=_checkpoint_stepfix is correct;-1for non-numeric suffixes keeps a straycheckpoint-finalfrom outranking numbered ones. - Routing thresholds through an explicit
thresholdsdict built withif os.environ.get(v)cleanly kills thefloat('')crash and makescheck_regression()testable withoutmonkeypatch— a genuine improvement over the env-var reads it replaces. - The
METRICS_MISSINGpath closes the real hole: with a threshold set, a missing orJSONDecodeError-ingtrainer_state.jsonnow exits 1 instead of warning and returning 0. main()returning an int viasys.exit(main()), and the log lines staying byte-identical on the pass path, are the right calls for not breaking existing log readers.- Tests are proportionate and each pins a distinct behaviour — the parametrized CLI test in particular covers the exit code and the JSON contract together, which is what a downstream gate actually depends on. Good restraint: no test restates the implementation.
Most impactful finding
[IMPORTANT] A crashed training run can still write pass: true. error_handler sets FAIL_EXIT=1 but does not exit, and dflash_online_training.sh has no set -e. So when accelerate launch dies, control falls through to the export loop and then to this gate, which grades the surviving checkpoint and can emit {"pass": true, "failure_class": null} for a job that failed. The Slurm job state is still correct (the EXIT trap exits 1), but the verdict file — the contract this PR adds for the specdec-release gate in [8/8] — is not. A crash in the last stretch of the 92K-step recipe leaves a checkpoint comfortably under MAX_FINAL_LOSS. Suggested fix inline: short-circuit on FAIL_EXIT and write a TRAIN_FAILED verdict, ideally via a --failure-class flag so the schema stays in one place.
Suggestions (non-blocking)
find_trainer_state()still prefers any numeric checkpoint over<output_dir>/trainer_state.json, so "latest" can still be the wrong state whenmax_stepsis not a multiple ofsave_steps. Comparingglobal_stepacross candidates would complete the fix the PR sets out to make.get_final_metrics()can raise outside(OSError, ValueError)—AttributeErroron a non-object JSON,TypeErroron"loss": null— and thefloat(os.environ[v])comprehension raises before the verdict dict exists on a malformed threshold. Both exit non-zero with noregression_check.json, and|| truediscards the signal, defeating the point of--json.REGRESSION_STRICTmatches only the literal1;"true"silently disables the gate.
Threshold recalibration
MAX_FINAL_LOSS 5.0 → 7.0 is justified by the five GPU runs in the description, and the 4.2% seed-to-seed margin is disclosed rather than hidden — not flagging it. One cross-PR note for [8/8]: hf_online_domino.yaml and hf_online_lilicorr.yaml keep the inherited MAX_FINAL_LOSS: "5.0" and will now write pass: false verdict files (harmless at the job level thanks to || true, but a file-reading gate would reject those runs). Worth confirming the [8/8] gate keys off thresholds/strictness and not pass alone.
Risk
Low. The blast radius is one launcher tool plus one example recipe; strict mode is opt-in and the default behaviour and log format are unchanged, so existing recipes and log readers are unaffected. The IMPORTANT finding is a gap in the new verdict's trustworthiness rather than a regression against main, but it is worth closing before [8/8] starts consuming the file.
🤖 Generated with Claude Code
…REGRESSION_STRICT
A diverged run logs NaN or Infinity. json.dump wrote those as bare NaN/Infinity,
which strict JSON readers reject, so a failed run's verdict could be unreadable.
A non-finite checked metric now fails as TRAIN_REGRESSION, every non-finite metric
is written as null, and the verdict is dumped with allow_nan=False.
trainer_state.json that parses but has the wrong shape ("loss": null, a top-level
list) raised TypeError/AttributeError past the read, so no verdict was written;
it now fails as METRICS_MISSING like an unreadable file.
REGRESSION_STRICT accepted only "1" and treated anything else, "true" included, as
warn-only. Any value but 0 or 1 now fails the task with a message.
Signed-off-by: Zhehao Hu <zhehaoh@nvidia.com>
cjluo-nv
left a comment
cjluo-nv
left a comment
There was a problem hiding this comment.
Bot review (gpt-6.1-sol) — DM the bot to share feedback.
Changes requested: the non-finite metric fix is sound, but failed training can still produce a passing verdict in this PR.
Needs action:
- 💬 Author deferred crash handling to [3/8] #2931 — move that safeguard into
dflash_online_training.shhere and test failed training with a surviving passing checkpoint; each merged PR must be safe independently. - Fix threshold parsing in
check_regression.pyto report malformed and non-finite values in a valid failed JSON verdict; add CLI tests for7,0, NaN, and infinity.
No action needed:
- ✔️ Resolved since the last review: non-finite metric serialization, wrong-shaped metric reads, and invalid
REGRESSION_STRICThandling. - Existing test edits preserve assertions and are justified by explicit threshold arguments. The checkpoint-ordering reply provides concrete final-save evidence.
- Tests were inspected, not executed.
| case "${REGRESSION_STRICT:-0}" in | ||
| 1) python3 common/check_regression.py "${REGRESSION_ARGS[@]}" ;; | ||
| 0) python3 common/check_regression.py "${REGRESSION_ARGS[@]}" || true ;; | ||
| # Fail rather than silently run warn-only on a typo such as "true". |
There was a problem hiding this comment.
Bot comment.
The reply explains the intended fix in #2931, but this independently merged PR still runs the checker after accelerate launch fails: error_handler only sets FAIL_EXIT, and execution continues. A surviving checkpoint with acceptable metrics therefore writes pass: true for failed training. Please bring the failure-status safeguard (including stale-verdict cleanup) forward into this PR and add a shell regression covering a failed launch with an otherwise passing checkpoint. A later series member does not resolve the current correctness gap.
There was a problem hiding this comment.
Moved into this PR in 4101423. A failed accelerate launch now skips export and the regression check, and an old regression_check.json is removed before training starts. The failure is still reported through error_handler, and the task exits 1. tests/test_dflash_online_training.py runs the script in a sandbox with a failing launch, a checkpoint that passes MAX_FINAL_LOSS, and an old passing verdict. It asserts no verdict and exit 1. A second case, with the launch succeeding, checks that the same sandbox does grade the checkpoint. Against the previous script, the failing-launch case finds pass: true. [3/8] #2931 no longer carries this block.
| return | ||
| def evaluate(output_dir): | ||
| """Run the check on output_dir. Returns the verdict dict that --json writes.""" | ||
| thresholds = {v: float(os.environ[v]) for v in THRESHOLD_VARS if os.environ.get(v)} |
There was a problem hiding this comment.
Bot comment.
Threshold values remain outside the validation/error-reporting path. MAX_FINAL_LOSS=7,0 raises before a verdict exists; NaN or Infinity parses successfully but stays in verdict['thresholds'], so json.dump(..., allow_nan=False) raises after opening/truncating the output file, leaving incomplete JSON. Please validate conversion and finiteness before evaluation, produce a JSON-safe failed configuration verdict, and cover these cases in the CLI tests. The new metric sanitization fixes metrics, not this remaining threshold path.
There was a problem hiding this comment.
Fixed in 4101423. parse_thresholds() reports a value that isn't a finite number (7,0, nan, inf) as a failed verdict with failure_class: METRICS_MISSING, and leaves it out of thresholds. The verdict is serialized before the file is opened, so a failed dump can't leave a truncated file. test_invalid_threshold_writes_failed_verdict covers all three values and parses the file as strict JSON.
|
/claude review |
| return | ||
| def evaluate(output_dir): | ||
| """Run the check on output_dir. Returns the verdict dict that --json writes.""" | ||
| thresholds = {v: float(os.environ[v]) for v in THRESHOLD_VARS if os.environ.get(v)} |
There was a problem hiding this comment.
[IMPORTANT Correctness] A malformed or non-finite threshold leaves no readable verdict file, which is the one thing --json exists to guarantee.
Two distinct paths, both reachable from a YAML typo:
- Malformed —
MAX_FINAL_LOSS: "7,0"(comma decimal) or"7.0 "with a stray char.float(os.environ[v])raisesValueErrorhere, beforeverdictexists.evaluate()propagates out ofmain()→ traceback, noregression_check.jsonwritten at all. - Non-finite —
MAX_FINAL_LOSS: "inf"/"nan"parses fine, sothresholds = {"MAX_FINAL_LOSS": inf}. Everything proceeds, thenjson.dump(verdict, f, indent=2, allow_nan=False)on line 224 raisesValueError: Out of range float values are not JSON compliant— butopen(args.json_path, "w")has already truncated the file andjson.dumpstreams chunks viaiterencode, so"thresholds"is the 5th of six keys: the file on disk is left truncated mid-object. A downstream reader gets aJSONDecodeErroron a file that exists.
Why it matters: under the default REGRESSION_STRICT=0 path the shell's || true swallows the non-zero exit, so Slurm reports COMPLETED and the specdec-release gate in [8/8] finds either a missing or a corrupt regression_check.json. Note case 2 is strictly worse than today's behaviour — main never wrote a file, so there was nothing to misparse; a truncated file is an existing-but-invalid artifact. The allow_nan=False hardening correctly covers metrics (line 196) but thresholds is unsanitized.
Suggested fix — validate thresholds into the verdict rather than letting them raise, so a bad threshold is reported as a failure class like any other:
def _parse_thresholds():
"""Returns ({env_var: value}, [error messages]). A bad value is reported, not raised."""
thresholds, errors = {}, []
for v in THRESHOLD_VARS:
raw = os.environ.get(v)
if not raw:
continue
try:
val = float(raw)
except ValueError:
errors.append(f"REGRESSION: {v}={raw!r} is not a number")
continue
if not math.isfinite(val):
errors.append(f"REGRESSION: {v}={raw!r} is not a finite number")
continue
thresholds[v] = val
return thresholds, errorsThen in evaluate(), if errors is non-empty, return {"pass": False, "failure_class": "THRESHOLD_INVALID", "messages": errors, ...} and skip the metric read. That keeps every exit path emitting one strict-JSON verdict, which is the contract [8/8] depends on. Worth adding "7,0", "inf" and "nan" threshold cases to the parametrized CLI test so the verdict file is asserted to parse in each.
There was a problem hiding this comment.
| An empty variable counts as unset. Once any threshold is set, a missing or | ||
| unreadable trainer_state.json, or one without training logs, also fails the check, | ||
| and so does a NaN or infinite metric. |
There was a problem hiding this comment.
[SUGGESTION] The docstring overstates the non-finite rule. Only a metric that has a threshold set fails on NaN/infinity — the math.isfinite check on line 148 sits after the env_var not in thresholds guard. An unthresholded non-finite metric passes; this PR's own test asserts exactly that ('{"log_history": [{"loss": 4.5, "acc": Infinity}]}' → exit 0, failure_class=None). As written a reader would expect that case to fail.
| An empty variable counts as unset. Once any threshold is set, a missing or | |
| unreadable trainer_state.json, or one without training logs, also fails the check, | |
| and so does a NaN or infinite metric. | |
| An empty variable counts as unset. Once any threshold is set, a missing or | |
| unreadable trainer_state.json, or one without training logs, also fails the check, | |
| and so does a NaN or infinite value for a metric that has a threshold. |
There was a problem hiding this comment.
Applied in 4101423, plus a sentence on invalid thresholds.
| if verdict["metrics"]: | ||
| print(f"Final metrics: {verdict['metrics']}") |
There was a problem hiding this comment.
[SUGGESTION] The human-readable log now prints the JSON-sanitized metrics, so a diverged run emits two contradictory lines:
=== Regression Check (/scratchspace/dflash2_bs16/checkpoint-2000/trainer_state.json) ===
Final metrics: {'loss': None}
REGRESSION: loss is nan, not a finite number
Regression check FAILED
{'loss': None} reads as "the trainer never logged a loss" — the METRICS_MISSING story — when the actual fact is that it logged nan. That is the single most diagnostic value in the log for a divergence, and the None substitution is an artifact of a constraint (strict JSON) that only applies to the file.
The sanitization is only needed for json.dump, so keep the raw dict for the log and sanitize at write time:
# in evaluate(): keep the raw values
verdict.update({"pass": passed, "metrics": metrics, "messages": messages})
# in main(), at the write:
if args.json_path:
# Strict JSON: NaN/Infinity would break non-Python readers of the verdict.
serializable = dict(verdict)
serializable["metrics"] = {
k: v if math.isfinite(v) else None for k, v in verdict["metrics"].items()
}
with open(args.json_path, "w") as f:
json.dump(serializable, f, indent=2, allow_nan=False)This also keeps Final metrics: byte-identical to main on every path, matching the PR's "existing log lines are unchanged" claim. The existing CLI test still passes, since it asserts on failure_class/pass/thresholds rather than on metrics.
There was a problem hiding this comment.
Applied in 4101423. evaluate() keeps the raw metrics, so the log prints {'loss': nan}, and main() writes non-finite values as null only in the file.
There was a problem hiding this comment.
Claude review — [1/8] launcher: make the regression gate machine-readable and opt-in strict
Scope: bare /claude review, so a full review. Re-reviewed at 9a5904c5 against my previous pass at 60fce3f. All 4 changed files read (169 additions): tools/launcher/common/check_regression.py, common/specdec/dflash_online_training.sh, examples/Qwen/Qwen3-8B/hf_online_dflash2.yaml, tests/test_check_regression.py. Nothing under modelopt/ is touched, so no mode-registration, modelopt_state or export-compat surfaces are in play. I also re-read common/service_utils.sh for the trap semantics, and grepped every consumer of the gate across tools/ and plugins/.
Findings: CRITICAL: 0 · IMPORTANT: 2 · SUGGESTION: 2
Resolved since 60fce3f ✅
- Non-finite metric handling —
math.isfinitecheck plusallow_nan=False, with NaN andInfinityCLI cases. Correct. - Wrong-shaped
trainer_state.json—TypeError/AttributeErroradded to theexcept, with"[]"and{"loss": null}test cases. REGRESSION_STRICTvalidation — thecasestatement now rejects"true"loudly viafalse+ theERRtrap instead of silently running warn-only. I confirmed the wiring: all four recipes that setMAX_FINAL_LOSS(hf_online_dflash2,hf_online_domino,hf_online_lilicorr,hf_offline_dflash) runcommon/specdec/dflash_online_training.sh, soREGRESSION_STRICT: "1"does take effect where it is set.
Most impactful findings
[IMPORTANT] A malformed or non-finite threshold leaves no readable verdict file (inline). float(os.environ[v]) on line 163 raises before verdict exists, so MAX_FINAL_LOSS: "7,0" produces a traceback and no regression_check.json. Worse, "inf"/"nan" parses fine and reaches json.dump(..., allow_nan=False) — which raises after open(..., "w") truncated the file, and since thresholds is the 5th of six keys, the file is left truncated mid-object. That case is strictly worse than main, which never wrote a file at all: a corrupt artifact is harder for a gate to handle than a missing one. Under the default || true both exit silently. The allow_nan=False hardening covers metrics but not thresholds. The inline comment has a _parse_thresholds() sketch that routes bad values into the verdict as a failure class.
[IMPORTANT] The speculative-decoding skill now documents this gate wrongly. This PR changes the behaviour the skill tree describes but not the descriptions, and per CLAUDE.md these files are loaded into agent context on every use — the [8/8] specdec-release gate reasons from them:
plugins/modelopt/skills/speculative-decoding/references/algorithms/dflash2.md:127-129— "the online example setsMAX_FINAL_LOSS=5.0andMIN_FINAL_ACC=0.15for its 2000-step convergence check." This PR sets it to7.0. An agent judging a run at loss 6.5 against the documented 5.0 would call a healthy run a regression — the exact misdiagnosis the recalibration exists to prevent.plugins/modelopt/skills/speculative-decoding/references/algorithms/dflash.md:140-142— "check_regression.pyis invoked with|| true, and it only warns when notrainer_state.jsonexists. A green Slurm exit is not proof the gate ran — confirm the=== Regression Check ===block is present in the log." All three clauses are now wrong for a strict recipe:|| trueis dropped, a missingtrainer_state.jsonwith a threshold set now fails, and a green exit is proof underREGRESSION_STRICT=1. It also still sends the agent to parse logs for a verdict that is now a file.
These need roughly four lines of edits across two files, plus a mention of regression_check.json and REGRESSION_STRICT. Keeping it in this PR means [1/8] stands on its own; deferring to [8/8] leaves the skill actively misleading in between. Not posted inline because neither file is in the diff.
Suggestions (non-blocking)
- The module docstring says a NaN or infinite metric fails, but only a thresholded metric does — this PR's own test asserts
"acc": Infinityexits0. inline Final metrics: {loss: None}is printed directly aboveREGRESSION: loss is nan— the log hides the single most diagnostic value behind a constraint that only applies to the JSON file. Sanitize at write time instead; that also keepsFinal metrics:byte-identical tomain. inline
Still open from my last pass
A crashed training run can still write pass: true. error_handler sets FAIL_EXIT=1 without exiting and the script has no set -e, so after accelerate launch dies control falls through to the export loop and then to this gate, which grades the surviving checkpoint. The Slurm state is right; the verdict file — the artifact this PR adds — is not. You deferred this to [3/8] #2931, and I accept that as a scoping call rather than re-litigating it. I will note only that per CLAUDE.md "every sub-PR stands on its own", and this is the PR that introduces the file [8/8] will trust. If a TRAIN_FAILED short-circuit is cheap here, it is worth the three lines; if it genuinely belongs with the rest of the crash handling in [3/8], say so in the description so a reader of the merged [1/8] knows the verdict is not yet crash-safe.
Checkpoint ordering vs. a final direct save — your final-save evidence addresses this. Withdrawn.
Threshold recalibration
MAX_FINAL_LOSS 5.0 to 7.0 is backed by five GPU runs (6.37–6.71), and the 4.2% seed-to-seed margin is disclosed rather than buried. Not flagging. One cross-PR note I confirmed by grep: hf_online_domino.yaml, hf_online_lilicorr.yaml and hf_offline_dflash.yaml keep the inherited MAX_FINAL_LOSS: "5.0" on the same script, so they will now start writing pass: false verdict files — harmless at the job level thanks to || true, but confirm the [8/8] gate keys off strictness and not pass alone, or those three runs get rejected on an uncalibrated threshold.
Assessment
Risk: low. Blast radius is one launcher tool and one recipe; strict mode is opt-in and the pass-path log format is unchanged, so existing recipes and log readers are unaffected. The three bugs the PR sets out to fix are genuinely fixed, and the test suite is proportionate — each case pins a distinct behaviour with no restatement of the implementation. Both IMPORTANT findings are about the new verdict file being trustworthy before [8/8] starts consuming it, rather than regressions against main — except the truncated-file case, which is a real if narrow step backwards.
🤖 Generated with Claude Code
dflash_online_training.sh ran the regression check after a failed accelerate
launch, because the ERR trap only records a failure. A checkpoint that survived
the crash could then be graded and written as pass: true, and a verdict left by
an earlier run in the same output_dir was read as this run's. A training
failure now skips export and the check, and an old regression_check.json is
removed before training starts. test_dflash_online_training.py runs the script
with a failing launch and a passing checkpoint.
check_regression.py also never leaves a missing or truncated verdict for a bad
threshold: a value that isn't a finite number ("7,0", "nan", "inf") fails as
METRICS_MISSING, and the verdict is serialized before the file is opened. The
log keeps the raw metric values; only the file writes non-finite ones as null.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Signed-off-by: Zhehao Hu <zhehaoh@nvidia.com>
cjluo-nv
left a comment
cjluo-nv
left a comment
There was a problem hiding this comment.
Bot review (gpt-6.1-sol) — DM the bot to share feedback.
Changes requested: the crash and threshold fixes resolve the main correctness concerns, but the previously flagged skill guidance still contradicts the gate.
Needs action:
- Update
plugins/modelopt/skills/speculative-decoding/references/algorithms/dflash2.mdto use the calibrated 7.0 loss threshold and document strict mode. - Replace the stale gate note in
plugins/modelopt/skills/speculative-decoding/references/algorithms/dflash.mdwith conditional strict-mode behavior, missing-metrics failure, and theregression_check.jsonartifact.
No action needed:
- ✔️ Resolved since the last review: failed-launch handling and stale-verdict cleanup, invalid thresholds, non-finite serialization, wrong-shaped metric reads, strict-option validation, and diagnostic log wording.
- Existing test edits are justified by explicit threshold arguments and preserve their assertions; new tests cover the crash and malformed-threshold paths.
- The new test header matches
LICENSE_HEADER. Tests were inspected, not executed.
Additional comments (outside the PR diff):
plugins/modelopt/skills/speculative-decoding/references/algorithms/dflash2.md:128— > Bot comment.
The previously flagged documentation mismatch remains: this says the online example uses MAX_FINAL_LOSS=5.0, but this PR calibrates it to 7.0. Replace these lines with concise guidance such as: The online convergence check sets MAX_FINAL_LOSS=7.0, MIN_FINAL_ACC=0.15, and REGRESSION_STRICT=1; read <training.output_dir>/regression_check.json for its verdict. Otherwise an agent following this sheet can reject a healthy run.
plugins/modelopt/skills/speculative-decoding/references/algorithms/dflash.md:140— > Bot comment.
This previously flagged note still says missing trainer_state.json only warns and the checker always runs with || true. Both become inaccurate with this PR. Replace it with: The script writes <output_dir>/regression_check.json. With thresholds set, missing or unreadable metrics fail the check. REGRESSION_STRICT=1 propagates checker failure to the task; unset/0 remains warn-only, so a green Slurm exit alone is not proof of a passing verdict.
… gate REGRESSION_STRICT was only checked after training and export, so a typo such as "true" cost a full training run before failing. The script now checks it right after removing any old regression_check.json, before installing or training. The old verdict goes first so the early exit cannot leave an earlier run's pass behind. test_dflash_online_training.py exited 1 in both cases: the sandbox had no torch, so the GPU count tripped the ERR trap. It now puts a stub torch on PYTHONPATH and asserts exit 0 for a successful run and 1 for a crash. New cases cover strict mode end to end (REGRESSION_STRICT=1 fails the task on a regression, 0 does not) and the up-front check (no training, no verdict). The DFlash skill sheets described the old gate: dflash2.md cited MAX_FINAL_LOSS=5.0 and dflash.md said the check only warns. They now give the calibrated 7.0, strict mode, the METRICS_MISSING and TRAIN_REGRESSION failures, and the regression_check.json verdict. Signed-off-by: Zhehao Hu <zhehaoh@nvidia.com>
|
Addressed in ad66c1c, replying to the latest review (its two comments are outside the diff, so they have no thread).
The same commit fixes two more issues found in review:
Launcher suite: 187 passed at the stack tip, 186 at [3/8]. |
There was a problem hiding this comment.
Bot review (gpt-6.1-sol) — DM the bot to share feedback.
Prior correctness concerns are resolved; the required design review still needs a comparison with the existing Day-0 verdict pattern.
Needs action:
- Explain in the PR body why the launcher-specific verdict contract is preferable to adapting
plugins/modelopt/skills/day0-release/scripts/gate_run.py; clarify whether their schemas intentionally remain independent.
No action needed:
- ✔️ Previous concerns are addressed, including failed-launch handling and stale-verdict cleanup, invalid thresholds, non-finite JSON, wrong-shaped metrics, strict-option validation, and stale skill guidance.
- Existing test edits are justified by the explicit-threshold signature change and preserve their assertions. Sandbox tests now cover successful training, crashes, strict regression failure, and early option validation.
- The new test header matches
LICENSE_HEADER. Tests were inspected, not executed.
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
yeyu-nvidia
left a comment
yeyu-nvidia
left a comment
There was a problem hiding this comment.
Human sign-off on the test_check_regression.py edits, as the bot asked — I wrote that file in #2776, so I can speak to whether the edits weakened it. They did not. Approving.
The test edits
Every change to my tests is a mechanical adaptation to check_regression(metrics) becoming check_regression(metrics, thresholds), with assertions intact:
test_threshold_with_no_metric_failsstill assertsnot passedand"never evaluated" in m. That is the exact invariant #2776 existed to pin — a threshold whose metric never reachedtrainer_state.jsonmust fail rather than warn — and it survives.test_unset_threshold_is_not_a_failureandtest_accuracy_threshold_is_enforced_when_presentare the same assertions against an explicit{}/{"MIN_FINAL_ACC": 0.15}.test_metrics_pick_up_any_accuracy_keyis untouched.
Taking the thresholds as an argument instead of reading the environment inside the check is the better shape anyway: the parsing failure modes now live in parse_thresholds where they can be tested directly, which is what makes the "7,0" case reachable at all.
What I checked beyond that
- The string-sort bug is real and the fix is right.
checkpoint-8000vscheckpoint-12000is the classic case, andhf_online_dflash2.yaml's ~92K-step recipe atsave_steps=4000crosses the digit boundary, so this was silently grading an old checkpoint on exactly the run that matters most. || trueis the right thing to remove. A gate whose verdict exists only as log text, behind a swallowed exit code, is the same silent-pass shape as #2288'sar_validateand the one #2776 fixed — the third instance of it in this area.- The strict-JSON assertion on the verdict is well judged.
parse_constant=pytest.failmeans a NaN or Infinity reaching the verdict file breaks the test rather than the downstream reader. I raised exactly this hazard on #2640 last week, where-Infinityin a published artifact round-trips in Python and is rejected by Go and Rust parsers. Good to see it pinned here rather than discovered later. - The skill-sheet updates are accurate. They are my files from #2602 and the new text matches the behaviour this PR implements — the
regression_check.jsonartifact row, thefailure_classvalues, andREGRESSION_STRICTbeing DFlash2-online only.
The MAX_FINAL_LOSS: 5.0 → 7.0 recalibration is worth noting for a reason beyond this PR: DFlash's own online example still carries 5.0, and loss scale moves with the objective as well as the variant. I hit the same class of thing on a DFlash run whose dpace loss floors near 0.39, where a 5.0 threshold passes trivially. A threshold calibrated against one objective is not portable to another, and nothing records which objective a given number was calibrated against.
One design note, not blocking
failure_class conflates two different causes. A malformed threshold — MAX_FINAL_LOSS="7,0", an operator typo in a YAML — reports METRICS_MISSING, the same class as a truncated trainer_state.json from a job killed mid-save. test_invalid_threshold_writes_failed_verdict pins that behaviour deliberately, so it is a choice rather than an oversight.
But the whole point of a machine-readable failure_class is letting automation branch on cause, and those two want opposite responses: a truncated checkpoint is worth a retry, a typo in the YAML never is. Something like CONFIG_INVALID would keep the distinction, and the thresholds == {} field already carries enough to tell them apart if you would rather not add a class — but a consumer has to know to look.
Not worth holding [1/8] for; worth settling before anything downstream in the series starts branching on the field.
|
…ction (#2930) > **[2/8] Day-0 speculative-decoding readiness flow.** Merge in order 1 → 8. Base: `main` ([1/8] has merged). > Series: [1/8] #2839 (merged) · [2/8] #2930 (this) · [3/8] #2931 · [4/8] #2932 · [5/8] #2933 · [6/8] #2934 · [7/8] #2935 · [8/8] #2936 ### What does this PR do? Type of change: Bug fix The vLLM smoke test (`common/specdec/vllm_smoke_test.sh`) is the serving check after draft training. It had six problems: 1. **It could serve the wrong checkpoint.** `ar_eval_mtbench.sh` and `vllm_smoke_test.sh` picked the newest checkpoint by sorting on a dash-separated field of the **whole path**, so a dash in a parent directory (e.g. `/data/my-drafts/...`) picked the wrong one. `quantize_drafter.sh` had a third, different sort. The smoke test also ignored `exported-checkpoint-final`. 2. **Every throughput figure printed `?`.** Timing used `bc`, which the vLLM image doesn't have. 3. **A drafter that never speculated passed.** With `MIN_ACCEPTANCE_LENGTH` set, a vLLM log with no acceptance length only warned. 4. **The acceptance length covered only the last few seconds.** It came from the log's "Mean acceptance length" line, which covers only vLLM's last logging interval. 5. **The verdict existed only as log text.** 6. **The server wait was fixed at 180 s.** Qwen3.8-27B with a DFlash2 drafter takes longer to start on GB300 (torch.compile, kernel JIT, CUDA graph capture, GDN warmups), so every smoke attempt stopped at `server_timeout`. The script's exit then tore the container down mid-warmup, and the engine's next exec failed with a ptxas ENOENT that read like the cause. Changes: - `service_utils.sh`: new `latest_checkpoint <dir> <prefix>`. It returns `<dir>/<prefix>final` if present, else the highest numeric step, keyed on the directory's own name. It replaces the three ad hoc sorts in `ar_eval_mtbench.sh`, `vllm_smoke_test.sh` and `quantize_drafter.sh`. - `vllm_smoke_test.sh`: - Prefers `exported-checkpoint-final`, as `quantize_drafter.sh` already did. - Times prompts with shell arithmetic and `python3` instead of `bc`. - Computes the mean acceptance length over the whole run from vLLM's cumulative `/metrics` counters: 1 + `spec_decode_num_accepted_tokens_total` / `spec_decode_num_drafts_total`. - With `MIN_ACCEPTANCE_LENGTH` set, fails when no speculative drafts were recorded. - Fails loudly if `service_utils.sh` is missing (it used to be optional, so the failure surfaced as a misleading `draft_missing`). Only a missing file is fatal: `source` returns the status of the file's last command, a rank-0 append that fails outside a launcher working directory. - `SERVER_TIMEOUT` sets the server wait in seconds (default 180) and appears in the timeout reason. - Checks `MIN_ACCEPTANCE_LENGTH` (a finite number), and `NUM_SPEC_TOKENS` and `SERVER_TIMEOUT` (positive integers without leading zeros, since `NUM_SPEC_TOKENS` goes into vLLM's JSON config), before any work, and fails with `stage: invalid_config` and a reason naming the value. Both checks run after the `EXIT` trap is set, so even these failures replace any earlier `SMOKE_JSON`. - Optional `SMOKE_JSON=<path>` writes the verdict on every exit: `{stage, pass, reason, method, target, draft, num_speculative_tokens, prompts passed/failed, tokens, average tok/s, mean_acceptance_length, min_acceptance_length}`. `stage` says how far it got (`invalid_config`, `starting`, `draft_missing`, `server_died`, `server_timeout`, `serving`, `completed`), so a downstream gate can tell "not servable" from "below the floor". The file's parent directory is created if needed, and it is strict JSON (a non-numeric value is written as `null`). The mean acceptance length is checked against the floor and written at full precision; only the log rounds it. ### Usage ```yaml task_2: script: common/specdec/vllm_smoke_test.sh environment: - HF_MODEL_CKPT: <<global_vars.hf_model>> - DRAFT_CKPT_DIR: <<global_vars.output_dir>> # newest export, final preferred - SPEC_METHOD: "dflash" - NUM_SPEC_TOKENS: "7" - MIN_ACCEPTANCE_LENGTH: "1.5" - SMOKE_JSON: <<global_vars.output_dir>>/smoke_check.json - SERVER_TIMEOUT: "900" # optional; default 180 s ``` ### Testing - `tools/launcher/tests/test_service_utils.py`: `latest_checkpoint` with a dash in a parent directory, `final` winning over steps, numeric rather than string order, and no checkpoint at all (3 cases). - `tools/launcher/tests/test_vllm_smoke_test.py` runs the script as rank 0 from a directory with no `modules/` (stubbed `pip` and `nvidia-smi`), up to the paths that stop before vLLM starts (6 cases). Against the script without these fixes, all 6 fail: - It reaches `draft_missing` rather than stopping at `source`. - `MIN_ACCEPTANCE_LENGTH` of `1,5` or `nan`, `NUM_SPEC_TOKENS` of `x` or `08` and `SERVER_TIMEOUT=15m` each fail as `invalid_config` with a strict-JSON `SMOKE_JSON`. Before, `NUM_SPEC_TOKENS=x` crashed in the `EXIT` trap and wrote no JSON. - `bash -n` on all four scripts. Launcher unit suite on this commit: 192 passed. - **GPU (OCI-JHB, GB300, vLLM v0.31.0)**, through `hf_online_dflash2.yaml`'s smoke task ([4/8]) and one smoke-only submission, six runs in all: - `latest_checkpoint` picked `exported-checkpoint-final` in every training run. - Per-prompt tok/s printed (e.g. 335 tok/s average), 8/8 prompts passed. - Mean acceptance length 1.68–1.79 against the 1.5 floor; `smoke_check.json` written each time. - **GPU, Qwen3.8-27B + DFlash2 (same cluster and image):** four smoke runs at the old fixed 180 s each wrote `smoke_check.json` with `stage: server_timeout`. A fifth with `SERVER_TIMEOUT=900` passed at AL 1.76. - Not run on GPU: - the `server_died` and no-drafts failure paths. They exit through the same `write_smoke_json` trap as the passing path; - creating `SMOKE_JSON`'s parent directory. A smoke-only run hit the missing directory before this was added; the rerun used an existing directory. ### Before your PR is "*Ready for review*" Make sure you read and follow [Contributor guidelines](https://github.com/NVIDIA/Model-Optimizer/blob/main/CONTRIBUTING.md) and your commits are signed (`git commit -s -S`). Make sure you read and follow the [Security Best Practices](https://github.com/NVIDIA/Model-Optimizer/blob/main/SECURITY.md#security-coding-practices-for-contributors) (e.g. avoiding hardcoded `trust_remote_code=True`, `torch.load(..., weights_only=False)`, `pickle`, etc.). - Is this change backward compatible?: ✅ (`SMOKE_JSON` is optional. With `MIN_ACCEPTANCE_LENGTH` set, a run with no speculative drafts now fails instead of passing; that was the bug.) - If you copied code from any other sources or added a new PIP dependency, did you follow guidance in `CONTRIBUTING.md`: N/A - Did you write any new necessary tests?: ✅ (`latest_checkpoint`; the smoke script itself is validated on GPU as above) - Did you update [Changelog](https://github.com/NVIDIA/Model-Optimizer/blob/main/CHANGELOG.rst)?: N/A - Did you get Claude approval on this PR?: ❌ - not yet run ### Additional Information - The `specdec-release` skill's `gate_smoke.py` ([6/8]) reads `SMOKE_JSON`. <!-- This is an auto-generated comment: release notes by coderabbit.ai --> <!-- coderabbit-ai-generated:v1 --> ## Summary by CodeRabbit * **Improvements** * Checkpoint selection now prefers a `final` checkpoint, or otherwise chooses the highest-numbered checkpoint. * Smoke tests validate numeric settings, use a configurable server timeout, and report run results, failure reasons, and performance metrics in an optional JSON file. * Acceptance-length checks now use cumulative draft and accepted-token counts; prompt failures and unmet acceptance thresholds result in a failed test. * **Tests** * Added coverage for checkpoint selection and smoke-test configuration failures. <!-- end of auto-generated comment: release notes by coderabbit.ai --> --------- Signed-off-by: Zhehao Hu <zhehaoh@nvidia.com> Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>
What does this PR do?
Type of change: Bug fix
common/check_regression.py gates DFlash-family training runs on final loss and accuracy. Four things stopped that gate from catching a bad run:
find_trainer_state()sortedcheckpoint-*as strings, socheckpoint-8000ranked aftercheckpoint-12000and an older checkpoint's metrics were checked. Any run whose saved steps differ in digit count was affected; the ~92K-step full DFlash2 recipe inhf_online_dflash2.yamlsaves every 4000 steps. It now sorts by step number.trainer_state.json, or one with no training logs, only warned and exited 0. An empty threshold variable (MAX_FINAL_LOSS="") crashed onfloat('').dflash_online_training.shruns the check ascheck_regression.py … || true, so Slurm reportsCOMPLETEDeven when the log saysREGRESSION. The verdict also existed only as log text.ERRtrap only records a failure, so after a failedaccelerate launchit went on to export and run the check, and a checkpoint that survived the crash could pass.Changes:
check_regression.py:trainer_state.jsonof the wrong shape. An empty threshold variable counts as unset. A threshold that isn't a finite number (7,0,nan) fails asMETRICS_MISSING.TRAIN_REGRESSION.--json <path>and writes{pass, failure_class, trainer_state, metrics, thresholds, messages}as strict JSON (non-finite metrics asnull).failure_classisnull,TRAIN_REGRESSIONorMETRICS_MISSING. The verdict is serialized before the file is opened, so a failure can't leave a truncated file.=== Regression Check (…) ===,PASS:/REGRESSION:,Regression check PASSED) are unchanged, so existing log readers keep working. A failure now also printsRegression check FAILED.dflash_online_training.shwrites the verdict to${OUTPUT_DIR}/regression_check.json, next to the exports.REGRESSION_STRICT=1drops the|| true, so a regression fails the task through the existingERRtrap andexit_handler. The default (0or unset) is unchanged (warn only), so the other examples with possibly uncalibrated thresholds don't start failing. Any other value fails the task at startup, before installing or training.dflash_online_training.shalso skips export and the check when training fails, still reporting the failure like anyERR-trap failure, and removes an oldregression_check.jsonfirst thing, so no early exit can leave an earlier run's verdict behind.examples/Qwen/Qwen3-8B/hf_online_dflash2.yaml:REGRESSION_STRICT: "1".MAX_FINAL_LOSSfrom 5.0 to 7.0. The 5.0 was inherited fromhf_online_dflash.yaml, whose own reference curve only drops below 5.0 at step ~2500 of a 12.5K-step run. DFlash2 adds a selector loss term and this example stops at 2000 steps, so 5.0 fails every healthy run.MIN_FINAL_ACCstays at 0.15.speculative-decodingskill'sdflash.mdanddflash2.mddescribe the new gate: judge a run byregression_check.json, the two failure classes, strict mode, and DFlash2's 7.0 threshold.Usage
Testing
tools/launcher/tests/test_check_regression.py: 4 new test functions (13 cases). Run againstmain'scheck_regression.py, 17 of the file's 18 cases fail:test_latest_checkpoint_is_chosen_by_step_numberguards the string-sort bug.test_empty_threshold_var_is_unsetguards thefloat('')crash.test_cli_exit_code_and_json_verdictis parametrized over a missingtrainer_state.json, a truncated one, two of the wrong shape, a NaN loss, an unchecked infinite metric, a regression and a pass. Each case checks the exit code and parses the verdict as strict JSON, so it covers the missing-metrics hole and the--jsoncontract.test_invalid_threshold_writes_failed_verdictcovers7,0,nanandinf.tools/launcher/tests/test_dflash_online_training.py(5 cases) runs the real script in a sandbox: stubbedpip,nvidia-smi,accelerate, exporter andtorch, a checkpoint, and an old passing verdict. Each case asserts the exact exit code. Againstmain's script, all 5 fail:ERR-trap error.REGRESSION_STRICT=1exits 1 and0exits 0, both writingTRAIN_REGRESSION.REGRESSION_STRICT=trueexits 1 without launching training and leaves no verdict.uv run --with pytest python3 -m pytest tests/, 183 passed.dflash_online_training.sh:bash -npasses.Before your PR is "Ready for review"
Make sure you read and follow Contributor guidelines and your commits are signed (
git commit -s -S).Make sure you read and follow the Security Best Practices (e.g. avoiding hardcoded
trust_remote_code=True,torch.load(..., weights_only=False),pickle, etc.).CONTRIBUTING.md: N/AAdditional Information
specdec-releaseskill ([8/8]) gates training onregression_check.json, so it needs a verdict it can read and a job state that reflects it.day0-release/scripts/gate_run.py? The two sit at different layers and check different things, so their schemas stay independent on purpose:gate_run.pyis an agent-side gate for an evaluation run. It reads a run summary the agent builds from NEL output, and checks per-task status, sample counts and judge errors. None of that exists for a training run.check_regression.pyruns inside the Slurm training task, on rank 0, and its verdict decides the task's exit status in strict mode. It readstrainer_state.jsonagainst thresholds from the launcher YAML, and has to work for launcher users who never load a skill. A launcher script importing a skill script would also invert the dependency: skills drive the launcher, not the other way round.gate_run.pyfor training isspecdec-release/scripts/gate_train.py([6/8] [6/8] specdec-release: add training and serving readiness gates for drafters #2934). It readsregression_check.jsonand returns the day-0 gate shape ({pass, failure_class, detail, …}), withfailure_classfrom the shared triage vocabulary. That is the adapter layer, socheck_regression.pydoesn't need to matchgate_run.py.pass, and afailure_classthat isnullon a pass. The launcher emits onlyMETRICS_MISSINGandTRAIN_REGRESSION, andgate_train.pytreats any other value asTRAIN_REGRESSION, so the launcher can add detail fields without breaking the gate.MAX_FINAL_LOSS: "5.0"(hf_online_domino.yaml,hf_online_lilicorr.yaml). They are left unchanged and non-strict because there is no calibration run for them.Summary by CodeRabbit
New Features
Bug Fixes
Documentation