[torch][ao] Do not crash numerics debugger if the shape of the tensors do not match (#149330)

dulinriley · dulinriley · commit c755508f157c · 2025-03-17T14:17:49.000-07:00
Summary: Pull Request resolved: #149330 Occasionally we see the loss function to crash because the shapes of `a` and `b` tensors are different. This diff avoids crashing is such scenarios and lets the comparison work for other nodes where the shapes match. Test Plan: - CI Reviewed By: jerryzh168 Differential Revision: D66245053
diff --git a/torch/ao/quantization/pt2e/_numeric_debugger.py b/torch/ao/quantization/pt2e/_numeric_debugger.py
@@ -318,10 +318,28 @@ def compare_results(
             )
             continue
         actual_name, actual_stack, actual_stats = actual_results[debug_handle]
+        for a, b in zip(actual_stats, ref_stats):
+            if (
+                isinstance(a, torch.Tensor)
+                and isinstance(b, torch.Tensor)
+                and a.shape != b.shape
+            ):
+                log.warning(
+                    "Cannot compare tensors with different shapes: actual=%s vs ref=%s",
+                    a.shape,
+                    b.shape,
+                )
         try:
             results = [
                 QuantizationComparisonResult(actual=a, ref=b)
                 for a, b in zip(actual_stats, ref_stats)
+                # Only compare objects if they're both collections, or both single
+                # tensors of the same shape.
+                # TODO: use torch.fx.node.map_aggregate for the collection of
+                # comparison results instead.
+                if not isinstance(a, torch.Tensor)
+                or not isinstance(b, torch.Tensor)
+                or a.shape == b.shape
             ]
         except Exception as e:
             # Add extra information for an exception from QuantizationComparisonResult
