executorch/backends/xnnpack/test/ops

Differential Revision: D72436120

Pull Request resolved: #9956

Author: gmagogsfm

backends/xnnpack/test/ops/test_check_quant_params.py

@@ -52,7 +52,7 @@ def _test_check_quant_message(self, ep_modifier, expected_message):
         torch._dynamo.reset()
         mod = torch.nn.Linear(10, 10)
         quantizer = XNNPACKQuantizer()
-        captured = export_for_training(mod, (torch.randn(1, 10),)).module()
+        captured = export_for_training(mod, (torch.randn(1, 10),), strict=True).module()
         quantizer.set_global(get_symmetric_quantization_config(is_per_channel=True))
         prepared = prepare_pt2e(captured, quantizer)
 
@@ -68,7 +68,6 @@ def _test_check_quant_message(self, ep_modifier, expected_message):
         self.assertEquals(str(context.exception), expected_message)
 
     def test_in_per_tensor_quant(self):
-
         for invalid_scale in [
             float("nan"),
             float("inf"),

examples/llm_manual/export_nanogpt.py

@@ -28,7 +28,7 @@
 # The torch.no_grad() call tells PyTorch to exclude training-specific logic.
 with sdpa_kernel([SDPBackend.MATH]), torch.no_grad():
     m = export_for_training(
-        model, example_inputs, dynamic_shapes=dynamic_shape
+        model, example_inputs, dynamic_shapes=dynamic_shape, strict=True
     ).module()
     traced_model = export(m, example_inputs, dynamic_shapes=dynamic_shape, strict=True)
 

examples/mediatek/aot_utils/oss_utils/utils.py

@@ -30,7 +30,9 @@ def build_executorch_binary(
         if quant_dtype not in Precision:
             raise AssertionError(f"No support for Precision {quant_dtype}.")
 
-        captured_model = torch.export.export_for_training(model, inputs).module()
+        captured_model = torch.export.export_for_training(
+            model, inputs, strict=True
+        ).module()
         annotated_model = prepare_pt2e(captured_model, quantizer)
         print("Quantizing the model...")
         # calibration

examples/mediatek/model_export_scripts/llama.py

@@ -319,7 +319,7 @@ def export_to_et_ir(
     )
     print("Getting pre autograd ATen Dialect Graph")
     pre_autograd_aten_dialect = torch.export.export_for_training(
-        model, example_inputs, dynamic_shapes=dynamic_shapes
+        model, example_inputs, dynamic_shapes=dynamic_shapes, strict=True
     ).module()  # NOTE: Will be replaced with export
     quantizer = NeuropilotQuantizer()
     quantizer.setup_precision(getattr(Precision, precision))

examples/models/phi-3-mini/export_phi-3-mini.py

@@ -65,7 +65,7 @@ def export(args) -> None:
         xnnpack_quantizer.set_global(xnnpack_quant_config)
 
         model = export_for_training(
-            model, example_inputs, dynamic_shapes=dynamic_shapes
+            model, example_inputs, dynamic_shapes=dynamic_shapes, strict=True
         ).module()
         model = prepare_pt2e(model, xnnpack_quantizer)  # pyre-fixme[6]
         model(*example_inputs)

examples/models/test/test_export.py

@@ -29,7 +29,9 @@ def collect_executorch_and_eager_outputs(
         Returns a tuple containing the outputs of the eager mode model and the executorch mode model.
         """
         eager_model = eager_model.eval()
-        model = torch.export.export_for_training(eager_model, example_inputs).module()
+        model = torch.export.export_for_training(
+            eager_model, example_inputs, strict=True
+        ).module()
         edge_model = export_to_edge(model, example_inputs)
 
         executorch_prog = edge_model.to_executorch()

examples/portable/scripts/export_and_delegate.py

@@ -61,7 +61,7 @@ def export_composite_module_with_lower_graph():
     m_compile_spec = m.get_compile_spec()
 
     # pre-autograd export. eventually this will become torch.export
-    m = torch.export.export_for_training(m, m_inputs).module()
+    m = torch.export.export_for_training(m, m_inputs, strict=True).module()
     edge = export_to_edge(m, m_inputs)
     logging.info(f"Exported graph:\n{edge.exported_program().graph}")
 
@@ -84,7 +84,7 @@ def forward(self, *args):
     m = CompositeModule()
     m = m.eval()
     # pre-autograd export. eventually this will become torch.export
-    m = torch.export.export_for_training(m, m_inputs).module()
+    m = torch.export.export_for_training(m, m_inputs, strict=True).module()
     composited_edge = export_to_edge(m, m_inputs)
 
     # The graph module is still runnerable
@@ -134,7 +134,7 @@ def get_example_inputs(self):
     m = Model()
     m_inputs = m.get_example_inputs()
     # pre-autograd export. eventually this will become torch.export
-    m = torch.export.export_for_training(m, m_inputs).module()
+    m = torch.export.export_for_training(m, m_inputs, strict=True).module()
     edge = export_to_edge(m, m_inputs)
     logging.info(f"Exported graph:\n{edge.exported_program().graph}")
 
@@ -171,7 +171,7 @@ def export_and_lower_the_whole_graph():
 
     m_inputs = m.get_example_inputs()
     # pre-autograd export. eventually this will become torch.export
-    m = torch.export.export_for_training(m, m_inputs).module()
+    m = torch.export.export_for_training(m, m_inputs, strict=True).module()
     edge = export_to_edge(m, m_inputs)
     logging.info(f"Exported graph:\n{edge.exported_program().graph}")
 

examples/xnnpack/aot_compiler.py

@@ -87,14 +87,14 @@
 
     model = model.eval()
     # pre-autograd export. eventually this will become torch.export
-    ep = torch.export.export_for_training(model, example_inputs)
+    ep = torch.export.export_for_training(model, example_inputs, strict=True)
     model = ep.module()
 
     if args.quantize:
         logging.info("Quantizing Model...")
         # TODO(T165162973): This pass shall eventually be folded into quantizer
         model = quantize(model, example_inputs, quant_type)
-        ep = torch.export.export_for_training(model, example_inputs)
+        ep = torch.export.export_for_training(model, example_inputs, strict=True)
 
     edge = to_edge_transform_and_lower(
         ep,

examples/xnnpack/quantization/example.py

@@ -60,7 +60,9 @@ def verify_xnnpack_quantizer_matching_fx_quant_model(model_name, model, example_
     m = model
 
     # 1. pytorch 2.0 export quantization flow (recommended/default flow)
-    m = torch.export.export_for_training(m, copy.deepcopy(example_inputs)).module()
+    m = torch.export.export_for_training(
+        m, copy.deepcopy(example_inputs), strict=True
+    ).module()
     quantizer = XNNPACKQuantizer()
     quantization_config = get_symmetric_quantization_config(is_per_channel=True)
     quantizer.set_global(quantization_config)
@@ -177,7 +179,9 @@ def main() -> None:
 
     model = model.eval()
     # pre-autograd export. eventually this will become torch.export
-    model = torch.export.export_for_training(model, example_inputs).module()
+    model = torch.export.export_for_training(
+        model, example_inputs, strict=True
+    ).module()
     start = time.perf_counter()
     quantized_model = quantize(model, example_inputs)
     end = time.perf_counter()

exir/backend/test/test_partitioner.py

@@ -76,7 +76,7 @@ def partition(
 
         mlp = MLP()
         example_inputs = mlp.get_random_inputs()
-        model = export_for_training(mlp, example_inputs).module()
+        model = export_for_training(mlp, example_inputs, strict=True).module()
         aten = export(model, example_inputs, strict=True)
         spec_key = "path"
         spec_value = "/a/b/c/d"
@@ -137,7 +137,7 @@ def partition(
 
         mlp = MLP()
         example_inputs = mlp.get_random_inputs()
-        model = export_for_training(mlp, example_inputs).module()
+        model = export_for_training(mlp, example_inputs, strict=True).module()
         aten = export(model, example_inputs, strict=True)
         edge = exir.to_edge(aten)
 
@@ -177,7 +177,7 @@ def partition(
 
         mlp = MLP()
         example_inputs = mlp.get_random_inputs()
-        model = export_for_training(mlp, example_inputs).module()
+        model = export_for_training(mlp, example_inputs, strict=True).module()
         edge = exir.to_edge(export(model, example_inputs, strict=True))
 
         with self.assertRaisesRegex(
@@ -229,7 +229,9 @@ def partition(
                     partition_tags=partition_tags,
                 )
 
-        model = export_for_training(self.AddConst(), (torch.ones(2, 2),)).module()
+        model = export_for_training(
+            self.AddConst(), (torch.ones(2, 2),), strict=True
+        ).module()
         edge = exir.to_edge(export(model, (torch.ones(2, 2),), strict=True))
         delegated = edge.to_backend(PartitionerNoTagData())
 
@@ -308,7 +310,9 @@ def partition(
                     partition_tags=partition_tags,
                 )
 
-        model = export_for_training(self.AddConst(), (torch.ones(2, 2),)).module()
+        model = export_for_training(
+            self.AddConst(), (torch.ones(2, 2),), strict=True
+        ).module()
         edge = exir.to_edge(export(model, (torch.ones(2, 2),), strict=True))
         delegated = edge.to_backend(PartitionerTagData())
 
@@ -383,7 +387,9 @@ def partition(
                     partition_tags=partition_tags,
                 )
 
-        model = export_for_training(self.AddConst(), (torch.ones(2, 2),)).module()
+        model = export_for_training(
+            self.AddConst(), (torch.ones(2, 2),), strict=True
+        ).module()
         edge = exir.to_edge(export(model, (torch.ones(2, 2),), strict=True))
         delegated = edge.to_backend(PartitionerTagData())
 
@@ -471,7 +477,9 @@ def partition(
                 )
 
         inputs = (torch.ones(2, 2),)
-        model = export_for_training(ReuseConstData(), (torch.ones(2, 2),)).module()
+        model = export_for_training(
+            ReuseConstData(), (torch.ones(2, 2),), strict=True
+        ).module()
         edge = exir.to_edge(export(model, (torch.ones(2, 2),), strict=True))
         exec_prog = edge.to_backend(PartitionerTagData()).to_executorch()
         executorch_module = _load_for_executorch_from_buffer(exec_prog.buffer)
@@ -531,7 +539,9 @@ def partition(
                     partition_tags=partition_tags,
                 )
 
-        model = export_for_training(ReuseConstData(), (torch.ones(2, 2),)).module()
+        model = export_for_training(
+            ReuseConstData(), (torch.ones(2, 2),), strict=True
+        ).module()
         edge = exir.to_edge(export(model, (torch.ones(2, 2),), strict=True))
         with self.assertRaises(RuntimeError) as error:
             _ = edge.to_backend(PartitionerTagData())

exir/backend/test/test_passes.py

@@ -28,7 +28,9 @@ def forward(self, x):
                 z = x - self.const
                 return y, z
 
-        model = export_for_training(ReuseConstData(), (torch.ones(2, 2),)).module()
+        model = export_for_training(
+            ReuseConstData(), (torch.ones(2, 2),), strict=True
+        ).module()
         edge = exir.to_edge(
             torch.export.export(model, (torch.ones(2, 2),), strict=True)
         )

exir/emit/test/test_emit.py

@@ -1751,8 +1751,8 @@ def forward(self, x):
         module_1(*example_inputs)
         module_2(*example_inputs)
 
-        ep1 = export_for_training(module_1, example_inputs)
-        ep2 = export_for_training(module_2, example_inputs)
+        ep1 = export_for_training(module_1, example_inputs, strict=True)
+        ep2 = export_for_training(module_2, example_inputs, strict=True)
 
         edge_program_manager = exir.to_edge(
             {"forward1": ep1, "forward2": ep2},

exir/tests/test_passes.py

@@ -1164,7 +1164,9 @@ def forward(self, query, key, value):
         value = torch.randn(32, 32, 32, 32)
 
         # Capture the model
-        m = torch.export.export_for_training(M(32), (query, key, value)).module()
+        m = torch.export.export_for_training(
+            M(32), (query, key, value), strict=True
+        ).module()
 
         # 8w16a quantization
         from torch.ao.quantization.observer import (
@@ -1405,8 +1407,7 @@ def quantize_model(
         ) -> Tuple[EdgeProgramManager, int, int]:
             # program capture
             m = torch.export.export_for_training(
-                m_eager,
-                example_inputs,
+                m_eager, example_inputs, strict=True
             ).module()
 
             quantizer = XNNPACKQuantizer()

exir/tests/test_quantization.py

@@ -52,7 +52,7 @@ def test_resnet(self) -> None:
             m_copy = copy.deepcopy(m)
             # program capture
             m = torch.export.export_for_training(
-                m, copy.deepcopy(example_inputs)
+                m, copy.deepcopy(example_inputs), strict=True
             ).module()
 
             quantizer = XNNPACKQuantizer()

exir/tests/test_quantize_io_pass.py

@@ -39,12 +39,14 @@ def _quantize(self, mod, example_inputs):
         operator_config = get_symmetric_quantization_config()
         quantizer.set_global(operator_config)
         m = torch.export.export_for_training(
-            mod, copy.deepcopy(example_inputs)
+            mod, copy.deepcopy(example_inputs), strict=True
         ).module()
         m = prepare_pt2e(m, quantizer)
         _ = m(*example_inputs)
         m = convert_pt2e(m)
-        exported_program = torch.export.export_for_training(m, example_inputs)
+        exported_program = torch.export.export_for_training(
+            m, example_inputs, strict=True
+        )
         return exported_program
 
     def _check_count(self, op, count, epm):

extension/export_util/utils.py

@@ -108,7 +108,7 @@ def export_to_exec_prog(
 ) -> ExecutorchProgramManager:
     m = model.eval()
     # pre-autograd export. eventually this will become torch.export
-    m = export_for_training(m, example_inputs).module()
+    m = export_for_training(m, example_inputs, strict=True).module()
 
     core_aten_ep = _to_core_aten(
         m,

extension/llm/export/builder.py

@@ -234,6 +234,7 @@ def _export(self, module: Optional[torch.nn.Module] = None) -> ExportedProgram:
                     self.example_inputs,
                     kwargs=self.example_kwarg_inputs,
                     dynamic_shapes=dynamic_shape,
+                    strict=True,
                 )
         return exported_module
 

extension/llm/export/test_export_passes.py

@@ -10,10 +10,7 @@
 
 class RemoveRedundantTransposesPassTest(unittest.TestCase):
     def _export(self, model, example_inputs):
-        exported_module = export_for_training(
-            model,
-            example_inputs,
-        )
+        exported_module = export_for_training(model, example_inputs, strict=True)
         return exported_module.module()
 
     def _check(self, model, example_inputs, key, before_count, after_count):