diff --git a/modelopt/onnx/autocast/precisionconverter.py b/modelopt/onnx/autocast/precisionconverter.py index 5474b57a062..3396045def4 100644 --- a/modelopt/onnx/autocast/precisionconverter.py +++ b/modelopt/onnx/autocast/precisionconverter.py @@ -289,6 +289,33 @@ 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 not tensor.shape: + return None + return list(tensor.shape) + + 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:] + graph = gs.import_onnx(model) traversed_tensors = [] @@ -398,7 +425,9 @@ def _propagate_cast_type_through_nodes(node, np_type, iter=1): # Set the output shape if not out.shape: - if isinstance(inp, gs.Constant): + if (shape := _infer_gathernd_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 diff --git a/modelopt/onnx/utils.py b/modelopt/onnx/utils.py index 3b8edf76a84..3f62376923c 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 4c0f19759e3..856e3b42366 100644 --- a/tests/unit/onnx/autocast/test_precisionconverter.py +++ b/tests/unit/onnx/autocast/test_precisionconverter.py @@ -1854,3 +1854,123 @@ 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] == []