diff --git a/tensorrt_llm/_torch/models/modeling_deepseekv4.py b/tensorrt_llm/_torch/models/modeling_deepseekv4.py index 3d65b1ee7595..5724c1052d8c 100644 --- a/tensorrt_llm/_torch/models/modeling_deepseekv4.py +++ b/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 diff --git a/tests/unittest/_torch/modeling/test_modeling_deepseekv4.py b/tests/unittest/_torch/modeling/test_modeling_deepseekv4.py index 106ca0e5c59d..b46ed3e477b9 100644 --- a/tests/unittest/_torch/modeling/test_modeling_deepseekv4.py +++ b/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"