From e891544116a27fe55fb90b16670f822c425ae996 Mon Sep 17 00:00:00 2001 From: h-guo18 <67671475+h-guo18@users.noreply.github.com> Date: Thu, 20 Aug 2026 06:16:15 +0000 Subject: [PATCH] fix(speculative): read rope_theta from rope_parameters first A Transformers 5 config can carry BOTH a top-level rope_theta and a rope_parameters dict holding a different value: the real base lives in rope_parameters while the config class default (10000.0 for Qwen3) stays visible as the flat attribute. Reading the flat field first therefore picked up 10000.0 for a Qwen3-8B target whose actual base is 1000000. Worse, the training-side enforcement in HFDFlashModel.modify guarded on hasattr(base_config, "rope_theta"), which is False for a pure Transformers 5 config, so the "enforce the base model's RoPE" loop silently skipped every attribute and left the draft on the Qwen3Config default. DFlash injects the target's KV into every draft layer, so a draft built this way trains, exports and loads without complaint while its RoPE base is 100x off the target's -- there is no error at any stage, only degraded acceptance. Prefer rope_parameters in both the exporter's _get_rope_theta and the training-side enforcement, and keep the draft's own rope_parameters dict in sync with the flat field it is derived from (Qwen3Config populates the dict at construction, so setattr on the flat field alone leaves the rotary module reading a stale value). Co-Authored-By: Claude Opus 5 (1M context) Signed-off-by: h-guo18 <67671475+h-guo18@users.noreply.github.com> --- CHANGELOG.rst | 1 + .../torch/export/plugins/hf_spec_export.py | 18 +++++++--- .../torch/speculative/plugins/hf_dflash.py | 21 ++++++++++-- .../torch/export/test_hf_spec_rope_export.py | 33 ++++++++++++++++++- 4 files changed, 65 insertions(+), 8 deletions(-) diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 4580a045f39..b4b96e0ef95 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -41,6 +41,7 @@ Changelog - Update HuggingFace checkpoint export to use name-based tied-weight deduplication instead of the previous address-based approach. The address-based deduplication could incorrectly drop an untied weight that happened to share memory with a tied one, producing an incomplete checkpoint (observed as a false positive on MiniMax-M2.7). - Fix EAGLE-3 training with context parallelism (``--cp_size > 1`` in ``examples/speculative_decoding``), which failed to start on ``accelerate >= 1.13`` and then raised ``got mixed torch.Tensor and DTensor``. +- Fix the DFlash draft inheriting the wrong RoPE base from a Transformers 5 target. Such a config can carry both a ``rope_parameters`` dict holding the model's real base and a top-level ``rope_theta`` left at the config-class default, and ModelOpt read the flat field first — so a Qwen3-8B draft trained and exported with ``rope_theta`` 10000 against a target using 1000000. Retrain and re-export any DFlash-family draft (DFlash, Domino, DSpark) built against such a target. 0.46 (2026-08-17) ^^^^^^^^^^^^^^^^^ diff --git a/modelopt/torch/export/plugins/hf_spec_export.py b/modelopt/torch/export/plugins/hf_spec_export.py index 255b1d9ab04..87d2decc2b5 100644 --- a/modelopt/torch/export/plugins/hf_spec_export.py +++ b/modelopt/torch/export/plugins/hf_spec_export.py @@ -31,11 +31,15 @@ def _get_rope_theta(config, default=None): - """Get RoPE theta from either legacy or Transformers 5 config fields.""" - rope_theta = getattr(config, "rope_theta", None) - if rope_theta is not None: - return rope_theta - + """Get RoPE theta from either legacy or Transformers 5 config fields. + + ``rope_parameters`` is checked FIRST. A config can carry both fields with + different values: Transformers 5 stores the real base under + ``rope_parameters`` while the class default (10000.0 for Qwen3) may still be + visible as a top-level ``rope_theta``. Reading ``rope_theta`` first silently + exports a draft whose RoPE base is 100x off the target's, which breaks + serving because DFlash injects the target's KV into every draft layer. + """ # Transformers 5 stores this under rope_parameters (and exposes the same # data through rope_scaling for backwards compatibility). for attr in ("rope_parameters", "rope_scaling"): @@ -43,6 +47,10 @@ def _get_rope_theta(config, default=None): if isinstance(rope_config, dict) and rope_config.get("rope_theta") is not None: return rope_config["rope_theta"] + rope_theta = getattr(config, "rope_theta", None) + if rope_theta is not None: + return rope_theta + return default diff --git a/modelopt/torch/speculative/plugins/hf_dflash.py b/modelopt/torch/speculative/plugins/hf_dflash.py index e0d63bde136..c6a99f5fce7 100644 --- a/modelopt/torch/speculative/plugins/hf_dflash.py +++ b/modelopt/torch/speculative/plugins/hf_dflash.py @@ -419,10 +419,21 @@ def modify(self, config): # overwrite any user value and warn. (rope_scaling is intentionally NOT inherited: # DFlash uses standard Qwen3 RotaryEmbedding; the long-context YaRN scaling is # added only at export via dflash_export_rope_scaling.) + # A config can carry BOTH a top-level rope_theta and a rope_parameters dict + # with different values: Transformers 5 keeps the real base in + # rope_parameters while the class default (10000.0 for Qwen3) stays visible + # as rope_theta. rope_parameters wins, otherwise the draft trains against a + # RoPE base 100x off the target's. + base_rope_params = getattr(base_config, "rope_parameters", None) + if not isinstance(base_rope_params, dict): + base_rope_params = {} for attr in ("rope_theta", "rope_type", "rope_interleaved"): - if not hasattr(base_config, attr): + if attr in base_rope_params: + base_val = base_rope_params[attr] + elif hasattr(base_config, attr): + base_val = getattr(base_config, attr) + else: continue - base_val = getattr(base_config, attr) user_val = getattr(self.dflash_config, attr, None) if user_val is not None and user_val != base_val: logger.warning( @@ -434,6 +445,12 @@ def modify(self, config): base_val, ) setattr(self.dflash_config, attr, base_val) + # Qwen3Config populates rope_parameters at construction, so a later + # setattr on the flat field alone would leave the dict — which is what + # the rotary module reads — holding the stale value. + draft_rope_params = getattr(self.dflash_config, "rope_parameters", None) + if isinstance(draft_rope_params, dict) and attr in draft_rope_params: + draft_rope_params[attr] = base_val self.dflash_config.head_dim = getattr( self.dflash_config, diff --git a/tests/unit/torch/export/test_hf_spec_rope_export.py b/tests/unit/torch/export/test_hf_spec_rope_export.py index fbeb218793e..b999b7b8541 100644 --- a/tests/unit/torch/export/test_hf_spec_rope_export.py +++ b/tests/unit/torch/export/test_hf_spec_rope_export.py @@ -20,7 +20,11 @@ import torch -from modelopt.torch.export.plugins.hf_spec_export import DFlashExporter, EagleExporter +from modelopt.torch.export.plugins.hf_spec_export import ( + DFlashExporter, + EagleExporter, + _get_rope_theta, +) DEFAULT_ROPE_SCALING = { "rope_type": "yarn", @@ -152,3 +156,30 @@ def test_dflash_rope_theta_inherits_base_rope_parameters(): config = exporter._export_config() assert config["rope_theta"] == 5000000.0 + + +def test_get_rope_theta_prefers_rope_parameters_over_flat_field(): + """rope_parameters wins when a config carries both fields with different values. + + A real Transformers 5 Qwen3-8B config keeps the true base (1e6) in + rope_parameters while the Qwen3Config class default (1e4) stays visible as a + top-level rope_theta. Reading the flat field first yields a draft whose RoPE + base is 100x off the target's. + """ + config = SimpleNamespace( + rope_theta=10000.0, # class default, NOT the model's real base + rope_parameters={"rope_theta": 1000000, "rope_type": "default"}, + ) + assert _get_rope_theta(config) == 1000000 + + +def test_get_rope_theta_falls_back_to_flat_field(): + """Legacy configs that only have the flat field still resolve.""" + config = SimpleNamespace(rope_theta=500000.0, rope_parameters=None) + assert _get_rope_theta(config) == 500000.0 + + +def test_get_rope_theta_default_when_absent(): + """No rope information anywhere returns the caller's default.""" + config = SimpleNamespace(rope_theta=None, rope_parameters=None, rope_scaling=None) + assert _get_rope_theta(config, default=1234) == 1234