`nemo_rl.models.automodel.draft.setup`#

Building the draft model and its optimizer for Automodel co-training.

Adapts a dspark/dflash/eagle3 checkpoint’s config onto the vendored draft models, validates it against the training options, and FSDP2-shards it. Also builds the composite policy+draft optimizer used for checkpoint I/O.

Module Contents#

Classes#

PolicyWithDraft

Composite module pairing the policy and draft for optimizer state I/O.

Functions#

`load_draft_hf_config`	Load a draft config for `algo`, adapting speculators-format checkpoints.
`_adapt_speculators_dspark_config`	Map a speculators-format dspark/dflash config onto the vendored layout.
`default_eagle3_aux_layer_ids_vllm`	Default vLLM EAGLE3 aux layers for a target, in vLLM capture indexing.
`_adapt_speculators_eagle3_config`	Map a speculators-format eagle3 config onto automodel’s LlamaEagle3DraftModel.
`_adapt_native_flat_eagle3_config`	Adapt a native (non-speculators) flat eagle3 config, e.g. SpecForge.
`validate_dspark_draft_config`	Validate a dspark/dflash checkpoint’s config against the training options.
`build_dspark_draft_model`	Build, validate, and FSDP2-shard a dspark/dflash draft from a checkpoint.
`build_eagle3_draft_model`	Build, validate, and FSDP2-shard an EAGLE3 draft from a checkpoint.
`_load_eagle3_weights`	Load an eagle3 checkpoint’s weights into automodel’s LlamaEagle3DraftModel.
`build_policy_and_draft_optimizer`	Build the optimizer with named [policy, draft] param groups; the draft group uses its own lr.

Data#

`DSPARK_REQUIRED_CONFIG_FIELDS`
`NATIVE_FLAT_EAGLE3_ARCHS`

API#

nemo_rl.models.automodel.draft.setup.DSPARK_REQUIRED_CONFIG_FIELDS#: (‘block_size’, ‘target_layer_ids’, ‘mask_token_id’)

nemo_rl.models.automodel.draft.setup.NATIVE_FLAT_EAGLE3_ARCHS#: (‘LlamaForCausalLMEagle3’,)

nemo_rl.models.automodel.draft.setup.load_draft_hf_config( model_name: str, algo: str = 'dspark', target_num_hidden_layers: Optional[int] = None, ) → Any#

Load a draft config for algo, adapting speculators-format checkpoints.

Two checkpoint families exist:

Flat qwen3-style configs (e.g. deepseek-ai/dspark_qwen3_8b_block7 and its dflash sibling): loadable via AutoConfig, drafter fields at the top level. Their target_layer_ids already use the trainer’s output-of-layer convention and their blocks use the next-token supervision layout (they were produced by the trainer this code is vendored from), so no convention shifts apply.
Speculators-format configs (e.g. RedHatAI/-speculator.): no top-level model_type (AutoConfig fails), the draft transformer fields nest under transformer_layer_config, and layer/block conventions follow vLLM/speculators indexing. These are adapted to the flat layout the vendored models expect.

nemo_rl.models.automodel.draft.setup._adapt_speculators_dspark_config( config_dict: dict[str, Any], algo: str = 'dspark', ) → Any#: Map a speculators-format dspark/dflash config onto the vendored layout.

nemo_rl.models.automodel.draft.setup.default_eagle3_aux_layer_ids_vllm( target_num_hidden_layers: int, ) → list[int]#

Default vLLM EAGLE3 aux layers for a target, in vLLM capture indexing.

