refactor(draft): add typed speculator contract

examples/configs/grpo_math_1B.yaml

@@ -314,6 +314,8 @@ policy:
     env_vars: null
 
   draft:
+    # Draft co-training speculator. Currently only eagle3 is supported.
+    speculator_type: eagle3
     enabled: false
     model_name: null
     loss_weight: 0.1

examples/nemo_gym/run_grpo_nemo_gym.py

@@ -167,7 +167,7 @@ def main() -> None:
             "A generation config is required for GRPO"
         )
         has_refit_draft_weights = (
-            "draft" in config.policy and config.policy["draft"]["enabled"]
+            "draft" in config.policy and config.policy["draft"].enabled
         )
         trains_mtp = (
             "megatron_cfg" in config.policy

examples/run_grpo.py

@@ -29,6 +29,7 @@
 from nemo_rl.data.utils import setup_response_data
 from nemo_rl.distributed.virtual_cluster import init_ray
 from nemo_rl.models.generation import configure_generation_config
+from nemo_rl.models.policy.draft_config import draft_refit_enabled
 from nemo_rl.utils.config import (
     load_config,
     parse_hydra_overrides,
@@ -116,7 +117,7 @@ def main() -> None:
         assert config.policy["generation"] is not None, (
             "A generation config is required for GRPO"
         )
-        has_refit_draft_weights = bool(config.policy["draft"]["enabled"])
+        has_refit_draft_weights = draft_refit_enabled(config.policy.get("draft"))
         megatron_cfg = config.policy.get("megatron_cfg") or {}
         trains_mtp = bool(megatron_cfg.get("mtp_num_layers"))
         config.policy["generation"] = configure_generation_config(

examples/run_grpo_single_controller.py

@@ -39,6 +39,7 @@
 from nemo_rl.distributed.virtual_cluster import init_ray
 from nemo_rl.environments.nemo_gym import setup_nemo_gym_config
 from nemo_rl.models.generation import configure_generation_config
+from nemo_rl.models.policy.draft_config import draft_refit_enabled
 from nemo_rl.utils.config import (
     load_config,
     parse_hydra_overrides,
@@ -121,7 +122,7 @@ def main() -> None:
     assert config.policy["generation"] is not None, (
         "A generation config is required for SC-driven async GRPO"
     )
-    has_refit_draft_weights = bool(config.policy["draft"]["enabled"])
+    has_refit_draft_weights = draft_refit_enabled(config.policy.get("draft"))
     megatron_cfg = config.policy.get("megatron_cfg") or {}
     trains_mtp = bool(megatron_cfg.get("mtp_num_layers"))
     config.policy["generation"] = configure_generation_config(

nemo_rl/models/megatron/draft/training.py

@@ -0,0 +1,81 @@
+# Copyright (c) 2026, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Protocol, cast
+
+from megatron.core.process_groups_config import ProcessGroupCollection
+from megatron.core.transformer import MegatronModule
+
+from nemo_rl.models.megatron.draft.utils import build_draft_model
+from nemo_rl.models.policy.draft_config import Eagle3DraftConfig
+
+if TYPE_CHECKING:
+    from megatron.bridge.models.model_provider import ModelProviderMixin
+
+
+class DraftSpeculator(Protocol):
+    """Model-build operation selected by a configured draft speculator."""
+
+    config: Eagle3DraftConfig
+
+    def build_model(
+        self,
+        *,
+        model_provider: ModelProviderMixin,
+        pg_collection: ProcessGroupCollection,
+        policy_model_chunk: MegatronModule,
+    ) -> MegatronModule | None:
+        """Build the draft model for the policy chunk that owns it."""
+
+
+@dataclass(frozen=True)
+class Eagle3Speculator:
+    """Current EAGLE-3 draft speculator."""
+
+    config: Eagle3DraftConfig
+
+    def build_model(
+        self,
+        *,
+        model_provider: ModelProviderMixin,
+        pg_collection: ProcessGroupCollection,
+        policy_model_chunk: MegatronModule,
+    ) -> MegatronModule | None:
+        """Build the EAGLE-3 draft model with the existing implementation."""
+        return build_draft_model(
+            model_provider=model_provider,
+            draft_config=self.config,
+            pg_collection=pg_collection,
+            policy_model_chunk=policy_model_chunk,
+        )
+
+
+_SPECULATOR_FACTORIES: dict[str, type[Eagle3Speculator]] = {
+    "eagle3": Eagle3Speculator,
+}
+
+
+def resolve_draft_speculator(
+    config: Eagle3DraftConfig | None,
+) -> DraftSpeculator | None:
+    """Resolve an enabled draft configuration to its speculator."""
+    if config is None or not config.enabled:
+        return None
+    return cast(
+        DraftSpeculator,
+        _SPECULATOR_FACTORIES[config.speculator_type](config),
+    )

nemo_rl/models/megatron/draft/utils.py

@@ -18,7 +18,7 @@
 import re
 from dataclasses import dataclass
 from pathlib import Path
-from typing import Any, Callable, Mapping
+from typing import Callable, Mapping, cast
 
 import torch
 import torch.distributed as dist
@@ -28,6 +28,8 @@
 from megatron.core.utils import unwrap_model
 from torch import Tensor
 
+from nemo_rl.models.policy.draft_config import Eagle3DraftConfig
+
 StateDict = dict[str, Tensor]
 CheckpointLoader = Callable[[Path], StateDict]
 
@@ -416,7 +418,7 @@ def _gather_tp_weight_if_needed(
     split_axis: int | None = None,
 ) -> Tensor:
     if split_axis is None:
-        tp_group = expected_shape_or_tp_group
+        tp_group = cast(dist.ProcessGroup | None, expected_shape_or_tp_group)
         if tp_group is None or not dist.is_available() or not dist.is_initialized():
             return local_weight
 
@@ -951,7 +953,10 @@ def load_hf_weights_to_eagle(
         config=unwrap_model(model).config,
     )
 
-    return model.load_state_dict(new_state, strict=False)
+    return cast(
+        tuple[list[str], list[str]],
+        model.load_state_dict(new_state, strict=False),
+    )
 
 
 def _require_state_tensor(
@@ -1227,12 +1232,12 @@ def register_draft_grad_norm_group() -> None:
 
 def build_draft_model(
     model_provider,
-    draft_config: dict[str, Any],
+    draft_config: Eagle3DraftConfig,
     pg_collection: ProcessGroupCollection,
     policy_model_chunk: MegatronModule,
 ) -> MegatronModule | None:
     """Build an Eagle draft model before parent mixed-precision/DDP wrapping."""
-    if not draft_config["enabled"]:
+    if not draft_config.enabled:
         return None
 
     from transformers import AutoConfig
@@ -1242,9 +1247,9 @@ def build_draft_model(
         get_eagle3_aux_hidden_state_layers,
     )
 
-    model_name = draft_config.get("model_name")
+    model_name = draft_config.model_name
     hf_config = AutoConfig.from_pretrained(model_name).to_dict() if model_name else {}
-    draft_num_layers = draft_config.get("num_layers")
+    draft_num_layers = draft_config.num_layers
     config = TransformerConfig(
         normalization="RMSNorm",
         activation_func=torch.nn.functional.silu,
@@ -1315,9 +1320,7 @@ def build_draft_model(
             "eagle_aux_hidden_state_layer_ids", []
         )
     else:
-        config.eagle_aux_hidden_state_layer_ids = (
-            draft_config.get("aux_layer_indices") or []
-        )
+        config.eagle_aux_hidden_state_layer_ids = draft_config.aux_layer_indices or []
     if (
         config.use_aux_hidden_state
         and len(config.eagle_aux_hidden_state_layer_ids) == 0

nemo_rl/models/megatron/setup.py

@@ -237,8 +237,8 @@ def _sync_distrib_opt(distrib_opt):
     ModelAndOptimizerState,
     RuntimeConfig,
 )
+from nemo_rl.models.megatron.draft.training import resolve_draft_speculator
 from nemo_rl.models.megatron.draft.utils import (
-    build_draft_model,
     find_draft_owner_chunk,
     get_attached_draft_model,
 )
@@ -1444,11 +1444,11 @@ def _create_draft_pre_wrap_hook(
     preload_policy_from_pretrained: bool,
 ) -> Callable[[list[MegatronModule]], list[MegatronModule]]:
     """Create the hook that attaches draft weights before mixed-precision/DDP wrapping."""
-    draft_cfg = policy_cfg["draft"]
+    draft_speculator = resolve_draft_speculator(policy_cfg.get("draft"))
 
     def draft_pre_wrap_hook(model: list[MegatronModule]) -> list[MegatronModule]:
         """Optionally preload the base policy, then attach the draft module to the owner chunk."""
-        if not draft_cfg["enabled"]:
+        if draft_speculator is None:
             return model
 
         # Base pretrained checkpoints do not contain draft weights, so load the
@@ -1483,9 +1483,8 @@ def draft_pre_wrap_hook(model: list[MegatronModule]) -> list[MegatronModule]:
             )
 
         pg_collection = get_pg_collection(model)
-        draft_model = build_draft_model(
-            megatron_cfg.model,
-            draft_config=draft_cfg,
+        draft_model = draft_speculator.build_model(
+            model_provider=megatron_cfg.model,
             pg_collection=pg_collection,
             policy_model_chunk=draft_owner,
         )
@@ -1682,7 +1681,7 @@ def setup_model_and_optimizer(
     pre_wrap_hook = []
 
     use_peft = policy_cfg["megatron_cfg"].get("peft", {}).get("enabled", False)
-    draft_enabled = "draft" in policy_cfg and policy_cfg["draft"]["enabled"]
+    draft_enabled = "draft" in policy_cfg and policy_cfg["draft"].enabled
     resume_checkpoint_exists = (
         megatron_cfg.checkpoint.load is not None
         and checkpoint_exists(megatron_cfg.checkpoint.load)

nemo_rl/models/megatron/train.py

@@ -558,7 +558,7 @@ def __call__(
                     loss_fn=loss_fn_wrapped,
                     prepare_fn=prepare_loss_input_wrapped,
                     data_dict=data_dict,
-                    loss_weight=float(self.cfg["draft"]["loss_weight"]),
+                    loss_weight=float(self.cfg["draft"].loss_weight),
                     vocab_parallel_rank=get_tensor_model_parallel_rank(),
                     vocab_parallel_group=get_tensor_model_parallel_group(),
                     context_parallel_group=get_context_parallel_group(),

nemo_rl/models/policy/__init__.py

@@ -15,6 +15,7 @@
 from typing import Any, Literal, NotRequired, TypedDict, Union
 
 from nemo_rl.models.generation.interfaces import GenerationConfig
+from nemo_rl.models.policy.draft_config import Eagle3DraftConfig
 from nemo_rl.utils.checkpoint import PretrainedCheckpointConfig
 
 
@@ -481,22 +482,6 @@ class MegatronConfig(TypedDict):
     fp8_cfg: NotRequired[Fp8Config]
 
 
-class DraftConfigDisabled(TypedDict):
-    """Configuration shape for the disabled draft-model training path."""
-
-    enabled: Literal[False]
-
-
-class DraftConfig(TypedDict):
-    """Configuration for Eagle draft-model training alongside the policy model."""
-
-    enabled: Literal[True]
-    model_name: NotRequired[str | None]
-    loss_weight: NotRequired[float]
-    num_layers: NotRequired[int | None]
-    aux_layer_indices: NotRequired[list[int] | None]
-
-
 class TokenizerConfig(TypedDict):
     name: str
     # None selects NeMo-RL's passthrough prompt/response template.
@@ -571,7 +556,7 @@ class PolicyConfig(TypedDict):
     reward_model_cfg: NotRequired[RewardModelConfig]
     dtensor_cfg: DTensorConfig | DTensorConfigDisabled
     megatron_cfg: NotRequired[MegatronConfig | MegatronConfigDisabled]
-    draft: NotRequired[DraftConfig | DraftConfigDisabled]
+    draft: NotRequired[Eagle3DraftConfig]
     pretrained_checkpoint: NotRequired[PretrainedCheckpointConfig]
     router_replay: NotRequired[RouterReplayConfig | RouterReplayConfigDisabled]
     hf_config_overrides: NotRequired[dict[str, Any]]

nemo_rl/models/policy/draft_config.py

@@ -0,0 +1,33 @@
+# Copyright (c) 2026, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from typing import Literal
+
+from pydantic import BaseModel
+
+
+class Eagle3DraftConfig(BaseModel, extra="allow"):
+    """Configuration for EAGLE-3 draft-model co-training with the policy."""
+
+    speculator_type: Literal["eagle3"] = "eagle3"
+    enabled: bool = False
+    model_name: str | None = None
+    loss_weight: float = 0.1
+    num_layers: int | None = None
+    aux_layer_indices: list[int] | None = None
+
+
+def draft_refit_enabled(config: Eagle3DraftConfig | None) -> bool:
+    """Return whether generation must accept refitted draft weights."""
+    return config is not None and config.enabled

nemo_rl/models/policy/lm_policy.py

@@ -117,7 +117,8 @@ def __init__(
 
         megatron_enable = bool(config.get("megatron_cfg", {}).get("enabled", False))
         dtensor_enable = bool(config.get("dtensor_cfg", {}).get("enabled", False))
-        draft_enabled = bool(config.get("draft", {}).get("enabled", False))
+        draft_config = config.get("draft")
+        draft_enabled = bool(draft_config is not None and draft_config.enabled)
         if megatron_enable and dtensor_enable:
             raise ValueError(
                 "Configure either Megatron (policy.megatron_cfg.enabled=true) or "

nemo_rl/models/policy/teacher_worker_group.py

@@ -161,8 +161,7 @@ def __init__(
             cfg["dtensor_cfg"]["enabled"] = False
         if "peft" in cfg["megatron_cfg"]:
             cfg["megatron_cfg"]["peft"]["enabled"] = False
-        if "draft" in cfg:
-            cfg["draft"]["enabled"] = False
+        cfg.pop("draft", None)
         # The teacher uses the plain Megatron worker, so a student-side quant_cfg
         # would be silently ignored. Drop it explicitly and warn instead.
         if cfg.get("quant_cfg") is not None:

nemo_rl/models/policy/workers/megatron_policy_worker.py

@@ -901,7 +901,7 @@ def train(
                     self._set_mtp_grad_scale_func(lambda: mtp_scale)
 
                     # Forward pass.
-                    draft_enabled = "draft" in self.cfg and self.cfg["draft"]["enabled"]
+                    draft_enabled = "draft" in self.cfg and self.cfg["draft"].enabled
                     use_router_replay = _should_use_router_replay(
                         enabled=self._router_replay_enabled,
                         data=batch,
@@ -1476,7 +1476,7 @@ def _train_microbatch_body(
         # hooks. The 1/N rescale happens once at finish.
         placeholder_n = torch.tensor(1.0, device="cuda")
 
-        draft_enabled = "draft" in self.cfg and self.cfg["draft"]["enabled"]
+        draft_enabled = "draft" in self.cfg and self.cfg["draft"].enabled
         use_router_replay = _should_use_router_replay(
             enabled=self._router_replay_enabled,
             data=data,

pyrefly.toml

@@ -214,8 +214,11 @@ project-includes = [
   "nemo_rl/models/huggingface/__init__.py",
   "nemo_rl/models/megatron/__init__.py",
   "nemo_rl/models/megatron/draft/__init__.py",
+  "nemo_rl/models/megatron/draft/training.py",
+  "nemo_rl/models/megatron/draft/utils.py",
   "nemo_rl/models/megatron/memory_saver.py",
   "nemo_rl/models/policy/__init__.py",
+  "nemo_rl/models/policy/draft_config.py",
   "nemo_rl/models/policy/interfaces.py",
   "nemo_rl/models/policy/utils.py",
   "nemo_rl/models/policy/workers/__init__.py",