qdq_graph#

Q/DQ graph policy for ONNX quantization.

Functions

`build_non_residual_input_map`	Builds a map of non-residual Add input name to the Add node name from the given graph.
`classify_partially_quantized_weighted_ops`	Ensures that the input of non-quantizable weighted nodes do not get quantized.
`classify_partition_nodes`	We should partially quantize the partition nodes with inputs outside of the partition.
`filter_quantizable_kgen_heads`	Returns the list of kgen head names if it follows a CASK partition.
`find_conv_to_layernorm_nodes`	Find LayerNormalization nodes whose input comes from a CASK (Conv) partition.
`get_concat_eliminated_tensors`	Find the input tensors and output tensor of concat that will be quantized.
`get_layer_info`	Generate a mapping of weight tensor names to their quantization configuration.
`get_layer_precision_mapping`	Generate a mapping of layer names to their quantization precision (4 bits or 8 bits) for an ONNX model.
`get_resize_scales`	Record Resize op's old scale value before converting to fp16.
`print_stat`	Collect and print stats of the quantized model.
`remove_partial_input_qdq`	Modifies the onnx model by removing QDQ nodes from the marked inputs, ex.
`should_quantize_to_8bit`	Check if layer should be quantized to 8 bits.
`validate_8bit_layers`	Validate the format of layers_8bit string.

build_non_residual_input_map(graph)#

Builds a map of non-residual Add input name to the Add node name from the given graph.

This assumes that the Add layer only has 2 inputs.

We will refer to a subgraph which has a Convolution node with a single output that is summed (element-wise) with another non-constant input-tensor as a “residual-add” subgraph, because it occurs in modern convnets that use residual connections.

Parameters:: graph (Graph) – Onnx model graph.
Returns:: Dictionary of Add node names vs their non-residual input name. List of partially-quantizable inputs with non-quantizable input info as (src, dst, input_name)
Return type:: tuple[dict[str, str], list[tuple[Node, Node, str]]]

classify_partially_quantized_weighted_ops(graph, nodes_to_exclude)#

Ensures that the input of non-quantizable weighted nodes do not get quantized.

Parameters:

graph (Graph)
nodes_to_exclude (list[str])

Return type:

list[tuple[Node, Node, str]]

classify_partition_nodes(partitions)#

We should partially quantize the partition nodes with inputs outside of the partition.

Parameters:: partitions (list[list[Node]]) – Partitions created by modelopt ptq algo.
Returns:: List of non-quantizable nodes. List of quantizable nodes. List of partially-quantizable inputs with non-quantizable input info as (src, dst, input_name)
Return type:: tuple[list[Node], list[Node], list[tuple[Node, Node, str]]]

filter_quantizable_kgen_heads(cask_fusible_partitions, kgen_partitions, quantizable_op_types, graph)#

Returns the list of kgen head names if it follows a CASK partition.

Parameters:

cask_fusible_partitions (list[list[Node]])
kgen_partitions (list[list[Node]])
quantizable_op_types (list[str])
graph (Graph)

Return type:

tuple[list[Node], list[tuple[Node, Node, str]]]

find_conv_to_layernorm_nodes(graph, cask_fusible_partitions)#

Find LayerNormalization nodes whose input comes from a CASK (Conv) partition.

When a Conv’s output feeds into a LayerNormalization, the Conv output should be quantized to enable faster INT8 kernels in TRT. This function detects such patterns and returns the LayerNormalization nodes that should be added to the quantizable nodes list so that Q/DQ pairs are inserted on their input (i.e. the Conv output).

Parameters:

graph (Graph) – ONNX model graph.
cask_fusible_partitions (list[list[Node]]) – List of CASK fusible partitions.

Returns:

List of LayerNormalization nodes that consume CASK partition outputs.

Return type:

list[Node]

get_concat_eliminated_tensors(onnx_model, nodes_to_quantize)#

Find the input tensors and output tensor of concat that will be quantized.

We can do some perf optimization for TRT.

For example, like the below pattern: (t1) q1 -> dq1 (t2) q2 -> dq2 -> concat -> q4 -> dq4 (t4) (t3) q3 -> dq3 /

In TRT, q4 will be propagated forward concat. It will be like: (t1) q1 -> dq1 -> q4 (t2) q2 -> dq2 -> q4 -> concat -> dq4 (t4) (t3) q3 -> dq3 -> q4 /

