diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 3906cc2ff84..ff846f92e2f 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -35,7 +35,8 @@ Changelog - The ``evaluation`` agent skill can stop looping ("doom loop") reasoning early: ask it to avoid doom loops and it adds vLLM's ``repetition_detection`` request field (vLLM 0.17.0 or later) to every task, ending a request once its output repeats. Use the same setting for baseline and candidate runs. - Add the ordered DLA compatibility pipeline through the ``modelopt.onnx.graph_surgery.make_dla_compatible`` - API and ``make-dla-compatible`` CLI command. Use ``--remove-all-qdq`` to remove Q/DQ pairs of every dtype. + API and ``make-dla-compatible`` CLI command. Use ``--remove-all-qdq`` to remove Q/DQ pairs of every dtype, + or ``--force-fp16-activations`` for FP16 internal computation with FP32 model boundaries. **Backward Breaking Changes** diff --git a/modelopt/onnx/graph_surgery/__main__.py b/modelopt/onnx/graph_surgery/__main__.py index 0c130960de4..de97becee5e 100644 --- a/modelopt/onnx/graph_surgery/__main__.py +++ b/modelopt/onnx/graph_surgery/__main__.py @@ -284,6 +284,14 @@ def main(): action="store_true", help="Remove QDQ pairs for every dtype, matching the C++ remove-all mode", ) + dla_parser.add_argument( + "--force-fp16-activations", + action="store_true", + help=( + "Use FP16 internal activations and weights with FP32 model boundaries; " + "L2 normalization intermediates can overflow or underflow" + ), + ) # Analyze subcommand analyze_parser = subparsers.add_parser( @@ -367,6 +375,7 @@ def main(): verbose=not args.quiet, qdq_quantized_dtypes=args.qdq_quantized_dtypes, remove_all_qdq=args.remove_all_qdq, + force_fp16_activations=args.force_fp16_activations, ) elif args.command == "analyze": diff --git a/modelopt/onnx/graph_surgery/dla_transforms/dla_force_fp16_activations.py b/modelopt/onnx/graph_surgery/dla_transforms/dla_force_fp16_activations.py new file mode 100644 index 00000000000..ce45ed3d6a7 --- /dev/null +++ b/modelopt/onnx/graph_surgery/dla_transforms/dla_force_fp16_activations.py @@ -0,0 +1,688 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Apply the C++ FP16 activation policy while retaining FP32 public boundaries.""" + +import copy +import math + +import numpy as np +from onnx import AttributeProto, TensorProto, defs, helper, numpy_helper + +from ..utils.graph_utils import ( + GraphCache, + _attribute, + _shape_before_boundary_promotion, + batch_replace_nodes, + convert_initializers_to_dtype, +) +from .dla_constants_to_initializers import _add_legacy_initializer_inputs +from .dla_graph_cleanup import _subgraph_captures +from .dla_indexing import _insert_support_nodes, _operator_nodes +from .dla_rank_conversion import ( + _apply_non4d_lpnorm, + _runtime_input_names, + _specialized_collapsed_axis, +) + +__all__ = [] + + +def _cast(source, output, dtype, cache): + """Allocate a boundary cast without colliding with existing graph owners.""" + name = source + "_to_fp16" if dtype == TensorProto.FLOAT16 else output + "_to_fp32" + return helper.make_node("Cast", [source], [output], name=cache.unique_name(name), to=dtype) + + +def _fp16_operator_nodes(model, operations): + """Select standard operations and supported contrib schemas for FP16 policy.""" + nodes = _operator_nodes(model, operations) + if any(op.domain == "com.microsoft" and op.version == 1 for op in model.opset_import): + nodes += _operator_nodes( + model, operations & {"QuantizeLinear", "DequantizeLinear", "Gelu"}, ("com.microsoft",) + ) + return nodes + + +def _is_cast(node, dtype): + """Match a standard Cast with the requested destination type.""" + target = _attribute(node, "to") if node is not None else None + return ( + node is not None + and node.domain in {"", "ai.onnx"} + and not node.overload + and node.op_type == "Cast" + and target is not None + and target.i == dtype + ) + + +def _apply_fp16_lpnorm(model, *, prepared_data_ops=None): + """Decompose last-axis L2 normalization followed by a static gamma multiply.""" + graph = model.graph + cache = GraphCache(graph) + standard = {id(node) for node in _operator_nodes(model, {"LpNormalization", "Mul", "Constant"})} + static = cache.init_names - _runtime_input_names(model) | { + node.output[0] for node in graph.node if node.op_type == "Constant" and id(node) in standard + } + opset = next((o.version for o in model.opset_import if o.domain in ("", "ai.onnx")), 1) + reshapes = {id(node) for node in _operator_nodes(model, {"Reshape"})} + mapped_axes = {} + for node in _operator_nodes(model, {"LpNormalization"}): + attribute = _attribute(node, "axis") + axis = attribute.i if attribute is not None else -1 + shape = _shape_before_boundary_promotion(cache, node.input[0]) + mapped_axes[id(node)] = ( + _specialized_collapsed_axis(cache, node, shape, axis, reshapes) + if shape and len(shape) == 5 + else -1 + ) + nodes = [] + for node in graph.node: + if node.op_type != "LpNormalization" or id(node) not in standard: + nodes.append(node) + continue + axis_attr, p_attr = _attribute(node, "axis"), _attribute(node, "p") + axis = axis_attr.i if axis_attr is not None else -1 + shape = cache.get_shape(node.input[0]) + consumers = cache.get_consumers(node.output[0]) + scaled = ( + len(consumers) == 1 + and consumers[0][0].op_type == "Mul" + and id(consumers[0][0]) in standard + and any(name in static for name in consumers[0][0].input if name != node.output[0]) + ) + if ( + (p_attr is not None and p_attr.i != 2) + or not scaled + or not (axis == -1 or (shape is not None and axis == len(shape) - 1)) + ): + nodes.append(node) + continue + base = node.name or node.output[0] + "_lpnorm" + axis = mapped_axes[id(node)] + axes = cache.unique_name(base + "_axes") + epsilon = cache.unique_name(base + "_epsilon") + if opset >= 13: + graph.initializer.append(numpy_helper.from_array(np.array([axis], np.int64), axes)) + graph.initializer.append(numpy_helper.from_array(np.float32(2**-14), epsilon)) + power, total, norm, safe = [ + cache.unique_name(base + suffix) + for suffix in ("_power_output", "_sum_output", "_norm_output", "_safe_norm_output") + ] + nodes.extend( + [ + helper.make_node( + "Mul", + [node.input[0], node.input[0]], + [power], + name=cache.unique_name(base + "_square"), + ), + helper.make_node( + "ReduceSum", + [power, axes] if opset >= 13 else [power], + [total], + name=cache.unique_name(base + "_sum"), + **({} if opset >= 13 else {"axes": [axis]}), + keepdims=1, + ), + helper.make_node("Sqrt", [total], [norm], name=cache.unique_name(base + "_sqrt")), + helper.make_node( + "Max", [norm, epsilon], [safe], name=cache.unique_name(base + "_safe_norm") + ), + helper.make_node( + "Div", + [node.input[0], safe], + list(node.output), + name=cache.unique_name(base + "_divide"), + ), + ] + ) + retained = {id(node) for node in nodes} + batch_replace_nodes( + graph, + {index: ([node], []) for index, node in enumerate(graph.node) if id(node) not in retained}, + ) + _insert_support_nodes(graph, nodes) + _add_legacy_initializer_inputs(model) + return _apply_non4d_lpnorm(model, prepared_data_ops=prepared_data_ops) + + +def _validate_float_weights(graph): + """Validate all FLOAT payloads before calling the existing weight converter.""" + for tensor in graph.initializer: + if tensor.data_type != TensorProto.FLOAT: + continue + count = math.prod(tensor.dims) + raw = tensor.HasField("raw_data") + if ( + count < 0 + or (raw and len(tensor.raw_data) != count * 4) + or (not raw and len(tensor.float_data) != count) + or (tensor.data_location == TensorProto.EXTERNAL and not raw) + ): + raise ValueError(f"Invalid or unloaded FLOAT initializer {tensor.name!r}") + + +def _fp32_control(node, port, native, required): + """Retain schema-mandated FLOAT inputs and inputs of native FP32 regions.""" + return id(node) in native or (id(node), port) in required + + +def _fp32_regions(model, cache): + """Preserve native floating dependencies and schema-required FLOAT bindings.""" + supported = _fp16_operator_nodes(model, {node.op_type for node in model.graph.node}) + contrib = {id(node) for node in supported if node.domain == "com.microsoft"} + standard = {id(node) for node in supported} + native = { + id(node) + for node in model.graph.node + if id(node) not in standard + or any( + attr.type in {AttributeProto.GRAPH, AttributeProto.GRAPHS} for attr in node.attribute + ) + or any( + (info := cache.vi_map.get(name)) is not None + and info.type.WhichOneof("value") not in {None, "tensor_type"} + for name in [*node.input, *node.output] + ) + } + version = next((op.version for op in model.opset_import if op.domain in {"", "ai.onnx"}), 1) + if _operator_nodes(model, {"QuantizeLinear", "DequantizeLinear"}): + version = max(version, 19) + pending = list(_subgraph_captures(model.graph)) + schemas, required, outputs = {}, set(), set() + for node in model.graph.node: + if id(node) in native: + pending.extend([*node.input, *node.output]) + if id(node) not in standard: + continue + schema_version = version + if id(node) in contrib: + # Contrib v1 supports HALF without upgrading the standard opset. + schema_version = 20 if node.op_type == "Gelu" else 19 + key = node.op_type, schema_version + if key not in schemas: + schema = defs.get_schema(node.op_type, schema_version) + only_fp32 = { + item.type_param_str + for item in schema.type_constraints + if "tensor(float)" in item.allowed_type_strs + and "tensor(float16)" not in item.allowed_type_strs + } + schemas[key] = schema, only_fp32 | {"tensor(float)"} + schema, only_fp32 = schemas[key] + for port, name in enumerate(node.input): + if name and ( + schema.inputs[min(port, len(schema.inputs) - 1)].type_str in only_fp32 + or (node.op_type == "Resize" and port in (1, 2)) + ): + required.add((id(node), port)) + pending.append(name) + float_outputs = {name for name in node.output if cache.get_dtype(name) == TensorProto.FLOAT} + dtype = _attribute(node, "dtype") + if (dtype is not None and dtype.i == TensorProto.FLOAT) or ( + node.op_type != "Cast" + and not any(cache.get_dtype(name) == TensorProto.FLOAT for name in node.input) + ): + outputs.update(float_outputs) + outputs.update( + name + for port, name in enumerate(node.output) + if name in float_outputs + and schema.outputs[min(port, len(schema.outputs) - 1)].type_str in only_fp32 + ) + values = set() + while pending: + name = pending.pop() + if name in values or cache.get_dtype(name) != TensorProto.FLOAT: + continue + values.add(name) + producer = cache.get_producer(name) + if producer is not None: + native.add(id(producer)) + pending.extend([*producer.input, *producer.output]) + return values | outputs, native, required + + +def _upgrade_fp16_quantizers(model, cache): + """Migrate inherited reduction and Split schemas when FP16 quantizers require opset 19.""" + if not _operator_nodes(model, {"QuantizeLinear", "DequantizeLinear"}): + return + opset = next(op for op in model.opset_import if op.domain in {"", "ai.onnx"}) + if opset.version >= 19: + return + functions = {(fn.domain, fn.name, fn.overload) for fn in model.functions} + pending = [model.graph] + while pending: + graph = pending.pop() + allocator = cache if graph is model.graph else GraphCache(graph) + allocator._names.update(_subgraph_captures(graph)) + for node in graph.node: + for attr in node.attribute: + if attr.type == AttributeProto.GRAPH: + pending.append(attr.g) + elif attr.type == AttributeProto.GRAPHS: + pending.extend(attr.graphs) + if ( + node.domain not in {"", "ai.onnx"} + or node.overload + or (node.domain, node.op_type, node.overload) in functions + ): + continue + if ( + node.op_type == "Split" + and (len(node.input) < 2 or not node.input[1]) + and _attribute(node, "num_outputs") is None + ): + # Split-18 requires an explicit choice; legacy equal splits were even. + del node.input[1:] + node.attribute.append(helper.make_attribute("num_outputs", len(node.output))) + continue + axes = _attribute(node, "axes") + if axes is None or not node.op_type.startswith("Reduce"): + continue + if any(value.name == "axes" for value in defs.get_schema(node.op_type, 19).inputs): + control = allocator.unique_name((node.name or node.output[0]) + "_axes") + length = len(axes.ints) + graph.initializer.append( + numpy_helper.from_array(np.asarray(axes.ints, np.int64), control) + ) + node.input.append(control) + node.attribute.remove(axes) + if model.ir_version < 4: + graph.input.append( + helper.make_tensor_value_info(control, TensorProto.INT64, [length]) + ) + opset.version = 19 + + +def _apply_force_fp16_activations(model): + """Run after the ordered DLA passes and their final symbolic shape inference.""" + graph = model.graph + cache = GraphCache(graph) + opset = next((op.version for op in model.opset_import if op.domain in {"", "ai.onnx"}), 1) + if opset < 13 and _operator_nodes(model, {"QuantizeLinear", "DequantizeLinear"}): + raise ValueError("FP16 quantizer migration requires ONNX opset 13 or newer") + fp32, fp16 = TensorProto.FLOAT, TensorProto.FLOAT16 + runtime = _runtime_input_names(model) + native_values, native_nodes, fp32_inputs = _fp32_regions(model, cache) + preserved_inits = { + name + for name in cache.init_names + if name in runtime + or name in native_values + or any( + _fp32_control(node, port, native_nodes, fp32_inputs) + for node, port in cache.get_consumers(name) + ) + } + _validate_float_weights(graph) + if any( + tensor.data_type == fp32 and tensor.name not in preserved_inits + for tensor in graph.initializer + ): + excluded = [ + tensor + for tensor in graph.initializer + if tensor.data_type == fp32 and tensor.name in preserved_inits + ] + for tensor in excluded: + graph.initializer.remove(tensor) + try: + convert_initializers_to_dtype(graph, "float16") + finally: + graph.initializer.extend(excluded) + cache = GraphCache(graph) + preserved = set(native_values) + terminal_concats = set() + for node in graph.node: + if not _is_cast(node, fp32) or cache.get_dtype(node.input[0]) in (fp32, fp16, None): + continue + for concat, _ in cache.get_consumers(node.output[0]): + if concat.op_type != "Concat" or len(concat.output) != 1: + continue + uses = cache.get_consumers(concat.output[0]) + if len(uses) == 1: + boundary, port = uses[0] + if ( + port == 0 + and boundary.op_type == "Squeeze" + and len(boundary.output) == 1 + and (boundary.output[0] in cache.graph_output_names) + ): + preserved.add(concat.output[0]) + terminal_concats.add(id(concat)) + for node in graph.node: + if _is_cast(node, fp32) and node.output[0] not in cache.graph_output_names | preserved: + _attribute(node, "to").i = fp16 + for value in graph.value_info: + if ( + value.type.HasField("tensor_type") + and value.type.tensor_type.elem_type == fp32 + and (value.name not in cache.graph_output_names | preserved | preserved_inits) + ): + value.type.tensor_type.elem_type = fp16 + for value in graph.input: + if ( + value.name in cache.init_names - preserved_inits + and value.type.tensor_type.elem_type == fp32 + ): + value.type.tensor_type.elem_type = fp16 + before, after = {}, {} + front, back = [], [] + for tensor in graph.initializer: + if tensor.name not in preserved_inits - runtime or tensor.data_type != fp32: + continue + consumers = [ + (node, port) + for node, port in cache.get_consumers(tensor.name) + if not _fp32_control(node, port, native_nodes, fp32_inputs) and not _is_cast(node, fp16) + ] + if consumers: + output = cache.unique_name(tensor.name + "_fp16") + front.append(_cast(tensor.name, output, fp16, cache)) + for node, port in consumers: + node.input[port] = output + for node in graph.node: + for name in node.output: + if name not in native_values: + continue + consumers = [ + (consumer, port) + for consumer, port in cache.get_consumers(name) + if not _fp32_control(consumer, port, native_nodes, fp32_inputs) + and not _is_cast(consumer, fp16) + ] + if consumers: + output = cache.unique_name(name + "_fp16") + after.setdefault(id(node), []).append(_cast(name, output, fp16, cache)) + for consumer, port in consumers: + consumer.input[port] = output + for node in graph.node: + if id(node) not in terminal_concats: + continue + for port, source in enumerate(node.input): + if source not in preserved: + output = cache.unique_name(source + "_concat_fp32") + before.setdefault(id(node), []).append(_cast(source, output, fp32, cache)) + node.input[port] = output + preserved.add(output) + for value in graph.input: + if value.name in cache.init_names - runtime or value.type.tensor_type.elem_type != fp32: + continue + direct = None + rank = len(value.type.tensor_type.shape.dim) + for consumer, port in cache.get_consumers(value.name): + if _fp32_control(consumer, port, native_nodes, fp32_inputs) or _is_cast(consumer, fp16): + continue + shape_info = cache.vi_map.get(consumer.output[0]) + adapter = ( + rank < 4 + and port == 0 + and consumer.op_type == "Unsqueeze" + and (shape_info is not None and len(shape_info.type.tensor_type.shape.dim) == 4) + ) + if adapter: + assert shape_info is not None + source = consumer.output[0] + output = cache.unique_name(source + "_fp16") + after.setdefault(id(consumer), []).append(_cast(source, output, fp16, cache)) + for downstream, slot in cache.get_consumers(source): + if not _is_cast(downstream, fp16): + downstream.input[slot] = output + shape_info.type.tensor_type.elem_type = fp32 + else: + if direct is None: + direct = cache.unique_name(value.name + "_fp16") + front.append(_cast(value.name, direct, fp16, cache)) + consumer.input[port] = direct + _upgrade_fp16_quantizers(model, cache) + for value in graph.output: + if value.type.tensor_type.elem_type != fp32 or value.name in native_values: + continue + producer = cache.get_producer(value.name) + public_cast = _is_cast(producer, fp32) + if producer is not None and (producer.op_type == "Squeeze" or public_cast): + if not public_cast: + source = producer.input[0] + if source not in preserved: + output = cache.unique_name(source + "_to_fp32_fp16") + before.setdefault(id(producer), []).append(_cast(source, output, fp32, cache)) + producer.input[0] = output + internal = [ + (consumer, port) + for consumer, port in cache.get_consumers(value.name) + if not _fp32_control(consumer, port, native_nodes, fp32_inputs) + and not _is_cast(consumer, fp16) + and (not public_cast or consumer.op_type != "Cast") + ] + if internal: + view = cache.unique_name(value.name + "_fp16") + after.setdefault(id(producer), []).append(_cast(value.name, view, fp16, cache)) + for consumer, port in internal: + consumer.input[port] = view + continue + tensor = cache.get_init(value.name) if producer is None else None + if producer is None and (tensor is None or tensor.data_type != fp16): + continue + internal = cache.unique_name(value.name + "_fp16") + if producer is None: + assert tensor is not None + tensor.name = internal + if model.ir_version < 4: + for declaration in graph.input: + if declaration.name == value.name: + declaration.name = internal + else: + for port, name in enumerate(producer.output): + if name == value.name: + producer.output[port] = internal + break + for annotation in graph.value_info: + if annotation.name == value.name: + annotation.name = internal + annotation.type.tensor_type.elem_type = fp16 + for consumer, port in cache.get_consumers(value.name): + consumer.input[port] = internal + back.append(_cast(internal, value.name, fp32, cache)) + nodes = list(front) + for node in graph.node: + nodes.extend(before.get(id(node), [])) + nodes.append(node) + nodes.extend(after.get(id(node), [])) + nodes.extend(back) + _insert_support_nodes(graph, nodes) + _rewrite_terminal_projections(model) + _add_legacy_initializer_inputs(model) + return model + + +def _rewrite_terminal_projections(model): + """Bound private fixed projection channels while preserving native owners.""" + graph = model.graph + cache = GraphCache(graph) + standard = { + id(node) for node in _fp16_operator_nodes(model, {node.op_type for node in graph.node}) + } + protected = cache.graph_output_names | _subgraph_captures(graph) + runtime = _runtime_input_names(model) + removed, removed_values, removed_inits = set(), set(), set() + replacements, additions = {}, [] + + def attributes(node, key, expected): + """Match one fixed operator attribute used by the projection pattern.""" + value = _attribute(node, key) + return ( + value is not None + and (list(value.ints) if isinstance(expected, list) else value.i) == expected + ) + + def axes(node, expected): + """Match fixed axes without changing their initializer owner.""" + value = cache.get_init_array(node.input[1]) if len(node.input) == 2 else None + return value is not None and value.reshape(-1).tolist() == expected + + for output in graph.output: + if output.type.tensor_type.elem_type != TensorProto.FLOAT: + continue + chain = [] + current = output.name + for op in ("Squeeze", "Cast", "Concat", "ReduceSum", "Concat", "Transpose", "Conv"): + node = cache.get_producer(current) + if node is None or id(node) not in standard or node.op_type != op: + break + chain.append(node) + current = node.input[0] + if len(chain) != 7: + continue + squeeze, cast, final, reduce, initial, transpose, conv = chain + if not ( + axes(squeeze, [0]) + and _is_cast(cast, TensorProto.FLOAT) + and len(final.input) == len(initial.input) == 1 + and axes(reduce, [0]) + and attributes(reduce, "keepdims", 1) + and attributes(transpose, "perm", [0, 3, 2, 1]) + and len(conv.input) == 2 + and len(conv.output) == 1 + and ((_attribute(conv, "group") is None) or _attribute(conv, "group").i == 1) + ): + continue + weight = cache.get_init(conv.input[1]) + dq = cache.get_producer(conv.input[1]) if weight is None else None + if dq is not None: + parameters = [cache.get_init(name) for name in dq.input[1:]] + if ( + dq.op_type != "DequantizeLinear" + or id(dq) not in standard + or len(dq.input) < 2 + or any( + parameter is None or math.prod(parameter.dims) != 1 for parameter in parameters + ) + ): + continue + weight = cache.get_init(dq.input[0]) + chain.append(dq) + if ( + weight is None + or len(weight.dims) != 4 + or weight.dims[0] <= 8192 + or not weight.HasField("raw_data") + ): + continue + if weight.data_type != (TensorProto.UINT8 if dq is not None else TensorProto.FLOAT16): + continue + chain_ids = {id(n) for n in chain} + if any( + name in protected | runtime for name in (weight.name, squeeze.input[1], reduce.input[1]) + ): + continue + if any(name in protected for item in chain[1:] for name in item.output): + continue + if any( + id(consumer) not in chain_ids + for name in (weight.name, squeeze.input[1], reduce.input[1]) + for consumer, _ in cache.get_consumers(name) + ): + continue + if any(len(cache.get_consumers(n.output[0])) != 1 for n in chain[1:]): + continue + channels = weight.dims[0] + if cache.get_shape(output.name) != [1, 1, channels] or cache.get_shape(conv.output[0]) != [ + 1, + channels, + 1, + 1, + ]: + continue + bytes_per_channel = math.prod(weight.dims[1:]) * (1 if dq is not None else 2) + payload = weight.raw_data + if bytes_per_channel <= 0 or len(payload) != channels * bytes_per_channel: + continue + nodes, outputs = [], [] + for index, start in enumerate(range(0, channels, 8192)): + count = min(8192, channels - start) + suffix = f"_dla_chunk_{index}" + chunk = TensorProto( + name=cache.unique_name(weight.name + suffix), + data_type=weight.data_type, + dims=[count, *weight.dims[1:]], + ) + chunk.raw_data = payload[ + start * bytes_per_channel : (start + count) * bytes_per_channel + ] + additions.append(chunk) + cloned_conv = copy.deepcopy(conv) + cloned_conv.name = cache.unique_name(conv.name + suffix) + if dq is not None: + cloned_dq = copy.deepcopy(dq) + cloned_dq.name = cache.unique_name(dq.name + suffix) + cloned_dq.input[0] = chunk.name + cloned_dq.output[0] = cache.unique_name(dq.output[0] + suffix) + nodes.append(cloned_dq) + cloned_conv.input[1] = cloned_dq.output[0] + else: + cloned_conv.input[1] = chunk.name + cloned_conv.output[0] = cache.unique_name(conv.output[0] + suffix) + nodes.append(cloned_conv) + shape = cache.unique_name(conv.output[0] + suffix + "_bounded_shape") + additions.append(numpy_helper.from_array(np.array([1, 1, count], np.int64), shape)) + bounded = cache.unique_name(conv.output[0] + suffix + "_bounded") + nodes.append( + helper.make_node( + "Reshape", + [cloned_conv.output[0], shape], + [bounded], + name=cache.unique_name(conv.name + suffix + "_bounded_reshape"), + ) + ) + result = cache.unique_name(conv.output[0] + suffix + "_fp32") + nodes.append(_cast(bounded, result, TensorProto.FLOAT, cache)) + outputs.append(result) + nodes.append( + helper.make_node( + "Concat", + outputs, + [output.name], + axis=2, + name=cache.unique_name(output.name + "_dla_terminal_concat"), + ) + ) + replacements[id(conv)] = nodes + removed.update(id(n) for n in chain) + removed_values.update(name for n in chain for name in n.output if name != output.name) + removed_inits.update([weight.name, squeeze.input[1], reduce.input[1]]) + if not replacements: + return + batch_replace_nodes( + graph, + { + index: ([node], replacements.get(id(node), [])) + for index, node in enumerate(graph.node) + if id(node) in removed + }, + ) + for index in range(len(graph.initializer) - 1, -1, -1): + if graph.initializer[index].name in removed_inits: + del graph.initializer[index] + graph.initializer.extend(additions) + for index in range(len(graph.value_info) - 1, -1, -1): + if graph.value_info[index].name in removed_values: + del graph.value_info[index] + if model.ir_version < 4: + for index in range(len(graph.input) - 1, -1, -1): + if graph.input[index].name in removed_inits: + del graph.input[index] diff --git a/modelopt/onnx/graph_surgery/make_dla_compatible.py b/modelopt/onnx/graph_surgery/make_dla_compatible.py index c14fbce5268..81d16d3bc7d 100644 --- a/modelopt/onnx/graph_surgery/make_dla_compatible.py +++ b/modelopt/onnx/graph_surgery/make_dla_compatible.py @@ -35,6 +35,10 @@ from .dla_transforms.dla_cast_to_fp32 import _apply_cast_to_fp32 from .dla_transforms.dla_constants_to_initializers import _constants_to_initializers_from_model from .dla_transforms.dla_decompose_lstm import _apply_decompose_lstm +from .dla_transforms.dla_force_fp16_activations import ( + _apply_force_fp16_activations, + _apply_fp16_lpnorm, +) from .dla_transforms.dla_graph_cleanup import _remove_unused_initializers from .dla_transforms.dla_greater import _apply_greater from .dla_transforms.dla_indexing import ( @@ -232,6 +236,7 @@ def _transform_make_dla_compatible( verbose: bool = True, qdq_quantized_dtypes: Iterable[int | str] | None = None, remove_all_qdq: bool = False, + force_fp16_activations: bool = False, ) -> onnx.ModelProto: """Apply the logged DLA sequence in memory, then infer shapes and validate.""" started = time.perf_counter() @@ -297,12 +302,22 @@ def _transform_make_dla_compatible( options = _boundary_stage_options( name, quant_types, verbose, prepared_input_ops, prepared_output_ops ) + if name == "Non4DLpNorm" and force_fp16_activations: + transform = _apply_fp16_lpnorm result = transform(model, **options) if result is not None: model = result model = infer_shapes(model) model = _apply_reorder_nodes(model) onnx.checker.check_model(model) + if force_fp16_activations: + if verbose: + logger.info("[DLA pipeline] Applying ForceFp16Activations") + model = _apply_force_fp16_activations(model) + # Match C++: symbolic inference stays before precision conversion. + model = onnx.shape_inference.infer_shapes(model, check_type=True, strict_mode=True) + # Strict typed inference already covers full_check's extra model-copy pass. + onnx.checker.check_model(model) if verbose: logger.info("[DLA pipeline] Done (%.1fs).", time.perf_counter() - started) return model @@ -317,6 +332,7 @@ def dla_make_dla_compatible( verbose: bool = True, qdq_quantized_dtypes: Iterable[int | str] | None = None, remove_all_qdq: bool = False, + force_fp16_activations: bool = False, ) -> onnx.ModelProto: """Load an ONNX model, apply the ordered DLA transforms, and save the result. @@ -325,6 +341,13 @@ def dla_make_dla_compatible( Masked multiplication can change the sign of zero even for finite inputs; validate sign-sensitive downstream operations before using this conversion. + With ``force_fp16_activations=True``, last-axis L2 normalization followed by a + static gamma multiply uses FP16 arithmetic with a denominator floor of ``2**-14``. + This decomposition does not rescale inputs: squares or their sum can overflow, + and small squares can underflow, even for finite FP32 inputs. Validate numerical + accuracy on the target runtime; leave this option disabled when its intermediate + computations need a wider floating-point range. + Args: model_path: Input file; external weights are loaded from its directory. output_path: Destination ONNX file. @@ -335,6 +358,8 @@ def dla_make_dla_compatible( Defaults to UINT16 and INT16. Accepts ONNX dtype IDs or names. remove_all_qdq: Remove QDQ pairs for every dtype, matching the C++ flag. Cannot be combined with qdq_quantized_dtypes. + force_fp16_activations: Use FP16 where supported, keeping FP32 public + boundaries and operations or controls that require FP32. Returns: The transformed model. Invalid models and failed transformations raise. @@ -344,6 +369,7 @@ def dla_make_dla_compatible( verbose=verbose, qdq_quantized_dtypes=qdq_quantized_dtypes, remove_all_qdq=remove_all_qdq, + force_fp16_activations=force_fp16_activations, ) output_dir = os.path.dirname(output_path) if output_dir: diff --git a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_force_fp16_activations.py b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_force_fp16_activations.py new file mode 100644 index 00000000000..826567d77ef --- /dev/null +++ b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_force_fp16_activations.py @@ -0,0 +1,1155 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Focused FP16 payload, boundary, and DLA projection assertions on small fixtures.""" + +import numpy as np +import onnx +import onnxruntime as ort +import pytest +from onnx import TensorProto as T +from onnx import helper as h +from onnx import numpy_helper as nh +from onnx.reference import ReferenceEvaluator + +from modelopt.onnx.graph_surgery import make_dla_compatible +from modelopt.onnx.graph_surgery.dla_transforms.dla_force_fp16_activations import ( + _apply_force_fp16_activations, + _apply_fp16_lpnorm, +) +from modelopt.onnx.graph_surgery.dla_transforms.dla_rank_conversion import _apply_non4d_reshape +from modelopt.onnx.graph_surgery.make_dla_compatible import _transform_make_dla_compatible + + +def _value(name, shape, dtype=T.FLOAT): + """Declare a typed tensor for a checked precision fixture.""" + return h.make_tensor_value_info(name, dtype, shape) + + +def _model(nodes, inputs, outputs, initializers, opset=19): + """Build a model with inferred intermediate types for a precision fixture.""" + model = h.make_model( + h.make_graph(nodes, "fp16", inputs, outputs, initializers), + opset_imports=[h.make_opsetid("", opset)], + ) + model.ir_version = 9 + return onnx.shape_inference.infer_shapes(model) + + +@pytest.mark.parametrize("enabled", [False, True]) +def test_public_option_preserves_boundaries_and_selects_weight_precision(tmp_path, enabled): + """Keep public FLOAT interfaces while selecting internal weight precision.""" + values = np.array([1.0003, 0.12345], np.float32) + model = _model( + [h.make_node("Mul", ["x", "w"], ["y"])], + [_value("x", [1, 1, 1, 2])], + [_value("y", [1, 1, 1, 2])], + [nh.from_array(values.reshape(1, 1, 1, 2), "w")], + ) + source, output = tmp_path / "source.onnx", tmp_path / "output.onnx" + onnx.save(model, source) + result = make_dla_compatible( + str(source), + str(output), + verbose=False, + use_external_data=False, + force_fp16_activations=enabled, + ) + assert result.graph.input[0].type.tensor_type.elem_type == T.FLOAT + assert result.graph.output[0].name == "y" + assert result.graph.output[0].type.tensor_type.elem_type == T.FLOAT + assert result.graph.initializer[0].data_type == (T.FLOAT16 if enabled else T.FLOAT) + data = np.array([[[[0.3333, 2.0003]]]], np.float32) + dtype = np.float16 if enabled else np.float32 + expected = (data.astype(dtype) * values.astype(dtype)).astype(np.float32) + actual = ReferenceEvaluator(onnx.load(output)).run(None, {"x": data})[0] + np.testing.assert_array_equal(actual, expected) + + +def test_initializer_rounding_preserves_subnormals_signed_zero_and_nan(): + """Retain the original converter rounding, signed zero and NaN behavior.""" + values = np.array([0, -0.0, 2**-24, 2**-25, 2**-14, 65504, 65520, 0, 0], np.float32) + values.view(np.uint32)[-2:] = [0x7F812345, 0xFF812345] + model = _model([], [], [_value("weight", [len(values)])], [nh.from_array(values, "weight")]) + result = _apply_force_fp16_activations(model) + actual = nh.to_array(result.graph.initializer[0]).view(np.uint16) + np.testing.assert_array_equal(actual[:-2], [0, 0x8000, 1, 0, 0x400, 0x7BFF, 0x7C00]) + assert np.isnan(actual[-2:].view(np.float16)).all() + assert result.graph.node[0].op_type == "Cast" + assert result.graph.node[0].output[0] == "weight" + + +def test_boundary_adapters_and_direct_input_consumers_get_their_own_casts(): + """Give boundary adapters and direct input uses their own typed views.""" + model = _model( + [ + h.make_node("Unsqueeze", ["x", "axes"], ["x4"]), + h.make_node("Relu", ["x4"], ["r4"]), + h.make_node("Squeeze", ["r4", "axes"], ["y"]), + h.make_node("Neg", ["x"], ["direct"]), + ], + [_value("x", [1, 2])], + [_value("y", [1, 2]), _value("direct", [1, 2])], + [nh.from_array(np.array([0, 1], np.int64), "axes")], + ) + result = _apply_force_fp16_activations(model) + producers = {name: n for n in result.graph.node for name in n.output} + relu = next(n for n in result.graph.node if n.op_type == "Relu") + assert producers[relu.input[0]].op_type == "Cast" + assert producers[producers[relu.input[0]].input[0]].op_type == "Unsqueeze" + squeeze = producers["y"] + assert squeeze.op_type == "Squeeze" + assert producers[squeeze.input[0]].op_type == "Cast" + actual = ReferenceEvaluator(result).run(None, {"x": np.array([[-2, 3]], np.float32)}) + np.testing.assert_array_equal(actual[0], [[0, 3]]) + np.testing.assert_array_equal(actual[1], [[2, -3]]) + + +@pytest.mark.parametrize("operation", ["Resize", "QuantizeLinear"]) +def test_schema_specific_float_parameters_and_qdq_opset(operation): + """Keep required FLOAT controls and migrate FP16 quantizer schemas.""" + if operation == "Resize": + nodes = [ + h.make_node("Mul", ["x", "weight"], ["weighted"]), + h.make_node("Resize", ["weighted", "roi", "scales"], ["y"]), + ] + initializers = [ + nh.from_array(np.full((1, 1, 1, 2), 0.5, np.float32), "weight"), + nh.from_array(np.array([], np.float32), "roi"), + nh.from_array(np.array([1, 1, 1, 2], np.float32), "scales"), + ] + output_shape = [1, 1, 1, 4] + else: + nodes = [ + h.make_node("QuantizeLinear", ["x", "scale", "zero"], ["q"]), + h.make_node("DequantizeLinear", ["q", "scale", "zero"], ["y"]), + ] + initializers = [nh.from_array(np.float32(0.1), "scale"), nh.from_array(np.int8(0), "zero")] + output_shape = [1, 1, 1, 2] + model = _model( + nodes, [_value("x", [1, 1, 1, 2])], [_value("y", output_shape)], initializers, opset=17 + ) + result = _apply_force_fp16_activations(model) + result = onnx.shape_inference.infer_shapes(result, check_type=True) + onnx.checker.check_model(result) + types = {t.name: t.data_type for t in result.graph.initializer} + if operation == "Resize": + assert types == {"weight": T.FLOAT16, "roi": T.FLOAT, "scales": T.FLOAT} + assert result.opset_import[0].version == 17 + else: + assert types == {"scale": T.FLOAT16, "zero": T.INT8} + assert result.opset_import[0].version == 19 + + +def test_scaled_l2_normalization_keeps_zero_slices_finite(): + """Stabilize zero-vector normalization before applying FP16 precision.""" + model = _model( + [ + h.make_node("LpNormalization", ["x"], ["norm"], p=2, axis=-1), + h.make_node("Mul", ["norm", "gamma"], ["y"]), + ], + [_value("x", [1, 1, 2, 2])], + [_value("y", [1, 1, 2, 2])], + [nh.from_array(np.ones((1, 1, 1, 2), np.float32), "gamma")], + ) + result = _apply_force_fp16_activations(_apply_fp16_lpnorm(model)) + assert not any(n.op_type == "LpNormalization" for n in result.graph.node) + assert any(n.op_type == "Max" for n in result.graph.node) + data = np.array([[[[0, 0], [3, 4]]]], np.float32) + actual = ReferenceEvaluator(result).run(None, {"x": data})[0] + assert np.isfinite(actual).all() + np.testing.assert_allclose(actual, [[[[0, 0], [0.6, 0.8]]]], atol=3e-4) + + +@pytest.mark.parametrize("shape", [[2, 3], [2, 3, 4]]) +@pytest.mark.parametrize("positive_axis", [False, True]) +@pytest.mark.parametrize("enabled", [False, True]) +@pytest.mark.parametrize("unary", [False, True]) +def test_scaled_low_rank_lpnorm_preserves_last_axis(shape, positive_axis, enabled, unary): + """Normalize the last runtime dimension after promoting a native positive axis.""" + axis = len(shape) - 1 if positive_axis else -1 + nodes = [h.make_node("Relu", ["x"], ["native"])] if unary else [] + nodes.extend( + [ + h.make_node("LpNormalization", ["native" if unary else "x"], ["norm"], axis=axis, p=2), + h.make_node("Mul", ["norm", "gamma"], ["y"]), + ] + ) + model = _model( + nodes, + [_value("x", shape)], + [_value("y", shape)], + [nh.from_array(np.linspace(0.5, 1.5, shape[-1], dtype=np.float32), "gamma")], + opset=17, + ) + feeds = {"x": np.linspace(0.25, 2, np.prod(shape), dtype=np.float32).reshape(shape)} + onnx.checker.check_model(model, full_check=True) + expected = ReferenceEvaluator(model).run(None, feeds)[0] + converted = _transform_make_dla_compatible(model, verbose=False, force_fp16_activations=enabled) + onnx.checker.check_model(converted, full_check=True) + actual = ReferenceEvaluator(converted).run(None, feeds)[0] + np.testing.assert_allclose(actual, expected, rtol=2e-3, atol=3e-4) + + +def _scaled_collapsed_lpnorm_model(shape, axis, prefix=False): + """Build checked native normalization around a real singleton Reshape collapse.""" + collapsed = shape.copy() + collapsed.remove(1) + nodes = [ + h.make_node("Reshape", ["x", "native_shape"], ["native"]), + h.make_node("LpNormalization", ["native"], ["norm"], axis=axis, p=2), + h.make_node("Mul", ["norm", "gamma"], ["scaled"]), + h.make_node("Reshape", ["scaled", "output_shape"], ["y"]), + ] + outputs = [_value("y", collapsed)] + initializers = [ + nh.from_array(np.asarray(shape, np.int64), "native_shape"), + nh.from_array(np.asarray(collapsed, np.int64), "output_shape"), + nh.from_array(np.ones(shape[-1], np.float32), "gamma"), + ] + if prefix: + nodes = [ + h.make_node("LpNormalization", ["x"], ["first_norm"], axis=-1, p=2), + h.make_node("Mul", ["first_norm", "first_gamma"], ["first"]), + *nodes, + ] + outputs.insert(0, _value("first", collapsed)) + initializers.append(nh.from_array(np.ones(collapsed[-1], np.float32), "first_gamma")) + return _model(nodes, [_value("x", collapsed)], outputs, initializers, opset=17) + + +@pytest.mark.parametrize("shape", [[1, 2, 3, 4, 5], [2, 3, 1, 4, 5]]) +@pytest.mark.parametrize("axis", [-1, 4]) +def test_fp16_scaled_lpnorm_uses_actual_collapsed_axis(shape, axis): + """Use the mapped last axis after the registered five-dimensional collapse.""" + model = _scaled_collapsed_lpnorm_model(shape, axis) + collapsed = shape.copy() + collapsed.remove(1) + feeds = {"x": np.linspace(0.25, 2, np.prod(collapsed), dtype=np.float32).reshape(collapsed)} + onnx.checker.check_model(model, full_check=True) + expected = ReferenceEvaluator(model).run(None, feeds)[0] + converted = _transform_make_dla_compatible(model, verbose=False, force_fp16_activations=True) + onnx.checker.check_model(converted, full_check=True) + actual = ReferenceEvaluator(converted).run(None, feeds)[0] + np.testing.assert_allclose(actual, expected, rtol=2e-3, atol=3e-4) + + +@pytest.mark.parametrize("axis", [-1, 4]) +@pytest.mark.parametrize("enabled", [False, True]) +def test_scaled_lpnorm_refuses_normalizing_over_the_removed_axis(axis, enabled): + """Give FP16 and FP32 the same refusal for a normalization axis removed by collapse.""" + model = _scaled_collapsed_lpnorm_model([2, 3, 4, 5, 1], axis) + feeds = {"x": np.ones([2, 3, 4, 5], np.float32)} + onnx.checker.check_model(model, full_check=True) + np.testing.assert_array_equal(ReferenceEvaluator(model).run(None, feeds)[0], feeds["x"]) + with pytest.raises(ValueError, match="cannot map axis"): + _transform_make_dla_compatible(model, verbose=False, force_fp16_activations=enabled) + + +@pytest.mark.parametrize("axis", [-1, 4]) +def test_fp16_lpnorm_validates_late_collapsed_axis_before_decomposition(axis): + """Refuse unsupported normalization before mutating an earlier eligible FP16 branch.""" + model = _scaled_collapsed_lpnorm_model([2, 3, 4, 5, 1], axis, prefix=True) + onnx.checker.check_model(model, full_check=True) + ReferenceEvaluator(model).run(None, {"x": np.ones([2, 3, 4, 5], np.float32)}) + _apply_non4d_reshape(model) + before = model.SerializeToString() + with pytest.raises(ValueError, match="cannot map axis"): + _apply_fp16_lpnorm(model) + assert model.SerializeToString() == before + + +def test_terminal_projection_splits_channels_and_preserves_values(): + """Bound private projection channels while preserving real output values.""" + channels = 8193 + model = _model( + [ + h.make_node("Conv", ["x", "weight"], ["conv"]), + h.make_node("Transpose", ["conv"], ["transposed"], perm=[0, 3, 2, 1]), + h.make_node("Concat", ["transposed"], ["concat1"], axis=0), + h.make_node("ReduceSum", ["concat1", "axes"], ["reduced"], keepdims=1), + h.make_node("Concat", ["reduced"], ["concat2"], axis=2), + h.make_node("Squeeze", ["concat2", "axes"], ["y"]), + ], + [_value("x", [1, 1, 1, 1])], + [_value("y", [1, 1, channels])], + [ + nh.from_array(np.full((channels, 1, 1, 1), np.float32(0.25)), "weight"), + nh.from_array(np.array([0], np.int64), "axes"), + ], + ) + result = _apply_force_fp16_activations(model) + assert sum(n.op_type == "Conv" for n in result.graph.node) == 2 + weights = [t for t in result.graph.initializer if t.data_type == T.FLOAT16] + assert [list(t.dims) for t in weights] == [[8192, 1, 1, 1], [1, 1, 1, 1]] + options = ort.SessionOptions() + options.intra_op_num_threads = 1 + options.inter_op_num_threads = 1 + actual = ort.InferenceSession( + result.SerializeToString(), options, providers=["CPUExecutionProvider"] + ).run(None, {"x": np.array([[[[2]]]], np.float32)})[0] + np.testing.assert_array_equal(actual, np.full((1, 1, channels), np.float32(0.5))) + + +def _checked_cpu_outputs(model, feeds): + """Execute a fully checked graph with bounded CPU runtime threading.""" + onnx.checker.check_model(model, full_check=True) + options = ort.SessionOptions() + options.intra_op_num_threads = 1 + options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_DISABLE_ALL + return ort.InferenceSession( + model.SerializeToString(), options, providers=["CPUExecutionProvider"] + ).run(None, feeds) + + +def test_fp16_runtime_initializer_default_retains_fp32_public_contract(): + """Keep FP32 default-input declarations and accept FP32 runtime overrides.""" + model = _model( + [h.make_node("Mul", ["x", "w"], ["y"])], + [_value("x", [2]), _value("w", [2])], + [_value("y", [2])], + [nh.from_array(np.asarray([0.5, 2], np.float32), "w")], + ) + before = model.graph.initializer[0].SerializeToString() + cases = [ + {"x": np.asarray([2, 4], np.float32)}, + {"x": np.asarray([2, 4], np.float32), "w": np.asarray([2, 0.25], np.float32)}, + ] + expected = [_checked_cpu_outputs(model, feeds)[0] for feeds in cases] + changed = _apply_force_fp16_activations(model) + assert ( + next(value for value in changed.graph.input if value.name == "w").type.tensor_type.elem_type + == T.FLOAT + ) + assert ( + next(value for value in changed.graph.initializer if value.name == "w").SerializeToString() + == before + ) + for feeds, wanted in zip(cases, expected): + np.testing.assert_array_equal(_checked_cpu_outputs(changed, feeds)[0], wanted) + + +def test_fp16_resize_parameter_shared_with_arithmetic_uses_private_precision_view(): + """Keep required FP32 scales while providing a separate FP16 arithmetic operand.""" + model = _model( + [ + h.make_node("Mul", ["x", "scales"], ["weighted"]), + h.make_node("Resize", ["weighted", "roi", "scales"], ["y"]), + ], + [_value("x", [1, 1, 1, 4])], + [_value("y", [1, 1, 1, 8])], + [ + nh.from_array(np.asarray([1, 1, 1, 2], np.float32), "scales"), + nh.from_array(np.asarray([], np.float32), "roi"), + ], + ) + feeds = {"x": np.asarray([[[[2, 4, 6, 8]]]], np.float32)} + expected = _checked_cpu_outputs(model, feeds)[0] + changed = _apply_force_fp16_activations(model) + np.testing.assert_array_equal(_checked_cpu_outputs(changed, feeds)[0], expected) + + +def test_fp16_conversion_preserves_unrelated_constant_object(): + """Avoid copying an unrelated integer Constant while inserting floating boundaries.""" + constant = h.make_node( + "Constant", [], ["integer_payload"], value=nh.from_array(np.arange(64, dtype=np.uint8)) + ) + model = _model( + [constant, h.make_node("Relu", ["x"], ["y"])], [_value("x", [2])], [_value("y", [2])], [] + ) + constant = model.graph.node[0] + before = constant.SerializeToString() + changed = _apply_force_fp16_activations(model) + assert next(node for node in changed.graph.node if node.op_type == "Constant") is constant + assert constant.SerializeToString() == before + _checked_cpu_outputs(changed, {"x": np.asarray([-1, 2], np.float32)}) + + +def test_fp16_malformed_late_weight_refuses_before_any_mutation(): + """Validate every floating payload before converting an earlier fixed initializer.""" + model = h.make_model( + h.make_graph( + [h.make_node("Mul", ["x", "valid"], ["y"])], + "invalid_late_weight", + [_value("x", [2])], + [_value("y", [2])], + [ + nh.from_array(np.asarray([1, 2], np.float32), "valid"), + T(name="invalid", data_type=T.FLOAT, dims=[2], float_data=[1]), + ], + ), + opset_imports=[h.make_opsetid("", 19)], + ) + model.ir_version = 9 + before = model.SerializeToString() + with pytest.raises(ValueError): + _apply_force_fp16_activations(model) + assert model.SerializeToString() == before + + +@pytest.mark.parametrize( + ("opset", "split"), + [(opset, split) for opset in [13, 17, 18, 19] for split in ["none", "equal", "sizes", "nested"]] + + [(13, "empty"), (17, "empty")], +) +def test_fp16_qdq_upgrade_keeps_legacy_reduce_axes_valid(opset, split): + """Migrate reduction axes and live equal splits while retaining explicit Split configurations.""" + initializers = [nh.from_array(np.float32(0.1), "scale"), nh.from_array(np.int8(0), "zero")] + mean_inputs, mean_attributes = ["x"], {"keepdims": 1} + if opset < 18: + mean_attributes["axes"] = [1] + else: + mean_inputs.append("axes") + initializers.append(nh.from_array(np.array([1], np.int64), "axes")) + nodes = [ + h.make_node("ReduceMean", mean_inputs, ["mean"], **mean_attributes), + h.make_node("QuantizeLinear", ["mean", "scale", "zero"], ["q"]), + h.make_node( + "DequantizeLinear", ["q", "scale", "zero"], ["y" if split == "none" else "decoded"] + ), + ] + outputs = [_value("y", [4, 1])] + if split != "none": + split_inputs = ["x" if split == "nested" else "decoded"] + attrs = {"axis": 0} + if split == "sizes": + split_inputs.append("sizes") + initializers.append(nh.from_array(np.array([1, 3], np.int64), "sizes")) + else: + if split == "empty": + split_inputs.append("") + if opset >= 18: + attrs["num_outputs"] = 2 + splitting = [ + h.make_node("Split", split_inputs, ["left", "right"], **attrs), + h.make_node( + "Concat", ["left", "right"], ["y" if split != "nested" else "branch_y"], axis=0 + ), + ] + if split == "nested": + branch = h.make_graph(splitting, "split_branch", [], [_value("branch_y", [4, 2])]) + initializers.append(nh.from_array(np.asarray(True), "condition")) + nodes.append( + h.make_node("If", ["condition"], ["y"], then_branch=branch, else_branch=branch) + ) + outputs = [_value("y", [4, 2]), _value("decoded", [4, 1])] + else: + nodes.extend(splitting) + model = _model( + nodes, + [_value("x", [4, 2])], + outputs, + initializers, + opset=opset, + ) + onnx.checker.check_model(model, full_check=True) + feeds = {"x": np.repeat(np.asarray([0.2, 0.4, 0.6, 0.8], np.float32), 2).reshape(4, 2)} + expected = _checked_cpu_outputs(model, feeds) + changed = _apply_force_fp16_activations(model) + onnx.checker.check_model(changed, full_check=True) + for actuals in [ + ReferenceEvaluator(changed).run(None, feeds), + _checked_cpu_outputs(changed, feeds), + ]: + for actual, reference in zip(actuals, expected): + np.testing.assert_allclose(actual, reference, rtol=2e-3, atol=1e-4) + + +def test_fp16_lpnorm_keeps_opaque_function_native(): + """Do not interpret an opaque LpNormalization name as standard normalization.""" + model = _model( + [ + h.make_node("LpNormalization", ["x"], ["norm"], p=2, axis=-1), + h.make_node("Mul", ["norm", "gamma"], ["y"]), + ], + [_value("x", [1, 1, 1, 2])], + [_value("y", [1, 1, 1, 2])], + [nh.from_array(np.ones([1, 1, 1, 2], np.float32), "gamma")], + ) + model.graph.node[0].domain = "opaque" + model.functions.append( + h.make_function( + "opaque", + "LpNormalization", + ["x"], + ["y"], + [h.make_node("Neg", ["x"], ["y"])], + [h.make_opsetid("", 19)], + attributes=["p", "axis"], + ) + ) + model.opset_import.append(h.make_opsetid("opaque", 1)) + before = model.SerializeToString() + changed = _apply_fp16_lpnorm(model) + assert changed.SerializeToString() == before + _checked_cpu_outputs(changed, {"x": np.asarray([[[[3, 4]]]], np.float32)}) + + +def test_fp16_boundary_cast_names_are_allocated(): + """Keep generated Cast node names distinct from existing nodes and tensor names.""" + model = _model( + [ + h.make_node( + "Constant", + [], + ["integer"], + name="x_to_fp16", + value=nh.from_array(np.asarray([1], np.int8)), + ), + h.make_node("Relu", ["x"], ["y"]), + ], + [_value("x", [2])], + [_value("y", [2])], + [], + ) + changed = _apply_force_fp16_activations(model) + names = [node.name for node in changed.graph.node if node.name] + assert len(names) == len(set(names)) + _checked_cpu_outputs(changed, {"x": np.asarray([-1, 2], np.float32)}) + + +def test_fp16_lpnorm_keeps_unrelated_constant_object(): + """Retain a large-payload owner while replacing a private normalization node.""" + model = _model( + [ + h.make_node( + "Constant", [], ["integer"], value=nh.from_array(np.arange(64, dtype=np.uint8)) + ), + h.make_node("LpNormalization", ["x"], ["norm"], p=2, axis=-1), + h.make_node("Mul", ["norm", "gamma"], ["y"]), + ], + [_value("x", [1, 1, 1, 2])], + [_value("y", [1, 1, 1, 2])], + [nh.from_array(np.ones([1, 1, 1, 2], np.float32), "gamma")], + ) + constant = model.graph.node[0] + changed = _apply_fp16_lpnorm(model) + assert next(node for node in changed.graph.node if node.op_type == "Constant") is constant + onnx.checker.check_model(changed, full_check=True) + + +def test_fp16_keeps_opaque_float_function_type_contract(): + """Retain a local function's FP32 region and cast only its downstream standard use.""" + model = _model( + [h.make_node("Identity", ["x"], ["middle"]), h.make_node("Relu", ["middle"], ["y"])], + [_value("x", [2])], + [_value("y", [2])], + [], + ) + model.graph.node[0].domain = "opaque" + model.functions.append( + h.make_function( + "opaque", + "Identity", + ["x"], + ["out"], + [ + h.make_node("Constant", [], ["bias"], value=nh.from_array(np.float32(0.25))), + h.make_node("Add", ["x", "bias"], ["out"]), + ], + [h.make_opsetid("", 19)], + ) + ) + model.opset_import.append(h.make_opsetid("opaque", 1)) + feeds = {"x": np.asarray([0.5, -1], np.float32)} + expected = _checked_cpu_outputs(model, feeds)[0] + changed = _apply_force_fp16_activations(model) + np.testing.assert_array_equal(_checked_cpu_outputs(changed, feeds)[0], expected) + + +def test_fp16_retains_captured_values_and_control_flow_types(): + """Keep a captured FP32 result available to unchanged nested graph bodies.""" + branch = h.make_graph( + [h.make_node("Identity", ["middle"], ["branch_value"])], + "capture", + [], + [_value("branch_value", [2])], + ) + model = _model( + [ + h.make_node("Relu", ["x"], ["middle"]), + h.make_node("Neg", ["middle"], ["y"]), + h.make_node("If", ["condition"], ["captured"], then_branch=branch, else_branch=branch), + ], + [_value("x", [2]), _value("condition", [], T.BOOL)], + [_value("y", [2]), _value("captured", [2])], + [], + ) + feeds = {"x": np.asarray([1, -2], np.float32), "condition": np.asarray(True)} + expected = _checked_cpu_outputs(model, feeds) + changed = _apply_force_fp16_activations(model) + for actual, wanted in zip(_checked_cpu_outputs(changed, feeds), expected): + np.testing.assert_array_equal(actual, wanted) + + +def test_fp16_float_constant_keeps_valid_consumer_types(): + """Give an unfurled FLOAT Constant a private FP16 data view.""" + model = _model( + [ + h.make_node( + "Constant", [], ["bias"], value=nh.from_array(np.asarray([0.5, 0.25], np.float32)) + ), + h.make_node("Add", ["x", "bias"], ["y"]), + ], + [_value("x", [2])], + [_value("y", [2])], + [], + ) + feeds = {"x": np.asarray([1, 2], np.float32)} + expected = _checked_cpu_outputs(model, feeds)[0] + changed = _apply_force_fp16_activations(model) + np.testing.assert_array_equal(_checked_cpu_outputs(changed, feeds)[0], expected) + + +def _terminal_model(channels=8193, *, groups=1, quantized=False, batch=1): + """Build the exact terminal projection pattern with independently checkable values.""" + data_channels = groups + weights = nh.from_array( + np.ones((channels, 1, 1, 1), np.uint8) + if quantized + else np.full((channels, 1, 1, 1), 0.25, np.float32), + "weight", + ) + nodes = [] + initializers = [weights, nh.from_array(np.asarray([0], np.int64), "axes")] + if quantized: + initializers.extend( + [nh.from_array(np.float32(0.25), "scale"), nh.from_array(np.uint8(0), "zero")] + ) + nodes.append(h.make_node("DequantizeLinear", ["weight", "scale", "zero"], ["dequantized"])) + nodes.extend( + [ + h.make_node( + "Conv", ["x", "dequantized" if quantized else "weight"], ["conv"], group=groups + ), + h.make_node("Transpose", ["conv"], ["transposed"], perm=[0, 3, 2, 1]), + h.make_node("Concat", ["transposed"], ["concat1"], axis=0), + h.make_node("ReduceSum", ["concat1", "axes"], ["reduced"], keepdims=1), + h.make_node("Concat", ["reduced"], ["concat2"], axis=2), + h.make_node("Squeeze", ["concat2", "axes"], ["y"]), + ] + ) + return _model( + nodes, + [_value("x", [batch, data_channels, 1, 1])], + [_value("y", [1, 1, channels])], + initializers, + ) + + +@pytest.mark.parametrize("quantized", [False, True]) +def test_terminal_projection_preserves_nontrivial_batch_reduction(quantized): + """Retain batch summation instead of replacing it with singleton projection chunks.""" + model = _terminal_model(batch=2, quantized=quantized) + feeds = {"x": np.asarray([2, 6], np.float32).reshape(2, 1, 1, 1)} + expected = _checked_cpu_outputs(model, feeds)[0] + converted = _apply_force_fp16_activations(model) + assert sum(node.op_type == "Conv" for node in converted.graph.node) == 1 + np.testing.assert_array_equal(_checked_cpu_outputs(converted, feeds)[0], expected) + + +@pytest.mark.parametrize( + "owner", ["grouped", "runtime_weight", "public_axes", "legacy", "constant"] +) +def test_terminal_projection_preserves_native_owners_and_group_layout(owner): + """Chunk only private fixed projections and preserve control owners and grouped channels.""" + grouped = owner == "grouped" + model = _terminal_model( + channels=8194 if grouped else 8193, + groups=2 if grouped else 1, + quantized=owner == "runtime_weight", + ) + if owner == "runtime_weight": + model.graph.input.append(_value("weight", [8193, 1, 1, 1], T.UINT8)) + if owner == "public_axes": + model.graph.output.append(_value("axes", [1], T.INT64)) + if owner == "legacy": + model.ir_version = 3 + model.graph.input.extend( + _value(t.name, list(t.dims), t.data_type) for t in model.graph.initializer + ) + constant = None + if owner == "constant": + model.graph.node.insert( + 0, + h.make_node( + "Constant", [], ["integer"], value=nh.from_array(np.arange(64, dtype=np.uint8)) + ), + ) + constant = model.graph.node[0] + data = ( + np.asarray([[[[2]], [[8]]]], np.float32) if grouped else np.asarray([[[[2]]]], np.float32) + ) + feeds = {"x": data} + onnx.checker.check_model(model, full_check=True) + expected = ReferenceEvaluator(model).run(None, feeds) + override = ( + {"x": data, "weight": np.full((8193, 1, 1, 1), 2, np.uint8)} + if owner == "runtime_weight" + else None + ) + override_expected = ( + ReferenceEvaluator(model).run(None, override) if override is not None else None + ) + changed = _apply_force_fp16_activations(model) + if owner in {"grouped", "runtime_weight", "public_axes"}: + assert sum(node.op_type == "Conv" for node in changed.graph.node) == 1 + if owner == "legacy": + assert not {"weight", "axes"} & {value.name for value in changed.graph.input} + if constant is not None: + assert next(node for node in changed.graph.node if node.op_type == "Constant") is constant + onnx.checker.check_model(changed, full_check=True) + for actual, wanted in zip(ReferenceEvaluator(changed).run(None, feeds), expected): + np.testing.assert_array_equal(actual, wanted) + if override is not None: + for actual, wanted in zip( + ReferenceEvaluator(changed).run(None, override), override_expected + ): + np.testing.assert_array_equal(actual, wanted) + + +def test_fp16_public_legacy_initializer_uses_private_declaration(): + """Rename the legacy initializer declaration together with its private FP16 payload.""" + model = _model( + [], + [_value("weight", [2])], + [_value("weight", [2])], + [nh.from_array(np.asarray([1, 2], np.float32), "weight")], + ) + model.ir_version = 3 + changed = _apply_force_fp16_activations(model) + onnx.checker.check_model(changed, full_check=True) + np.testing.assert_array_equal( + _checked_cpu_outputs(changed, {})[0], np.asarray([1, 2], np.float32) + ) + + +@pytest.mark.parametrize("legacy", [False, True]) +def test_fp16_qdq_upgrade_migrates_nested_reduce_axes(legacy): + """Migrate axes in nested graphs that inherit the upgraded main operator set.""" + branch = h.make_graph( + [h.make_node("ReduceMean", ["x"], ["branch_value"], axes=[1], keepdims=1)], + "nested_reduce", + [], + [_value("branch_value", [2, 1])], + ) + model = _model( + [ + h.make_node("QuantizeLinear", ["x", "scale", "zero"], ["q"]), + h.make_node("DequantizeLinear", ["q", "scale", "zero"], ["y"]), + h.make_node("If", ["condition"], ["reduced"], then_branch=branch, else_branch=branch), + ], + [_value("x", [2, 2]), _value("condition", [], T.BOOL)], + [_value("y", [2, 2]), _value("reduced", [2, 1])], + [nh.from_array(np.float32(0.1), "scale"), nh.from_array(np.int8(0), "zero")], + opset=17, + ) + if legacy: + model.ir_version = 3 + model.graph.input.extend( + _value(t.name, list(t.dims), t.data_type) for t in model.graph.initializer + ) + feeds = {"x": np.asarray([[0.2, 0.2], [0.4, 0.4]], np.float32), "condition": np.asarray(True)} + expected = _checked_cpu_outputs(model, feeds) + changed = _apply_force_fp16_activations(model) + onnx.checker.check_model(changed, full_check=True) + for got, wanted in zip(ReferenceEvaluator(changed).run(None, feeds), expected): + np.testing.assert_allclose(got, wanted, rtol=2e-3, atol=1e-4) + + +def test_fp16_legacy_qdq_schema_refuses_before_mutation(): + """Refuse legacy operator-set migration that needs adapters beyond the supported range.""" + model = _model( + [ + h.make_node("Clip", ["x"], ["clipped"], min=0.0, max=1.0), + h.make_node("QuantizeLinear", ["clipped", "scale", "zero"], ["y"]), + ], + [_value("x", [2])], + [_value("y", [2], T.INT8)], + [nh.from_array(np.float32(0.1), "scale"), nh.from_array(np.int8(0), "zero")], + opset=10, + ) + onnx.checker.check_model(model, full_check=True) + before = model.SerializeToString() + with pytest.raises(ValueError, match="opset"): + _apply_force_fp16_activations(model) + assert model.SerializeToString() == before + + +@pytest.mark.parametrize("operation", ["Range", "NonMaxSuppression"]) +def test_fp16_preserves_operator_schemas_without_half_inputs(operation): + """Keep mandatory FP32 schema bindings valid instead of forcing unsupported HALF inputs.""" + if operation == "Range": + model = _model( + [h.make_node("Range", ["start", "limit", "delta"], ["y"])], + [_value("start", [])], + [_value("y", [4])], + [nh.from_array(np.float32(4), "limit"), nh.from_array(np.float32(1), "delta")], + ) + feeds = {"start": np.asarray(0, np.float32)} + else: + model = _model( + [h.make_node("NonMaxSuppression", ["boxes", "scores", "count", "iou"], ["y"])], + [_value("boxes", [1, 2, 4]), _value("scores", [1, 1, 2])], + [_value("y", [None, 3], T.INT64)], + [nh.from_array(np.int64(1), "count"), nh.from_array(np.float32(0.5), "iou")], + ) + feeds = { + "boxes": np.asarray([[[0, 0, 1, 1], [2, 2, 3, 3]]], np.float32), + "scores": np.asarray([[[0.8, 0.5]]], np.float32), + } + expected = _checked_cpu_outputs(model, feeds)[0] + changed = _apply_force_fp16_activations(model) + np.testing.assert_array_equal(_checked_cpu_outputs(changed, feeds)[0], expected) + + +@pytest.mark.parametrize("operation", ["ConstantOfShape", "RandomNormal", "EyeLike"]) +def test_fp16_preserves_explicit_generated_float_types(operation): + """Keep generated FLOAT values typed correctly and provide private HALF consumer views.""" + inputs = [] + initializers = [] + feeds = {} + shape = [2, 2] + if operation == "ConstantOfShape": + initializers = [nh.from_array(np.asarray(shape, np.int64), "shape")] + generated = h.make_node( + operation, ["shape"], ["generated"], value=nh.from_array(np.asarray([0.25], np.float32)) + ) + elif operation == "RandomNormal": + generated = h.make_node(operation, [], ["generated"], shape=shape, dtype=T.FLOAT, seed=42.0) + else: + inputs = [_value("x", shape)] + feeds = {"x": np.zeros(shape, np.float32)} + generated = h.make_node(operation, ["x"], ["generated"], dtype=T.FLOAT) + model = _model( + [generated, h.make_node("Relu", ["generated"], ["y"])], + inputs, + [_value("y", shape)], + initializers, + ) + expected_cpu = _checked_cpu_outputs(model, feeds)[0] + expected = ReferenceEvaluator(model).run(None, feeds)[0].astype(np.float16).astype(np.float32) + changed = _apply_force_fp16_activations(model) + onnx.checker.check_model(changed, full_check=True) + np.testing.assert_array_equal(ReferenceEvaluator(changed).run(None, feeds)[0], expected) + # CPU fallback can retain FP32 arithmetic even when the graph requests HALF. + np.testing.assert_allclose( + _checked_cpu_outputs(changed, feeds)[0], expected_cpu, rtol=1e-3, atol=1e-3 + ) + + +def test_fp16_lpnorm_allocates_support_node_names(): + """Reserve normalization support-node names across existing graph owners.""" + model = _model( + [ + h.make_node( + "Constant", + [], + ["integer"], + name="normalize_square", + value=nh.from_array(np.asarray([1], np.int8)), + ), + h.make_node("LpNormalization", ["x"], ["norm"], name="normalize", p=2, axis=-1), + h.make_node("Mul", ["norm", "gamma"], ["y"]), + ], + [_value("x", [1, 1, 1, 2])], + [_value("y", [1, 1, 1, 2])], + [nh.from_array(np.ones([1, 1, 1, 2], np.float32), "gamma")], + ) + changed = _apply_fp16_lpnorm(model) + names = [node.name for node in changed.graph.node if node.name] + assert len(names) == len(set(names)) + onnx.checker.check_model(changed, full_check=True) + + +@pytest.mark.parametrize("dtype", [T.UINT8, T.FLOAT]) +def test_fp16_skips_conversion_for_preserved_runtime_defaults_only(dtype): + """Skip weight conversion when the model contains only preserved defaults.""" + numpy_dtype = np.uint8 if dtype == T.UINT8 else np.float32 + model = _model( + [h.make_node("Cast", ["weight"], ["y"], to=T.FLOAT)], + [_value("weight", [2], dtype)], + [_value("y", [2])], + [nh.from_array(np.asarray([1, 2], numpy_dtype), "weight")], + ) + weight = model.graph.initializer[0] + changed = _apply_force_fp16_activations(model) + assert next(t for t in changed.graph.initializer if t.name == "weight") is weight + np.testing.assert_array_equal( + _checked_cpu_outputs(changed, {})[0], np.asarray([1, 2], np.float32) + ) + np.testing.assert_array_equal( + _checked_cpu_outputs(changed, {"weight": np.asarray([2, 3], numpy_dtype)})[0], + np.asarray([2, 3], np.float32), + ) + + +@pytest.mark.parametrize("operation", ["QuantizeLinear", "DequantizeLinear"]) +def test_fp16_supported_contrib_quantizers_keep_half_interface(operation): + """Retain the supported contrib quantizers' HALF scales without treating them as opaque.""" + initializers = [nh.from_array(np.float32(0.5), "scale"), nh.from_array(np.int8(0), "zero")] + if operation == "QuantizeLinear": + nodes = [h.make_node(operation, ["x", "scale", "zero"], ["y"])] + inputs = [_value("x", [2])] + outputs = [_value("y", [2], T.INT8)] + feeds = {"x": np.asarray([-1, 2], np.float32)} + else: + initializers.append(nh.from_array(np.asarray([-2, 4], np.int8), "quantized")) + nodes = [ + h.make_node(operation, ["quantized", "scale", "zero"], ["middle"]), + h.make_node("Relu", ["middle"], ["y"]), + ] + inputs = [] + outputs = [_value("y", [2])] + feeds = {} + model = _model(nodes, inputs, outputs, initializers, opset=17) + model.graph.node[0].domain = "com.microsoft" + model.opset_import.append(h.make_opsetid("com.microsoft", 1)) + expected = _checked_cpu_outputs(model, feeds)[0] + changed = _apply_force_fp16_activations(model) + assert next(t for t in changed.graph.initializer if t.name == "scale").data_type == T.FLOAT16 + assert next(op for op in changed.opset_import if not op.domain).version == 17 + np.testing.assert_array_equal(_checked_cpu_outputs(changed, feeds)[0], expected) + + +def test_terminal_projection_chunks_supported_contrib_dequantization(): + """Keep the supported contrib dequantizer on each private bounded weight chunk.""" + model = _terminal_model(quantized=True) + model.graph.node[0].domain = "com.microsoft" + model.opset_import.append(h.make_opsetid("com.microsoft", 1)) + feeds = {"x": np.asarray([[[[2]]]], np.float32)} + expected = _checked_cpu_outputs(model, feeds)[0] + changed = _apply_force_fp16_activations(model) + assert sum(node.op_type == "Conv" for node in changed.graph.node) == 2 + np.testing.assert_array_equal(_checked_cpu_outputs(changed, feeds)[0], expected) + + +def test_fp16_supported_contrib_gelu_keeps_upstream_half_precision(): + """Keep a supported contrib activation and its standard normalization in HALF.""" + model = h.make_model( + h.make_graph( + [ + h.make_node("LayerNormalization", ["x", "scale", "bias"], ["normalized"], axis=-1), + h.make_node("Gelu", ["normalized"], ["y"], domain="com.microsoft"), + ], + "contrib_gelu", + [_value("x", [2, 2])], + [_value("y", [2, 2])], + [ + nh.from_array(np.ones(2, np.float32), "scale"), + nh.from_array(np.zeros(2, np.float32), "bias"), + ], + ), + opset_imports=[h.make_opsetid("", 17), h.make_opsetid("com.microsoft", 1)], + ) + model.ir_version = 9 + model = onnx.shape_inference.infer_shapes(model) + feeds = {"x": np.asarray([[1, 3], [2, 4]], np.float32)} + expected = _checked_cpu_outputs(model, feeds)[0] + changed = _apply_force_fp16_activations(model) + assert all(t.data_type == T.FLOAT16 for t in changed.graph.initializer) + assert next(op for op in changed.opset_import if not op.domain).version == 17 + np.testing.assert_allclose( + _checked_cpu_outputs(changed, feeds)[0], expected, rtol=2e-3, atol=1e-4 + ) + + +@pytest.mark.parametrize("operation", ["Mul", "Add"]) +@pytest.mark.parametrize("control_reader", [False, True]) +def test_public_squeeze_has_private_fp16_readers_and_native_controls(operation, control_reader): + """Preserve public FP32 Squeeze values while keeping internal arithmetic well typed.""" + nodes = [ + h.make_node("Relu", ["x"], ["r4"]), + h.make_node("Squeeze", ["r4", "axes"], ["public"]), + h.make_node(operation, ["public", "weight"], ["result"]), + h.make_node("Cast", ["public"], ["half"], to=T.FLOAT16), + h.make_node("Cast", ["public"], ["integer"], to=T.INT32), + ] + inputs = [_value("x", [1, 1, 1, 2])] + outputs = [ + _value("public", [2]), + _value("result", [2]), + _value("half", [2], T.FLOAT16), + _value("integer", [2], T.INT32), + ] + if control_reader: + nodes.append(h.make_node("Resize", ["image", "", "public"], ["resized"], mode="nearest")) + inputs.append(_value("image", [1, 2])) + outputs.append(_value("resized", [1, 4])) + model = _model( + nodes, + inputs, + outputs, + [ + nh.from_array(np.array([0, 1, 2], np.int64), "axes"), + nh.from_array(np.array([2, 3], np.float32), "weight"), + ], + ) + feeds = {"x": np.array([[[[1, 2]]]], np.float32)} + if control_reader: + feeds["image"] = np.array([[4, 5]], np.float32) + expected = ReferenceEvaluator(model).run(None, feeds) + changed = _apply_force_fp16_activations(model) + onnx.checker.check_model(changed, full_check=True) + options = ort.SessionOptions() + options.intra_op_num_threads = 1 + actual = ort.InferenceSession( + changed.SerializeToString(), options, providers=["CPUExecutionProvider"] + ).run(None, feeds) + for left, right in zip(actual, expected): + np.testing.assert_array_equal(left, right) + assert actual[0].dtype == np.float32 + casts = {node.output[0]: node for node in changed.graph.node if node.op_type == "Cast"} + assert casts["half"].input[0] == "public" + if control_reader: + resize = next(node for node in changed.graph.node if node.op_type == "Resize") + assert resize.input[2] == "public" + + +@pytest.mark.parametrize("existing_output", [False, True]) +def test_fp16_public_conversion_refuses_invalid_control_types_before_saving( + tmp_path, existing_output +): + """Reject half-precision Resize scales without creating or replacing an output file.""" + model = h.make_model( + h.make_graph( + [h.make_node("Resize", ["x", "", "scales"], ["y"], mode="nearest")], + "invalid_resize_scales", + [_value("x", [1, 1, 1, 2])], + [_value("y", [1, 1, 1, 4])], + [nh.from_array(np.array([1, 1, 1, 2], np.float16), "scales")], + ), + opset_imports=[h.make_opsetid("", 19)], + ir_version=9, + ) + onnx.checker.check_model(model) + with pytest.raises((onnx.shape_inference.InferenceError, onnx.checker.ValidationError)): + onnx.checker.check_model(model, full_check=True) + source, output = tmp_path / "source.onnx", tmp_path / "output.onnx" + onnx.save(model, source) + original = source.read_bytes() + previous = b"existing output must survive validation failure" + if existing_output: + output.write_bytes(previous) + with pytest.raises((onnx.shape_inference.InferenceError, onnx.checker.ValidationError)): + make_dla_compatible( + str(source), + str(output), + use_external_data=False, + verbose=False, + force_fp16_activations=True, + ) + assert source.read_bytes() == original + assert output.read_bytes() == previous if existing_output else not output.exists() + + +@pytest.mark.parametrize("source_type", [T.INT32, T.INT64]) +@pytest.mark.parametrize("operation", ["Mul", "Add"]) +@pytest.mark.parametrize("control_reader", [False, True]) +def test_public_float_cast_keeps_boundaries_and_typed_internal_readers( + source_type, operation, control_reader +): + """Share a half view for arithmetic while preserving public values and explicit Cast readers.""" + nodes = [ + h.make_node("Cast", ["x"], ["public"], to=T.FLOAT), + h.make_node(operation, ["public", "weight"], ["result"]), + h.make_node("Cast", ["public"], ["half"], to=T.FLOAT16), + h.make_node("Cast", ["public"], ["alias"], to=T.FLOAT), + h.make_node("Cast", ["public"], ["integer"], to=T.INT64), + ] + inputs = [_value("x", [2], source_type)] + outputs = [ + _value("public", [2]), + _value("result", [2]), + _value("half", [2], T.FLOAT16), + _value("alias", [2]), + _value("integer", [2], T.INT64), + ] + dtype = np.int32 if source_type == T.INT32 else np.int64 + feeds = {"x": np.array([1, 2] if control_reader else [2049, 2051], dtype)} + if control_reader: + nodes.append(h.make_node("Resize", ["image", "", "public"], ["resized"], mode="nearest")) + inputs.append(_value("image", [1, 2])) + outputs.append(_value("resized", [1, 4])) + feeds["image"] = np.array([[4, 5]], np.float32) + weight = np.array([2, 3], np.float32) + model = _model(nodes, inputs, outputs, [nh.from_array(weight, "weight")]) + expected = ReferenceEvaluator(model).run(None, feeds) + expected[1] = ( + (expected[2] * weight.astype(np.float16)) + if operation == "Mul" + else (expected[2] + weight.astype(np.float16)) + ).astype(np.float32) + changed = _apply_force_fp16_activations(model) + onnx.checker.check_model(changed, full_check=True) + actual = ReferenceEvaluator(changed).run(None, feeds) + # ORT CPU promotes FP16 arithmetic; use exact values for the runtime smoke check. + cpu_feeds = dict(feeds, x=np.array([1, 2], dtype)) + for left, right in zip( + _checked_cpu_outputs(changed, cpu_feeds), + ReferenceEvaluator(changed).run(None, cpu_feeds), + ): + np.testing.assert_array_equal(left, right) + for left, right in zip(actual, expected): + np.testing.assert_array_equal(left, right) + assert actual[0].dtype == np.float32 + casts = {node.output[0]: node for node in changed.graph.node if node.op_type == "Cast"} + assert casts["half"].input[0] == "public" + if not control_reader: + assert casts["alias"].input[0] == "public" + assert casts["integer"].input[0] == "public" + else: + resize = next(node for node in changed.graph.node if node.op_type == "Resize") + assert resize.input[2] == "public" + + +@pytest.mark.parametrize("source_type", [T.INT32, T.INT64]) +def test_public_fp16_conversion_supports_shared_float_cast_outputs(tmp_path, source_type): + """Convert and reload a valid shared integer-to-FLOAT graph through the real public API.""" + shape = [1, 1, 1, 2] + model = _model( + [ + h.make_node("Cast", ["x"], ["public"], to=T.FLOAT), + h.make_node("Mul", ["public", "weight"], ["result"]), + ], + [_value("x", shape, source_type)], + [_value("public", shape), _value("result", shape)], + [nh.from_array(np.array([2, 3], np.float32).reshape(shape), "weight")], + ) + source, output = tmp_path / "source.onnx", tmp_path / "output.onnx" + onnx.save(model, source) + make_dla_compatible( + str(source), + str(output), + use_external_data=False, + verbose=False, + force_fp16_activations=True, + ) + changed = onnx.load(output) + onnx.checker.check_model(changed, full_check=True) + dtype = np.int32 if source_type == T.INT32 else np.int64 + values = np.array([2049, 2051], dtype).reshape(shape) + actual = ReferenceEvaluator(changed).run(None, {"x": values}) + cpu_feeds = {"x": np.array([1, 2], dtype).reshape(shape)} + for left, right in zip( + _checked_cpu_outputs(changed, cpu_feeds), + ReferenceEvaluator(changed).run(None, cpu_feeds), + ): + np.testing.assert_array_equal(left, right) + np.testing.assert_array_equal(actual[0], values.astype(np.float32)) + expected = (values.astype(np.float16) * np.array([2, 3], np.float16).reshape(shape)).astype( + np.float32 + ) + np.testing.assert_array_equal(actual[1], expected)