fix torch.export 0/1 specializing

Author: titaiwangms

_unittests/ut_tasks/test_tasks.py

@@ -43,7 +43,9 @@ def test_text_generation(self):
         model, inputs, ds = data["model"], data["inputs"], data["dynamic_shapes"]
         model(**inputs)
         model(**data["inputs2"])
-        with torch_export_patches(patch_transformers=True, verbose=10):
+        with torch_export_patches(
+            patch_transformers=True, verbose=10
+        ), torch.fx.experimental._config.patch(backed_size_oblivious=True):
             torch.export.export(
                 model, (), kwargs=inputs, dynamic_shapes=use_dyn_not_str(ds), strict=False
             )

_unittests/ut_torch_export_patches/test_dynamic_class.py

@@ -307,7 +307,9 @@ def test_phi2_export_module(self):
             str_inputs, string_type(inputs_copied, with_shape=True, with_min_max=True)
         )
 
-        with torch_export_patches(patch_transformers=True):
+        with torch_export_patches(
+            patch_transformers=True
+        ), torch.fx.experimental._config.patch(backed_size_oblivious=True):
             ep = torch.export.export(
                 model,
                 (),
@@ -346,7 +348,9 @@ def test_phi2_export_interpreter(self):
             str_inputs, string_type(inputs_copied, with_shape=True, with_min_max=True)
         )
 
-        with torch_export_patches(patch_transformers=True, verbose=1):
+        with torch_export_patches(
+            patch_transformers=True, verbose=1
+        ), torch.fx.experimental._config.patch(backed_size_oblivious=True):
             if masking_utils is not None:
                 self.assertEqual(
                     masking_utils.ALL_MASK_ATTENTION_FUNCTIONS["sdpa"],

onnx_diagnostic/tasks/text_generation.py

@@ -221,7 +221,7 @@ def get_inputs(
                 "input_ids": {0: batch, 1: seq_length},
                 "attention_mask": {
                     0: batch,
-                    1: "cache+seq",  # past_seq_length + seq_length
+                    1: "past_seq_length+seq_length",  # past_seq_length + seq_length
                 },
                 "position_ids": {
                     0: batch,

onnx_diagnostic/torch_models/validate.py

@@ -1105,16 +1105,24 @@ def call_torch_export_export(
         print("[call_torch_export_export] export...")
 
     model = data["model"]
+
+    def _run_torch_export():
+        with torch.fx.experimental._config.patch(backed_size_oblivious=True):
+            ep = torch.export.export(
+                model,
+                args,
+                kwargs=kwargs,
+                dynamic_shapes=dse,
+                strict=strict,
+            )
+        return ep
+
     ep = _quiet_or_not_quiet(
         quiet,
         "export_export",
         summary,
         data,
-        (
-            lambda m=model, args=args, kws=kwargs, dse=dse, s=strict: (
-                torch.export.export(m, args, kwargs=kws, dynamic_shapes=dse, strict=s)
-            )
-        ),
+        _run_torch_export,
     )
     if "ERR_export_export" in summary:
         return summary, data
@@ -1715,23 +1723,24 @@ def call_torch_export_custom(
         kws["target_opset"] = opset
     if output_names:
         kws["output_names"] = output_names
-
-    epo, opt_stats = _quiet_or_not_quiet(
-        quiet,
-        "export_export_onnx_c",
-        summary,
-        data,
-        (
-            lambda m=model, args=args, kwargs=kwargs, kws=kws: (
-                to_onnx(
-                    model,
-                    args,
-                    kwargs=kwargs,
-                    **kws,
+    # anti-specializing 0/1 during torch.export.export
+    with torch.fx.experimental._config.patch(backed_size_oblivious=True):
+        epo, opt_stats = _quiet_or_not_quiet(
+            quiet,
+            "export_export_onnx_c",
+            summary,
+            data,
+            (
+                lambda m=model, args=args, kwargs=kwargs, kws=kws: (
+                    to_onnx(
+                        model,
+                        args,
+                        kwargs=kwargs,
+                        **kws,
+                    )
                 )
-            )
-        ),
-    )
+            ),
+        )
     if "ERR_export_onnx_c" in summary:
         return summary, data