PARQ quantizer support for torchao's weight-only configs

[lisjin]

This is the first step in supporting torchao.quantize_ for PARQ trained models. I target only Int4WeightOnlyConfig and IntxWeightOnlyConfig for now since PARQ does not have activation quantization.

Instead of converting the state (e.g., scale, zero point) from PARQ's existing quantizers to torchao format, I decided to create a new quantizer UnifTorchaoQuantizer. This quantizer calls torchao's quantization primitives choose_qparams_affine, quantize_affine, dequantize_affine to ensure parity between the two QAT methods.

@metascroy It would be great if you could check the correctness of how the quantizer in TestUnifTorchaoQuantizer.test_intx_weight_only is initialized. I'm not sure if I missed any subtleties with int8.

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/ao/2091

• 📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

✅ No Failures

As of commit 9e70d6d with merge base cdced21 ():
💚 Looks good so far! There are no failures yet. 💚

This comment was automatically generated by Dr. CI and updates every 15 minutes.

[metascroy]

@lisjin can you give a little code snippet of our QAT prepare/convert would work for this API?

I'm having trouble following. Here are some example code snippets from other APIs: https://fb.workplace.com/groups/pytorch.edge2.team/permalink/1186139489308568/

[lisjin]

Added parametrized test cases that loop over (int2, int3, int4, int8) data types.

[lisjin]

@metascroy Here is a new optimizer-level method within PARQ that (1) loads the latent full-precision weights and (2) recursively quantizes all modules using torchao.quantize_. Just to double-check, will it also work on modules besides torch.nn.Linear?

[metascroy]

Usually to do non-linear, we need a third filter_fn argument, e.g., for embedding:

quantize_(model, IntxWeightOnlyConfig(torch.int4), lambda m, fqn: isinstance(m, torch.nn.Embedding))

By default the filter function in quantize_ is torch.nn.Linear, which is why you don't usually see it specified.

[andrewor14]

Hi @lisjin, do you mind adding a code snippet on the main README on what the end-to-end flow would look like? My understanding is you can just replace the LSBQuantizer there with your new UnifTorchaoQuantizer. Then what happens after training? Do we call quantize_(model, Int4WeightOnlyConfig) as before? Would be good to clarify

[lisjin]

@andrewor14 Thanks for the feedback—I removed config from UnifTorchaoQuantizer. In the README, I've also added a side-by-side comparison of PARQ vs. torchao prepare and convert steps. After PARQ training, we call optimizer.torchao_quantize_(model, config). Let me know if there's anything missing.

[andrewor14]

Looks great, thanks @lisjin! The README is very clear.

One thing I want to discuss is whether we can just use a new PARQConfig instead so the PARQ flow looks more like the existing torchao QAT flow. This is current convert flow in the PR now:

config = IntXWeightOnlyConfig(weight_dtype=torch.int4, granularity=PerGroup(32))
optimizer.torchao_quantize_(model, config)

What do you think about something like this instead?

inner_config = IntXWeightOnlyConfig(weight_dtype=torch.int4, granularity=PerGroup(32))
parq_config = PARQConfig(optimizer, inner_config)
quantize_(model, parq_config)

Also curious if @metascroy has any thoughts on this

[lisjin]

@andrewor14 I think the quantize_(model, PARQConfig(optimizer, inner_config)) interface looks cleaner, but I see a potential issue with applying this inner_quantize_ function. PARQ uses param_set to filter which modules to quantize. However, quantize_ has its own recursive function that quantizes parameters based off filter_fn mapping nn.Module -> bool. I'm not sure how to reconcile the two.

[andrewor14]

I see, yeah doesn't seem trivial to reconcile these two. Maybe this is OK then

[andrewor14]

Another idea is:

quantize_(model, config, optimizer.filter_fn)

where optimizer.filter_fn checks against the param_set. Then you don't need PARQConfig. What do you think?

[andrewor14]

By the way this call itself is also recursive, so I feel this will end up quantizing the same submodules multiple times? Maybe we should express this as part of a filter_fn, something like:

def _filter_fn(module: torch.nn.Module) -> bool:
    for param in module.parameters(recurse=False):
        if param.data_ptr() in param_set:
            return True
    return False

quantize_(model, config, _filter_fn)

[lisjin]

Good point! I'll add this change in

[andrewor14]

Looks good to me other than the recursion comment. @metascroy any other thoughts?

[metascroy]

Looks good to me! Thanks @lisjin!

Can we add an end-to-end test_intx_weight_only_e2e for intx (with various x-values), similar to test_int4_weight_only_e2e?

[metascroy]

Nit: pass mapping type directly?

[metascroy]

Just make torch.int8 the default in init?