[6058907] Fix ShapeInferenceError in ONNX int8+fp16 quantization of weakly-typed models

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)
 ^^^^^^^^^^^^^^^^^
 

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(

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.")

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,53 @@ 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 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 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}")
+
+        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"):
+                    _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:
+                    _clear(vi, clear_shapes)
+        else:
+            for vi in g.value_info:
+                _clear(vi, clear_shape=True)
+
+        for out in g.output:
+            _clear(out, clear_shapes)
+
+    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,

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,121 @@ 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 _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
+        # 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 = _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 = _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).
+        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
+
+
 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 +2014,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