From f2449a27475eff3ece32ae8167fd61b60f32116f Mon Sep 17 00:00:00 2001 From: h-guo18 <67671475+h-guo18@users.noreply.github.com> Date: Thu, 1 Oct 2026 12:47:06 +0000 Subject: [PATCH] fix(speculative): keep the DFlash draft's RoPE frequencies in fp32 on a fresh convert restore_draft_precision() is documented as a no-op on a freshly converted model, and examples/speculative_decoding/main.py calls it unconditionally on that basis. It was not one. modify() has already cast the draft to the base dtype and then built the rotary buffer -- after the cast on purpose, so inv_freq stays fp32 -- and restore_draft_precision() cast the draft again. Module.to() casts floating buffers too, so inv_freq came back in bf16. The rounding hits the frequencies rather than the cos/sin values, so the phase error grows with position: up to ~0.4% per frequency (measured max 3.65e-3), several radians in the highest-frequency pairs by position 4096. vLLM computes the angles from fp32 frequencies and only stores cos/sin in the model dtype, so a draft trained this way meets different RoPE phases at serve time. This affects every DFlash-family draft trained from scratch through main.py since #2483; the from_pretrained resume path was already correct. modify() now records whether it placed the draft, and restore_draft_precision() returns early when it did. On the from_pretrained path (base on meta) it still places the draft and builds the buffer. The two new tests pin both sides: the first fails on the previous code (inv_freq comes back bf16), the second fails if restore_draft_precision() always returns. Signed-off-by: h-guo18 <67671475+h-guo18@users.noreply.github.com> Co-Authored-By: Claude Opus 5.5 (1M context) --- .../torch/speculative/plugins/hf_dflash.py | 7 ++++-- .../speculative/plugins/test_hf_dflash.py | 23 +++++++++++++++++++ 2 files changed, 28 insertions(+), 2 deletions(-) diff --git a/modelopt/torch/speculative/plugins/hf_dflash.py b/modelopt/torch/speculative/plugins/hf_dflash.py index 084040a5e06..a3f4477ea0c 100644 --- a/modelopt/torch/speculative/plugins/hf_dflash.py +++ b/modelopt/torch/speculative/plugins/hf_dflash.py @@ -555,7 +555,8 @@ def modify(self, config): # lists. After the cast, because `Module.to` casts float buffers and building first # would round the RoPE frequencies. base_device = self._base_device() - if base_device.type != "meta": + self._draft_placed = base_device.type != "meta" + if self._draft_placed: self.dflash_module.to(device=base_device, dtype=self._base_model.dtype) self.dflash_module._maybe_init_rotary_emb(device=base_device) @@ -583,10 +584,12 @@ def restore_draft_precision(self): placement entirely. A no-op on a freshly converted model. """ base_device = self._base_device() - if base_device.type == "meta": + # Placing an already-placed draft again would round its fp32 rotary buffer to bf16. + if self._draft_placed or base_device.type == "meta": return self.dflash_module.to(device=base_device, dtype=self._base_model.dtype) self.dflash_module._maybe_init_rotary_emb(device=base_device) + self._draft_placed = True # Draft-module entries that legitimately come from the base model rather than the # exported draft checkpoint, so their absence (or presence) is not an error. diff --git a/tests/unit/torch/speculative/plugins/test_hf_dflash.py b/tests/unit/torch/speculative/plugins/test_hf_dflash.py index 76050c0e568..40f27e05f40 100644 --- a/tests/unit/torch/speculative/plugins/test_hf_dflash.py +++ b/tests/unit/torch/speculative/plugins/test_hf_dflash.py @@ -530,6 +530,29 @@ def test_rotary_emb_init_is_idempotent(self): assert dflash_mod.rotary_emb is first assert not any(b.is_meta for b in dflash_mod.rotary_emb.buffers()) + def test_restore_draft_precision_keeps_fp32_rope_after_convert(self): + """On a fresh convert the draft is already placed; re-casting would round RoPE to bf16.""" + model = get_tiny_llama(num_hidden_layers=4) + mtsp.convert(model, [("dflash", get_dflash_config())]) + inv_freq = model.dflash_module.rotary_emb.inv_freq.clone() + assert inv_freq.dtype == torch.float32 + + model.restore_draft_precision() + + assert model.dflash_module.rotary_emb.inv_freq.dtype == torch.float32 + assert torch.equal(model.dflash_module.rotary_emb.inv_freq, inv_freq) + + def test_restore_draft_precision_places_draft_converted_on_meta(self): + """After a meta-device convert (the from_pretrained path) it still places the draft.""" + model = get_tiny_llama(num_hidden_layers=4).to("meta") + mtsp.convert(model, [("dflash", get_dflash_config())]) + model.to_empty(device="cpu") # stands in for from_pretrained loading the weights + + model.restore_draft_precision() + + assert model.dflash_module.rotary_emb.inv_freq.dtype == torch.float32 + assert all(p.dtype == torch.bfloat16 for p in model.dflash_module.parameters()) + class TestBuildTargetLayerIds: """Test target layer selection."""