diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 8782d849aa7..23223e607c3 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -41,6 +41,7 @@ Changelog **Bug Fixes** - The default PTQ recipes now leave multi-token-prediction (MTP) layers unquantized in Megatron vision-language models too (``language_model.mtp.*``), as they already did at the model root. +- Fix vanilla DSpark exports to use the Qwen3 DSpark architecture, head names and configuration expected by vLLM. Re-export existing training checkpoints to use this format. - Fix ONNX quantization placing a per-tensor Q/DQ pair on a rank-0 (scalar) constant input, which TensorRT's ONNX parser rejects (``convertAxis`` asserts the quantization axis lies in ``[0, nbDims]`` but ``nbDims == 0``). Rank-0 constant operands are now kept unquantized in the int8, fp8, and autotune export paths; the existing single-element Mul/Div/Pad stripping stays gated under ``target_dla``. - Fix PTQ of native Transformers FP8 experts to apply source weight scales during calibration and export. Re-quantize affected checkpoints from the original FP8 source. - Fix ONNX preprocessing with custom TensorRT plugin libraries on TensorRT 11. diff --git a/modelopt/torch/export/plugins/hf_spec_export.py b/modelopt/torch/export/plugins/hf_spec_export.py index 38ee13c7aaa..92ce9121b53 100644 --- a/modelopt/torch/export/plugins/hf_spec_export.py +++ b/modelopt/torch/export/plugins/hf_spec_export.py @@ -25,6 +25,7 @@ import torch.nn as nn from safetensors.torch import save_file +from ..quant_aware_conversion import revert_quant_config_names from .hf_spec_configs import kimik2_eagle_template_config, llama_eagle_template_config ALL_SPEC_MODES = ["eagle", "dflash"] @@ -344,12 +345,15 @@ class DFlashExporter(SpeculativeDecodingExporter): - config.json: Qwen3-style config with dflash_config field """ + def _map_export_name(self, name: str) -> str: + return name.split("dflash_module.", 1)[-1] + def _extract_state_dict(self, full_state_dict: dict): """Extract DFlash module weights, stripping the dflash_module prefix.""" export_sd = {} for key, value in full_state_dict.items(): if "dflash_module." in key: - export_key = key.split("dflash_module.", 1)[1] + export_key = self._map_export_name(key) # Skip rotary embedding buffers (not needed, recomputed) if "rotary_emb" in export_key: continue @@ -465,7 +469,14 @@ def export(self, export_dir: Path | str, dtype: torch.dtype | None = None): if has_quant_opt(self.model): from ..unified_export_hf import _export_transformers_checkpoint + module_names = tuple(name for name, _ in self.model.named_modules()) full_sd, hf_quant_config = _export_transformers_checkpoint(self.model, dtype) + if hf_quant_config is not None: + revert_quant_config_names( + hf_quant_config.get("quantization", {}), + self._map_export_name, + module_names=module_names, + ) else: full_sd, hf_quant_config = self.model.state_dict(), None @@ -582,14 +593,24 @@ def _export_config(self): class DSparkExporter(DFlashExporter): """Draft model exporter for DSpark (DFlash backbone + sequential Markov head). - Same z-lab-compatible format as DFlash, plus the DSpark head weights - (``markov_w1.*`` / ``markov_w2.*`` / ``gate_proj.*`` / ``joint_proj.*`` / - ``confidence_proj.*``, already captured by the inherited ``dflash_module.`` - stripping) and the extra config fields the loader needs to rebuild the head - (``projector_type``, ``markov_rank``, ``markov_head_type``, - ``use_confidence_head``, ``shift_label``). + Vanilla heads use Qwen3DSparkModel to select vLLM's loader. Gated and recurrent heads keep + the native layout because that runtime does not implement those variants. """ + def _map_export_name(self, name: str) -> str: + is_draft = "dflash_module." in name + name = super()._map_export_name(name) + if is_draft and self.model.dflash_module.markov_head_type == "vanilla": + prefixes = { + "markov_w1": "markov_head.markov_w1", + "markov_w2": "markov_head.markov_w2", + "confidence_proj": "confidence_head.proj", + } + for old_prefix, new_prefix in prefixes.items(): + if name == old_prefix or name.startswith(old_prefix + "."): + return new_prefix + name[len(old_prefix) :] + return name + def _export_config(self): """Extend the DFlash config with the DSpark head fields.""" config = super()._export_config() @@ -600,10 +621,24 @@ def _export_config(self): "projector_type": getattr(draft_config, "projector_type", "dspark"), "shift_label": getattr(draft_config, "shift_label", True), "markov_rank": draft_config.markov_rank, - "markov_head_type": getattr(draft_config, "markov_head_type", "vanilla"), + "markov_head_type": self.model.dflash_module.markov_head_type, "use_confidence_head": bool(getattr(draft_config, "use_confidence_head", False)), } ) + if config["dflash_config"]["markov_head_type"] == "vanilla": + config["architectures"] = ["Qwen3DSparkModel"] + # vLLM 0.32's DSparkSpeculator samples the anchor for next-token-aligned drafts. + # Qwen3DSparkModel's confidence head also consumes ModelOpt's Markov features. + config.update( + { + "mask_token_id": config["dflash_config"]["mask_token_id"], + "target_layer_ids": config["dflash_config"]["target_layer_ids"], + "markov_rank": config["dflash_config"]["markov_rank"], + "sample_from_anchor": True, + "enable_confidence_head": config["dflash_config"]["use_confidence_head"], + "confidence_head_with_markov": config["dflash_config"]["use_confidence_head"], + } + ) return config diff --git a/modelopt/torch/speculative/plugins/modeling_draft_checkpoint_compat.py b/modelopt/torch/speculative/plugins/modeling_draft_checkpoint_compat.py index d2243433661..ec1920c5064 100644 --- a/modelopt/torch/speculative/plugins/modeling_draft_checkpoint_compat.py +++ b/modelopt/torch/speculative/plugins/modeling_draft_checkpoint_compat.py @@ -13,13 +13,11 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Per-released-checkpoint key remaps for loading published drafters. +"""Checkpoint key remaps for loading exported and published drafters. -ModelOpt's draft modules define one canonical parameter layout, which is also the export -format. Some published drafters ship the same tensors under different key names; loading -those requires rewriting keys, which is a property of a specific release rather than of the -architecture. Keeping those remaps here — instead of in the module that defines the -architecture — means a new checkpoint quirk adds an entry to this file only. +ModelOpt's draft modules define one canonical parameter layout. Serving exports and +published drafters can use different key names; the load hooks here translate those +names back to the module layout. Each entry is a ``load_state_dict`` pre-hook body: it mutates ``state_dict`` in place to move released-layout keys onto the canonical names. @@ -29,17 +27,20 @@ logger = logging.getLogger(__name__) -# Head submodules that some published DSpark checkpoints nest under a `markov_head.` parent. -# ModelOpt keeps them flat (``markov_w1`` / ``markov_w2`` / ...), matching the upstream -# DeepSpec layout; ``nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-BF16-DSpark`` nests them. +# DSpark modules keep Markov heads flat; serving checkpoints can nest them. _DSPARK_NESTED_HEAD_PREFIX = "markov_head." def remap_dspark_nested_head_keys(state_dict, prefix, *args, **kwargs): - """Accept DSpark head weights nested under ``markov_head.`` as well as flat.""" + """Accept DSpark Markov and confidence heads in serving or native layouts.""" nested = prefix + _DSPARK_NESTED_HEAD_PREFIX - for key in [k for k in state_dict if k.startswith(nested)]: - flat = prefix + key[len(nested) :] + confidence = prefix + "confidence_head.proj." + for key in [k for k in state_dict if k.startswith((nested, confidence))]: + flat = ( + prefix + key[len(nested) :] + if key.startswith(nested) + else prefix + "confidence_proj." + key[len(confidence) :] + ) nested_value = state_dict.pop(key) # A flat key already present wins. The Markov tables are ~14% of the draft's # parameters and loading the wrong copy is invisible at runtime, so say which was diff --git a/modelopt/torch/speculative/plugins/modeling_dspark.py b/modelopt/torch/speculative/plugins/modeling_dspark.py index 2a87b5a0d9e..e6464380d17 100644 --- a/modelopt/torch/speculative/plugins/modeling_dspark.py +++ b/modelopt/torch/speculative/plugins/modeling_dspark.py @@ -62,8 +62,8 @@ The head uses its OWN embedding table (``markov_w1``), not the base model's, so the bias computation is fully self-contained on this module. Submodule names (``markov_w1`` / ``markov_w2`` / ``gate_proj`` / ``joint_proj`` / -``confidence_proj``) match the upstream DeepSpec checkpoint layout so exported -checkpoints stay portable. +``confidence_proj``) match the upstream DeepSpec layout; ``DSparkExporter`` adapts +the serving checkpoint format. """ import logging diff --git a/tests/unit/torch/speculative/plugins/test_hf_dspark.py b/tests/unit/torch/speculative/plugins/test_hf_dspark.py index 3686788ad90..3898f6b9211 100644 --- a/tests/unit/torch/speculative/plugins/test_hf_dspark.py +++ b/tests/unit/torch/speculative/plugins/test_hf_dspark.py @@ -18,14 +18,13 @@ DSpark reuses the DFlash mode/pipeline and adds a lightweight sequential (Markov) head plus an optional confidence head. These tests cover conversion routing for the three head variants, the three-term training forward (CE + TVD + confidence -BCE), and the export format (head weights + config) against the z-lab-compatible -layout (``markov_w1.*`` / ``markov_w2.*`` / ``gate_proj.*`` / ``joint_proj.*`` / -``confidence_proj.*``). +BCE), vLLM vanilla exports, and native gated/recurrent exports. """ import json import shutil from copy import deepcopy +from fnmatch import fnmatchcase import pytest import torch @@ -34,6 +33,7 @@ from transformers import AutoModelForCausalLM import modelopt.torch.opt as mto +import modelopt.torch.quantization as mtq import modelopt.torch.speculative as mtsp from modelopt.torch.speculative.config import DFLASH_DEFAULT_CFG from modelopt.torch.speculative.plugins.hf_dflash import HFDFlashModel @@ -267,7 +267,7 @@ def spy(*args, **kwargs): class TestDSparkExporter: - """Test the DSpark checkpoint export format (z-lab-compatible layout).""" + """Test vanilla serving exports and the native gated/recurrent formats.""" def _export(self, tmp_path, head_type="vanilla", use_confidence_head=False): model = get_tiny_llama(num_hidden_layers=4) @@ -286,38 +286,115 @@ def _export(self, tmp_path, head_type="vanilla", use_confidence_head=False): model.get_exporter().export(export_dir) return export_dir - @pytest.mark.parametrize("head_type", HEAD_TYPES) - def test_export_weight_keys_match_reference(self, tmp_path, head_type): - """Exported weights carry the head tensors under reference names, no prefix.""" - sd = load_file(str(self._export(tmp_path, head_type=head_type) / "model.safetensors")) + @pytest.mark.parametrize("head_type", [*HEAD_TYPES, "VANILLA"]) + @pytest.mark.parametrize("use_confidence_head", [False, True]) + def test_export_weight_keys_match_reference(self, tmp_path, head_type, use_confidence_head): + """Only vanilla heads use the Qwen3 DSpark serving tensor names.""" + export_dir = self._export( + tmp_path, head_type=head_type, use_confidence_head=use_confidence_head + ) + sd = load_file(str(export_dir / "model.safetensors")) for key in sd: assert "dflash_module." not in key assert "rotary_emb" not in key - assert "markov_w1.weight" in sd - assert "markov_w2.weight" in sd + nested = head_type.lower() == "vanilla" + for name in ("markov_w1.weight", "markov_w2.weight"): + assert (name in sd) == (not nested) + assert (f"markov_head.{name}" in sd) == nested assert ("gate_proj.weight" in sd) == (head_type == "gated") assert ("joint_proj.weight" in sd) == (head_type == "rnn") - - def test_export_includes_confidence_weights(self, tmp_path): - """The confidence head weights are exported when enabled.""" - sd = load_file(str(self._export(tmp_path, use_confidence_head=True) / "model.safetensors")) - assert "confidence_proj.weight" in sd - - def test_export_config_has_dspark_fields(self, tmp_path): - """config.json carries the dflash_config DSpark head fields.""" - export_dir = self._export(tmp_path, head_type="gated") + for suffix in ("weight", "bias"): + assert (f"confidence_head.proj.{suffix}" in sd) == (nested and use_confidence_head) + assert (f"confidence_proj.{suffix}" in sd) == (not nested and use_confidence_head) + + @pytest.mark.parametrize("head_type", [*HEAD_TYPES, "VANILLA"]) + @pytest.mark.parametrize("use_confidence_head", [False, True]) + def test_export_config_has_dspark_fields(self, tmp_path, head_type, use_confidence_head): + """Vanilla exports select the runtime and preserve the trained head settings.""" + export_dir = self._export( + tmp_path, head_type=head_type, use_confidence_head=use_confidence_head + ) with open(export_dir / "config.json") as f: cfg = json.load(f) - assert cfg["architectures"] == ["DFlashDraftModel"] + vanilla = head_type.lower() == "vanilla" + assert cfg["architectures"] == ["Qwen3DSparkModel" if vanilla else "DFlashDraftModel"] dc = cfg["dflash_config"] assert dc["projector_type"] == "dspark" assert dc["markov_rank"] == MARKOV_RANK - assert dc["markov_head_type"] == "gated" - assert dc["use_confidence_head"] is False + assert dc["markov_head_type"] == head_type.lower() + assert dc["use_confidence_head"] is use_confidence_head assert dc["shift_label"] is True assert "mask_token_id" in dc assert "target_layer_ids" in dc + if vanilla: + assert cfg["model_type"] == "qwen3" + assert cfg["markov_rank"] == MARKOV_RANK + assert cfg["sample_from_anchor"] is True + assert cfg["enable_confidence_head"] is use_confidence_head + assert cfg["confidence_head_with_markov"] is use_confidence_head + assert cfg["mask_token_id"] == dc["mask_token_id"] + assert cfg["target_layer_ids"] == dc["target_layer_ids"] + else: + assert "enable_confidence_head" not in cfg + + @pytest.mark.parametrize("head_type", HEAD_TYPES) + @pytest.mark.parametrize("quantize_draft", [False, True]) + def test_quantized_export_metadata_matches_weights(self, tmp_path, head_type, quantize_draft): + """Packed exports keep mixed-precision keys and wildcard exclusions aligned.""" + model = get_tiny_llama(num_hidden_layers=4) + mtsp.convert( + model, + [ + ( + "dflash", + _get_dspark_config(head_type=head_type, use_confidence_head=True, block_size=2), + ) + ], + ) + quantizers = [{"quantizer_name": "*", "enable": False}] + layers = [("model.layers.0.mlp.down_proj", None)] + if quantize_draft: + layers += [("dflash_module.fc", 0), ("dflash_module.markov_w2", None)] + for name, axis in layers: + quantizers.append( + { + "quantizer_name": name + ".weight_quantizer", + "cfg": {"num_bits": (4, 3), "axis": axis}, + "enable": True, + } + ) + mtq.quantize( + model, + {"quant_cfg": quantizers, "algorithm": "max"}, + forward_loop=lambda m: m(torch.ones((1, 2), dtype=torch.long)), + ) + model.get_exporter().export(tmp_path, dtype=torch.bfloat16) + cfg = json.loads((tmp_path / "config.json").read_text()) + quant_cfg = json.loads((tmp_path / "hf_quant_config.json").read_text()) + assert cfg["quantization_config"] == quant_cfg + quantization = quant_cfg["quantization"] + assert quantization["quant_algo"] == ("MIXED_PRECISION" if quantize_draft else "FP8") + prefix = "markov_head." if head_type == "vanilla" else "" + confidence = "confidence_head.proj" if head_type == "vanilla" else "confidence_proj" + for name, excluded in [ + (prefix + "markov_w1", True), + (prefix + "markov_w2", not quantize_draft), + (confidence, True), + ("layers.0.self_attn.q_proj", True), + ("fc", not quantize_draft), + ("model.layers.0.mlp.down_proj", False), + ("model.layers.1.self_attn.q_proj", True), + ]: + assert any(fnmatchcase(name, p) for p in quantization["exclude_modules"]) == excluded + if quantize_draft: + assert quantization["quantized_layers"] == { + "model.layers.0.mlp.down_proj": {"quant_algo": "FP8"}, + "fc": {"quant_algo": "FP8_PER_CHANNEL_PER_TOKEN"}, + prefix + "markov_w2": {"quant_algo": "FP8"}, + } + sd = load_file(str(tmp_path / "model.safetensors")) + assert sd[prefix + "markov_w2.weight"].dtype == torch.float8_e4m3fn class TestDraftAttentionPattern: @@ -533,15 +610,17 @@ def test_export_omits_sink_when_disabled(self, tmp_path): class TestMarkovHeadKeyRemap: """Head weights load from either the flat or the nested `markov_head.` layout. - ModelOpt (and the upstream DeepSpec reference) keeps the head tensors flat on the - module and exports them that way. Some released drafters — notably + ModelOpt (and the upstream DeepSpec reference) keeps head tensors flat on the + module. Serving exports and some released drafters — notably nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-BF16-DSpark — nest them under a `markov_head.` parent, so loading accepts both. """ - def _module(self): + def _module(self, use_confidence_head=False): model = get_tiny_llama(num_hidden_layers=4) - mtsp.convert(model, [("dflash", _get_dspark_config())]) + mtsp.convert( + model, [("dflash", _get_dspark_config(use_confidence_head=use_confidence_head))] + ) return model.dflash_module def test_nested_keys_load(self): @@ -560,8 +639,8 @@ def test_nested_keys_load(self): ) def test_flat_keys_still_load(self): - """The exported (flat) layout keeps working unchanged.""" - module = self._module() + """The legacy flat layout keeps working unchanged.""" + module = self._module(use_confidence_head=True) flat = dict(module.state_dict()) flat["markov_w1.weight"] = torch.full_like(flat["markov_w1.weight"], 0.25) res = module.load_state_dict(flat, strict=True) @@ -570,16 +649,21 @@ def test_flat_keys_still_load(self): module.markov_w1.weight, torch.full_like(module.markov_w1.weight, 0.25) ) - def test_flat_key_wins_over_nested(self): + @pytest.mark.parametrize( + ("flat_key", "nested_key"), + [ + ("markov_w1.weight", "markov_head.markov_w1.weight"), + ("confidence_proj.weight", "confidence_head.proj.weight"), + ], + ) + def test_flat_key_wins_over_nested(self, flat_key, nested_key): """An explicit flat key is never clobbered by a nested duplicate.""" - module = self._module() + module = self._module(use_confidence_head=True) sd = dict(module.state_dict()) - sd["markov_w1.weight"] = torch.full_like(sd["markov_w1.weight"], 1.0) - sd["markov_head.markov_w1.weight"] = torch.full_like(sd["markov_w1.weight"], 9.0) + sd[flat_key] = torch.full_like(sd[flat_key], 1.0) + sd[nested_key] = torch.full_like(sd[flat_key], 9.0) module.load_state_dict(sd, strict=True) - assert torch.allclose( - module.markov_w1.weight, torch.full_like(module.markov_w1.weight, 1.0) - ) + assert torch.allclose(module.state_dict()[flat_key], torch.full_like(sd[flat_key], 1.0)) class TestInitCheckpoint: @@ -605,10 +689,41 @@ def _export(self, tmp_path, **overrides): model.get_exporter().export(export_dir) return export_dir - def test_warm_start_loads_exported_weights(self, tmp_path): - """Every draft parameter comes from the checkpoint, not a fresh init.""" - export_dir = self._export(tmp_path) - model = self._make_model(init_checkpoint=export_dir) + @pytest.mark.parametrize("use_confidence_head", [False, True]) + @pytest.mark.parametrize("legacy_layout", [False, True]) + def test_warm_start_loads_exported_weights(self, tmp_path, use_confidence_head, legacy_layout): + """New serving exports and legacy exports warm-start every draft parameter.""" + architecture = { + **_get_dspark_config()["dflash_architecture_config"], + "use_confidence_head": use_confidence_head, + } + export_dir = self._export(tmp_path, dflash_architecture_config=architecture) + if legacy_layout: + path = export_dir / "model.safetensors" + sd = load_file(str(path)) + sd = { + key.removeprefix("markov_head.").replace( + "confidence_head.proj.", "confidence_proj." + ): value + for key, value in sd.items() + } + save_file(sd, str(path)) + config_path = export_dir / "config.json" + config = json.loads(config_path.read_text()) + config["architectures"] = ["DFlashDraftModel"] + for field in ( + "mask_token_id", + "target_layer_ids", + "markov_rank", + "sample_from_anchor", + "enable_confidence_head", + "confidence_head_with_markov", + ): + config.pop(field) + config_path.write_text(json.dumps(config)) + model = self._make_model( + init_checkpoint=export_dir, dflash_architecture_config=architecture + ) for name, p in model.dflash_module.named_parameters(): assert torch.allclose(p, torch.full_like(p, 0.125)), f"{name} was not warm-started" @@ -695,20 +810,26 @@ def test_restore_does_not_replay_warm_start(self, tmp_path): # than returning False. What matters here is the values, not the load dtype. assert torch.allclose(p, ref.to(p.dtype)), f"{name} differs after restore" - def test_nested_head_shape_mismatch_reported(self, tmp_path): - """A wrong-shaped tensor is caught even under the nested `markov_head.` layout. + @pytest.mark.parametrize( + "head_key", ["markov_head.markov_w1.weight", "confidence_head.proj.weight"] + ) + def test_nested_head_shape_mismatch_reported(self, tmp_path, head_key): + """A wrong-shaped tensor is caught even under the nested serving layout. The shape check has to resolve the module's load hooks first; otherwise a remapped key skips it and fails later with a far less obvious error. """ - export_dir = self._export(tmp_path) + architecture = { + **_get_dspark_config()["dflash_architecture_config"], + "use_confidence_head": True, + } + export_dir = self._export(tmp_path, dflash_architecture_config=architecture) path = export_dir / "model.safetensors" sd = load_file(str(path)) - sd["markov_head.markov_w1.weight"] = torch.zeros(3, MARKOV_RANK) - del sd["markov_w1.weight"] + sd[head_key] = torch.zeros(3, sd[head_key].shape[1]) save_file(sd, str(path)) with pytest.raises(ValueError, match="shape mismatch"): - self._make_model(init_checkpoint=export_dir) + self._make_model(init_checkpoint=export_dir, dflash_architecture_config=architecture) class TestExplicitTargetLayerIds: