[fix] Fix DeepSeek w4a8 weight loading

Signed-off-by: Jinyang Yuan <[email protected]>

Author: jinyangyuan-nvidia

tensorrt_llm/_torch/models/modeling_deepseekv3.py

@@ -25,6 +25,7 @@
 # SOFTWARE.
 # --------------------------------------------------
 
+import copy
 import math
 import os
 import warnings
@@ -1102,6 +1103,17 @@ class DeepseekV3ForCausalLM(DecoderModelForCausalLM[DeepseekV3Model,
                                                     PretrainedConfig]):
 
     def __init__(self, model_config: ModelConfig[PretrainedConfig]):
+        # Rename some keys of quant_config_dict to support legacy checkpoints
+        model_config = copy.deepcopy(model_config)
+        quant_config_dict = {}
+        for key, val in model_config.quant_config_dict.items():
+            key_split = key.split(".")
+            if key_split[-1] == "fused_a":
+                key = ".".join(key_split[:-1] + ["kv_a_proj_with_mqa"])
+            quant_config_dict[key] = val
+        model_config._frozen = False
+        model_config.quant_config_dict = quant_config_dict
+        model_config._frozen = True
         super().__init__(DeepseekV3Model(model_config),
                          config=model_config,
                          hidden_size=model_config.pretrained_config.hidden_size,

tensorrt_llm/_torch/models/modeling_utils.py

@@ -458,11 +458,11 @@ def __post_init__(self):
                         if name + '.q_proj' in n:
                             module.quant_config = q
                             break
-                elif hasattr(module, 'fused_a'):
+                elif hasattr(module, 'kv_a_proj_with_mqa'):
                     # DeepseekV3Attention
                     for n, q in quant_config_dict.items():
                         # reuse q_proj quant config as the attention quant config
-                        if name + '.fused_a' in n:
+                        if name + '.kv_a_proj_with_mqa' in n:
                             module.quant_config = q
                             break