Skip to content

enforce AOBaseConfig type in quantize_'s config argument #1861

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 88 commits into from
Mar 12, 2025
Merged
Changes from all commits
Commits
Show all changes
88 commits
Select commit Hold shift + click to select a range
73ea701
Update
vkuzo Mar 7, 2025
4f2c69d
Update
vkuzo Mar 7, 2025
75fb698
Update
vkuzo Mar 7, 2025
5796efc
Update
vkuzo Mar 7, 2025
3b016e7
Update
vkuzo Mar 7, 2025
11bce08
Update
vkuzo Mar 7, 2025
3a91a06
Update
vkuzo Mar 7, 2025
91c350f
Update
vkuzo Mar 7, 2025
125546b
Update
vkuzo Mar 7, 2025
130a52e
Update
vkuzo Mar 7, 2025
70e0e53
Update
vkuzo Mar 7, 2025
049d750
Update
vkuzo Mar 7, 2025
7b8fb8d
Update
vkuzo Mar 7, 2025
74965d3
Update
vkuzo Mar 7, 2025
1ea0e2f
Update
vkuzo Mar 7, 2025
03ea2e4
Update
vkuzo Mar 7, 2025
9d7210b
Update
vkuzo Mar 7, 2025
cf3ad33
Update
vkuzo Mar 7, 2025
19ac99d
Update
vkuzo Mar 7, 2025
5deed22
Update
vkuzo Mar 7, 2025
2a1f7b2
Update
vkuzo Mar 7, 2025
5fa0e27
Update
vkuzo Mar 7, 2025
160cc29
Update
vkuzo Mar 7, 2025
7e40b15
Update
vkuzo Mar 7, 2025
b0a32c3
Update
vkuzo Mar 7, 2025
0ecb02d
Update
vkuzo Mar 7, 2025
7cb810c
Update
vkuzo Mar 7, 2025
eb567cd
Update
vkuzo Mar 7, 2025
cd97b30
Update
vkuzo Mar 7, 2025
de38b6e
Update
vkuzo Mar 7, 2025
bfba1d9
Update
vkuzo Mar 7, 2025
9ac2334
Update
vkuzo Mar 7, 2025
c042458
Update
vkuzo Mar 7, 2025
6f3d127
Update
vkuzo Mar 8, 2025
cda5d18
Update
vkuzo Mar 8, 2025
96d74a3
Update
vkuzo Mar 8, 2025
c83c029
Update
vkuzo Mar 8, 2025
6f1c92d
Update
vkuzo Mar 8, 2025
95be23e
Update
vkuzo Mar 8, 2025
0776629
Update
vkuzo Mar 8, 2025
52a31a8
Update
vkuzo Mar 8, 2025
fdb292e
Update
vkuzo Mar 8, 2025
706ff1f
Update
vkuzo Mar 8, 2025
ac2314e
Update
vkuzo Mar 8, 2025
8002c39
Update
vkuzo Mar 8, 2025
a4dfaa1
Update
vkuzo Mar 8, 2025
ecdab3b
Update
vkuzo Mar 8, 2025
7679470
Update
vkuzo Mar 8, 2025
0506d32
Update
vkuzo Mar 8, 2025
237a72a
Update
vkuzo Mar 8, 2025
7183e83
Update
vkuzo Mar 8, 2025
10d0dff
Update
vkuzo Mar 8, 2025
be63b3c
Update
vkuzo Mar 8, 2025
fc44c79
Update
vkuzo Mar 8, 2025
2ec7827
Update
vkuzo Mar 8, 2025
3f10bc5
Update
vkuzo Mar 8, 2025
fa8c0f1
Update
vkuzo Mar 8, 2025
e8ee9a1
Update
vkuzo Mar 8, 2025
025bd67
Update
vkuzo Mar 8, 2025
50a7f9f
Update
vkuzo Mar 8, 2025
8c8d7e4
Update
vkuzo Mar 8, 2025
5b62372
Update
vkuzo Mar 8, 2025
89f6763
Update
vkuzo Mar 8, 2025
32a93b4
Update
vkuzo Mar 10, 2025
34ce5f4
Update
vkuzo Mar 10, 2025
c038451
Update
vkuzo Mar 10, 2025
c3815dc
Update
vkuzo Mar 10, 2025
9280bc9
Update
vkuzo Mar 10, 2025
fd5cfc0
Update
vkuzo Mar 10, 2025
4f195b6
Update
vkuzo Mar 10, 2025
0d388c8
Update
vkuzo Mar 10, 2025
42f323f
Update
vkuzo Mar 10, 2025
0a9fc76
Update
vkuzo Mar 10, 2025
db3e3d3
Update
vkuzo Mar 12, 2025
5f742de
Update
vkuzo Mar 12, 2025
6500a26
Update
vkuzo Mar 12, 2025
8d16b60
Update
vkuzo Mar 12, 2025
3d32932
Update
vkuzo Mar 12, 2025
bdf1aea
Update
vkuzo Mar 12, 2025
c2577a2
Update
vkuzo Mar 12, 2025
902f8fa
Update
vkuzo Mar 12, 2025
a7369a7
Update
vkuzo Mar 12, 2025
82b6281
Update
vkuzo Mar 12, 2025
34bfd94
Update
vkuzo Mar 12, 2025
32d29f0
Update
vkuzo Mar 12, 2025
b63f186
Update
vkuzo Mar 12, 2025
4abb545
Update
vkuzo Mar 12, 2025
2c6104a
Update
vkuzo Mar 12, 2025
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
17 changes: 3 additions & 14 deletions torchao/quantization/quant_api.py
Original file line number Diff line number Diff line change
Expand Up @@ -486,7 +486,7 @@ def insert_subclass(lin):

def quantize_(
model: torch.nn.Module,
config: Union[AOBaseConfig, Callable[[torch.nn.Module], torch.nn.Module]],
config: AOBaseConfig,
filter_fn: Optional[Callable[[torch.nn.Module, str], bool]] = None,
set_inductor_config: Optional[bool] = None,
device: Optional[torch.types.Device] = None,
Expand All @@ -495,7 +495,7 @@ def quantize_(

Args:
model (torch.nn.Module): input model
config (Union[AOBaseConfig, Callable[[torch.nn.Module], torch.nn.Module]]): either (1) a workflow configuration object or (2) a function that applies tensor subclass conversion to the weight of a module and return the module (e.g. convert the weight tensor of linear to affine quantized tensor). Note: (2) will be deleted in a future release.
config (AOBaseConfig): a workflow configuration object.
filter_fn (Optional[Callable[[torch.nn.Module, str], bool]]): function that takes a nn.Module instance and fully qualified name of the module, returns True if we want to run `config` on
the weight of the module
set_inductor_config (bool, optional): Whether to automatically use recommended inductor config settings (defaults to None)
Expand Down Expand Up @@ -546,21 +546,10 @@ def quantize_(
)

else:
# old behavior, keep to avoid breaking BC
warnings.warn(
raise AssertionError(
"""Passing a generic Callable to `quantize_` is no longer recommended and will be deprecated at a later release. Please see https://github.com/pytorch/ao/issues/1690 for instructions on how to pass in workflow configuration instead."""
)

# make the variable name make sense
apply_tensor_subclass = config

_replace_with_custom_fn_if_matches_filter(
model,
apply_tensor_subclass,
_is_linear if filter_fn is None else filter_fn,
device=device,
)


def _int8_asymm_per_token_quant(x: torch.Tensor) -> torch.Tensor:
"""This is defined here instead of local function to support serialization"""
Expand Down
Loading