Update on "Do not import transformers when import torch._dynamo"

soulitzer · soulitzer · commit 27d40a822943 · 2024-04-22T14:57:40.000-04:00
Fixes #123954 [ghstack-poisoned]
diff --git a/torch/onnx/_internal/fx/patcher.py b/torch/onnx/_internal/fx/patcher.py
@@ -1,17 +1,18 @@
 import copy
+import functools
 import io
 from typing import List, Union
-import functools
 
 import torch
 
+
 # TODO: Remove after https://github.com/huggingface/safetensors/pull/318
-@functools.cache
+@functools.lru_cache(None)
 def has_safetensors_and_transformers():
     try:
         # safetensors is not an exporter requirement, but needed for some huggingface models
         import safetensors  # type: ignore[import]  # noqa: F401
-        import transformers  # type: ignore[import]
+        import transformers  # type: ignore[import]  # noqa: F401
 
         from safetensors import torch as safetensors_torch  # noqa: F401
 
@@ -66,6 +67,8 @@ def torch_load_wrapper(f, *args, **kwargs):
         self.torch_load_wrapper = torch_load_wrapper
 
         if has_safetensors_and_transformers():
+            import safetensors
+            import transformers
 
             def safetensors_load_file_wrapper(filename, device="cpu"):
                 # Record path for later serialization into ONNX proto
@@ -114,6 +117,9 @@ def __enter__(self):
         torch.fx._symbolic_trace._wrapped_methods_to_patch = desired_wrapped_methods
 
         if has_safetensors_and_transformers():
+            import safetensors
+            import transformers
+
             safetensors.torch.load_file = self.safetensors_torch_load_file_wrapper
             transformers.modeling_utils.safe_load_file = (
                 self.safetensors_torch_load_file_wrapper
@@ -125,6 +131,9 @@ def __exit__(self, exc_type, exc_value, traceback):
             self.torch_fx__symbolic_trace__wrapped_methods_to_patch
         )
         if has_safetensors_and_transformers():
+            import safetensors
+            import transformers
+
             safetensors.torch.load_file = self.safetensors_torch_load_file
             transformers.modeling_utils.safe_load_file = (
                 self.transformers_modeling_utils_safe_load_file
