coderabbit

Signed-off-by: h-guo18 <67671475+h-guo18@users.noreply.github.com>

Author: h-guo18

examples/speculative_decoding/main.py

@@ -268,10 +268,10 @@ def train():
             mtsp.convert(model, [("medusa", config)])
         elif training_args.mode == "eagle3":
             # Validate and rewrite eagle config fields
-            EagleConfig.model_validate(
+            eagle_cfg = EagleConfig.model_validate(
                 eagle_cfg,
                 context={"training_args": training_args, "data_args": data_args},
-            )
+            ).model_dump()
             mtsp.convert(model, [("eagle", eagle_cfg)])
 
             # Load draft vocab cache if the draft model uses a compressed vocabulary

modelopt/torch/export/plugins/hf_spec_export.py

@@ -187,14 +187,15 @@ def _get_config_from_draft_or_base(key: str, model: nn.Module):
                     new_value = str(new_value).replace("torch.", "")
                 template_config[key] = new_value
 
-        # Inject export rope scaling override (validated at config time to require
-        # training rope_type == "default").
+        # Inject export rope scaling override when training rope_type is "default".
+        rope_cfg = self.model.eagle_config.rope_scaling or {}
+        training_rope_type = rope_cfg.get("rope_type") or rope_cfg.get("type")
         eagle_export_rope_scaling = getattr(self.model, "eagle_export_rope_scaling", None)
-        if eagle_export_rope_scaling:
+        if eagle_export_rope_scaling and training_rope_type == "default":
             template_config["rope_scaling"] = eagle_export_rope_scaling
 
         # In transformers 5.x, rope_theta is under rope_scaling, not the main config.
-        rope_cfg = self.model.eagle_config.rope_scaling
+        # Always source from the training rope config (rope_theta is not in export overrides).
         if template_config.get("rope_theta") is None and rope_cfg:
             template_config["rope_theta"] = rope_cfg.get("rope_theta")
 

modelopt/torch/speculative/config.py

@@ -148,8 +148,8 @@ def _derive_eagle_offline(cls, data: Any, info: ValidationInfo) -> Any:
     def _check_rope_scaling_consistency(self) -> "EagleConfig":
         if not self.eagle_export_rope_scaling:
             return self
-        rope_cfg = self.eagle_architecture_config.get("rope_scaling", {})
-        rope_type = rope_cfg.get("rope_type")
+        rope_cfg = self.eagle_architecture_config.get("rope_scaling", {}) or {}
+        rope_type = rope_cfg.get("rope_type") or rope_cfg.get("type")
         if rope_type is not None and rope_type != "default":
             raise ValueError(
                 f"eagle_export_rope_scaling is set but eagle_architecture_config has "

tests/unit/torch/speculative/test_eagle_config.py

@@ -36,6 +36,16 @@ def test_rope_consistency_error_non_default_rope_type():
         EagleConfig.model_validate(cfg)
 
 
+def test_rope_consistency_error_non_default_rope_type_alt_key():
+    """Error when rope_scaling uses 'type' key instead of 'rope_type' (kimik2-style)."""
+    cfg = {
+        "eagle_export_rope_scaling": {"rope_type": "yarn", "factor": 32.0},
+        "eagle_architecture_config": {"rope_scaling": {"type": "yarn"}},
+    }
+    with pytest.raises(ValidationError, match="rope_type='yarn'"):
+        EagleConfig.model_validate(cfg)
+
+
 def test_rope_consistency_ok_default_rope_type():
     """No error when training rope_type is 'default'."""
     cfg = {