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."""