fix: lm_head quantization divisibility uses TP (not PP); no misleading warning when already disabled

modelopt/torch/export/model_config_export.py

@@ -319,8 +319,11 @@ def torch_to_tensorrt_llm_checkpoint(
                 ):
                     config.share_embedding_table = True
                 else:
-                    # This will update lm_head quantization config according to constraints from TRT-LLM
-                    update_lm_head_quantization(config, module, inference_pipeline_parallel)
+                    # This will update lm_head quantization config according to constraints from TRT-LLM.
+                    # The lm_head (column linear) vocab dim is sharded by TP, not PP, so the AWQ
+                    # block divisibility check needs the inference TP size (0 means keeping the
+                    # calibration parallelism, i.e. no extra split).
+                    update_lm_head_quantization(config, module, max(inference_tensor_parallel, 1))
                     config.lm_head = build_linear_config(module, "column")
             elif is_conv(module) and decoder_type == "whisper":
                 if config.conv1 is None:

modelopt/torch/export/postprocess.py

@@ -746,7 +746,8 @@ def update_lm_head_quantization(
         input_quantizer.disable()
         print("Disable lm_head quantization for TRT-LLM export due to deployment limitations.")
 
-    else:
+    elif weight_quantizer.is_enabled:
+        # Only warn when lm_head quantization is actually enabled and kept.
         warn(
             "Enable lm_head quantization. lm_head quantization may lead to additional accuracy loss."
         )

tests/unit/torch/export/test_export_weight.py

@@ -20,6 +20,8 @@
 from _test_utils.torch.export.utils import ToyModel, partial_fp8_config, partial_w4a8_config
 
 import modelopt.torch.quantization as mtq
+from modelopt.torch.export.model_config import ModelConfig
+from modelopt.torch.export.postprocess import update_lm_head_quantization
 from modelopt.torch.export.unified_export_hf import (
     _export_quantized_weight,
     _process_quantized_modules,
@@ -102,6 +104,69 @@ def test_export_per_block_quantized_weight():
     assert not hasattr(model.linears[2], quantizer_attrs.output_scale)
 
 
+def _build_quantized_lm_head(vocab_size=256, hidden_size=256, block_size=128):
+    """Returns an INT4 block-quantized lm_head QuantLinear."""
+    model = nn.Sequential(nn.Linear(hidden_size, vocab_size, bias=False))
+    quant_cfg = {
+        "quant_cfg": [
+            {"quantizer_name": "*", "enable": False},
+            {
+                "quantizer_name": "*weight_quantizer",
+                "cfg": {"num_bits": 4, "block_sizes": {-1: block_size, "type": "static"}},
+                "enable": True,
+            },
+            {
+                "quantizer_name": "*input_quantizer",
+                "cfg": {"num_bits": 8, "axis": None},
+                "enable": True,
+            },
+        ],
+        "algorithm": "max",
+    }
+    mtq.quantize(model, quant_cfg, lambda m: m(torch.randn(2, hidden_size)))
+    return model[0]
+
+
+def test_update_lm_head_quantization_disabled_quantizer_no_warning(recwarn):
+    """A disabled lm_head quantizer must not trigger the 'Enable lm_head quantization' warning."""
+    lm_head = _build_quantized_lm_head()
+    lm_head.weight_quantizer.disable()
+    lm_head.input_quantizer.disable()
+
+    update_lm_head_quantization(ModelConfig(), lm_head, 1)
+
+    assert len(recwarn) == 0
+    # The function must not silently (re-)enable quantization either.
+    assert not lm_head.weight_quantizer.is_enabled
+    assert not lm_head.input_quantizer.is_enabled
+
+
+def test_update_lm_head_quantization_enabled_warns():
+    """An enabled, deployable lm_head quantizer stays enabled and warns about accuracy loss."""
+    lm_head = _build_quantized_lm_head()
+
+    with pytest.warns(UserWarning, match="Enable lm_head quantization"):
+        update_lm_head_quantization(ModelConfig(), lm_head, 1)
+
+    assert lm_head.weight_quantizer.is_enabled
+    assert lm_head.input_quantizer.is_enabled
+
+
+def test_update_lm_head_quantization_disables_on_tp_indivisible_vocab(recwarn):
+    """vocab_size divisible by block but not block * TP: quantization is disabled, no warning.
+
+    vocab 384 with block 128 deploys fine on TP=1, but TP=2 shards the vocab (column
+    linear) dim to 192, which no longer holds an integer number of AWQ blocks per rank.
+    """
+    lm_head = _build_quantized_lm_head(vocab_size=384, block_size=128)
+
+    update_lm_head_quantization(ModelConfig(), lm_head, 2)
+
+    assert not lm_head.weight_quantizer.is_enabled
+    assert not lm_head.input_quantizer.is_enabled
+    assert len(recwarn) == 0
+
+
 class QuantMoELinear(nn.Module):
     def __init__(self):
         super().__init__()