migration of quantize_ workflow configuration from callables to configs

[vkuzo]

summary

This issue tracks the migration of quantize_ per-workflow configuration from Callables to configs..

Motivation: passing direct configuraton is intuintive and widely used in similar contexts across various projects. Passing configuration wrapped in a callable is not intuitive, hard to understand and debug, and we have evidence that it pushes a portion of users from building on top of torchao.

We will keep the old callable syntax supported by quantize_ for one release cycle, and delete it afterwards. We will keep the old names as aliases for new names going forward (example: int4_weight_only as an alias of Int4WeightOnlyConfig) to keep existing callsites working without changes.

impact on API users

If you are just using the torchao quantize_ API as specified in the README, this is not BC-breaking. For example, the following syntax will keep working.

quantize_(model, int8_weight_only())

Note that the type of the object created by int8_weight_only() will change from a Callable to a config. You have the option to migrate to the explicit config creation, as follows:

quantize_(model, Int8WeightOnlyConfig())

user facing API changes

signature of quantize_

#
# before
#
def quantize(
    model: torch.nn.Module,
    apply_tensor_subclass: Callable[[torch.nn.Module], torch.nn.Module],
    ...,
): ...

#
# after - intermediate state, support both old and new for one release
#
def quantize(
    model: torch.nn.Module,
    config: Union[AOBaseConfig, Callable[[torch.nn.Module], torch.nn.Module]],
    ...,
): ...

#
# after - long term state
#
def quantize(
    model: torch.nn.Module,
    config: AOBaseConfig,
    ...,
): ...

usage example

An example for int4_weight_only

#
# before
#
quantize_(m, int4_weight_only(group_size=32))

#
# after, with new user facing names
#
quantize_(m, Int4WeightOnlyConfig(group_size=32))

#
# AND, after, with BC names
#
quantize_(m, int4_weight_only(group_size=32))

developer facing changes

See the PR details for examples, but they can be summarized as:

#
# old
#

# quantize_ calls the instance of calling this function on each module of the model
def int4_weight_only(group_size: int, ...) -> Callable:

    def new_callable(weight: torch.Tensor):
        # configuration is captured here via local variables
        ...
        
    # return type is a Callable
    return _get_linear_subclass_inserter(new_callable)

#
# new
#

# config base class
class AOBaseConfig(abc.ABC):
    pass

# user facing configuration of a workflow
@dataclass
class Int4WeightOnlyConfig(AOBaseConfig):
    group_size: int = 128
    ...

# not user facing transform of a module according to a worfklow's configuration
@register_quantize_module_handler(Int4WeightOnlyConfig)
def _int4_weight_only_transform(
    module: torch.nn.Module, 
    config: Int4WeightOnlyConfig,
) -> torch.nn.Module:
    # map to AQT, not user facing
    ...

migration status

workflow configuration

• int4_weight_only - [bc-breaking] enable direct configuration in quantize_ #1595
• qat: intx_quantization_aware_training - [bc-breaking] enable direct configuration in quantize_ #1595
• qat: from_intx_quantization_aware_training - [bc-breaking] enable direct configuration in quantize_ #1595
• float8_dynamic_activation_float8_weight - config migration: float8* #1694
• float8_static_activation_float8_weight - config migration: float8* #1694
• float8_weight_only - config migration: float8* #1694
• fpx_weight_only - config migration: fpx, gemlite, uintx #1697
• gemlite_uintx_weight_only - config migration: fpx, gemlite, uintx #1697
• int4_dynamic_activation_int4_weight - config migration: int* #1696
• int8_dynamic_activation_int4_weight - config migration: int* #1696
• int8_dynamic_activation_int8_semi_sparse_weight - marked for deprecation, confirmed with @jcaip we can delete this
• int8_dynamic_activation_int8_weight - config migration: int* #1696
• int8_weight_only - config migration: int* #1696
• uintx_weight_only - config migration: fpx, gemlite, uintx #1697

tutorials (replace with new registration API)

• calibration_flow/static_quant.py - migrate static quant tutorials to direct configuration #1710
• calibration_flow/gptq_like.py - migrate static quant tutorials to direct configuration #1710
• calibration_flow/awq_like.py - migrate static quant tutorials to direct configuration #1710

replace docblocks and public facing descriptions with new names

• README.md - update torchao READMEs with new configuration APIs #1711
• QAT README.md - update torchao READMEs with new configuration APIs #1711
• quantization README.md - update torchao READMEs with new configuration APIs #1711

verify partner integrations still work

• HF callsite: https://github.com/huggingface/transformers/blob/1feebb5b4150882deabddd190a541f336f3be817/src/transformers/quantizers/quantizer_torchao.py#L199
• SGLANG callsite: https://github.com/sgl-project/sglang/blob/2f47d710ae9cb1bdbbe0fe2392a0634827d257b3/python/sglang/srt/layers/torchao_utils.py#L39
• Diffusers callsite: https://github.com/huggingface/diffusers/blob/7fb481f840b5d73982cafd1affe89f21a5c0b20b/src/diffusers/quantizers/torchao/torchao_quantizer.py#L234

confirmed two out of three here: vkuzo/pytorch_scripts#28

delete old path (one version after migration)

• not done, upcoming in version TBD