From 876e973ac1ab49b6ec49e15e258e03a6b7b0fbd9 Mon Sep 17 00:00:00 2001 From: Wanli Jiang <35160485+Wanli-Jiang@users.noreply.github.com> Date: Fri, 24 Jul 2026 07:53:36 +0000 Subject: [PATCH 1/2] [None][fix] Wire lm_head quant from global config for homogeneous checkpoints DecoderModelForCausalLM only resolved the lm_head quant entry from the per-layer quant_config_dict (MIXED_PRECISION checkpoints). A homogeneous checkpoint with a single global quant_algo that quantizes lm_head (and does not exclude it) built lm_head unquantized, so loading its packed NVFP4 weight failed with a shape mismatch (e.g. Nemotron-Nano-3.5 NVFP4: 'size of tensor a (2688) must match tensor b (1344)'). Resolve lm_head_quant_config from config.quant_config when quant_config_dict is None and the global config quantizes, and share the exclude_modules / tie_word_embeddings / lm-head-TP-in-ADP guards across both sources. No-op for existing models: the MIXED_PRECISION path is unchanged, and homogeneous checkpoints that keep lm_head in FP16 list it in exclude_modules, so the exclude guard keeps it unquantized as before. Verified on Nemotron-Nano-3.5 NVFP4 (W4A16_NVFP4 + FP8 KV) on H100: single-GPU and 4-GPU attention-DP + EP now load and generate coherently (with and without MTP). Signed-off-by: Wanli Jiang <35160485+Wanli-Jiang@users.noreply.github.com> --- tensorrt_llm/_torch/models/modeling_utils.py | 83 ++++++++++++-------- 1 file changed, 50 insertions(+), 33 deletions(-) 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( From d8b4bb66c805d79c8106505ba06c8c271d2c3137 Mon Sep 17 00:00:00 2001 From: Wanli Jiang <35160485+Wanli-Jiang@users.noreply.github.com> Date: Fri, 24 Jul 2026 08:32:39 +0000 Subject: [PATCH 2/2] [None][fix] Split mamba in_proj NVFP4 scales for the whole NVFP4 family NemotronHHfWeightMapper gated the structural (z/x/B/C/dt) TP split of the mamba in_proj block scales on quant_algo == "NVFP4", which misses W4A16_NVFP4 (and W4A8_NVFP4_FP8). For those checkpoints the packed in_proj weight was TP-split structurally but its group scales were passed through and even-split by the base loader, so weight and scale ended up misaligned and TP>1 produced garbage (TP1 / attention-DP use tp_size=1, so the split is identity and were unaffected). Use quant_config.quant_mode.has_nvfp4() so all NVFP4 variants take the structural scale split. Verified on Nemotron-Nano-3.5 NVFP4 (W4A16_NVFP4 + FP8 KV) on 4x H100: TP4 and TEP4 (with and without MTP) now generate coherently; previously all produced garbage. Signed-off-by: Wanli Jiang <35160485+Wanli-Jiang@users.noreply.github.com> --- .../_torch/models/checkpoints/hf/nemotron_h_weight_mapper.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) 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