From b2d9c0f729725dbbd530064b54f750a6eed0448a Mon Sep 17 00:00:00 2001 From: h-guo18 <67671475+h-guo18@users.noreply.github.com> Date: Thu, 3 Sep 2026 01:28:06 -0700 Subject: [PATCH 1/2] feat(speculative): optional fp32 master weights for the DFlash draft A bf16 parameter initialised at exactly 1.0 cannot move once the learning rate falls below half the downward ULP there (2**-9 = 0.00195): every Adam step rounds straight back. The largest step Adam can take is the learning rate, so any decaying schedule eventually freezes every RMSNorm weight in the draft while the loss keeps falling and nothing reports an error. Measured on a Gemma-4-E4B DSpark run: at step 56,000, 78% of the draft's q_norm / k_norm entries were still EXACTLY 1.0, and the furthest any had moved was 34 of the 245 ULPs it needed to reach its target value. dflash_fp32_master_weights keeps the draft's parameters and Adam moments in fp32 under bf16 autocast. Matmuls still run in bf16 on tensor cores, so the cost is memory (12 B/param instead of 6), not speed. The frozen base keeps the target's dtype either way. Off by default: existing recipes train exactly as before. Signed-off-by: h-guo18 <67671475+h-guo18@users.noreply.github.com> --- modelopt/torch/speculative/config.py | 18 +++++++ .../torch/speculative/dflash/dflash_model.py | 1 + .../torch/speculative/plugins/hf_dflash.py | 8 ++++ .../speculative/plugins/test_hf_dflash.py | 47 +++++++++++++++++++ 4 files changed, 74 insertions(+) diff --git a/modelopt/torch/speculative/config.py b/modelopt/torch/speculative/config.py index 08e93b6f6a7..706e6c58635 100644 --- a/modelopt/torch/speculative/config.py +++ b/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.", diff --git a/modelopt/torch/speculative/dflash/dflash_model.py b/modelopt/torch/speculative/dflash/dflash_model.py index 3ce06afeda8..5e5f6f016e2 100644 --- a/modelopt/torch/speculative/dflash/dflash_model.py +++ b/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 diff --git a/modelopt/torch/speculative/plugins/hf_dflash.py b/modelopt/torch/speculative/plugins/hf_dflash.py index ceb76c51f91..9f439a346b5 100644 --- a/modelopt/torch/speculative/plugins/hf_dflash.py +++ b/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: diff --git a/tests/unit/torch/speculative/plugins/test_hf_dflash.py b/tests/unit/torch/speculative/plugins/test_hf_dflash.py index ab5c5a57d21..1d2f8c06a97 100644 --- a/tests/unit/torch/speculative/plugins/test_hf_dflash.py +++ b/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).""" From 2ff5f2b47729f519fae257eacf6242b48202ca60 Mon Sep 17 00:00:00 2001 From: h-guo18 <67671475+h-guo18@users.noreply.github.com> Date: Thu, 3 Sep 2026 01:28:07 -0700 Subject: [PATCH 2/2] fix(speculative): keep fp32 master weights across an HF-format resume The HF-format resume path builds the model with load_vlm_or_llm() and never calls mtsp.convert(), so HFDFlashModel.modify() -- where dflash_fp32_master_weights installs the fp32 master copy -- does not run. With dtype="auto" reading bfloat16 out of the checkpoint config, a run that trained with fp32 master weights silently continues in pure bf16 from the first resume onward. Nothing errors. The loss keeps falling; only the draft's norms stop learning, for the ULP reason described in the preceding commit. Observed on a run whose original job reported 3 bf16 / 86 fp32 draft tensors and whose resumed job reported 89 bf16 / 0 fp32 -- two checkpoints of that run are consequently not comparable with the rest. Re-apply the flag on this path so the option survives a resume, which is the only way it is usable on any job long enough to need one. Signed-off-by: h-guo18 <67671475+h-guo18@users.noreply.github.com> --- examples/speculative_decoding/main.py | 17 +++++++++++++++++ 1 file changed, 17 insertions(+) diff --git a/examples/speculative_decoding/main.py b/examples/speculative_decoding/main.py index 46848203606..babe5701639 100644 --- a/examples/speculative_decoding/main.py +++ b/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(