Skip to content

Fix drafter initialization and preserve draft weights - #72

Open
vx120 wants to merge 2 commits into
verl-project:mainfrom
vx120:fix_drafter_publish
Open

Fix drafter initialization and preserve draft weights#72
vx120 wants to merge 2 commits into
verl-project:mainfrom
vx120:fix_drafter_publish

Conversation

@vx120

@vx120 vx120 commented Aug 26, 2026

Copy link
Copy Markdown
Collaborator

Fix: Ensure drafter initialization before first serving request

Summary

Fixes the issue where online drafter weights would be lost when the serving worker experiences a level-2 sleep/wake-up cycle. Implements proper state tracking and initialization flow for speculative decoding drafters.

Changes

  1. Initial Drafter Initialization (_speco_ensure_initial_draft_weights)

    • Load base drafter from checkpoint before admitting first request
    • Prevent weight loss during the cold-start phase using async lock
  2. Weight Source Tracking (_speco_draft_weight_source)

    • Track whether current drafter weights come from "checkpoint" or "online" publication
    • Preserve weight source during level-2 snapshot/restore cycles
    • Validate state transitions with runtime checks
  3. Wake-up Recovery

    • Add fallback when online drafter snapshot is unavailable
    • Restore checkpoint-based weights and validate restoration
    • Handle uninitialized drafter states gracefully

Affected Files

  • verl_speco/integration/vllm_runtime.py: Main implementation

vx120 added 2 commits August 25, 2026 15:04
Signed-off-by: vx120 <893600387@qq.com>
Signed-off-by: vx120 <893600387@qq.com>
@tpx818

tpx818 commented Aug 31, 2026

Copy link
Copy Markdown
Collaborator

1.测试一下rollout开多DP的情况;
2.部分会触发vllm sleep level=1的情况下,online train时,权重恢复是否有问题

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.

2 participants