fp32 master weights for the DFlash draft, and keep them across a resume

examples/speculative_decoding/main.py

@@ -222,6 +222,23 @@ def train():
         tokenizer = transformers.AutoTokenizer.from_pretrained(
             checkpoint, trust_remote_code=recipe.model.trust_remote_code
         )
+        # This branch does NOT call mtsp.convert(), so HFDFlashModel.modify() -- which is
+        # where dflash_fp32_master_weights installs the fp32 master copy -- never runs on
+        # resume. Combined with dtype="auto" reading bfloat16 out of the checkpoint config,
+        # a run that trained with fp32 master weights silently continues in pure bf16.
+        #
+        # That is not a small difference: bf16 spacing at a weight of 1.0 is 0.0039, an
+        # update must clear 0.00195 to round anywhere, and the largest step Adam can take
+        # is the learning rate -- so once the decayed lr falls below ~2e-3 every RMSNorm
+        # weight is frozen. The run keeps going and the loss keeps falling; only the norms
+        # stop learning, which is why this was invisible for two full runs.
+        if isinstance(recipe, ModelOptDFlashRecipe) and recipe.dflash.dflash_fp32_master_weights:
+            model.dflash_fp32_master_weights = True
+            model.dflash_module.float()
+            print_rank_0(
+                "Resume: reapplied fp32 master weights to the draft "
+                f"({sum(1 for _ in model.dflash_module.parameters())} params)"
+            )
     else:
         if checkpoint:
             print_rank_0(

modelopt/torch/speculative/config.py

@@ -144,6 +144,24 @@ class DFlashConfig(ModeloptBaseConfig):
         default={}, description="Config for the DFlash draft module architecture."
     )
 
+    dflash_fp32_master_weights: bool = ModeloptField(
+        default=False,
+        description=(
+            "Keep the draft's parameters in fp32 while training under bf16 autocast, i.e. "
+            "classic mixed precision with fp32 master weights. Matmuls still run in bf16 on "
+            "tensor cores, so the cost is memory (12 bytes/param for weight+Adam moments "
+            "instead of 6), not speed.\n\n"
+            "This exists because a pure-bf16 parameter initialised at exactly 1.0 cannot "
+            "move: the downward ULP there is 2**-8 = 0.0039, so an update must exceed "
+            "0.00195 to round to a new value, and the maximum Adam step is the learning "
+            "rate. Under the Gemma-4 recipe's linear decay from 2e-3 that threshold is "
+            "crossed around step 2,600, after which every RMSNorm weight still sitting at "
+            "1.0 is frozen for the rest of the run. Measured on a 56k-step Gemma-4 draft: "
+            "78% of the q_norm/k_norm entries were still exactly 1.0, and the furthest any "
+            "had travelled was 34 ULPs of the 245 needed to reach head_dim**-0.5."
+        ),
+    )
+
     dflash_use_torch_compile: bool = ModeloptField(
         default=True,
         description="Whether to use torch.compile on DFlash forward/loss methods.",

modelopt/torch/speculative/dflash/dflash_model.py

@@ -48,6 +48,7 @@ def modify(self, config):
         self.dflash_num_anchors = config.dflash_num_anchors
         self.dflash_report_acc = config.dflash_report_acc
         self.dflash_use_torch_compile = config.dflash_use_torch_compile
+        self.dflash_fp32_master_weights = config.dflash_fp32_master_weights
         self.dflash_swa_window_size = config.dflash_swa_window_size
         self.dflash_draft_attention = config.dflash_draft_attention
         self.dflash_attention_sink = config.dflash_attention_sink

modelopt/torch/speculative/plugins/hf_dflash.py

@@ -509,6 +509,14 @@ def modify(self, config):
             base_device = next(self._base_model.layers[-1].parameters()).device
         if base_device.type != "meta":
             self.dflash_module.to(self._base_model.dtype).to(base_device)
+            if self.dflash_fp32_master_weights:
+                # Deliberately AFTER the base-dtype match: that call also moves the module
+                # to the right device, and matching first keeps this a pure dtype change.
+                # Autocast (HF Trainer under bf16=True) casts activations and weights to
+                # bf16 at each op, so compute is unchanged; only the master copy and the
+                # optimizer moments stay fp32. See DFlashConfig.dflash_fp32_master_weights
+                # for why bf16 masters silently freeze the RMSNorm weights.
+                self.dflash_module.float()
 
         # Delete base model layers for offline training (save memory)
         if self.dflash_offline:

tests/unit/torch/speculative/plugins/test_hf_dflash.py

@@ -576,6 +576,53 @@ def test_no_sliding_window_without_config(self):
         assert attn.sliding_window is None
 
 
+class TestDFlashFp32MasterWeights:
+    """The draft can hold fp32 parameters while the frozen base stays bf16.
+
+    Motivation is numerical, not cosmetic: a bf16 parameter initialised at exactly 1.0
+    cannot move once the learning rate drops below half the downward ULP there
+    (2**-9 = 0.00195), because every Adam step rounds back. On a 56k-step Gemma-4 run
+    that left 78% of the draft's RMSNorm weights still exactly 1.0.
+    """
+
+    @staticmethod
+    def _bf16_base():
+        model = get_tiny_llama(num_hidden_layers=4)
+        return model.to(torch.bfloat16)
+
+    @staticmethod
+    def _draft_dtypes(model):
+        return {p.dtype for n, p in model.named_parameters() if "dflash_module" in n}
+
+    def test_draft_is_fp32_while_base_stays_bf16(self):
+        """The flag lifts ONLY the draft; the frozen base keeps the target's dtype."""
+        model = self._bf16_base()
+        config = get_dflash_config()
+        config["dflash_fp32_master_weights"] = True
+        mtsp.convert(model, [("dflash", config)])
+        assert self._draft_dtypes(model) == {torch.float32}
+        base = {p.dtype for n, p in model.named_parameters() if "dflash_module" not in n}
+        assert base == {torch.bfloat16}
+
+    def test_default_keeps_the_draft_in_the_base_dtype(self):
+        """Off by default: existing recipes must keep training exactly as before."""
+        model = self._bf16_base()
+        mtsp.convert(model, [("dflash", get_dflash_config())])
+        assert self._draft_dtypes(model) == {torch.bfloat16}
+
+    def test_flag_survives_save_restore(self):
+        """The flag lives in DFlashConfig, so a restored checkpoint must still carry it.
+
+        A checkpoint written with the flag and restored by code that dropped it fails
+        pydantic validation with extra_forbidden -- which is how every eval job for this
+        run failed until the eval container was moved to matching code.
+        """
+        model = self._bf16_base()
+        config = get_dflash_config()
+        config["dflash_fp32_master_weights"] = True
+        mtsp.convert(model, [("dflash", config)])
+        assert model.dflash_fp32_master_weights is True
+
 class TestDFlashSwaMask:
     """Test all-layer non-causal sliding-window attention mask (MiMo-style)."""