Skip to content

Make QuantAbsorption's input-side absorption opt-in #24780

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
Open

Make QuantAbsorption's input-side absorption opt-in #24780

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
23 changes: 18 additions & 5 deletions backends/transforms/fused_quant/optimization_passes/quant_absorption.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -567,11 +567,16 @@ class QuantAbsorptionPass(ExportedProgramPassBase):
feeds into a dequant→quant pair, this pass composes the producer, dequantize,
and downstream quantize qparams into the fused op's output qparams. When the
producer and dequantize qparams match, this reduces to the downstream target.
This is always on because it is never less accurate: the value still ends up
on the downstream quantize's grid, rounded once instead of twice, so each
element is at least as close to the float result. It is not bit-exact with
the q/dq graph, though, and a QAT model was trained with both roundings.

On the input side, when a dequant→quant pair feeds any fused_quant input with
per-tensor qparams, this pass composes the dequantize, quantize, and consumer
input qparams, then bypasses the pair. This handles graph inputs, which have
no fused producer into which the quantization can be absorbed.
no fused producer into which the quantization can be absorbed. It is opt-in
(``absorb_into_inputs``) because it is not always more accurate: see below.

It also absorbs a *bare* downstream quantize into a fused op whose output is
currently unquantized (null out qparams), turning that op into a directly
Expand All @@ -581,10 +586,14 @@ class QuantAbsorptionPass(ExportedProgramPassBase):
Args:
absorb_with_fork: When True, also absorb dequant→quant pairs where
the dequant has additional users (a "forked" dequant). The dequant
is kept for the other users but its qparams are updated to match
the new output qparams. This changes (not necessarily worsens) the
numerics of the float path through the dequant, but should improve
numerics on the int path since it removes a redundant dq-q pair.
is kept for the other users, with the downstream quantize's qparams.
Off by default: those float users then see values on the quantize's
grid instead of the producer's, which can be coarser.
absorb_into_inputs: When True, also absorb a dequant→quant pair into
the input qparams of the fused op it feeds. Off by default: the
consumer then reads values on the dequantize's grid instead of the
quantize's. Dropping the second rounding is more accurate on
average, but not for every value.
symmetric_quant_input_ops: Ops whose inputs must stay symmetrically
quantized (zero point 0) because their backend kernel ignores input
zero points. An input absorption that would give them a nonzero input
Expand All @@ -594,10 +603,12 @@ class QuantAbsorptionPass(ExportedProgramPassBase):
def __init__(
self,
absorb_with_fork: bool = False,
absorb_into_inputs: bool = False,
symmetric_quant_input_ops: frozenset[EdgeOpOverload] = frozenset(),
) -> None:
super().__init__()
self._absorb_with_fork = absorb_with_fork
self._absorb_into_inputs = absorb_into_inputs
self._symmetric_quant_input_ops = symmetric_quant_input_ops

def call(self, exported_program: ExportedProgram) -> ExportedProgramPassResult:
Expand All @@ -615,6 +626,8 @@ def call(self, exported_program: ExportedProgram) -> ExportedProgramPassResult:
modified |= _try_absorb_quant_into_unquantized_out(
exported_program, node, graph
)
if not self._absorb_into_inputs:
continue
assert isinstance(node.target, EdgeOpOverload)
input_names, _ = split_fused_arg_names(node.target)
keep_symmetric = node.target in self._symmetric_quant_input_ops
Expand Down
33 changes: 21 additions & 12 deletions backends/transforms/fused_quant/test/test_quant_absorption.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1416,10 +1416,17 @@ def _build_dequant_quant_add(self) -> tuple[ExportedProgram, torch.Tensor]:
builder.output([add])
return builder.get_program(), inp

def test_does_not_absorb_into_consumer_input_by_default(self) -> None:
ep, _ = self._build_dequant_quant_pool()

result = QuantAbsorptionPass()(ep)

self.assertFalse(result.modified)

def test_absorbs_dequant_quant_into_fused_consumer_input(self) -> None:
ep, inp = self._build_dequant_quant_pool()

result = QuantAbsorptionPass()(ep)
result = QuantAbsorptionPass(absorb_into_inputs=True)(ep)

self.assertTrue(result.modified)
(actual,) = result.exported_program.module()(inp)
Expand Down Expand Up @@ -1458,7 +1465,7 @@ def test_absorbs_dequant_quant_into_fused_consumer_input(self) -> None:
def test_absorbs_consumer_input_through_view_copy(self) -> None:
ep, inp = self._build_dequant_quant_pool(input_passthrough="view")

