Second phase of torch 2.6 cleanups

@@ -17,7 +17,7 @@ import os from typing import TYPE_CHECKING, Any -from ..utils import is_torch_available, is_torch_distributed_available, is_torch_greater_or_equal, logging, strtobool+from ..utils import is_torch_available, is_torch_distributed_available, logging, strtobool from ..utils.quantization_config import QuantizationMethod from .tensor_parallel import _get_parameter_plan, replace_layer_number_by_wildcard from .utils import _is_torch_distributed_initialized@@ -32,7 +32,7 @@ if is_torch_available():     import torch -if is_torch_distributed_available() and is_torch_greater_or_equal("2.6"):+if is_torch_distributed_available():     from torch.distributed._composable.fsdp import fully_shard     from torch.distributed.fsdp import CPUOffloadPolicy, MixedPrecisionPolicy 

@@ -203,9 +203,6 @@ def initialize_tensor_parallelism(     if tp_plan is not None and device_map is not None:         raise ValueError("`tp_plan` and `device_map` are mutually exclusive. Choose either one for parallelization.")     if device_mesh is None:-        if not is_torch_greater_or_equal("2.5"):-            raise OSError("Tensor parallel is only supported for `torch>=2.5`.")-         # Detect the accelerator on the machine. If no accelerator is available, it returns CPU.         device_type = torch._C._get_accelerator().type         if device_type == "mps":

@@ -27,10 +27,6 @@ from ..configuration_utils import get_head_shapes from ..generation.configuration_utils import GenerationConfig from ..modeling_utils import PreTrainedModel-from ..pytorch_utils import (-    is_torch_greater_or_equal,-    is_torch_greater_or_equal_than_2_6,-)   class TorchExportableModuleForVLM:@@ -788,32 +784,13 @@ def convert_and_export_with_cache(             else torch.tensor([0], dtype=torch.long, device=model.device)         ) -        if is_torch_greater_or_equal("2.6.0"):-            exported_program = torch.export.export(-                TorchExportableModuleWithStaticCache(model),-                args=(),-                kwargs={"input_ids": example_input_ids, "cache_position": example_cache_position},-                dynamic_shapes=dynamic_shapes,-                strict=strict if strict is not None else True,-            )-        else:-            if dynamic_shapes is not None:-                logging.warning(-                    "Dynamic shapes spec will be ignored by convert_and_export_with_cache for torch < 2.6.0."-                )-            if strict is not None:-                logging.warning("The strict flag will be ignored by convert_and_export_with_cache for torch < 2.6.0.")-            # 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=(),-                kwargs={"input_ids": example_input_ids, "cache_position": example_cache_position},-                pre_dispatch=False,-                strict=True,-            )+        exported_program = torch.export.export(+            TorchExportableModuleWithStaticCache(model),+            args=(),+            kwargs={"input_ids": example_input_ids, "cache_position": example_cache_position},+            dynamic_shapes=dynamic_shapes,+            strict=strict if strict is not None else True,+        )         return exported_program  @@ -1105,9 +1082,6 @@ def _get_cache_dict(cache: DynamicCache):             "This pytree flattening function should be applied to DynamicCache containing only `DynamicLayer` and `DynamicSlidingWindowLayer`"         ) -    if not is_torch_greater_or_equal_than_2_6:-        logging.warning("DynamicCache + torch.export is tested on torch 2.6.0+ and may not work on earlier versions.")-     return {         "key_cache": [layer.keys for layer in cache.layers if layer.keys is not None],         "value_cache": [layer.values for layer in cache.layers if layer.values is not None],

@@ -35,7 +35,6 @@ from ..utils.import_utils import (     get_torch_version,     is_torch_greater_or_equal,-    is_torch_less_or_equal,     is_torchdynamo_compiling, ) @@ -78,12 +77,10 @@ def __init__(self, training):         """         if not self._is_flex_compiled or training != self.training:             self.training = training-            if is_torch_less_or_equal("2.5.1"):-                self._compiled_flex_attention = torch.compile(flex_attention, dynamic=False)             # In PyTorch 2.6.0, there's a known issue with flex attention compilation which may             # cause errors. The suggested fix is to compile with "max-autotune-no-cudagraphs"             # see https://github.com/pytorch/pytorch/issues/146260 for training-            elif version.parse(get_torch_version()).base_version == "2.6.0" and training:+            if version.parse(get_torch_version()).base_version == "2.6.0" and training:                 self._compiled_flex_attention = torch.compile(                     flex_attention, dynamic=False, mode="max-autotune-no-cudagraphs"                 )@@ -242,8 +239,7 @@ def mask_mod(batch_idx, head_idx, q_idx, kv_idx):         Q_LEN=query_length,         KV_LEN=key_length,         device=device,-        # compiling the mask is not BC with older torch-        _compile=not is_torch_less_or_equal("2.5.1"),+        _compile=True,     )  

@@ -16,32 +16,27 @@  import torch import torch.nn.functional as F+from torch._dynamo._trace_wrapped_higher_order_op import TransformGetItemToIndex  from .cache_utils import Cache from .configuration_utils import PreTrainedConfig from .utils import is_torch_xpu_available, logging from .utils.generic import GeneralInterface, is_flash_attention_requested from .utils.import_utils import (     is_torch_flex_attn_available,-    is_torch_greater_or_equal,     is_torchdynamo_exporting,     is_tracing, )   if is_torch_flex_attn_available():-    from torch.nn.attention.flex_attention import _DEFAULT_SPARSE_BLOCK_SIZE as flex_default_block_size     from torch.nn.attention.flex_attention import BlockMask, create_block_mask else:     # Register a fake type to avoid crashing for annotations and `isinstance` checks     BlockMask = torch.Tensor -_is_torch_greater_or_equal_than_2_6 = is_torch_greater_or_equal("2.6", accept_dev=True) _is_torch_xpu_available = is_torch_xpu_available() -if _is_torch_greater_or_equal_than_2_6:-    from torch._dynamo._trace_wrapped_higher_order_op import TransformGetItemToIndex-  logger = logging.get_logger(__name__) @@ -394,7 +389,6 @@ def sdpa_mask(     """     Create a 4D boolean mask of shape `(batch_size, 1, query_length, kv_length)` where a value of True indicates that     the element should take part in the attention computation, and False that it should not.-    This function can only be used with torch>=2.5, as the context manager is otherwise not available.      Args:         batch_size (`int`):@@ -523,21 +517,14 @@ def sdpa_mask(         # Expand the mask to match batch size and query length if they weren't used in the mask function         attention_mask = attention_mask.expand(batch_size, -1, q_length, kv_length) -    # Option 2: Vmap mask creation (torch>=2.6 and custom patterns)-    elif _is_torch_greater_or_equal_than_2_6:+    # Option 2: Vmap mask creation (custom patterns)+    else:         # This creates the 4D mask easily. Note that we need this context manager as vmap cannot handle slicing a tensor from         # scalar tensor (it internally calls `.item()` which vmap does not allow, but this context works around it         # We don't need to add an offset to the mask_function either, as we vmap directly the correct indices for k and kv indices         with TransformGetItemToIndex():             attention_mask = _vmap_expansion_sdpa(mask_function)(batch_arange, head_arange, q_arange, kv_arange) -    # Option 3: Error out since it indicates that the user did something custom, which they shouldn't have (torch<2.6)-    else:-        raise ValueError(-            "The vmap functionality for mask creation is only supported from torch>=2.6. "-            "Please update your torch version or use `use_vmap=False` with index-based masks."-        )-     return attention_mask  @@ -688,13 +675,6 @@ def flex_attention_mask(     """     # Potentially add the padding 2D mask     if attention_mask is not None and not fast_all(attention_mask):-        # Older torch (2.5.x) cannot handle sequences not in multiples of 128 (default block size)-        # Hence we pad to multiples of this as a minimum to ensure this-        pad_len = ((attention_mask.shape[1] // flex_default_block_size) + 1) * flex_default_block_size-        pad_len = pad_len - attention_mask.shape[1]-        if not _is_torch_greater_or_equal_than_2_6 and pad_len > 0:-            attention_mask = torch.nn.functional.pad(attention_mask, value=0, pad=(0, pad_len))-         padding_mask = prepare_padding_mask(attention_mask, kv_length, kv_offset)         mask_function = and_masks(mask_function, padding_mask_function(padding_mask)) @@ -709,7 +689,7 @@ def flex_attention_mask(         Q_LEN=q_length,         KV_LEN=kv_length,         device=device,-        _compile=_is_torch_greater_or_equal_than_2_6,+        _compile=True,     )     return block_mask @@ -961,14 +941,10 @@ def create_causal_mask(     # Note that it is very important to apply this before any other deviations of the mask (such as packed sequence mask,     # padding mask, etc) as the resulting mask may otherwise not be correct!     if or_mask_function is not None:-        if not _is_torch_greater_or_equal_than_2_6:-            raise ValueError("Using `or_mask_function` or `and_mask_function` arguments require torch>=2.6")         mask_factory_function = or_masks(mask_factory_function, or_mask_function)         allow_is_causal_skip = False         use_vmap = True     if and_mask_function is not None:-        if not _is_torch_greater_or_equal_than_2_6:-            raise ValueError("Using `or_mask_function` or `and_mask_function` arguments require torch>=2.6")         mask_factory_function = and_masks(mask_factory_function, and_mask_function)         allow_is_causal_skip = False         use_vmap = True@@ -1071,14 +1047,10 @@ def create_bidirectional_mask(     # Note that it is very important to apply this before any other deviations of the mask (such as packed sequence mask,     # padding mask, etc) as the resulting mask may otherwise not be correct!     if or_mask_function is not None:-        if not _is_torch_greater_or_equal_than_2_6:-            raise ValueError("Using `or_mask_function` or `and_mask_function` arguments require torch>=2.6")         mask_factory_function = or_masks(mask_factory_function, or_mask_function)         allow_is_bidirectional_skip = False         use_vmap = True     if and_mask_function is not None:-        if not _is_torch_greater_or_equal_than_2_6:-            raise ValueError("Using `or_mask_function` or `and_mask_function` arguments require torch>=2.6")         mask_factory_function = and_masks(mask_factory_function, and_mask_function)         allow_is_bidirectional_skip = False         use_vmap = True@@ -1201,14 +1173,10 @@ def create_sliding_window_causal_mask(     # Note that it is very important to apply this before any other deviations of the mask (such as packed sequence mask,     # padding mask, etc) as the resulting mask may otherwise not be correct!     if or_mask_function is not None:-        if not _is_torch_greater_or_equal_than_2_6:-            raise ValueError("Using `or_mask_function` or `and_mask_function` arguments require torch>=2.6")         mask_factory_function = or_masks(mask_factory_function, or_mask_function)         allow_is_causal_skip = False         use_vmap = True     if and_mask_function is not None:-        if not _is_torch_greater_or_equal_than_2_6:-            raise ValueError("Using `or_mask_function` or `and_mask_function` arguments require torch>=2.6")         mask_factory_function = and_masks(mask_factory_function, and_mask_function)         allow_is_causal_skip = False         use_vmap = True@@ -1310,14 +1278,10 @@ def create_bidirectional_sliding_window_mask(     use_vmap = False      if or_mask_function is not None:-        if not _is_torch_greater_or_equal_than_2_6:-            raise ValueError("Using `or_mask_function` or `and_mask_function` arguments require torch>=2.6")         mask_factory_function = or_masks(mask_factory_function, or_mask_function)         allow_is_bidirectional_skip = False         use_vmap = True     if and_mask_function is not None:-        if not _is_torch_greater_or_equal_than_2_6:-            raise ValueError("Using `or_mask_function` or `and_mask_function` arguments require torch>=2.6")         mask_factory_function = and_masks(mask_factory_function, and_mask_function)         allow_is_bidirectional_skip = False         use_vmap = True@@ -1435,14 +1399,10 @@ def create_chunked_causal_mask(     # Note that it is very important to apply this before any other deviations of the mask (such as packed sequence mask,     # padding mask, etc) as the resulting mask may otherwise not be correct!     if or_mask_function is not None:-        if not _is_torch_greater_or_equal_than_2_6:-            raise ValueError("Using `or_mask_function` or `and_mask_function` arguments require torch>=2.6")         mask_factory_function = or_masks(mask_factory_function, or_mask_function)         allow_is_causal_skip = False         use_vmap = True     if and_mask_function is not None:-        if not _is_torch_greater_or_equal_than_2_6:-            raise ValueError("Using `or_mask_function` or `and_mask_function` arguments require torch>=2.6")         mask_factory_function = and_masks(mask_factory_function, and_mask_function)         allow_is_causal_skip = False         use_vmap = True

@@ -1704,10 +1704,7 @@ def _flex_attn_can_dispatch(self, is_init_check: bool = False) -> bool:                 ' Example: `model = AutoModel.from_pretrained("openai/whisper-tiny", attn_implementation="eager")`'             )         if not is_torch_flex_attn_available():-            raise ImportError(-                "PyTorch Flex Attention requirements in Transformers are not met. Please install torch>=2.5.0 and"-                " run on a device other than TPU."-            )+            raise ImportError("PyTorch Flex Attention is not supported on TPU.")          # If no error raise by this point, we can return `True`         return True

@@ -28,7 +28,7 @@ from ..utils import is_torch_available, is_torchao_available, logging  -MIN_TORCH_VERSION = "2.5.0"+MIN_TORCH_VERSION = "2.6.0"   if is_torch_available():

@@ -81,7 +81,6 @@ ) from .optimization import GreedyLR, get_scheduler from .processing_utils import ProcessorMixin-from .pytorch_utils import is_torch_greater_or_equal_than_2_6 from .tokenization_utils_base import BatchEncoding, PreTrainedTokenizerBase from .trainer_callback import (     CallbackHandler,@@ -1027,10 +1026,8 @@ def _get_dataloader(             "persistent_workers": self.args.dataloader_persistent_workers,             "multiprocessing_context": self.args.dataloader_multiprocessing_context,             "prefetch_factor": self.args.dataloader_prefetch_factor,+            "in_order": self.args.dataloader_in_order,         }-        # `in_order` was added in torch 2.6; on older versions the loader always behaves as `in_order=True`.-        if is_torch_greater_or_equal_than_2_6:-            dataloader_params["in_order"] = self.args.dataloader_in_order          sampler = None         if not isinstance(dataset, torch.utils.data.IterableDataset):

@@ -455,12 +455,6 @@ def _get_torchao_optimizer(ctx: OptimizerContext) -> tuple[Any, dict[str, Any]]:             "Install it with `pip install torchao` or follow the instructions here: "             "https://github.com/pytorch/ao"         )-    if version.parse(importlib.metadata.version("torch")) <= version.parse("2.4"):-        raise ImportError(-            "You need to have `torch>2.4` in order to use torch 4-bit optimizers. "-            "Install it with `pip install --upgrade torch` it is available on pipy. "-            "Otherwise, you need to install torch nightly."-        )      if version.parse(importlib.metadata.version("torchao")) >= version.parse("0.11.0"):         from torchao.optim import AdamW4bit, AdamW8bit

