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. diff --git a/modelopt/torch/export/quant_utils.py b/modelopt/torch/export/quant_utils.py index a40934261d9..e0ab6dd6d1f 100755 --- a/modelopt/torch/export/quant_utils.py +++ b/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, diff --git a/tests/gpu/torch/export/test_export.py b/tests/gpu/torch/export/test_export.py index 55137a64639..36e6cbbdecc 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 ( @@ -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"), [ 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"