diff --git a/modelopt/onnx/autocast/utils.py b/modelopt/onnx/autocast/utils.py index 218034a6413..8483fe016bc 100644 --- a/modelopt/onnx/autocast/utils.py +++ b/modelopt/onnx/autocast/utils.py @@ -122,7 +122,7 @@ def clear_types_and_shapes_recursive( """Recursively clear type/shape information for a graph and all its subgraphs. Resets intermediate (``value_info``) and output tensor types to ``UNDEFINED`` and, when - ``clear_shapes`` is True, replaces concrete dims with a symbolic ``"unk"`` so a subsequent + ``clear_shapes`` is True, clears concrete dims so a subsequent :func:`modelopt.onnx.utils.infer_types` re-derives them from the operator graph. For subgraphs, input types/shapes are cleared too so they propagate from the parent graph. This does not change tensor *rank*, so it cannot repair a stale rank (see ``_reconcile_stale_output_shapes``). @@ -137,8 +137,8 @@ def _clear(value_info: onnx.ValueInfoProto, clear_shape: bool) -> None: value_info.type.tensor_type.elem_type = onnx.TensorProto.UNDEFINED if clear_shape: for dim in value_info.type.tensor_type.shape.dim: - if dim.dim_value: - dim.dim_param = "unk" + if dim.HasField("dim_value"): + dim.ClearField("dim_value") def _clear_callback(g: onnx.GraphProto, parent: onnx.NodeProto, is_sub: bool) -> None: logger.debug(f"Clearing types/shapes in {'subgraph' if is_sub else 'main graph'}: {g.name}") diff --git a/tests/unit/onnx/autocast/test_precisionconverter.py b/tests/unit/onnx/autocast/test_precisionconverter.py index d5ed09215fc..c5f0f09c7b1 100644 --- a/tests/unit/onnx/autocast/test_precisionconverter.py +++ b/tests/unit/onnx/autocast/test_precisionconverter.py @@ -15,6 +15,7 @@ import numpy as np import onnx +import onnxruntime as ort import pytest from onnx import TensorProto, helper, numpy_helper @@ -2746,3 +2747,83 @@ def test_loop_subgraph_high_precision_capture( ) onnx.checker.check_model(converted_model) onnx.shape_inference.infer_shapes(converted_model, strict_mode=True, check_type=True) + + +def test_dynamic_rope_shapes_do_not_alias(): + """AutoCast must not alias unrelated unknown dimensions with one symbolic name.""" + x = helper.make_tensor_value_info("X", TensorProto.FLOAT, ["batch", 2, "sequence", 64]) + y = helper.make_tensor_value_info("Y", TensorProto.FLOAT, ["batch", 2, "sequence", 64]) + + def int64_initializer(name, values): + return numpy_helper.from_array(np.array(values, dtype=np.int64), name=name) + + initializers = [ + int64_initializer("width_index", 3), + int64_initializer("two", 2), + int64_initializer("unsqueeze_axis", [0]), + int64_initializer("zero", [0]), + int64_initializer("slice_axis", [3]), + numpy_helper.from_array(np.ones((1, 1, 1, 64), dtype=np.float32), name="scale"), + ] + nodes = [ + helper.make_node("Shape", ["X"], ["shape"], name="shape"), + helper.make_node("Gather", ["shape", "width_index"], ["width"], name="gather_width"), + helper.make_node("Div", ["width", "two"], ["half_width"], name="half_width"), + helper.make_node("Unsqueeze", ["width", "unsqueeze_axis"], ["width_1d"], name="width_1d"), + helper.make_node( + "Unsqueeze", + ["half_width", "unsqueeze_axis"], + ["half_width_1d"], + name="half_width_1d", + ), + helper.make_node( + "Slice", + ["X", "half_width_1d", "width_1d", "slice_axis"], + ["upper_half"], + name="slice_upper_half", + ), + helper.make_node( + "Slice", + ["X", "zero", "half_width_1d", "slice_axis"], + ["lower_half"], + name="slice_lower_half", + ), + helper.make_node("Neg", ["upper_half"], ["negated_half"], name="negate"), + helper.make_node( + "Concat", + ["negated_half", "lower_half"], + ["rotated"], + name="rotate_half", + axis=3, + ), + helper.make_node("Mul", ["rotated", "scale"], ["Y"], name="rope_mul"), + ] + value_info = [ + helper.make_tensor_value_info(name, TensorProto.FLOAT, ["batch", 2, "sequence", width]) + for name, width in [ + ("upper_half", 32), + ("lower_half", 32), + ("negated_half", 32), + ("rotated", 64), + ] + ] + graph = helper.make_graph(nodes, "dynamic_rope", [x], [y], initializers, value_info=value_info) + model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", 20)]) + model.ir_version = LATEST_IR_VERSION_SUPPORTED_BY_ORT + + converted_model = convert_to_f16(model, keep_io_types=True) + + onnx.checker.check_model(converted_model) + assert all( + dim.dim_param != "unk" + for value in [*converted_model.graph.value_info, *converted_model.graph.output] + for dim in value.type.tensor_type.shape.dim + ) + session = ort.InferenceSession( + converted_model.SerializeToString(), providers=["CPUExecutionProvider"] + ) + output = session.run(None, {"X": np.ones((1, 2, 5, 64), dtype=np.float32)})[0] + + assert output.shape == (1, 2, 5, 64) + np.testing.assert_array_equal(output[..., :32], -1) + np.testing.assert_array_equal(output[..., 32:], 1)