10000 Do not import transformers when import torch._dynamo by soulitzer · Pull Request #124634 · pytorch/pytorch · GitHub
[go: up one dir, main page]
Skip to content

Do not import transformers when import torch._dynamo #124634

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
wants to merge 3 commits into from
Closed

Do not import transformers when import torch._dynamo #124634

Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
e392c65
Do not import transformers when import torch._dynamo
soulitzer Apr 22, 2024
fe887a0
Update on "Do not import transformers when import torch._dynamo"
soulitzer Apr 22, 2024
27d40a8
Update on "Do not import transformers when import torch._dynamo"
soulitzer Apr 22, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
35 changes: 24 additions & 11 deletions torch/onnx/_internal/fx/patcher.py
View file
Open in desktop
708E
Original file line number Diff line number Diff line change
@@ -1,19 +1,24 @@
import copy
import functools
import io
from typing import List, Union

import torch


# TODO: Remove after https://github.com/huggingface/safetensors/pull/318
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]
from safetensors import torch as safetensors_torch # noqa: F401
@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] # noqa: F401

from safetensors import torch as safetensors_torch # noqa: F401

has_safetensors_and_transformers = True
except ImportError:
has_safetensors_and_transformers = False
return True
except ImportError:
return False


class ONNXTorchPatcher:
Expand Down Expand Up @@ -61,7 +66,9 @@ def torch_load_wrapper(f, *args, **kwargs):
# Wrapper or modified version of torch functions.
self.torch_load_wrapper = torch_load_wrapper

if has_safetensors_and_transformers:
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
Expand Down Expand Up @@ -109,7 +116,10 @@ def __enter__(self):
desired_wrapped_methods.append((torch.Tensor, "__getitem__"))
torch.fx._symbolic_trace._wrapped_methods_to_patch = desired_wrapped_methods

if has_safetensors_and_transformers:
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
Expand All @@ -120,7 +130,10 @@ def __exit__(self, exc_type, exc_value, traceback):
torch.fx._symbolic_trace._wrapped_methods_to_patch = (
self.torch_fx__symbolic_trace__wrapped_methods_to_patch
)
if has_safetensors_and_transformers:
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
Expand Down
0