From d1bbd177f3daacb5858e1fc0c6bee93650023c66 Mon Sep 17 00:00:00 2001 From: Fanrong Li Date: Wed, 15 Jul 2026 13:40:04 +0000 Subject: [PATCH 1/5] [None][fix] Load DeepSeek V4 mixed-precision NVFP4 checkpoints Signed-off-by: Fanrong Li --- .../_torch/models/modeling_deepseekv4.py | 9 +-- .../modeling/test_modeling_deepseekv4.py | 61 +++++++++++++++++++ 2 files changed, 66 insertions(+), 4 deletions(-) diff --git a/tensorrt_llm/_torch/models/modeling_deepseekv4.py b/tensorrt_llm/_torch/models/modeling_deepseekv4.py index 2fb611ecb999..a88c5853b9e8 100644 --- a/tensorrt_llm/_torch/models/modeling_deepseekv4.py +++ b/tensorrt_llm/_torch/models/modeling_deepseekv4.py @@ -1756,11 +1756,12 @@ def __init__( post_mult_value=2.0, ) - # FIXME: incompatible with mixed quantization mode quant_config = self._get_decoder_layer_quant_config(model_config, layer_idx) - self.is_nvfp4 = quant_config.layer_quant_mode.has_nvfp4() - assert quant_config.quant_algo is not QuantAlgo.MIXED_PRECISION, ( - "MIXED_PRECISION is ambiguous" + # MIXED_PRECISION is resolved per module. Routed experts use their + # layer-specific config, while layer-level NVFP4 fusions stay disabled. + self.is_nvfp4 = ( + quant_config.quant_algo != QuantAlgo.MIXED_PRECISION + and quant_config.layer_quant_mode.has_nvfp4() ) self.allreduce = None diff --git a/tests/unittest/_torch/modeling/test_modeling_deepseekv4.py b/tests/unittest/_torch/modeling/test_modeling_deepseekv4.py index 106ca0e5c59d..c253c5e71be0 100644 --- a/tests/unittest/_torch/modeling/test_modeling_deepseekv4.py +++ b/tests/unittest/_torch/modeling/test_modeling_deepseekv4.py @@ -5,6 +5,7 @@ import textwrap import weakref from copy import deepcopy +from types import SimpleNamespace import pytest import torch @@ -30,6 +31,7 @@ DeepseekV4DecoderLayer, DeepseekV4ForCausalLM, DeepseekV4Gate, + DeepseekV4MoE, DeepseekV4MTP, _copy_deepseek_v4_fused_a_weight_scale, _deepseek_v4_pos_embd_params, @@ -442,6 +444,65 @@ def test_deepseek_v4_moe_auto_backend_on_blackwell(monkeypatch): assert ModelConfig.resolve_moe_backend("AUTO", "DeepseekV4ForCausalLM") == "TRTLLM" +def test_deepseek_v4_decoder_accepts_mixed_precision_experts(monkeypatch): + config = DeepseekV4Config( + hidden_size=16, + moe_intermediate_size=8, + n_routed_experts=4, + n_shared_experts=1, + num_experts_per_tok=2, + hc_mult=2, + hc_sinkhorn_iters=1, + ) + config.torch_dtype = torch.bfloat16 + quant_config = QuantConfig( + quant_algo=QuantAlgo.MIXED_PRECISION, + group_size=16, + exclude_modules=["*.attn.*", "*.ffn.shared_experts.*", "head", "mtp.*"], + ) + experts_quant_config = QuantConfig(quant_algo=QuantAlgo.NVFP4, group_size=16) + model_config = SimpleNamespace( + pretrained_config=config, + mapping=Mapping(world_size=1, rank=0, tp_size=1), + quant_config=quant_config, + quant_config_dict={"model.layers.0.mlp.experts": experts_quant_config}, + allreduce_strategy=None, + ) + captured = {} + + def fake_moe(**kwargs): + captured.update(kwargs) + return torch.nn.Identity() + + monkeypatch.setattr( + "tensorrt_llm._torch.models.modeling_deepseekv4.mHC", + lambda *args, **kwargs: torch.nn.Identity(), + ) + monkeypatch.setattr( + "tensorrt_llm._torch.models.modeling_deepseekv4.DeepseekV4Attention", + lambda *args, **kwargs: torch.nn.Identity(), + ) + monkeypatch.setattr("tensorrt_llm._torch.models.modeling_deepseekv4.DeepseekV4MoE", fake_moe) + monkeypatch.setattr( + "tensorrt_llm._torch.models.modeling_deepseekv4.RMSNorm", + lambda *args, **kwargs: torch.nn.Identity(), + ) + monkeypatch.setattr( + "tensorrt_llm._torch.models.modeling_deepseekv4.can_access_peer", + lambda _mapping: False, + ) + + layer = DeepseekV4DecoderLayer( + model_config, + layer_idx=0, + aux_stream_dict={AuxStreamType.Attention: None, AuxStreamType.MoeShared: None}, + ) + + assert layer.is_nvfp4 is False + assert captured["override_quant_config"] is quant_config + assert DeepseekV4MoE._get_experts_quant_config(model_config, 0) is experts_quant_config + + 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" From 4be177fe1a8e26679dc1fff2844edd324b39e7b4 Mon Sep 17 00:00:00 2001 From: Fanrong Li Date: Tue, 21 Jul 2026 13:21:32 +0000 Subject: [PATCH 2/5] [None][fix] Resolve DeepSeek V4 mixed-precision base layers Treat the non-expert weights in DeepSeek V4 NVFP4 checkpoints as FP8 block-scale weights while retaining the per-layer NVFP4 expert configuration. Signed-off-by: Fanrong Li --- .../_torch/models/modeling_deepseekv4.py | 33 +++++++++++++++ .../modeling/test_modeling_deepseekv4.py | 41 +++++++++++++++++++ 2 files changed, 74 insertions(+) diff --git a/tensorrt_llm/_torch/models/modeling_deepseekv4.py b/tensorrt_llm/_torch/models/modeling_deepseekv4.py index fa7f56c385b7..51babd6832fb 100644 --- a/tensorrt_llm/_torch/models/modeling_deepseekv4.py +++ b/tensorrt_llm/_torch/models/modeling_deepseekv4.py @@ -296,6 +296,38 @@ def _resolve_enable_fused_hc(config: PretrainedConfig) -> bool: return bool(getattr(config, "enable_fused_hc", True)) +def _normalize_deepseek_v4_mixed_precision_config( + model_config: ModelConfig[PretrainedConfig], +) -> ModelConfig[PretrainedConfig]: + """Resolve DeepSeek-V4 NVFP4 checkpoints' FP8 base-layer config.""" + quant_config = model_config.quant_config + hf_quant_config = getattr(model_config.pretrained_config, "quantization_config", None) + if ( + quant_config.quant_algo != QuantAlgo.MIXED_PRECISION + 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 + + fp8_quant_config = quant_config.model_copy( + deep=True, + update={ + "quant_algo": QuantAlgo.FP8_BLOCK_SCALES, + "group_size": 128, + "exclude_modules": ["*kv_b_proj*", "*k_b_proj*", "*eh_proj*"], + }, + ) + fp8_quant_config.__dict__.pop("quant_mode", None) + fp8_quant_config.__dict__.pop("layer_quant_mode", None) + + normalized_config = copy.deepcopy(model_config) + normalized_config._frozen = False + normalized_config.quant_config = fp8_quant_config + normalized_config._frozen = True + return normalized_config + + def _copy_deepseek_v4_fused_a_weight_scale( module: Linear, fused_a: torch.Tensor, fused_a_scale: torch.Tensor ) -> None: @@ -2484,6 +2516,7 @@ def get_model_defaults(cls, llm_args: "TorchLlmArgs") -> dict: } def __init__(self, model_config: ModelConfig[PretrainedConfig]): + model_config = _normalize_deepseek_v4_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 c253c5e71be0..6585ffb5c713 100644 --- a/tests/unittest/_torch/modeling/test_modeling_deepseekv4.py +++ b/tests/unittest/_torch/modeling/test_modeling_deepseekv4.py @@ -35,6 +35,7 @@ DeepseekV4MTP, _copy_deepseek_v4_fused_a_weight_scale, _deepseek_v4_pos_embd_params, + _normalize_deepseek_v4_mixed_precision_config, _remap_deepseek_v4_checkpoint_keys, _resolve_enable_fused_hc, ) @@ -444,6 +445,46 @@ def test_deepseek_v4_moe_auto_backend_on_blackwell(monkeypatch): assert ModelConfig.resolve_moe_backend("AUTO", "DeepseekV4ForCausalLM") == "TRTLLM" +def test_deepseek_v4_mixed_precision_uses_fp8_base_config(): + config = DeepseekV4Config() + config.quantization_config = { + "quant_method": "fp8", + "weight_block_size": [128, 128], + } + 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_mixed_precision_config(model_config) + + assert normalized_config is not model_config + assert model_config.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 == [ + "*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_decoder_accepts_mixed_precision_experts(monkeypatch): config = DeepseekV4Config( hidden_size=16, From 5762b27e0515e674afa1379d71a4eca2cf14993f Mon Sep 17 00:00:00 2001 From: Fanrong Li Date: Tue, 21 Jul 2026 15:20:45 +0000 Subject: [PATCH 3/5] [None][fix] Normalize DeepSeek V4 quant config before model init Signed-off-by: Fanrong Li --- tensorrt_llm/_torch/model_config.py | 34 +++++++ .../_torch/models/modeling_deepseekv4.py | 43 +++------ .../modeling/test_modeling_deepseekv4.py | 88 +++++++++++++++---- 3 files changed, 115 insertions(+), 50 deletions(-) diff --git a/tensorrt_llm/_torch/model_config.py b/tensorrt_llm/_torch/model_config.py index 8f5641920c03..be705650f406 100644 --- a/tensorrt_llm/_torch/model_config.py +++ b/tensorrt_llm/_torch/model_config.py @@ -639,6 +639,36 @@ def _has_deepseek_v4_layer_only_modelopt_quant_config( return (quantization_config.get('quant_algo', None) is None and quantization_config.get('quantized_layers', None) is not None) + @staticmethod + def _normalize_deepseek_v4_mixed_precision_base_quant_config( + pretrained_config: transformers.PretrainedConfig, + quant_config: QuantConfig) -> QuantConfig: + """Resolve FP8 base layers in DeepSeek-V4 mixed checkpoints.""" + hf_quant_config = getattr(pretrained_config, "quantization_config", + None) + if (quant_config.quant_algo != QuantAlgo.MIXED_PRECISION + 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 quant_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) + return fp8_quant_config + @staticmethod def _set_deepseek_v4_routed_moe_quant_config(pretrained_config, checkpoint_dir: str, @@ -1175,6 +1205,10 @@ def _recursive_update_config(config: transformers.PretrainedConfig, quant_config, layer_quant_config = cls.load_quant_config_from_dtypes_json( quant_config_file, moe_backend_hint) + if architecture in _DEEPSEEK_V4_ARCHITECTURES: + quant_config = cls._normalize_deepseek_v4_mixed_precision_base_quant_config( + pretrained_config, quant_config) + kwargs['moe_backend'] = cls.resolve_moe_backend( requested_moe_backend, architecture, diff --git a/tensorrt_llm/_torch/models/modeling_deepseekv4.py b/tensorrt_llm/_torch/models/modeling_deepseekv4.py index 51babd6832fb..0b03e988a882 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,38 +299,6 @@ def _resolve_enable_fused_hc(config: PretrainedConfig) -> bool: return bool(getattr(config, "enable_fused_hc", True)) -def _normalize_deepseek_v4_mixed_precision_config( - model_config: ModelConfig[PretrainedConfig], -) -> ModelConfig[PretrainedConfig]: - """Resolve DeepSeek-V4 NVFP4 checkpoints' FP8 base-layer config.""" - quant_config = model_config.quant_config - hf_quant_config = getattr(model_config.pretrained_config, "quantization_config", None) - if ( - quant_config.quant_algo != QuantAlgo.MIXED_PRECISION - 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 - - fp8_quant_config = quant_config.model_copy( - deep=True, - update={ - "quant_algo": QuantAlgo.FP8_BLOCK_SCALES, - "group_size": 128, - "exclude_modules": ["*kv_b_proj*", "*k_b_proj*", "*eh_proj*"], - }, - ) - fp8_quant_config.__dict__.pop("quant_mode", None) - fp8_quant_config.__dict__.pop("layer_quant_mode", None) - - normalized_config = copy.deepcopy(model_config) - normalized_config._frozen = False - normalized_config.quant_config = fp8_quant_config - normalized_config._frozen = True - return normalized_config - - def _copy_deepseek_v4_fused_a_weight_scale( module: Linear, fused_a: torch.Tensor, fused_a_scale: torch.Tensor ) -> None: @@ -2516,7 +2487,13 @@ def get_model_defaults(cls, llm_args: "TorchLlmArgs") -> dict: } def __init__(self, model_config: ModelConfig[PretrainedConfig]): - model_config = _normalize_deepseek_v4_mixed_precision_config(model_config) + # ModelConfig.from_pretrained resolves this before backend selection. + # Keep direct ModelConfig construction consistent with that path. + model_config.quant_config = ( + ModelConfig._normalize_deepseek_v4_mixed_precision_base_quant_config( + model_config.pretrained_config, model_config.quant_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 6585ffb5c713..31b956707acd 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 @@ -35,7 +38,6 @@ DeepseekV4MTP, _copy_deepseek_v4_fused_a_weight_scale, _deepseek_v4_pos_embd_params, - _normalize_deepseek_v4_mixed_precision_config, _remap_deepseek_v4_checkpoint_keys, _resolve_enable_fused_hc, ) @@ -450,6 +452,7 @@ def test_deepseek_v4_mixed_precision_uses_fp8_base_config(): config.quantization_config = { "quant_method": "fp8", "weight_block_size": [128, 128], + "modules_to_not_convert": ["lm_head", "*eh_proj*"], } mixed_quant_config = QuantConfig( quant_algo=QuantAlgo.MIXED_PRECISION, @@ -458,30 +461,81 @@ def test_deepseek_v4_mixed_precision_uses_fp8_base_config(): ) 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}, + normalized_config = ModelConfig._normalize_deepseek_v4_mixed_precision_base_quant_config( + config, mixed_quant_config ) - model_config._frozen = True - normalized_config = _normalize_deepseek_v4_mixed_precision_config(model_config) + assert normalized_config is not mixed_quant_config + assert mixed_quant_config.quant_algo == QuantAlgo.MIXED_PRECISION + assert normalized_config.quant_algo == QuantAlgo.FP8_BLOCK_SCALES + assert normalized_config.layer_quant_mode.has_fp8_block_scales() + assert normalized_config.group_size == 128 + assert normalized_config.mamba_ssm_cache_dtype == torch.bfloat16 + assert normalized_config.exclude_modules == [ + "lm_head", + "*eh_proj*", + "*kv_b_proj*", + "*k_b_proj*", + ] + + +def test_deepseek_v4_model_config_resolves_mixed_precision_base(tmp_path, monkeypatch): + config = DeepseekV4Config( + architectures=["DeepseekV4ForCausalLM"], + num_hidden_layers=1, + compress_ratios=[1], + ) + config.quantization_config = { + "quant_method": "fp8", + "weight_block_size": [128, 128], + "modules_to_not_convert": ["lm_head"], + } + monkeypatch.setattr( + "tensorrt_llm._torch.model_config.load_pretrained_config", + lambda *args, **kwargs: config, + ) + (tmp_path / "hf_quant_config.json").write_text( + json.dumps( + { + "producer": { + "name": "modelopt", + "version": "dsv4-nvfp4-experts", + }, + "quantization": { + "quant_algo": "MIXED_PRECISION", + "group_size": 16, + "exclude_modules": ["*.attn.*", "*.ffn.shared_experts.*"], + "quantized_layers": { + "layers.0.ffn.experts": { + "quant_algo": "NVFP4", + "group_size": 16, + } + }, + }, + } + ) + ) + tensor_name = "layers.0.ffn.experts.0.w1.weight" + shard_name = "model-00001-of-00001.safetensors" + _write_safetensors_header(tmp_path / shard_name, tensor_name, "U8", [2, 2]) + (tmp_path / "model.safetensors.index.json").write_text( + json.dumps({"weight_map": {tensor_name: shard_name}}) + ) + + model_config = ModelConfig.from_pretrained( + str(tmp_path), attn_backend="TRTLLM", moe_backend="TRTLLM" + ) - assert normalized_config is not model_config - assert model_config.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 == [ + assert model_config.quant_config.quant_algo == QuantAlgo.FP8_BLOCK_SCALES + assert model_config.quant_config.group_size == 128 + assert model_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 + model_config.quant_config_dict["model.layers.0.mlp.experts"].quant_algo == QuantAlgo.NVFP4 ) From 5b266eec087476c019292ef6587649564e5fed7d Mon Sep 17 00:00:00 2001 From: Fanrong Li Date: Tue, 21 Jul 2026 15:39:16 +0000 Subject: [PATCH 4/5] [None][refactor] Scope DeepSeek V4 quant normalization to model Signed-off-by: Fanrong Li --- tensorrt_llm/_torch/model_config.py | 34 -------- .../_torch/models/modeling_deepseekv4.py | 45 ++++++++-- .../modeling/test_modeling_deepseekv4.py | 87 ++++--------------- 3 files changed, 57 insertions(+), 109 deletions(-) diff --git a/tensorrt_llm/_torch/model_config.py b/tensorrt_llm/_torch/model_config.py index be705650f406..8f5641920c03 100644 --- a/tensorrt_llm/_torch/model_config.py +++ b/tensorrt_llm/_torch/model_config.py @@ -639,36 +639,6 @@ def _has_deepseek_v4_layer_only_modelopt_quant_config( return (quantization_config.get('quant_algo', None) is None and quantization_config.get('quantized_layers', None) is not None) - @staticmethod - def _normalize_deepseek_v4_mixed_precision_base_quant_config( - pretrained_config: transformers.PretrainedConfig, - quant_config: QuantConfig) -> QuantConfig: - """Resolve FP8 base layers in DeepSeek-V4 mixed checkpoints.""" - hf_quant_config = getattr(pretrained_config, "quantization_config", - None) - if (quant_config.quant_algo != QuantAlgo.MIXED_PRECISION - 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 quant_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) - return fp8_quant_config - @staticmethod def _set_deepseek_v4_routed_moe_quant_config(pretrained_config, checkpoint_dir: str, @@ -1205,10 +1175,6 @@ def _recursive_update_config(config: transformers.PretrainedConfig, quant_config, layer_quant_config = cls.load_quant_config_from_dtypes_json( quant_config_file, moe_backend_hint) - if architecture in _DEEPSEEK_V4_ARCHITECTURES: - quant_config = cls._normalize_deepseek_v4_mixed_precision_base_quant_config( - pretrained_config, quant_config) - kwargs['moe_backend'] = cls.resolve_moe_backend( requested_moe_backend, architecture, diff --git a/tensorrt_llm/_torch/models/modeling_deepseekv4.py b/tensorrt_llm/_torch/models/modeling_deepseekv4.py index 0b03e988a882..da0d58c5a66f 100644 --- a/tensorrt_llm/_torch/models/modeling_deepseekv4.py +++ b/tensorrt_llm/_torch/models/modeling_deepseekv4.py @@ -299,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: @@ -2487,13 +2524,7 @@ def get_model_defaults(cls, llm_args: "TorchLlmArgs") -> dict: } def __init__(self, model_config: ModelConfig[PretrainedConfig]): - # ModelConfig.from_pretrained resolves this before backend selection. - # Keep direct ModelConfig construction consistent with that path. - model_config.quant_config = ( - ModelConfig._normalize_deepseek_v4_mixed_precision_base_quant_config( - model_config.pretrained_config, model_config.quant_config - ) - ) + 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 31b956707acd..9e8338a08add 100644 --- a/tests/unittest/_torch/modeling/test_modeling_deepseekv4.py +++ b/tests/unittest/_torch/modeling/test_modeling_deepseekv4.py @@ -38,6 +38,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, ) @@ -447,12 +448,12 @@ def test_deepseek_v4_moe_auto_backend_on_blackwell(monkeypatch): assert ModelConfig.resolve_moe_backend("AUTO", "DeepseekV4ForCausalLM") == "TRTLLM" -def test_deepseek_v4_mixed_precision_uses_fp8_base_config(): +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", "*eh_proj*"], + "modules_to_not_convert": ["lm_head"], } mixed_quant_config = QuantConfig( quant_algo=QuantAlgo.MIXED_PRECISION, @@ -461,81 +462,31 @@ def test_deepseek_v4_mixed_precision_uses_fp8_base_config(): ) mixed_quant_config.mamba_ssm_cache_dtype = torch.bfloat16 assert not mixed_quant_config.layer_quant_mode.has_fp8_block_scales() - normalized_config = ModelConfig._normalize_deepseek_v4_mixed_precision_base_quant_config( - config, mixed_quant_config - ) - - assert normalized_config is not mixed_quant_config - assert mixed_quant_config.quant_algo == QuantAlgo.MIXED_PRECISION - assert normalized_config.quant_algo == QuantAlgo.FP8_BLOCK_SCALES - assert normalized_config.layer_quant_mode.has_fp8_block_scales() - assert normalized_config.group_size == 128 - assert normalized_config.mamba_ssm_cache_dtype == torch.bfloat16 - assert normalized_config.exclude_modules == [ - "lm_head", - "*eh_proj*", - "*kv_b_proj*", - "*k_b_proj*", - ] - - -def test_deepseek_v4_model_config_resolves_mixed_precision_base(tmp_path, monkeypatch): - config = DeepseekV4Config( - architectures=["DeepseekV4ForCausalLM"], - num_hidden_layers=1, - compress_ratios=[1], - ) - config.quantization_config = { - "quant_method": "fp8", - "weight_block_size": [128, 128], - "modules_to_not_convert": ["lm_head"], - } - monkeypatch.setattr( - "tensorrt_llm._torch.model_config.load_pretrained_config", - lambda *args, **kwargs: config, - ) - (tmp_path / "hf_quant_config.json").write_text( - json.dumps( - { - "producer": { - "name": "modelopt", - "version": "dsv4-nvfp4-experts", - }, - "quantization": { - "quant_algo": "MIXED_PRECISION", - "group_size": 16, - "exclude_modules": ["*.attn.*", "*.ffn.shared_experts.*"], - "quantized_layers": { - "layers.0.ffn.experts": { - "quant_algo": "NVFP4", - "group_size": 16, - } - }, - }, - } - ) - ) - tensor_name = "layers.0.ffn.experts.0.w1.weight" - shard_name = "model-00001-of-00001.safetensors" - _write_safetensors_header(tmp_path / shard_name, tensor_name, "U8", [2, 2]) - (tmp_path / "model.safetensors.index.json").write_text( - json.dumps({"weight_map": {tensor_name: shard_name}}) + 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 - model_config = ModelConfig.from_pretrained( - str(tmp_path), attn_backend="TRTLLM", moe_backend="TRTLLM" - ) + normalized_config = _normalize_deepseek_v4_nvfp4_mixed_precision_config(model_config) - assert model_config.quant_config.quant_algo == QuantAlgo.FP8_BLOCK_SCALES - assert model_config.quant_config.group_size == 128 - assert model_config.quant_config.exclude_modules == [ + 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 ( - model_config.quant_config_dict["model.layers.0.mlp.experts"].quant_algo == QuantAlgo.NVFP4 + normalized_config.quant_config_dict["model.layers.0.mlp.experts"].quant_algo + == QuantAlgo.NVFP4 ) From 981165230461f3c5d097471ed15d9eb2de034e4d Mon Sep 17 00:00:00 2001 From: Fanrong Li Date: Tue, 21 Jul 2026 15:42:01 +0000 Subject: [PATCH 5/5] [None][test] Remove redundant DeepSeek V4 mixed-precision test Signed-off-by: Fanrong Li --- .../_torch/models/modeling_deepseekv4.py | 9 ++- .../modeling/test_modeling_deepseekv4.py | 61 ------------------- 2 files changed, 4 insertions(+), 66 deletions(-) diff --git a/tensorrt_llm/_torch/models/modeling_deepseekv4.py b/tensorrt_llm/_torch/models/modeling_deepseekv4.py index da0d58c5a66f..5724c1052d8c 100644 --- a/tensorrt_llm/_torch/models/modeling_deepseekv4.py +++ b/tensorrt_llm/_torch/models/modeling_deepseekv4.py @@ -1789,12 +1789,11 @@ def __init__( post_mult_value=2.0, ) + # FIXME: incompatible with mixed quantization mode quant_config = self._get_decoder_layer_quant_config(model_config, layer_idx) - # MIXED_PRECISION is resolved per module. Routed experts use their - # layer-specific config, while layer-level NVFP4 fusions stay disabled. - self.is_nvfp4 = ( - quant_config.quant_algo != QuantAlgo.MIXED_PRECISION - and quant_config.layer_quant_mode.has_nvfp4() + self.is_nvfp4 = quant_config.layer_quant_mode.has_nvfp4() + assert quant_config.quant_algo is not QuantAlgo.MIXED_PRECISION, ( + "MIXED_PRECISION is ambiguous" ) self.allreduce = None diff --git a/tests/unittest/_torch/modeling/test_modeling_deepseekv4.py b/tests/unittest/_torch/modeling/test_modeling_deepseekv4.py index 9e8338a08add..b46ed3e477b9 100644 --- a/tests/unittest/_torch/modeling/test_modeling_deepseekv4.py +++ b/tests/unittest/_torch/modeling/test_modeling_deepseekv4.py @@ -8,7 +8,6 @@ import textwrap import weakref from copy import deepcopy -from types import SimpleNamespace import pytest import torch @@ -34,7 +33,6 @@ DeepseekV4DecoderLayer, DeepseekV4ForCausalLM, DeepseekV4Gate, - DeepseekV4MoE, DeepseekV4MTP, _copy_deepseek_v4_fused_a_weight_scale, _deepseek_v4_pos_embd_params, @@ -490,65 +488,6 @@ def test_deepseek_v4_nvfp4_mixed_precision_config(): ) -def test_deepseek_v4_decoder_accepts_mixed_precision_experts(monkeypatch): - config = DeepseekV4Config( - hidden_size=16, - moe_intermediate_size=8, - n_routed_experts=4, - n_shared_experts=1, - num_experts_per_tok=2, - hc_mult=2, - hc_sinkhorn_iters=1, - ) - config.torch_dtype = torch.bfloat16 - quant_config = QuantConfig( - quant_algo=QuantAlgo.MIXED_PRECISION, - group_size=16, - exclude_modules=["*.attn.*", "*.ffn.shared_experts.*", "head", "mtp.*"], - ) - experts_quant_config = QuantConfig(quant_algo=QuantAlgo.NVFP4, group_size=16) - model_config = SimpleNamespace( - pretrained_config=config, - mapping=Mapping(world_size=1, rank=0, tp_size=1), - quant_config=quant_config, - quant_config_dict={"model.layers.0.mlp.experts": experts_quant_config}, - allreduce_strategy=None, - ) - captured = {} - - def fake_moe(**kwargs): - captured.update(kwargs) - return torch.nn.Identity() - - monkeypatch.setattr( - "tensorrt_llm._torch.models.modeling_deepseekv4.mHC", - lambda *args, **kwargs: torch.nn.Identity(), - ) - monkeypatch.setattr( - "tensorrt_llm._torch.models.modeling_deepseekv4.DeepseekV4Attention", - lambda *args, **kwargs: torch.nn.Identity(), - ) - monkeypatch.setattr("tensorrt_llm._torch.models.modeling_deepseekv4.DeepseekV4MoE", fake_moe) - monkeypatch.setattr( - "tensorrt_llm._torch.models.modeling_deepseekv4.RMSNorm", - lambda *args, **kwargs: torch.nn.Identity(), - ) - monkeypatch.setattr( - "tensorrt_llm._torch.models.modeling_deepseekv4.can_access_peer", - lambda _mapping: False, - ) - - layer = DeepseekV4DecoderLayer( - model_config, - layer_idx=0, - aux_stream_dict={AuxStreamType.Attention: None, AuxStreamType.MoeShared: None}, - ) - - assert layer.is_nvfp4 is False - assert captured["override_quant_config"] is quant_config - assert DeepseekV4MoE._get_experts_quant_config(model_config, 0) is experts_quant_config - - 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"