From 8d81e56166da05eb6387466e8b476629c56e6b31 Mon Sep 17 00:00:00 2001 From: Tianlei Wu Date: Tue, 14 Jan 2025 06:01:55 +0000 Subject: [PATCH] support bfloat16 --- .../python/tools/transformers/float16.py | 199 +++++++++++++----- 1 file changed, 148 insertions(+), 51 deletions(-) diff --git a/onnxruntime/python/tools/transformers/float16.py b/onnxruntime/python/tools/transformers/float16.py index 74adc951c4..f1062ad6e5 100644 --- a/onnxruntime/python/tools/transformers/float16.py +++ b/onnxruntime/python/tools/transformers/float16.py @@ -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: