[None][fix] Fix nemotron-h quant and loading config

tensorrt_llm/_torch/models/checkpoints/hf/nemotron_h_weight_mapper.py

@@ -42,7 +42,10 @@ def _split_mamba2_mixer_in_proj(w: torch.Tensor) -> torch.Tensor:
             w = torch.concat(w).contiguous()
             return w
 
-        is_nvfp4 = self.config.quant_config.quant_algo == "NVFP4"
+        # Covers the whole NVFP4 family (NVFP4, W4A16_NVFP4, W4A8_NVFP4_FP8, …);
+        # a bare ``== "NVFP4"`` check misses W4A16_NVFP4, leaving the mamba
+        # in_proj block scales unsplit while the weight is TP-split -> garbage.
+        is_nvfp4 = self.config.quant_config.quant_mode.has_nvfp4()
         n_groups = config.n_groups
         d_state = config.ssm_state_size
         nheads = config.mamba_num_heads

tensorrt_llm/_torch/models/modeling_utils.py

@@ -378,6 +378,53 @@ class DecoderModelForCausalLM(nn.Module,
                               Generic[TModel, TConfig],
                               metaclass=PostInitCaller):
 
+    @staticmethod
+    def _resolve_lm_head_quant_config(
+            config: ModelConfig[TConfig]) -> Optional[QuantConfig]:
+        """Resolve the quant config for lm_head, or None to keep it unquantized.
+
+        lm_head is quantized when the checkpoint quantizes it: MIXED_PRECISION
+        checkpoints carry a per-layer entry in ``quant_config_dict``; homogeneous
+        checkpoints (a single global ``quant_algo``) use the global
+        ``quant_config`` — the same quantize-by-default-then-exclude rule the
+        model applies to its own Linear layers. Several cases force it back to
+        unquantized.
+        """
+        if config.quant_config_dict is not None:
+            quant_config = config.quant_config_dict.get("lm_head")
+        elif (config.quant_config is not None
+              and config.quant_config.quant_algo is not None):
+            quant_config = config.quant_config
+        else:
+            quant_config = None
+        if quant_config is None:
+            return None
+
+        # exclude_modules always wins (an FP16 lm_head is listed there).
+        if (config.quant_config is not None
+                and config.quant_config.is_module_excluded_from_quantization(
+                    "lm_head")):
+            return None
+
+        # Tied embeddings replace lm_head.weight with the dense bf16 embedding
+        # weight, which would silently clash with a quantized (packed) weight.
+        if getattr(config.pretrained_config, 'tie_word_embeddings', False):
+            logger.info("Ignoring lm_head quant entry: tie_word_embeddings "
+                        "shares the dense embedding weight, so lm_head stays "
+                        "unquantized")
+            return None
+
+        # lm_head TP in ADP slices the dense weight at forward time for the
+        # spec-decoding head (see LMHead.forward), which is incompatible with
+        # quantized (packed) weights — LMHead rejects that at construction.
+        if (config.mapping.enable_attention_dp
+                and config.mapping.enable_lm_head_tp_in_adp):
+            logger.info("Ignoring lm_head quant entry: lm_head TP in ADP "
+                        "slices the dense weight, so lm_head stays unquantized")
+            return None
+
+        return quant_config
+
     def __init__(self, model: TModel, *, config: ModelConfig[TConfig],
                  hidden_size: int, vocab_size: int):
         super().__init__()
@@ -387,39 +434,9 @@ def __init__(self, model: TModel, *, config: ModelConfig[TConfig],
         self.pp_size = config.mapping.pp_size
         self.has_custom_lm_head = False
 
-        # Per-layer quant entry for lm_head (e.g. ModelOpt MIXED_PRECISION
-        # checkpoints that quantize lm_head to NVFP4). Model-specific
-        # config normalizers opt in by keeping/synthesizing this entry;
-        # exclude_modules still wins. Applies to both the attention-DP
-        # (replicated) and TP lm_head below.
-        lm_head_quant_config = None
-        if config.quant_config_dict is not None:
-            lm_head_quant_config = config.quant_config_dict.get("lm_head")
-            if (lm_head_quant_config is not None
-                    and config.quant_config is not None and config.quant_config.
-                    is_module_excluded_from_quantization("lm_head")):
-                lm_head_quant_config = None
-            if (lm_head_quant_config is not None and getattr(
-                    config.pretrained_config, 'tie_word_embeddings', False)):
-                # Tied embeddings replace lm_head.weight with the dense
-                # bf16 embedding weight below, which would silently clash
-                # with a quantized (packed) weight and its quant method.
-                logger.info(
-                    "Ignoring lm_head quant entry: tie_word_embeddings "
-                    "shares the dense embedding weight, so lm_head stays "
-                    "unquantized")
-                lm_head_quant_config = None
-            if (lm_head_quant_config is not None
-                    and config.mapping.enable_attention_dp
-                    and config.mapping.enable_lm_head_tp_in_adp):
-                # lm_head TP in ADP slices the dense weight at forward time
-                # for the spec-decoding head (see LMHead.forward), which is
-                # incompatible with quantized (packed) weights — LMHead
-                # rejects that combination at construction.
-                logger.info(
-                    "Ignoring lm_head quant entry: lm_head TP in ADP "
-                    "slices the dense weight, so lm_head stays unquantized")
-                lm_head_quant_config = None
+        # Quant config for lm_head (applies to both the attention-DP replicated
+        # and TP lm_head below); None keeps it unquantized.
+        lm_head_quant_config = self._resolve_lm_head_quant_config(config)
 
         if config.mapping.enable_attention_dp and not config.mapping.enable_lm_head_tp_in_adp:
             self.lm_head = LMHead(