Dump nodes with potential overflow in half conversion

[tianleiwu]

Description

Add a tool to generate node_block_list used in float16 conversion tool.

Previously, we have a feature to dump statistics data (like min, max) of each node input/output. However, it is time consuming to generate a list of nodes that need to be kept in float32 when model is large.

This could help speed up the process by outputting a list of nodes that have potential overflow in float-to-half conversion.

Usage is to build onnxruntime from source with --cmake_extra_defines onnxruntime_DEBUG_NODE_INPUTS_OUTPUTS=1, then set some environment variables before running float32 optimized onnx model like:

export ORT_DEBUG_NODE_IO_DUMP_HALF_CONVERSION_OVERFLOW=1
export ORT_DEBUG_NODE_IO_HALF_OVERFLOW_THRESHOLD=50000

python benchmark.py -e optimum --height 1024 --width 1024 --steps 3 -b 1 -v Flux.1D -p flux1_dev_onnx/fp32_opt --skip_warmup

The threshold ORT_DEBUG_NODE_IO_HALF_OVERFLOW_THRESHOLD shall be <= 65504. The default value is 50000 if the environment variable is not set. It is better to leave some margin if number of samples are not large enough in the test.

As a demo, we add an option --skip_warmup to benchmark.py for Flux, so that we can reduce the time on dumping warm-up runs.

Example snippet of stdout (each inference session has such a summary when session ended):

Total counter in node dumping: 141
Found 2 nodes cannot be converted to half precision due to potential input/output overflow.
# Example python script for float16 conversion
# For details, search `node_block_list` in https://github.com/microsoft/onnxruntime/blob/main/onnxruntime/python/tools/transformers/float16.py
# -------
from onnxruntime.transformers.onnx_model import OnnxModel
m = OnnxModel(onnx.load('flux1_dev_onnx/fp32_opt/vae_decoder/model.onnx'))
node_block_list = [
  '/decoder/mid_block/attentions.0/Softmax',
  '/decoder/mid_block/attentions.0/MatMul',
]
m.convert_float_to_float16(keep_io_types=False, node_block_list=node_block_list)
m.save_model_to_file('fp16/optimized.onnx', use_external_data_format=False)

Then you can use the python script to convert corresponding model to float16.

Motivation and Context

It is a tool used to generate node_block_list used in float16 conversion of stable diffusion 3.x and flux models in #22986.

In stable diffusion or Flux pipeline, there are multiple models and there could be multiple session runs for each model. Without a proper tool, it is time consuming to get node_block_list for each model.

[github-actions]

You can commit the suggested changes from lintrunner.

[jiafatom]

What is the use case for skip_warmup=True? For some one-time test for a certain input?

[tianleiwu]

Having an option to skip the warm-up runs is useful. You do not need run a model multiple times to collect data dumping or nvtx tracing.

[jiafatom]

Seems that we eventually need compare with half, then why not have a half type for kMaxHalfThreshold?

[tianleiwu]

It is not used to compare with half. One purpose is to check the integer value of the threshold parsed from environ variable.

[jiafatom]

In order to detect half overflow, so the whole dump analysis should be conducted on a fp32 model, right? Do we have some constraints on the use of this function, to ensure that users would not make fp16/int4 model as input by accident?

[tianleiwu]

The design is to analyze float tensor only. It will ignore other data type safely. It is compatible with half conversion script float16.py behavior.

[jiafatom]

Seems that this blocklist is on the node level. Why not on the node type level? We only benchmark a few data, this skiplayer node type has potential to break with some data that is not in benchmark. Why not block this node type?

[tianleiwu]

Once you have node level block list, you can decide whether it is better to extend to node type block list.
Block list on node level has it pros and cons. The pros: it could achieve better performance when only a subset (like 50% SkipLayerNorm nodes) has overflow. The cons: the model has potential to overflow with new data (the risk is smaller if you test more samples and use lower threshold).