Reject unsupported partial-block INT4/W4A8 AWQ export

CHANGELOG.rst

@@ -131,6 +131,7 @@ Changelog
 
 **Bug Fixes**
 
+- Reject INT4 and W4A8 AWQ checkpoint export when a weight's input dimension is not divisible by the configured block size. Choose an ``awq_block_size`` that evenly divides every quantized weight's input dimension.
 - Fix NemotronH dense MLP quantization with the ``nvfp4_mlp_only`` and ``nvfp4_omlp_only`` recipe families. NemotronH registers these projections as ``mixer.up_proj`` / ``mixer.down_proj``, which the previous ``*mlp*`` selector missed, producing checkpoints with a null ``quant_algo``.
 - Fix ``ShapeInferenceError`` during ONNX INT8 + FP16 quantization (``--high_precision_dtype fp16``) of weakly-typed models (e.g. TensorFlow exports) that carry stale rank-0 ``graph.output`` shapes or ops such as ``TopK`` that ONNX's static shape inference cannot resolve. Stale output shapes are now reconciled via symbolic shape inference, and AutoCast falls back to schema-based type inference so unresolved ops no longer leave tensors untyped.
 - Fix fused MoE expert auto-detection (``register_fused_experts_on_the_fly``) skipping modules without an ``act_fn`` attribute. Modules applying a custom gated activation between the two ``F.linear`` calls (e.g. ``MiniMaxM3VLExperts``) were silently skipped, leaving routed experts unquantized and failing HF export. Enables NVFP4/FP8 quantization and export for MiniMax-M2 / MiniMax-M3.

modelopt/torch/export/quant_utils.py

@@ -792,12 +792,22 @@ def process_layer_quant_config(layer_config_dict):
     return per_layer_config
 
 
-def pack_int4_in_uint8(weight, weights_scaling_factor):
+def _validate_int4_block_size(in_dim, block_size):
+    if not isinstance(block_size, int) or block_size <= 0:
+        raise ValueError(f"Block size must be a positive integer, got {block_size}.")
+    if in_dim % block_size != 0:
+        raise NotImplementedError(
+            f"Cannot pack weight with input dimension {in_dim} and block size {block_size}: "
+            "partial blocks are not supported."
+        )
+
+
+def pack_int4_in_uint8(weight, weights_scaling_factor, block_size):
     """Packs the INT4 weights into uint8 tensor."""
     out_dim = weight.shape[-2]
     assert out_dim % 2 == 0, f"Cannot pack weight. Out dimension {out_dim} is not an even number."
     in_dim = weight.shape[-1]
