Skip to content

Insert the .default overload of the affine quant ops in convert (#5101) #5101

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.

Sign up for GitHub

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

Open
mcremon-meta wants to merge 1 commit into
base: main
Choose a base branch
from
+125 −22
Open

Insert the .default overload of the affine quant ops in convert (#5101) #5101

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
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
34 changes: 17 additions & 17 deletions test/quantization/pt2e/test_quantize_pt2e.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -4218,12 +4218,12 @@ def forward(self, x):
return self.linear(x)

node_occurrence = {
torch.ops.torchao.quantize_affine: 1,
torch.ops.torchao.dequantize_affine: 2,
torch.ops.torchao.quantize_affine.default: 1,
torch.ops.torchao.dequantize_affine.default: 2,
}
node_list = [
torch.ops.torchao.quantize_affine,
torch.ops.torchao.dequantize_affine,
torch.ops.torchao.quantize_affine.default,
torch.ops.torchao.dequantize_affine.default,
]
example_inputs = (torch.randn(5, 128),)
self._test_quantizer(
Expand Down Expand Up @@ -4304,18 +4304,18 @@ def forward(self, x):
return self.linear(x)

node_occurrence = {
torch.ops.torchao.choose_qparams_affine: 1,
torch.ops.torchao.choose_qparams_affine.default: 1,
operator.getitem: 2,
torch.ops.torchao.quantize_affine: 1,
torch.ops.torchao.dequantize_affine: 1,
torch.ops.torchao.quantize_affine.default: 1,
torch.ops.torchao.dequantize_affine.default: 1,
torch.ops.quantized_decomposed.dequantize_per_channel.default: 1,
}
node_list = [
torch.ops.quantized_decomposed.dequantize_per_channel.default,
torch.ops.torchao.choose_qparams_affine,
torch.ops.torchao.choose_qparams_affine.default,
operator.getitem,
torch.ops.torchao.quantize_affine,
torch.ops.torchao.dequantize_affine,
torch.ops.torchao.quantize_affine.default,
torch.ops.torchao.dequantize_affine.default,
]
example_inputs = (torch.randn(5, 128),)
self._test_quantizer(
Expand Down Expand Up @@ -4396,17 +4396,17 @@ def forward(self, x):
return self.linear(x)

node_occurrence = {
torch.ops.torchao.choose_qparams_affine: 1,
torch.ops.torchao.choose_qparams_affine.default: 1,
operator.getitem: 2,
torch.ops.torchao.quantize_affine: 1,
torch.ops.torchao.dequantize_affine: 2,
torch.ops.torchao.quantize_affine.default: 1,
torch.ops.torchao.dequantize_affine.default: 2,
}
node_list = [
torch.ops.torchao.dequantize_affine,
torch.ops.torchao.choose_qparams_affine,
torch.ops.torchao.dequantize_affine.default,
torch.ops.torchao.choose_qparams_affine.default,
operator.getitem,
torch.ops.torchao.quantize_affine,
torch.ops.torchao.dequantize_affine,
torch.ops.torchao.quantize_affine.default,
torch.ops.torchao.dequantize_affine.default,
]
example_inputs = (torch.randn(5, 128),)
self._test_quantizer(
Expand Down
96 changes: 95 additions & 1 deletion test/quantization/pt2e/test_representation.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -20,8 +20,17 @@
)
from torch.testing._internal.common_utils import run_tests

from torchao.quantization.pt2e._affine_quantization import (
AffineQuantizedMinMaxObserver,
AffineQuantizedPlaceholderObserver,
)
from torchao.quantization.pt2e.observer import MappingType, PerGroup, PerToken
from torchao.quantization.pt2e.quantize_pt2e import convert_pt2e, prepare_pt2e
from torchao.quantization.pt2e.quantizer import Quantizer
from torchao.quantization.pt2e.quantizer import (
QuantizationAnnotation,
QuantizationSpec,
Quantizer,
)
from torchao.testing.pt2e._xnnpack_quantizer import (
XNNPACKQuantizer,
get_symmetric_quantization_config,
Expand Down Expand Up @@ -277,6 +286,91 @@ def forward(self, x):
output_scale_idx=2,
)

def test_dynamic_linear_4bit_groupwise(self):
"""Affine observers' q/dq nodes match the 4-bit groupwise linear rewrite."""
# The pattern is traced with groups of 8, and 32 spans the whole row.
for group_size in (8, 16, 32):
with self.subTest(group_size=group_size):
self._test_dynamic_linear_4bit_groupwise(group_size)

def _test_dynamic_linear_4bit_groupwise(self, group_size):
eps = torch.finfo(torch.float32).eps

class BackendAQuantizer(Quantizer):
def annotate(self, model: torch.fx.GraphModule) -> torch.fx.GraphModule:
for node in model.graph.find_nodes(
op="call_function", target=torch.ops.aten.linear.default
):
act_qspec = QuantizationSpec(
dtype=torch.int8,
quant_min=-128,
quant_max=127,
qscheme=None,
is_dynamic=True,
observer_or_fake_quant_ctr=AffineQuantizedPlaceholderObserver.with_args(
target_dtype=torch.int8,
mapping_type=MappingType.ASYMMETRIC,
granularity=PerToken(),
quant_min=-128,
quant_max=127,
eps=eps,
scale_dtype=torch.float32,
zero_point_dtype=torch.int32,
),
)
weight_qspec = QuantizationSpec(
dtype=torch.int8,
quant_min=-8,
quant_max=7,
qscheme=None,
is_dynamic=False,
observer_or_fake_quant_ctr=AffineQuantizedMinMaxObserver.with_args(
target_dtype=torch.int8,
mapping_type=MappingType.SYMMETRIC,
granularity=PerGroup(group_size=group_size),
quant_min=-8,
quant_max=7,
eps=eps,
scale_dtype=torch.float32,
zero_point_dtype=torch.int8,
),
)
node.meta["quantization_annotation"] = QuantizationAnnotation(
input_qspec_map={
node.args[0]: act_qspec,
node.args[1]: weight_qspec,
},
_annotated=True,
)
return model

def validate(self, model: torch.fx.GraphModule) -> None:
pass

class M(torch.nn.Module):
def __init__(self) -> None:
super().__init__()
self.linear = torch.nn.Linear(32, 16)

def forward(self, x):
return self.linear(x)

self._test_representation(
M().eval(),
(torch.randn(4, 32),),
BackendAQuantizer(),
ref_node_occurrence={
ns.call_function(torch.ops.torchao.reference_dqlinear_int4.default): 1,
ns.call_function(torch.ops.torchao.dequantize_affine.default): 0,
},
non_ref_node_occurrence={
ns.call_function(torch.ops.torchao.choose_qparams_affine.default): 1,
ns.call_function(torch.ops.torchao.quantize_affine.default): 1,
ns.call_function(torch.ops.torchao.dequantize_affine.default): 2,
},
fixed_output_tol=1e-2,
)

def test_qdq(self):
"""Test representation for quantize and dequantize op"""

Expand Down
2 changes: 2 additions & 0 deletions torchao/quantization/pt2e/constant_fold.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -180,6 +180,8 @@ def is_woq_int8_pattern(node: torch.fx.node.Node) -> bool:
torch.ops.quantized_decomposed.dequantize_per_tensor.default,
torch.ops.quantized_decomposed.dequantize_per_tensor.tensor,
torch.ops.quantized_decomposed.convert_element_type.no_fuse,
torch.ops.torchao.dequantize_affine.default,
# Observers that override convert() may still insert the OpOverloadPacket.
torch.ops.torchao.dequantize_affine,
]

Expand Down
6 changes: 3 additions & 3 deletions torchao/quantization/pt2e/observer.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1673,7 +1673,7 @@ def convert(self, model: torch.fx.GraphModule, observer_node: Node):
)
if hasattr(self, "is_dynamic") and self.is_dynamic:
choose_qparams_affine = model.graph.call_function(
torch.ops.torchao.choose_qparams_affine,
torch.ops.torchao.choose_qparams_affine.default,
(
observer_node.args[0],
self.mapping_type.name,
Expand Down Expand Up @@ -1710,7 +1710,7 @@ def convert(self, model: torch.fx.GraphModule, observer_node: Node):
)

q_node = model.graph.call_function(
torch.ops.torchao.quantize_affine,
torch.ops.torchao.quantize_affine.default,
(
observer_node.args[0],
self.block_size,
Expand All @@ -1723,7 +1723,7 @@ def convert(self, model: torch.fx.GraphModule, observer_node: Node):
{},
)
dq_node = model.graph.call_function(
torch.ops.torchao.dequantize_affine,
torch.ops.torchao.dequantize_affine.default,
(
q_node,
self.block_size,
Expand Down
2 changes: 2 additions & 0 deletions torchao/quantization/pt2e/quantize_pt2e.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -249,6 +249,8 @@ def train_loop(model, train_data):
torch.ops.quantized_decomposed.quantize_per_tensor.default,
torch.ops.quantized_decomposed.quantize_per_tensor.tensor,
torch.ops.quantized_decomposed.quantize_per_channel.default,
torch.ops.torchao.quantize_affine.default,
# Observers that override convert() may still insert the OpOverloadPacket.
torch.ops.torchao.quantize_affine,
]

Expand Down
4 changes: 4 additions & 0 deletions torchao/quantization/pt2e/quantizer/port_metadata_pass.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -33,19 +33,23 @@
torch.ops.quantized_decomposed.quantize_per_tensor.default,
torch.ops.quantized_decomposed.quantize_per_tensor.tensor,
torch.ops.quantized_decomposed.quantize_per_channel.default,
torch.ops.torchao.quantize_affine.default,
# Observers that override convert() may still insert the OpOverloadPackets.
torch.ops.torchao.quantize_affine,
]

_DEQUANTIZE_OPS = [
torch.ops.quantized_decomposed.dequantize_per_tensor.default,
torch.ops.quantized_decomposed.dequantize_per_tensor.tensor,
torch.ops.quantized_decomposed.dequantize_per_channel.default,
torch.ops.torchao.dequantize_affine.default,
torch.ops.torchao.dequantize_affine,
]

_CHOOSE_QPARAMS_OPS = [
torch.ops.quantized_decomposed.choose_qparams.tensor,
torch.ops.quantized_decomposed.choose_qparams_symmetric.tensor,
torch.ops.torchao.choose_qparams_affine.default,
torch.ops.torchao.choose_qparams_affine,
]

Expand Down
3 changes: 2 additions & 1 deletion torchao/quantization/pt2e/reference_representation_rewrite.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -360,7 +360,8 @@ def _reference_dqlinear_int4(
out_shape[-1] = out_features

if weight_scale.ndim == 1:
weight_scale = weight_scale.unsqueeze(0)
# A group that spans the whole row leaves one scale per output channel.
weight_scale = weight_scale.unsqueeze(1)

for group_idx in range(num_groups):
# Extract current group
Expand Down
Loading