Is it possible to export a QAT model in AWQ Format?

[ambroser53]

I'm new to torchao and QAT but I'm pretty comfortable with PTQ techniques like AWQ and GPTQ. My deployment pipeline requires AWQ format (safetensors supported by autoawq or gptqmodel's new AWQ integration, needs to be in uint32 like Int4PackingFormat.PLAIN_INT32). I want to train a model with Int4WeightOnlyConfig and but it's confusing as to how I convert the final model into AWQ format, as AWQ format is supported but is this only for PTQ? Unless I'm missing something, you can save to roughly the same format (PLAIN_INT32 but only on xpu?) AND have AWQ support but there's no way to export to this format? If wrap my Int4WeightOnlyConfig in an AWQConfig, will it be trainable or only able to calibrate? Could I otherwise use something along the lines to the converter defined in this project?

[supriyar]

cc @andrewor14

[andrewor14]

Hi @ambroser53, looks like today for QAT we only support Int4PackingFormat.PLAIN and PRESHUFFLED:

ao/torchao/quantization/qat/fake_quantize_config.py

Lines 389 to 392 in eb2a8fc

	supported_packing_formats = [
	Int4PackingFormat.PLAIN,
	Int4PackingFormat.PRESHUFFLED,
	]

But I don't see why it wouldn't work for PLAIN_INT32, maybe just need to add it to that list. Actually I think it may already work today if you just do the following, since PLAIN and PLAIN_INT32 should have similar (if not the same) PTQ numerics, just in a different layout. Also cc @jerryzh168 to confirm:

quantize_(
    model,
    QATConfig(Int4WeightOnlyConfig(), step="prepare"),  # defaults to PLAIN
)

train(model)

quantize_(
    model,
    QATConfig(
        Int4WeightOnlyConfig(int4_packing_format=Int4PackingFormat.PLAIN_INT32),
        step="convert"
    )
)

[ambroser53]

Hi thanks @andrewor14 this works, the initial issue is that PLAIN_INT32 crashes on an assert if it's not got device type "xpu", commenting that out makes it work fine with cuda. The real issue here now is that the AWQ format requires the zero_point tensors to be int32 packed int4s in the same way as the main tensors. right now torchao PLAIN_INT32 put the zero_point in int8. How would one go about doing that within Int4PlainInt32Tensor.from_hp?

Just to illustrate, here is the qzeros from my AWQ representation a layer in a model:

here are the same qzeros from the torchao version of the exact same layer in the same model:

I'm also assuming there's no way to save Int4PlainInt32Tensor with safe_serialization? It's a relatively simple fix if not but if there's something I'm missing it'll save me writing a converter.