Replace export_for_training with torch.export.export (#2724)

* Deprecate old TORCH_VERSION variables

**Summary:** This commit deprecates the following variables:

```
TORCH_VERSION_AT_LEAST_2_5
TORCH_VERSION_AT_LEAST_2_4
TORCH_VERSION_AT_LEAST_2_3
TORCH_VERSION_AT_LEAST_2_2
TORCH_VERSION_AFTER_2_5
TORCH_VERSION_AFTER_2_4
TORCH_VERSION_AFTER_2_3
TORCH_VERSION_AFTER_2_2
```

As of this commit, the latest released version of PyTorch is 2.8,
which means we can drop support for 2.5 and before since we only
support 3 of the latest releases.

The next commit will remove usages of all of these variables
from within torchao.

**Test Plan:**
```
python test/test_utils.py -k torch_version_deprecation
```

[ghstack-poisoned]

* Update on "Deprecate old TORCH_VERSION variables"


**Summary:** This commit deprecates the following variables:

```
TORCH_VERSION_AT_LEAST_2_5
TORCH_VERSION_AT_LEAST_2_4
TORCH_VERSION_AT_LEAST_2_3
TORCH_VERSION_AT_LEAST_2_2
TORCH_VERSION_AFTER_2_5
TORCH_VERSION_AFTER_2_4
TORCH_VERSION_AFTER_2_3
TORCH_VERSION_AFTER_2_2
```

As of this commit, the latest released version of PyTorch is 2.8,
which means we can drop support for 2.5 and before since we only
support 3 of the latest releases.

The next commit will remove usages of all of these variables
from within torchao.

**Test Plan:**
```
python test/test_utils.py -k torch_version_deprecation
```

[ghstack-poisoned]

* Drop support for PyTorch 2.5 and before

**Summary:** We gate on the PyTorch version throughout the repo.
Recently PyTorch 2.8 was released, so the oldest PyTorch version
we need to support is 2.6. After this commit, we assume the user
is running PyTorch 2.6+, and remove all references to the following
variables, which are deprecated.

```
TORCH_VERSION_AT_LEAST_2_6
TORCH_VERSION_AT_LEAST_2_5
TORCH_VERSION_AT_LEAST_2_4
TORCH_VERSION_AT_LEAST_2_3
TORCH_VERSION_AT_LEAST_2_2
TORCH_VERSION_AFTER_2_5
TORCH_VERSION_AFTER_2_4
TORCH_VERSION_AFTER_2_3
TORCH_VERSION_AFTER_2_2
```

**Test Plan:** CI

[ghstack-poisoned]

* Remove old `change_linear_weights_to_*` APIs

**Summary:** This commit removes these super old quantization
APIs that aren't even accessible by the user:

```
change_linear_weights_to_int8_dqtensors
change_linear_weights_to_int8_woqtensors
change_linear_weights_to_int4_woqtensors
```

**Test Plan:** CI

[ghstack-poisoned]

* Update base for Update on "Remove old `change_linear_weights_to_*` APIs"


**Summary:** This commit removes these super old quantization
APIs that aren't even accessible by the user:

```
change_linear_weights_to_int8_dqtensors
change_linear_weights_to_int8_woqtensors
change_linear_weights_to_int4_woqtensors
```

**Test Plan:** CI

[ghstack-poisoned]

* Update base for Update on "Remove old `change_linear_weights_to_*` APIs"


**Summary:** This commit removes these super old quantization
APIs that aren't even accessible by the user:

```
change_linear_weights_to_int8_dqtensors
change_linear_weights_to_int8_woqtensors
change_linear_weights_to_int4_woqtensors
```

**Test Plan:** CI

[ghstack-poisoned]

* Update base for Update on "Remove old `change_linear_weights_to_*` APIs"


**Summary:** This commit removes these super old quantization
APIs that aren't even accessible by the user:

```
change_linear_weights_to_int8_dqtensors
change_linear_weights_to_int8_woqtensors
change_linear_weights_to_int4_woqtensors
```

**Test Plan:** CI

[ghstack-poisoned]

* Update base for Update on "Remove old `change_linear_weights_to_*` APIs"


**Summary:** This commit removes these super old quantization
APIs that aren't even accessible by the user:

```
change_linear_weights_to_int8_dqtensors
change_linear_weights_to_int8_woqtensors
change_linear_weights_to_int4_woqtensors
```

**Test Plan:** CI

[ghstack-poisoned]

* Replace `export_for_training` with `torch.export.export`

**Summary:** Bypasses the following deprecation warning:

```
`torch.export.export_for_training` is deprecated and will be removed in PyTorch 2.10.
Please use `torch.export.export` instead, which is functionally equivalent.
```

Bonus: remove some references to `capture_pre_autograd_graph`,
which is even older.

**Test Plan:** CI

[ghstack-poisoned]

* Update base for Update on "Replace `export_for_training` with `torch.export.export`"


**Summary:** Bypasses the following deprecation warning:

```
`torch.export.export_for_training` is deprecated and will be removed in PyTorch 2.10.
Please use `torch.export.export` instead, which is functionally equivalent.
```

Bonus: remove some references to `capture_pre_autograd_graph`,
which is even older.

**Test Plan:** CI

[ghstack-poisoned]

* Update base for Update on "Replace `export_for_training` with `torch.export.export`"


**Summary:** Bypasses the following deprecation warning:

```
`torch.export.export_for_training` is deprecated and will be removed in PyTorch 2.10.
Please use `torch.export.export` instead, which is functionally equivalent.
```

Bonus: remove some references to `capture_pre_autograd_graph`,
which is even older.

**Test Plan:** CI

[ghstack-poisoned]

* Update base for Update on "Replace `export_for_training` with `torch.export.export`"


**Summary:** Bypasses the following deprecation warning:

```
`torch.export.export_for_training` is deprecated and will be removed in PyTorch 2.10.
Please use `torch.export.export` instead, which is functionally equivalent.
```

Bonus: remove some references to `capture_pre_autograd_graph`,
which is even older.

**Test Plan:** CI

[ghstack-poisoned]

* Update base for Update on "Replace `export_for_training` with `torch.export.export`"


**Summary:** Bypasses the following deprecation warning:

```
`torch.export.export_for_training` is deprecated and will be removed in PyTorch 2.10.
Please use `torch.export.export` instead, which is functionally equivalent.
```

Bonus: remove some references to `capture_pre_autograd_graph`,
which is even older.

**Test Plan:** CI

[ghstack-poisoned]

* Update base for Update on "Replace `export_for_training` with `torch.export.export`"


**Summary:** Bypasses the following deprecation warning:

```
`torch.export.export_for_training` is deprecated and will be removed in PyTorch 2.10.
Please use `torch.export.export` instead, which is functionally equivalent.
```

Bonus: remove some references to `capture_pre_autograd_graph`,
which is even older.

**Test Plan:** CI

[ghstack-poisoned]

Author: andrewor14

docs/source/tutorials_source/pt2e_quant_ptq.rst

@@ -362,7 +362,7 @@ Here is how you can use ``torch.export`` to export the model:
       {0: torch.export.Dim("dim")} if i == 0 else None
       for i in range(len(example_inputs))
     )
