diff --git a/tensorrt_llm/_torch/models/checkpoints/hf/nemotron_h_weight_mapper.py b/tensorrt_llm/_torch/models/checkpoints/hf/nemotron_h_weight_mapper.py index 23462d725259..513146e8b254 100644 --- a/tensorrt_llm/_torch/models/checkpoints/hf/nemotron_h_weight_mapper.py +++ b/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 diff --git a/tensorrt_llm/_torch/models/modeling_utils.py b/tensorrt_llm/_torch/models/modeling_utils.py index 94d290d4f192..0e9a1ee325a9 100755 --- a/tensorrt_llm/_torch/models/modeling_utils.py +++ b/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(