[BUG] FP8 real_quantization doesnt work with block_sizes

[ishan-modi]

Describe the bug

Following is the part of code that is causing issue here

The issue occurs because the amax that is set during the calibration step doesnt take into consideration block_sizes here

And when we try to compress it the previously calculated amax is passed as scales here

This results into following error

/usr/local/lib/python3.11/dist-packages/modelopt/torch/quantization/qtensor/fp8_tensor.py in quantize(cls, input, scales, axis, block_sizes)
     99                 expanded_scales = expanded_scales.reshape(expected_shape)
    100 
--> 101             assert scales.shape == tuple(expected_shape), (
    102                 f"Mismatch in expected scale shape: {scales.shape} vs {tuple(expected_shape)}"
    103             )

AssertionError: Mismatch in expected scale shape: torch.Size([]) vs (1152, 18)

System information

nvidia_modelopt - 0.29.0

[ishan-modi]

@meenchen @realAsma, any updates on this ?

[meenchen]

Hi @ishan-modi, is this error still happening for the latest release? I think we have tests for 2D block real quant. If it still fails, could you share how to reproduce it?

[ishan-modi]

Hi @meenchen, the error is still there

from diffusers import SanaTransformer2DModel
import modelopt.torch.quantization as mtq

checkpoint = "Efficient-Large-Model/Sana_600M_1024px_diffusers"
model = SanaTransformer2DModel.from_pretrained(checkpoint, subfolder="transformer")

FP8_DEFAULT_CFG = {
    "quant_cfg": {
        "*weight_quantizer": {
            "fake_quant": False,
            "num_bits": (4, 3),
            "block_sizes": {-1: 128},
        },
        "*input_quantizer": {"enable": False},
        "*output_quantizer": {"enable": False},
        "*q_bmm_quantizer": {"enable": False},
        "*k_bmm_quantizer": {"enable": False},
        "*v_bmm_quantizer": {"enable": False},
        "*softmax_quantizer": {"enable": False},
        "nn.BatchNorm1d": {"*": {"enable": False}},
        "nn.BatchNorm2d": {"*": {"enable": False}},
        "nn.BatchNorm3d": {"*": {"enable": False}},
        "nn.LeakyReLU": {"*": {"enable": False}},
        "*lm_head*": {"enable": False},
        "*proj_out.*": {"enable": False},
        "*block_sparse_moe.gate*": {"enable": False},
        "*router*": {"enable": False},
        "*mlp.gate.*": {"enable": False},
        "*mlp.shared_expert_gate.*": {"enable": False},
        "*output_layer*": {"enable": False},
        "output.*": {"enable": False},
        "default": {"enable": False},
        "*conv*": {"enable": False},
        "*proj*": {"enable": False}
    },
    "algorithm": {"method": "max"},
}
mtq.quantize(model, FP8_DEFAULT_CFG)
mtq.compress(model)

[meenchen]

PR for the bug fix: #522

[meenchen]

@ishan-modi, sorry for the delay. Please try the PR and see if that fixes your issue.