[None][fix] Load DeepSeek V4 mixed-precision NVFP4 checkpoints

tensorrt_llm/_torch/models/modeling_deepseekv4.py

@@ -1,3 +1,6 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
 # --------------------------------------------------
 # Portions of this code were derived from DeepSeek‑V3:
 #   https://github.com/deepseek-ai/DeepSeek-V3
@@ -296,6 +299,43 @@ def _resolve_enable_fused_hc(config: PretrainedConfig) -> bool:
     return bool(getattr(config, "enable_fused_hc", True))
 
 
+def _normalize_deepseek_v4_nvfp4_mixed_precision_config(
+    model_config: ModelConfig[PretrainedConfig],
+) -> ModelConfig[PretrainedConfig]:
+    """Resolve FP8 base layers in DeepSeek-V4 NVFP4 checkpoints."""
+    quant_config = model_config.quant_config
+    hf_quant_config = getattr(model_config.pretrained_config, "quantization_config", None)
+    layer_quant_configs = model_config.quant_config_dict or {}
+    has_nvfp4_experts = any(
+        name.endswith(".mlp.experts") and config.quant_algo == QuantAlgo.NVFP4
+        for name, config in layer_quant_configs.items()
+    )
+    if (
+        quant_config.quant_algo != QuantAlgo.MIXED_PRECISION
+        or not has_nvfp4_experts
+        or not isinstance(hf_quant_config, dict)
+        or hf_quant_config.get("quant_method") != "fp8"
+        or tuple(hf_quant_config.get("weight_block_size", ())) != (128, 128)
+    ):
+        return model_config
+
+    default_exclude = ["*kv_b_proj*", "*k_b_proj*", "*eh_proj*"]
+    hf_exclude_modules = hf_quant_config.get("modules_to_not_convert") or []
+    exclude_modules = list(dict.fromkeys(list(hf_exclude_modules) + default_exclude))
+    fp8_quant_config = quant_config.model_copy(
+        deep=True,
+        update={
+            "quant_algo": QuantAlgo.FP8_BLOCK_SCALES,
+            "group_size": 128,
+            "exclude_modules": exclude_modules,
+        },
+    )
+    fp8_quant_config.__dict__.pop("quant_mode", None)
+    fp8_quant_config.__dict__.pop("layer_quant_mode", None)
+    model_config.quant_config = fp8_quant_config
+    return model_config
+
+
 def _copy_deepseek_v4_fused_a_weight_scale(
     module: Linear, fused_a: torch.Tensor, fused_a_scale: torch.Tensor
 ) -> None:
@@ -2483,6 +2523,7 @@ def get_model_defaults(cls, llm_args: "TorchLlmArgs") -> dict:
         }
 
     def __init__(self, model_config: ModelConfig[PretrainedConfig]):
+        model_config = _normalize_deepseek_v4_nvfp4_mixed_precision_config(model_config)
         self.mapping_with_cp = None
         # Note: Currently the usage of mapping is all over the place making its usage brittle
         # in this file. As a temporary WAR, we hold on to an original copy of mapping when CP

tests/unittest/_torch/modeling/test_modeling_deepseekv4.py

@@ -1,3 +1,6 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
 import ast
 import inspect
 import json
@@ -33,6 +36,7 @@
     DeepseekV4MTP,
     _copy_deepseek_v4_fused_a_weight_scale,
     _deepseek_v4_pos_embd_params,
+    _normalize_deepseek_v4_nvfp4_mixed_precision_config,
     _remap_deepseek_v4_checkpoint_keys,
     _resolve_enable_fused_hc,
 )
@@ -442,6 +446,48 @@ def test_deepseek_v4_moe_auto_backend_on_blackwell(monkeypatch):
     assert ModelConfig.resolve_moe_backend("AUTO", "DeepseekV4ForCausalLM") == "TRTLLM"
 
 
+def test_deepseek_v4_nvfp4_mixed_precision_config():
+    config = DeepseekV4Config()
+    config.quantization_config = {
+        "quant_method": "fp8",
+        "weight_block_size": [128, 128],
+        "modules_to_not_convert": ["lm_head"],
+    }
+    mixed_quant_config = QuantConfig(
+        quant_algo=QuantAlgo.MIXED_PRECISION,
+        group_size=16,
+        exclude_modules=["*.attn.*", "*.ffn.shared_experts.*", "head", "mtp.*"],
+    )
+    mixed_quant_config.mamba_ssm_cache_dtype = torch.bfloat16
+    assert not mixed_quant_config.layer_quant_mode.has_fp8_block_scales()
+    experts_quant_config = QuantConfig(quant_algo=QuantAlgo.NVFP4, group_size=16)
+    model_config = ModelConfig(
+        pretrained_config=config,
+        quant_config=mixed_quant_config,
+        quant_config_dict={"model.layers.0.mlp.experts": experts_quant_config},
+    )
+    model_config._frozen = True
+
+    normalized_config = _normalize_deepseek_v4_nvfp4_mixed_precision_config(model_config)
+
+    assert normalized_config is model_config
+    assert mixed_quant_config.quant_algo == QuantAlgo.MIXED_PRECISION
+    assert normalized_config.quant_config.quant_algo == QuantAlgo.FP8_BLOCK_SCALES
+    assert normalized_config.quant_config.layer_quant_mode.has_fp8_block_scales()
+    assert normalized_config.quant_config.group_size == 128
+    assert normalized_config.quant_config.mamba_ssm_cache_dtype == torch.bfloat16
+    assert normalized_config.quant_config.exclude_modules == [
+        "lm_head",
+        "*kv_b_proj*",
+        "*k_b_proj*",
+        "*eh_proj*",
+    ]
+    assert (
+        normalized_config.quant_config_dict["model.layers.0.mlp.experts"].quant_algo
+        == QuantAlgo.NVFP4
+    )
+
+
 def test_deepseek_v4_routed_moe_quant_config_from_mxfp4_header(tmp_path, monkeypatch):
     monkeypatch.setattr("tensorrt_llm._torch.model_config.get_sm_version", lambda: 100)
     tensor_name = "layers.0.ffn.experts.0.w1.weight"