graph_selection#

Node selection and runtime probes for calibrated ONNX quantization.

Functions

find_nodes_from_convs_to_exclude

Find unsupported Conv nodes to exclude from quantization.

find_nodes_from_matmul_to_exclude

Find MatMul nodes that meet gemv or small-gemm conditions and should be excluded.

find_nodes_from_mha_to_exclude

Find MatMul nodes in MHA pattern to exclude.

find_nodes_to_exclude

Find the node names from the ONNX graph which matches user's exclusion patterns.

get_extended_model_outputs

Run one inference step on an onnx model which has some intermediate tensor marked as model outputs.

get_input_shapes

Returns the input shapes of the given ONNX model.

validate_op_types_spelling

Validate spelling in op types.
find_nodes_from_convs_to_exclude(graph, quantize_mode='int8')#

Find unsupported Conv nodes to exclude from quantization.

  • The input and output channels should be >= 16. The exception is for Conv layers in INT8 quantization mode, which supports it if the input or output channel % 8.

  • The filter size for FP8 conv kernels should be less than 32.

  • For FP8 mode, Conv nodes with input or output channels <= _MIN_CHANNELS_FP8 are excluded. Small-channel convolutions do not benefit from FP8 quantization.

Parameters:

  • graph (Graph) – Onnx model graph.

  • quantize_mode (str) – Quantize mode (int8 or fp8).

Returns:

List of Conv nodes.

find_nodes_from_matmul_to_exclude(onnx_path, use_external_data_format=False, intermediate_generated_files=None, calibration_data_reader=None, calibration_eps=['cpu', 'cuda:0', 'trt'], calibration_shapes=None, input_shapes_profile=None, trt_rtx_backend='legacy')#

Find MatMul nodes that meet gemv or small-gemm conditions and should be excluded.

A MatMul is excluded if either:

  • m or n in the output is 1 (GEMV): cannot utilize TensorCores; or

  • K or N is smaller than _MIN_MATMUL_DIM (16): both INT8 and FP8 Tensor Core kernels need K/N >= 16 to be efficient, and adding Q/DQ layers on such small GEMMs causes TRT perf regressions.

Parameters:

  • onnx_path (str) – Path to the onnx model.

  • use_external_data_format (bool) – If True, external data path will be used to store the weights of the intermediate model.

  • intermediate_generated_files (list[str] | None) – List of intermediate generated files that will be deleted after quantization.

  • calibration_data_reader (CalibrationDataReader) – Calibration data reader for running inference.

  • calibration_shapes (str | dict | None) – Model input shapes for inference. If provided, symbolic shape inference will be used instead of calibration_data_reader.

  • calibration_eps (list[str]) – Priority order for the execution providers (EP) to calibrate the model. Any subset of [‘cuda:x’, ‘cpu’, ‘trt’], where ‘x’ is the device id.

  • input_shapes_profile (Sequence[dict[str, str]] | None)

  • trt_rtx_backend (str)

Returns:

List of Nodes to exclude from quantization.

Return type:

list[str]

find_nodes_from_mha_to_exclude(onnx_path, use_external_data_format=False, nodes_to_exclude=None, disable_mha_qdq=False, quantize_mode='int8', intermediate_generated_files=None, calibration_data_reader=None, calibration_eps=['cpu', 'cuda:0', 'trt'], input_shapes_profile=None, trt_rtx_backend='legacy')#

Find MatMul nodes in MHA pattern to exclude.

If disable_mha_qdq is set, don’t add Q/DQ layers to MatMuls in MHA pattern. else when quantize_mode == “fp8”, if head_size > 256 or head_size <= 8 or mha doesn’t meet fp8 fMHA v2 pattern, don’t add Q/DQ layers to MatMuls in MHA pattern. else when quantize_mode == “int8”, if seq_len > 512, don’t add Q/DQ layers to MatMuls in MHA pattern.

Parameters:

  • onnx_path (str) – Path to the onnx model.

  • use_external_data_format (bool) – If True, external data path will be used to store the weights of the intermediate model.

  • nodes_to_exclude (list[str] | None) – List of Nodes to exclude from quantization.

  • disable_mha_qdq (bool) – If True, all MHA’s BMM1 and BMM2 will be added to nodes_to_exclude. Else, each MHA will be checked whether to enable QDQ or not when is_fp8fp16 is True.

  • quantize_mode (str) – Quantization mode. One of ‘int8’ (default), ‘int4’ and ‘fp8’.

  • intermediate_generated_files (list[str] | None) – List of intermediate generated files that will be deleted after quantization.

  • calibration_data_reader (CalibrationDataReader) – Calibration data reader for running inference.

  • calibration_eps (list[str]) – Priority list of execution providers (EP) for calibration.

  • input_shapes_profile (Sequence[dict[str, str]] | None)

  • trt_rtx_backend (str)

Returns:

List of Nodes to exclude from quantization.

Return type:

list[str]

find_nodes_to_exclude(graph, nodes_to_exclude, op_types_to_exclude)#

Find the node names from the ONNX graph which matches user’s exclusion patterns.

Parameters:

  • graph (Graph)

  • nodes_to_exclude (list[str])

  • op_types_to_exclude (list[str])

get_extended_model_outputs(onnx_path, extended_model, use_external_data_format, intermediate_generated_files, calibration_data_reader, calibration_eps, input_shapes_profile=None, trt_rtx_backend='legacy')#

Run one inference step on an onnx model which has some intermediate tensor marked as model outputs.

The first calibration data is used for the dummy inference. This is useful when we want to know the shape of an intermediate tensor given the calibration data.

Parameters:

  • onnx_path (str) – Path to the original onnx model, used for saving the extended model nearby if it is larger than 2GB.

  • extended_model (ModelProto) – The onnx model with some intermediate tensors marked as model outputs.

  • use_external_data_format (bool) – If True, external data path will be used to store the weights of the intermediate model.

  • intermediate_generated_files (list[str]) – List of intermediate generated files that will be deleted after quantization.

  • calibration_data_reader (CalibrationDataReader) – Calibration data reader for running inference.

  • calibration_eps (list[str]) – Priority order for the execution providers (EP) to calibrate the model. Any subset of [‘cuda:x’, ‘cpu’, ‘trt’], where ‘x’ is the device id.

  • input_shapes_profile (Sequence[dict[str, str]] | None)

  • trt_rtx_backend (str)

Return type:

dict[str, ndarray]

Returns: a map with each output name pointed to the corresponding output numpy ndarray.

get_input_shapes(onnx_path)#

Returns the input shapes of the given ONNX model.

Parameters:

onnx_path (str)

Return type:

dict[str, list[int]]

validate_op_types_spelling(onnx_path, op_types_to_quantize, op_types_to_exclude)#

Validate spelling in op types.

Return type:

None