diff --git a/modelopt/torch/export/model_config_export.py b/modelopt/torch/export/model_config_export.py index 1230702e320..d3884ce665f 100644 --- a/modelopt/torch/export/model_config_export.py +++ b/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: diff --git a/modelopt/torch/export/postprocess.py b/modelopt/torch/export/postprocess.py index 376a52a4134..8a0fd6dbaa2 100644 --- a/modelopt/torch/export/postprocess.py +++ b/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." ) diff --git a/tests/unit/torch/export/test_export_weight.py b/tests/unit/torch/export/test_export_weight.py index 6fc17d982e8..ec6f6da58e0 100644 --- a/tests/unit/torch/export/test_export_weight.py +++ b/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__()