Update ONNX's IO Adapter to support FakeTensor with ExportedProgram (#114407) (#115578)

Currently, the ONNX exporter using torch.nn.Module as input can support
FakeTensor because the ONNX model stores all initializers

When using torch.export.ExportedProgram as input, the initializers are
lifted as inputs. In order to execute the ONNX model, we need to pass a
reference to the non-fake model to the
ONNXProgram.adapt_torch_inputs_to_onnx API, so that initializers can be
fetched from the model and fed to the ONNX model as input

ps: #115461 will track the API revision for the cases where additional `model_with_state_dict` are required to produce complete ONNX files exported with fake support. This is also tracked by the umbrella fake tensor issue #105464 FYI @BowenBao
Pull Request resolved: #114407
Approved by: https://github.com/BowenBao

Author: Thiago Crepaldi

test/onnx/onnx_test_common.py

@@ -439,9 +439,9 @@ def _compare_pytorch_onnx_with_ort(
     # ONNXProgram holds a reference (not copy) to the original ref_model, including its state_dict.
     # Thus, ONNXProgram() must run before ref_model() to prevent ref_model.forward() from changing the state_dict.
     # Otherwise, the ref_model can change buffers on state_dict which would be used by ONNXProgram.__call__()
-    ort_outputs = onnx_program(*input_args, **input_kwargs)
+    ort_outputs = onnx_program(*input_args, model=ref_model, **input_kwargs)
     ref_outputs = ref_model(*ref_input_args, **ref_input_kwargs)
-    ref_outputs = onnx_program.adapt_torch_outputs_to_onnx(ref_outputs)
+    ref_outputs = onnx_program.adapt_torch_outputs_to_onnx(ref_model, ref_outputs)
 
     if len(ref_outputs) != len(ort_outputs):
         raise AssertionError(

test/onnx/pytorch_test_common.py

@@ -188,11 +188,13 @@ def wrapper(self, *args, **kwargs):
     return skip_dec
 
 
-def skip_dynamic_fx_test(reason: str):
+def skip_dynamic_fx_test(reason: str, skip_model_type=None):
     """Skip dynamic exporting test.
 
     Args:
         reason: The reason for skipping dynamic exporting test.
+        skip_model_type (onnx_test_common.TorchModelType): The model type to skip dynamic exporting test for.
+            When None, model type is not used to skip dynamic tests.
 
     Returns:
         A decorator for skipping dynamic exporting test.
@@ -201,7 +203,9 @@ def skip_dynamic_fx_test(reason: str):
     def skip_dec(func):
         @functools.wraps(func)
         def wrapper(self, *args, **kwargs):
-            if self.dynamic_shapes:
+            if self.dynamic_shapes and (
+                not skip_model_type or self.model_type == skip_model_type
+            ):
                 raise unittest.SkipTest(
                     f"Skip verify dynamic shapes test for FX. {reason}"
                 )

test/onnx/test_fx_to_onnx_with_onnxruntime.py

@@ -198,15 +198,23 @@ def func(x, b=1.0):
             ),
         )
         onnx_test_common.assert_dynamic_shapes(onnx_program, self.dynamic_shapes)
-        onnx_format_args = onnx_program.adapt_torch_inputs_to_onnx(tensor_x, 8.0)
-        ref_outputs = onnx_program.adapt_torch_outputs_to_onnx(func(tensor_x, 8.0))
+        onnx_format_args = onnx_program.adapt_torch_inputs_to_onnx(
+            tensor_x, model=func, b=8.0
+        )
+        ref_outputs = onnx_program.adapt_torch_outputs_to_onnx(
+            func, func(tensor_x, 8.0)
+        )
         ort_outputs = onnx_test_common.run_ort(onnx_program, onnx_format_args)
         for ref_output, ort_output in zip(ref_outputs, ort_outputs):
             torch.testing.assert_close(ref_output, torch.tensor(ort_output))
 
         # test on different non-tensor input - xfail
-        onnx_format_args = onnx_program.adapt_torch_inputs_to_onnx(tensor_x, 9.0)
-        ref_outputs = onnx_program.adapt_torch_outputs_to_onnx(func(tensor_x, 9.0))
+        onnx_format_args = onnx_program.adapt_torch_inputs_to_onnx(
+            tensor_x, model=func, b=9.0
+        )
+        ref_outputs = onnx_program.adapt_torch_outputs_to_onnx(
+            func, func(tensor_x, 9.0)
+        )
         _ = onnx_test_common.run_ort(onnx_program, onnx_format_args)
         for ref_output, ort_output in zip(ref_outputs, ort_outputs):
             torch.testing.assert_close(ref_output, torch.tensor(ort_output))
@@ -538,7 +546,7 @@ def forward(self, x):
             additional_test_inputs=[((x2,),)],
         )
 
-    @pytorch_test_common.xfail(
+    @pytorch_test_common.xfail_if_model_type_is_not_exportedprogram(
         "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."
     )
@@ -831,10 +839,10 @@ def _test_fx_symbolic_tracer_large_scale_exporter(
             kwargs = create_pytorch_only_kwargs()
             # Original outputs.
             ref_outputs = onnx_program.adapt_torch_outputs_to_onnx(
-                model(*args, **kwargs)
+                model, model(*args, **kwargs)
             )
             # ORT outputs.
-            args_not_none = onnx_program.adapt_torch_inputs_to_onnx(*args)
+            args_not_none = onnx_program.adapt_torch_inputs_to_onnx(*args, model=model)
 
             # Drop Parameters and buffers added by fx_serialization.save_model_with_external_data
             args_not_none = args_not_none[: len(args) - len(kwargs)]
@@ -923,14 +931,24 @@ def create_pytorch_only_extra_kwargs():
 def _parameterized_class_attrs_and_values_with_fake_options():
     input_values = []
     input_values.extend(
-        itertools.product((True, False), (True, False), (True, False), (True, False))
+        itertools.product(
+            (True, False),
+            (True, False),
+            (True, False),
+            (True, False),
+            (
+                onnx_test_common.TorchModelType.TORCH_NN_MODULE,
+                onnx_test_common.TorchModelType.TORCH_EXPORT_EXPORTEDPROGRAM,
+            ),
+        )
     )
     return {
         "attrs": [
             "op_level_debug",
             "dynamic_shapes",
             "load_checkpoint_during_init",
             "export_within_fake_mode",
+            "model_type",
         ],
         "input_values": input_values,
     }
@@ -950,6 +968,7 @@ class TestFxToOnnxFakeTensorWithOnnxRuntime(onnx_test_common._TestONNXRuntime):
     dynamic_shapes: bool
     load_checkpoint_during_init: bool
     export_within_fake_mode: bool
+    model_type: onnx_test_common.TorchModelType
 
     def setUp(self):
         super().setUp()
@@ -964,6 +983,7 @@ def _test_fake_tensor_mode_exporter(
         create_kwargs: Callable,
         load_checkpoint_during_init: bool,
         export_within_fake_mode: bool,
+        model_type: onnx_test_common.TorchModelType,
     ):
         """Test helper for FakeTensorMode-enabled exporter.
 
@@ -975,6 +995,8 @@ def _test_fake_tensor_mode_exporter(
             load_checkpoint_during_init: Whether to load a checkpoint during model initialization.
                 (after or during model creation, but before exporting starts)
             export_within_fake_mode: Whether to call torch.onnx._dynamo_export within torch._subclasses.FakeTensorMode
+            model_type: Type of user model. Used to determine whether the user model must be exported to
+                torch.export.ExportedProgram before passing it to torch.onnx.dynamo_export
 
         This test contains several steps.
 
@@ -990,13 +1012,17 @@ def _test_fake_tensor_mode_exporter(
 
         # Create the toy model with real weight.
         real_model = create_model()
+        state_dict = real_model.state_dict()  # concrete (non-fake) state_dict
+        if model_type == onnx_test_common.TorchModelType.TORCH_EXPORT_EXPORTEDPROGRAM:
+            real_model = torch.export.export(
+                real_model, args=create_args(), kwargs=create_kwargs()
+            )
 
         with tempfile.NamedTemporaryFile(
             prefix=model_name, suffix=".pt"
         ) as tmp_checkpoint_file:
             # Dump state_dict to a file to simulate how HuggingFace model is initialized.
             # The file will be loaded via .load_state_dict(...)
-            state_dict = real_model.state_dict()
             torch.save(state_dict, tmp_checkpoint_file.name)
 
             with torch.onnx.enable_fake_mode() as fake_context:
@@ -1014,6 +1040,13 @@ def _test_fake_tensor_mode_exporter(
                 )
 
                 if export_within_fake_mode:
+                    if (
+                        model_type
+                        == onnx_test_common.TorchModelType.TORCH_EXPORT_EXPORTEDPROGRAM
+                    ):
+                        fake_model = torch.export.export(
+                            fake_model, args=fake_args, kwargs=fake_kwargs
+                        )
                     onnx_program = torch.onnx.dynamo_export(
                         fake_model,
                         *fake_args,
@@ -1022,6 +1055,13 @@ def _test_fake_tensor_mode_exporter(
                     )
 
             if not export_within_fake_mode:
+                if (
+                    model_type
+                    == onnx_test_common.TorchModelType.TORCH_EXPORT_EXPORTEDPROGRAM
+                ):
+                    fake_model = torch.export.export(
+                        fake_model, args=fake_args, kwargs=fake_kwargs
+                    )
                 onnx_program = torch.onnx.dynamo_export(
                     fake_model, *fake_args, **fake_kwargs, export_options=export_options
                 )
@@ -1038,10 +1078,12 @@ def _test_fake_tensor_mode_exporter(
                 kwargs = create_kwargs()
                 # Original outputs.
                 ref_outputs = onnx_program.adapt_torch_outputs_to_onnx(
-                    real_model(*args, **kwargs)
+                    fake_model, real_model(*args, **kwargs)
                 )
                 # ORT outputs.
-                args_not_none = onnx_program.adapt_torch_inputs_to_onnx(*args, **kwargs)
+                args_not_none = onnx_program.adapt_torch_inputs_to_onnx(
+                    *args, model=real_model, **kwargs
+                )
 
                 ort_outputs = onnx_test_common.run_ort(
                     tmp_onnx_file.name,
@@ -1053,6 +1095,10 @@ def _test_fake_tensor_mode_exporter(
                 for ref_output, ort_output in zip(ref_outputs, ort_outputs):
                     torch.testing.assert_close(ref_output, torch.tensor(ort_output))
 
+    @pytorch_test_common.skip_dynamic_fx_test(
+        "AssertionError: Dynamic shape check failed for graph inputs",
+        skip_model_type=onnx_test_common.TorchModelType.TORCH_EXPORT_EXPORTEDPROGRAM,
+    )
     def test_fake_tensor_mode_simple(self):
         def create_model() -> nn.Module:
             class Model(torch.nn.Module):
@@ -1079,14 +1125,19 @@ def create_kwargs():
             create_kwargs,
             load_checkpoint_during_init=self.load_checkpoint_during_init,
             export_within_fake_mode=self.export_within_fake_mode,
+            model_type=self.model_type,
         )
 
-    @pytorch_test_common.xfail(
+    @pytorch_test_common.xfail_if_model_type_is_not_exportedprogram(
         "[ONNXRuntimeError] : 1 : FAIL : Type Error: Data in initializer 'h_0_attn_bias' "
         "has element type tensor(uint8) but usage of initializer in graph expects tensor(bool)"
         "https://github.com/huggingface/transformers/issues/21013"
         "This can be addressed by using GPT2Config, but it is not now supported by FakeTensor exporting."
     )
+    @pytorch_test_common.skip_dynamic_fx_test(
+        "AssertionError: Dynamic shape check failed for graph inputs",
+        skip_model_type=onnx_test_common.TorchModelType.TORCH_EXPORT_EXPORTEDPROGRAM,
+    )
     def test_large_scale_exporter_with_tiny_gpt2(self):
         model_name = "sshleifer/tiny-gpt2"
         device = "cpu"
@@ -1111,8 +1162,13 @@ def create_kwargs():
             create_kwargs,
             load_checkpoint_during_init=self.load_checkpoint_during_init,
             export_within_fake_mode=self.export_within_fake_mode,
+            model_type=self.model_type,
         )
 
+    @pytorch_test_common.skip_dynamic_fx_test(
+        "AssertionError: Dynamic shape check failed for graph inputs",
+        skip_model_type=onnx_test_common.TorchModelType.TORCH_EXPORT_EXPORTEDPROGRAM,
+    )
     def test_large_scale_exporter_with_toy_mlp(self):
         class MLPModel(nn.Module):
             def __init__(self):
@@ -1148,8 +1204,13 @@ def create_kwargs():
             create_kwargs,
             load_checkpoint_during_init=self.load_checkpoint_during_init,
             export_within_fake_mode=self.export_within_fake_mode,
+            model_type=self.model_type,
         )
 
+    @pytorch_test_common.skip_dynamic_fx_test(
+        "AssertionError: Dynamic shape check failed for graph inputs",
+        skip_model_type=onnx_test_common.TorchModelType.TORCH_EXPORT_EXPORTEDPROGRAM,
+    )
     def test_fake_tensor_mode_huggingface_google_t5(self):
         config = transformers.T5Config(
             vocab_size=8096, d_model=64, num_layers=2, num_heads=2
@@ -1179,8 +1240,13 @@ def create_model():
             create_kwargs,
             load_checkpoint_during_init=self.load_checkpoint_during_init,
             export_within_fake_mode=self.export_within_fake_mode,
+            model_type=self.model_type,
         )
 
+    @pytorch_test_common.skip_dynamic_fx_test(
+        "AssertionError: Dynamic shape check failed for graph inputs",
+        skip_model_type=onnx_test_common.TorchModelType.TORCH_EXPORT_EXPORTEDPROGRAM,
+    )
     def test_fake_tensor_mode_huggingface_openai_whisper(self):
         config = transformers.WhisperConfig(
             vocab_size=8096,
@@ -1231,6 +1297,7 @@ def create_kwargs():
             create_kwargs,
             load_checkpoint_during_init=self.load_checkpoint_during_init,
             export_within_fake_mode=self.export_within_fake_mode,
+            model_type=self.model_type,
         )
 
     @pytorch_test_common.xfail(
@@ -1260,6 +1327,7 @@ def create_model():
             create_kwargs,
             load_checkpoint_during_init=self.load_checkpoint_during_init,
             export_within_fake_mode=self.export_within_fake_mode,
+            model_type=self.model_type,
         )
 
     @pytorch_test_common.skip_dynamic_fx_test(
@@ -1287,6 +1355,7 @@ def create_model():
             create_kwargs,
             load_checkpoint_during_init=self.load_checkpoint_during_init,
             export_within_fake_mode=self.export_within_fake_mode,
+            model_type=self.model_type,
         )
 
 

test/onnx/torch_export/test_torch_export_with_onnxruntime.py

@@ -31,10 +31,12 @@ def _compare_onnx_and_torch_exported_program(
         # NOTE: ONNXProgram holds a reference (not copy) to the original ref_model, including its state_dict.
         # Thus, ONNXProgram() must run before ref_model() to prevent ref_model.forward() from changing the state_dict.
         # Otherwise, the ref_model can change buffers on state_dict which would be used by ONNXProgram.__call__()
-        onnx_outputs = onnx_exported_program(*input_args, **input_kwargs)
+        onnx_outputs = onnx_exported_program(
+            *input_args, model=torch_exported_program, **input_kwargs
+        )
         torch_outputs = torch_exported_program(*input_args, **input_kwargs)
         torch_outputs_onnx_format = onnx_exported_program.adapt_torch_outputs_to_onnx(
-            torch_outputs
+            torch_exported_program, torch_outputs
         )
         if len(torch_outputs_onnx_format) != len(onnx_outputs):
             raise AssertionError(