[bc-breaking] enable direct configuration in quantize_

[vkuzo]

summary

This PR enables passing per-workflow arguments to quantize_ directly, without wrapping them in a Callable.

Motivation: passing direct configuraton is intuintive and widely used in similar contexts across various projects. Passing configuration wrapped in a callable is IMO 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.

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
    ...

current status

The current PR migrates three user facing workflows:

• PTQ's int4_weight_only
• QAT's intx_quantization_aware_training and from_intx_quantization_aware_training

I've chosen to migrate one PTQ and two QAT workflows to prove generality of the new flow, but avoid a high LOC in this PR to make it easier to review. We will migrate the rest of the workflows in future PRs, detailed below:

• int8_dynamic_activation_int4_weight
• int8_dynamic_activation_int8_weight
• int8_dynamic_activation_int8_semi_sparse_weight
• int8_weight_only
• float8_weight_only
• float8_dynamic_activation_float8_weight
• float8_static_activation_float8_weight
• uintx_weight_only
• fpx_weight_only
• gemlite_uintx_weight_only
• callsites from the prototype folder

After a release cycle, we will delete the old callable syntax.

Test Plan:

pytest test/quantization/test_quant_api.py -s -x -k test_int4_weight_only_numerics
pytest test/quantization/test_qat.py -s -x -k test_quantize_api_standalone
pytest test/quantization/test_qat.py -s -x -k test_quantize_api_convert_path

Reviewers:

Subscribers:

Tasks:

Tags:

[vkuzo]

Stack from ghstack (oldest at bottom):

• -> [bc-breaking] enable direct configuration in quantize_ #1595

[pytorch-bot]

🔗 Helpful Links

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

• 📄 Preview Python docs built from this PR

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

❗ 1 Active SEVs

There are 1 currently active SEVs. If your PR is affected, please view them below:

• ROCM Infra failures during checkout of PyTorch

✅ No Failures

As of commit 26850da with merge base 8afd10e ():
💚 Looks good so far! There are no failures yet. 💚

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

[andrewor14]

Looks great! Mostly just minor doc nits.

[andrewor14]

remove?

[andrewor14]

once we migrate all functions to configs we won't need this check anymore right? Should we add a TODO to remove it?

[andrewor14]

I feel we can just add this to torchao/config.py without making a new core directory. No strong preference though

[andrewor14]

slightly stronger preference is I feel "core" shouldn't appear in the import, so users should be able to do this:

from torchao.config import AOBaseConfig

but we can do that by adding this to __init__.py

[andrewor14]

do we need this change?

[andrewor14]

need to update the docstring here in the previous line

[andrewor14]

I can't comment up there but need to update the docstring in L282

[andrewor14]

should we add a paragraph here or under quantize_ about how this is related to register_quantize_module_handler, so users who wish to add their own configs know how to do it?

[andrewor14]

alternatively we can pass in a lambda, then we don't need to add extra_args or pass in config:

replace_fn = lambda mod: handler(mod, config)

seems simpler

[vkuzo]

I'm really not a fan of passing callables around, it's easy when the callable is simple but easy for future people to tack ugly stuff on and increase complexity. Non-callable args make it harder to make the code ugly in the future.

[andrewor14]

oh sorry, I meant pass in replace_fn instead of handler, like:

replace_fn = lambda mod: handler(mod, config)
_replace_with_custom_fn_if_matches_filter(
            model,
            replace_fn,
            _is_linear if filter_fn is None else filter_fn,
            device=device,
)

either way you're passing a callable

[andrewor14]

nit: add some docstrings here to explain how this is related to quantize_ and AOBaseConfig?