@@ -15,7 +15,6 @@ Torch utilities for the Trainer class. """ -import contextlib import copy import datetime import io@@ -1647,9 +1646,6 @@ def safe_globals():     See: https://pytorch.org/docs/stable/notes/serialization.html#torch.serialization.add_safe_globals     See: https://github.com/huggingface/accelerate/pull/3036     """-    if version.parse(torch.__version__).release < version.parse("2.6").release:-        return contextlib.nullcontext()-     np_core = np._core if version.parse(np.__version__) >= version.parse("2.0.0") else np.core     allowlist = [np_core.multiarray._reconstruct, np.ndarray, np.dtype]     # numpy >1.25 defines numpy.dtypes.UInt32DType, but below works for

@@ -515,7 +515,7 @@ def sample_indices_fn(metadata, **kwargs):   # `torchvision.io.read_video` removed in `torchvision==0.26` (https://github.com/pytorch/vision/releases#release-v0.26.0),-# which ships with `torch==2.11`. As `transformers` supports `torch>=2.5`, we cannot rely on the pinned version and have to check+# which ships with `torch==2.11`. As `transformers` supports `torch>=2.6`, we cannot rely on the pinned version and have to check # it at runtime instead. Once the minimum supported `torch` version is bumped to 2.11, the whole `torchvision` video # decoding backend can be deleted. TORCHVISION_VIDEO_DECODING_REMOVED_VERSION = "0.26.0"

@@ -50,7 +50,6 @@     require_torch,     require_torch_accelerator,     require_torch_gpu,-    require_torch_greater_or_equal,     require_torch_multi_accelerator,     set_config_for_less_flaky_test,     set_model_for_less_flaky_test,@@ -1722,7 +1721,6 @@ def test_generate_with_quant_cache(self):      @pytest.mark.generate     @pytest.mark.torch_compile_test-    @require_torch_greater_or_equal("2.6")  # Uses torch.compiler.set_stance     def test_generate_compile_model_forward_fullgraph(self):         """         Tests that `.generate` is compatible with torch.compile, keeping the same results. Also confirms that

@@ -16,7 +16,6 @@ import unittest  import pytest-from packaging import version from parameterized import parameterized from pytest import mark @@ -183,9 +182,6 @@ def test_model_flash_attn(self):      @pytest.mark.torch_export_test     def test_export_static_cache(self):-        if version.parse(torch.__version__) < version.parse("2.5.0"):-            self.skipTest(reason="This test requires torch >= 2.5 to run.")-         from transformers.integrations.executorch import (             TorchExportableModuleWithStaticCache,             convert_and_export_with_cache,

@@ -16,7 +16,6 @@ import unittest  import pytest-from packaging import version from parameterized import parameterized from pytest import mark @@ -205,9 +204,6 @@ def test_model_9b_flash_attn(self):     @pytest.mark.torch_export_test     @slow     def test_export_static_cache(self):-        if version.parse(torch.__version__) < version.parse("2.5.0"):-            self.skipTest(reason="This test requires torch >= 2.5 to run.")-         from transformers.integrations.executorch import (             TorchExportableModuleWithStaticCache,         )@@ -280,10 +276,6 @@ def test_export_static_cache(self):     @pytest.mark.torch_export_test     def test_export_hybrid_cache(self):         from transformers.integrations.executorch import TorchExportableModuleForDecoderOnlyLM-        from transformers.pytorch_utils import is_torch_greater_or_equal--        if not is_torch_greater_or_equal("2.6.0"):-            self.skipTest(reason="This test requires torch >= 2.6 to run.")          model_id = "google/gemma-2-2b"         model = AutoModelForCausalLM.from_pretrained(model_id)

@@ -65,7 +65,6 @@     )     from transformers.cache_utils import StaticCache     from transformers.models.gemma3.modeling_gemma3 import create_masks_for_vision_model-    from transformers.pytorch_utils import is_torch_greater_or_equal   class Gemma3TextModelTester(CausalLMModelTester):@@ -836,9 +835,6 @@ def test_generation_beyond_sliding_window(self, attn_implementation: str):      @pytest.mark.torch_export_test     def test_export_text_only(self):-        if not is_torch_greater_or_equal("2.6.0"):-            self.skipTest(reason="This test requires torch >= 2.6 to run.")-         from transformers.integrations.executorch import TorchExportableModuleForDecoderOnlyLM          model_id = "google/gemma-3-1b-it"

@@ -154,10 +154,6 @@ def test_model_8b_long_prompt(self):     @pytest.mark.torch_export_test     def test_export_text_with_hybrid_cache(self):         # TODO: Exportability is not working-        from transformers.testing_utils import is_torch_greater_or_equal--        if not is_torch_greater_or_equal("2.6.0"):-            self.skipTest(reason="This test requires torch >= 2.6 to run.")          from transformers.integrations.executorch import TorchExportableModuleForDecoderOnlyLM 

@@ -16,7 +16,6 @@ import unittest  import pytest-from packaging import version from parameterized import parameterized  from transformers import (@@ -113,9 +112,6 @@ def test_model_pipeline_bf16(self):     @pytest.mark.torch_export_test     @slow     def test_export_static_cache(self):-        if version.parse(torch.__version__) < version.parse("2.5.0"):-            self.skipTest(reason="This test requires torch >= 2.5 to run.")-         from transformers.integrations.executorch import (             TorchExportableModuleWithStaticCache,         )

@@ -28,7 +28,7 @@ from transformers.distributed.configuration_utils import DistributedConfig from transformers.distributed.pipeline_parallel import PipelineIdentityLayer, PipelineStage, apply_pipeline_parallelism from transformers.modeling_utils import LoadStateDictConfig-from transformers.testing_utils import TestCasePlus, require_torch_greater_or_equal+from transformers.testing_utils import TestCasePlus from transformers.utils.loading_report import log_state_dict_report  @@ -300,7 +300,6 @@ def test_init_from_device_mesh(self, mock_get_backend):         mock_get_backend.assert_called_once_with(stage.pp_group)  -@require_torch_greater_or_equal("2.5") class TestPipelineParallelLoadReport(TestCasePlus):     @parameterized.expand([(pp_size, tie_word_embeddings) for pp_size in [2] for tie_word_embeddings in [True, False]])     def test_pp_loading_report_table(self, pp_size, tie_word_embeddings):@@ -314,7 +313,6 @@ def test_pp_loading_report_table(self, pp_size, tie_word_embeddings):         )  -@require_torch_greater_or_equal("2.5") class TestPipelineParallelSplit(TestCasePlus):     @parameterized.expand([(pp_size, tie_word_embeddings) for pp_size in [2] for tie_word_embeddings in [True, False]])     def test_pp_split(self, pp_size, tie_word_embeddings):@@ -328,7 +326,6 @@ def test_pp_split(self, pp_size, tie_word_embeddings):         )  -@require_torch_greater_or_equal("2.5") class TestPipelineParallelWeightLoading(TestCasePlus):     @parameterized.expand([(pp_size,) for pp_size in [2]])     def test_pp_weight_loading(self, pp_size):@@ -342,7 +339,6 @@ def test_pp_weight_loading(self, pp_size):         )  -@require_torch_greater_or_equal("2.5") class TestPipelineParallelGeneration(TestCasePlus):     @parameterized.expand([(pp_size,) for pp_size in [2]])     def test_pp_generation(self, pp_size):

@@ -39,7 +39,6 @@     backend_device_count,     require_accelerate,     require_torch,-    require_torch_greater_or_equal,     torch_device, ) from transformers.tokenization_utils_base import BatchEncoding@@ -267,7 +266,6 @@ def test_multiprocessing_context_is_forwarded(self):         # DataLoader turns the string into a multiprocessing context object.         self.assertEqual(dataloader.multiprocessing_context._name, "spawn") -    @require_torch_greater_or_equal("2.6")     def test_in_order_is_forwarded(self):         trainer = self._build_iterable_trainer(dataloader_num_workers=2, dataloader_in_order=False)         dataloader = trainer.get_train_dataloader()

@@ -1307,9 +1307,6 @@ def test_hybrid_cache_exportability(self):         """         Tests that static cache works with `torch.export()`         """-        if not is_torch_greater_or_equal("2.6"):-            self.skipTest(reason="This test requires torch >= 2.6 to run.")-         from transformers.integrations.executorch import TorchExportableModuleForDecoderOnlyLM          set_seed(42)