nemo_rl.models.generation.vllm.draft_utils

vLLM-side draft co-training helpers; vllm imports stay function-local.

Module Contents

Functions

draft_module_sharing_disable_required	Whether this worker’s engine needs the draft module-sharing disable.
disable_draft_module_sharing	Keep the drafter’s embed_tokens/lm_head separate from the target’s.
_is_full_draft_stream	Whether a draft.* stream is the Automodel FULL drafter.
_draft_owns_speculator	Whether this rank should own the co-trained speculator.
_expected_draft_keys	Expected incoming draft.* keys derived from the drafter’s layout.
_validate_draft_refit_info	Hard-error when the trainer’s draft.* manifest mismatches the drafter.
_assert_drafter_owns_modules	Refuse to load draft embed/lm_head through target-shared modules.

Data

COTRAINED_SPECULATIVE_METHODS
DRAFT_DISABLE_MODULE_SHARING_ENV
_DRAFT_SKIPPED_KEY_SUBSTRINGS

API

nemo_rl.models.generation.vllm.draft_utils.COTRAINED_SPECULATIVE_METHODS

(‘dspark’, ‘dflash’, ‘eagle3’)

nemo_rl.models.generation.vllm.draft_utils.DRAFT_DISABLE_MODULE_SHARING_ENV

‘NRL_DRAFT_DISABLE_MODULE_SHARING’

nemo_rl.models.generation.vllm.draft_utils.draft_module_sharing_disable_required(

config: dict[str, Any],

) → bool

Whether this worker’s engine needs the draft module-sharing disable.

Full-stream draft co-training refits the drafter’s trained embed_tokens/lm_head. Under load_format=”dummy” the pinned vLLM would alias those drafter modules to the target model’s (no checkpoint load ever marks them as owned), and the draft refit would then overwrite the policy’s serving weights through the alias. The method name alone doesn’t imply a full stream – Megatron block-drafter paths can use the same dspark/dflash method names with a headless exporter that relies on module sharing, just like the megatron eagle3 path does – so gate on _draft_full_refit (true only for Automodel co-training, which always streams the drafter’s entire state_dict) for all three methods.

nemo_rl.models.generation.vllm.draft_utils.disable_draft_module_sharing() → None

Keep the drafter’s embed_tokens/lm_head separate from the target’s.

The pinned vLLM’s load_dspark_model / load_dflash_model / load_eagle_model alias the drafter’s embed_tokens and lm_head modules to the target model’s whenever the drafter’s load_weights has not marked them as owned – which is always the case under load_format="dummy", where no checkpoint weights are read at startup. With draft co-training the refit stream carries trained draft.embed_tokens / draft.lm_head weights, and loading them through the aliased modules silently replaces the POLICY model’s serving embed and lm_head with the draft’s. Forcing _should_share to False makes the drafter keep the modules it built in __init__; refit fills them before the first generation, and CUDA graphs capture drafter-owned storage.

_should_share is defined in the eagle utils module and imported by value into the dspark/dflash utils modules, so each module’s global must be rebound individually.

A second, independent embed_tokens-sharing decision lives in SpecDecodeBaseProposer._maybe_share_embeddings (used by the eagle3 proposer): it shares whenever self.model.has_own_embed_tokens is falsy, a flag only ever set (to True) by process_eagle_weight when a weight literally named “embed_tokens” is loaded – which never happens under load_format="dummy". This is not reachable through _should_share at all, so it needs its own no-op patch; skipping it entirely leaves each drafter holding the embed_tokens it built in __init__, matching the _should_share override above.

Must run before engine creation (drafter load and CUDA-graph capture).

nemo_rl.models.generation.vllm.draft_utils._DRAFT_SKIPPED_KEY_SUBSTRINGS

None

nemo_rl.models.generation.vllm.draft_utils._is_full_draft_stream(draft_keys: set[str] | list[str]) → bool

Whether a draft.* stream is the Automodel FULL drafter.

The Automodel co-training path streams the drafter’s entire state_dict (always including per-layer keys like layers.0.*, and normally embed_tokens too), regardless of method (dspark/dflash/eagle3). The megatron eagle3 exporter instead streams a PARTIAL drafter – it intentionally omits embed_tokens and aliases its single collapsed layer as midlayer.* instead of layers.0.*, relying on drafter module sharing at serve time.

Neither signal alone is reliable, so the stream is treated as FULL unless it carries neither: a one-key batch slice carrying only lm_head.weight – the genuine megatron partial path’s minimal payload – is read as partial, while a FULL stream missing just embed_tokens is still caught by its remaining layers.* keys.

nemo_rl.models.generation.vllm.draft_utils._draft_owns_speculator(

model_runner: Any,

method: Optional[str],

) → bool

Whether this rank should own the co-trained speculator.

vLLM keeps the drafter on the last pipeline stage, so earlier stages legitimately have no speculator and must skip draft payloads; every rank owns it in single-stage layouts.

nemo_rl.models.generation.vllm.draft_utils._expected_draft_keys(draft_model: Any, method: str) → set[str]

Expected incoming draft.* keys derived from the drafter’s layout.

Inverts the drafter load_weights name handling shared by Qwen3DSparkForCausalLM / DFlashQwen3ForCausalLM / Eagle3Qwen3ForCausalLM: trainer names are prefixed with model. (except lm_head.*, and d2t which maps to draft_id_to_target_id), and fused parameters load from their stacked components (qkv_proj <- q/k/v_proj, gate_up_proj <- gate/up_proj). Parameters the loader never feeds from the stream (see _DRAFT_SKIPPED_KEY_SUBSTRINGS) are excluded.

nemo_rl.models.generation.vllm.draft_utils._validate_draft_refit_info(

model_runner: Any,

draft_model: Any,

method: Optional[str],

state_dict_info: dict[str, Any],

) → None

Hard-error when the trainer’s draft.* manifest mismatches the drafter.

Owning ranks require exactly the drafter’s loadable key set (plus keys the loader intentionally skips); non-owning ranks ignore draft payloads entirely and are never required to have a drafter.

nemo_rl.models.generation.vllm.draft_utils._assert_drafter_owns_modules(

model_runner: Any,

draft_model: Any,

draft_weights: list[tuple[str, Any]],

) → None

Refuse to load draft embed/lm_head through target-shared modules.

vLLM’s drafter loaders may alias the drafter’s embed_tokens and lm_head to the target model’s modules as a weight-sharing optimization; loading refit draft weights through such an alias would overwrite the policy’s serving weights with the draft’s (see disable_draft_module_sharing). Guard every refit so any path that reintroduces the sharing fails loudly instead of silently corrupting generation.