Repository navigation
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. |
Walkthrough
Merge Risk:
 |
| Check name | Status | Explanation | Resolution |
|---|---|---|---|
| Docstring Coverage |  |
Docstring coverage is 76.92% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff with cyclomatic complexity of at least 2. Documenta… | Write docstrings for the functions missing them to satisfy the coverage threshold. |
-
Check name Status Explanation Description Check Check skipped - CodeRabbit’s high-level summary is enabled. Title check The title clearly describes the main change: re-exporting partial or stale DFlash checkpoints after training. The prefix is extra metadata but does not make the title unclear or misleading. 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 PASS. The PR changes only one launcher test Python file and one shell script; it adds no modelopt-package or examples Python security-sensitive code. The added Python code contains no torch.load(..., … Full details: Security Anti-Patterns
-
Explanation
PASS. The PR changes only one launcher test Python file and one shell script; it adds no modelopt-package or examples Python security-sensitive code. The added Python code contains no torch.load(..., weights_only=False), numpy.load(..., allow_pickle=True), hardcoded trust_remote_code=True, eval(), exec(), or # nosec. No pyproject.toml or requirements.txt dependency changes are present. The existing shell --trust_remote_code invocation is unchanged.
-
Full details: Docstring Coverage
-
Explanation
Docstring coverage is 76.92% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff with cyclomatic complexity of at least 2. Documentation-only edits and deletions within surviving functions are included; unchanged diff context and fully deleted functions are excluded. Analyzed 13 functions across 2 files.
Resolution
Write docstrings for the functions missing them to satisfy the coverage threshold.
Finishing Touches: 💡 1
-
📝 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.
6b6b999 to
56f75fc
Compare
3ac8822 to
9a9b3cc
Compare
56f75fc to
66c70e1
Compare
a2d6478 to
dbddd39
Compare
…strict (#2839) > **[1/8] Day-0 speculative-decoding readiness flow.** Merge in order 1 → 8. Base: `main`. > Series: [1/8] #2839 (this) · [2/8] #2930 · [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 [common/check_regression.py](https://github.com/NVIDIA/Model-Optimizer/blob/main/tools/launcher/common/check_regression.py) gates DFlash-family training runs on final loss and accuracy. Four things stopped that gate from catching a bad run: 1. **It read the wrong checkpoint on long runs.** `find_trainer_state()` sorted `checkpoint-*` as strings, so `checkpoint-8000` ranked after `checkpoint-12000` and an older checkpoint's metrics were checked. Any run whose saved steps differ in digit count was affected; the ~92K-step full DFlash2 recipe in `hf_online_dflash2.yaml` saves every 4000 steps. It now sorts by step number. 2. **An unreadable run passed.** With thresholds set, a missing `trainer_state.json`, or one with no training logs, only warned and exited 0. An empty threshold variable (`MAX_FINAL_LOSS=""`) crashed on `float('')`. 3. **A regression never failed the job.** `dflash_online_training.sh` runs the check as `check_regression.py … || true`, so Slurm reports `COMPLETED` even when the log says `REGRESSION`. The verdict also existed only as log text. 4. **A crashed run was still graded.** The script's `ERR` trap only records a failure, so after a failed `accelerate launch` it went on to export and run the check, and a checkpoint that survived the crash could pass. Changes: - `check_regression.py`: - Sorts checkpoints by step number. - Fails when thresholds are set but there are no metrics to check, including a `trainer_state.json` of the wrong shape. An empty threshold variable counts as unset. A threshold that isn't a finite number (`7,0`, `nan`) fails as `METRICS_MISSING`. - A NaN or infinite value for a metric that has a threshold fails as `TRAIN_REGRESSION`. - Takes `--json <path>` and writes `{pass, failure_class, trainer_state, metrics, thresholds, messages}` as strict JSON (non-finite metrics as `null`). `failure_class` is `null`, `TRAIN_REGRESSION` or `METRICS_MISSING`. The verdict is serialized before the file is opened, so a failure can't leave a truncated file. - Its existing log lines (`=== Regression Check (…) ===`, `PASS:` / `REGRESSION:`, `Regression check PASSED`) are unchanged, so existing log readers keep working. A failure now also prints `Regression check FAILED`. - `dflash_online_training.sh` writes the verdict to `${OUTPUT_DIR}/regression_check.json`, next to the exports. `REGRESSION_STRICT=1` drops the `|| true`, so a regression fails the task through the existing `ERR` trap and `exit_handler`. **The default (`0` or 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.sh` also skips export and the check when training fails, still reporting the failure like any `ERR`-trap failure, and removes an old `regression_check.json` first thing, so no early exit can leave an earlier run's verdict behind. - `examples/Qwen/Qwen3-8B/hf_online_dflash2.yaml`: - Sets `REGRESSION_STRICT: "1"`. - Recalibrates `MAX_FINAL_LOSS` from 5.0 to 7.0. The 5.0 was inherited from `hf_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_ACC` stays at 0.15. - The `speculative-decoding` skill's `dflash.md` and `dflash2.md` describe the new gate: judge a run by `regression_check.json`, the two failure classes, strict mode, and DFlash2's 7.0 threshold. ### Usage ```bash MAX_FINAL_LOSS=7.0 MIN_FINAL_ACC=0.15 \ python3 common/check_regression.py /scratchspace/dflash2_bs16 --json /scratchspace/dflash2_bs16/regression_check.json echo $? # 0 pass, 1 regression or nothing to check ``` ```yaml environment: - MAX_FINAL_LOSS: "7.0" - REGRESSION_STRICT: "1" # fail the training task on a regression ``` ### Testing - `tools/launcher/tests/test_check_regression.py`: 4 new test functions (13 cases). Run against `main`'s `check_regression.py`, 17 of the file's 18 cases fail: - `test_latest_checkpoint_is_chosen_by_step_number` guards the string-sort bug. - `test_empty_threshold_var_is_unset` guards the `float('')` crash. - `test_cli_exit_code_and_json_verdict` is parametrized over a missing `trainer_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 `--json` contract. - `test_invalid_threshold_writes_failed_verdict` covers `7,0`, `nan` and `inf`. - `tools/launcher/tests/test_dflash_online_training.py` (5 cases) runs the real script in a sandbox: stubbed `pip`, `nvidia-smi`, `accelerate`, exporter and `torch`, a checkpoint, and an old passing verdict. Each case asserts the exact exit code. Against `main`'s script, all 5 fail: - A failing launch leaves no verdict and exits 1; a succeeding one grades the checkpoint and exits 0 with no `ERR`-trap error. - With a checkpoint over the threshold, `REGRESSION_STRICT=1` exits 1 and `0` exits 0, both writing `TRAIN_REGRESSION`. - `REGRESSION_STRICT=true` exits 1 without launching training and leaves no verdict. - Full launcher suite on this commit: `uv run --with pytest python3 -m pytest tests/`, 183 passed. - `dflash_online_training.sh`: `bash -n` passes. - **GPU (OCI-JHB, GB300, global batch 8):** five 2000-step runs of this example with this change passed the strict gate. Final loss was 6.37–6.71 and accuracy 0.167–0.177, against 7.0 / 0.15. The calibration run that motivated 7.0 (before this change) finished at 6.51 / 0.181. The seed-to-seed spread puts the loss margin at about 4%; see the variance note under Additional Information. - Not run: the strict-failure path and a real training crash on GPU (covered by the sandbox test above). ### 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?: ✅ (strict mode is opt-in; log output unchanged) - 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?: ✅ - 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 - Part of the Day-0 speculative-decoding workflow: the `specdec-release` skill ([8/8]) gates training on `regression_check.json`, so it needs a verdict it can read and a job state that reflects it. - **Why not reuse `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.py` is 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.py` runs *inside* the Slurm training task, on rank 0, and its verdict decides the task's exit status in strict mode. It reads `trainer_state.json` against 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. - The skill-side counterpart of `gate_run.py` for training is `specdec-release/scripts/gate_train.py` ([6/8] #2934). It reads `regression_check.json` and returns the day-0 gate shape (`{pass, failure_class, detail, …}`), with `failure_class` from the shared triage vocabulary. That is the adapter layer, so `check_regression.py` doesn't need to match `gate_run.py`. - What the two share is the convention: a boolean `pass`, and a `failure_class` that is `null` on a pass. The launcher emits only `METRICS_MISSING` and `TRAIN_REGRESSION`, and `gate_train.py` treats any other value as `TRAIN_REGRESSION`, so the launcher can add detail fields without breaking the gate. - Other Qwen3-8B examples share the inherited `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. - Variance: a seed-1 retrain landed at loss 6.705, 4.2% under the 7.0 floor. Accuracy has more headroom (11%). Recalibrating from more seeds is planned for the release gates, not this PR. <!-- This is an auto-generated comment: release notes by coderabbit.ai --> <!-- coderabbit-ai-generated:v1 --> ## Summary by CodeRabbit * **New Features** * Training runs with an output directory now produce a regression-check verdict file reporting whether configured quality thresholds passed. * Regression checks account for missing or invalid metrics and select checkpoints by numeric step. * **Bug Fixes** * Failed training no longer proceeds to checkpoint export or regression checks. * Strict mode fails the task when regression checks fail; invalid strict-mode values stop the run before training starts. * **Documentation** * Updated DFlash2 guidance and example settings for the final-loss threshold and strict-mode behavior. <!-- end of auto-generated comment: release notes by coderabbit.ai --> --------- Signed-off-by: Zhehao Hu <[email protected]> Co-authored-by: Claude Opus 5.5 <[email protected]>
dbddd39 to
6d46dab
Compare
a2bec7b to
6941971
Compare
6d46dab to
d217cd2
Compare
…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 <[email protected]> Co-authored-by: Claude Opus 5.5 <[email protected]>
dflash_online_training.sh skipped any exported-checkpoint-* directory that existed, so an export interrupted mid-write was never redone, and a resumed run kept the previous exported-checkpoint-final. An export now counts as done only once its config.json exists (exporters write weights first); otherwise the directory is removed and exported again. exported-checkpoint-final is also redone when modelopt_state.pth is newer than its config.json. Signed-off-by: Zhehao Hu <[email protected]>
d217cd2 to
84467d2
Compare
Codecov Report✅ All modified and coverable lines are covered by tests. Additional details and impacted files@@ Coverage Diff @@
## main #2931 +/- ##
==========================================
+ Coverage 73.58% 78.83% +5.24%
==========================================
Files 661 664 +3
Lines 75682 82662 +6980
==========================================
+ Hits 55693 65167 +9474
+ Misses 19989 17495 -2494
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:
|
shengliangxu
left a comment
shengliangxu
left a comment
There was a problem hiding this comment.
Bot review (gpt-6.1-sol) — DM the bot to share feedback.
Comment: config.json is created before serialization finishes, so interrupted exports can still be mistaken for completed exports.
Needs action:
- Fix the completion checks in
tools/launcher/common/specdec/dflash_online_training.shto reject empty or truncated configurations; see the inline comment. - Extend
tools/launcher/tests/test_dflash_online_training.pywith partial-export, stale/current-final, and repeat-run cases. These filesystem decisions do not require interrupting a GPU job.
| STEP=$(basename "$CKPT" | sed 's/checkpoint-//') | ||
| EXPORT_DIR="${OUTPUT_DIR}/exported-checkpoint-${STEP}" | ||
| if [ -d "$EXPORT_DIR" ]; then | ||
| if [ -f "${EXPORT_DIR}/config.json" ]; then |
There was a problem hiding this comment.
Bot comment.
DFlashExporter.export writes the weights first, but then opens config.json directly with "w" before calling json.dump. A preemption or write failure during serialization leaves an empty/truncated file that passes this check, so subsequent runs permanently skip that saved-step export. The final-export predicate has the same completion-marker weakness. Could we validate the configuration before treating the directory as complete, or publish completion atomically across the export writers? Please cover an existing empty/truncated config as well as an absent config.
There was a problem hiding this comment.
Fixed in c515154.
- Completion check: an export now counts as complete only when its
config.jsonparses to a non-empty object. This covers both the step exports and the final export, and works for either writer. - Atomic swap: each export is written to
<dir>.tmpand moved into place only on success. A failed re-export keeps the previous export. - Sharded checkpoints: if an FSDP2-sharded checkpoint has an incomplete export, that export is removed rather than re-exported. The full-HF exporter cannot load shards.
- Tests:
test_dflash_online_training.pycovers absent, empty, truncated and valid configs, sharded checkpoints, stale/current/failed final exports, and a repeat run.
shengliangxu
left a comment
shengliangxu
left a comment
There was a problem hiding this comment.
Bot review (gpt-6.1-sol) — DM the bot to share feedback.
Comment: incomplete FSDP exports can now let a successful task evaluate an older draft.
Needs action:
- Fix latest-checkpoint handling in
dflash_online_training.shso an incomplete FSDP export is repaired or fails the task; test a complete older export alongside an incomplete newest export (inline). - Preserve the previous export when publication fails in
export_draft, and add a publication-failure regression test (inline). - Update the PR description and Testing checklist to reflect JSON validation, scratch exporting, FSDP handling, and the added filesystem tests; record current validation results.
No action needed:
- ✔️ Resolved since the last review: empty/truncated-config detection and the requested partial-export, stale/current-final, and repeat-run coverage.
| if compgen -G "${CKPT}/pytorch_model_fsdp_*" >/dev/null; then | ||
| echo "WARNING: ${CKPT} is FSDP-sharded and has no complete export; removing ${EXPORT_DIR}." >&2 | ||
| rm -rf "${EXPORT_DIR}" | ||
| continue |
There was a problem hiding this comment.
Bot comment.
[P2] Do not report success while silently selecting an older FSDP draft. With a complete exported-checkpoint-1600, an incomplete exported-checkpoint-2000, and no root modelopt_state.pth, this branch deletes the newest export and continues successfully. check_regression.py::find_trainer_state still grades checkpoint 2000, but vllm_smoke_test.sh calls latest_checkpoint and selects the older 1600 draft. The real DFlashFSDP2ShardedSDExportCallback.on_save catches write failures, so training can succeed in precisely this state. Previously the unsupported full-HF export attempt failed the task; this change makes it pass with mismatched training and serving artifacts. Could we repair through the sharded draft-export path, or fail when the latest required export is incomplete, while still allowing obsolete partial checkpoints to be skipped? Extend the sandbox test with complete-older/incomplete-newest exports.
There was a problem hiding this comment.
Fixed in 5a934fc. An incomplete export of the newest FSDP-sharded checkpoint now fails the task, unless a final export will be served. Incomplete exports of older checkpoints are still skipped with a warning. test_sharded_checkpoint_is_never_reexported covers a complete 1600 export alongside an incomplete 2000 export, both with and without a final export.
| --model_path "$1" \ | ||
| --export_path "$2.tmp" \ | ||
| --trust_remote_code || return | ||
| rm -rf "$2" && mv "$2.tmp" "$2" |
There was a problem hiding this comment.
Bot comment.
[P2] Keep the previous export until publication succeeds. rm -rf followed by mv is not an atomic swap: a preemption between them, or a failing mv, removes the usable previous export even though the replacement has not been published. Readers can also encounter a missing or partially deleted directory during removal. The new failure test only fails the exporter subprocess, before this destructive step, so it does not cover the stated preservation guarantee. Could we retain the old version with rollback/recovery, or publish through an atomic version pointer, and add a sandbox case where exporting succeeds but publication fails?
There was a problem hiding this comment.
Fixed in 5a934fc. export_draft now sets the previous export aside as <dir>.old and deletes it only after the new export is in place. A failed publish restores it, and the next run restores it if the job stopped mid-swap. New sandbox cases: the exporter succeeds but the publishing mv fails (MV_FAIL), and a leftover .old with or without a published export (test_interrupted_publish_is_recovered).
shengliangxu
left a comment
shengliangxu
left a comment
There was a problem hiding this comment.
Bot review (gpt-6.1-sol) — DM the bot to share feedback.
Comment: JSON validation and rollback are fixed, but leftover root metadata can still allow success while serving an older FSDP draft.
Needs action:
- 💬 Author added the newest-FSDP guard — still require a fresh final checkpoint before bypassing it; see the inline comment in
dflash_online_training.sh. - Extend
test_sharded_checkpoint_is_never_reexportedwith leftover root metadata and an older final export, and require failure when the newest sharded export is missing.
No action needed:
- ✔️ Resolved since the last review: empty/truncated-config detection, publication rollback and mid-swap recovery, and the description/testing updates.
| # its export is written during training by DFlashFSDP2ShardedSDExportCallback. | ||
| if compgen -G "${CKPT}/pytorch_model_fsdp_*" >/dev/null; then | ||
| rm -rf "${EXPORT_DIR}" | ||
| if [ "$CKPT" = "$LATEST_CKPT" ] && [ ! -f "${OUTPUT_DIR}/modelopt_state.pth" ]; then |
There was a problem hiding this comment.
Bot comment.
[P2] A leftover root modelopt_state.pth does not establish that a current final checkpoint will be served. Consider a reused output directory containing an older root model and its valid exported-checkpoint-final (whose config is newer than the root sidecar), followed by newer FSDP-sharded checkpoints with no export for the newest step. This condition only warns because the old sidecar exists; the final-export condition also skips because the old final config parses and is newer than that sidecar. The task therefore succeeds, find_trainer_state grades the newer step, and latest_checkpoint selects the old final. With the newest export absent, the original script instead attempted the unsupported HF export and failed. Could we establish that the final source belongs to the current run/latest checkpoint before suppressing this failure, or repair through the sharded export path? Please extend the existing sandbox case with leftover root metadata and an older final, rather than only a newly seeded current final.
5a934fc to
c4bd17a
Compare
shengliangxu
left a comment
shengliangxu
left a comment
There was a problem hiding this comment.
Bot review (gpt-6.1-sol) — DM the bot to share feedback.
Comment: the previous fixes are present, but a complete newest FSDP export can still lose to a stale final export.
Needs action:
- Exclude stale final exports from serving selection in
dflash_online_training.sh; extend the sharded test with a complete newest step and leftover final metadata (inline). - Stop export processing if
.oldrecovery fails, preserving the backup; add a recovery-failure case intest_dflash_online_training.py(inline).
No action needed:
- ✔️ Previous concerns addressed: JSON completion validation, newest-incomplete-FSDP failure including leftover metadata, publication rollback/recovery, and description/testing updates.
| --export_path "${OUTPUT_DIR}/exported-checkpoint-final" \ | ||
| --trust_remote_code | ||
| FINAL_EXPORT="${OUTPUT_DIR}/exported-checkpoint-final" | ||
| if [ -f "${OUTPUT_DIR}/modelopt_state.pth" ] && { ! export_complete "$FINAL_EXPORT" || \ |
There was a problem hiding this comment.
Bot comment.
[P2] The new FINAL_IS_CURRENT check only gates an incomplete newest FSDP step; it does not prevent a stale final from winning serving selection. With complete exports for steps 1600 and 2000, plus the same leftover root sidecar and old valid final used by the new test, both step exports are skipped and this predicate leaves the old final untouched. The task succeeds, find_trainer_state() grades step 2000, but vllm_smoke_test.sh calls latest_checkpoint, which always prefers exported-checkpoint-final and serves the old draft. Could the freshness result also retire/exclude obsolete final exports, allowing the complete newest step to be served, rather than only suppressing the missing-step bypass? Extend the existing sharded parametrization with a complete newest export and final="leftover", and check the actual serving selection.
| rm -rf "$OLD" | ||
| else | ||
| echo "Restoring ${OLD%.old} from an interrupted publish" | ||
| mv -T "$OLD" "${OLD%.old}" |
There was a problem hiding this comment.
Bot comment.
[P2] A failed recovery does not stop this loop: the ERR trap only sets FAIL_EXIT. If restoring <dir>.old fails, the subsequent export_draft call begins by unconditionally deleting that same .old backup. A later exporter/publication failure therefore loses the previous usable export despite the preservation guarantee. Please stop export processing when recovery fails (before reaching the destructive cleanup), and extend the mid-swap test to fail the restoring mv and verify that the backup weights remain intact.
c4bd17a to
1a782b1
Compare
… and never serve a stale final An interrupted export can leave config.json empty or truncated, which the existence check counted as done. An export now counts as complete only when its config.json parses to a non-empty object, and each export is checked again after it is written; a failed or invalid export fails the task. The full-HF exporter cannot load FSDP2 shards, so a sharded checkpoint whose export is incomplete now fails the task instead of being re-exported. latest_checkpoint always serves exported-checkpoint-final, so a final left in a reused output_dir would win over the newest step. The final model is exported only when the root model was saved no earlier than the newest checkpoint; otherwise an existing final export is removed once the newest step export is complete. Co-Authored-By: Claude Opus 5.5 <[email protected]> Signed-off-by: Zhehao Hu <[email protected]>
1a782b1 to
abe2e9f
Compare
shengliangxu
left a comment
shengliangxu
left a comment
There was a problem hiding this comment.
Bot review (gpt-6.1-sol) — DM the bot to share feedback.
Comment: stale-final selection is fixed, but the latest push removes export preservation and rejects harmless older FSDP failures.
Needs action:
- 💬 Author previously added rollback/recovery — restore preservation in
export_draftand its failure/recovery tests; the current helper deletes the previous export before exporting (inline). - Limit the FSDP refusal in
dflash_online_training.shto required artifacts; test an incomplete older export alongside a complete newest export or current final (inline). - Update the PR description and Testing section: scratch publishing,
.oldrecovery, selective FSDP skipping, and their claimed tests are absent from this head.
No action needed:
- ✔️ Resolved since the last review: empty/truncated JSON detection, stale/current-final and repeat-run coverage, and stale-final serving selection.
| local source=$1 destination=$2 | ||
|
|
||
| # Callers use ||, which suppresses ERR handling inside this function. | ||
| rm -rf "$destination" || return 1 |
There was a problem hiding this comment.
Bot comment.
[P2] The current helper no longer contains the scratch publication/rollback mechanism described in the PR and earlier replies. Both callers pass the public export directory directly, so this rm -rf deletes a usable stale exported-checkpoint-final before the exporter starts. If loading/exporting fails or the job is preempted, the previous export is lost and readers can see a partially written replacement. The current failure test does not seed a previous final, so it cannot catch this. Could we restore scratch exporting with publication rollback and checked recovery (or an atomic version pointer), and retain regression cases for exporter failure with an existing final, publication failure, and failed .old recovery?
| --model_path "${CKPT}" \ | ||
| --export_path "${EXPORT_DIR}" \ | ||
| --trust_remote_code | ||
| if compgen -G "${CKPT}/pytorch_model_fsdp_*" >/dev/null; then |
There was a problem hiding this comment.
Bot comment.
[P2] This now fails on every incomplete sharded export, before checking whether a usable newest-step or current final export exists. For example, MiniMax's FSDP2 SHARDED_STATE_DICT flow can contain an empty-config exported-checkpoint-1600 from a caught callback write failure and a complete exported-checkpoint-2000. The base script skipped both existing directories and serving selected 2000; this version exits 1 at 1600 and never writes the regression verdict, despite the correct draft being available. The PR description promises obsolete incomplete exports are only warned/skipped. Could we determine the required serving artifact first, skip obsolete incomplete sharded exports, and fail only when no current required export can be served? Extend the sharded test with incomplete-older/complete-newest and incomplete-step/current-final cases.
…blished, and skip only unneeded FSDP exports export_draft deleted the previous export before exporting, so a failed export left nothing to serve. Each export is now written to <dir>.tmp, checked, and swapped in while the previous export waits as <dir>.old; a failed export or publish keeps the previous one, and the next run restores a .old left by a job stopped mid-swap. An incomplete export of an FSDP-sharded checkpoint no longer fails the task by itself. It is skipped with a warning, and the task fails only when the newest checkpoint lacks a complete export and no current final export will be served. Co-Authored-By: Claude Opus 5.5 <[email protected]> Signed-off-by: Zhehao Hu <[email protected]>
shengliangxu
left a comment
shengliangxu
left a comment
There was a problem hiding this comment.
Bot review (gpt-6.1-sol) — DM the bot to share feedback.
LGTM: previous concerns are resolved; human sign-off remains for the justified launcher-test fixture changes.
Needs action:
- Sign off on the fixture changes in
tools/launcher/tests/test_dflash_online_training.py: completed exports now use non-empty configs, while existing training-failure and regression assertions remain intact.
No action needed:
- ✔️ Six prior concerns addressed, including export preservation/recovery and selective FSDP handling.
![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.
Reviewed the shell changes. Careful work on a failure mode that is genuinely hard to see — approving, with one question about the staleness oracle.
Verified rather than taken on trust
export_complete's premise holds. It treats a parseable non-emptyconfig.jsonas proof of completion, which only works if config is written after the weights. It is, in both writers:hf_spec_export.py:292→:298for DFlash and:489→:498for DSpark,save_filethenjson.dump. So the check is sound rather than approximately sound..tmpand.oldreally are invisible tolatest_checkpoint.service_utils.shstrips the prefix and requires^[0-9]+$, soexported-checkpoint-2000.tmpyields2000.tmpand is skipped;...final.tmpyieldsfinal.tmp, which also misses the= finalearly return. Worth stating in the PR body, since the whole scratch-and-publish scheme depends on it and it is not obvious from either file alone.- Failing when the newest checkpoint is sharded with no complete export is the right call. The alternative is serving an older draft while everything reports success, which is the exact shape this series keeps finding.
Replacing ls -d ... | sort -t- -k2 -n with a glob plus [ -d ] is a quiet improvement too — that pipeline broke on any path with whitespace and silently produced nothing on an empty directory.
The question
-nt as the staleness oracle worries me more than the equal-timestamp case the comment addresses.
! [ "$LATEST_SAVE" -nt "${OUTPUT_DIR}/modelopt_state.pth" ]On a shared filesystem these two files are written by different processes, and on lustre neither clock skew between nodes nor attribute caching guarantees that mtime ordering matches logical ordering. Equal timestamps are handled; an inverted pair is not, and it fails in the bad direction — a stale exported-checkpoint-final is kept and served as current.
Is a step number available at the root save? If OUTPUT_DIR carries a trainer_state.json (or anything recording global_step), comparing step numbers would be exact where mtime is heuristic. If it genuinely does not — which I suspect, since the end-of-training save is save_model rather than a checkpoint write — then mtime may be the only signal available, and the right outcome is a comment saying so, so the next reader knows it is a constraint rather than a choice.
Smaller, related: when trainer_state.json is missing, LATEST_SAVE falls back to the checkpoint directory. Directory mtime updates when entries are added or removed, not when file contents are written, so the fallback compares a different quantity than the primary path. Probably still ordered correctly in practice, but it is a second semantics hiding behind one comparison.
Minor
mv -T and compgen -G are GNU/bash-specific. The shebang is #!/bin/bash and this only ever runs in the training container, so no action — noting it only because the file is otherwise portable shell.
Tests not executed.
What does this PR do?
Type of change: Bug fix
After training,
dflash_online_training.shexports each saved checkpoint, but it skipped anyexported-checkpoint-*directory that already existed. Two cases left a bad export in place:config.json, or an empty or truncated one (the exporters open it before serializing), and every later run skipped it.exported-checkpoint-final, even though training had moved on.Changes:
config.jsonparses to a non-empty object. Both writers (DFlashExporter.exportandDFlashFSDP2ShardedSDExportCallback) write it after the weights, so a valid config means the export finished.<dir>.tmpand published only on success, so readers never see a partial export. The previous export is set aside as<dir>.olduntil the new one is in place: a failed export or publish keeps it, and the next run restores it if a job stopped mid-swap.latest_checkpointignores both suffixes.latest_checkpointwould serve an older draft than the one the regression check graded.exported-checkpoint-finalis exported only when the root model was saved no earlier than the newest checkpoint, and redone whenmodelopt_state.pthis newer than itsconfig.json. Otherwiselatest_checkpointwould serve an obsolete final over the newest step, so it is removed once the newest step's export is complete.Usage
No new options. Re-running or resuming a training task now repairs partial or stale exports.
Testing
bash -npasses. Launcher unit suite on this commit: 229 passed.test_dflash_online_training.pyruns the real script on seeded output dirs: absent, empty, truncated and valid configs; an incomplete sharded newest export beside a leftover obsolete final (the task fails), and an incomplete sharded export beside a complete newest export or a current final (the task passes); stale, current and obsolete final exports, including whetherlatest_checkpointthen serves the final or the newest step; a failed export or publish that keeps the old final; a run stopped mid-swap; and a repeat run that exports nothing new. The failed-publish, mid-swap recovery and selective-FSDP cases fail on the previous commit.hf_online_dflash2.yamlruns exportedcheckpoint-2000andfinalthrough this path as first posted (before the review fixes), and the smoke test served the result each time. Not rerun on GPU since.finalafter a resume). The sandbox tests cover the decisions they make.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]) freezes and benchmarksexported-checkpoint-final, so a stale or partialfinalwould be measured as the new drafter.Summary by CodeRabbit