diff --git a/modelopt/onnx/autocast/precisionconverter.py b/modelopt/onnx/autocast/precisionconverter.py index da45cd3ecde..4ea05bcd57f 100644 --- a/modelopt/onnx/autocast/precisionconverter.py +++ b/modelopt/onnx/autocast/precisionconverter.py @@ -139,6 +139,10 @@ def __init__( self.original_network_io.update( {io.name: io.type.tensor_type.elem_type for io in self.model.graph.output} ) + self.original_network_io_metadata = { + "input": [deepcopy(io) for io in self.model.graph.input], + "output": [deepcopy(io) for io in self.model.graph.output], + } self.min_opset = min_opset self.max_ir_version = max_ir_version self.trt_plugins = trt_plugins @@ -276,6 +280,8 @@ def convert( # Remove redundant casts self._cleanup() + self._restore_original_io_metadata() + self._sanity_check() return self.model @@ -289,6 +295,120 @@ def _ensure_types_are_defined(self): def _propagate_types_shapes_custom_ops(self, model): """Propagate types and shapes after insertion of 'Cast' nodes or other graph modifications.""" logger.info("Propagating tensor shapes and types in model with custom ops.") + + def _get_shape(tensor): + if isinstance(tensor, gs.Constant): + return list(tensor.values.shape) + if tensor.shape is None: + return None + return list(tensor.shape) + + def _get_const_values(tensor): + if isinstance(tensor, gs.Constant): + return tensor.values + if tensor.inputs and tensor.inputs[0].op == "Constant": + return tensor.inputs[0].attrs["value"].values + return None + + def _get_int_attr(node, attr_name, default): + value = node.attrs.get(attr_name, default) + return value if isinstance(value, int) else None + + def _infer_gathernd_op_shape(node): + if node.op != "GatherND" or len(node.inputs) < 2: + return None + + data_shape = _get_shape(node.inputs[0]) + indices_shape = _get_shape(node.inputs[1]) + if not data_shape or not indices_shape: + return None + + index_rank = indices_shape[-1] + batch_dims = node.attrs.get("batch_dims", 0) + if not isinstance(index_rank, int) or not isinstance(batch_dims, int): + return None + + suffix_start = batch_dims + index_rank + if suffix_start > len(data_shape): + return None + return indices_shape[:-1] + data_shape[suffix_start:] + + def _infer_gather_op_shape(node): + if node.op != "Gather" or len(node.inputs) < 2: + return None + + data_shape = _get_shape(node.inputs[0]) + indices_shape = _get_shape(node.inputs[1]) + if data_shape is None or indices_shape is None: + return None + + axis = _get_int_attr(node, "axis", 0) + if axis is None: + return None + if axis < 0: + axis += len(data_shape) + if axis < 0 or axis >= len(data_shape): + return None + + return data_shape[:axis] + indices_shape + data_shape[axis + 1 :] + + def _infer_unsqueeze_op_shape(node): + if node.op != "Unsqueeze" or len(node.inputs) < 2: + return None + + data_shape = _get_shape(node.inputs[0]) + axes = _get_const_values(node.inputs[1]) + if data_shape is None or axes is None: + return None + + axes = [int(axis) for axis in np.asarray(axes).flatten()] + output_rank = len(data_shape) + len(axes) + normalized_axes = [] + for axis in axes: + if axis < 0: + axis += output_rank + if axis < 0 or axis >= output_rank: + return None + normalized_axes.append(axis) + + output_shape = list(data_shape) + for axis in sorted(normalized_axes): + output_shape.insert(axis, 1) + return output_shape + + def _infer_shape_op_shape(node): + if node.op != "Shape" or not node.inputs: + return None + + data_shape = _get_shape(node.inputs[0]) + if data_shape is None: + return None + + rank = len(data_shape) + start = _get_int_attr(node, "start", 0) + end = _get_int_attr(node, "end", rank) + if start is None or end is None: + return None + if start < 0: + start += rank + if end < 0: + end += rank + start = min(max(start, 0), rank) + end = min(max(end, 0), rank) + return [max(end - start, 0)] + + def _infer_standard_op_shape(node): + for infer_shape in ( + _infer_gathernd_op_shape, + _infer_gather_op_shape, + _infer_unsqueeze_op_shape, + _infer_shape_op_shape, + ): + shape = infer_shape(node) + if shape is not None: + return shape + return None + graph = gs.import_onnx(model) traversed_tensors = [] @@ -397,12 +517,14 @@ def _propagate_cast_type_through_nodes(node, np_type, iter=1): out.dtype = np_type # Set the output shape - if not out.shape: - if isinstance(inp, gs.Constant): + if out.shape is None: + if (shape := _infer_standard_op_shape(node)) is not None: + out.shape = shape + elif isinstance(inp, gs.Constant): out.shape = inp.values.shape elif inp.inputs and inp.inputs[0].op == "Constant": out.shape = inp.inputs[0].attrs["value"].values.shape - elif inp.shape: + elif node.op in self.custom_ops and inp.shape: out.shape = inp.shape # Propagate tensor types to the children nodes (until another Cast or Q node is met) @@ -410,6 +532,32 @@ def _propagate_cast_type_through_nodes(node, np_type, iter=1): return gs.export_onnx(graph) + def _restore_original_io_metadata(self) -> None: + """Preserve complete public I/O metadata when keep_io_types=True.""" + if not self.keep_io_types: + return + + for io_field in ("input", "output"): + current_values = list(getattr(self.model.graph, io_field)) + original_values = self.original_network_io_metadata[io_field] + current_names = [value.name for value in current_values] + original_names = [value.name for value in original_values] + if current_names != original_names: + raise RuntimeError( + f"Cannot restore public graph {io_field} metadata because names changed: " + f"{original_names} -> {current_names}" + ) + for current_value, original_value in zip(current_values, original_values, strict=True): + if ( + current_value.type.tensor_type.elem_type + != original_value.type.tensor_type.elem_type + ): + raise RuntimeError( + f"Cannot restore public graph {io_field} metadata for {current_value.name}: " + "element type changed" + ) + current_value.CopyFrom(original_value) + def _is_bf16(self, type: PrecisionTypes = None) -> bool: if type is None: type = self.low_precision_type diff --git a/modelopt/onnx/utils.py b/modelopt/onnx/utils.py index bc37c4a9333..04be94a743c 100644 --- a/modelopt/onnx/utils.py +++ b/modelopt/onnx/utils.py @@ -1532,6 +1532,21 @@ def _convert_constant_values(constant_node: onnx.NodeProto, cast_node: onnx.Node break +def _sync_value_info_elem_type(graph: onnx.GraphProto, tensor_name: str, elem_type: int) -> None: + """Synchronize declarations for a tensor whose producer dtype changed.""" + for value_info in list(graph.value_info) + list(graph.input) + list(graph.output): + if value_info.name == tensor_name and value_info.type.HasField("tensor_type"): + value_info.type.tensor_type.elem_type = elem_type + + for node in graph.node: + for attr in node.attribute: + if attr.type == onnx.AttributeProto.GRAPH: + _sync_value_info_elem_type(attr.g, tensor_name, elem_type) + elif attr.type == onnx.AttributeProto.GRAPHS: + for subgraph in attr.graphs: + _sync_value_info_elem_type(subgraph, tensor_name, elem_type) + + def remove_redundant_casts(onnx_model: onnx.ModelProto) -> onnx.ModelProto: """Removes both sequential casts and casts that don't change precision. @@ -1571,6 +1586,9 @@ def remove_redundant_casts(onnx_model: onnx.ModelProto) -> onnx.ModelProto: assert len(cast_producers) == 1 and cast_producers[0].op_type == "Constant" constant_producer = cast_producers[0] _convert_constant_values(constant_producer, node) + _sync_value_info_elem_type( + onnx_model.graph, constant_producer.output[0], get_cast_to_type(node) + ) _bypass_cast_node(onnx_model, node) logger.debug(f"Found foldable Constant->Cast pattern, removing {node.name}") diff --git a/tests/unit/onnx/autocast/test_precisionconverter.py b/tests/unit/onnx/autocast/test_precisionconverter.py index 9872e0c4f2b..c3da9743564 100644 --- a/tests/unit/onnx/autocast/test_precisionconverter.py +++ b/tests/unit/onnx/autocast/test_precisionconverter.py @@ -1955,3 +1955,229 @@ def test_if_subgraph_outer_scope_type_preservation( assert len(else_x_info) > 0, "X value_info should be preserved in else branch" assert then_x_info[0].type.tensor_type.elem_type != onnx.TensorProto.UNDEFINED assert else_x_info[0].type.tensor_type.elem_type != onnx.TensorProto.UNDEFINED + + +@pytest.mark.parametrize("value_info_elem_type", [TensorProto.FLOAT, TensorProto.UNDEFINED]) +def test_folded_constant_cast_updates_value_info_type(value_info_elem_type): + const_tensor = numpy_helper.from_array( + np.array([1.0, 2.0], dtype=np.float32), name="const_value" + ) + const_node = helper.make_node( + "Constant", [], ["const_out"], name="const_node", value=const_tensor + ) + cast_node = helper.make_node( + "Cast", ["const_out"], ["cast_out"], name="cast_to_fp16", to=TensorProto.FLOAT16 + ) + identity_node = helper.make_node("Identity", ["cast_out"], ["Y"], name="identity") + + graph = helper.make_graph( + [const_node, cast_node, identity_node], + "constant_cast_value_info", + [], + [helper.make_tensor_value_info("Y", TensorProto.FLOAT16, [2])], + [], + value_info=[helper.make_tensor_value_info("const_out", value_info_elem_type, [2])], + ) + model = helper.make_model(graph, producer_name="constant_cast_value_info") + model.opset_import[0].version = 19 + model.ir_version = 10 + + folded = onnx_utils.remove_redundant_casts(model) + + assert [node.op_type for node in folded.graph.node] == ["Constant", "Identity"] + const_out = next(vi for vi in folded.graph.value_info if vi.name == "const_out") + assert const_out.type.tensor_type.elem_type == TensorProto.FLOAT16 + onnx.shape_inference.infer_shapes(folded, strict_mode=True, check_type=True) + + +def test_custom_op_mode_uses_schema_shape_for_standard_gathernd(): + data = helper.make_tensor_value_info("data", TensorProto.FLOAT, [1, 4, 2]) + plugin_in = helper.make_tensor_value_info("plugin_in", TensorProto.FLOAT, [1, 4, 2]) + indices_init = numpy_helper.from_array( + np.array([[[0, 0], [3, 1]]], dtype=np.int64), name="indices" + ) + custom_node = helper.make_node( + "FakeTensorRTPlugin", ["plugin_in"], ["plugin_out"], name="fake_plugin" + ) + gather_node = helper.make_node( + "GatherND", + ["data", "indices"], + ["last_token_embed"], + name="shape_changing_gathernd", + batch_dims=1, + ) + graph = helper.make_graph( + [custom_node, gather_node], + "custom_op_gathernd_shape", + [data, plugin_in], + [ + helper.make_tensor_value_info("plugin_out", TensorProto.FLOAT, [1, 4, 2]), + helper.make_tensor_value_info("last_token_embed", TensorProto.FLOAT, None), + ], + [indices_init], + ) + model = helper.make_model(graph, producer_name="custom_op_gathernd_shape") + model.opset_import[0].version = 19 + model.ir_version = 10 + value_info_map, initializer_map, node_to_init_map = utils.setup_mappings(model) + + converter = PrecisionConverter( + model, + value_info_map, + initializer_map, + node_to_init_map, + keep_io_types=True, + custom_ops={"FakeTensorRTPlugin"}, + ) + propagated = converter._propagate_types_shapes_custom_ops(model) + + output = next(vi for vi in propagated.graph.output if vi.name == "last_token_embed") + assert [dim.dim_value for dim in output.type.tensor_type.shape.dim] == [1, 2] + + +def test_custom_op_mode_preserves_scalar_gathernd_shape(): + data = helper.make_tensor_value_info("data", TensorProto.FLOAT, [4]) + plugin_in = helper.make_tensor_value_info("plugin_in", TensorProto.FLOAT, [4]) + indices_init = numpy_helper.from_array(np.array([2], dtype=np.int64), name="indices") + custom_node = helper.make_node( + "FakeTensorRTPlugin", ["plugin_in"], ["plugin_out"], name="fake_plugin" + ) + gather_node = helper.make_node( + "GatherND", + ["data", "indices"], + ["selected_scalar"], + name="scalar_gathernd", + ) + graph = helper.make_graph( + [custom_node, gather_node], + "custom_op_scalar_gathernd_shape", + [data, plugin_in], + [ + helper.make_tensor_value_info("plugin_out", TensorProto.FLOAT, [4]), + helper.make_tensor_value_info("selected_scalar", TensorProto.FLOAT, None), + ], + [indices_init], + ) + model = helper.make_model(graph, producer_name="custom_op_scalar_gathernd_shape") + model.opset_import[0].version = 19 + model.ir_version = 10 + value_info_map, initializer_map, node_to_init_map = utils.setup_mappings(model) + + converter = PrecisionConverter( + model, + value_info_map, + initializer_map, + node_to_init_map, + keep_io_types=True, + custom_ops={"FakeTensorRTPlugin"}, + ) + propagated = converter._propagate_types_shapes_custom_ops(model) + + output = next(vi for vi in propagated.graph.output if vi.name == "selected_scalar") + assert [dim.dim_value for dim in output.type.tensor_type.shape.dim] == [] + + +def test_custom_op_mode_uses_schema_shapes_for_standard_rank_changes(): + gather_data = helper.make_tensor_value_info("gather_data", TensorProto.FLOAT, [4]) + unsqueeze_data = helper.make_tensor_value_info("unsqueeze_data", TensorProto.FLOAT, [4]) + plugin_in = helper.make_tensor_value_info("plugin_in", TensorProto.FLOAT, [1]) + gather_index = numpy_helper.from_array(np.array(0, dtype=np.int64), "gather_index") + unsqueeze_axes = numpy_helper.from_array(np.array([0], dtype=np.int64), "unsqueeze_axes") + gather_node = helper.make_node( + "Gather", ["gather_data", "gather_index"], ["gather_y_pre_cast"], name="gather" + ) + gather_cast = helper.make_node( + "Cast", ["gather_y_pre_cast"], ["gather_y"], name="gather_cast", to=TensorProto.FLOAT + ) + unsqueeze_node = helper.make_node( + "Unsqueeze", + ["unsqueeze_data", "unsqueeze_axes"], + ["unsqueeze_y_pre_cast"], + name="unsqueeze", + ) + unsqueeze_cast = helper.make_node( + "Cast", + ["unsqueeze_y_pre_cast"], + ["unsqueeze_y"], + name="unsqueeze_cast", + to=TensorProto.FLOAT, + ) + custom_node = helper.make_node( + "FakePlugin", ["plugin_in"], ["plugin_y"], name="plugin", domain="test.plugins" + ) + graph = helper.make_graph( + [gather_node, gather_cast, unsqueeze_node, unsqueeze_cast, custom_node], + "custom_op_standard_rank_changes", + [gather_data, unsqueeze_data, plugin_in], + [ + helper.make_tensor_value_info("gather_y", TensorProto.FLOAT, []), + helper.make_tensor_value_info("unsqueeze_y", TensorProto.FLOAT, [1, 4]), + helper.make_tensor_value_info("plugin_y", TensorProto.FLOAT, [1]), + ], + [gather_index, unsqueeze_axes], + ) + model = helper.make_model( + graph, + producer_name="custom_op_standard_rank_changes", + opset_imports=[helper.make_opsetid("", 19), helper.make_opsetid("test.plugins", 1)], + ir_version=10, + ) + value_info_map, initializer_map, node_to_init_map = utils.setup_mappings(model) + + converter = PrecisionConverter( + model, + value_info_map, + initializer_map, + node_to_init_map, + keep_io_types=True, + custom_ops={"FakePlugin"}, + ) + propagated = converter._propagate_types_shapes_custom_ops(model) + + value_infos = {vi.name: vi for vi in [*propagated.graph.value_info, *propagated.graph.output]} + assert [ + dim.dim_value for dim in value_infos["gather_y_pre_cast"].type.tensor_type.shape.dim + ] == [] + assert [ + dim.dim_value for dim in value_infos["unsqueeze_y_pre_cast"].type.tensor_type.shape.dim + ] == [1, 4] + assert [dim.dim_value for dim in value_infos["gather_y"].type.tensor_type.shape.dim] == [] + assert [dim.dim_value for dim in value_infos["unsqueeze_y"].type.tensor_type.shape.dim] == [ + 1, + 4, + ] + onnx.checker.check_model(propagated, full_check=True) + + +def test_custom_op_mode_preserves_known_scalar_custom_op_shape(): + plugin_in = helper.make_tensor_value_info("plugin_in", TensorProto.FLOAT, [4]) + custom_node = helper.make_node( + "FakePlugin", ["plugin_in"], ["plugin_scalar"], name="plugin", domain="test.plugins" + ) + graph = helper.make_graph( + [custom_node], + "custom_op_known_scalar_shape", + [plugin_in], + [helper.make_tensor_value_info("plugin_scalar", TensorProto.FLOAT, [])], + [], + ) + model = helper.make_model( + graph, + producer_name="custom_op_known_scalar_shape", + opset_imports=[helper.make_opsetid("", 19), helper.make_opsetid("test.plugins", 1)], + ir_version=10, + ) + value_info_map, initializer_map, node_to_init_map = utils.setup_mappings(model) + + converter = PrecisionConverter( + model, + value_info_map, + initializer_map, + node_to_init_map, + keep_io_types=True, + custom_ops={"FakePlugin"}, + ) + propagated = converter._propagate_types_shapes_custom_ops(model) + + output = next(vi for vi in propagated.graph.output if vi.name == "plugin_scalar") + assert [dim.dim_value for dim in output.type.tensor_type.shape.dim] == []