Avoid normalization layers in HF's quantization_config

[lisjin]

This is a follow-up to #3015. I found that setting model.config.quantization_config with a "_default" key will quantize not just linear layer weights, but all weights, including normalization layers.

The culprit is still this line from TorchAoHfQuantizer.create_quantized_param:

quantize_(module, c, filter_fn=lambda x, fqn: True)

In order to avoid quantizing non-linear weights, I had to manually add all module names for normalization layers to modules_to_not_convert, which is an argument to TorchAoConfig.

@andrewor14 Do you know if this is the intended behavior? I don't see why they aren't using the default filter_fn for quantize_.

[pytorch-bot]

🔗 Helpful Links

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

• 📄 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 1572b34 with merge base 8525185 ():
💚 Looks good so far! There are no failures yet. 💚

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

[andrewor14]

I think that's the intended behavior but just wanted to cc @jerryzh168 to confirm.

There are two ways to skip quantizing certain layers today:

1. ModuleFqnToConfig, more flexible but also a bit more verbose
2. modules_to_not_convert, just specify a list of module names to not convert

However, neither of these accept regex today so you'll have to specify all the modules you want to skip manually, which may be a bit brittle

[jerryzh168]

oh I meant just using some specific model defined in transformers, and use the public APIs, just making sure, would the tests work for existing models in transformers?

[lisjin]

Yes, existing transformers models also inherit from PreTrainedModel. AutoModelForCausalLM.from_pretrained(..., quantization_config=quantization_config) can be tested in the same way

[jerryzh168]

lg, see comments inline, just want to make sure that the APIs we are using are applicable to all huggingface models

[metascroy]

I think that's the intended behavior but just wanted to cc @jerryzh168 to confirm.

This is not at all what I would expect for how default_ to behave since the "default" in quantize_ is to only apply to linear.

@jerryzh168 this behavior means we're quantizing more than we expect in our hummingbird models?

[jerryzh168]

This is not at all what I would expect for how default_ to behave since the "default" in quantize_ is to only apply to linear.

yeah this is the intended behavior since now user can do the filtering manually when they generate ModuleFqnToConfig

it will still fail if the module does not have .weight though, due to the implementation details of quantize_

for the models we quantize, I think it's still just linears that are quantized, otherwise these are likely to fail:

ao/torchao/quantization/quant_api.py

Lines 965 to 977 in 3d48174

	new_weight, new_bias = _int8_dynamic_activation_intx_weight_quantize_tensor(
	module.weight,
	module.bias,
	config,
	custom_scale=custom_scale,
	custom_zero_point=custom_zero_point,
	)
	module.weight = torch.nn.Parameter(new_weight, requires_grad=False)
	if new_bias is None:
	module.bias = None
	if isinstance(module, nn.Linear):
	module.extra_repr = types.MethodType(_linear_extra_repr, module)
	return module

although it does make sense to have better error messages when it's not what we expect