From 9359aa0b74b746dc9c37026aa34f73414c788aec Mon Sep 17 00:00:00 2001 From: ajrasane <131806219+ajrasane@users.noreply.github.com> Date: Thu, 4 Jun 2026 18:19:57 +0000 Subject: [PATCH 1/4] [6058907][ONNX] Fix ShapeInferenceError in int8+fp16 quantization of weakly-typed models Weakly-typed ONNX models (e.g. TensorFlow exports) can carry graph.output entries whose stored rank disagrees with the graph topology -- most commonly a leftover rank-0 (scalar) annotation on a tensor that is really rank-2+. A stale rank-0 passes onnx.checker (a scalar is valid) but poisons downstream shape inference: ORT fails while augmenting the model for INT8 calibration ("axis must be in [-rank, rank-1]. Input rank was 0") and onnx.shape_inference(strict_mode=True) raises "Inferred shape and existing shape differ in rank" during fp16 autocast. Such models can also contain ops (e.g. TopK) that ONNX's static shape inference cannot resolve, leaving downstream tensors untyped and breaking AutoCast's type lookups. Fixes: - modelopt/onnx/utils.py: clear_stale_value_info now reconciles stale graph.output shapes. It snapshots and clears each output shape, re-derives shapes from the operator graph via ORT symbolic shape inference (falling back to the size-aware infer_shapes wrapper), and overwrites a declaration only when it is genuinely stale -- a rank mismatch or a conflicting concrete dim. Outputs that merely differ in symbolic dim_param names (e.g. a re-derived unk__0 vs a declared batch) keep their declaration, so healthy models with dynamic dims are untouched. A graph output is never left without a shape field (onnx.checker requires it). - modelopt/onnx/utils.py: infer_types falls back to the schema-based standalone type inferencer when ONNX shape inference fails (e.g. raises in strict mode on an op it cannot resolve), so the returned model is still fully typed. - modelopt/onnx/autocast/convert.py: convert_to_mixed_precision and convert_to_f16 now run the initial infer_types in strict mode, so an unresolvable op surfaces as an exception (triggering the fallback above) instead of silently leaving tensors untyped. Healthy models are unaffected: strict shape inference succeeds for them and output-shape reconciliation preserves their declared (incl. dynamic) shapes. Adds unit tests for stale rank-0 output reconciliation, preservation of a valid output shape, preservation of named dynamic dims, the infer_types fallback on a TopK-overflow model, and an AutoCast-level convert_to_f16 fallback test. Co-Authored-By: Claude Opus 4.8 (1M context) Signed-off-by: ajrasane <131806219+ajrasane@users.noreply.github.com> --- CHANGELOG.rst | 4 + modelopt/onnx/autocast/convert.py | 12 +- modelopt/onnx/utils.py | 145 ++++++++++++++++++++-- tests/unit/onnx/autocast/test_autocast.py | 34 +++++ tests/unit/onnx/test_onnx_utils.py | 92 ++++++++++++++ 5 files changed, 272 insertions(+), 15 deletions(-) diff --git a/CHANGELOG.rst b/CHANGELOG.rst index bc809ff66d0..b47aa32e51d 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -34,6 +34,10 @@ Changelog - ``hf_ptq.py`` also unwraps ``ModelOutput`` dataclasses from ``.generate()`` so the preview decode works on diffusion models. Non-tied models see no behavioral change. - Add **context-parallel (CP)** and **data-parallel (DP)** support to the shared Megatron-Core inference/calibration utilities. Under CP, ``get_megatron_calibration_forward_loop`` and ``megatron_mmlu`` partition each sequence across CP ranks (zigzag load-balanced), ``megatron_prefill`` accepts a CP-partitioned ``position_ids`` and lets the CP-aware causal attention build the mask, and MMLU gathers per-rank logits back to the full sequence for last-token scoring. Under DP, calibration shards the dataset across data-parallel ranks (``DistributedSampler``; amax is max-reduced across the DP group inside ``mtq``) and ``megatron_mmlu`` shards whole batches across DP ranks and all-reduces the per-subject counts. DP is implicit (``world_size / (tp * pp * cp)``); ``examples/megatron_bridge/quantize.py`` gains a ``--cp_size`` flag. +**Bug Fixes** + +- Fix ``ShapeInferenceError`` during ONNX INT8 + FP16 quantization (``--high_precision_dtype fp16``) of weakly-typed models (e.g. TensorFlow exports) that carry stale rank-0 ``graph.output`` shapes or ops such as ``TopK`` that ONNX's static shape inference cannot resolve. ``clear_stale_value_info`` now reconciles stale output shapes via symbolic shape inference (keeping every output's shape field populated), and AutoCast runs ONNX shape inference in strict mode and falls back to schema-based standalone type inference when it fails, so unresolved ops no longer leave tensors untyped. + 0.45 (2026-07-02) ^^^^^^^^^^^^^^^^^ diff --git a/modelopt/onnx/autocast/convert.py b/modelopt/onnx/autocast/convert.py index ec66ec11edd..0288c729dd4 100644 --- a/modelopt/onnx/autocast/convert.py +++ b/modelopt/onnx/autocast/convert.py @@ -136,8 +136,10 @@ def convert_to_mixed_precision( graph_sanitizer.sanitize() model = graph_sanitizer.model - # Setup internal mappings - model = onnx_utils.infer_types(model, use_standalone_type_inference) + # Setup internal mappings. Use strict shape inference so an op ONNX cannot resolve surfaces + # as an exception (triggering infer_types' standalone type-inference fallback) instead of + # silently leaving tensors untyped, which would break later type lookups. + model = onnx_utils.infer_types(model, use_standalone_type_inference, strict_mode=True) value_info_map, initializer_map, node_to_init_map = utils.setup_mappings(model) # Automatically add 'trt' to list of providers if custom ops are detected @@ -267,8 +269,10 @@ def convert_to_f16( sanitizer.convert_fp64_to_fp32() model = sanitizer.model - # Setup internal mappings - model = onnx_utils.infer_types(model, use_standalone_type_inference) + # Setup internal mappings. Use strict shape inference so an op ONNX cannot resolve surfaces + # as an exception (triggering infer_types' standalone type-inference fallback) instead of + # silently leaving tensors untyped, which would break later type lookups. + model = onnx_utils.infer_types(model, use_standalone_type_inference, strict_mode=True) value_info_map, initializer_map, node_to_init_map = utils.setup_mappings(model) precision_converter = PrecisionConverter( diff --git a/modelopt/onnx/utils.py b/modelopt/onnx/utils.py index 64b1fd1c414..c45c1ccd6b0 100644 --- a/modelopt/onnx/utils.py +++ b/modelopt/onnx/utils.py @@ -1205,19 +1205,32 @@ def infer_types( When use_standalone_type_inference is True, uses a standalone type inference implementation that only infers types. Otherwise, uses ONNX's infer_shapes which infers both types and shapes. + ONNX's ``infer_shapes`` can fail on weakly-typed models -- with ``strict_mode=True`` it raises + on an op it cannot resolve (e.g. a ``TopK`` whose axis it resolves to a stale dimension) + instead of silently leaving that node's outputs untyped. On any shape-inference failure this + falls back to the standalone type inferencer, which derives types from operator schemas + regardless of shapes, so downstream type lookups (e.g. in AutoCast) do not fail. Callers that + need a fully typed graph should pass ``strict_mode=True`` so incomplete inference surfaces as + an exception that triggers the fallback. + Args: model: ONNX model to infer types/shapes for. use_standalone_type_inference: If True, use standalone type inference (_infer_types_only). If False, use ONNX's shape inference (infer_shapes). - **kwargs: Additional arguments passed to infer_shapes when not using standalone type inference. + **kwargs: Additional arguments passed to infer_shapes when not using standalone type + inference (e.g. ``strict_mode``, ``check_type``, ``data_prop``). Returns: onnx.ModelProto: Model with inferred types (and shapes if not using standalone type inference). """ if use_standalone_type_inference: return _infer_types_only(model) - else: + + try: return infer_shapes(model, **kwargs) + except Exception as e: + logger.debug("ONNX shape inference failed (%s); using standalone type inference.", e) + return _infer_types_only(model) def onnx_type_str_to_enum(dtype: str) -> int: @@ -1862,21 +1875,126 @@ def change_casts_to_fp16(model: onnx.ModelProto, target_op_types: list[str]) -> return model +def _reconcile_stale_output_shapes(model: onnx.ModelProto) -> int: + """Re-derive stale ``graph.output`` shapes from the operator graph. + + Weakly-typed models (e.g. exported from TensorFlow) can declare an output rank + that conflicts with the graph topology -- most commonly a leftover rank-0 + (scalar) annotation on a tensor that is really rank-2+. Such a stale rank poisons + downstream shape inference: ORT fails while augmenting the model for INT8 + calibration (``axis must be in [-rank, rank-1]. Input rank was 0``), and + ``onnx.shape_inference`` with ``strict_mode=True`` raises ``Inferred shape and + existing shape differ in rank`` during fp16 autocast. + + Strategy: snapshot the declared output shapes, clear them, and re-derive them from + the operator graph -- preferring ORT's symbolic shape inference (it resolves ops + such as ``TopK`` that ONNX's static inference gives up on) and falling back to the + size-aware ``infer_shapes`` wrapper. A declared shape is only overwritten when it is + genuinely stale -- a rank mismatch (the rank-0-vs-rank-N bug) or a conflicting + concrete dimension. Outputs that merely differ in symbolic ``dim_param`` names (e.g. + a re-derived ``unk__0`` vs a declared ``batch``) keep their original declaration, so + healthy models -- including dynamic batch/sequence dims -- are left untouched. A + graph output is never left without a shape (``onnx.checker`` requires the field). + + Args: + model: Loaded in-memory onnx ModelProto, ideally with ``value_info`` already + cleared so re-inference derives shapes from the operator graph. + + Returns: + Number of graph outputs whose shape was changed. + """ + outputs = model.graph.output + if not outputs: + return 0 + + def _sig(shape: onnx.TensorShapeProto | None) -> bytes | None: + return None if shape is None else shape.SerializeToString() + + def _is_stale(declared: onnx.TensorShapeProto | None, inferred: onnx.TensorShapeProto | None): + # Only treat a declaration as stale when inference contradicts it: a different + # rank, or a concrete dim that disagrees with an inferred concrete dim. A missing + # declaration is "stale" (adopt whatever was inferred); a missing inference is not + # (keep the declaration). Symbolic dim_param renames are intentionally ignored. + if inferred is None: + return False + if declared is None: + return True + if len(declared.dim) != len(inferred.dim): + return True + return any( + d.HasField("dim_value") and i.HasField("dim_value") and d.dim_value != i.dim_value + for d, i in zip(declared.dim, inferred.dim) + ) + + # Snapshot declared shapes, then clear them so re-inference starts from the + # topology instead of being biased by the stale annotations. + declared: dict[str, onnx.TensorShapeProto | None] = {} + for o in outputs: + tt = o.type.tensor_type + if tt.HasField("shape"): + snapshot = onnx.TensorShapeProto() + snapshot.CopyFrom(tt.shape) + declared[o.name] = snapshot + else: + declared[o.name] = None + tt.ClearField("shape") + + # Re-derive output shapes from the cleared model (neither inference call mutates it): + # prefer ORT symbolic shape inference, then fall back to the size-aware infer_shapes + # wrapper if it is unavailable or yields nothing. + inferred: dict[str, onnx.TensorShapeProto] = {} + try: + from onnxruntime.tools.symbolic_shape_infer import SymbolicShapeInference + + inferred_model = SymbolicShapeInference.infer_shapes(model, auto_merge=True) + inferred = { + o.name: o.type.tensor_type.shape + for o in inferred_model.graph.output + if o.type.tensor_type.HasField("shape") + } + except Exception as e: + logger.debug("Symbolic shape inference unavailable/failed: %s", e) + if not inferred: + try: + inferred_model = infer_shapes(model, strict_mode=False, data_prop=True) + inferred = { + o.name: o.type.tensor_type.shape + for o in inferred_model.graph.output + if o.type.tensor_type.HasField("shape") + } + except Exception as e: + logger.debug("ONNX shape inference for output reconciliation failed: %s", e) + + changed = 0 + for o in outputs: + decl = declared[o.name] + # Adopt the inferred shape only when the declaration is genuinely stale; otherwise + # restore the declared shape (never leaving a graph output shapeless). + new_shape = inferred.get(o.name) if _is_stale(decl, inferred.get(o.name)) else decl + if new_shape is not None: + o.type.tensor_type.shape.CopyFrom(new_shape) + if _sig(decl) != _sig(new_shape): + changed += 1 + return changed + + def clear_stale_value_info(model: onnx.ModelProto) -> int: - """Clear stale type metadata that would otherwise trip ORT's type checker. + """Clear stale type/shape metadata that would otherwise trip ORT's type checker. - Walks every ``Cast`` node and forces the ``elem_type`` of any - ``graph.output`` entry produced by that Cast to match the Cast's ``to`` - attribute (the spec-defined contract for a Cast's output dtype). Then - clears ``value_info`` so ORT/shape-inference re-derives intermediate-tensor - types from the operator graph during session setup -- except entries for - outputs of ``trt.plugins`` custom-op nodes, whose types ORT cannot infer. + Walks every ``Cast`` node and forces the ``elem_type`` of any ``graph.output`` + entry produced by that Cast to match the Cast's ``to`` attribute (the spec-defined + contract for a Cast's output dtype). Clears ``value_info`` so ORT/shape-inference + re-derives intermediate-tensor types from the operator graph during session setup + -- except entries for outputs of ``trt.plugins`` custom-op nodes, whose types ORT + cannot infer. Finally, reconciles stale ``graph.output`` *shapes* (e.g. a leftover + rank-0 scalar on a tensor that is really rank-2+) which would otherwise propagate a + wrong rank into downstream shape inference. Args: model: Loaded in-memory onnx ModelProto. Returns: - Number of Cast outputs reconciled plus value_info entries cleared. + Total number of entries reconciled or cleared. """ cast_to_by_output = { node.output[0]: get_cast_to_type(node) @@ -1901,4 +2019,9 @@ def clear_stale_value_info(model: onnx.ModelProto) -> int: if n_cleared: del model.graph.value_info[:] model.graph.value_info.extend(preserved) - return fixed_outputs + n_cleared + + # Reconcile output shapes after value_info is cleared so the re-inference inside + # the helper derives shapes cleanly from the operator graph. + fixed_shapes = _reconcile_stale_output_shapes(model) + + return fixed_outputs + fixed_shapes + n_cleared diff --git a/tests/unit/onnx/autocast/test_autocast.py b/tests/unit/onnx/autocast/test_autocast.py index f761e1e9e3c..ccad0c7201b 100644 --- a/tests/unit/onnx/autocast/test_autocast.py +++ b/tests/unit/onnx/autocast/test_autocast.py @@ -26,6 +26,7 @@ import modelopt.onnx.utils as onnx_utils from modelopt.onnx.autocast import convert_to_mixed_precision from modelopt.onnx.autocast.__main__ import get_parser, main +from modelopt.onnx.autocast.convert import convert_to_f16 from modelopt.onnx.autocast.logging_config import configure_logging configure_logging("DEBUG") @@ -321,3 +322,36 @@ def test_opset_parser_argument(): # Test parsing without opset (should be None) args = parser.parse_args(["--onnx_path", "test.onnx"]) assert args.opset is None + + +@pytest.fixture +def weakly_typed_topk_model(): + # TopK k (5) exceeds the static axis size (3), so ONNX shape inference cannot resolve it. + x = onnx.helper.make_tensor_value_info("X", onnx.TensorProto.FLOAT, [1, 3]) + out = onnx.helper.make_tensor_value_info("out", onnx.TensorProto.FLOAT, [1, 5]) + weight = onnx.numpy_helper.from_array(np.ones((1, 3), dtype=np.float32), name="weight") + k = onnx.numpy_helper.from_array(np.array([5], dtype=np.int64), name="k") + nodes = [ + onnx.helper.make_node("Add", ["X", "weight"], ["a"], name="add"), + onnx.helper.make_node("TopK", ["a", "k"], ["vals", "inds"], axis=1, name="topk"), + onnx.helper.make_node("Cast", ["inds"], ["out"], to=onnx.TensorProto.FLOAT, name="cast"), + ] + graph = onnx.helper.make_graph(nodes, "weakly_typed_topk", [x], [out], [weight, k]) + model = onnx.helper.make_model(graph, opset_imports=[onnx.helper.make_opsetid("", 17)]) + model.ir_version = 10 + return model + + +def test_convert_to_f16_falls_back_on_unresolvable_op(weakly_typed_topk_model): + """A weakly-typed graph ONNX shape inference cannot resolve must still convert. + + The TopK ``k`` (5) exceeds the static axis size (3), so ONNX shape inference raises + in strict mode (the same failure class as NVBug 6058907). ``convert_to_f16`` -- the + path used by INT8 + ``--high_precision_dtype fp16`` quantization -- runs infer_types + in strict mode and must fall back to standalone type inference instead of crashing, + typing the TopK's int64 indices output that feeds the downstream Cast. + """ + converted_model = convert_to_f16(weakly_typed_topk_model, keep_io_types=True) + + onnx.checker.check_model(converted_model) + assert any(n.op_type == "TopK" for n in converted_model.graph.node) diff --git a/tests/unit/onnx/test_onnx_utils.py b/tests/unit/onnx/test_onnx_utils.py index ec1c5ef75a5..36face35b90 100644 --- a/tests/unit/onnx/test_onnx_utils.py +++ b/tests/unit/onnx/test_onnx_utils.py @@ -33,6 +33,7 @@ clear_stale_value_info, get_input_names_from_bytes, get_output_names_from_bytes, + infer_types, randomize_weights_onnx_bytes, remove_node_training_mode, remove_weights_data, @@ -364,3 +365,94 @@ def test_clear_stale_value_info(output_elem_type, with_value_info, expected_coun assert model.graph.output[0].type.tensor_type.elem_type == onnx.TensorProto.FLOAT assert len(model.graph.value_info) == 0 assert count == expected_count + + +def _make_matmul_model(output_shape): + """Build an X[3,4] @ W[4,5] -> Y model with Y declared using ``output_shape``.""" + weights = make_tensor("W", onnx.TensorProto.FLOAT, [4, 5], np.zeros(20, dtype=np.float32)) + nodes = [make_node("MatMul", ["X", "W"], ["Y"], name="matmul")] + inputs = [make_tensor_value_info("X", onnx.TensorProto.FLOAT, [3, 4])] + outputs = [make_tensor_value_info("Y", onnx.TensorProto.FLOAT, output_shape)] + graph = make_graph(nodes, "matmul_graph", inputs, outputs, initializer=[weights]) + return make_model(graph, producer_name="modelopt test", opset_imports=[make_opsetid("", 17)]) + + +def test_clear_stale_value_info_reconciles_stale_rank0_output(): + # Y is really rank-2 [3, 5] but the model declares it as a rank-0 scalar (stale + # metadata typical of weakly-typed exports). This is the rank-(N)-vs-(0) class of + # conflict that crashes downstream shape inference (NVBug 6058907). + model = _make_matmul_model(output_shape=[]) + assert len(model.graph.output[0].type.tensor_type.shape.dim) == 0 # stale rank-0 + + clear_stale_value_info(model) + + out_type = model.graph.output[0].type.tensor_type + assert out_type.HasField("shape") # shape field must remain (onnx.checker requires it) + assert [d.dim_value for d in out_type.shape.dim] == [3, 5] # reconciled to the real shape + onnx.checker.check_model(model) + + +def test_clear_stale_value_info_preserves_valid_output_shape(): + # A correct output shape must be left untouched (no-op for healthy models). + model = _make_matmul_model(output_shape=[3, 5]) + + clear_stale_value_info(model) + + out_type = model.graph.output[0].type.tensor_type + assert [d.dim_value for d in out_type.shape.dim] == [3, 5] + + +def _make_dynamic_dim_model(): + """Build an X[batch,4] -> Relu -> Y[my_batch,4] model (output declares a different dim_param).""" + nodes = [make_node("Relu", ["X"], ["Y"], name="relu")] + inputs = [make_tensor_value_info("X", onnx.TensorProto.FLOAT, ["batch", 4])] + outputs = [make_tensor_value_info("Y", onnx.TensorProto.FLOAT, ["my_batch", 4])] + graph = make_graph(nodes, "dyn_graph", inputs, outputs) + return make_model(graph, producer_name="modelopt test", opset_imports=[make_opsetid("", 17)]) + + +def test_clear_stale_value_info_preserves_dynamic_dim_names(): + # A healthy output with a named dynamic dim must not be rewritten just because + # symbolic shape inference re-derives a different dim_param. Y is declared with + # "my_batch" while the graph would infer "batch" from the input: same rank, no + # concrete-dim conflict, so the declaration (incl. its dim_param) must be preserved. + model = _make_dynamic_dim_model() + + clear_stale_value_info(model) + + out_dims = model.graph.output[0].type.tensor_type.shape.dim + assert [d.dim_param or d.dim_value for d in out_dims] == ["my_batch", 4] + onnx.checker.check_model(model) + + +def _make_topk_overflow_model(): + """Build a model whose TopK ``k`` (5) exceeds the static axis dim (3). + + ONNX shape inference raises "Axis has less than the requested k elements" on this + model (the same failure class seen in NVBug 6058907), while standalone type + inference can still derive the output types (values float, indices int64). + """ + k = make_tensor("k", onnx.TensorProto.INT64, [1], [5]) + nodes = [ + make_node("TopK", ["X", "k"], ["vals", "inds"], axis=1, name="topk"), + make_node("Cast", ["inds"], ["out"], to=onnx.TensorProto.FLOAT, name="cast_inds"), + ] + inputs = [make_tensor_value_info("X", onnx.TensorProto.FLOAT, [1, 3])] + outputs = [make_tensor_value_info("out", onnx.TensorProto.FLOAT, [1, 5])] + graph = make_graph(nodes, "topk_overflow", inputs, outputs, initializer=[k]) + return make_model(graph, producer_name="modelopt test", opset_imports=[make_opsetid("", 17)]) + + +def test_infer_types_falls_back_to_standalone_when_onnx_fails(): + # ONNX shape inference cannot resolve this model's TopK. With strict_mode=True it raises + # (instead of silently leaving the TopK outputs untyped), so infer_types catches the + # error and falls back to standalone type inference, which still types every tensor. + model = _make_topk_overflow_model() + + inferred = infer_types(model, strict_mode=True) + + value_info_types = {vi.name: vi.type.tensor_type.elem_type for vi in inferred.graph.value_info} + output_types = {o.name: o.type.tensor_type.elem_type for o in inferred.graph.output} + assert value_info_types.get("vals") == onnx.TensorProto.FLOAT + assert value_info_types.get("inds") == onnx.TensorProto.INT64 # TopK indices + assert output_types.get("out") == onnx.TensorProto.FLOAT From 26fe4a66e9e1a5d138a668a33e2d98fdc5125794 Mon Sep 17 00:00:00 2001 From: ajrasane <131806219+ajrasane@users.noreply.github.com> Date: Tue, 16 Jun 2026 15:56:37 +0000 Subject: [PATCH 2/4] [6058907][ONNX] Extract clear_types_and_shapes_recursive to autocast utils Move the recursive type/shape clearing out of PrecisionConverter into modelopt/onnx/autocast/utils.py as a shared clear_types_and_shapes_recursive() next to walk_subgraphs_recursive; PrecisionConverter now delegates to it. Behavior is unchanged (clear_shapes = not use_standalone_type_inference). The standalone type-inference fallback on shape-inference failure already lives in infer_types(), which the AutoCast path exercises via strict_mode=True. Co-Authored-By: Claude Opus 4.8 (1M context) Signed-off-by: ajrasane <131806219+ajrasane@users.noreply.github.com> --- modelopt/onnx/autocast/precisionconverter.py | 64 +----------------- modelopt/onnx/autocast/utils.py | 68 ++++++++++++++++++++ 2 files changed, 71 insertions(+), 61 deletions(-) diff --git a/modelopt/onnx/autocast/precisionconverter.py b/modelopt/onnx/autocast/precisionconverter.py index 1603a05a2f2..853fa3d1cb3 100644 --- a/modelopt/onnx/autocast/precisionconverter.py +++ b/modelopt/onnx/autocast/precisionconverter.py @@ -255,7 +255,9 @@ def convert( self.model = self._propagate_types_shapes_custom_ops(self.model) else: # Clear type/shape information for intermediates and outputs (including subgraphs) - self._clear_types_and_shapes_recursive(self.model.graph) + utils.clear_types_and_shapes_recursive( + self.model.graph, clear_shapes=not self.use_standalone_type_inference + ) # Populate type information with inferred types self.model = onnx_utils.infer_types( self.model, self.use_standalone_type_inference, strict_mode=True, check_type=False @@ -284,66 +286,6 @@ def _ensure_types_are_defined(self): if vi.type.tensor_type.elem_type == onnx.TensorProto.UNDEFINED: vi.type.tensor_type.elem_type = self.low_precision_type.onnx_type - def _clear_types_and_shapes_recursive( - self, graph: onnx.GraphProto, is_subgraph: bool = False - ) -> None: - """Recursively clear type/shape information for a graph and all its subgraphs. - - If use_standalone_type_inference is True, we clear only types, not shapes. - For subgraphs, input types/shapes are cleared, so that the input types/shapes are propagated - from the main graph. - - Args: - graph: The ONNX graph to clear types and shapes for. - is_subgraph: Whether this is a subgraph (True) or the main graph (False). - """ - - 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}" - ) - - # Clear type/shape information for inputs (only for subgraphs, not main graph inputs) - if is_sub: - for inp in g.input: - if inp.type.HasField("tensor_type"): - inp.type.tensor_type.elem_type = onnx.TensorProto.UNDEFINED - if not self.use_standalone_type_inference: - for idx, d in enumerate(inp.type.tensor_type.shape.dim): - if d.dim_value: - inp.type.tensor_type.shape.dim[idx].dim_param = "unk" - - if is_sub: - # Identify which tensors are produced by nodes in this subgraph - subgraph_outputs = set() - for node in g.node: - subgraph_outputs.update(node.output) - - # Clear value_info only for intermediates produced by nodes in this subgraph - for vi in g.value_info: - if vi.name in subgraph_outputs: - vi.type.tensor_type.elem_type = onnx.TensorProto.UNDEFINED - if not self.use_standalone_type_inference: - for idx, d in enumerate(vi.type.tensor_type.shape.dim): - if d.dim_value: - vi.type.tensor_type.shape.dim[idx].dim_param = "unk" - else: - for vi in g.value_info: - vi.type.tensor_type.elem_type = onnx.TensorProto.UNDEFINED - for idx, d in enumerate(vi.type.tensor_type.shape.dim): - if d.dim_value: - vi.type.tensor_type.shape.dim[idx].dim_param = "unk" - - # Clear outputs for both main graph and subgraphs - for out in g.output: - out.type.tensor_type.elem_type = onnx.TensorProto.UNDEFINED - if not self.use_standalone_type_inference: - for idx, d in enumerate(out.type.tensor_type.shape.dim): - if d.dim_value: - out.type.tensor_type.shape.dim[idx].dim_param = "unk" - - utils.walk_subgraphs_recursive(graph, _clear_callback, is_subgraph=is_subgraph) - 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.") diff --git a/modelopt/onnx/autocast/utils.py b/modelopt/onnx/autocast/utils.py index a9ad5484c0c..95773764fb2 100644 --- a/modelopt/onnx/autocast/utils.py +++ b/modelopt/onnx/autocast/utils.py @@ -28,6 +28,7 @@ import onnx import modelopt.onnx.utils as onnx_utils +from modelopt.onnx.autocast.logging_config import logger from modelopt.onnx.utils import get_opset_version @@ -115,6 +116,73 @@ def walk_subgraphs_recursive( walk_subgraphs_recursive(subgraph, callback, parent_node=node, is_subgraph=True) +def clear_types_and_shapes_recursive( + graph: onnx.GraphProto, clear_shapes: bool = True, is_subgraph: bool = False +) -> None: + """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 dimensions with a symbolic ``"unk"`` placeholder -- + so that a subsequent :func:`modelopt.onnx.utils.infer_types` re-derives types (and shapes) from + the operator graph instead of trusting stale annotations. For subgraphs, input types/shapes are + cleared as well so that they are propagated from the parent graph. + + Note: this clears in place but does not change tensor *rank*. It cannot repair a stale rank (e.g. + a leftover rank-0 scalar on a tensor that is really rank-2+); see + :func:`modelopt.onnx.utils._reconcile_stale_output_shapes` for that. + + Args: + graph: The ONNX graph to clear types and shapes for. + clear_shapes: If True, also clear shape information. Pass False to keep shapes when only + types will be re-inferred (mirrors standalone type inference). + is_subgraph: Whether this is a subgraph (True) or the main graph (False). + """ + + 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}") + + # Clear type/shape information for inputs (only for subgraphs, not main graph inputs) + if is_sub: + for inp in g.input: + if inp.type.HasField("tensor_type"): + inp.type.tensor_type.elem_type = onnx.TensorProto.UNDEFINED + if clear_shapes: + for idx, d in enumerate(inp.type.tensor_type.shape.dim): + if d.dim_value: + inp.type.tensor_type.shape.dim[idx].dim_param = "unk" + + if is_sub: + # Identify which tensors are produced by nodes in this subgraph + subgraph_outputs = set() + for node in g.node: + subgraph_outputs.update(node.output) + + # Clear value_info only for intermediates produced by nodes in this subgraph + for vi in g.value_info: + if vi.name in subgraph_outputs: + vi.type.tensor_type.elem_type = onnx.TensorProto.UNDEFINED + if clear_shapes: + for idx, d in enumerate(vi.type.tensor_type.shape.dim): + if d.dim_value: + vi.type.tensor_type.shape.dim[idx].dim_param = "unk" + else: + for vi in g.value_info: + vi.type.tensor_type.elem_type = onnx.TensorProto.UNDEFINED + for idx, d in enumerate(vi.type.tensor_type.shape.dim): + if d.dim_value: + vi.type.tensor_type.shape.dim[idx].dim_param = "unk" + + # Clear outputs for both main graph and subgraphs + for out in g.output: + out.type.tensor_type.elem_type = onnx.TensorProto.UNDEFINED + if clear_shapes: + for idx, d in enumerate(out.type.tensor_type.shape.dim): + if d.dim_value: + out.type.tensor_type.shape.dim[idx].dim_param = "unk" + + walk_subgraphs_recursive(graph, _clear_callback, is_subgraph=is_subgraph) + + def get_op_types_not_supported_in_low_precision( model: onnx.ModelProto, min_opset: int, From d4d2667193f0879c8ec3e589ab1bae0f6d61ad67 Mon Sep 17 00:00:00 2001 From: ajrasane <131806219+ajrasane@users.noreply.github.com> Date: Tue, 16 Jun 2026 16:16:44 +0000 Subject: [PATCH 3/4] [6058907][ONNX] Simplify clear_types_and_shapes_recursive Trim the docstring to the module's density and factor the repeated elem_type/shape clearing into a local helper. No functional change. Co-Authored-By: Claude Opus 4.8 (1M context) Signed-off-by: ajrasane <131806219+ajrasane@users.noreply.github.com> --- modelopt/onnx/autocast/utils.py | 60 +++++++++++---------------------- 1 file changed, 20 insertions(+), 40 deletions(-) diff --git a/modelopt/onnx/autocast/utils.py b/modelopt/onnx/autocast/utils.py index 95773764fb2..218034a6413 100644 --- a/modelopt/onnx/autocast/utils.py +++ b/modelopt/onnx/autocast/utils.py @@ -121,64 +121,44 @@ def clear_types_and_shapes_recursive( ) -> None: """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 dimensions with a symbolic ``"unk"`` placeholder -- - so that a subsequent :func:`modelopt.onnx.utils.infer_types` re-derives types (and shapes) from - the operator graph instead of trusting stale annotations. For subgraphs, input types/shapes are - cleared as well so that they are propagated from the parent graph. - - Note: this clears in place but does not change tensor *rank*. It cannot repair a stale rank (e.g. - a leftover rank-0 scalar on a tensor that is really rank-2+); see - :func:`modelopt.onnx.utils._reconcile_stale_output_shapes` for that. + 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 + :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``). Args: graph: The ONNX graph to clear types and shapes for. - clear_shapes: If True, also clear shape information. Pass False to keep shapes when only - types will be re-inferred (mirrors standalone type inference). + clear_shapes: If True, also clear shapes (False keeps shapes for type-only inference). is_subgraph: Whether this is a subgraph (True) or the main graph (False). """ + 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" + 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}") - # Clear type/shape information for inputs (only for subgraphs, not main graph inputs) if is_sub: + # Subgraph inputs are cleared so they propagate from the parent graph. for inp in g.input: if inp.type.HasField("tensor_type"): - inp.type.tensor_type.elem_type = onnx.TensorProto.UNDEFINED - if clear_shapes: - for idx, d in enumerate(inp.type.tensor_type.shape.dim): - if d.dim_value: - inp.type.tensor_type.shape.dim[idx].dim_param = "unk" - - if is_sub: - # Identify which tensors are produced by nodes in this subgraph - subgraph_outputs = set() - for node in g.node: - subgraph_outputs.update(node.output) - - # Clear value_info only for intermediates produced by nodes in this subgraph + _clear(inp, clear_shapes) + # Only clear value_info for intermediates produced within this subgraph. + subgraph_outputs = {out for node in g.node for out in node.output} for vi in g.value_info: if vi.name in subgraph_outputs: - vi.type.tensor_type.elem_type = onnx.TensorProto.UNDEFINED - if clear_shapes: - for idx, d in enumerate(vi.type.tensor_type.shape.dim): - if d.dim_value: - vi.type.tensor_type.shape.dim[idx].dim_param = "unk" + _clear(vi, clear_shapes) else: for vi in g.value_info: - vi.type.tensor_type.elem_type = onnx.TensorProto.UNDEFINED - for idx, d in enumerate(vi.type.tensor_type.shape.dim): - if d.dim_value: - vi.type.tensor_type.shape.dim[idx].dim_param = "unk" + _clear(vi, clear_shape=True) - # Clear outputs for both main graph and subgraphs for out in g.output: - out.type.tensor_type.elem_type = onnx.TensorProto.UNDEFINED - if clear_shapes: - for idx, d in enumerate(out.type.tensor_type.shape.dim): - if d.dim_value: - out.type.tensor_type.shape.dim[idx].dim_param = "unk" + _clear(out, clear_shapes) walk_subgraphs_recursive(graph, _clear_callback, is_subgraph=is_subgraph) From c3a8dae6e79185572b034f7a23bbcc0873963f33 Mon Sep 17 00:00:00 2001 From: ajrasane <131806219+ajrasane@users.noreply.github.com> Date: Tue, 16 Jun 2026 18:36:41 +0000 Subject: [PATCH 4/4] [6058907][ONNX] Simplify _reconcile_stale_output_shapes Drop the _sig byte-comparison helper and derive the change count directly from _is_stale, which already encodes the overwrite policy. Extract a small _outputs_with_shapes helper shared by the two inference attempts. No behavior change. Co-Authored-By: Claude Opus 4.8 (1M context) Signed-off-by: ajrasane <131806219+ajrasane@users.noreply.github.com> --- modelopt/onnx/utils.py | 31 +++++++++++++------------------ 1 file changed, 13 insertions(+), 18 deletions(-) diff --git a/modelopt/onnx/utils.py b/modelopt/onnx/utils.py index c45c1ccd6b0..3b8edf76a84 100644 --- a/modelopt/onnx/utils.py +++ b/modelopt/onnx/utils.py @@ -1907,8 +1907,12 @@ def _reconcile_stale_output_shapes(model: onnx.ModelProto) -> int: if not outputs: return 0 - def _sig(shape: onnx.TensorShapeProto | None) -> bytes | None: - return None if shape is None else shape.SerializeToString() + def _outputs_with_shapes(m: onnx.ModelProto) -> dict[str, onnx.TensorShapeProto]: + return { + o.name: o.type.tensor_type.shape + for o in m.graph.output + if o.type.tensor_type.HasField("shape") + } def _is_stale(declared: onnx.TensorShapeProto | None, inferred: onnx.TensorShapeProto | None): # Only treat a declaration as stale when inference contradicts it: a different @@ -1946,35 +1950,26 @@ def _is_stale(declared: onnx.TensorShapeProto | None, inferred: onnx.TensorShape try: from onnxruntime.tools.symbolic_shape_infer import SymbolicShapeInference - inferred_model = SymbolicShapeInference.infer_shapes(model, auto_merge=True) - inferred = { - o.name: o.type.tensor_type.shape - for o in inferred_model.graph.output - if o.type.tensor_type.HasField("shape") - } + inferred = _outputs_with_shapes(SymbolicShapeInference.infer_shapes(model, auto_merge=True)) except Exception as e: logger.debug("Symbolic shape inference unavailable/failed: %s", e) if not inferred: try: - inferred_model = infer_shapes(model, strict_mode=False, data_prop=True) - inferred = { - o.name: o.type.tensor_type.shape - for o in inferred_model.graph.output - if o.type.tensor_type.HasField("shape") - } + inferred = _outputs_with_shapes(infer_shapes(model, strict_mode=False, data_prop=True)) except Exception as e: logger.debug("ONNX shape inference for output reconciliation failed: %s", e) changed = 0 for o in outputs: decl = declared[o.name] + inf = inferred.get(o.name) # Adopt the inferred shape only when the declaration is genuinely stale; otherwise # restore the declared shape (never leaving a graph output shapeless). - new_shape = inferred.get(o.name) if _is_stale(decl, inferred.get(o.name)) else decl - if new_shape is not None: - o.type.tensor_type.shape.CopyFrom(new_shape) - if _sig(decl) != _sig(new_shape): + if _is_stale(decl, inf): + o.type.tensor_type.shape.CopyFrom(inf) changed += 1 + elif decl is not None: + o.type.tensor_type.shape.CopyFrom(decl) return changed