saymrwulf/onnxruntime
1
0
Fork 0
mirror of https://github.com/saymrwulf/onnxruntime.git synced 2026-07-29 20:14:01 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

support bfloat16

Browse source
This commit is contained in:
Tianlei Wu 2025-01-14 06:01:55 +00:00
parent 80d8931f1d
commit 8d81e56166
1 changed files with 148 additions and 51 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

199
onnxruntime/python/tools/transformers/float16.py
View file

@ -108,6 +108,41 @@ def convert_tensor_float_to_float16(tensor, min_positive_val=5.96e-08, max_finit
return tensor
def convert_tensor_float_to_bfloat16(tensor, min_positive_val=9.2e-41, max_finite_val=3.38953139e38):
"""Convert tensor float to bfloat16.
Args:
tensor (TensorProto): the tensor to convert.
min_positive_val (float, optional): minimal positive value. Defaults to 9.2e-41.
max_finite_val (float, optional): maximal finite value. Defaults to 3.38953139e38.
Raises:
ValueError: input type is not TensorProto.
Returns:
TensorProto: the converted tensor.
"""
if not isinstance(tensor, TensorProto):
raise ValueError(f"Expected input type is an ONNX TensorProto but got {type(tensor)}")
if tensor.data_type == TensorProto.FLOAT:
values = numpy_helper.to_array(tensor).copy()
values[(values > max_finite_val) & (values < float("inf"))] = max_finite_val
values[(values < min_positive_val) & (values > 0)] = min_positive_val
values[(values > -min_positive_val) & (values < 0)] = -min_positive_val
values[(values < -max_finite_val) & (values > float("-inf"))] = -max_finite_val
tensor.CopyFrom(helper.make_tensor(
name=tensor.name + "_casted_to_bf16",
data_type=TensorProto.BFLOAT16,
dims=values.shape if len(values.shape) != 0 else [],
vals=values, #values.tostring(),
raw=True,
))
return tensor
def make_value_info_from_tensor(tensor):
shape = numpy_helper.to_array(tensor).shape
return helper.make_tensor_value_info(tensor.name, tensor.data_type, shape)
@ -144,25 +179,31 @@ DEFAULT_OP_BLOCK_LIST = [
"Upsample",
]
# Some operators has data type fixed as float for some inputs. Key is op_type, value is list of input indices
# Note that DirectML allows float16 gamma and beta in GroupNorm. Use force_fp16_inputs parameter could overwrite this.
ALWAYS_FLOAT_INPUTS = {"Resize": [2], "GroupNorm": [1, 2], "SkipGroupNorm": [1, 2]}
# Some operators do not support bfloat16 in CUDA. This is not a full list, just some common operators in transformers.
DEFAULT_BF16_OP_BLACK_LIST = ["SkipSimplifiedLayerNormalization", "Attention", "MultiHeadAttention"] + list(ALWAYS_FLOAT_INPUTS.keys())
class InitializerTracker:
"""Class for keeping track of initializer."""
def __init__(self, initializer: TensorProto):
self.initializer = initializer
self.bf16_nodes = []
self.fp32_nodes = []
self.fp16_nodes = []
def add_node(self, node: NodeProto, is_node_blocked):
if is_node_blocked:
def add_node(self, node: NodeProto, dtype: int):
if dtype == TensorProto.FLOAT:
self.fp32_nodes.append(node)
else:
elif dtype == TensorProto.BFLOAT16:
self.bf16_nodes.append(node)
elif dtype == TensorProto.FLOAT16:
self.fp16_nodes.append(node)
else:
raise ValueError("Invalid dtype")
def convert_float_to_float16(
@ -175,26 +216,37 @@ def convert_float_to_float16(
node_block_list=None,
force_fp16_initializers=False,
force_fp16_inputs=None,
use_bfloat16_as_blocked_nodes_dtype=False,
*,
use_bf16_as_blocked_dtype=False,
bf16_op_black_list=None,
):
"""Convert tensor float type in the input ONNX model to tensor float16.
The operators or nodes that are in block lists will not be converted to float16. They will keep the original data
type or be converted to bfloat16 if use_bf16_as_blocked_dtype is True.
Args:
model (ModelProto or str): The ONNX model or path of the model to convert.
min_positive_val (float, optional): minimal positive value. Defaults to 5.96e-08.
max_finite_val (float, optional): maximal finite value of float16. Defaults to 65504.
keep_io_types (Union[bool, List[str]], optional): It could be boolean or a list of float32 input/output names.
If True, model inputs/outputs should be left as float32.
If True, model inputs/outputs keep the original data types.
Defaults to False.
disable_shape_infer (bool, optional): Skips running onnx shape/type inference.
Useful if shape inference has been done. Defaults to False.
op_block_list (List[str], optional): List of op types to leave as float32.
op_block_list (List[str], optional): List of op types to run in float32 or bfloat16.
Defaults to None, which will use `float16.DEFAULT_OP_BLOCK_LIST`.
node_block_list (List[str], optional): List of node names to leave as float32. Defaults to None.
node_block_list (List[str], optional): List of node names to to run in float32 or bfloat16. Defaults to None.
force_fp16_initializers(bool): force converting all float initializers to float16.
Default to false, which will convert only the one needed to avoid precision loss.
force_fp16_inputs(Dict[str, List[int]]): Force the conversion of the inputs of some operators to float16, even if
this script's preference it to keep them in float32.
use_bf16_as_blocked_dtype(bool): use bfloat16 as the data type for blocked nodes or operators.
Default to False.
bf16_op_black_list(List[str], optional): List of operators that do not support bfloat16.
The list is enabled only when use_bf16_as_blocked_dtype is True.
Defaults to None, which will use `float16.DEFAULT_BF16_OP_BLACK_LIST`.
Raises:
ValueError: input type is not ModelProto.
@ -236,6 +288,11 @@ def convert_float_to_float16(
op_block_list = DEFAULT_OP_BLOCK_LIST
if node_block_list is None:
node_block_list = []
bf16_black_list: list[str] = (bf16_op_black_list or DEFAULT_BF16_OP_BLACK_LIST) if use_bf16_as_blocked_dtype else []
def use_bf16_for_blocked_node(n: NodeProto) -> bool:
return use_bf16_as_blocked_dtype and n.op_type not in bf16_black_list
op_block_list = set(op_block_list)
node_block_list = set(node_block_list)
@ -311,6 +368,36 @@ def convert_float_to_float16(
# if q is model, push q.graph (GraphProto)
if isinstance(q, ModelProto):
next_level.append(q.graph)
continue
# if q is model.graph.node.attribute, push q.g and q.graphs (GraphProto)
# and process node.attribute.t and node.attribute.tensors (TensorProto)
if isinstance(q, AttributeProto):
next_level.append(q.g)
for n in q.graphs:
next_level.append(n) # noqa: PERF402
q.t.CopyFrom(convert_tensor_float_to_float16(q.t, min_positive_val, max_finite_val))
for n in q.tensors:
n = convert_tensor_float_to_float16(n, min_positive_val, max_finite_val) # noqa: PLW2901
continue
# if q is graph, process input, output and value_info (ValueInfoProto)
if isinstance(q, GraphProto):
# Note that float initializers tracked by fp32_initializers will be processed later.
# for all ValueInfoProto with tensor(float) type in input, output and value_info, convert them to
# tensor(float16) except map and seq(map). And save them in value_info_list for further processing
for n in itertools.chain(q.input, q.output, q.value_info):
if n.type.tensor_type.elem_type == TensorProto.FLOAT:
if n.name not in graph_io_to_skip:
n.type.tensor_type.elem_type = TensorProto.FLOAT16
value_info_list.append(n)
if n.type.HasField("sequence_type"):
if n.type.sequence_type.elem_type.tensor_type.elem_type == TensorProto.FLOAT:
if n.name not in graph_io_to_skip:
n.type.sequence_type.elem_type.tensor_type.elem_type = TensorProto.FLOAT16
value_info_list.append(n)
continue
# if q is model.graph, push q.node.attribute (AttributeProto)
if isinstance(q, GraphProto):
for n in q.initializer: # TensorProto type
@ -319,8 +406,6 @@ def convert_float_to_float16(
fp32_initializers[n.name] = InitializerTracker(n)
for n in q.node:
# if n is in the block list (doesn't support float16), no conversion for the node,
# and save the node for further processing
if n.name in io_casts:
continue
for i in range(len(n.input)):
@ -330,16 +415,23 @@ def convert_float_to_float16(
if n.output[i] in name_mapping:
n.output[i] = name_mapping[n.output[i]]
# Update tracker for fp32 initializers used by this node.
is_node_blocked = n.op_type in op_block_list or n.name in node_block_list
for i, input_name in enumerate(n.input):
if input_name in fp32_initializers:
# For Resize/GroupNorm, only the first input can be float16
use_fp32_weight = is_node_blocked or (
i in ALWAYS_FLOAT_INPUTS.get(n.op_type, [])
and i not in force_fp16_inputs_dict.get(n.op_type, [])
)
fp32_initializers[input_name].add_node(n, use_fp32_weight)
if i in ALWAYS_FLOAT_INPUTS.get(n.op_type, []) and i not in force_fp16_inputs_dict.get(
n.op_type, []
):
dtype = TensorProto.FLOAT
elif is_node_blocked:
dtype = TensorProto.BFLOAT16 if use_bf16_for_blocked_node(n) else TensorProto.FLOAT
else:
dtype = TensorProto.FLOAT16
fp32_initializers[input_name].add_node(n, dtype)
# if n is in the block list (doesn't support float16), no conversion for the node,
# and save the node for further processing
if is_node_blocked:
node_list.append(n)
else:
@ -371,49 +463,32 @@ def convert_float_to_float16(
if (n.op_type in ["RandomNormal", "RandomUniform", "SequenceEmpty"]) and not has_dtype:
n.attribute.extend([helper.make_attribute("dtype", TensorProto.FLOAT16)])
# For Resize/GroupNorm, attribute data type cannot be changed
if n.op_type not in ALWAYS_FLOAT_INPUTS or n.op_type in force_fp16_inputs_dict:
for attr in n.attribute:
next_level.append(attr) # noqa: PERF402
else:
mixed_float_type_node_list.append(n)
for attr in n.attribute:
next_level.append(attr) # noqa: PERF402
# if q is model.graph.node.attribute, push q.g and q.graphs (GraphProto)
# and process node.attribute.t and node.attribute.tensors (TensorProto)
if isinstance(q, AttributeProto):
next_level.append(q.g)
for n in q.graphs:
next_level.append(n) # noqa: PERF402
q.t.CopyFrom(convert_tensor_float_to_float16(q.t, min_positive_val, max_finite_val))
for n in q.tensors:
n = convert_tensor_float_to_float16(n, min_positive_val, max_finite_val) # noqa: PLW2901
# if q is graph, process input, output and value_info (ValueInfoProto)
if isinstance(q, GraphProto):
# Note that float initializers tracked by fp32_initializers will be processed later.
# for all ValueInfoProto with tensor(float) type in input, output and value_info, convert them to
# tensor(float16) except map and seq(map). And save them in value_info_list for further processing
for n in itertools.chain(q.input, q.output, q.value_info):
if n.type.tensor_type.elem_type == TensorProto.FLOAT:
if n.name not in graph_io_to_skip:
n.type.tensor_type.elem_type = TensorProto.FLOAT16
value_info_list.append(n)
if n.type.HasField("sequence_type"):
if n.type.sequence_type.elem_type.tensor_type.elem_type == TensorProto.FLOAT:
if n.name not in graph_io_to_skip:
n.type.sequence_type.elem_type.tensor_type.elem_type = TensorProto.FLOAT16
value_info_list.append(n)
if n.op_type in ALWAYS_FLOAT_INPUTS and n.op_type not in force_fp16_inputs_dict:
mixed_float_type_node_list.append(n)
queue = next_level
bf16_initializers: Dict[str, TensorProto] = {} # map from float intiailizer name to bfloat16 initializer
for value in fp32_initializers.values():
# By default, to avoid precision loss, do not convert an initializer to fp16 when it is used only by fp32 nodes.
if force_fp16_initializers or value.fp16_nodes:
value.initializer = convert_tensor_float_to_float16(value.initializer, min_positive_val, max_finite_val)
value_info_list.append(make_value_info_from_tensor(value.initializer))
if value.fp32_nodes and not force_fp16_initializers:
if (value.fp32_nodes or value.bf16_nodes) and not force_fp16_initializers:
logger.info(
f"initializer is used by both fp32 and fp16 nodes. Consider add these nodes to block list:{value.fp16_nodes}"
f"initializer is used by both fp32/bf16 and fp16 nodes. Consider add these nodes to block list:{value.fp16_nodes}"
)
elif value.bf16_nodes:
if value.fp32_nodes:
logger.info(
f"initializer is used by both fp32 and bf16 nodes. Consider add these nodes to block list:{value.bf16_nodes}"
)
# If a float initializer is used by bfloat16 nodes, we can convert it to bfloat16.
bf16_initializers[value.initializer.name] = convert_tensor_float_to_bfloat16(value.initializer)
continue
# Some operators have data type fixed as float for some input. Add a float16 to float cast for those inputs.
for node in mixed_float_type_node_list:
@ -436,14 +511,35 @@ def convert_float_to_float16(
node.input[i] = output_name
break
accuracy_type = TensorProto.BFLOAT16 if use_bfloat16_as_blocked_nodes_dtype else TensorProto.FLOAT
# process the nodes in block list that doesn't support tensor(float16)
for node in node_list:
# if input's name is in the value_info_list meaning input is tensor(float16) type,
# insert a float16 to float Cast node before the node,
# insert a float16 to target type (float or bfloat16) Cast node before the node,
# change current node's input name and create new value_info for the new name
use_bf16 = use_bf16_for_blocked_node(node)
accuracy_type = TensorProto.BFLOAT16 if use_bf16 else TensorProto.FLOAT
for i in range(len(node.input)):
input_name = node.input[i]
# It is an intializer that has been converted to bfloat16.
if input_name in bf16_initializers:
if use_bf16:
node.input[i] = bf16_initializers[input_name].name
else: # Add a cast node to convert it from bfloat16 to float.
output_name = node.name + "_input_cast_" + str(i)
value_info = helper.make_tensor_value_info(
name=output_name, elem_type=accuracy_type, shape=bf16_initializers[input_name].dims
)
new_value_info = model.graph.value_info.add()
new_value_info.CopyFrom(value_info)
node_name = node.name + "_input_cast" + str(i)
new_node = [helper.make_node("Cast", [bf16_initializers[input_name].name], [output_name], to=accuracy_type, name=node_name)]
model.graph.node.extend(new_node)
node.input[i] = output_name
continue
# When the name is in value_info_list, its data type has been converted to float16.
# We need to Add a cast node to convert it to float or bfloat16.
for value_info in value_info_list:
if input_name == value_info.name:
# create new value_info for current node's new input name
@ -459,8 +555,9 @@ def convert_float_to_float16(
# change current node's input name
node.input[i] = output_name
break
# if output's name is in the value_info_list meaning output is tensor(float16) type, insert a float to
# float16 Cast node after the node, change current node's output name and create new value_info for the new name
# if output's name is in the value_info_list meaning output is tensor(float16) type, insert a Cast (to float16)
# node after it, change current node's output name and create new value_info for the new name.
for i in range(len(node.output)):
output = node.output[i]
for value_info in value_info_list: