From 8c9001cccbe6c224d688ca9b97fe8ffd5eeec0c1 Mon Sep 17 00:00:00 2001 From: realAsma Date: Wed, 2 Sep 2026 20:49:25 +0000 Subject: [PATCH 1/6] Fix INT4 packing for partial AWQ blocks Signed-off-by: realAsma --- modelopt/torch/export/quant_utils.py | 13 ++++++++++--- tests/gpu/torch/export/test_export.py | 18 ++++++++++++++++++ 2 files changed, 28 insertions(+), 3 deletions(-) diff --git a/modelopt/torch/export/quant_utils.py b/modelopt/torch/export/quant_utils.py index a40934261d9..6b422bbc98a 100755 --- a/modelopt/torch/export/quant_utils.py +++ b/modelopt/torch/export/quant_utils.py @@ -792,12 +792,19 @@ def process_layer_quant_config(layer_config_dict): return per_layer_config -def pack_int4_in_uint8(weight, weights_scaling_factor): +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] + if block_size is None or block_size <= 0: + raise ValueError(f"Block size must be a positive integer, got {block_size}.") + expected_scale_count = (in_dim + block_size - 1) // block_size + if weights_scaling_factor.shape[-1] != expected_scale_count: + raise ValueError( + f"Expected {expected_scale_count} weight scaling factors for input dimension {in_dim} " + f"and block size {block_size}, got {weights_scaling_factor.shape[-1]}." + ) # Scale, round, and clamp to the signed 4-bit range [-8..7]. int8_tensor = ( @@ -912,7 +919,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, diff --git a/tests/gpu/torch/export/test_export.py b/tests/gpu/torch/export/test_export.py index 55137a64639..a70fe866472 100644 --- a/tests/gpu/torch/export/test_export.py +++ b/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 ( @@ -393,6 +394,23 @@ def test_get_weight_block_size(config, expected_block_size): assert block_size == 0 +def test_to_quantized_weight_int4_partial_block(): + block_size = 128 + in_dim = 2 * block_size + 1 + scales = torch.tensor([[1.0, 2.0, 4.0]] * 4, device="cuda") + quantized_values = torch.arange(1, 5, device="cuda")[:, None] + weight = scales.repeat_interleave(block_size, dim=-1)[..., :in_dim] * quantized_values + + packed = to_quantized_weight(weight, scales, QUANTIZATION_W4A8_AWQ, 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")) + + with pytest.raises(ValueError, match="Expected 3 weight scaling factors"): + to_quantized_weight(weight, scales[..., :-1], QUANTIZATION_W4A8_AWQ, block_size=block_size) + + @pytest.mark.parametrize( ("config", "maxbound", "expected_amax"), [ From 69054641954666a7efa89b89879012847d6cc19e Mon Sep 17 00:00:00 2001 From: realAsma Date: Wed, 2 Sep 2026 21:24:20 +0000 Subject: [PATCH 2/6] Reject partial INT4 AWQ blocks Signed-off-by: realAsma --- modelopt/torch/export/quant_utils.py | 7 ++++++- tests/gpu/torch/export/test_export.py | 19 ++++++++++++------- 2 files changed, 18 insertions(+), 8 deletions(-) diff --git a/modelopt/torch/export/quant_utils.py b/modelopt/torch/export/quant_utils.py index 6b422bbc98a..5b06712b0ca 100755 --- a/modelopt/torch/export/quant_utils.py +++ b/modelopt/torch/export/quant_utils.py @@ -799,7 +799,12 @@ def pack_int4_in_uint8(weight, weights_scaling_factor, block_size): in_dim = weight.shape[-1] if block_size is None or block_size <= 0: raise ValueError(f"Block size must be a positive integer, got {block_size}.") - expected_scale_count = (in_dim + block_size - 1) // 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." + ) + expected_scale_count = in_dim // block_size if weights_scaling_factor.shape[-1] != expected_scale_count: raise ValueError( f"Expected {expected_scale_count} weight scaling factors for input dimension {in_dim} " diff --git a/tests/gpu/torch/export/test_export.py b/tests/gpu/torch/export/test_export.py index a70fe866472..f8a61c08ac3 100644 --- a/tests/gpu/torch/export/test_export.py +++ b/tests/gpu/torch/export/test_export.py @@ -394,21 +394,26 @@ def test_get_weight_block_size(config, expected_block_size): assert block_size == 0 -def test_to_quantized_weight_int4_partial_block(): +@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 + 1 - scales = torch.tensor([[1.0, 2.0, 4.0]] * 4, device="cuda") + 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)[..., :in_dim] * quantized_values + weight = scales.repeat_interleave(block_size, dim=-1) * quantized_values - packed = to_quantized_weight(weight, scales, QUANTIZATION_W4A8_AWQ, block_size=block_size) + 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")) - with pytest.raises(ValueError, match="Expected 3 weight scaling factors"): - to_quantized_weight(weight, scales[..., :-1], QUANTIZATION_W4A8_AWQ, block_size=block_size) + with pytest.raises(ValueError, match="Expected 2 weight scaling factors"): + to_quantized_weight(weight, scales[..., :-1], quantization, block_size=block_size) + + partial_weight = torch.cat((weight, quantized_values), dim=-1) + with pytest.raises(NotImplementedError, match="partial blocks are not supported"): + to_quantized_weight(partial_weight, scales, quantization, block_size=block_size) @pytest.mark.parametrize( From 3f2c679c174e47e627dad56cc910e5372737b219 Mon Sep 17 00:00:00 2001 From: realAsma Date: Wed, 2 Sep 2026 22:33:43 +0000 Subject: [PATCH 3/6] Validate INT4 AWQ block sizes Signed-off-by: realAsma --- modelopt/torch/export/quant_utils.py | 2 +- tests/gpu/torch/export/test_export.py | 10 ++++++++++ 2 files changed, 11 insertions(+), 1 deletion(-) diff --git a/modelopt/torch/export/quant_utils.py b/modelopt/torch/export/quant_utils.py index 5b06712b0ca..b1032efc6f5 100755 --- a/modelopt/torch/export/quant_utils.py +++ b/modelopt/torch/export/quant_utils.py @@ -797,7 +797,7 @@ def pack_int4_in_uint8(weight, weights_scaling_factor, block_size): 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] - if block_size is None or block_size <= 0: + 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( diff --git a/tests/gpu/torch/export/test_export.py b/tests/gpu/torch/export/test_export.py index f8a61c08ac3..a9e471f89e9 100644 --- a/tests/gpu/torch/export/test_export.py +++ b/tests/gpu/torch/export/test_export.py @@ -416,6 +416,16 @@ def test_to_quantized_weight_int4_block_size(quantization): to_quantized_weight(partial_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"), [ From 1580415775f1be12b71ff052569f95ccd46099dc Mon Sep 17 00:00:00 2001 From: realAsma Date: Wed, 2 Sep 2026 23:39:10 +0000 Subject: [PATCH 4/6] Fix AWQ export test dimensions Signed-off-by: realAsma --- tests/gpu/torch/export/test_fsdp2_export.py | 6 +++--- .../test_unified_hf_export_and_check_safetensors.py | 9 ++++++++- 2 files changed, 11 insertions(+), 4 deletions(-) diff --git a/tests/gpu/torch/export/test_fsdp2_export.py b/tests/gpu/torch/export/test_fsdp2_export.py index 5f7c186c0ed..65f33abd41c 100644 --- a/tests/gpu/torch/export/test_fsdp2_export.py +++ b/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) diff --git a/tests/gpu/torch/export/test_unified_hf_export_and_check_safetensors.py b/tests/gpu/torch/export/test_unified_hf_export_and_check_safetensors.py index 486d0fdbed5..b45943e6e5e 100644 --- a/tests/gpu/torch/export/test_unified_hf_export_and_check_safetensors.py +++ b/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" From 89bdde581f6d675f7b913eddad0ec5b3109afe63 Mon Sep 17 00:00:00 2001 From: realAsma Date: Fri, 4 Sep 2026 04:43:05 +0000 Subject: [PATCH 5/6] Reject partial-block compressed AWQ export Signed-off-by: realAsma --- modelopt/torch/export/quant_utils.py | 22 +++++++++++----------- tests/gpu/torch/export/test_export.py | 12 ++++++++---- 2 files changed, 19 insertions(+), 15 deletions(-) diff --git a/modelopt/torch/export/quant_utils.py b/modelopt/torch/export/quant_utils.py index b1032efc6f5..e0ab6dd6d1f 100755 --- a/modelopt/torch/export/quant_utils.py +++ b/modelopt/torch/export/quant_utils.py @@ -792,11 +792,7 @@ def process_layer_quant_config(layer_config_dict): return per_layer_config -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] +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: @@ -804,12 +800,14 @@ def pack_int4_in_uint8(weight, weights_scaling_factor, block_size): f"Cannot pack weight with input dimension {in_dim} and block size {block_size}: " "partial blocks are not supported." ) - expected_scale_count = in_dim // block_size - if weights_scaling_factor.shape[-1] != expected_scale_count: - raise ValueError( - f"Expected {expected_scale_count} weight scaling factors for input dimension {in_dim} " - f"and block size {block_size}, got {weights_scaling_factor.shape[-1]}." - ) + + +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] + _validate_int4_block_size(in_dim, block_size) # Scale, round, and clamp to the signed 4-bit range [-8..7]. int8_tensor = ( @@ -864,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: diff --git a/tests/gpu/torch/export/test_export.py b/tests/gpu/torch/export/test_export.py index a9e471f89e9..36e6cbbdecc 100644 --- a/tests/gpu/torch/export/test_export.py +++ b/tests/gpu/torch/export/test_export.py @@ -70,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( @@ -408,13 +409,16 @@ def test_to_quantized_weight_int4_block_size(quantization): 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")) - with pytest.raises(ValueError, match="Expected 2 weight scaling factors"): - to_quantized_weight(weight, scales[..., :-1], quantization, block_size=block_size) - - partial_weight = torch.cat((weight, quantized_values), dim=-1) + 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]) From ffc0c3a27ed9ec23d2edb68436d6e67381bf7046 Mon Sep 17 00:00:00 2001 From: realAsma Date: Fri, 4 Sep 2026 05:40:03 +0000 Subject: [PATCH 6/6] Document partial-block AWQ export rejection Signed-off-by: realAsma --- CHANGELOG.rst | 1 + 1 file changed, 1 insertion(+) diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 68927078e83..5b5b50542d2 100755 --- a/CHANGELOG.rst +++ b/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.