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