[RFC] Support draftmodel training in vime

[momo609]

VIME External Draft Training Design

1. Background

VIME can already load an external EAGLE/DFlash/Dspark.. Draft model through vLLM's SpeculativeConfig for speculative decoding. However, online training currently only covers the MTP layers inside the Target model. As Target parameters change throughout RL training, an external Draft model gradually becomes misaligned. This lowers Draft-token acceptance and average acceptance length, and may
eventually make the additional Draft and verification cost larger than the speculative-decoding benefit.

This design follows the implementation ideas in verl-SpeCo. It turns hiddenstates of the Target model on RL data into Draft training samples, trains an independent Draft model during every Actor update or at a fixed interval, and hot-updates the new weights into the external Draft model used by vLLM.

VIME already provides the following building blocks:

• vLLM external Draft inference configuration;
• old-policy/current-policy log-prob forward in the Megatron Actor;
• Ray training Actors and placement groups;
• online Target-weight synchronization to vLLM;
• the start_draft_weight_update Draft-weight update entry point;
• rollout pause, cache flush, weight update, and generation-resume flow;
• speculative acceptance-rate and acceptance-length fields in Samples.

The core capabilities to add are:

1. Collect Draft-training features from the Megatron Target forward pass;
2. Provide an independent external Draft model, optimizer, and checkpoint;
3. Coordinate Target and Draft weight publication;
4. Manage Target/Draft/Feature version consistency;
5. Expose complete Draft-training and speculative-decoding metrics.

2. Design Goals

2.1 Functional goals

• Train an external Draft model online during VIME RL training;
• Collect auxiliary-layer hidden states and final hidden states from the Actor's
  old-policy forward pass;
• Ensure that hidden states, the frozen Target LM Head, and the supervision
  distribution come from the same Target version;
• Periodically hot-update new Draft weights to every vLLM rollout engine;
• Save and restore the Draft model, optimizer, scheduler, and version state;
• Provide collect-only feature persistence for offline Draft training and
  diagnosis;
• Preserve the exact Target distribution during speculative decoding and avoid
  changing the PPO sample objective distribution.

2.2 Performance goals

• Collect only selected samples and token windows instead of storing full-batch,
  full-sequence hidden states;
• Transfer large feature tensors through Ray ObjectRefs or an equivalent
  point-to-point channel rather than through the Python driver;
• Make Draft-training impact on rollout time configurable and observable;
• In production, update both Target and Draft weights within one rollout pause
  window;
• Achieve higher end-to-end rollout tokens/s after online training than the
  no-speculation baseline.

2.3 Why the reference orchestration is not copied directly

verl-SpeCo's Target/rollout is based on the verl worker system, and its Draft
Trainer uses Hugging Face Transformers and FSDP. VIME uses a Megatron Actor,
packed sequences, and TP, PP, CP, and VPP layouts. Therefore:

• Draft model, loss, data-alignment rules, and feature schema can be ported;
• verl workers, dispatch, configuration, and process-group code should not be
  copied directly;
• Target hidden collection must be implemented specifically for the Megatron
  pipeline and packed sequences;
• Target-to-vLLM conversion should continue using VIME's existing
  Megatron-to-HF path;
• Draft-to-vLLM publication should use an independent HF Draft parameter
  iterator.

4. VIME Current State and Extension Points

4.1 Main training loop

The current train.py timeline is:

rollout_manager.generate
    -> actor_model.async_train
    -> actor_model.save_model
    -> actor_model.update_weights

The suitable extension is:

rollout_manager.generate
    -> actor_model.async_train and collect Draft features
    -> draft_model.collect/train
    -> save Actor/Draft
    -> jointly publish Target/Draft

4.2 Actor forward pass

MegatronTrainRayActor already
calls forward_only through compute_log_prob. This is the first-phase feature
collection entry point because:

• it runs before Actor parameters are updated;
• the current Actor weights normally match the rollout Target weights for this
  round;
• the input already contains the complete prompt/response tokens required by
  PPO training;
• there is no need to extend the vLLM generation response protocol to return
  large hidden-state payloads.

The existing custom_megatron_before_log_prob_hook can only run before the
forward pass. It cannot access per-layer outputs, batch token positions, or the
final hidden state. A structured collector must therefore be added inside
forward_only instead of relying only on the existing hook.

4.3 vLLM Draft weight update

VLLMEngine exposes:

• start_weight_update;
• start_draft_weight_update;
• finish_weight_update.

During a Draft update session, VIME's vLLM patch switches the weight-transfer
target to the Draft model inside the external DraftModelSpeculator. The
existing MTP update already follows this timeline:

pause generation
flush cache
start target update
send target weights
finish target update
start draft update
send draft weights
finish draft update
continue generation

External Draft training should reuse this lifecycle, but the Draft weights must
come from the independent Draft Worker rather than sending Actor weights again
as in the MTP path.

5. Overall Architecture

flowchart TB
    subgraph Rollout["Rollout resources"]
        VLLM["vLLM Target + External Draft"]
    end

    subgraph Actor["Actor Megatron resources"]
        Train["PPO training"]
        Collect["DraftFeatureCollector"]
        Head["Export Target LM Head"]
    end

    subgraph Draft["Independent Draft resources"]
        Queue["Versioned Feature Queue"]
        DTrain["HF Draft + FSDP/DDP"]
        DCkpt["Draft Checkpoint"]
        DPublish["Draft Weight Exporter"]
    end

    Driver["ExternalDraftCoordinator"]

    VLLM -->|"rollout tokens"| Train
    Train --> Collect
    Collect -->|"Ray ObjectRef"| Queue
    Head -->|"same-version LM Head"| Queue
    Queue --> DTrain
    DTrain --> DCkpt
    DTrain --> DPublish
    Driver --> Train
    Driver --> DTrain
    Driver --> VLLM
    DPublish -->|"draft update session"| VLLM
    Train -->|"target update session"| VLLM

Loading

Suppoort Plan

GPU

• eagle3 draftmodel support
• dflash draftmodel support
• dspark draftmodel support

Ascend

• eagle3 draftmodel support: https://github.com/momo609/vime/tree/ascend
• dflash draftmodel support
• dspark draftmodel support

TEST RESULT

we have test Qwen3-VL-4B-Instruct with AngelSlim/Qwen3-VL-4B-Instruct_eagle3 weight on Ascend: