Fix NameError in accelerate.py and clarify error for PyTorch < 2.6 (#50089)

Co-authored-by: Meghana Zagabathuni <mz4156@nyu.edu>

Author: megacodes24

src/transformers/integrations/accelerate.py

@@ -16,6 +16,8 @@
 and simplicity/ease of use.
 """
 
+from __future__ import annotations
+
 import copy
 import functools
 import inspect
@@ -93,7 +95,7 @@ def get_module_size_with_ties(
     return module_size_with_ties, tied_module_names, tied_modules
 
 
-def check_and_set_device_map(device_map: "torch.device | int | str | dict | None") -> dict | str | None:
+def check_and_set_device_map(device_map: torch.device | int | str | dict | None) -> dict | str | None:
     from ..modeling_utils import get_torch_context_manager_or_global_device
 
     # Potentially detect context manager or global device, and use it (only if no device_map was provided)
@@ -142,8 +144,8 @@ def check_and_set_device_map(device_map: "torch.device | int | str | dict | None
 
 
 def compute_module_sizes(
-    model: "PreTrainedModel",
-    hf_quantizer: "HfQuantizer | None" = None,
+    model: PreTrainedModel,
+    hf_quantizer: HfQuantizer | None = None,
     buffers_only: bool = False,
     only_modules: bool = True,
 ) -> tuple[dict[str, int], dict[str, int]]:
@@ -191,7 +193,7 @@ def all_tensors():
     return all_module_sizes, leaves_module_sizes
 
 
-def compute_module_total_buffer_size(model: nn.Module, hf_quantizer: "HfQuantizer | None" = None):
+def compute_module_total_buffer_size(model: nn.Module, hf_quantizer: HfQuantizer | None = None):
     """
     Compute the total size of buffers in each submodule of a given model.
     """
@@ -238,10 +240,10 @@ def get_max_memory(max_memory: dict[int | str, int | str] | None = None):
 
 
 def get_balanced_memory(
-    model: "PreTrainedModel",
+    model: PreTrainedModel,
     max_memory: dict[int | str, int | str] | None = None,
     no_split_module_classes: set[str] | None = None,
-    hf_quantizer: "HfQuantizer | None" = None,
+    hf_quantizer: HfQuantizer | None = None,
     low_zero: bool = False,
 ):
     """
@@ -336,10 +338,10 @@ def get_balanced_memory(
 
 
 def _get_device_map(
-    model: "PreTrainedModel",
+    model: PreTrainedModel,
     device_map: dict | str | None,
     max_memory: dict | None,
-    hf_quantizer: "HfQuantizer | None",
+    hf_quantizer: HfQuantizer | None,
 ) -> dict:
     """Compute the final `device_map` to use if we passed a value in ['auto', 'balanced', 'balanced_low_0', 'sequential'].
     Otherwise, we check for any device inconsistencies in the device_map.
@@ -457,7 +459,7 @@ def get_device(device_map: dict | None, param_name: str, valid_torch_device: boo
 
 
 def accelerate_disk_offload(
-    model: "PreTrainedModel",
+    model: PreTrainedModel,
     disk_offload_folder: str | None,
     checkpoint_files: list[str] | None,
     device_map: dict,
@@ -543,7 +545,7 @@ def offload_weight(weight: torch.Tensor, weight_name: str, offload_folder: str |
     return offload_index
 
 
-def load_offloaded_parameter(model: "PreTrainedModel", param_name: str) -> torch.Tensor:
+def load_offloaded_parameter(model: PreTrainedModel, param_name: str) -> torch.Tensor:
     """Load `param_name` from disk, if it was offloaded due to the device_map, and thus lives as a meta parameter
     inside `model`.
     This is needed when resaving a model, when some parameters were offloaded (we need to load them from disk, to
@@ -574,7 +576,7 @@ def _init_infer_auto_device_map(
     max_memory: dict[int | str, int | str] | None = None,
     no_split_module_classes: set[str] | None = None,
     tied_parameters: list[list[str]] | None = None,
-    hf_quantizer: "HfQuantizer | None" = None,
+    hf_quantizer: HfQuantizer | None = None,
 ) -> tuple[
     list[int | str],
     dict[int | str, int | str],
@@ -647,7 +649,7 @@ def infer_auto_device_map(
     clean_result: bool = True,
     offload_buffers: bool = False,
     tied_parameters: list[list[str]] | None = None,
-    hf_quantizer: "HfQuantizer | None" = None,
+    hf_quantizer: HfQuantizer | None = None,
     no_placement_params: set[str] | None = None,
 ):
     """

src/transformers/utils/import_utils.py

@@ -2090,8 +2090,9 @@ def is_training_run_on_sagemaker() -> bool:
 
 # docstyle-ignore
 PYTORCH_IMPORT_ERROR = """
-{0} requires the PyTorch library but it was not found in your environment. Check out the instructions on the
-installation page: https://pytorch.org/get-started/locally/ and follow the ones that match your environment.
+{0} requires PyTorch >= 2.6, but PyTorch was not found or an older version is installed.
+Please install or upgrade PyTorch by following the instructions on the installation page:
+https://pytorch.org/get-started/locally/
 Please note that you may need to restart your runtime after installation.
 """