-    block_size = weight.shape[-1] // weights_scaling_factor.shape[-1]
+    _validate_int4_block_size(in_dim, block_size)
 
     # Scale, round, and clamp to the signed 4-bit range [-8..7].
     int8_tensor = (
@@ -852,6 +862,8 @@ def to_quantized_weight(
 
     # For compressed weights, we directly return the data from wrapper
     if isinstance(weight, QTensorWrapper):
+        if quantization in [QUANTIZATION_INT4_AWQ, QUANTIZATION_W4A8_AWQ]:
+            _validate_int4_block_size(weight.metadata["shape"][-1], block_size)
         return weight.data
 
     if quantization == QUANTIZATION_FP8:
@@ -912,7 +924,7 @@ def to_quantized_weight(
         return (weight / weights_scaling_factor[:, None]).to(torch.float8_e4m3fn)
 
     if quantization in [QUANTIZATION_INT4_AWQ, QUANTIZATION_W4A8_AWQ]:
-        return pack_int4_in_uint8(weight, weights_scaling_factor)
+        return pack_int4_in_uint8(weight, weights_scaling_factor, block_size)
 
     if quantization in [
         QUANTIZATION_NVFP4,

tests/gpu/torch/export/test_export.py

@@ -56,6 +56,7 @@
     get_weight_block_size,
     postprocess_state_dict,
     process_layer_quant_config,
+    to_quantized_weight,
 )
 from modelopt.torch.export.unified_export_hf import export_hf_checkpoint
 from modelopt.torch.quantization.config import (
@@ -69,6 +70,7 @@
     W4A8_AWQ_BETA_CFG,
 )
 from modelopt.torch.quantization.nn import SequentialQuantizer, TensorQuantizer
+from modelopt.torch.quantization.qtensor import INT4QTensor, QTensorWrapper
 
 
 @pytest.mark.parametrize(
@@ -393,6 +395,41 @@ def test_get_weight_block_size(config, expected_block_size):
             assert block_size == 0
 
 
+@pytest.mark.parametrize("quantization", [QUANTIZATION_INT4_AWQ, QUANTIZATION_W4A8_AWQ])
+def test_to_quantized_weight_int4_block_size(quantization):
+    block_size = 128
+    in_dim = 2 * block_size
+    scales = torch.tensor([[1.0, 2.0]] * 4, device="cuda")
+    quantized_values = torch.arange(1, 5, device="cuda")[:, None]
+    weight = scales.repeat_interleave(block_size, dim=-1) * quantized_values
+
+    packed = to_quantized_weight(weight, scales, quantization, block_size=block_size)
+
+    assert packed.shape == (2, in_dim)
+    assert torch.equal(packed[0], torch.full((in_dim,), 0x21, dtype=torch.uint8, device="cuda"))
+    assert torch.equal(packed[1], torch.full((in_dim,), 0x43, dtype=torch.uint8, device="cuda"))
+
+    partial_weight = torch.cat((weight, quantized_values.repeat(1, 2)), dim=-1)
+    with pytest.raises(NotImplementedError, match="partial blocks are not supported"):
+        to_quantized_weight(partial_weight, scales, quantization, block_size=block_size)
+
+    compressed_weight, _ = INT4QTensor.quantize(partial_weight, block_size)
+    with pytest.raises(NotImplementedError, match="partial blocks are not supported"):
+        to_quantized_weight(
+            QTensorWrapper(compressed_weight), scales, quantization, block_size=block_size
+        )
+
+
+@pytest.mark.parametrize("quantization", [QUANTIZATION_INT4_AWQ, QUANTIZATION_W4A8_AWQ])
+@pytest.mark.parametrize("block_size", [None, 0, -1, 2.0])
+def test_to_quantized_weight_invalid_int4_block_size(quantization, block_size):
+    weight = torch.ones((4, 4), device="cuda")
+    scales = torch.ones((4, 2), device="cuda")
+
+    with pytest.raises(ValueError, match="Block size must be a positive integer"):
+        to_quantized_weight(weight, scales, quantization, block_size=block_size)
+
+
 @pytest.mark.parametrize(
     ("config", "maxbound", "expected_amax"),
     [

tests/gpu/torch/export/test_fsdp2_export.py

@@ -161,15 +161,15 @@ def calib_fn(x):
 def _export_quantized_weight_test(rank, size, quant_config, bias):
     with patch_fsdp_mp_dtypes():
         # Initialize model
-        model = SmallQKVModel(dim=32, bias=bias).to("cuda")
-        non_fsdp_model = SmallQKVModel(dim=32, bias=bias).to("cuda")
+        model = SmallQKVModel(dim=128, bias=bias).to("cuda")
+        non_fsdp_model = SmallQKVModel(dim=128, bias=bias).to("cuda")
         non_fsdp_model.load_state_dict(copy.deepcopy(model.state_dict()))
         model.eval()
         non_fsdp_model.eval()
         _compare_parameters_and_buffers(model, non_fsdp_model)
 
         # Create calibration data ONCE
-        calib_data = torch.randn(1, 32, device="cuda")
+        calib_data = torch.randn(1, 128, device="cuda")
 
         def calib_fn(x):
             return x(calib_data)

tests/gpu/torch/export/test_unified_hf_export_and_check_safetensors.py

@@ -84,7 +84,14 @@ def test_unified_hf_export_and_check_safetensors(
     elif expected_suffix.startswith("tiny_gpt_oss"):
         tiny_model_dir = create_tiny_gpt_oss_dir(tmp_path, with_tokenizer=True, num_hidden_layers=1)
     else:
-        tiny_model_dir = create_tiny_llama_dir(tmp_path, with_tokenizer=True, num_hidden_layers=1)
+        model_dims = (
+            {"hidden_size": 128, "intermediate_size": 128}
+            if qformat in {"int4_awq", "w4a8_awq_beta"}
+            else {}
+        )
+        tiny_model_dir = create_tiny_llama_dir(
+            tmp_path, with_tokenizer=True, num_hidden_layers=1, **model_dims
+        )
 
     # Create an output directory in tmp_path
     # We'll replicate the naming convention, e.g. "tiny_llama-fp8"