Repository: huggingface/transformers
Pull request: Use public export API on torch 2.5 and future

# What does this PR do?

This is a follow-up of a TODO work. Due to issue https://github.com/pytorch/pytorch/issues/128394, we're using an internal API to export models if torch version is less than 2.5. Since the fix is included in 2.5, switch to use public export API for 2.5+.

Test
Tested the models exported using this integration points with torch 2.5, 2.6 and 2.7.0.dev

## Before submitting
- [x] Did you read the [contributor guideline](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md#create-a-pull-request),
      Pull Request section?

## Who can review?

Anyone in the community is free to review the PR once the tests have passed. Feel free to tag
members/contributors who may be interested in your PR.

@ArthurZucker @amyeroberts @qubvel


Changed files:
filename: src/transformers/integrations/executorch.py
patch: @@ -19,7 +19,7 @@
 
 if is_torch_available():
     from transformers import PreTrainedModel, StaticCache
-    from transformers.pytorch_utils import is_torch_greater_or_equal_than_2_3
+    from transformers.pytorch_utils import is_torch_greater_or_equal, is_torch_greater_or_equal_than_2_3
 
 
 class TorchExportableModuleWithStaticCache(torch.nn.Module):
@@ -193,7 +193,6 @@ def convert_and_export_with_cache(
     Returns:
         Exported program (`torch.export.ExportedProgram`): The exported program generated via `torch.export`.
     """
-
     if not is_torch_greater_or_equal_than_2_3:
         raise ImportError("torch >= 2.3 is required.")
 
@@ -208,15 +207,25 @@ def convert_and_export_with_cache(
             example_cache_position if example_cache_position is not None else torch.tensor([0], dtype=torch.long)
         )
 
-        # Due to issue https://github.com/pytorch/pytorch/issues/128394, we need to switch to use an internal
-        # export API and pre_dispatch=False. Switch to use the public API once the issue is included in 2.5 release.
-        exported_program = torch.export._trace._export(
-            TorchExportableModuleWithStaticCache(model),
-            args=(example_input_ids,),
-            kwargs={"cache_position": example_cache_position},
-            pre_dispatch=False,
-            strict=True,
-        )
+        if is_torch_greater_or_equal("2.5.0"):
+            exported_program = torch.export.export(
+                TorchExportableModuleWithStaticCache(model),
+                args=(example_input_ids,),
+                kwargs={"cache_position": example_cache_position},
+                strict=True,
+            )
+        else:
+            # We have to keep this path for BC.
+            #
+            # Due to issue https://github.com/pytorch/pytorch/issues/128394, we need to switch to use an internal
+            # export API and pre_dispatch=False. Switch to use the public API once the issue is included in 2.5 release.
+            exported_program = torch.export._trace._export(
+                TorchExportableModuleWithStaticCache(model),
+                args=(example_input_ids,),
+                kwargs={"cache_position": example_cache_position},
+                pre_dispatch=False,
+                strict=True,
+            )
         return exported_program