How to Disable Concat Quantization Nodes, and how to Share the Same Quantization Parameters Across Multiple Layers

[piantou]

(1) As shown in the figure below, I want to disable quantization for the specified Concat node. How should I write the corresponding code or configuration rules?

(2) Taking the quantization tool TinyQ as an example, it can explicitly bind multiple operators through surgeon.connect to share the same set of quantization parameters (scale/zero_point). During operator fusion, only one set of QDQ nodes is retained, which reduces computation and accelerates inference.
Taking the Add operation as an example (shown in the figure below), to enable operator fusion and acceleration in the subsequent inference stage, it is required that the quantization parameters (scale) of all its input branches must be exactly the same. How can this be implemented in AIMET?

Looking forward to your reply to clarify my questions. Thank you.

[michaelgtuttle]

Hi piantou, are you working in aimet-onnx or aimet-torch here?

[piantou]

Hi piantou, are you working in aimet-onnx or aimet-torch here?

aimet_torch dear

[michaelgtuttle]

Hi piantou, the most straightforward way to tie the output quantizers of two layers in aimet-torch is to assign the same quantizer object to each layer, for example:

# Output encodings of layer_1 will always match output encodings of layer_2
sim.model.layer_1.output_quantizers[0] = sim.model.layer_2.output_quantizers[0]

To manually disable a layer's output quantizer you can set it to None, e.g.,

# Disables output quantizer of model.silu_1
sim.model.silu_1.output_quantizers[0] = None

For your specific use-case of tying the output quantizers of a layer to the inputs, we have a util which will do this graph traversal for you and tie the quantizers automatically:

from aimet_torch.experimental import quantsim_utils
from aimet_torch.nn.modules import custom

# To tie one specific layer's output quantizer to inputs:
quantsim_utils.propagate_output_encodings(sim, sim.model.add)

# To tie for all layers of type custom.Add or custom.Concat (change as necessary for your model)
quantsim_utils.propagate_output_encodings(
    sim, lambda layer: isinstance(layer, (custom.Add, custom.Concat))
)

Note that for Concat layers, tying the output quantizer to input quantizers is exactly the same as disabling the input quantizers from an accuracy standpoint, but is generally easier for downstream tools to ingest.