-    exported_model = torch.export.export_for_training(model_to_quantize, example_inputs, dynamic_shapes=dynamic_shapes).module()
+    exported_model = torch.export.export(model_to_quantize, example_inputs, dynamic_shapes=dynamic_shapes).module()
 
     # for pytorch 2.5 and before
     # dynamic_shape API may vary as well
@@ -501,7 +501,7 @@ Now we can compare the size and model accuracy with baseline model.
     # Quantized model size and accuracy
     print("Size of model after quantization")
     # export again to remove unused weights
-    quantized_model = torch.export.export_for_training(quantized_model, example_inputs).module()
+    quantized_model = torch.export.export(quantized_model, example_inputs).module()
     print_size_of_model(quantized_model)
 
     top1, top5 = evaluate(quantized_model, criterion, data_loader_test)

docs/source/tutorials_source/pt2e_quant_qat.rst

@@ -13,7 +13,6 @@ to the post training quantization (PTQ) flow for the most part:
 .. code:: python
 
   import torch
-  from torch._export import capture_pre_autograd_graph
   from torchao.quantization.pt2e.quantize_pt2e import (
     prepare_qat_pt2e,
     convert_pt2e,
@@ -434,7 +433,6 @@ prepared. For example:
 
 .. code:: python
 
-    from torch._export import capture_pre_autograd_graph
     from executorch.backends.xnnpack.quantizer.xnnpack_quantizer import (
         get_symmetric_quantization_config,
         XNNPACKQuantizer,
@@ -443,7 +441,7 @@ prepared. For example:
 
     example_inputs = (torch.rand(2, 3, 224, 224),)
     float_model = resnet18(pretrained=False)
-    exported_model = capture_pre_autograd_graph(float_model, example_inputs)
+    exported_model = torch.export.export(float_model, example_inputs).module()
     quantizer = XNNPACKQuantizer()
     quantizer.set_global(get_symmetric_quantization_config(is_qat=True))
     prepared_model = prepare_qat_pt2e(exported_model, quantizer)

docs/source/tutorials_source/pt2e_quant_x86_inductor.rst

@@ -105,7 +105,7 @@ We will start by performing the necessary imports, capturing the FX Graph from t
         exported_model = export(
             model,
             example_inputs
-        )
+        ).module()
 
 
 Next, we will have the FX Module to be quantized.
@@ -243,12 +243,10 @@ The PyTorch 2 Export QAT flow is largely similar to the PTQ flow:
 .. code:: python
 
   import torch
-  from torch._export import capture_pre_autograd_graph
   from torchao.quantization.pt2e.quantize_pt2e import (
     prepare_qat_pt2e,
     convert_pt2e,
   )
-  from torch.export import export
   import torchao.quantization.pt2e.quantizer.x86_inductor_quantizer as xiq
   from torchao.quantization.pt2e.quantizer.x86_inductor_quantizer import X86InductorQuantizer
 
@@ -264,9 +262,7 @@ The PyTorch 2 Export QAT flow is largely similar to the PTQ flow:
   m = M()
 
   # Step 1. program capture
-  # NOTE: this API will be updated to torch.export API in the future, but the captured
-  # result shoud mostly stay the same
-  exported_model = export(m, example_inputs)
+  exported_model = torch.export.export(m, example_inputs).module()
   # we get a model with aten ops
 
   # Step 2. quantization-aware training

examples/sam2_amg_server/compile_export_utils.py

@@ -118,10 +118,7 @@ def aot_compile(
             "max_autotune": True,
             "triton.cudagraphs": True,
         }
-
-    from torch.export import export_for_training
-
-    exported = export_for_training(fn, sample_args, sample_kwargs, strict=True)
+    exported = torch.export.export(fn, sample_args, sample_kwargs, strict=True)
     exported.run_decompositions()
     output_path = torch._inductor.aoti_compile_and_package(
         exported,

examples/sam2_vos_example/compile_export_utils.py

@@ -81,10 +81,7 @@ def aot_compile(
             "max_autotune": True,
             "triton.cudagraphs": True,
         }
-
-    from torch.export import export_for_training
-
-    exported = export_for_training(fn, sample_args, sample_kwargs, strict=True)
+    exported = torch.export.export(fn, sample_args, sample_kwargs, strict=True)
     exported.run_decompositions()
     output_path = torch._inductor.aoti_compile_and_package(
         exported,

test/dtypes/test_uint4.py

@@ -242,10 +242,7 @@ def forward(self, x):
 
         # program capture
         m = copy.deepcopy(m_eager)
-        m = torch.export.texport_for_training(
-            m,
-            example_inputs,
-        ).module()
+        m = torch.export.export(m, example_inputs).module()
 
         m = prepare_pt2e(m, quantizer)
         # Calibrate

test/integration/test_integration.py

@@ -1953,9 +1953,7 @@ def forward(self, x):
         # TODO: export changes numerics right now, this is because of functionalization according to Zhengxu
         # we can re-enable this after non-functional IR is enabled in export
         # model = torch.export.export(model, example_inputs).module()
-        model = torch.export.export_for_training(
-            model, example_inputs, strict=True
-        ).module()
+        model = torch.export.export(model, example_inputs, strict=True).module()
         after_export = model(x)
         self.assertTrue(torch.equal(after_export, ref))
         if api is _int8da_int4w_api:

test/prototype/inductor/test_int8_sdpa_fusion.py

@@ -157,8 +157,6 @@ def _check_common(
     )
     @config.patch({"freezing": True})
     def _test_sdpa_int8_rewriter(self):
-        from torch.export import export_for_training
-
         import torchao.quantization.pt2e.quantizer.x86_inductor_quantizer as xiq
         from torchao.quantization.pt2e.quantize_pt2e import convert_pt2e, prepare_pt2e
         from torchao.quantization.pt2e.quantizer.x86_inductor_quantizer import (
@@ -199,11 +197,7 @@ def _test_sdpa_int8_rewriter(self):
                 quantizer.set_function_type_qconfig(
                     torch.matmul, quantizer.get_global_quantization_config()
                 )
-                export_model = export_for_training(
-                    mod,
-                    inputs,
-                    strict=True,
-                ).module()
+                export_model = torch.export.export(mod, inputs, strict=True).module()
                 prepare_model = prepare_pt2e(export_model, quantizer)
                 prepare_model(*inputs)
                 convert_model = convert_pt2e(prepare_model)

test/quantization/pt2e/test_arm_inductor_quantizer.py

@@ -14,7 +14,6 @@
 
 import torch
 import torch.nn as nn
-from torch.export import export_for_training
 from torch.testing._internal.common_quantization import (
     NodeSpec as ns,
 )
@@ -315,10 +314,7 @@ def _test_quantizer(
 
         # program capture
         m = copy.deepcopy(m_eager)
-        m = export_for_training(
-            m,
-            example_inputs,
-        ).module()
+        m = torch.export.export(m, example_inputs).module()
 
         # QAT Model failed to deepcopy
         export_model = m if is_qat else copy.deepcopy(m)
@@ -576,7 +572,7 @@ def _test_linear_unary_helper(
         Test pattern of linear with unary post ops (e.g. relu) with ArmInductorQuantizer.
         """
         use_bias_list = [True, False]
-        # TODO test for inplace add after refactoring of export_for_training
+        # TODO test for inplace add after refactoring of export
         inplace_list = [False]
         if post_op_algo_list is None:
             post_op_algo_list = [None]
@@ -716,7 +712,7 @@ def _test_linear_binary_helper(self, is_qat=False, is_dynamic=False):
         Currently, only add as binary post op is supported.
         """
         linear_pos_list = [NodePosType.left, NodePosType.right, NodePosType.both]
-        # TODO test for inplace add after refactoring of export_for_training
+        # TODO test for inplace add after refactoring of export
         inplace_add_list = [False]
         example_inputs = (torch.randn(2, 16),)
         quantizer = ArmInductorQuantizer().set_global(
@@ -1078,7 +1074,7 @@ def forward(self, x):
         )
         example_inputs = (torch.randn(2, 2),)
         m = M().eval()
-        m = export_for_training(m, example_inputs).module()
+        m = torch.export.export(m, example_inputs).module()
         m = prepare_pt2e(m, quantizer)
         # Use a linear count instead of names because the names might change, but
         # the order should be the same.

test/quantization/pt2e/test_duplicate_dq.py

@@ -110,7 +110,7 @@ def _test_duplicate_dq(
 
         # program capture
         m = copy.deepcopy(m_eager)
-        m = export_for_training(m, example_inputs, strict=True).module()
+        m = torch.export.export(m, example_inputs, strict=True).module()
 
         m = prepare_pt2e(m, quantizer)
         # Calibrate