diff --git a/ci/run_example_test.sh b/ci/run_example_test.sh index 10de4c6..531562a 100644 --- a/ci/run_example_test.sh +++ b/ci/run_example_test.sh @@ -15,6 +15,9 @@ case "${platform}/${backend}/${drafter}" in gpu/vllm/dspark) example="examples/run_qwen3-8b_drafter_dspark_vllm.sh" ;; + gpu/vllm/peagle|gpu/vllm/domino) + example="examples/run_qwen3-8b_drafter_domino_peagle_separate_training.sh" + ;; gpu/sglang/eagle3) example="examples/run_qwen3-8b_drafter_eagle3_sglang.sh" ;; @@ -40,7 +43,7 @@ case "${platform}/${backend}/${drafter}" in example="examples/run_qwen3-8b_drafter_dflash_sglang.sh" ;; *) - echo "usage: $0 {gpu|npu} {vllm|sglang} {eagle3|megatron-eagle3|dflash|dspark}" >&2 + echo "usage: $0 {gpu|npu} {vllm|sglang} {eagle3|megatron-eagle3|dflash|dspark|peagle|domino}" >&2 exit 2 ;; esac @@ -58,6 +61,15 @@ for name in "${required_vars[@]}"; do fi done +# Domino and P-EAGLE are not engine-level speculative algorithms, so neither can +# be trained inside the rollout loop. They use the two-stage separate-training +# workflow instead: stage 1 rolls out with the engine algorithm whose +# hidden-state layout the drafter consumes and writes a feature store, stage 2 +# trains the drafter offline from that same store. +# The example owns the collect-to-train algorithm pairing; the runner only needs +# the collect side of it to pick the matching cached draft model. +separate_training="false" + case "${drafter}" in eagle3|megatron-eagle3) draft_model="${SPECO_EAGLE3_DRAFT_MODEL:-}" @@ -71,9 +83,19 @@ case "${drafter}" in draft_model="${SPECO_DSPARK_DRAFT_MODEL:-}" draft_algorithm="DSPARK" ;; + peagle) + draft_model="${SPECO_EAGLE3_DRAFT_MODEL:-}" + draft_algorithm="EAGLE3" + separate_training="true" + ;; + domino) + draft_model="${SPECO_DFLASH_DRAFT_MODEL:-}" + draft_algorithm="DFLASH" + separate_training="true" + ;; esac if [[ -z "${draft_model}" ]]; then - echo "required ${drafter} draft model environment variable is not set" >&2 + echo "required ${drafter} collect-stage draft model environment variable is not set" >&2 exit 2 fi @@ -216,9 +238,62 @@ if [[ "${drafter}" == "dspark" ]]; then ) fi +train_overrides=() +if [[ "${separate_training}" == "true" ]]; then + feature_store_dir="${SPECO_CKPT_DIR}/${drafter}_features" + # A draft init path that does not exist yet cold-starts the drafter from the + # target config, so this lane needs no extra model in the CI cache. + draft_init_path="${SPECO_CKPT_DIR}/${drafter}_draft_init" + + # The example already sets the two-stage shape (collect_only/offline modes, + # the standalone launcher, the feature-store flags and the per-algorithm + # hyperparameters). Override only what CI has to control: where the stages + # meet on disk, and the sizes that keep a smoke run cheap. + # A configured SPECO_CKPT_DIR persists between jobs on a self-hosted runner, + # so drop any earlier shards: otherwise a collect stage that produced nothing + # still trains, and the job goes green on stale features. + rm -rf "${feature_store_dir}" "${draft_init_path}" + + overrides+=( + "actor_rollout_ref.rollout.drafter.training.feature_store.path=${feature_store_dir}" + "actor_rollout_ref.rollout.drafter.training.feature_store.max_samples_per_shard=32" + ) + + train_overrides=( + "actor_rollout_ref.model.path=${SPECO_TARGET_MODEL}" + "actor_rollout_ref.rollout.drafter.model_path=${draft_init_path}" + "actor_rollout_ref.rollout.drafter.checkpoint_path=${SPECO_CKPT_DIR}/${drafter}_draft_ckpts" + "actor_rollout_ref.rollout.drafter.training.max_steps=1" + "actor_rollout_ref.rollout.drafter.training.save_interval_steps=1" + "actor_rollout_ref.rollout.drafter.training.batch_size_per_gpu=${SPECO_DRAFTER_BATCH_SIZE_PER_GPU:-1}" + "actor_rollout_ref.rollout.drafter.training.feature_store.path=${feature_store_dir}" + ) + + case "${drafter}" in + peagle) + train_overrides+=( + "actor_rollout_ref.rollout.drafter.training.peagle_num_draft_layers=1" + "actor_rollout_ref.rollout.drafter.training.peagle_num_depths=2" + ) + ;; + domino) + train_overrides+=( + "actor_rollout_ref.rollout.drafter.training.domino_block_size=4" + "actor_rollout_ref.rollout.drafter.training.domino_num_anchors=8" + "actor_rollout_ref.rollout.drafter.training.domino_max_window=64" + "actor_rollout_ref.rollout.drafter.training.domino_emb_dim=64" + "actor_rollout_ref.rollout.drafter.training.domino_gru_hidden_dim=128" + "actor_rollout_ref.rollout.drafter.training.domino_lambda_base_decay_steps=10" + ) + ;; + esac +fi + if [[ -n "${SPECO_EXTRA_HYDRA_ARGS:-}" ]]; then while IFS= read -r extra_arg; do [[ -z "${extra_arg}" ]] && continue + # Stage 2 runs a different entrypoint with a different config tree, so these + # rollout-oriented overrides stay on the collect stage. overrides+=("${extra_arg}") done <<< "${SPECO_EXTRA_HYDRA_ARGS}" fi @@ -229,10 +304,26 @@ if [[ "${SPECO_DRY_RUN:-false}" == "true" ]]; then echo "drafter=${drafter}" echo "example=${example}" echo "draft_algorithm=${draft_algorithm}" + echo "separate_training=${separate_training}" echo "ASCEND_RT_VISIBLE_DEVICES=${ASCEND_RT_VISIBLE_DEVICES:-}" printf 'Hydra overrides:\n' printf ' %q\n' "${overrides[@]}" + if [[ "${separate_training}" == "true" ]]; then + printf 'Hydra train overrides:\n' + printf ' %q\n' "${train_overrides[@]}" + fi exit 0 fi -bash "${example}" "${overrides[@]}" +if [[ "${separate_training}" == "true" ]]; then + DRAFT_ALGO="${drafter}" RUN_STAGE=collect bash "${example}" "${overrides[@]}" + if [[ "${enable_training}" == "true" ]]; then + DRAFT_ALGO="${drafter}" RUN_STAGE=train DRAFT_TRAIN_GPUS_PER_NODE="${accelerator_count}" bash "${example}" "${train_overrides[@]}" + else + # Generation-only runs capture no hidden states, so the feature store the + # offline trainer would read is empty. + echo "SPECO_ENABLE_TRAINING is not true; skipping the offline train stage" + fi +else + bash "${example}" "${overrides[@]}" +fi diff --git a/examples/run_qwen3-8b_drafter_domino_peagle_separate_training.sh b/examples/run_qwen3-8b_drafter_domino_peagle_separate_training.sh index b585bca..f1ff83e 100644 --- a/examples/run_qwen3-8b_drafter_domino_peagle_separate_training.sh +++ b/examples/run_qwen3-8b_drafter_domino_peagle_separate_training.sh @@ -27,7 +27,10 @@ RUN_STAGE=${RUN_STAGE:-both} gen_tp=2 train_sp=1 ppo_gpus_per_node=8 -draft_train_gpus_per_node=8 +# The standalone launcher resolves its device count from the FIRST matching +# override, so appending one after "$@" cannot lower it. Read it from the +# environment instead, which is what a smaller box (or CI) needs. +draft_train_gpus_per_node=${DRAFT_TRAIN_GPUS_PER_NODE:-8} MODEL_PATH=/path/to/model CKPTS_DIR=/path/to/checkpoint diff --git a/tests/examples/test_ci_example_runner.py b/tests/examples/test_ci_example_runner.py index 00d5f32..9e4f934 100644 --- a/tests/examples/test_ci_example_runner.py +++ b/tests/examples/test_ci_example_runner.py @@ -13,10 +13,12 @@ # limitations under the License. from __future__ import annotations +import functools import os import shlex import subprocess import shutil +import tempfile from pathlib import Path import pytest @@ -38,6 +40,7 @@ def _workflow_source(name: str) -> str: return (WORKFLOWS / name).read_text(encoding="utf-8") +@functools.lru_cache(maxsize=None) def _require_working_bash() -> str: bash = shutil.which("bash") if bash is None: @@ -58,6 +61,7 @@ def _bash_path(path: Path, bash: str) -> str: return path.as_posix() +@functools.lru_cache(maxsize=None) def _runner_script() -> str: return "\n".join(RUNNER.read_text(encoding="utf-8").splitlines()) + "\n" @@ -289,6 +293,15 @@ def test_example_runner_covers_gpu_and_npu_backend_matrix() -> None: assert "examples/run_qwen3-8b_drafter_dflash_sglang.sh" in source assert "examples/run_qwen3-8b_drafter_dspark_vllm.sh" in source assert "examples/run_qwen3-8b_drafter_dspark_vllm_npu.sh" in source + assert "gpu/vllm/peagle" in source + assert "gpu/vllm/domino" in source + assert ( + "examples/run_qwen3-8b_drafter_domino_peagle_separate_training.sh" in source + ) + # P-EAGLE attends through torch flex_attention, which the NPU runtime does + # not provide, so the separate-training lane stays GPU-only. + assert "npu/vllm/peagle" not in source + assert "npu/vllm/domino" not in source def test_example_runner_exposes_required_hydra_overrides() -> None: @@ -319,28 +332,7 @@ def test_example_runner_exposes_required_hydra_overrides() -> None: def test_example_runner_dry_run_covers_npu_dspark() -> None: - bash = _require_working_bash() - env = { - "SPECO_DRY_RUN": "true", - "SPECO_TARGET_MODEL": "/models/target", - "SPECO_DSPARK_DRAFT_MODEL": "/models/dspark", - "SPECO_TRAIN_FILE": "/data/train.parquet", - "SPECO_TEST_FILE": "/data/test.parquet", - "SPECO_CKPT_DIR": "/tmp/speco", - "SPECO_ACCELERATOR_COUNT": "1", - } - script = "".join( - f"export {name}={shlex.quote(value)}\n" for name, value in env.items() - ) - script += _runner_script() - result = subprocess.run( - [bash, "-s", "--", "npu", "vllm", "dspark"], - env=os.environ.copy(), - input=script.encode("utf-8"), - capture_output=True, - check=True, - ) - stdout = result.stdout.decode("utf-8", errors="replace") + stdout = _dry_run("dspark", platform="npu") assert "example=examples/run_qwen3-8b_drafter_dspark_vllm_npu.sh" in stdout assert "draft_algorithm=DSPARK" in stdout @@ -350,29 +342,148 @@ def test_example_runner_dry_run_covers_npu_dspark() -> None: def test_example_runner_dry_run_omits_ulysses_overrides_for_npu_megatron() -> None: + stdout = _dry_run("megatron-eagle3", platform="npu", accelerator_count="8") + + assert "example=examples/run_qwen3-4b_actor_megatron_drafter_eagle3_vllm_npu.sh" in stdout + assert "draft_algorithm=EAGLE3" in stdout + assert "ulysses_sequence_parallel_size" not in stdout + + +@functools.lru_cache(maxsize=None) +def _dry_run( + drafter: str, + platform: str = "gpu", + backend: str = "vllm", + accelerator_count: str = "1", + extra_hydra_args: str | None = None, +) -> str: + """Dump the Hydra overrides the runner would pass, without launching a job.""" bash = _require_working_bash() env = { "SPECO_DRY_RUN": "true", "SPECO_TARGET_MODEL": "/models/target", "SPECO_EAGLE3_DRAFT_MODEL": "/models/eagle3", + "SPECO_DFLASH_DRAFT_MODEL": "/models/dflash", + "SPECO_DSPARK_DRAFT_MODEL": "/models/dspark", "SPECO_TRAIN_FILE": "/data/train.parquet", "SPECO_TEST_FILE": "/data/test.parquet", "SPECO_CKPT_DIR": "/tmp/speco", - "SPECO_ACCELERATOR_COUNT": "8", + "SPECO_ACCELERATOR_COUNT": accelerator_count, } + if extra_hydra_args is not None: + env["SPECO_EXTRA_HYDRA_ARGS"] = extra_hydra_args script = "".join( f"export {name}={shlex.quote(value)}\n" for name, value in env.items() ) script += _runner_script() - result = subprocess.run( - [bash, "-s", "--", "npu", "vllm", "megatron-eagle3"], - env=os.environ.copy(), - input=script.encode("utf-8"), - capture_output=True, - check=True, + + # Run from a file rather than piping the script into `bash -s`: the runner + # exits early in dry-run mode, and on a script this size that races the + # writer filling the stdin pipe and can deadlock. + with tempfile.TemporaryDirectory() as tmp: + entry = Path(tmp) / "dry_run.sh" + entry.write_text(script, encoding="utf-8") + result = subprocess.run( + [bash, _bash_path(entry, bash), platform, backend, drafter], + cwd=ROOT, + env=os.environ.copy(), + stdin=subprocess.DEVNULL, + capture_output=True, + check=True, + ) + return result.stdout.decode("utf-8", errors="replace") + + +def _sections(stdout: str) -> tuple[str, str]: + """Split the dry-run dump into its collect and offline-train override blocks.""" + marker = "Hydra train overrides:" + assert marker in stdout, stdout + collect, _, train = stdout.partition(marker) + _, _, collect = collect.partition("Hydra overrides:") + return collect, train + + +@pytest.mark.parametrize( + ("drafter", "collect_algorithm", "shrunk_key"), + ( + ("peagle", "EAGLE3", "training.peagle_num_depths=2"), + ("domino", "DFLASH", "training.domino_block_size=4"), + ), +) +def test_example_runner_dry_run_covers_separate_training_stages( + drafter: str, collect_algorithm: str, shrunk_key: str +) -> None: + stdout = _dry_run(drafter) + + assert ( + "example=examples/run_qwen3-8b_drafter_domino_peagle_separate_training.sh" + in stdout ) - stdout = result.stdout.decode("utf-8", errors="replace") + assert "separate_training=true" in stdout + # Stage 1 rolls out under the engine-servable algorithm whose hidden-state + # layout this drafter consumes, not under the drafter's own algorithm. + assert f"draft_algorithm={collect_algorithm}" in stdout - assert "example=examples/run_qwen3-4b_actor_megatron_drafter_eagle3_vllm_npu.sh" in stdout - assert "draft_algorithm=EAGLE3" in stdout - assert "ulysses_sequence_parallel_size" not in stdout + collect, train = _sections(stdout) + + feature_store = ( + "actor_rollout_ref.rollout.drafter.training.feature_store.path=" + f"/tmp/speco/{drafter}_features" + ) + assert ( + "actor_rollout_ref.rollout.drafter.speculative_algorithm=" + f"{collect_algorithm}" in collect + ) + assert feature_store in collect + + # Stage 2 reads that same store through the standalone offline trainer. + assert feature_store in train + assert ( + f"actor_rollout_ref.rollout.drafter.model_path=/tmp/speco/{drafter}_draft_init" + in train + ) + assert ( + "actor_rollout_ref.rollout.drafter.checkpoint_path=" + f"/tmp/speco/{drafter}_draft_ckpts" in train + ) + assert "actor_rollout_ref.rollout.drafter.training.max_steps=1" in train + assert shrunk_key in train + + # The standalone launcher takes the first matching override, so the device + # count has to reach the example through the environment instead. + assert "speco.draft_training.num_gpus_per_node" not in train + assert "DRAFT_TRAIN_GPUS_PER_NODE" in _runner_script() + + +def test_example_runner_dry_run_keeps_algorithm_knobs_out_of_the_other_lane() -> None: + peagle = _sections(_dry_run("peagle"))[1] + domino = _sections(_dry_run("domino"))[1] + + assert "domino_" not in peagle + assert "peagle_" not in domino + + +def test_example_runner_dry_run_keeps_extra_hydra_args_on_the_collect_stage() -> None: + collect, train = _sections(_dry_run("peagle", extra_hydra_args="trainer.nnodes=1")) + + assert "trainer.nnodes=1" in collect + # Stage 2 is a different entrypoint with a different config tree. + assert "trainer.nnodes=1" not in train + + +def test_example_runner_skips_the_offline_stage_without_training() -> None: + source = _runner_script() + + assert 'if [[ "${enable_training}" == "true" ]]; then' in source + assert "skipping the offline train stage" in source + assert 'rm -rf "${feature_store_dir}"' in source + + +def test_separate_training_example_takes_its_device_count_from_the_env() -> None: + example = ( + ROOT + / "examples" + / "run_qwen3-8b_drafter_domino_peagle_separate_training.sh" + ).read_text(encoding="utf-8") + + assert "draft_train_gpus_per_node=${DRAFT_TRAIN_GPUS_PER_NODE:-8}" in example