If the scaling factor of dq1 and q4 are different, it will cause the dq-q compute latency. If they are the same, then the dq-q pairs can be eliminated in TRT, and no extra dq-q compute latency. However, it will sacrifice the accuracy.

Thus, this function will collect which tensors should have the same scaling factors. For the above example, we want the scaling factor of dq1, dq2, dq3, q4 be the same. This function will return like { t1: {t1,t2,t3,t4}, t2: {t1,t2,t3,t4}, t3: {t1,t2,t3,t4}, t4: {t1,t2,t3,t4}, } This format is convenient for calibrator to assign the same scaling factor.

Returns:

set of tensors that should share the same scaling factor}

Return type:

{current tensor name

Parameters:

onnx_model (ModelProto)
nodes_to_quantize (list[str])

get_layer_info(onnx_model, nodes_to_exclude=['/lm_head'], block_size=128, quantize_axis=0, **kwargs)#

Generate a mapping of weight tensor names to their quantization configuration.

This function determines the quantization configuration (precision, block_size, axis) for each weight tensor in the ONNX model, based on the provided configuration. If mixed quantization is enabled, it uses the layer precision mapping; otherwise, it returns None.

Parameters:

onnx_model (onnx.ModelProto) – The ONNX model to analyze.
nodes_to_exclude (list[str] | None) – List of node name patterns to exclude from quantization.
**kwargs (Any) – Additional keyword arguments, such as: - enable_mixed_quant (bool): Whether to enable mixed quantization. - layers_8bit (str): Comma-separated list of layer patterns to quantize to 8 bit. - block_size (int): Default block size for quantization. - quantize_axis (int): Default quantization axis. - gather_block_size (int): Default block size for gather quantization. - gather_quantize_axis (int): Default quantization axis for gather.
block_size (int)
quantize_axis (int)
**kwargs

Returns:

A mapping from weight tensor names to their quantization configuration (with keys: precision, block_size, axis), or None if mixed quantization is not enabled.

Return type:

dict[str, dict[str, Any]] | None

get_layer_precision_mapping(onnx_model, precision_pattern_8bit=None, nodes_to_exclude=['/lm_head'], block_size=128, quantize_axis=0)#

Generate a mapping of layer names to their quantization precision (4 bits or 8 bits) for an ONNX model.

Parameters:

onnx_model (onnx.ModelProto) – The ONNX model to analyze.
precision_pattern_8bit (str, optional) – Comma-separated string of layer patterns to quantize to 8 bits. If None, a default set of patterns is used to select layers for 8 bits quantization.
nodes_to_exclude (list[str], optional) – List of node name patterns to exclude from quantization. Defaults to [r”/lm_head”].
block_size (int)
quantize_axis (int)

Returns:

A mapping from layer names to their quantization precision (e.g., {“layer_name”: “8”}).

Return type:

dict

get_resize_scales(onnx_model)#

Record Resize op’s old scale value before converting to fp16.

Because low precision scale will lead to wrong shape. For example, if 7 is resized to 6, fp32 scale should be 6/7 = 0.85714. After converting to fp16, it becomes 0.85693 but 7 * 0.85693 = 5.9985 < 6.

print_stat(graph)#

Collect and print stats of the quantized model.

Parameters:: graph (Graph)
Return type:: None

remove_partial_input_qdq(graph, no_quantize_inputs)#

Modifies the onnx model by removing QDQ nodes from the marked inputs, ex. non-residual inputs etc.

Parameters:

graph (Graph) – Onnx model graph.
no_quantize_inputs (list[tuple[Node, Node, str]]) – List non-quantizable input info as (src, dst, input_name)

Return type:

None

should_quantize_to_8bit(layer_name, layers_8bit)#

Check if layer should be quantized to 8 bits.

The layers_8bit list contains ONNX node names like ‘/model/layers.13/attn/qkv_proj/MatMul’. The layer_name argument is an ONNX initializer name like ‘model.layers.13.attn.qkv_proj.MatMul.weight’.

To match these, we:

Remove the leading slash from the node name.
Replace all ‘/’ with ‘.’ to match the naming convention of the initializer.

This allows us to correctly identify which weights should be quantized to 8 bits.

Parameters:

layer_name (str)
layers_8bit (list[str])

validate_8bit_layers(layers_str)#

Validate the format of layers_8bit string.

Parameters:: layers_str (str)
Return type:: bool