fix: post-release cleanup; torch amp cleanup (#4098) [skip ci]

* chore: drop torch<2.4 amp fallback in kernel utils

* chore: share kernel amp decorators with gradient checkpoint offload

* chore: remove torch version gates below the 2.13 minimum

src/axolotl/cli/merge_sharded_fsdp_weights.py

@@ -14,7 +14,6 @@ from accelerate import PartialState
 from accelerate.utils import (
     SAFE_WEIGHTS_INDEX_NAME,
     SAFE_WEIGHTS_NAME,
-    is_torch_version,
 )
 from huggingface_hub import split_torch_state_dict_into_shards
 from safetensors.torch import save_file as safe_save_file
@@ -125,13 +124,10 @@ def merge_fsdp_weights(
             Whether to remove the checkpoint directory after merging.
 
     Raises:
-        ValueError: If torch version < 2.3.0, or if `checkpoint_dir` does not exist.
+        ValueError: If `checkpoint_dir` does not exist.
     """
     checkpoint_dir_ = Path(checkpoint_dir)
 
-    if not is_torch_version(">=", "2.3.0"):
-        raise ValueError("`merge_fsdp_weights` requires PyTorch >= 2.3.0`")
-
     # Verify that the checkpoint directory exists
     if not checkpoint_dir_.exists():
         model_path_exists = (checkpoint_dir_ / "pytorch_model_fsdp_0").exists()

src/axolotl/kernels/utils.py

@@ -1,19 +1,12 @@
 """Utilities for `axolotl.kernels` submodules."""
 
 import torch
-from packaging.version import Version
 
 # Detect the actual accelerator so the AMP decorators are device-agnostic.
-_accelerator = (
-    torch.accelerator.current_accelerator() if hasattr(torch, "accelerator") else None
-)
+_accelerator = torch.accelerator.current_accelerator()
 # torch.amp.custom_fwd/bwd expect a device type string ("cuda", "npu", ...);
 # str(device) would be the invalid device type "None" on builds without one.
 _amp_device_type = _accelerator.type if _accelerator is not None else "cuda"
 
-if Version(torch.__version__) < Version("2.4.0"):
-    torch_amp_custom_fwd = torch.cuda.amp.custom_fwd
-    torch_amp_custom_bwd = torch.cuda.amp.custom_bwd
-else:
-    torch_amp_custom_fwd = torch.amp.custom_fwd(device_type=_amp_device_type)
-    torch_amp_custom_bwd = torch.amp.custom_bwd(device_type=_amp_device_type)
+torch_amp_custom_fwd = torch.amp.custom_fwd(device_type=_amp_device_type)
+torch_amp_custom_bwd = torch.amp.custom_bwd(device_type=_amp_device_type)

src/axolotl/monkeypatch/gradient_checkpointing/offload_cpu.py

@@ -14,35 +14,9 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import inspect
-
 import torch
-from packaging import version
-from torch.utils.checkpoint import (
-    set_device_states,
-)
 
-# support different pytorch versions
-has_device_type = "device_type" in inspect.signature(set_device_states).parameters
-
-torch_version = version.parse(torch.__version__)
-
-# Detect the actual accelerator (cuda/npu/xpu/mps/...) so the AMP custom_fwd/bwd
-# decorators are device-agnostic. Falls back to "cuda" when torch has no
-# `accelerator` API (old PyTorch).
-_accelerator = (
-    torch.accelerator.current_accelerator() if hasattr(torch, "accelerator") else None
-)
-# torch.amp.custom_fwd/bwd expect a device type string ("cuda", "npu", ...);
-# str(device) would be the invalid device type "None" on builds without one.
-_amp_device_type = _accelerator.type if _accelerator is not None else "cuda"
-
-if torch_version < version.parse("2.4.0"):
-    torch_cuda_amp_custom_fwd = torch.cuda.amp.custom_fwd
-    torch_cuda_amp_custom_bwd = torch.cuda.amp.custom_bwd
-else:
-    torch_cuda_amp_custom_fwd = torch.amp.custom_fwd(device_type=_amp_device_type)
-    torch_cuda_amp_custom_bwd = torch.amp.custom_bwd(device_type=_amp_device_type)
+from axolotl.kernels.utils import torch_amp_custom_bwd, torch_amp_custom_fwd
 
 
 class CPU_Offloaded_Gradient_Checkpointer(torch.autograd.Function):
@@ -52,7 +26,7 @@ class CPU_Offloaded_Gradient_Checkpointer(torch.autograd.Function):
     """
 
     @staticmethod
-    @torch_cuda_amp_custom_fwd
+    @torch_amp_custom_fwd
     def forward(ctx, forward_function, hidden_states, *args):
         ctx._device = (
             hidden_states.device
@@ -66,7 +40,7 @@ class CPU_Offloaded_Gradient_Checkpointer(torch.autograd.Function):
         return output
 
     @staticmethod
-    @torch_cuda_amp_custom_bwd
+    @torch_amp_custom_bwd
     def backward(ctx, dY):
         (hidden_states,) = ctx.saved_tensors
         hidden_states = hidden_states.to(ctx._device, non_blocking=True).detach()

src/axolotl/monkeypatch/gradient_checkpointing/offload_disk.py

@@ -31,19 +31,9 @@ from typing import Dict
 
 import torch
 
+from axolotl.kernels.utils import torch_amp_custom_bwd, torch_amp_custom_fwd
 from axolotl.utils.logging import get_logger
 
-# Detect the actual accelerator (cuda/npu/xpu/...) so the AMP custom_fwd/bwd
-# decorators are device-agnostic.
-_accelerator = (
-    torch.accelerator.current_accelerator() if hasattr(torch, "accelerator") else None
-)
-# torch.amp.custom_fwd/bwd expect a device type string ("cuda", "npu", ...);
-# str(device) would be the invalid device type "None" on builds without one.
-_amp_device_type = _accelerator.type if _accelerator is not None else "cuda"
-torch_cuda_amp_custom_fwd = torch.amp.custom_fwd(device_type=_amp_device_type)
-torch_cuda_amp_custom_bwd = torch.amp.custom_bwd(device_type=_amp_device_type)
-
 # Setup logger
 logger = get_logger(__name__)
 
@@ -451,7 +441,7 @@ class Disco(torch.autograd.Function):
         return Disco._manager
 
     @staticmethod
-    @torch_cuda_amp_custom_fwd
+    @torch_amp_custom_fwd
     def forward(
         ctx,
         forward_function,
@@ -488,7 +478,7 @@ class Disco(torch.autograd.Function):
         return output
 
     @staticmethod
-    @torch_cuda_amp_custom_bwd
+    @torch_amp_custom_bwd
     def backward(ctx, *grad_outputs):
         """Backward pass that loads activations from disk with prefetching"""
         # Get the manager

src/axolotl/utils/config/__init__.py

@@ -10,7 +10,6 @@ from pydantic import ValidationError
 from pydantic_core import PydanticUndefined, PydanticUndefinedType
 from transformers.utils import is_torch_bf16_gpu_available
 from transformers.utils.import_utils import (
-    is_torch_greater_or_equal,
     is_torch_npu_available,
 )
 
@@ -181,13 +180,9 @@ def resolve_dtype(cfg):
     else:
         if cfg.tf32 is True:
             torch.set_float32_matmul_precision("high")
-            if is_torch_greater_or_equal("2.9.0"):
-                torch.backends.fp32_precision = "tf32"
-                torch.backends.cuda.matmul.fp32_precision = "tf32"
-                torch.backends.cudnn.fp32_precision = "tf32"
-            else:
-                torch.backends.cuda.matmul.allow_tf32 = True
-                torch.backends.cudnn.allow_tf32 = True
+            torch.backends.fp32_precision = "tf32"
+            torch.backends.cuda.matmul.fp32_precision = "tf32"
+            torch.backends.cudnn.fp32_precision = "tf32"
         if cfg.bf16:
             cfg.fp16 = False
 

src/axolotl/utils/quantization.py

@@ -3,7 +3,6 @@ Utilities for quantization including QAT and PTQ using torchao.
 """
 
 import torch
-from packaging import version
 from torchao.core.config import AOBaseConfig
 from torchao.quantization import quantize_
 from torchao.quantization.granularity import PerGroup
@@ -31,38 +30,37 @@ quantization_config_to_str = {
     Float8DynamicActivationInt4WeightConfig: "fp8int4",
 }
 
-if version.parse(torch.__version__) >= version.parse("2.8.0"):
-    try:
-        from torchao.prototype.mx_formats import NVFP4WeightOnlyConfig
+try:
+    from torchao.prototype.mx_formats import NVFP4WeightOnlyConfig
 
-        quantization_config_to_str[NVFP4WeightOnlyConfig] = "nvfp4"
-    except (ImportError, RuntimeError):
-        pass
+    quantization_config_to_str[NVFP4WeightOnlyConfig] = "nvfp4"
+except (ImportError, RuntimeError):
+    pass
 
-    try:
-        from torchao.prototype.mx_formats import NVFP4DynamicActivationNVFP4WeightConfig
+try:
+    from torchao.prototype.mx_formats import NVFP4DynamicActivationNVFP4WeightConfig
 
-        quantization_config_to_str[NVFP4DynamicActivationNVFP4WeightConfig] = (
-            "nvfp4-dynamic"
-        )
-    except (ImportError, RuntimeError):
-        pass
+    quantization_config_to_str[NVFP4DynamicActivationNVFP4WeightConfig] = (
+        "nvfp4-dynamic"
+    )
+except (ImportError, RuntimeError):
+    pass
 
-    try:
-        from torchao.quantization.quant_api import Int4WeightOnlyConfig
+try:
+    from torchao.quantization.quant_api import Int4WeightOnlyConfig
 
-        quantization_config_to_str[Int4WeightOnlyConfig] = "int4"
-    except (ImportError, RuntimeError):
-        pass
+    quantization_config_to_str[Int4WeightOnlyConfig] = "int4"
+except (ImportError, RuntimeError):
+    pass
 
-    try:
-        from torchao.prototype.mx_formats import (
-            MXDynamicActivationMXWeightConfig as MXLinearConfig,
-        )
+try:
+    from torchao.prototype.mx_formats import (
+        MXDynamicActivationMXWeightConfig as MXLinearConfig,
+    )
 
-        quantization_config_to_str[MXLinearConfig] = "mxfp4"
-    except (ImportError, RuntimeError):
-        pass
+    quantization_config_to_str[MXLinearConfig] = "mxfp4"
+except (ImportError, RuntimeError):
+    pass
 
 
 def get_quantization_config(

src/axolotl/utils/schemas/config.py

@@ -1904,7 +1904,7 @@ class AxolotlConfigWCapabilities(AxolotlInputConfig):
 
     @model_validator(mode="after")
     def check_fp8_attention_preflight(self):
-        """fp8 attention requires SM90+ and torch >= 2.11 (torchao >= 0.17 is pinned)."""
+        """fp8 attention requires SM90+ (torchao >= 0.17 is pinned)."""
         if self.attn_implementation != "fp8":
             return self
 
@@ -1917,19 +1917,6 @@ class AxolotlConfigWCapabilities(AxolotlInputConfig):
                     f"higher (Hopper+). Detected {cc!r}."
                 )
 
-        torch_version = (
-            self.env_capabilities.torch_version if self.env_capabilities else None
-        )
-        if torch_version is None:
-            import torch
-
-            torch_version = str(torch.__version__).split("+", maxsplit=1)[0]
-        if version.parse(torch_version) < version.parse("2.11.0"):
-            raise ValueError(
-                f"attn_implementation=fp8 requires PyTorch >= 2.11.0. "
-                f"Detected {torch_version}."
-            )
-
         return self
 
     @model_validator(mode="before")
@@ -1990,62 +1977,20 @@ class AxolotlConfigWCapabilities(AxolotlInputConfig):
 
         return data
 
-    @model_validator(mode="before")
-    @classmethod
-    def check_adopt_torch_version(cls, data):
-        if (data.get("optimizer") is not None) and ("adopt" in data.get("optimizer")):
-            env_capabilities = data.get("env_capabilities", {})
-            torch_version = env_capabilities.get("torch_version")
-
-            if torch_version is None:
-                import torch
-
-                torch_version = str(torch.__version__).split("+", maxsplit=1)[0]
-
-            if version.parse(torch_version) < version.parse("2.5.1"):
-                raise ValueError(
-                    "ADOPT optimizer is incompatible with torch version < 2.5.1"
-                )
-        return data
-
-    @model_validator(mode="after")
-    def check_flex_torch_version(self):
-        if self.attn_implementation == "flex_attention":
-            torch_version = (
-                self.env_capabilities.torch_version if self.env_capabilities else None
-            )
-            if torch_version is None:
-                import torch
-
-                torch_version = str(torch.__version__).split("+", maxsplit=1)[0]
-
-            if version.parse(torch_version) < version.parse("2.6.0"):
-                raise ValueError(
-                    "Flex attention is not supported on torch version < 2.6.0"
-                )
-        return self
-
     @model_validator(mode="before")
     @classmethod
     def check_torch_compile_auto(cls, data):
         if data.get("torch_compile") == "auto":
             env_capabilities = data.get("env_capabilities", {})
             if env_capabilities.get("torch_version"):
-                if version.parse(
-                    env_capabilities.get("torch_version")
-                ) >= version.parse("2.5.1"):
-                    LOG.info(
-                        "torch.compile is available, setting torch_compile to True"
-                    )
-                    data["torch_compile"] = True
-                else:
-                    data["torch_compile"] = False
+                LOG.info("torch.compile is available, setting torch_compile to True")
+                data["torch_compile"] = True
             else:
                 data["torch_compile"] = False
             if data["torch_compile"] is False and data.get("torch_compile_options"):
                 LOG.warning(
                     "torch_compile: auto resolved to False on this environment "
-                    "(torch < 2.5.1); ignoring torch_compile_options."
+                    "(torch version unknown); ignoring torch_compile_options."
                 )
                 data["torch_compile_options"] = None
         return data
@@ -2062,9 +2007,9 @@ class AxolotlConfigWCapabilities(AxolotlInputConfig):
     def check_min_torch_version(self):
         if self.env_capabilities and self.env_capabilities.torch_version:
             torch_version = self.env_capabilities.torch_version
-            if version.parse(torch_version) < version.parse("2.6.0"):
+            if version.parse(torch_version) < version.parse("2.13.0"):
                 LOG.warning(
-                    f"torch=={torch_version} not be supported. Please upgrade to torch>=2.6.0."
+                    f"torch=={torch_version} may not be supported. Please upgrade to torch>=2.13.0."
                 )
 
         return self
@@ -2090,34 +2035,6 @@ class AxolotlConfigWCapabilities(AxolotlInputConfig):
         ) not in (None, "int8"):
             raise ValueError("Ternary QAT only supports activation_dtype: int8.")
 
-        env_capabilities = data.get("env_capabilities", {})
-        torch_version = env_capabilities.get("torch_version")
-
-        if torch_version is None:
-            import torch
-
-            torch_version = str(torch.__version__).split("+", maxsplit=1)[0]
-
-        if version.parse(torch_version) < version.parse("2.6.0"):
-            raise ValueError("QAT is not supported on torch version < 2.6.0")
-
-        return data
-
-    @model_validator(mode="before")
-    @classmethod
-    def check_fsdp_torch_version(cls, data):
-        env_capabilities = data.get("env_capabilities", {})
-        torch_version = env_capabilities.get("torch_version")
-
-        if torch_version is None:
-            import torch
-
-            torch_version = str(torch.__version__).split("+", maxsplit=1)[0]
-
-        if data.get("fsdp_config"):
-            if version.parse(torch_version) < version.parse("2.7.0"):
-                raise ValueError("FSDP2 is not supported on torch version < 2.7.0")
-
         return data
 
     @model_validator(mode="before")

src/axolotl/utils/schemas/validation.py

@@ -7,7 +7,6 @@ from collections.abc import Callable
 from pathlib import Path
 from typing import Any
 
-from packaging import version as packaging_version
 from pydantic import (
     field_validator,
     model_validator,
@@ -119,30 +118,6 @@ def lora_kernels_auto_enable(get: Callable[[str], Any]) -> bool:
     ):
         return False
 
-    # MoE without native grouped_mm (torch < 2.9): transformers' fallback uses
-    # torch.mm(out=) which bypasses autocast and fails on mixed dtypes in eval
-    torch_version = (get("env_capabilities") or {}).get("torch_version")
-    if torch_version is None:
-        import torch
-
-        torch_version = str(torch.__version__).split("+", maxsplit=1)[0]
-    if packaging_version.parse(torch_version) >= packaging_version.parse("2.9.0"):
-        return True
-    model_type = get("model_config_type") or ""
-    if "moe" in model_type.lower():
-        return False
-    base_model = get("base_model")
-    if base_model:
-        try:
-            from transformers import AutoConfig
-
-            auto_cfg = AutoConfig.from_pretrained(base_model, trust_remote_code=False)
-            if getattr(auto_cfg, "num_local_experts", None) or getattr(
-                auto_cfg, "num_experts", None
-            ):
-                return False
-        except Exception:  # pylint: disable=broad-exception-caught
-            pass
     return True
 
 

tests/monkeypatch/test_offload_amp_device_type.py

@@ -1,20 +1,14 @@
 """AMP decorators in the gradient-checkpointing offload paths."""
 
-import pytest
 import torch
 
 from axolotl.kernels.utils import torch_amp_custom_fwd
-from axolotl.monkeypatch.gradient_checkpointing import offload_disk
 from axolotl.monkeypatch.gradient_checkpointing.offload_cpu import (
     CPU_Offloaded_Gradient_Checkpointer,
 )
 
 
-@pytest.mark.parametrize(
-    "custom_fwd",
-    [torch_amp_custom_fwd, offload_disk.torch_cuda_amp_custom_fwd],
-)
-def test_amp_decorators_resolve_a_valid_device_type(custom_fwd):
+def test_amp_decorators_resolve_a_valid_device_type():
     """The decorators must not end up with the invalid device type "None".
 
     ``str(torch.accelerator.current_accelerator())`` is the string ``"None"``
@@ -22,7 +16,7 @@ def test_amp_decorators_resolve_a_valid_device_type(custom_fwd):
     a decorated function then raises a RuntimeError.
     """
 
-    @custom_fwd
+    @torch_amp_custom_fwd
     def double(x):
         return x * 2
 

tests/patched/test_validation.py

@@ -1363,38 +1363,6 @@ class TestValidation(BaseValidation):
                 in self._caplog.records[0].message
             )
 
-    def test_torch_version_adopt_req(self, minimal_cfg):
-        cfg = (
-            DictDefault(
-                {
-                    "optimizer": "adopt_adamw",
-                }
-            )
-            | minimal_cfg
-        )
-
-        with pytest.raises(
-            ValueError,
-            match=r".*ADOPT optimizer is incompatible with torch version*",
-        ):
-            env_capabilities = {"torch_version": "2.3.0"}
-            capabilities = {"bf16": False}
-            _ = validate_config(
-                cfg, capabilities=capabilities, env_capabilities=env_capabilities
-            )
-
-        env_capabilities = {"torch_version": "2.6.0"}
-        capabilities = {"bf16": False}
-        _ = validate_config(
-            cfg, capabilities=capabilities, env_capabilities=env_capabilities
-        )
-
-        env_capabilities = {"torch_version": "2.5.2"}
-        capabilities = {"bf16": False}
-        _ = validate_config(
-            cfg, capabilities=capabilities, env_capabilities=env_capabilities
-        )
-
 
 class TestTorchCompileValidation(BaseValidation):
     """
@@ -1419,14 +1387,6 @@ class TestTorchCompileValidation(BaseValidation):
 
         assert updated_cfg.torch_compile is True
 
-        env_capabilities = {"torch_version": "2.4.1"}
-        capabilities = {"bf16": True}
-        updated_cfg = validate_config(
-            cfg, capabilities=capabilities, env_capabilities=env_capabilities
-        )
-
-        assert updated_cfg.torch_compile is False
-
         env_capabilities = {}
         capabilities = {"bf16": True}
         updated_cfg = validate_config(
@@ -1538,7 +1498,7 @@ class TestTorchCompileValidation(BaseValidation):
             )
             | minimal_cfg
         )
-        env_capabilities = {"torch_version": "2.4.0"}
+        env_capabilities = {}
         with capture_axolotl_warnings(caplog):
             updated_cfg = validate_config(
                 cfg, capabilities={"bf16": True}, env_capabilities=env_capabilities

tests/utils/schemas/validation/test_lora_kernels_auto_enable.py

@@ -52,17 +52,6 @@ def test_blocked_by_model_capability(monkeypatch):
     assert lora_kernels_auto_enable(_get(model_config_type="x")) is True
 
 
-def test_blocked_for_moe_without_grouped_mm():
-    get = _get(
-        model_config_type="qwen3_moe", env_capabilities={"torch_version": "2.8.0"}
-    )
-    assert lora_kernels_auto_enable(get) is False
-    get = _get(
-        model_config_type="qwen3_moe", env_capabilities={"torch_version": "2.9.0"}
-    )
-    assert lora_kernels_auto_enable(get) is True
-
-
 def test_lora_kernel_enabled_follows_the_predicate(monkeypatch):
     cfg = SimpleNamespace(**BASE, lora_mlp_kernel=None, model_config_type="x")
     _fake_support(monkeypatch, Unsupported("no fused kernels"))