[ONNX] Update xfail reasons in fx runtime tests (#107257)

1. Update xfail reasons in fx runtime
2. Enable bloom-560m in runtime test. However, it's blocked by the unsupported constant tensor case. The previous error was because the when the model loads with external data, it surpasses 2GB, and couldn't be inlined. The fix is to inline the model it self, and then replace the original one. Pointing ORT to the path allows it to load with external data into model in runtime.
Pull Request resolved: #107257
Approved by: https://github.com/justinchuby

test/onnx/onnx_test_common.py

@@ -359,19 +359,23 @@ def run_ort(
     # with nested functions.
     # Ref: https://github.com/microsoft/onnxruntime/issues/15849
     try:
-        import onnx.inliner
+        import onnx.inliner  # type: ignore[import]
     except ImportError:
         warnings.warn("Cannot import onnx.inliner. Skip inlining model.")
     else:
         if isinstance(ort_model, bytes):
             buffer = io.BytesIO(ort_model)
+            model_proto = onnx.load(buffer)
+            inlined_model_proto = onnx.inliner.inline_local_functions(model_proto)
+            buffer = inlined_model_proto.SerializeToString()
+            ort_model = buffer
         else:
             assert isinstance(ort_model, str)
-            buffer = ort_model
-
-        model_proto = onnx.load(buffer)
-        inlined_model_proto = onnx.inliner.inline_local_functions(model_proto)
-        ort_model = inlined_model_proto.SerializeToString()
+            # NOTE: inline_local_functions doesn't work with >2GB models,
+            # so we need to load the model without external data to inline.
+            model_proto = onnx.load(ort_model, load_external_data=False)
+            inlined_model_proto = onnx.inliner.inline_local_functions(model_proto)
+            onnx.save(inlined_model_proto, ort_model)
 
     # Suppress floods of warnings from ONNX Runtime
     session_options = onnxruntime.SessionOptions()

test/onnx/test_fx_to_onnx_with_onnxruntime.py

@@ -269,15 +269,16 @@ def test_resnet18(self):
             (dummy_input,),
         )
 
-    @pytorch_test_common.xfail(
-        "RuntimeError: Unknown call_function target: aten.mean.dim"
+    @pytorch_test_common.skip_dynamic_fx_test(
+        "[ONNXRuntimeError] : 2 : INVALID_ARGUMENT : "
+        "Got invalid dimensions for input: arg0 for the following indices index: 0 Got: 3 Expected: 1"
     )
     @skip_if_no_torchvision
     def test_shufflenet_v2(self):
         # TODO(bowbao): see Note [training vs eval in dynamo_export]
         model = torchvision.models.shufflenet_v2_x0_5(pretrained=False).eval()
-        dummy_input = torch.randn(1, 3, 224, 224, requires_grad=True)
-        test_inputs = torch.randn(3, 3, 224, 224, requires_grad=True)
+        dummy_input = torch.randn(1, 3, 224, 224, requires_grad=False)
+        test_inputs = torch.randn(3, 3, 224, 224, requires_grad=False)
 
         self.run_test_with_fx_to_onnx_exporter_and_onnx_runtime(
             model,
@@ -374,11 +375,13 @@ def forward(self, x):
             additional_test_inputs=[((y,),)],
         )
 
-    @pytorch_test_common.xfail("torch._dynamo.exc.TorchRuntimeError")
+    @pytorch_test_common.xfail(
+        "torch._dynamo.exc.Unsupported: guard on data-dependent symbolic int/float"
+    )
     def test_squeeze_runtime_dim(self):
         class Squeeze(torch.nn.Module):
             def forward(self, d1, d2):
-                t = torch.zeros(d1[0], d2[0])
+                t = torch.zeros(d1[0], d2[0])  # problematic user code for dynamo
                 return t.squeeze(0)
 
         d1 = torch.tensor([1])
@@ -435,8 +438,8 @@ def forward(self, input):
         )
 
     @pytorch_test_common.xfail(
-        "fx.graph: torch._subclasses.fake_tensor.DataDependentOutputException: "
-        "aten._local_scalar_dense.default"
+        "[ONNXRuntimeError] : 1 : FAIL : Non-zero status code returned while running Slice node. Name:'n13__5' Status Message:"
+        "slice.cc:193 FillVectorsFromInput Starts must be a 1-D array"
     )
     def test_expand_as_fill_zero(self):
         class Model(torch.nn.Module):
@@ -453,7 +456,8 @@ def forward(self, x):
         )
 
     @pytorch_test_common.xfail(
-        "RuntimeError: Unknown call_function target: aten.lift_fresh_copy.default"
+        "[ONNXRuntimeError] : 1 : FAIL : Type Error: Type (tensor(float)) of output arg (copy) "
+        "of node (n0__4) does not match expected type (tensor(int64))"
     )
     def test_expand_as_fill_tensor(self):
         class Model(torch.nn.Module):
@@ -470,7 +474,8 @@ def forward(self, x):
         )
 
     @pytorch_test_common.xfail(
-        "Unknown call_function target: aten.lift_fresh_copy.default"
+        "RuntimeError: at::functionalization::impl::isFunctionalTensor(self_) INTERNAL ASSERT FAILED "
+        "at '/path/to/pytorch/torch/csrc/autograd/python_torch_functions_manual.cpp':514, please report a bug to PyTorch."
     )
     def test_expand_as_fill_seperate_tensor(self):
         class Model(torch.nn.Module):
@@ -998,16 +1003,15 @@ def create_kwargs():
         )
 
     @pytorch_test_common.xfail(
-        "ValueError: Message onnx.ModelProto exceeds maximum protobuf size of 2GB"
-        "[ONNXRuntimeError] : Status Message: Indices vs updates dimensions differs at position=1 0 vs 3"
+        "Constant tensor is not supported in FakeTensorMode export."
     )
     @pytorch_test_common.skip_dynamic_fx_test(
         "RuntimeError:: SymIntArrayRef expected to contain only concrete integers"
     )
     @pytorch_test_common.skip_load_checkpoint_after_model_creation(
         "HF Bloom model does not need `model.load_state_dict` to work."
     )
-    def test_fake_tensor_mode_huggingface_bigscience__bloom_560m(self):
+    def test_fake_tensor_mode_huggingface_bigscience_bloom_560m(self):
         from transformers import AutoModel, AutoTokenizer  # type: ignore[import]
 
         model_name = "bigscience/bloom-560m"