[quantization] feat: support aobaseconfig classes in TorchAOConfig

[sayakpaul]

What does this PR do?

The AOBaseConfig classes introduced in torchao (since 0.9.0) are more flexible. Similar to Transformers, this PR adds support for allowing them in Diffusers:

from diffusers import DiffusionPipeline, TorchAoConfig, PipelineQuantizationConfig
from torchao.quantization import Int8WeightOnlyConfig
import torch 

ckpt_id = "black-forest-labs/FLUX.1-dev"
pipeline_quant_config = PipelineQuantizationConfig(
    quant_mapping={"transformer": TorchAoConfig(Int8WeightOnlyConfig())}
)
pipe = DiffusionPipeline.from_pretrained(
    ckpt_id, quantization_config=pipeline_quant_config, torch_dtype=torch.bfloat16
).to("cuda")
_ = pipe("dog", num_inference_steps=2)

@stevhliu, would it be possible for you to propagate the relevant changes to our TorchAO docs from Transformers? Can happen in a later PR.

[HuggingFaceDocBuilderDev]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

[jerryzh168]

I think we can deprecate this one since this is less scalable than AOBaseConfig

[sayakpaul]

Yes, will do so after this PR. Meanwhile, if you could review the PR, it'd be helpful.

[SunMarc]

Thanks for updating this !

[jerryzh168]

nit: merge this with the branch in L521 to keep relevant things relevant together?

[sayakpaul]

Done!

[jerryzh168]

we should probably clean this up in the future..

[sayakpaul]

Maybe not even needed since we will deprecate this codepath after this PR is merged.

[jerryzh168]

looks good, would be good to put up a deprecation plan for torchao version support and clean up the old code for str support I think, I don't think there is any reason why people can't upgrade to the most recent version currently

[sayakpaul]

@jerryzh168 thanks for your reviews! I agree that it will make the library codebase cleaner if we started a deprecation cycle to promote the AoBaseConfig more and more. Will let @DN6 also review and then merge.

[DN6]

Think it would be cleaner to run the checks in the begining and exit/fail early

def post_init(self):
    if not isinstance(self.quant_type, str):
        if not is_torchao_version(">=", "0.9.0"):
            raise ValueError(
                f"torchao <= 0.9.0 only supports string quant_type, got {type(self.quant_type).__name__}. "
                f"Upgrade to torchao > 0.9.0 to use AOBaseConfig."
            )
        
        from torchao.quantization.quant_api import AOBaseConfig
        if not isinstance(self.quant_type, AOBaseConfig):
            raise TypeError(
                f"quant_type must be a string or AOBaseConfig instance, got {type(self.quant_type).__name__}"
            )
        return
    
    TORCHAO_QUANT_TYPE_METHODS = self._get_torchao_quant_type_to_method()
    
    if self.quant_type not in TORCHAO_QUANT_TYPE_METHODS:
        # remaining str type validation

[sayakpaul]

if not isinstance(self.quant_type, str):
        if not is_torchao_version(">=", "0.9.0"):
            raise ValueError(
                f"torchao <= 0.9.0 only supports string quant_type, got {type(self.quant_type).__name__}. "
                f"Upgrade to torchao > 0.9.0 to use AOBaseConfig."
            )

This will be a breaking change. I think we should introduce a deprecation cycle before enforcing this.

[DN6]

Why breaking? The error is only raised if quant_type is not a string and torchao<=0.9.0? Can change the second check to if is_torchao_version("<=", "0.9.0"): if that is more clear?

[sayakpaul]

Fixed.

[DN6]

Thanks 👍🏽