Mirrors SupportsEagle3.get_eagle3_default_aux_hidden_state_layers: (2, N // 2, N - 3). The trainer capture uses these minus 1 (output-of-layer convention).

nemo_rl.models.automodel.draft.setup._adapt_speculators_eagle3_config( config_dict: dict[str, Any], model_name: str, target_num_hidden_layers: Optional[int], ) → Any#: Map a speculators-format eagle3 config onto automodel’s LlamaEagle3DraftModel.

nemo_rl.models.automodel.draft.setup._adapt_native_flat_eagle3_config( config_dict: dict[str, Any], model_name: str, target_num_hidden_layers: Optional[int], ) → Any#

Adapt a native (non-speculators) flat eagle3 config, e.g. SpecForge.

Unlike the speculators family, these configs load through AutoConfig directly (draft_vocab_size etc. sit at the top level, model_type is a real HF value); what they don’t record is the aux capture layer ids or the first-layer residual convention, so both are filled in the same way as the speculators adapter does when a checkpoint pins neither.

class nemo_rl.models.automodel.draft.setup.PolicyWithDraft(policy: torch.nn.Module, draft: torch.nn.Module)#

Bases: torch.nn.Module

Composite module pairing the policy and draft for optimizer state I/O.

The Automodel checkpointer pairs one model with one optimizer for optimizer state save/load; the single optimizer here owns param groups from both the policy and the draft, so this composite is the module handed to the checkpointer’s optimizer paths. Training, refit, and inference keep referencing the policy module directly.

Initialization

forward(*args: Any, **kwargs: Any) → Any#

nemo_rl.models.automodel.draft.setup.validate_dspark_draft_config( draft_hf_config: Any, dspark_options: dict[str, Any], algo: str = 'dspark', ) → None#

Validate a dspark/dflash checkpoint’s config against the training options.

DFlash is the markov-free, confidence-free subset of DSpark: the same vendored model runs both, but a dflash run must not silently pick up markov/confidence behavior (or vice versa expect heads that don’t exist).

nemo_rl.models.automodel.draft.setup.build_dspark_draft_model( model_name: str, dspark_options: dict[str, Any], torch_dtype: torch.dtype, mesh: Any, algo: str = 'dspark', ) → nemo_rl.models.automodel.draft.draft_qwen3.Qwen3DSparkModel#: Build, validate, and FSDP2-shard a dspark/dflash draft from a checkpoint.

nemo_rl.models.automodel.draft.setup.build_eagle3_draft_model( model_name: str, eagle3_options: nemo_rl.models.policy.Eagle3DraftConfig, torch_dtype: torch.dtype, mesh: Any, target_num_hidden_layers: int, policy_model: Optional[torch.nn.Module] = None, ) → nemo_rl.models.automodel.draft.eagle3_llama.Eagle3DraftModel#

Build, validate, and FSDP2-shard an EAGLE3 draft from a checkpoint.

policy_model is only required for native-flat (SpecForge-style) checkpoints, which ship no embed_tokens of their own – see _load_eagle3_weights.

nemo_rl.models.automodel.draft.setup._load_eagle3_weights( draft_model: nemo_rl.models.automodel.draft.eagle3_llama.Eagle3DraftModel, model_name: str, policy_model: Optional[torch.nn.Module], needs_embed_tokens_from_policy: bool, ) → None#

Load an eagle3 checkpoint’s weights into automodel’s LlamaEagle3DraftModel.

Checkpoints publish a FLAT layout (no model. prefix – the vendored single-file trainer this replaces used the same flat layout) with the lone decoder layer under either layers.0. (speculators/RedHatAI) or midlayer. (SGLang SpecForge/lmsys); both remap onto automodel’s model.layers.0. (lm_head/d2t/t2d stay top-level in both layouts, matching automodel’s).

Some checkpoints (SpecForge) ship no embed_tokens at all – the drafter is meant to share the target model’s embedding table. True weight tying (aliasing the same nn.Parameter) isn’t possible here: the draft and policy are fully_shard-ed on different device meshes, so needs_embed_tokens_from_policy copies the policy’s embedding VALUES once at init instead (matching how the megatron eagle path backfills a checkpoint’s missing lm_head from the policy) and lets the copy train independently thereafter, governed by train_embed_and_head like the rest of the embedding/head.

nemo_rl.models.automodel.draft.setup.build_policy_and_draft_optimizer( optimizer_cls: Any, optimizer_kwargs: dict[str, Any], model: torch.nn.Module, draft_model: torch.nn.Module, draft_config: Any, ) → torch.optim.Optimizer#: Build the optimizer with named [policy, draft] param groups; the draft group uses its own lr.

nemo_rl.models.automodel.draft.setup#