feat(speco): add standalone TransferQueue (TQ) based draft feature transport and training - #73
Open
xxyyrr598 wants to merge 64 commits into
Open
feat(speco): add standalone TransferQueue (TQ) based draft feature transport and training#73xxyyrr598 wants to merge 64 commits into
xxyyrr598 wants to merge 64 commits into
Conversation
Co-authored-by: Codex <noreply@openai.com> Signed-off-by: 755651978 <755651978@qq.com>
Co-authored-by: Codex <noreply@openai.com> Signed-off-by: 755651978 <755651978@qq.com>
Co-authored-by: Codex <noreply@openai.com> Signed-off-by: 755651978 <755651978@qq.com>
…dalone TQ workflow
…r target model configurations such as Qwen3.
Signed-off-by: vx120 <893600387@qq.com>
Signed-off-by: vx120 <893600387@qq.com>
- Move hidden_state_vllm script from examples/ to tools/, drop redundant tools scripts - Adjust example script naming checks - Fix mypy errors (None guards, annotations, missing return) and doc Last updated info
…rate training script
# Conflicts: # examples/run_qwen3-8b_drafter_separate_training.sh # verl_speco/config/speco_base.yaml # verl_speco/trainer/draft_dataset.py # verl_speco/trainer/draft_training_loop.py # verl_speco/workers/speco_worker.py
xxyyrr598
temporarily deployed
to
speco-npu-ci
September 1, 2026 08:35 — with
GitHub Actions
Inactive
xxyyrr598
temporarily deployed
to
speco-npu-ci
September 1, 2026 08:35 — with
GitHub Actions
Inactive
xxyyrr598
temporarily deployed
to
speco-npu-ci
September 1, 2026 08:35 — with
GitHub Actions
Inactive
xxyyrr598
temporarily deployed
to
speco-npu-ci
September 1, 2026 08:35 — with
GitHub Actions
Inactive
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.
Summary
This PR introduces a standalone TQ (TransferQueue) training workflow for
draft-model feature transport and training. It decouples the drafter feature
producer and consumer through a queue-based protocol so that hidden states and
target logprobs are transferred off the driver, reducing driver-side memory
pressure and enabling independent, decoupled draft training.
Background / Motivation
Previously, draft-feature transport was tightly coupled to the training loop and
carried heavy
SampleMetadata-based re-descriptions of the hidden-state layout,which:
DraftFeatureSample.metadatafrom round-tripping cleanly;This PR reworks the transport layer around a single
DraftFeatureSampleprotocol and adds a standalone producer/consumer workflow.
Key Changes
Transport & protocol
verl_speco/transport/drafter_sample_protocol.py:DraftFeatureSamplecodecwith ready tags, EOS records and sample keys (schema v2).
verl_speco/integration/transferqueue_bridge.py: TransferQueue bridge forrestoring offloaded tensors on the consumer side.
verl_speco/integration/mooncake_hidden_states_connector.pyandverl_speco/trainer/mooncake_transfer.py: Mooncake-based hidden-statetransfer.
Producer
verl_speco/producer/input_reader.py: producer-side input reading.verl_speco/producer/vllm_feature_client.py: vLLM feature client pool withretries/backoff.
verl_speco/standalone_tq_producer.pyandverl_speco/tq_owner.py:standalone producer entry points and TQ ownership.
Consumer / training
verl_speco/standalone_tq_training_launcher.py: standalone training launcher.verl_speco/trainer/tq_sample_source.py,trainer/tq_feature_store.py:queue-backed sample source and feature store.
verl_speco/trainer/target_feature_replay.py,trainer/target_feature_pipeline.py: target feature replay/pipeline.verl_speco/trainer/standalone_checkpoint.py: standalone checkpoint export(including sharded
lm_headexport).trainer/draft_training_loop.pyandtrainer/feature_store.pytosupport the queue-based flow.
Scripts / examples
tools/run_qwen3-8b_drafter_hidden_state_vllm.sh,tools/wait_for_vllm_endpoints.py.examples/run_qwen3-8b_drafter_dspark_separate_training.sh,examples/run_qwen3-8b_drafter_eagle3_separate_training.sh.Docs & tests
docs/(TQ protocol refactor, standalone producer /consumer / training, transferqueue integration plan).
bridge and vLLM feature client.
Architecture
Hidden states and large tensors are offloaded from the driver and restored on
the consumer via TQ keys, so only small ready tags cross the driver.
Testing
tests/unit/test_*for the protocol, producer input reader,target feature replay, TQ consumer/producer, transferqueue bridge, vLLM
feature client, and standalone launcher.
pre-commit run --all-filespasses (ruff, ruff-format, mypy, docs/license/naming sanity checks, compileall).
Notes
VeOmni-based co-training workflow.
transferqueuepackage (seewheels/) and a vLLM instanceserving hidden states with the target layer configuration.
Checklist