Skip to content

feat(speco): add standalone TransferQueue (TQ) based draft feature transport and training - #73

Open
xxyyrr598 wants to merge 64 commits into
verl-project:mainfrom
755651978:main
Open

feat(speco): add standalone TransferQueue (TQ) based draft feature transport and training#73
xxyyrr598 wants to merge 64 commits into
verl-project:mainfrom
755651978:main

Conversation

@xxyyrr598

Copy link
Copy Markdown
Collaborator

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:

  • prevented DraftFeatureSample.metadata from round-tripping cleanly;
  • made the codec aware of DSpark/hidden-state layout internals;
  • required touching the TQ protocol whenever a new algorithm was added;
  • caused inconsistent stats between producer and consumer ready-tag filtering.

This PR reworks the transport layer around a single DraftFeatureSample
protocol and adds a standalone producer/consumer workflow.

Key Changes

Transport & protocol

  • verl_speco/transport/drafter_sample_protocol.py: DraftFeatureSample codec
    with ready tags, EOS records and sample keys (schema v2).
  • verl_speco/integration/transferqueue_bridge.py: TransferQueue bridge for
    restoring offloaded tensors on the consumer side.
  • verl_speco/integration/mooncake_hidden_states_connector.py and
    verl_speco/trainer/mooncake_transfer.py: Mooncake-based hidden-state
    transfer.

Producer

  • verl_speco/producer/input_reader.py: producer-side input reading.
  • verl_speco/producer/vllm_feature_client.py: vLLM feature client pool with
    retries/backoff.
  • verl_speco/standalone_tq_producer.py and verl_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_head export).
  • Reworked trainer/draft_training_loop.py and trainer/feature_store.py to
    support 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

  • Design docs under docs/ (TQ protocol refactor, standalone producer /
    consumer / training, transferqueue integration plan).
  • Unit tests for protocol, producer, TQ sample source, replay, transferqueue
    bridge and vLLM feature client.

Architecture

vLLM (hidden states)
        │  feature payloads
        ▼
 Producer (input_reader / vllm_feature_client)
        │  enqueue (TransferQueue / Mooncake)
        ▼
 Consumer (tq_sample_source / target_feature_replay)
        │
        ▼
 Draft trainer (draft_training_loop / feature_store)

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

  • New unit tests: 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-files passes (ruff, ruff-format, mypy, docs/license/
    naming sanity checks, compileall).

Notes

  • This is an independent standalone training path, separate from the
    VeOmni-based co-training workflow.
  • Requires the transferqueue package (see wheels/) and a vLLM instance
    serving hidden states with the target layer configuration.

Checklist

  • pre-commit passes
  • unit tests added
  • CI verified

755651978 and others added 30 commits July 9, 2026 13:59
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>
xxyyrr598 and others added 7 commits August 25, 2026 09:49
…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
# 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
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

4 participants