From 84467d29960f1616f0e99fc47dd48456d8bc7084 Mon Sep 17 00:00:00 2001 From: Zhehao Hu Date: Fri, 16 Oct 2026 21:37:12 +0000 Subject: [PATCH 1/5] launcher: re-export partial or stale DFlash checkpoints after training 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 --- .../common/specdec/dflash_online_training.sh | 16 ++++++++++------ 1 file changed, 10 insertions(+), 6 deletions(-) diff --git a/tools/launcher/common/specdec/dflash_online_training.sh b/tools/launcher/common/specdec/dflash_online_training.sh index e70800e069a..5b8ed3fb236 100644 --- a/tools/launcher/common/specdec/dflash_online_training.sh +++ b/tools/launcher/common/specdec/dflash_online_training.sh @@ -156,13 +156,15 @@ elif [ "${SLURM_PROCID:-0}" = "0" ]; then if [ -n "$OUTPUT_DIR" ]; then # Export all checkpoints that don't already have an exported version EXPORTED=0 - for CKPT in $(ls -d ${OUTPUT_DIR}/checkpoint-* 2>/dev/null | sort -t- -k2 -n); do + for CKPT in "${OUTPUT_DIR}"/checkpoint-*; do + [ -d "$CKPT" ] || continue 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 echo "Already exported: ${EXPORT_DIR}" continue fi + rm -rf "${EXPORT_DIR}" echo "=== Exporting: ${CKPT} → ${EXPORT_DIR} ===" CUDA_VISIBLE_DEVICES=0 python3 modules/Model-Optimizer/examples/speculative_decoding/scripts/export_hf_checkpoint.py \ --model_path "${CKPT}" \ @@ -170,12 +172,14 @@ elif [ "${SLURM_PROCID:-0}" = "0" ]; then --trust_remote_code EXPORTED=$((EXPORTED + 1)) done - # Also export final model if saved directly to output_dir - if [ -f "${OUTPUT_DIR}/modelopt_state.pth" ] && [ ! -d "${OUTPUT_DIR}/exported-checkpoint-final" ]; then - echo "=== Exporting: ${OUTPUT_DIR} → ${OUTPUT_DIR}/exported-checkpoint-final ===" + FINAL_EXPORT="${OUTPUT_DIR}/exported-checkpoint-final" + if [ -f "${OUTPUT_DIR}/modelopt_state.pth" ] && { [ ! -f "${FINAL_EXPORT}/config.json" ] || \ + [ "${OUTPUT_DIR}/modelopt_state.pth" -nt "${FINAL_EXPORT}/config.json" ]; }; then + rm -rf "${FINAL_EXPORT}" + echo "=== Exporting: ${OUTPUT_DIR} → ${FINAL_EXPORT} ===" CUDA_VISIBLE_DEVICES=0 python3 modules/Model-Optimizer/examples/speculative_decoding/scripts/export_hf_checkpoint.py \ --model_path "${OUTPUT_DIR}" \ - --export_path "${OUTPUT_DIR}/exported-checkpoint-final" \ + --export_path "${FINAL_EXPORT}" \ --trust_remote_code EXPORTED=$((EXPORTED + 1)) fi From abe2e9f2ec308a923fa6119794f1fb0ade977179 Mon Sep 17 00:00:00 2001 From: Zhehao Hu Date: Thu, 29 Oct 2026 04:52:04 +0000 Subject: [PATCH 2/5] launcher: count a DFlash export complete only once its config parses, 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 Signed-off-by: Zhehao Hu --- .../common/specdec/dflash_online_training.sh | 84 ++++++-- .../tests/test_dflash_online_training.py | 183 +++++++++++++++++- 2 files changed, 242 insertions(+), 25 deletions(-) diff --git a/tools/launcher/common/specdec/dflash_online_training.sh b/tools/launcher/common/specdec/dflash_online_training.sh index 5b8ed3fb236..77bebc78c31 100644 --- a/tools/launcher/common/specdec/dflash_online_training.sh +++ b/tools/launcher/common/specdec/dflash_online_training.sh @@ -134,6 +134,38 @@ if [ -n "${ACCELERATE_CONFIG:-}" ] && [ -f "${ACCELERATE_CONFIG}" ]; then echo "Using accelerate config: ${ACCELERATE_CONFIG}" fi +# Config creation precedes serialization, so existence alone does not prove completion. +export_complete() { + python3 - "$1/config.json" <<'PY' +import json +import sys + +try: + with open(sys.argv[1]) as f: + config = json.load(f) +except (OSError, ValueError): + sys.exit(1) + +sys.exit(0 if isinstance(config, dict) and config else 1) +PY +} + +export_draft() { + local source=$1 destination=$2 + + # Callers use ||, which suppresses ERR handling inside this function. + rm -rf "$destination" || return 1 + echo "=== Exporting: ${source} → ${destination} ===" + CUDA_VISIBLE_DEVICES=0 python3 modules/Model-Optimizer/examples/speculative_decoding/scripts/export_hf_checkpoint.py \ + --model_path "$source" \ + --export_path "$destination" \ + --trust_remote_code || return 1 + if ! export_complete "$destination"; then + echo "ERROR: export produced an invalid config: ${destination}" >&2 + return 1 + fi +} + set -x start_time=$(date +%s) # The ERR trap only records a failure and lets the script go on, so test the status here: @@ -160,28 +192,48 @@ elif [ "${SLURM_PROCID:-0}" = "0" ]; then [ -d "$CKPT" ] || continue STEP=$(basename "$CKPT" | sed 's/checkpoint-//') EXPORT_DIR="${OUTPUT_DIR}/exported-checkpoint-${STEP}" - if [ -f "${EXPORT_DIR}/config.json" ]; then + if export_complete "$EXPORT_DIR"; then echo "Already exported: ${EXPORT_DIR}" continue fi - rm -rf "${EXPORT_DIR}" - echo "=== Exporting: ${CKPT} → ${EXPORT_DIR} ===" - CUDA_VISIBLE_DEVICES=0 python3 modules/Model-Optimizer/examples/speculative_decoding/scripts/export_hf_checkpoint.py \ - --model_path "${CKPT}" \ - --export_path "${EXPORT_DIR}" \ - --trust_remote_code + if compgen -G "${CKPT}/pytorch_model_fsdp_*" >/dev/null; then + echo "ERROR: ${CKPT} needs re-exporting, but the HF exporter cannot load FSDP shards." >&2 + exit 1 + fi + export_draft "$CKPT" "$EXPORT_DIR" || exit 1 EXPORTED=$((EXPORTED + 1)) done + + LATEST_CKPT=$(latest_checkpoint "$OUTPUT_DIR" checkpoint-) + LATEST_SAVE="${LATEST_CKPT}/trainer_state.json" + if [ ! -f "$LATEST_SAVE" ]; then + LATEST_SAVE="$LATEST_CKPT" + fi + + FINAL_IS_CURRENT=0 + if [ -f "${OUTPUT_DIR}/modelopt_state.pth" ] && { + [ -z "$LATEST_CKPT" ] || + ! [ "$LATEST_SAVE" -nt "${OUTPUT_DIR}/modelopt_state.pth" ] + }; then + FINAL_IS_CURRENT=1 + fi + FINAL_EXPORT="${OUTPUT_DIR}/exported-checkpoint-final" - if [ -f "${OUTPUT_DIR}/modelopt_state.pth" ] && { [ ! -f "${FINAL_EXPORT}/config.json" ] || \ - [ "${OUTPUT_DIR}/modelopt_state.pth" -nt "${FINAL_EXPORT}/config.json" ]; }; then - rm -rf "${FINAL_EXPORT}" - echo "=== Exporting: ${OUTPUT_DIR} → ${FINAL_EXPORT} ===" - CUDA_VISIBLE_DEVICES=0 python3 modules/Model-Optimizer/examples/speculative_decoding/scripts/export_hf_checkpoint.py \ - --model_path "${OUTPUT_DIR}" \ - --export_path "${FINAL_EXPORT}" \ - --trust_remote_code - EXPORTED=$((EXPORTED + 1)) + if [ "$FINAL_IS_CURRENT" = 1 ]; then + if ! export_complete "$FINAL_EXPORT" || + [ "${OUTPUT_DIR}/modelopt_state.pth" -nt "${FINAL_EXPORT}/config.json" ]; then + export_draft "$OUTPUT_DIR" "$FINAL_EXPORT" || exit 1 + EXPORTED=$((EXPORTED + 1)) + fi + elif [ -e "$FINAL_EXPORT" ]; then + # Retire an obsolete final only once its replacement is available. + if [ -z "$LATEST_CKPT" ] || + ! export_complete "${OUTPUT_DIR}/exported-${LATEST_CKPT##*/}"; then + echo "ERROR: obsolete final export has no complete newest-step replacement." >&2 + exit 1 + fi + echo "Removing obsolete final export: ${FINAL_EXPORT}" + rm -rf "$FINAL_EXPORT" || exit 1 fi if [ "$EXPORTED" -eq 0 ]; then echo "No new checkpoints to export in ${OUTPUT_DIR}" diff --git a/tools/launcher/tests/test_dflash_online_training.py b/tools/launcher/tests/test_dflash_online_training.py index 1a065f48abb..205bf87e1cd 100644 --- a/tools/launcher/tests/test_dflash_online_training.py +++ b/tools/launcher/tests/test_dflash_online_training.py @@ -13,7 +13,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""dflash_online_training.sh: the regression verdict, its strict mode and its validation.""" +"""dflash_online_training.sh: export validaion and regression checks.""" import json import os @@ -25,37 +25,42 @@ _COMMON = Path(__file__).resolve().parents[1] / "common" _EXPORT = "modules/Model-Optimizer/examples/speculative_decoding/scripts/export_hf_checkpoint.py" +_CONFIG = '{"model_type": "qwen3"}' -def _run(tmp_path, train_exit=0, loss=4.0, **env_overrides): +def _run(tmp_path, train_exit=0, loss=4.0, checkpoint_time=None, **env_overrides): """Run the script in a sandbox: stub pip/nvidia-smi/accelerate/export/torch, real gate. Returns (process, verdict or None, whether training was launched). """ - (tmp_path / "common" / "specdec").mkdir(parents=True) + (tmp_path / "common" / "specdec").mkdir(parents=True, exist_ok=True) for name in ("service_utils.sh", "check_regression.py", "specdec/dflash_online_training.sh"): shutil.copy(_COMMON / name, tmp_path / "common" / name) (tmp_path / "modules/Model-Optimizer/modelopt").mkdir( - parents=True + parents=True, exist_ok=True ) # service_utils appends here export = tmp_path / _EXPORT - export.parent.mkdir(parents=True) + export.parent.mkdir(parents=True, exist_ok=True) export.write_text( "import os, sys\n" "d = sys.argv[sys.argv.index('--export_path') + 1]\n" "os.makedirs(d, exist_ok=True)\n" - "open(os.path.join(d, 'config.json'), 'w').write('{}')\n" + "open(os.path.join(d, 'model.safetensors'), 'w').write('new')\n" + "if os.environ.get('EXPORT_FAIL'):\n" + " sys.exit(1)\n" + f"config = os.environ.get('EXPORT_CONFIG', {_CONFIG!r})\n" + "open(os.path.join(d, 'config.json'), 'w').write(config)\n" ) # The script counts GPUs with torch; without it that line trips the ERR trap and every # run exits 1, which would hide the exit code under test. pystub = tmp_path / "pystub" - pystub.mkdir() + pystub.mkdir(exist_ok=True) (pystub / "torch.py").write_text( "class cuda:\n @staticmethod\n def device_count():\n return 1\n" ) launched = tmp_path / "launched" bin_dir = tmp_path / "bin" - bin_dir.mkdir() + bin_dir.mkdir(exist_ok=True) for tool, body in ( ("pip", "exit 0"), ("nvidia-smi", "echo 1"), @@ -66,10 +71,15 @@ def _run(tmp_path, train_exit=0, loss=4.0, **env_overrides): # A checkpoint with the given loss, and a passing verdict from an earlier run. out = tmp_path / "out" - (out / "checkpoint-2000").mkdir(parents=True) + (out / "checkpoint-2000").mkdir(parents=True, exist_ok=True) (out / "checkpoint-2000" / "trainer_state.json").write_text( json.dumps({"log_history": [{"loss": loss}]}) ) + if checkpoint_time is not None: + os.utime( + out / "checkpoint-2000" / "trainer_state.json", + (checkpoint_time, checkpoint_time), + ) (out / "regression_check.json").write_text('{"pass": true}') env = { @@ -91,6 +101,161 @@ def _run(tmp_path, train_exit=0, loss=4.0, **env_overrides): return proc, json.loads(verdict.read_text()) if verdict.exists() else None, launched.exists() +def _seed_export(path, config): + path.mkdir(parents=True) + (path / "model.safetensors").write_text("old") + if config is not None: + (path / "config.json").write_text(config) + + +@pytest.mark.parametrize("final", [False, True]) +@pytest.mark.parametrize("config", [None, "", '{"model_ty', "{}", "[]", _CONFIG]) +def test_incomplete_export_is_repaired(tmp_path, final, config): + """Repair invalid exports and leave complete exports untouched.""" + out = tmp_path / "out" + if final: + _seed_export(out / "exported-checkpoint-2000", _CONFIG) + (out / "modelopt_state.pth").write_text("") + export = out / f"exported-checkpoint-{'final' if final else '2000'}" + _seed_export(export, config) + proc, verdict, _ = _run(tmp_path, checkpoint_time=1) + assert proc.returncode == 0, proc.stderr + assert verdict["pass"] is True + assert (export / "config.json").read_text() == _CONFIG + assert (export / "model.safetensors").read_text() == ("old" if config == _CONFIG else "new") + + +@pytest.mark.parametrize("stale", [False, True]) +def test_final_export_tracks_root_model(tmp_path, stale): + """Refresh a valid final only when the root model is newer.""" + out = tmp_path / "out" + _seed_export(out / "exported-checkpoint-2000", _CONFIG) + export = out / "exported-checkpoint-final" + _seed_export(export, _CONFIG) + state = out / "modelopt_state.pth" + state.write_text("") + t = (export / "config.json").stat().st_mtime + (10 if stale else -10) + os.utime(state, (t, t)) + proc, _, _ = _run(tmp_path, checkpoint_time=1) + assert proc.returncode == 0, proc.stderr + assert (export / "model.safetensors").read_text() == ("new" if stale else "old") + + +@pytest.mark.parametrize("final", [False, True]) +@pytest.mark.parametrize("failure", [{"EXPORT_FAIL": "1"}, {"EXPORT_CONFIG": ""}]) +def test_failed_export_writes_no_verdict(tmp_path, final, failure): + """Exporter errors and invalid output stop the script before grading training.""" + out = tmp_path / "out" + out.mkdir() + (out / "modelopt_state.pth").write_text("") + if final: + _seed_export(out / "exported-checkpoint-2000", _CONFIG) + proc, verdict, _ = _run(tmp_path, checkpoint_time=1, **failure) + assert proc.returncode == 1, proc.stderr + assert verdict is None + if not final: + assert not (out / "exported-checkpoint-final").exists() + + +def test_repeat_run_keeps_complete_exports(tmp_path): + """A repeated run does not rewrite completed step or final exports.""" + out = tmp_path / "out" + out.mkdir() + (out / "modelopt_state.pth").write_text("") + proc, _, _ = _run(tmp_path, checkpoint_time=1) + assert proc.returncode == 0, proc.stderr + weights = [ + out / f"exported-checkpoint-{step}" / "model.safetensors" for step in ("2000", "final") + ] + for path in weights: + os.utime(path, (1, 1)) + proc, _, _ = _run(tmp_path, checkpoint_time=1) + assert proc.returncode == 0, proc.stderr + assert all(path.stat().st_mtime == 1 for path in weights) + + +@pytest.mark.parametrize("config", [None, ""]) +def test_sharded_export_failure_preserves_leftover_final(tmp_path, config): + """A leftover final cannot bypass an unexportable newest checkpoint.""" + out = tmp_path / "out" + checkpoint = out / "checkpoint-2000" + (checkpoint / "pytorch_model_fsdp_0").mkdir(parents=True) + + state = out / "modelopt_state.pth" + state.write_text("") + os.utime(state, (100, 100)) + + final = out / "exported-checkpoint-final" + _seed_export(final, _CONFIG) + os.utime(final / "config.json", (110, 110)) + + step_export = out / "exported-checkpoint-2000" + if config is not None: + _seed_export(step_export, config) + + # The stub exporter would succeed, so failure must come from the shard guard. + proc, verdict, _ = _run(tmp_path, checkpoint_time=200) + + assert proc.returncode == 1, proc.stderr + assert verdict is None + assert (final / "model.safetensors").read_text() == "old" + assert (final / "config.json").read_text() == _CONFIG + if config is None: + assert not step_export.exists() + else: + assert (step_export / "model.safetensors").read_text() == "old" + assert (step_export / "config.json").read_text() == config + + +@pytest.mark.parametrize("existing_final", [False, True]) +@pytest.mark.parametrize("root_time", [100, 200, 300]) +def test_final_selection_tracks_latest_checkpoint(tmp_path, existing_final, root_time): + """Only a current root model may supply the preferred final export.""" + out = tmp_path / "out" + step_export = out / "exported-checkpoint-2000" + _seed_export(step_export, _CONFIG) + + state = out / "modelopt_state.pth" + state.write_text("") + current_root = root_time >= 200 + os.utime(state, (root_time, root_time)) + + final = out / "exported-checkpoint-final" + if existing_final: + _seed_export(final, _CONFIG) + os.utime(final / "config.json", (root_time + 10, root_time + 10)) + + proc, verdict, _ = _run(tmp_path, checkpoint_time=200) + + assert proc.returncode == 0, proc.stderr + assert verdict["pass"] is True + assert verdict["trainer_state"].endswith("checkpoint-2000/trainer_state.json") + assert final.exists() == current_root + assert (step_export / "model.safetensors").read_text() == "old" + if current_root: + assert (final / "model.safetensors").read_text() == ("old" if existing_final else "new") + + selected = subprocess.run( + [ + "bash", + "-c", + "source common/service_utils.sh >/dev/null; " + 'latest_checkpoint "$1" exported-checkpoint-', + "bash", + str(out), + ], + cwd=tmp_path, + env={ + **os.environ, + "PATH": f"{tmp_path / 'bin'}:{os.environ['PATH']}", + }, + capture_output=True, + text=True, + check=True, + ) + assert selected.stdout.strip() == str(final if current_root else step_export) + + @pytest.mark.parametrize("train_exit", [0, 1]) def test_failed_training_writes_no_verdict(tmp_path, train_exit): """A crash must not let a surviving checkpoint, or an old verdict, pass the gate.""" From a2127efd68978024771709d2fa47ef57011e13fb Mon Sep 17 00:00:00 2001 From: Zhehao Hu Date: Thu, 29 Oct 2026 21:40:00 +0000 Subject: [PATCH 3/5] launcher: keep the previous DFlash export until its replacement is published, 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 .tmp, checked, and swapped in while the previous export waits as .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 Signed-off-by: Zhehao Hu --- .../common/specdec/dflash_online_training.sh | 48 +++++- .../tests/test_dflash_online_training.py | 142 +++++++++++++++--- 2 files changed, 164 insertions(+), 26 deletions(-) diff --git a/tools/launcher/common/specdec/dflash_online_training.sh b/tools/launcher/common/specdec/dflash_online_training.sh index 77bebc78c31..83ae850bfa0 100644 --- a/tools/launcher/common/specdec/dflash_online_training.sh +++ b/tools/launcher/common/specdec/dflash_online_training.sh @@ -150,20 +150,43 @@ sys.exit(0 if isinstance(config, dict) and config else 1) PY } +recover_export() { + local destination=$1 backup="$1.old" + + [ -e "$backup" ] || return 0 + if [ -e "$destination" ]; then + # Only a complete publication supersedes the backup. + export_complete "$destination" || return 1 + rm -rf "$backup" || return 1 + else + mv -T "$backup" "$destination" || return 1 + fi +} + export_draft() { local source=$1 destination=$2 + local scratch="$2.tmp" backup="$2.old" # Callers use ||, which suppresses ERR handling inside this function. - rm -rf "$destination" || return 1 + recover_export "$destination" || return 1 + rm -rf "$scratch" || return 1 echo "=== Exporting: ${source} → ${destination} ===" CUDA_VISIBLE_DEVICES=0 python3 modules/Model-Optimizer/examples/speculative_decoding/scripts/export_hf_checkpoint.py \ --model_path "$source" \ - --export_path "$destination" \ + --export_path "$scratch" \ --trust_remote_code || return 1 - if ! export_complete "$destination"; then - echo "ERROR: export produced an invalid config: ${destination}" >&2 + if ! export_complete "$scratch"; then + echo "ERROR: export produced an invalid config: ${scratch}" >&2 + return 1 + fi + if [ -e "$destination" ]; then + mv -T "$destination" "$backup" || return 1 + fi + if ! mv -T "$scratch" "$destination"; then + recover_export "$destination" || return 1 return 1 fi + rm -rf "$backup" || return 1 } set -x @@ -188,6 +211,10 @@ elif [ "${SLURM_PROCID:-0}" = "0" ]; then if [ -n "$OUTPUT_DIR" ]; then # Export all checkpoints that don't already have an exported version EXPORTED=0 + for BACKUP in "${OUTPUT_DIR}"/exported-checkpoint-*.old; do + [ -d "$BACKUP" ] || continue + recover_export "${BACKUP%.old}" || exit 1 + done for CKPT in "${OUTPUT_DIR}"/checkpoint-*; do [ -d "$CKPT" ] || continue STEP=$(basename "$CKPT" | sed 's/checkpoint-//') @@ -197,19 +224,19 @@ elif [ "${SLURM_PROCID:-0}" = "0" ]; then continue fi if compgen -G "${CKPT}/pytorch_model_fsdp_*" >/dev/null; then - echo "ERROR: ${CKPT} needs re-exporting, but the HF exporter cannot load FSDP shards." >&2 - exit 1 + echo "WARNING: cannot re-export FSDP shards in ${CKPT}; checking the required export below." >&2 + continue fi export_draft "$CKPT" "$EXPORT_DIR" || exit 1 EXPORTED=$((EXPORTED + 1)) done - LATEST_CKPT=$(latest_checkpoint "$OUTPUT_DIR" checkpoint-) LATEST_SAVE="${LATEST_CKPT}/trainer_state.json" if [ ! -f "$LATEST_SAVE" ]; then LATEST_SAVE="$LATEST_CKPT" fi + # Equal timestamps cannot establish ordering; only a newer checkpoint makes the root stale. FINAL_IS_CURRENT=0 if [ -f "${OUTPUT_DIR}/modelopt_state.pth" ] && { [ -z "$LATEST_CKPT" ] || @@ -218,6 +245,13 @@ elif [ "${SLURM_PROCID:-0}" = "0" ]; then FINAL_IS_CURRENT=1 fi + # Skipped sharded exports are safe only if the current serving artifact is available. + if [ "$FINAL_IS_CURRENT" != 1 ] && [ -n "$LATEST_CKPT" ] && + ! export_complete "${OUTPUT_DIR}/exported-${LATEST_CKPT##*/}"; then + echo "ERROR: newest checkpoint has no complete export: ${LATEST_CKPT}" >&2 + exit 1 + fi + FINAL_EXPORT="${OUTPUT_DIR}/exported-checkpoint-final" if [ "$FINAL_IS_CURRENT" = 1 ]; then if ! export_complete "$FINAL_EXPORT" || diff --git a/tools/launcher/tests/test_dflash_online_training.py b/tools/launcher/tests/test_dflash_online_training.py index 205bf87e1cd..4ded06482da 100644 --- a/tools/launcher/tests/test_dflash_online_training.py +++ b/tools/launcher/tests/test_dflash_online_training.py @@ -13,7 +13,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""dflash_online_training.sh: export validaion and regression checks.""" +"""dflash_online_training.sh: export validation and regression checks.""" import json import os @@ -26,6 +26,7 @@ _COMMON = Path(__file__).resolve().parents[1] / "common" _EXPORT = "modules/Model-Optimizer/examples/speculative_decoding/scripts/export_hf_checkpoint.py" _CONFIG = '{"model_type": "qwen3"}' +_MV = shutil.which("mv") def _run(tmp_path, train_exit=0, loss=4.0, checkpoint_time=None, **env_overrides): @@ -43,7 +44,9 @@ def _run(tmp_path, train_exit=0, loss=4.0, checkpoint_time=None, **env_overrides export.parent.mkdir(parents=True, exist_ok=True) export.write_text( "import os, sys\n" - "d = sys.argv[sys.argv.index('--export_path') + 1]\n" + "src = sys.argv[sys.argv.index('--model_path') + 1]\n" + f"open({str(tmp_path / 'exports.log')!r}, 'a').write(src + '\\n')\n" + "d =sys.argv[sys.argv.index('--export_path') + 1]\n" "os.makedirs(d, exist_ok=True)\n" "open(os.path.join(d, 'model.safetensors'), 'w').write('new')\n" "if os.environ.get('EXPORT_FAIL'):\n" @@ -65,6 +68,14 @@ def _run(tmp_path, train_exit=0, loss=4.0, checkpoint_time=None, **env_overrides ("pip", "exit 0"), ("nvidia-smi", "echo 1"), ("accelerate", f"touch {launched}; exit {train_exit}"), + ( + "mv", + ( + '[ -n "$PUBLISH_FAIL" ] && [[ "$2" == *.tmp ]] && exit 1\n' + '[ -n "$RECOVER_FAIL" ] && [[ "$2" == *.old ]] && exit 1\n' + f'exec {_MV} "$@"' + ), + ), ): (bin_dir / tool).write_text(f"#!/bin/bash\n{body}\n") (bin_dir / tool).chmod(0o755) @@ -180,26 +191,23 @@ def test_sharded_export_failure_preserves_leftover_final(tmp_path, config): out = tmp_path / "out" checkpoint = out / "checkpoint-2000" (checkpoint / "pytorch_model_fsdp_0").mkdir(parents=True) - state = out / "modelopt_state.pth" state.write_text("") os.utime(state, (100, 100)) - final = out / "exported-checkpoint-final" _seed_export(final, _CONFIG) os.utime(final / "config.json", (110, 110)) - step_export = out / "exported-checkpoint-2000" if config is not None: _seed_export(step_export, config) - # The stub exporter would succeed, so failure must come from the shard guard. + # The stub exporter would succeed, so the script must reject the missing required export. proc, verdict, _ = _run(tmp_path, checkpoint_time=200) - assert proc.returncode == 1, proc.stderr assert verdict is None assert (final / "model.safetensors").read_text() == "old" assert (final / "config.json").read_text() == _CONFIG + assert not (tmp_path / "exports.log").exists() if config is None: assert not step_export.exists() else: @@ -207,26 +215,119 @@ def test_sharded_export_failure_preserves_leftover_final(tmp_path, config): assert (step_export / "config.json").read_text() == config +@pytest.mark.parametrize( + ("incomplete_step", "current_final"), + [("1600", None), ("2000", "existing"), ("2000", "export")], +) +def test_incomplete_shards_allow_current_export(tmp_path, incomplete_step, current_final): + """An obsolete shard export must not block a complete newest step or current final.""" + out = tmp_path / "out" + for step in ("1600", "2000"): + (out / f"checkpoint-{step}" / "pytorch_model_fsdp_0").mkdir(parents=True) + _seed_export( + out / f"exported-checkpoint-{step}", "" if step == incomplete_step else _CONFIG + ) + final = out / "exported-checkpoint-final" + if current_final: + state = out / "modelopt_state.pth" + state.write_text("") + os.utime(state, (300, 300)) + if current_final == "existing": + _seed_export(final, _CONFIG) + os.utime(final / "config.json", (310, 310)) + + proc, verdict, _ = _run(tmp_path, checkpoint_time=200) + assert proc.returncode == 0, proc.stderr + assert verdict["pass"] is True + required = final if current_final else out / "exported-checkpoint-2000" + assert json.loads((required / "config.json").read_text()) + log = tmp_path / "exports.log" + assert (log.read_text().splitlines() if log.exists() else []) == ( + [str(out)] if current_final == "export" else [] + ) + assert _latest_export(tmp_path) == str(required) + + +@pytest.mark.parametrize( + "failure", + [ + {"EXPORT_FAIL": "1"}, + {"EXPORT_CONFIG": ""}, + {"PUBLISH_FAIL": "1"}, + {"PUBLISH_FAIL": "1", "RECOVER_FAIL": "1"}, + ], +) +def test_failed_replacement_preserves_final(tmp_path, failure): + """Export and publication failures retain the previous usable final.""" + out = tmp_path / "out" + _seed_export(out / "exported-checkpoint-2000", _CONFIG) + final = out / "exported-checkpoint-final" + _seed_export(final, _CONFIG) + os.utime(final / "config.json", (250, 250)) + state = out / "modelopt_state.pth" + state.write_text("") + os.utime(state, (300, 300)) + proc, verdict, _ = _run(tmp_path, checkpoint_time=200, **failure) + assert proc.returncode == 1, proc.stderr + assert verdict is None + backup = final.with_name(final.name + ".old") + preserved = backup if failure.get("RECOVER_FAIL") else final + assert (preserved / "model.safetensors").read_text() == "old" + assert (preserved / "config.json").read_text() == _CONFIG + assert backup.exists() == bool(failure.get("RECOVER_FAIL")) + + +@pytest.mark.parametrize( + ("published", "recovery_fails"), [(False, False), (True, False), (False, True)] +) +def test_interrupted_publication_recovery(tmp_path, published, recovery_fails): + """Recover a swap, stopping without deleting the backup if restoration fails.""" + out = tmp_path / "out" + _seed_export(out / "exported-checkpoint-2000", _CONFIG) + state = out / "modelopt_state.pth" + state.write_text("") + os.utime(state, (300, 300)) + final = out / "exported-checkpoint-final" + backup = final.with_name(final.name + ".old") + _seed_export(backup, _CONFIG) + os.utime(backup / "config.json", (310, 310)) + if published: + _seed_export(final, _CONFIG) + (final / "model.safetensors").write_text("published") + os.utime(final / "config.json", (310, 310)) + proc, verdict, _ = _run( + tmp_path, checkpoint_time=200, **({"RECOVER_FAIL": "1"} if recovery_fails else {}) + ) + assert proc.returncode == int(recovery_fails), proc.stderr + assert not (tmp_path / "exports.log").exists() + if recovery_fails: + assert verdict is None + assert not final.exists() + assert (backup / "model.safetensors").read_text() == "old" + assert (backup / "config.json").read_text() == _CONFIG + else: + assert verdict["pass"] is True + assert not backup.exists() + assert (final / "model.safetensors").read_text() == ("published" if published else "old") + + @pytest.mark.parametrize("existing_final", [False, True]) @pytest.mark.parametrize("root_time", [100, 200, 300]) def test_final_selection_tracks_latest_checkpoint(tmp_path, existing_final, root_time): - """Only a current root model may supply the preferred final export.""" + """Retire finals only when the checkpoint is strictly newer than the root.""" out = tmp_path / "out" step_export = out / "exported-checkpoint-2000" _seed_export(step_export, _CONFIG) - state = out / "modelopt_state.pth" state.write_text("") current_root = root_time >= 200 os.utime(state, (root_time, root_time)) - final = out / "exported-checkpoint-final" if existing_final: _seed_export(final, _CONFIG) os.utime(final / "config.json", (root_time + 10, root_time + 10)) proc, verdict, _ = _run(tmp_path, checkpoint_time=200) - assert proc.returncode == 0, proc.stderr assert verdict["pass"] is True assert verdict["trainer_state"].endswith("checkpoint-2000/trainer_state.json") @@ -235,25 +336,28 @@ def test_final_selection_tracks_latest_checkpoint(tmp_path, existing_final, root if current_root: assert (final / "model.safetensors").read_text() == ("old" if existing_final else "new") + assert _latest_export(tmp_path) == str(final if current_root else step_export) + + +def _latest_export(tmp_path): selected = subprocess.run( [ "bash", "-c", - "source common/service_utils.sh >/dev/null; " - 'latest_checkpoint "$1" exported-checkpoint-', + ( + "source common/service_utils.sh >/dev/null; " + 'latest_checkpoint "$1" exported-checkpoint-' + ), "bash", - str(out), + str(tmp_path / "out"), ], cwd=tmp_path, - env={ - **os.environ, - "PATH": f"{tmp_path / 'bin'}:{os.environ['PATH']}", - }, + env={**os.environ, "PATH": f"{tmp_path / 'bin'}:{os.environ['PATH']}"}, capture_output=True, text=True, check=True, ) - assert selected.stdout.strip() == str(final if current_root else step_export) + return selected.stdout.strip() @pytest.mark.parametrize("train_exit", [0, 1]) From 385a267503ce95fb2dd10ad7142dc35925c4d868 Mon Sep 17 00:00:00 2001 From: Zhehao Hu Date: Fri, 16 Oct 2026 23:54:43 +0000 Subject: [PATCH 4/5] launcher: clean up the Qwen3-8B DFlash2 launcher examples hf_online_dflash2.yaml: - Add task_2: vLLM smoke test of the newest export (DRAFT_CKPT_DIR), gated on MIN_ACCEPTANCE_LENGTH 1.5 with 7 spec tokens, a floor that a healthy 2000-step run of this config clears. - Run task_0 and task_1 on the same TRT-LLM image (1.3.0rc20); task_0 was pinned to rc10 while task_1 took the factory default. Both examples: - Route every artifact path through global_vars.output_dir (training output, exports, SMOKE_JSON), so one override moves the producer and its consumer together, e.g. to rerun the smoke task against an earlier experiment. - Pin vllm/vllm-openai:v0.31.0 (multi-arch, contains vllm#52816 for DFlash2) instead of latest/nightly, and set TRITON_CACHE_DIR node-local. - hf_streaming_dflash2.yaml exports to /export via EXPORT_PATH. Signed-off-by: Zhehao Hu --- .../Qwen/Qwen3-8B/hf_online_dflash2.yaml | 34 +++++++++++++++++-- .../Qwen/Qwen3-8B/hf_streaming_dflash2.yaml | 27 ++++++++++----- 2 files changed, 50 insertions(+), 11 deletions(-) diff --git a/tools/launcher/examples/Qwen/Qwen3-8B/hf_online_dflash2.yaml b/tools/launcher/examples/Qwen/Qwen3-8B/hf_online_dflash2.yaml index 6d08abba99c..5cb5dbb6ff7 100644 --- a/tools/launcher/examples/Qwen/Qwen3-8B/hf_online_dflash2.yaml +++ b/tools/launcher/examples/Qwen/Qwen3-8B/hf_online_dflash2.yaml @@ -9,9 +9,10 @@ # See the dflash2.yaml recipe and # modelopt/torch/speculative/plugins/{modeling,hf}_dflash2.py. # -# 2-step pipeline: +# 3-step pipeline: # task_0: Build training conversations (Daring-Anteater multi-turn SFT, 50K) # task_1: Online DFlash2 training + export of the drafter checkpoint +# task_2: vLLM smoke test of the exported drafter, gated on acceptance length # # As configured this is a short convergence check (max_steps=2000), matching the # other Qwen3-8B online examples so it finishes on one node. To reproduce the @@ -35,6 +36,10 @@ job_name: Qwen3-8B_DFlash2_online pipeline: global_vars: hf_model: /hf-local/Qwen/Qwen3-8B + # Training output: checkpoints, exported-checkpoint-*, and the gate verdicts. task_1 + # writes it and task_2 reads it, so override this one key to move both, e.g. to + # /cicd//dflash2_bs16 to rerun task_2 against an earlier experiment. + output_dir: /scratchspace/dflash2_bs16 # Step 1: Build input conversations. example_data_config.yaml enables only the # daring-anteater source (train: 50000) — multi-turn SFT with real assistant @@ -50,7 +55,8 @@ pipeline: nodes: 1 ntasks_per_node: 1 gpus_per_node: 1 - container: nvcr.io/nvidia/tensorrt-llm/release:1.3.0rc10 + # Same image as task_1, so both tasks see the same Python stack. + container: nvcr.io/nvidia/tensorrt-llm/release:1.3.0rc20 # Step 2: Online DFlash2 training (the script exports the drafter at the end). # Consumes the conversations built in task_0 (shared via /scratchspace). @@ -61,7 +67,7 @@ pipeline: - model.model_name_or_path=<> - data.data_path=/scratchspace/data/train.jsonl - data.chat_template=examples/Qwen/Qwen3-8B/chat_template_train.jinja - - training.output_dir=/scratchspace/dflash2_bs16 + - training.output_dir=<> - training.per_device_train_batch_size=1 - training.num_train_epochs=1 - training.max_steps=2000 @@ -108,3 +114,25 @@ pipeline: nodes: 1 ntasks_per_node: 1 gpus_per_node: 8 + container: nvcr.io/nvidia/tensorrt-llm/release:1.3.0rc20 + + # Step 3: Serve Qwen3-8B with the newest exported drafter and check it is accepted. + # DRAFT_CKPT_DIR picks exported-checkpoint-final (else the highest step) from task_1. + # The method stays "dflash": vLLM selects the DFlash2 path from the checkpoint's + # architectures: ["DFlash2DraftModel"] (vllm#52816, in v0.31.0). + task_2: + script: common/specdec/vllm_smoke_test.sh + environment: + - HF_MODEL_CKPT: <> + - DRAFT_CKPT_DIR: <> + - SPEC_METHOD: "dflash" + - NUM_SPEC_TOKENS: "7" + - MIN_ACCEPTANCE_LENGTH: "1.5" + - SMOKE_JSON: <>/smoke_check.json + - TRITON_CACHE_DIR: /tmp/triton-cache + slurm_config: + _factory_: "slurm_factory" + nodes: 1 + ntasks_per_node: 1 + gpus_per_node: 1 + container: vllm/vllm-openai:v0.31.0 diff --git a/tools/launcher/examples/Qwen/Qwen3-8B/hf_streaming_dflash2.yaml b/tools/launcher/examples/Qwen/Qwen3-8B/hf_streaming_dflash2.yaml index 89c06eeceed..4dd20a6c66f 100644 --- a/tools/launcher/examples/Qwen/Qwen3-8B/hf_streaming_dflash2.yaml +++ b/tools/launcher/examples/Qwen/Qwen3-8B/hf_streaming_dflash2.yaml @@ -15,10 +15,10 @@ # task_1: Streaming train — 1 serve node (TP=4) + 1 trainer node (4-rank DDP) # task_2: vLLM smoke test with the exported drafter # -# Site-specific settings left unset on purpose: an aarch64 cluster needs its own image -# (the default is x86), slurm_config.time where the queue's default is short, the NIXL and -# NCCL HCA pinning an InfiniBand fabric needs, and container-visible cache paths (srun -# runs with --no-container-mount-home, and TRITON_CACHE_DIR must be node-local). +# Site-specific settings left unset on purpose: slurm_config.time where the queue's default +# is short, the NIXL and NCCL HCA pinning an InfiniBand fabric needs, and container-visible +# cache paths such as HF_HOME (srun runs with --no-container-mount-home). The pinned images +# are multi-arch (x86_64 and aarch64), and TRITON_CACHE_DIR is set node-local below. # # Usage: # uv run launch.py --yaml examples/Qwen/Qwen3-8B/hf_streaming_dflash2.yaml --yes @@ -31,6 +31,10 @@ pipeline: global_vars: hf_model: /hf-local/Qwen/Qwen3-8B + # Training output: checkpoints, the exported drafter (export/), and the smoke verdict. + # task_1 writes it and task_2 reads it, so override this one key to move both, e.g. to + # /cicd//dflash2 to rerun task_2 against an earlier experiment. + output_dir: /scratchspace/dflash2 # Step 1: Build input conversations task_0: @@ -65,7 +69,7 @@ pipeline: # peels the last one off as the (unused) KD target and the draft's fc gets # 4x4096 where it wants 5x4096. - data.final_aux_is_base_hidden=true - - training.output_dir=/scratchspace/dflash2 + - training.output_dir=<> - training.training_seq_len=4096 - training.disable_tqdm=true # The serve prefills the corpus conversation rather than generating, so the @@ -98,12 +102,17 @@ pipeline: - STREAMING_NUM_WORKERS: "4" # DFlash2 uses a custom modeling file; export must trust remote code. - EXPORT_EXTRA_ARGS: "--trust_remote_code" + # Where the drafter is exported (train_eagle_streaming.sh defaults to /scratchspace/export). + - EXPORT_PATH: <>/export + # Node-local: a shared Triton cache breaks when jobs on different nodes write it. + - TRITON_CACHE_DIR: /tmp/triton-cache slurm_config: _factory_: "slurm_factory" nodes: 2 ntasks_per_node: 1 gpus_per_node: 4 - container: vllm/vllm-openai:latest + # A release with DFlash2 serving (vllm#52816). + container: vllm/vllm-openai:v0.31.0 # Step 3: vLLM smoke test (uses the exported checkpoint from training). # The method stays "dflash": vLLM has no separate dflash2 method and selects the @@ -112,13 +121,15 @@ pipeline: script: common/specdec/vllm_smoke_test.sh environment: - HF_MODEL_CKPT: <> - - DRAFT_MODEL: /scratchspace/export + - DRAFT_MODEL: <>/export - SPEC_METHOD: "dflash" - NUM_SPEC_TOKENS: "7" - MIN_ACCEPTANCE_LENGTH: "1.2" + - SMOKE_JSON: <>/smoke_check.json + - TRITON_CACHE_DIR: /tmp/triton-cache slurm_config: _factory_: "slurm_factory" - container: vllm/vllm-openai:nightly + container: vllm/vllm-openai:v0.31.0 nodes: 1 ntasks_per_node: 1 gpus_per_node: 1 From 59191c0ed1034727667762048b983024fbcf878f Mon Sep 17 00:00:00 2001 From: Zhehao Hu Date: Fri, 23 Oct 2026 04:40:12 +0000 Subject: [PATCH 5/5] launcher: train the Qwen3-8B DFlash2 examples on the block vLLM serves vLLM serves DFlash2 with a block of 1 + num_speculative_tokens (8 with the examples' NUM_SPEC_TOKENS 7) and ignores the checkpoint's block_size, while both examples trained on 16-token blocks. A dev-budget rerun of the online example with block 8 gave the same acceptance length as block 16 (1.75 at concurrency 1/8/32), so the served block costs nothing and removes a train/serve mismatch. It also matches the library default and the published DFlash2 drafters. The online example drops dflash_loss_decay_factor: the recipe's dpace objective ignores it, as the training log warns. Co-Authored-By: Claude Opus 5.5 Signed-off-by: Zhehao Hu --- .../examples/Qwen/Qwen3-8B/hf_online_dflash2.yaml | 9 +++++---- .../examples/Qwen/Qwen3-8B/hf_streaming_dflash2.yaml | 2 ++ 2 files changed, 7 insertions(+), 4 deletions(-) diff --git a/tools/launcher/examples/Qwen/Qwen3-8B/hf_online_dflash2.yaml b/tools/launcher/examples/Qwen/Qwen3-8B/hf_online_dflash2.yaml index 5cb5dbb6ff7..9eeb83f2da8 100644 --- a/tools/launcher/examples/Qwen/Qwen3-8B/hf_online_dflash2.yaml +++ b/tools/launcher/examples/Qwen/Qwen3-8B/hf_online_dflash2.yaml @@ -80,11 +80,12 @@ pipeline: - training.logging_steps=100 - training.disable_tqdm=true - training.answer_only_loss=true - # Draft backbone — identical to the DFlash / DSpark / Domino arms so the - # only difference between them is the correction head. - - dflash.dflash_block_size=16 + # Draft backbone — as in the DFlash / DSpark / Domino arms, except the block size. + # Train on the block vLLM serves, 1 + NUM_SPEC_TOKENS (task_2); vLLM ignores the + # checkpoint's block_size. A 16-token block gave the same AL at the dev budget. + # dflash_loss_decay_factor is not set: the recipe's dpace objective ignores it. + - dflash.dflash_block_size=8 - dflash.dflash_num_anchors=512 - - dflash.dflash_loss_decay_factor=7 - dflash.dflash_mask_token_id=151669 - dflash.dflash_self_logit_distillation=false - dflash.dflash_architecture_config.num_hidden_layers=5 diff --git a/tools/launcher/examples/Qwen/Qwen3-8B/hf_streaming_dflash2.yaml b/tools/launcher/examples/Qwen/Qwen3-8B/hf_streaming_dflash2.yaml index 4dd20a6c66f..c414ffb0b1e 100644 --- a/tools/launcher/examples/Qwen/Qwen3-8B/hf_streaming_dflash2.yaml +++ b/tools/launcher/examples/Qwen/Qwen3-8B/hf_streaming_dflash2.yaml @@ -71,6 +71,8 @@ pipeline: - data.final_aux_is_base_hidden=true - training.output_dir=<> - training.training_seq_len=4096 + # The recipe's 16; train on the block vLLM serves, 1 + NUM_SPEC_TOKENS (task_2). + - dflash.dflash_block_size=8 - training.disable_tqdm=true # The serve prefills the corpus conversation rather than generating, so the # assistant turn is present; Qwen3-8B's stock template just has no