result = QuantAbsorptionPass()(ep)
result = QuantAbsorptionPass(absorb_into_inputs=True)(ep)

self.assertTrue(result.modified)
(actual,) = result.exported_program.module()(inp)
Expand Down Expand Up @@ -1501,7 +1508,7 @@ def test_absorbs_consumer_input_through_view_copy(self) -> None:
def test_absorbs_consumer_input_through_permute_copy(self) -> None:
ep, inp = self._build_dequant_quant_pool(input_passthrough="permute")

result = QuantAbsorptionPass()(ep)
result = QuantAbsorptionPass(absorb_into_inputs=True)(ep)

self.assertTrue(result.modified)
(actual,) = result.exported_program.module()(inp)
Expand Down Expand Up @@ -1546,7 +1553,7 @@ def test_does_not_absorb_through_shared_input_passthrough(self) -> None:
input_passthrough="view", expose_passthrough=True
)

result = QuantAbsorptionPass()(ep)
result = QuantAbsorptionPass(absorb_into_inputs=True)(ep)

self.assertFalse(result.modified)
self.assertEqual(
Expand All @@ -1562,7 +1569,7 @@ def test_does_not_absorb_through_shared_input_passthrough(self) -> None:
def test_absorbs_all_quantized_inputs_of_binary_op(self) -> None:
ep, inp = self._build_dequant_quant_add()

result = QuantAbsorptionPass()(ep)
result = QuantAbsorptionPass(absorb_into_inputs=True)(ep)

self.assertTrue(result.modified)
(actual,) = result.exported_program.module()(inp)
Expand Down Expand Up @@ -1628,7 +1635,7 @@ def test_composes_mismatched_consumer_input_qparams(self) -> None:
)
expected = ep.module()(inp)

result = QuantAbsorptionPass()(ep)
result = QuantAbsorptionPass(absorb_into_inputs=True)(ep)

self.assertTrue(result.modified)
actual = result.exported_program.module()(inp)
Expand Down Expand Up @@ -1668,7 +1675,7 @@ def test_does_not_absorb_unrepresentable_consumer_input_qparams(self) -> None:
consumer_input_zero_point=-4,
)

result = QuantAbsorptionPass()(ep)
result = QuantAbsorptionPass(absorb_into_inputs=True)(ep)

self.assertFalse(result.modified)

Expand All @@ -1683,9 +1690,10 @@ def test_symmetric_quant_input_op_keeps_requantize_for_nonzero_zero_point(
)

result = QuantAbsorptionPass(
absorb_into_inputs=True,
symmetric_quant_input_ops=frozenset(
{exir_ops.edge.fused_quant.avg_pool2d_channels_last.default}
)
),
)(ep)

self.assertFalse(result.modified)
Expand All @@ -1696,9 +1704,10 @@ def test_symmetric_quant_input_op_absorbs_when_zero_point_stays_zero(
ep, inp = self._build_dequant_quant_pool()

result = QuantAbsorptionPass(
absorb_into_inputs=True,
symmetric_quant_input_ops=frozenset(
{exir_ops.edge.fused_quant.avg_pool2d_channels_last.default}
)
),
)(ep)

self.assertTrue(result.modified)
Expand Down Expand Up @@ -1746,7 +1755,7 @@ def test_does_not_absorb_per_channel_consumer_input(self) -> None:
builder.output([relu])
ep = builder.get_program()

result = QuantAbsorptionPass()(ep)
result = QuantAbsorptionPass(absorb_into_inputs=True)(ep)

self.assertFalse(result.modified)
self.assertEqual(
Expand All @@ -1762,7 +1771,7 @@ def test_does_not_absorb_per_channel_consumer_input(self) -> None:
def test_does_not_absorb_shared_quantize(self) -> None:
ep, _ = self._build_dequant_quant_pool(expose_quant=True)

result = QuantAbsorptionPass()(ep)
result = QuantAbsorptionPass(absorb_into_inputs=True)(ep)

self.assertFalse(result.modified)
self.assertEqual(
Expand All @@ -1779,7 +1788,7 @@ def test_absorbs_quantize_and_preserves_shared_dequantize(self) -> None:
ep, inp = self._build_dequant_quant_pool(expose_dequant=True)
_, expected_dequant = ep.module()(inp)

result = QuantAbsorptionPass()(ep)
result = QuantAbsorptionPass(absorb_into_inputs=True)(ep)

self.assertTrue(result.modified)
actual_pool, actual_dequant = result.exported_program.module()(inp)
Expand Down
Loading