test(ci): wire P-EAGLE and Domino into the example test runner - #74
Merged
tpx818 merged 1 commit intoAug 31, 2026
Merged
Conversation
Neither algorithm is servable by the rollout engines, so both are trained through the two-stage separate-training workflow: stage 1 rolls out under the engine algorithm whose hidden-state layout the drafter consumes and writes a feature store, stage 2 trains the drafter offline from that store. The example script for that workflow existed but no CI lane invoked it, so `ci/run_example_test.sh` accepted only the single-invocation drafters. Add `peagle` and `domino` cases that run the example twice, once per stage, and cover both lanes with dry-run assertions. The runner overrides only what CI must control (where the stages meet on disk, and the sizes that keep a smoke run cheap); the two-stage shape and the per-algorithm hyperparameters stay owned by the example. Two supporting fixes: * The standalone launcher resolves its device count from the first matching override, not the last, so an appended `num_gpus_per_node` could never lower the example's value and stage 2 would launch with the example's device count regardless of the runner. The example now reads that count from `DRAFT_TRAIN_GPUS_PER_NODE`. * The dry-run tests fed the runner to `bash -s` over stdin, which races the script's early exit once the script grows past a few kilobytes. They now run it from a file. Signed-off-by: khazic <khazzz1c@gmail.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What
ci/run_example_test.shcoverseagle3,megatron-eagle3,dflashanddspark. P-EAGLE and Domino have unit and integration coverage plus a standalone GPU smoke script, but no example lane, so the workflow they actually ship with is never exercised end to end.Neither is an engine-level speculative algorithm, so neither can be trained inside the rollout loop. Both go through the two-stage separate-training workflow that
examples/run_qwen3-8b_drafter_domino_peagle_separate_training.shalready implements:verl_speco.main,mode=collect_onlyverl_speco.draft_train_launcher,mode=offlinePEAGLE/DOMINOThis PR adds
gpu/vllm/peagleandgpu/vllm/dominoto the runner. They invoke that example once per stage, withRUN_STAGE=collectthenRUN_STAGE=train, since the two stages need different Hydra overrides and the example applies"$@"to both.The runner overrides only what CI has to control: where the two stages meet on disk, and the sizes that keep a smoke run cheap. The two-stage shape, the collect-to-train algorithm pairing and the per-algorithm hyperparameters stay owned by the example, so the lane tests the example as written instead of a private copy of it.
Stage 1 collects with a drafter that is already in the CI model cache (the EAGLE3 and DFlash drafters the existing lanes use), and stage 2 cold-starts its drafter from the target config, so no new model needs to be provisioned.
Two supporting fixes
The stage-2 device count could not be overridden.
draft_train_launcher._find_overridereturns the first matching override rather than the last, the opposite of Hydra's last-wins rule that every other override in this lane relies on. The example passesspeco.draft_training.num_gpus_per_nodebefore"$@", so an appended value was silently ignored and stage 2 would have launched with the example's device count on a runner that does not have that many. The example now reads the count fromDRAFT_TRAIN_GPUS_PER_NODE, which the runner sets fromSPECO_ACCELERATOR_COUNT. (Changing the launcher to last-wins would be the deeper fix, but it changes behavior for existing callers, so it is left out of this PR.)The dry-run tests could deadlock. They fed the runner into
bash -sover stdin. In dry-run mode the runner exits early, which races the writer still filling the stdin pipe, and past roughly 14 KB of script the test hangs instead of failing. They now write the script to a temporary file and execute that. This reproduced on the runner as it stands in this PR, so it is fixed here rather than left for whoever next grows the file.Scope
torch.nn.attention.flex_attention, which the NPU runtime does not provide.gpu_vllm_unit_tests.ymlonce you are happy with the wiring, and I am glad to do it here if you prefer.Test
pytest tests/examples tests/integration tests/compat tests/config: 170 passed, 49 skipped.bash -n ci/run_example_test.sh examples/*.sh.