Migrate to training IR in executorch tests (#5875)

Summary:

as title

Differential Revision: D63845664

Author: yushangdi

.ci/docker/ci_commit_pins/pytorch.txt

@@ -1 +1 @@
-5ba404f68775bb06a1125a100687f86b6d6de6a8
+b2979f4382fba5b5743793bbac2f2c92a46ae901

examples/apple/mps/scripts/mps_example.py

@@ -166,7 +166,7 @@ def get_model_config(args):
 
     # pre-autograd export. eventually this will become torch.export
     with torch.no_grad():
-        model = torch._export.capture_pre_autograd_graph(model, example_inputs)
+        model = torch.export.export_for_training(model, example_inputs).module()
         edge: EdgeProgramManager = export_to_edge(
             model,
             example_inputs,

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

@@ -15,13 +15,13 @@
 from executorch.backends.xnnpack.partition.xnnpack_partitioner import XnnpackPartitioner
 from executorch.backends.xnnpack.utils.configs import get_xnnpack_edge_compile_config
 from executorch.exir import to_edge
-from torch._export import capture_pre_autograd_graph
 from torch.ao.quantization.quantize_pt2e import convert_pt2e, prepare_pt2e
 
 from torch.ao.quantization.quantizer.xnnpack_quantizer import (
     get_symmetric_quantization_config,
     XNNPACKQuantizer,
 )
+from torch.export import export_for_training
 
 from transformers import Phi3ForCausalLM
 
@@ -64,9 +64,9 @@ def export(args) -> None:
         xnnpack_quantizer = XNNPACKQuantizer()
         xnnpack_quantizer.set_global(xnnpack_quant_config)
 
-        model = capture_pre_autograd_graph(
+        model = export_for_training(
             model, example_inputs, dynamic_shapes=dynamic_shapes
-        )
+        ).module()
         model = prepare_pt2e(model, xnnpack_quantizer)  # pyre-fixme[6]
         model(*example_inputs)
         model = convert_pt2e(model)

exir/tests/test_passes.py

@@ -1413,10 +1413,10 @@ def quantize_model(
             m_eager: torch.nn.Module, example_inputs: Tuple[torch.Tensor]
         ) -> Tuple[EdgeProgramManager, int, int]:
             # program capture
-            m = torch._export.capture_pre_autograd_graph(
+            m = torch.export.export_for_training(
                 m_eager,
                 example_inputs,
-            )
+            ).module()
 
             quantizer = XNNPACKQuantizer()
             quantization_config = get_symmetric_quantization_config()