unslothai--unsloth/unsloth/import_fixes.py

# Copyright 2023-present Daniel Han-Chen & the Unsloth team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import os
import importlib.abc
import importlib.machinery
import importlib.util
from pathlib import Path
from importlib.metadata import version as importlib_version
from packaging.version import Version as TrueVersion
import re
import logging
import textwrap
import warnings
import sys
import functools
import inspect

# We cannot do from unsloth_zoo.log import logger since FBGEMM might cause seg faults.
UNSLOTH_ENABLE_LOGGING = os.environ.get("UNSLOTH_ENABLE_LOGGING", "0") in (
    "1",
    "True",
    "true",
)
logger = logging.getLogger(__name__)
if UNSLOTH_ENABLE_LOGGING:
    logging.basicConfig(level = logging.INFO, format = "[%(name)s|%(levelname)s]%(message)s")
    logger.setLevel(logging.INFO)
else:
    logging.basicConfig(level = logging.WARNING, format = "[%(name)s|%(levelname)s]%(message)s")
    logger.setLevel(logging.WARNING)

_AMDGPU_IDS_MISSING_TEXT = "amdgpu.ids: No such file or directory"


def Version(version):
    try:
        new_version = str(version)
        new_version = re.match(r"[0-9\.]{1,}", new_version)
        if new_version is None:
            raise ValueError(f"Could not parse version: {version}")
        new_version = new_version.group(0).rstrip(".")
        if new_version != version:
            new_version += ".1"  # Add .1 for dev / alpha / beta / rc
        return TrueVersion(new_version)
    except:
        from inspect import getframeinfo, stack
        caller = getframeinfo(stack()[1][0])
        raise RuntimeError(
            f"Unsloth: Could not get version for `{version}`\n"
            f"File name = [{caller.filename}] Line number = [{caller.lineno}]"
        )


# Ignore logging messages
class HideLoggingMessage(logging.Filter):
    __slots__ = ("text",)

    def __init__(self, text):
        self.text = text

    def filter(self, x):
        return not (self.text in x.getMessage())


class HidePrintMessage:
    def __init__(self, original_stream):
        self._original_stream = original_stream
        self._hidden_texts = []

    def add_filter(self, text):
        self._hidden_texts.append(text)

    def write(self, message):
        if not any(text in message for text in self._hidden_texts):
            self._original_stream.write(message)

    def flush(self):
        self._original_stream.flush()

    def __getattr__(self, name):
        return getattr(self._original_stream, name)


import contextlib
import ctypes

try:
    _libc = ctypes.CDLL(None)
except Exception:
    _libc = None


@contextlib.contextmanager
def suppress_cuda_printf():
    """Suppress CUDA device-side printf by redirecting stdout/stderr fds to /dev/null.

    CUDA device printf (e.g. CUTLASS "Arch conditional MMA" errors on Blackwell)
    writes to fd 1 at the C level, bypassing Python's sys.stdout, so the
    HidePrintMessage filter can't catch it. Redirect fd 1 and 2 at the OS level,
    sync CUDA, then restore.
    """
    sys.stdout.flush()
    sys.stderr.flush()
    saved_fds = {}
    try:
        for fd in (1, 2):
            saved_fds[fd] = os.dup(fd)
            devnull = os.open(os.devnull, os.O_WRONLY)
            os.dup2(devnull, fd)
            os.close(devnull)
        yield
    finally:
        try:
            import torch
            if torch.cuda.is_available():
                torch.cuda.synchronize()
        except Exception:
            pass
        if _libc is not None:
            try:
                _libc.fflush(None)
            except Exception:
                pass
        for fd, saved in saved_fds.items():
            os.dup2(saved, fd)
            os.close(saved)


if not UNSLOTH_ENABLE_LOGGING:
    import sys

    # Apply to stderr for FBGEMM and CUTLASS errors
    sys.stderr = HidePrintMessage(sys.stderr)
    # https://github.com/pytorch/FBGEMM/blob/d99cd96490ec4aabac2ee95b1e76ea4dcfcfa628/fbgemm_gpu/experimental/gemm/triton_gemm/utils.py#L43-L52
    sys.stderr.add_filter("TMA benchmarks will be running")
    # CUTLASS/FBGEMM MMA instruction error on SM90 vs SM100 (Blackwell) GPUs
    # https://github.com/NVIDIA/cutlass/blob/main/include/cutlass/gemm/kernel/sm90_gemm_tma_warpspecialized.hpp
    sys.stderr.add_filter("Arch conditional MMA instruction used without targeting")
    # CUTLASS arch conditional errors for various architectures
    sys.stderr.add_filter("CUTE_INVALID_CONTROL_PATH")
    # CUTLASS TMA-related errors when not targeting correct architecture
    sys.stderr.add_filter("Trying to use tma without CUTE_ARCH_TMA")
    # torchao logs a cosmetic "Skipping import of cpp extensions" WARNING on torch < 2.11. The
    # bnb-4bit / Unsloth paths don't use torchao's cpp kernels, so drop only that record rather
    # than raising the whole torchao logger to ERROR.
    logging.getLogger("torchao").addFilter(
        HideLoggingMessage("Skipping import of cpp extensions due to incompatible torch version")
    )
    # torch >= 2.11 path: torchao dlopens each prebuilt _C*.so and logs "Failed to load
    # .../_C*.so" when one can't (ABI tag mismatch in the wheel, e.g. a cp310 .so under a
    # cp312 runtime on Colab, or an arch-specific kernel the GPU lacks). It falls back to
    # non-cpp paths and Unsloth doesn't use these kernels, so drop the cosmetic record.
    logging.getLogger("torchao").addFilter(HideLoggingMessage("Failed to load "))
    # SyntaxWarning: invalid escape sequence '\.'
    warnings.filterwarnings("ignore", message = "invalid escape sequence", category = SyntaxWarning)
    # PYTORCH_CUDA_ALLOC_CONF is deprecated warning from torch
    warnings.filterwarnings("ignore", message = "PYTORCH_CUDA_ALLOC_CONF is deprecated")
    # TF32 precision deprecation warning from torch
    warnings.filterwarnings("ignore", message = "Please use the new API settings to control TF32")
    # Deprecation warnings from torchao
    warnings.filterwarnings("ignore", message = "`int4_weight_only` is deprecated")
    warnings.filterwarnings("ignore", message = "`int8_weight_only` is deprecated")
    # torch._check_is_size FutureWarning (called by bitsandbytes 4-bit dequant)
    warnings.filterwarnings(
        "ignore", message = r"_check_is_size will be removed", category = FutureWarning
    )

    # TorchAO deprecated import paths (https://github.com/pytorch/ao/issues/2752)
    warnings.filterwarnings(
        "ignore",
        message = r"Importing.*from torchao\.dtypes.*is deprecated",
        category = DeprecationWarning,
    )
    warnings.filterwarnings(
        "ignore",
        message = r"Importing BlockSparseLayout from torchao\.dtypes is deprecated",
        category = DeprecationWarning,
    )

    # SWIG builtin type warnings (from bitsandbytes/triton SWIG bindings)
    warnings.filterwarnings(
        "ignore",
        message = r"builtin type Swig.*has no __module__ attribute",
        category = DeprecationWarning,
    )

    # Triton autotuner deprecation (https://github.com/triton-lang/triton/pull/4496)
    warnings.filterwarnings(
        "ignore",
        message = r"warmup, rep, and use_cuda_graph parameters are deprecated",
        category = DeprecationWarning,
    )

    # Python 3.12+ multiprocessing fork warning in multi-threaded processes
    warnings.filterwarnings(
        "ignore",
        message = r".*multi-threaded.*use of fork\(\) may lead to deadlocks",
        category = DeprecationWarning,
    )

    # Resource warnings from internal socket/file operations
    warnings.filterwarnings("ignore", message = r"unclosed.*socket", category = ResourceWarning)
    warnings.filterwarnings("ignore", message = r"unclosed file.*dev/null", category = ResourceWarning)

    # torch 2.9+ pin_memory/is_pinned device arg deprecation
    warnings.filterwarnings(
        "ignore",
        message = r"The `device` argument is deprecated",
        category = DeprecationWarning,
    )
    warnings.filterwarnings(
        "ignore",
        message = r".*pin_memory.*device.*deprecated",
        category = DeprecationWarning,
    )
    warnings.filterwarnings(
        "ignore",
        message = r".*is_pinned.*device.*deprecated",
        category = DeprecationWarning,
    )

    # vllm "Level is deprecated" stderr noise
    sys.stderr.add_filter("Level is deprecated")

    # PydanticSerializationUnexpectedValue warning
    warnings.filterwarnings(
        "ignore",
        message = r".*PydanticSerializationUnexpectedValue",
    )
    warnings.filterwarnings(
        "ignore",
        message = r"Expected.*but got.*with value.*is not.*subclass",
    )

    # Triton "df: No such file or directory" stderr noise
    sys.stderr.add_filter("df: No such file")
    # ROCm/libdrm missing ids table stderr noise on some AMD setups
    sys.stderr.add_filter(_AMDGPU_IDS_MISSING_TEXT)
    # Apex ROCm fused RoPE backend selection warning when Aiter is enabled.
    warnings.filterwarnings(
        "ignore",
        message = r"^Aiter backend is selected for fused RoPE\.?",
        category = UserWarning,
        module = r"^apex\.transformer\.functional\.fused_rope$",
    )


def fix_torch_check_is_size():
    """Shim torch._check_is_size if a future torch removes it (bitsandbytes 4-bit
    dequant calls it). The FutureWarning is silenced in suppress_cuda_printf."""
    try:
        import torch

        if hasattr(torch, "_check_is_size"):
            return

        def _check_is_size(
            i,
            message = None,
            *,
            max = None,
        ):
            torch._check(i >= 0, message)
            if max is not None:
                torch._check(i <= max, message)

        torch._check_is_size = _check_is_size
    except Exception:
        return


# Fix up AttributeError: 'MessageFactory' object has no attribute 'GetPrototype'
# MUST do this at the start primarily due to tensorflow causing issues
def fix_message_factory_issue():
    try:
        import google.protobuf.message_factory

        class MessageFactory:
            def CreatePrototype(self, *args, **kwargs):
                return

            def GetMessages(self, *args, **kwargs):
                return

            def GetPrototype(self, *args, **kwargs):
                return

        if not hasattr(google.protobuf.message_factory, "MessageFactory"):
            logger.info("Unsloth: Patching protobuf.MessageFactory as it doesn't exist")
            google.protobuf.message_factory.MessageFactory = MessageFactory
        elif (
            hasattr(google.protobuf.message_factory, "MessageFactory")
            and not hasattr(google.protobuf.message_factory.MessageFactory, "GetPrototype")
            and not hasattr(google.protobuf.message_factory, "GetMessageClass")
        ):
            google.protobuf.message_factory.MessageFactory = MessageFactory
            logger.info("Unsloth: Patching protobuf.MessageFactory as it doesn't exist")
        elif (
            hasattr(google.protobuf.message_factory, "MessageFactory")
            and not hasattr(google.protobuf.message_factory.MessageFactory, "GetPrototype")
            and hasattr(google.protobuf.message_factory, "GetMessageClass")
        ):
            GetMessageClass = google.protobuf.message_factory.GetMessageClass

            def GetPrototype(self, descriptor):
                return GetMessageClass(descriptor)

            google.protobuf.message_factory.MessageFactory.GetPrototype = GetPrototype
            logger.info("Unsloth: Patching protobuf.MessageFactory.GetPrototype")
        pass
    except:
        pass


# Fix Xformers performance issues since 0.0.25
def fix_xformers_performance_issue():
    spec = importlib.util.find_spec("xformers")
    if spec is None:
        return
    xformers_version = importlib_version("xformers")
    if Version(xformers_version) < Version("0.0.29"):
        xformers_location = spec.origin
        if xformers_location is None:
            xformers_location = spec.submodule_search_locations[0]
        else:
            xformers_location = os.path.split(xformers_location)[0]
        cutlass = Path(xformers_location) / "ops" / "fmha" / "cutlass.py"
        try:
            if cutlass.exists():
                with open(cutlass, "r+", encoding = "utf-8") as f:
                    text = f.read()
                    # See https://github.com/facebookresearch/xformers/issues/1176#issuecomment-2545829591
                    if "num_splits_key=-1," in text:
                        text = text.replace(
                            "num_splits_key=-1,",
                            "num_splits_key=None,",
                        )
                        f.seek(0)
                        f.write(text)
                        f.truncate()
                        logger.info("Unsloth: Patching Xformers to fix some performance issues.")
        except Exception as e:
            logger.info(f"Unsloth: Failed patching Xformers with error = {str(e)}")


def patch_vllm_for_notebooks():
    import sys

    ipython = None
    try:
        from IPython import get_ipython as _get_ipython
    except Exception:
        _get_ipython = None

    if _get_ipython is not None:
        try:
            ipython = _get_ipython()
        except Exception:
            ipython = None

    if ipython is None:
        try:
            import builtins
            _get_ipython = getattr(builtins, "get_ipython", None)
            if callable(_get_ipython):
                ipython = _get_ipython()
        except Exception:
            ipython = None

    if ipython is None:
        return

    try:
        shell = ipython.__class__.__name__
        is_notebook = shell == "ZMQInteractiveShell" or "google.colab" in str(type(ipython))
    except Exception:
        return

    if not is_notebook:
        return

    if not hasattr(sys.stdout, "fileno"):
        return

    needs_patch = False
    try:
        fd = sys.stdout.fileno()
        if not isinstance(fd, int) or fd < 0:
            needs_patch = True
    except Exception:
        needs_patch = True

    if not needs_patch:
        return

    logger.info(
        "Unsloth: Notebook detected - Patching sys.stdout.fileno for newer `vllm>=0.12.0` versions"
    )
    sys.stdout.fileno = lambda: 1


# ValueError: 'aimv2' is already used by a Transformers config, pick another name.
def fix_vllm_aimv2_issue():
    spec = importlib.util.find_spec("vllm")
    if spec is None:
        return
    vllm_version = importlib_version("vllm")
    if Version(vllm_version) < Version("0.10.1"):
        vllm_location = spec.origin
        if vllm_location is None:
            vllm_location = spec.submodule_search_locations[0]
        else:
            vllm_location = os.path.split(vllm_location)[0]
        ovis_config = Path(vllm_location) / "transformers_utils" / "configs" / "ovis.py"
        try:
            if ovis_config.exists():
                with open(ovis_config, "r+", encoding = "utf-8") as f:
                    text = f.read()
                    # See https://github.com/vllm-project/vllm-ascend/issues/2046
                    if 'AutoConfig.register("aimv2", AIMv2Config)' in text:
                        text = text.replace(
                            'AutoConfig.register("aimv2", AIMv2Config)',
                            "",
                        )
                        text = text.replace(
                            """backbone_config.pop('model_type')
                backbone_config = AutoConfig.for_model(model_type,
                                                       **backbone_config)""",
                            """if model_type != "aimv2":
                    backbone_config.pop('model_type')
                    backbone_config = AutoConfig.for_model(model_type, **backbone_config)
                else:
                    backbone_config = AIMv2Config(**backbone_config)""",
                        )
                        f.seek(0)
                        f.write(text)
                        f.truncate()
                        logger.info(
                            "Unsloth: Patching vLLM to fix `'aimv2' is already used by a Transformers config, pick another name.`"
                        )
        except Exception as e:
            logger.info(f"Unsloth: Failed patching vLLM with error = {str(e)}")


# vLLM >= 0.22 (PR #35024) deleted `vllm.transformers_utils.tokenizer`, but an
# older unsloth_zoo still imports it unguarded and crashes (issue #6385). Supply
# a stub via a meta path finder appended AFTER the real finders, so it only
# activates when vLLM no longer ships the module.
_VLLM_LORA_TOKENIZER_MODULE = "vllm.transformers_utils.tokenizer"
_VLLM_TOKENIZER_STUB_SENTINEL = "__unsloth_vllm_tokenizer_stub__"


def _unsloth_return_no_lora_tokenizer(*args, **kwargs):
    # None -> vLLM uses the base tokenizer for LoRA (matches unsloth_zoo).
    return None


class _VllmLoraTokenizerStubLoader(importlib.abc.Loader):
    __slots__ = ("module_name",)

    def __init__(self, module_name):
        self.module_name = module_name

    def create_module(self, spec):
        import types

        module = types.ModuleType(self.module_name)
        module.__file__ = f"<unsloth stub: {self.module_name}>"
        module.__package__ = self.module_name.rpartition(".")[0]
        setattr(module, _VLLM_TOKENIZER_STUB_SENTINEL, True)
        module.get_lora_tokenizer = _unsloth_return_no_lora_tokenizer
        module.get_lora_tokenizer_async = _unsloth_return_no_lora_tokenizer
        return module

    def exec_module(self, module):
        return None


class _VllmLoraTokenizerStubFinder(importlib.abc.MetaPathFinder):
    __slots__ = (_VLLM_TOKENIZER_STUB_SENTINEL,)

    def __init__(self):
        setattr(self, _VLLM_TOKENIZER_STUB_SENTINEL, True)

    def find_spec(
        self,
        fullname,
        path = None,
        target = None,
    ):
        if fullname != _VLLM_LORA_TOKENIZER_MODULE:
            return None
        return importlib.machinery.ModuleSpec(
            name = fullname,
            loader = _VllmLoraTokenizerStubLoader(fullname),
            is_package = False,
        )


def fix_vllm_lora_tokenizer_module():
    if importlib.util.find_spec("vllm") is None:
        return
    for finder in sys.meta_path:
        if getattr(finder, _VLLM_TOKENIZER_STUB_SENTINEL, False):
            return
    # Appended, not inserted at 0, so a real module on older vLLM always wins.
    sys.meta_path.append(_VllmLoraTokenizerStubFinder())
    logger.info(
        "Unsloth: Installed `vllm.transformers_utils.tokenizer` compatibility "
        "stub for newer vLLM versions"
    )


def fix_vllm_guided_decoding_params():
    def _maybe_raise_vllm_transformers_mismatch(error):
        error_text = str(error)
        if "ALLOWED_LAYER_TYPES" in error_text or "transformers.configuration_utils" in error_text:
            try:
                vllm_version = importlib_version("vllm")
            except Exception:
                vllm_version = "unknown"
            raise RuntimeError(
                "Unsloth: vLLM with version "
                f"{vllm_version} does not yet support transformers>=5.0.0. "
                "Please downgrade to transformers==4.57.3 via "
                'pip install --force-reinstall "transformers==4.57.3". '
                f"Original error: {error}"
            ) from error

    if importlib.util.find_spec("vllm") is None:
        return
    # GuidedDecodingParmas is renamed to StructuredOutputsParams in vLLM
    # https://github.com/vllm-project/vllm/pull/22772/files
    # trl still wants to use GuidedDecodingParams. This is a temporary patch till trl updates
    try:
        import vllm
    except (ImportError, OSError) as e:
        _maybe_raise_vllm_transformers_mismatch(e)
        if disable_broken_vllm(e):
            return
        raise

    try:
        from vllm.sampling_params import GuidedDecodingParams
    except (ImportError, OSError) as e:
        _maybe_raise_vllm_transformers_mismatch(e)
        if disable_broken_vllm(e):
            return
        if not hasattr(vllm, "sampling_params") or not hasattr(
            vllm.sampling_params, "StructuredOutputsParams"
        ):
            raise
        vllm.sampling_params.GuidedDecodingParams = vllm.sampling_params.StructuredOutputsParams


def fix_trl_vllm_ascend():
    # transformers >= 4.48's `_is_package_available(name)` returns a tuple
    # (bool, version_or_None). TRL caches that tuple in module-level
    # `_*_available` flags and the matching `is_*_available()` accessors
    # return the tuple directly. A non-empty tuple is always truthy, so
    # `if is_X_available():` fires even when X is absent, triggering an
    # unconditional `import X` that fails. The surfaced case is
    # `vllm_ascend` (blocks `from trl import GRPOConfig, GRPOTrainer`
    # outside Huawei Ascend hosts); `llm_blender`, `deepspeed`, `joblib`
    # share the same shape. Coerce every tuple-cached flag in
    # trl.import_utils to bool; the existing accessors that just return
    # the cached value then naturally yield a bool.
    if importlib.util.find_spec("trl") is None:
        return
    try:
        import trl.import_utils as tiu
    except Exception:
        return
    for attr in list(vars(tiu)):
        if not (attr.startswith("_") and attr.endswith("_available")):
            continue
        cached = getattr(tiu, attr)
        if isinstance(cached, tuple):
            setattr(tiu, attr, bool(cached and cached[0]))


def ignore_logger_messages():
    # Ignore Environment variable `HF_TOKEN` is set
    try:
        from huggingface_hub._login import logger as huggingface_hub_logger
        huggingface_hub_logger.addFilter(HideLoggingMessage("`HF_TOKEN`"))
        del huggingface_hub_logger
    except:
        pass


def patch_ipykernel_hf_xet():
    # HF-XET == 1.1.10 and ipykernel == 7.0.0 / 7.0.1 causes issues
    # See https://github.com/huggingface/xet-core/issues/526
    # 2025-10-13T20:37:33.028737Z ERROR  Python exception updating progress:, error: PyErr { type: <class 'LookupError'>, value: LookupError(<ContextVar name='shell_parent' at 0x7535b4cebd80>), traceback: Some(<traceback object at 0x753408489f40>) }, caller: "src/progress_update.rs:313"
    # at /home/runner/work/xet-core/xet-core/error_printer/src/lib.rs:28
    if importlib.util.find_spec("hf_xet") is None:
        return
    if importlib.util.find_spec("ipykernel") is None:
        return
    if importlib.util.find_spec("huggingface_hub") is None:
        return

    ipykernel_version = Version(importlib_version("ipykernel"))
    if (
        (Version(importlib_version("hf_xet")) == Version("1.1.10"))
        and (
            (ipykernel_version == Version("7.0.0"))
            or (
                ipykernel_version == Version("7.0.1")
            )  # 7.0.1 seems to also break with LookupError: <ContextVar name='shell_parent' at 0x7a9775143ec0>
        )
    ):
        print(
            "#### Unsloth: `hf_xet==1.1.10` and `ipykernel==7.0.0` or `ipykernel==7.0.1` breaks progress bars. Using ASCII progress bars.\n"
            "#### Unsloth: To re-enable progress bars, please upgrade to `ipykernel>=7.1.0` or wait for a fix to\n"
            "https://github.com/huggingface/xet-core/issues/526"
        )
        from huggingface_hub.utils import disable_progress_bars
        disable_progress_bars()


def patch_trackio():
    # Set some environment variables to customize the Trackio dashboard for experiment tracking
    # See https://github.com/unslothai/notebooks/pull/110
    os.environ["TRACKIO_LOGO_LIGHT_URL"] = (
        "https://raw.githubusercontent.com/unslothai/unsloth/main/images/unsloth%20logo%20black%20text.png"
    )
    os.environ["TRACKIO_LOGO_DARK_URL"] = (
        "https://raw.githubusercontent.com/unslothai/unsloth/main/images/unsloth%20logo%20white%20text.png"
    )
    os.environ["TRACKIO_PLOT_ORDER"] = "train/reward"


def patch_datasets():
    # Datasets 4.4.0 and 4.4.1 weirdly have some weird `_thread.RLock_recursion_count` issues
    if importlib.util.find_spec("datasets") is None:
        return

    datasets_version = Version(importlib_version("datasets"))
    if (datasets_version <= Version("4.5.0")) and (datasets_version >= Version("4.4.0")):
        raise NotImplementedError(
            f"#### Unsloth: Using `datasets = {str(datasets_version)}` will cause recursion errors.\n"
            "Please downgrade datasets to `datasets==4.3.0"
        )


def check_fbgemm_gpu_version():
    if importlib.util.find_spec("fbgemm_gpu") is None:
        return
    try:
        fbgemm_gpu_version = importlib_version("fbgemm_gpu_genai")
    except:
        return
    # We noticed some SegFault or bad alloc errors on lower versions of fbgemm_gpu.
    # Instead of raising an error, disable FBGEMM and fall back to Triton kernels.
    if Version(fbgemm_gpu_version) < Version("1.4.0"):
        os.environ["UNSLOTH_HAS_FBGEMM"] = "0"
        logger.info(
            f"Unsloth: fbgemm_gpu_genai=={fbgemm_gpu_version} is old and may cause issues. "
            f"Disabling FBGEMM - using Triton kernels instead."
        )
        return

    logger.info(f"Unsloth: fbgemm_gpu_genai=={fbgemm_gpu_version} detected.")


def patch_enable_input_require_grads():
    """Patch PreTrainedModel.enable_input_require_grads to tolerate vision models
    that raise NotImplementedError from get_input_embeddings()."""
    import inspect
    from transformers import PreTrainedModel

    # Only patch the new variant that iterates over self.modules().
    # Ref: https://github.com/huggingface/transformers/pull/41993/files#diff-6b72b98c4c2dcfc6cc606843917733f5d858374fbc22a735ff483bbc0c1e63eaL1979-R1996
    try:
        original_source = inspect.getsource(PreTrainedModel.enable_input_require_grads)
    except:
        return

    if "for module in self.modules()" not in original_source:
        return

    def _patched_enable_input_require_grads(self):
        def make_inputs_require_grads(module, input, output):
            output.requires_grad_(True)

        hooks = []
        seen_modules = set()

        for module in self.modules():
            if not (
                isinstance(module, PreTrainedModel) and hasattr(module, "get_input_embeddings")
            ):
                continue

            try:
                input_embeddings = module.get_input_embeddings()
            except NotImplementedError:
                # Vision models may not implement get_input_embeddings (e.g. GLM
                # V4.6 skips only `self.visual`); skip them
                continue

            if input_embeddings is None:
                continue

            embedding_id = id(input_embeddings)
            if embedding_id in seen_modules:
                continue

            seen_modules.add(embedding_id)
            hooks.append(input_embeddings.register_forward_hook(make_inputs_require_grads))

        self._require_grads_hooks = hooks
        if hooks:
            self._require_grads_hook = hooks[0]

    PreTrainedModel.enable_input_require_grads = _patched_enable_input_require_grads

    logger.info("Unsloth: Patched enable_input_require_grads for vision model compatibility")


def patch_unsafe_trainer_rng_load():
    """Harden Trainer._load_rng_state against CVE-2026-1839 (RCE from a malicious
    rng_state.pth on resume). Hardens only the rng torch.load, via a thread-local
    flag, so it forces weights_only=True (defeats TORCH_FORCE_NO_WEIGHTS_ONLY_LOAD)
    and refuses torch < 2.6 (CVE-2025-32434), while rng-less resumes and unrelated
    torch.load calls are untouched. No-op if transformers is absent or already
    guards the load (>= 5.0.0rc3)."""
    if importlib.util.find_spec("transformers") is None:
        return
    try:
        from transformers.trainer import Trainer
    except Exception:
        return
    load_rng_state = getattr(Trainer, "_load_rng_state", None)
    if load_rng_state is None or getattr(load_rng_state, "_unsloth_safe_rng_load", False):
        return
    try:
        source = inspect.getsource(load_rng_state)
    except Exception:
        return
    if "torch.load" not in source or "check_torch_load_is_safe" in source:
        return

    import threading, torch

    try:
        # Older supported transformers (>= 4.51.3) may not export the helper.
        from transformers.utils.import_utils import check_torch_load_is_safe
    except Exception:

        def check_torch_load_is_safe():
            if TrueVersion(torch.__version__.split("+")[0]) < TrueVersion("2.6"):
                raise RuntimeError(
                    "Unsloth: refusing to load checkpoint RNG state on torch < 2.6 "
                    "(CVE-2026-1839 / CVE-2025-32434); upgrade to torch >= 2.6."
                )

    # Install one process-wide torch.load shim that stays inert unless the calling
    # thread is inside _load_rng_state, so we gate only at the real rng load with
    # no global-swap race and no effect on other torch.load callers.
    if not getattr(torch.load, "_unsloth_rng_guard", False):
        _orig_load = torch.load
        _rng_active = threading.local()

        @functools.wraps(_orig_load)
        def _guarded_torch_load(*args, **kwargs):
            if getattr(_rng_active, "on", False):
                check_torch_load_is_safe()  # raises on torch < 2.6 (CVE-2025-32434)
                kwargs.setdefault("weights_only", True)
            return _orig_load(*args, **kwargs)

        _guarded_torch_load._unsloth_rng_guard = True
        _guarded_torch_load._unsloth_rng_flag = _rng_active
        torch.load = _guarded_torch_load
    _rng_active = torch.load._unsloth_rng_flag

    @functools.wraps(load_rng_state)
    def _unsloth_safe_load_rng_state(self, checkpoint):
        _rng_active.on = True
        try:
            return load_rng_state(self, checkpoint)
        finally:
            _rng_active.on = False

    _unsloth_safe_load_rng_state._unsloth_safe_rng_load = True
    Trainer._load_rng_state = _unsloth_safe_load_rng_state
    logger.info("Unsloth: Hardened Trainer._load_rng_state rng loading (CVE-2026-1839).")


def _is_custom_torch_build(raw_version_str):
    """Check if a raw version string indicates a custom or source build.

    Operates on the raw importlib_version() string (our Version() strips local
    identifiers). Standard releases use +cu124/+rocm6.3/+cpu/+xpu; custom builds
    use +gitXXXX or other suffixes.
    """
    if "+" not in raw_version_str:
        return False
    local = raw_version_str.split("+", 1)[1]
    if not local:
        return False
    # Use fullmatch so the entire local identifier must match, not just a prefix.
    # cu/rocm require a trailing digit (e.g. cu124, rocm6.3). cpu/xpu are exact.
    # Case-insensitive since some builds may use uppercase.
    return not re.fullmatch(r"cu\d[\d.]*|rocm\d[\d.]*|cpu|xpu", local, re.IGNORECASE)


def _infer_required_torchvision(torch_major, torch_minor):
    """Infer the minimum required torchvision minor version from torch version.

    The torch -> torchvision minor version mapping follows a consistent formula:
      torch 1.x  ->  torchvision 0.(x + 1)   (verified: torch 1.7 through 1.13)
      torch 2.x  ->  torchvision 0.(x + 15)  (verified: torch 2.0 through 2.9)

    Returns (tv_major, tv_minor) or None if the major version is unrecognized.
    """
    if torch_major == 1 and torch_minor >= 7:
        return (0, torch_minor + 1)
    if torch_major == 2:
        return (0, torch_minor + 15)
    return None


def torchvision_compatibility_check():
    # Allow skipping via environment variable for custom environments
    if os.environ.get("UNSLOTH_SKIP_TORCHVISION_CHECK", "0").lower() in ("1", "true"):
        return

    if importlib.util.find_spec("torch") is None:
        raise ImportError("Unsloth: torch not found. Please install torch first.")
    if importlib.util.find_spec("torchvision") is None:
        return

    try:
        torch_version_raw = importlib_version("torch")
        torchvision_version_raw = importlib_version("torchvision")
    except Exception:
        return

    try:
        torch_v = Version(torch_version_raw)
        tv_v = Version(torchvision_version_raw)
    except Exception:
        return

    # Known compatibility table (ground truth, takes precedence over formula).
    # See https://pytorch.org/get-started/previous-versions/
    TORCH_TORCHVISION_COMPAT = {
        (2, 9): (0, 24),
        (2, 8): (0, 23),
        (2, 7): (0, 22),
        (2, 6): (0, 21),
        (2, 5): (0, 20),
        (2, 4): (0, 19),
    }

    torch_release = torch_v.release
    if len(torch_release) < 2:
        return
    torch_major, torch_minor = torch_release[0], torch_release[1]

    # Known table first, then the formula for forward compatibility
    required = TORCH_TORCHVISION_COMPAT.get((torch_major, torch_minor))

    if required is None:
        required = _infer_required_torchvision(torch_major, torch_minor)

    if required is None:
        return

    required_tv_str = f"{required[0]}.{required[1]}.0"

    if tv_v >= Version(required_tv_str):
        logger.info(
            f"Unsloth: torch=={torch_version_raw} and "
            f"torchvision=={torchvision_version_raw} are compatible."
        )
        return

    # Version mismatch detected
    message = (
        f"Unsloth: torch=={torch_version_raw} requires "
        f"torchvision>={required_tv_str}, "
        f"but found torchvision=={torchvision_version_raw}. "
        f'Try updating torchvision via `pip install --upgrade "torchvision>={required_tv_str}"`. '
        f"Please refer to https://pytorch.org/get-started/previous-versions/ "
        f"for more information."
    )

    is_custom = _is_custom_torch_build(torch_version_raw) or _is_custom_torch_build(
        torchvision_version_raw
    )

    # Detect nightly/dev/alpha/beta/rc builds from the raw version string.
    # These often have version mismatches that are expected.
    _pre_tags = (".dev", "a0", "b0", "rc", "alpha", "beta", "nightly")
    is_prerelease = any(t in torch_version_raw for t in _pre_tags) or any(
        t in torchvision_version_raw for t in _pre_tags
    )

    # Only downgrade to warning for custom/source or prerelease builds.
    # Stable mismatches should fail fast to prevent runtime operator errors.
    if is_custom or is_prerelease:
        reason = "custom/source build" if is_custom else "pre-release build"
        logger.warning(
            f"{message}\n"
            f"Detected a {reason}. "
            f"Continuing with a warning. "
            f"Set UNSLOTH_SKIP_TORCHVISION_CHECK=1 to silence this."
        )
        return

    raise ImportError(message)


# Fix TRL OpenEnv 0.26 NameError: name 'SamplingParams' is not defined
def fix_openenv_no_vllm():
    spec = importlib.util.find_spec("trl")
    if spec is None:
        return
    trl_location = spec.origin
    if trl_location is None:
        trl_location = spec.submodule_search_locations[0]
    else:
        trl_location = os.path.split(trl_location)[0]
    openenv = Path(trl_location) / "experimental" / "openenv" / "utils.py"
    if not openenv.exists():
        return

    try:
        with open(openenv, "r+", encoding = "utf-8") as f:
            text = f.read()
            bad = (
                "if is_vllm_available():\n"
                "    from vllm import SamplingParams\n"
                "    from vllm.sampling_params import GuidedDecodingParams\n"
            )
            replace_with = bad + (
                "else:\n"
                "    from typing import Any\n"
                "    SamplingParams = Any\n"
                "    GuidedDecodingParams = Any\n"
                "\n"
            )
            if bad + "\n" + "\n" in text and replace_with not in text:
                text = text.replace(bad + "\n" + "\n", replace_with)
                f.seek(0)
                f.write(text)
                f.truncate()
                logger.info("Unsloth: Patching TRL OpenEnv to fix SamplingParams not defined")
    except Exception as e:
        logger.info(f"Unsloth: Failed patching TRL OpenEnv with error = {str(e)}")


# Fix Exeuctorch needing get_mapped_key
def fix_executorch():
    spec = importlib.util.find_spec("executorch")
    if spec is None:
        return
    executorch_location = spec.origin
    if executorch_location is None:
        executorch_location = spec.submodule_search_locations[0]
    else:
        executorch_location = os.path.split(executorch_location)[0]
    executorch = Path(executorch_location) / "examples" / "models" / "__init__.py"
    if not executorch.exists():
        return

    try:
        what = r"""
        import sys
        import types
        import re
        from typing import Any, Optional
        def get_mapped_key(key: str, mapping_dict: dict[str, str]) -> str:
            try:
                # Checks if there is a layer # in the key
                if any(k.isdigit() for k in key.split(".")):
                    # Replace layer number with "{}" to create key for lookup
                    abstract_key = re.sub(r"(\.\d+)", ".{}", key)
                    layer_num = re.search(r"\d+", key).group(0)
                    new_key = mapping_dict[abstract_key]
                    new_key = new_key.format(layer_num)
                else:
                    new_key = mapping_dict[key]
            except KeyError as e:
                raise Exception(
                    f'Error converting the state dict. Found unexpected key: "{key}". '
                    "Please make sure you're loading a checkpoint with the right format. "
                ) from e

            return new_key

        torchtune = types.ModuleType("torchtune")
        torchtune.__path__ = []
        models = types.ModuleType("torchtune.models")
        models.__path__ = []
        convert_weights = types.ModuleType("torchtune.models.convert_weights")
        convert_weights.get_mapped_key = get_mapped_key
        torchtune.models = models
        models.convert_weights = convert_weights
        sys.modules["torchtune"] = torchtune
        sys.modules["torchtune.models"] = models
        sys.modules["torchtune.models.convert_weights"] = convert_weights
        """
        what = textwrap.dedent(what)

        with open(executorch, "r+", encoding = "utf-8") as f:
            text = f.read()
            bad = "from enum import Enum\n"
            if bad in text and what not in text:
                text = text.replace(bad + "\n", bad + "\n" + what)
                f.seek(0)
                f.write(text)
                f.truncate()
                logger.info("Unsloth: Patching Executorch to fix get_mapped_key")
    except Exception as e:
        logger.info(f"Unsloth: Failed Executorch with error = {str(e)}")


def fix_diffusers_warnings():
    # Silence Flax classes are deprecated and will be removed in Diffusers v1.0.0.
    os.environ["DIFFUSERS_VERBOSITY"] = "error"


def fix_huggingface_hub():
    # huggingface_hub.is_offline_mode got removed, so add it back
    import huggingface_hub
    if not hasattr(huggingface_hub, "is_offline_mode"):
        huggingface_hub.is_offline_mode = lambda: huggingface_hub.constants.HF_HUB_OFFLINE


def fix_triton_compiled_kernel_missing_attrs():
    """
    Triton 3.6.0+ removed direct `num_ctas` and `cluster_dims` attributes from
    CompiledKernel, but torch 2.9.x Inductor still expects them in
    torch/_inductor/runtime/triton_heuristics.py make_launcher() (line ~1757).

    The scope dict eagerly evaluates:
        binary.metadata.num_ctas, *binary.metadata.cluster_dims
    when hasattr(binary, "metadata") is True, but metadata lacks cluster_dims.
    This crashes before reaching the new launch path that doesn't need cta_args.

    Upstream fix: pytorch/pytorch@97bd4db added hasattr guards.
    We monkey-patch CompiledKernel.__init__ to inject the missing attributes
    so the older hasattr(binary, "num_ctas") branch succeeds instead.
    """
    try:
        import torch
    except (ImportError, ModuleNotFoundError):
        return

    try:
        import triton
        import triton.compiler.compiler as triton_compiler
    except (ImportError, ModuleNotFoundError):
        return

    # Only needed when the CompiledKernel class lacks num_ctas as a direct attr
    # but has metadata (triton >= 3.6.0 with torch < 2.10)
    _ck_cls = triton_compiler.CompiledKernel
    if hasattr(_ck_cls, "num_ctas"):
        return  # Old triton with direct attrs -- no patch needed

    _orig_init = _ck_cls.__init__

    def _patched_init(self, *args, **kwargs):
        _orig_init(self, *args, **kwargs)
        if not hasattr(self, "num_ctas"):
            self.num_ctas = getattr(self.metadata, "num_ctas", 1)
        if not hasattr(self, "cluster_dims") and not hasattr(self, "clusterDims"):
            self.cluster_dims = (1, 1, 1)

    _ck_cls.__init__ = _patched_init
    logger.info(
        "Unsloth: Patched triton CompiledKernel with num_ctas/cluster_dims "
        "for torch.compile compatibility."
    )


def fix_dynamo_config_thread_visibility():
    """torch 2.12 made torch._dynamo/_inductor config overrides thread-local
    (ContextVars), so `config.recompile_limit = 1024` set on the main thread is
    invisible to the autograd worker threads that run backward. Gradient
    checkpointing recompiles fullgraph gpt-oss kernels there against the default
    limit of 8, raising FailOnRecompileLimitHit at step 0. Mirror direct config
    assignments into the process-global entry default (torch <= 2.11 semantics).
    config.patch(...) and config.load_config(...) also assign via __setattr__ but
    are thread-local by design, so skip mirroring while inside one (tracked per
    thread). No-op below torch 2.12 and on any torch without this internal layout.
    """
    try:
        import torch

        if Version(torch.__version__) < Version("2.12.0"):
            return
        import torch._dynamo.config as _dynamo_config
        from torch.utils._config_module import ConfigModule
        from contextvars import ContextVar
    except Exception:
        return

    try:
        probe = getattr(_dynamo_config, "_config", {}).get("recompile_limit", None)
        if probe is None or not isinstance(getattr(probe, "user_override", None), ContextVar):
            # Overrides are not context-local on this torch; nothing to fix.
            return
        original_setattr = ConfigModule.__setattr__
        if getattr(original_setattr, "__unsloth_patched__", False):
            return
    except Exception:
        return

    mirrored_modules = ("torch._dynamo.config", "torch._inductor.config")

    # config.patch(...) and config.load_config(...) also assign via __setattr__, but
    # their writes are thread-local by design; a per-thread depth counter marks them
    # so they are not mirrored into the process-global default.
    import threading

    _scoped_depth = threading.local()

    def _in_scoped_write():
        return getattr(_scoped_depth, "n", 0) > 0

    def _bump(delta):
        _scoped_depth.n = getattr(_scoped_depth, "n", 0) + delta

    original_patch = ConfigModule.patch
    if not getattr(original_patch, "__unsloth_patched__", False):

        @functools.wraps(original_patch)
        def _patched_patch(self, *args, **kwargs):
            ctx = original_patch(self, *args, **kwargs)
            try:
                cls = type(ctx)  # patch() builds a fresh ConfigPatch class each call
                if not getattr(cls, "__unsloth_patch_wrapped__", False):
                    _enter0, _exit0 = cls.__enter__, cls.__exit__

                    def _enter(s, _e = _enter0):
                        _bump(1)
                        try:
                            return _e(s)
                        finally:
                            _bump(-1)

                    def _exit(
                        s,
                        *a,
                        _x = _exit0,
                    ):
                        _bump(1)
                        try:
                            return _x(s, *a)
                        finally:
                            _bump(-1)

                    cls.__enter__, cls.__exit__ = _enter, _exit
                    cls.__unsloth_patch_wrapped__ = True
            except Exception:
                pass
            return ctx

        _patched_patch.__unsloth_patched__ = True
        ConfigModule.patch = _patched_patch

    # load_config restores a saved config by calling setattr per key (thread-local).
    original_load_config = getattr(ConfigModule, "load_config", None)
    if callable(original_load_config) and not getattr(
        original_load_config, "__unsloth_patched__", False
    ):

        @functools.wraps(original_load_config)
        def _patched_load_config(self, *args, **kwargs):
            _bump(1)
            try:
                return original_load_config(self, *args, **kwargs)
            finally:
                _bump(-1)

        _patched_load_config.__unsloth_patched__ = True
        ConfigModule.load_config = _patched_load_config

    @functools.wraps(original_setattr)
    def _patched_setattr(self, name, value):
        original_setattr(self, name, value)
        if _in_scoped_write():
            return  # transient patch / load_config write: keep it thread-local
        # Aliases (cache_size_limit -> recompile_limit) re-enter with the real name.
        if self.__dict__.get("__name__", None) in mirrored_modules:
            try:
                entry = self.__dict__["_config"].get(name, None)
                if entry is not None and entry.alias is None:
                    entry.default = value
            except Exception:
                pass

    _patched_setattr.__unsloth_patched__ = True
    ConfigModule.__setattr__ = _patched_setattr

    # No replay of existing overrides: unsloth installs this before it sets any
    # dynamo/inductor config, so the wrapper mirrors every later assignment. Replaying
    # would also bake a still-active config.patch override into the global default.
    logger.info(
        "Unsloth: Patched torch config modules so dynamo/inductor settings "
        "(e.g. recompile_limit) apply across threads on torch >= 2.12."
    )


def patch_trunc_normal_precision_issue():
    """
    Patch torch.nn.init.trunc_normal_ for low precision tensors to run init in fp32.

    torch.nn.init.trunc_normal_ can saturate at truncation bounds in fp16/bf16 on
    some versions/backends. This was observed in TorchTitan investigations where
    low-precision truncation produced boundary-heavy initialization behavior:
    https://github.com/pytorch/torchtitan/pull/2342

    To avoid that failure mode, initialize into a temporary fp32 tensor, then copy
    back to the original dtype.
    """
    try:
        import torch
    except (ImportError, ModuleNotFoundError):
        return

    if getattr(torch.nn.init, "_unsloth_trunc_normal_patched", False):
        return

    original_trunc_normal = torch.nn.init.trunc_normal_
    if getattr(original_trunc_normal, "__unsloth_trunc_normal_patched__", False):
        torch.nn.init._unsloth_trunc_normal_patched = True
        return

    low_precision_dtypes = {torch.float16, torch.bfloat16}

    def _call_original(target, mean, std, a, b, generator):
        if generator is None:
            return original_trunc_normal(target, mean = mean, std = std, a = a, b = b)
        try:
            return original_trunc_normal(target, mean = mean, std = std, a = a, b = b, generator = generator)
        except TypeError as exc:
            # Older torch versions may not accept a generator keyword argument.
            msg = str(exc).lower()
            if "unexpected keyword argument" in msg and "generator" in msg:
                return original_trunc_normal(target, mean = mean, std = std, a = a, b = b)
            raise

    try:
        from torch.distributed._tensor import DTensor
    except Exception:
        DTensor = None

    @torch.no_grad()
    def _patched_trunc_normal_(
        tensor,
        mean: float = 0.0,
        std: float = 1.0,
        a: float = -2.0,
        b: float = 2.0,
        generator = None,
    ):
        if DTensor is not None and isinstance(tensor, DTensor):
            local_tensor = getattr(tensor, "_local_tensor", None)
            if local_tensor is None:
                return _call_original(tensor, mean, std, a, b, generator)
            if local_tensor.dtype in low_precision_dtypes:
                local_fp32 = local_tensor.float()
                _call_original(local_fp32, mean, std, a, b, generator)
                local_tensor.copy_(local_fp32.to(dtype = local_tensor.dtype))
                return tensor
            return _call_original(tensor, mean, std, a, b, generator)

        if tensor.dtype in low_precision_dtypes:
            tensor_fp32 = tensor.float()
            _call_original(tensor_fp32, mean, std, a, b, generator)
            tensor.copy_(tensor_fp32.to(dtype = tensor.dtype))
            return tensor

        return _call_original(tensor, mean, std, a, b, generator)

    _patched_trunc_normal_.__unsloth_trunc_normal_patched__ = True
    _patched_trunc_normal_._unsloth_original = original_trunc_normal
    torch.nn.init._unsloth_trunc_normal_original = original_trunc_normal
    torch.nn.init.trunc_normal_ = _patched_trunc_normal_
    torch.nn.init._unsloth_trunc_normal_patched = True
    logger.info("Unsloth: Patched torch.nn.init.trunc_normal_ for fp16/bf16 stability.")


def check_vllm_torch_sm100_compatibility():
    """
    Check for incompatible vLLM + torch < 2.9.0 + SM100 (Blackwell) combination.

    vLLM's distributed module (device_communicators) crashes with std::bad_alloc
    when imported on SM100 GPUs (B200/B100) with torch < 2.9.0. This is due to
    C++ code in vLLM's NCCL/distributed layer being incompatible with older
    torch versions on the newer Blackwell architecture.

    This check runs early (before vLLM import) to provide a helpful error message
    instead of a cryptic std::bad_alloc crash.
    """
    # vLLM installed? (without importing it)
    if importlib.util.find_spec("vllm") is None:
        return

    try:
        torch_version = Version(importlib_version("torch"))
        if torch_version >= Version("2.9.0"):
            return  # torch >= 2.9.0 is compatible
    except Exception:
        return  # Can't determine torch version, skip check

    # Any SM100 (Blackwell) GPU?
    try:
        import torch

        if not torch.cuda.is_available():
            return

        has_sm100 = False
        sm100_gpu_name = None
        for i in range(torch.cuda.device_count()):
            major, minor = torch.cuda.get_device_capability(i)
            if major == 10:
                has_sm100 = True
                sm100_gpu_name = torch.cuda.get_device_name(i)
                break

        if not has_sm100:
            return
    except Exception:
        return

    try:
        vllm_version = importlib_version("vllm")
    except Exception:
        vllm_version = "unknown"

    # Incompatible combination: raise a helpful error
    raise RuntimeError(
        f"Unsloth: Incompatible configuration detected.\n\n"
        f"  GPU: {sm100_gpu_name} (SM100 / Blackwell architecture)\n"
        f"  torch version: {torch_version}\n"
        f"  vLLM version: {vllm_version}\n\n"
        f"vLLM's distributed module crashes with std::bad_alloc on SM100 GPUs "
        f"(B200/B100/Blackwell) when using torch < 2.9.0.\n\n"
        f"To fix this, please upgrade torch:\n"
        f"  pip install --upgrade torch>=2.9.0\n\n"
        f"Alternatively, if you don't need vLLM:\n"
        f"  pip uninstall vllm"
    )


def fix_vllm_pdl_blackwell():
    """
    Fix vLLM PDL (Programmatic Dependent Launch) bug on Blackwell GPUs (SM100).

    The issue: vLLM's LoRA Triton kernels use tl.extra.cuda.gdc_wait() for PDL
    optimization on SM90+ GPUs. This fails on SM100 (B200/B100) during CUDA graph
    capture because Triton's pipeliner can't handle gdc_wait in complex kernels.

    See: https://github.com/vllm-project/vllm/issues/30872
    """
    if importlib.util.find_spec("vllm") is None:
        return

    # Any SM100 (Blackwell) GPU? Fix applies globally via env var + monkey-patch.
    try:
        import torch

        if not torch.cuda.is_available():
            return

        has_sm100 = False
        sm100_gpu_name = None
        for i in range(torch.cuda.device_count()):
            major, minor = torch.cuda.get_device_capability(i)
            if major == 10:
                has_sm100 = True
                sm100_gpu_name = torch.cuda.get_device_name(i)
                break

        if not has_sm100:
            return
    except Exception:
        return

    def _spec_exists(name):
        try:
            return importlib.util.find_spec(name) is not None
        except (ImportError, OSError, ModuleNotFoundError, ValueError):
            return False

    # PDL-related modules present?
    has_utils = _spec_exists("vllm.lora.ops.triton_ops.utils")
    has_expand_op = _spec_exists("vllm.lora.ops.triton_ops.lora_expand_op")
    has_shrink_op = _spec_exists("vllm.lora.ops.triton_ops.lora_shrink_op")

    if not has_utils and not has_expand_op and not has_shrink_op:
        # Old vLLM version without PDL support - nothing to patch
        return

    # vLLM version already includes the fix?
    VLLM_PDL_FIX_VERSION = "0.15.0"
    try:
        vllm_version = Version(importlib_version("vllm"))
        if vllm_version >= Version(VLLM_PDL_FIX_VERSION):
            logger.info(
                f"Unsloth: SM100 ({sm100_gpu_name}) detected but vLLM {vllm_version} "
                f"should include PDL fix - skipping workaround"
            )
            return
    except Exception as e:
        logger.debug(f"Unsloth: vLLM version check failed ({e}), applying PDL workaround.")

    # Apply the PDL fix
    os.environ["TRITON_DISABLE_PDL"] = "1"

    def fake_supports_pdl(*args, **kwargs):
        return False

    patched = []
    patched_names = set()

    def _record_patch(name):
        if name not in patched_names:
            patched.append(name)
            patched_names.add(name)

    # Patch the source module (utils.py) where supports_pdl is defined. It uses
    # @lru_cache, so clear the cache to avoid stale results.
    try:
        utils_module = importlib.import_module("vllm.lora.ops.triton_ops.utils")
        if hasattr(utils_module, "supports_pdl"):
            original_fn = utils_module.supports_pdl
            if hasattr(original_fn, "cache_clear"):
                original_fn.cache_clear()
            utils_module.supports_pdl = fake_supports_pdl
            _record_patch("utils")
    except (ImportError, ModuleNotFoundError, AttributeError):
        pass

    # Also patch consumer modules that imported supports_pdl before this ran.
    consumer_modules = {
        "lora_expand_op": "vllm.lora.ops.triton_ops.lora_expand_op",
        "lora_shrink_op": "vllm.lora.ops.triton_ops.lora_shrink_op",
        "fused_moe_lora_op": "vllm.lora.ops.triton_ops.fused_moe_lora_op",
    }
    for name, path in consumer_modules.items():
        try:
            module = importlib.import_module(path)
            if hasattr(module, "supports_pdl"):
                module.supports_pdl = fake_supports_pdl
                _record_patch(name)
        except (ImportError, ModuleNotFoundError, AttributeError):
            pass

    # Patch any additional already-loaded triton ops consumers that expose supports_pdl.
    for module_name, module in tuple(sys.modules.items()):
        if not module_name.startswith("vllm.lora.ops.triton_ops."):
            continue
        if module is None or not hasattr(module, "supports_pdl"):
            continue
        module.supports_pdl = fake_supports_pdl
        _record_patch(module_name.rsplit(".", 1)[-1])

    if patched:
        logger.info(
            f"Unsloth: Applied PDL fix for SM100 ({sm100_gpu_name}) - patched: {', '.join(patched)}"
        )
    else:
        # Just set the env var - vLLM might be an older version without supports_pdl
        logger.info(f"Unsloth: Set TRITON_DISABLE_PDL=1 for SM100 ({sm100_gpu_name})")


def patch_openspiel_env_async():
    """Apply nest_asyncio for OpenEnv EnvClient async compatibility.

    OpenEnv's EnvClient uses async methods (reset/step). In Jupyter notebooks
    these work via top-level await, but converted scripts need
    asyncio.get_event_loop().run_until_complete() wrappers. Applying nest_asyncio
    ensures nested event loop calls work in all contexts without replacing the
    original async methods (which would break scripts that already have their own
    sync wrappers).
    """
    try:
        import inspect
        from openenv.core.env_client import EnvClient

        if not inspect.iscoroutinefunction(EnvClient.reset):
            return  # Already sync, nothing to do

        try:
            import nest_asyncio
            nest_asyncio.apply()
            logger.info("Unsloth: Applied nest_asyncio for OpenEnv EnvClient async compatibility")
        except ImportError:
            logger.info(
                "Unsloth: nest_asyncio not installed, OpenEnv async methods may need manual wrapping"
            )
    except (ImportError, AttributeError):
        pass  # openenv not installed


def patch_torchcodec_audio_decoder():
    """Call unsloth_zoo's AudioDecoder patch."""
    try:
        from unsloth_zoo.dataset_utils import patch_torchcodec_audio_decoder as _patch
        _patch()
    except (ImportError, AttributeError, RuntimeError):
        pass


def disable_torchcodec_if_broken():
    """Make broken torchcodec behave as if uninstalled (#5446).

    transformers and datasets both detect torchcodec via find_spec, which
    returns True even when the native libs cannot dlopen. We flip their
    flags and seat a sys.modules sentinel so downstream imports fall through
    their existing except ImportError handlers cleanly.
    """
    try:
        import importlib.util
        if importlib.util.find_spec("torchcodec") is None:
            return  # absent or already disabled

        # RuntimeError on dlopen failure; OSError covers chained libavutil.so misses.
        from torchcodec.decoders import AudioDecoder
    except (ImportError, RuntimeError, OSError):
        # transformers: flip flag (<5) and/or rebind lru_cache'd func (>=5).
        try:
            import transformers.utils.import_utils as tf_import_utils

            try:
                tf_import_utils._torchcodec_available = False
            except AttributeError:
                pass

            is_avail = getattr(tf_import_utils, "is_torchcodec_available", None)
            if is_avail is not None:
                try:
                    is_avail.cache_clear()
                except AttributeError:
                    pass
                tf_import_utils.is_torchcodec_available = lambda: False
        except ImportError:
            pass

        # datasets >= 4.0: own flag gating audio/video/features/formatters.
        try:
            import datasets.config as datasets_config
            if hasattr(datasets_config, "TORCHCODEC_AVAILABLE"):
                datasets_config.TORCHCODEC_AVAILABLE = False
        except ImportError:
            pass

        # Drop half-loaded entries and seat the absence sentinel. After this,
        # import torchcodec raises ModuleNotFoundError and find_spec returns None.
        for _stale in [
            n
            for n in list(sys.modules)
            if n == "torchcodec"
            or n.startswith("torchcodec.")
            or n == "datasets.features._torchcodec"
        ]:
            sys.modules.pop(_stale, None)
        sys.modules["torchcodec"] = None


def disable_broken_wandb():
    """Disable wandb if it's installed but cannot actually import.

    wandb can fail to import when there's a protobuf version mismatch
    (e.g., wandb < 0.19.11 with protobuf >= 6.0). This causes cascading
    import failures through trl -> transformers/accelerate -> wandb that
    crash unsloth's import chain.

    There are two separate is_wandb_available() functions used by trl:
      - transformers.integrations.integration_utils.is_wandb_available
        (used by most trl trainers)
      - accelerate.utils.imports.is_wandb_available
        (used by trl/trainer/callbacks.py)

    Both must be patched to fully prevent broken wandb imports.
    """
    if importlib.util.find_spec("wandb") is None:
        return  # wandb not installed, nothing to do

    try:
        import wandb
    except Exception:
        # wandb is installed but broken - patch all checkers to skip it
        logger.info(
            "Unsloth: wandb is installed but broken (likely a protobuf version mismatch). "
            "Disabling wandb to prevent import errors. To fix, run: pip install --upgrade wandb"
        )
        _wandb_false = lambda: False
        # Patch transformers' is_wandb_available (used by most trl trainers)
        try:
            import transformers.integrations.integration_utils as tf_integration
            tf_integration.is_wandb_available = _wandb_false
        except (ImportError, AttributeError):
            pass
        # Patch accelerate's is_wandb_available. Patch both the source module and
        # the re-export namespace, since `from accelerate.utils import
        # is_wandb_available` reads accelerate.utils, not accelerate.utils.imports.
        try:
            import accelerate.utils.imports as acc_imports
            acc_imports.is_wandb_available = _wandb_false
        except (ImportError, AttributeError):
            pass
        try:
            import accelerate.utils as acc_utils
            acc_utils.is_wandb_available = _wandb_false
        except (ImportError, AttributeError):
            pass
        # Set env var as additional fallback
        os.environ["WANDB_DISABLED"] = "true"


# ---------------------------------------------------------------------------
# peft 0.19.x + transformers 4.x drift
# ---------------------------------------------------------------------------
# peft 0.19.x's ``peft/utils/transformers_weight_conversion.py`` unconditionally
# imports ``transformers.conversion_mapping`` and ``transformers.core_model_loading``
# at module top. Neither submodule exists on transformers <5, so the import
# explodes with ModuleNotFoundError -- silently swallowed by the bare except
# in ``patch_peft_weight_converter_compatibility`` below. Fix: when (and only
# when) the import is broken, stub the two missing submodules with the symbols
# peft pulls at module top. The stubs are inert at runtime because peft itself
# only calls into them behind ``if is_transformers_ge_v5:`` gates.
# ---------------------------------------------------------------------------

# Stamped on stub modules so a second call is a strict no-op and so third
# parties can introspect ``__unsloth_stub__`` to detect our patch.
_UNSLOTH_STUB_SENTINEL = "__unsloth_stub__"
_PEFT_TENSOR_PARALLEL_FALLBACK_SYMBOLS = (
    "ALL_PARALLEL_STYLES",
    "ColwiseParallel",
    "EmbeddingParallel",
    "RowwiseParallel",
)


def _extract_peft_tensor_parallel_imported_symbols():
    """Return names PEFT imports from ``transformers.integrations.tensor_parallel``.

    Parsed from ``peft.utils.save_and_load._maybe_shard_state_dict_for_tp`` to
    avoid a stale hard-coded symbol list.
    """
    try:
        import peft.utils.save_and_load as _save_and_load
    except Exception:
        return ()
    try:
        sharding_fn = _save_and_load._maybe_shard_state_dict_for_tp
    except AttributeError:
        return ()

    try:
        source = inspect.getsource(sharding_fn)
    except Exception as exc:
        logger.debug("Failed to inspect PEFT tensor-parallel imports: %r", exc)
        return _PEFT_TENSOR_PARALLEL_FALLBACK_SYMBOLS

    import_pattern = re.compile(
        r"from\s+transformers\.integrations\.tensor_parallel\s+import\s*\((.*?)\)",
        re.S,
    )
    import_pattern_single = re.compile(
        r"from\s+transformers\.integrations\.tensor_parallel\s+import\s+([A-Za-z_][A-Za-z0-9_\s,]*)",
        re.S,
    )
    matches = import_pattern.findall(source)
    if not matches:
        matches = import_pattern_single.findall(source)

    symbols = []
    seen = set()
    for match in matches:
        pieces = re.split(r"[,\n]", match)
        for piece in pieces:
            candidate = piece.strip()
            if not candidate:
                continue
            if candidate.endswith(")"):
                candidate = candidate[:-1].strip()
            if not candidate.isidentifier():
                continue
            if candidate in seen:
                continue
            symbols.append(candidate)
            seen.add(candidate)
    return tuple(symbols) or _PEFT_TENSOR_PARALLEL_FALLBACK_SYMBOLS


def _raise_on_peft_tensor_parallel_symbol_use(symbol_name):
    raise NotImplementedError(
        f"Unsloth: cannot use unsupported "
        f"`transformers.integrations.tensor_parallel.{symbol_name}` on this "
        f"transformers installation. Please upgrade transformers before "
        f"using PEFT tensor-parallel adapter sharding features."
    )


def fix_peft_transformers_tensor_parallel_import_compat():
    """Add placeholders to ``transformers.integrations.tensor_parallel`` for symbols
    PEFT expects but this transformers build omits, keeping existing objects.

    Returns ``True`` when patched, ``False`` when no patch is needed, ``None``
    when transformers / PEFT context is absent.
    """
    try:
        tensor_parallel_spec = importlib.util.find_spec("transformers.integrations.tensor_parallel")
    except ModuleNotFoundError:
        return None
    if tensor_parallel_spec is None:
        return None

    required_symbols = _extract_peft_tensor_parallel_imported_symbols()
    if not required_symbols:
        return None

    try:
        tp_mod = importlib.import_module("transformers.integrations.tensor_parallel")
    except ModuleNotFoundError as exc:
        if exc.name not in {
            "transformers",
            "transformers.integrations",
            "transformers.integrations.tensor_parallel",
        }:
            raise
        return None
    missing = [symbol for symbol in required_symbols if not hasattr(tp_mod, symbol)]
    if not missing:
        return False

    def _install_symbol_placeholder(symbol_name):
        if symbol_name == "ALL_PARALLEL_STYLES":

            class _UnslothTensorParallelStyles(dict):
                def __getitem__(self, key):
                    _raise_on_peft_tensor_parallel_symbol_use(symbol_name)

                def get(self, *args, **kwargs):
                    _raise_on_peft_tensor_parallel_symbol_use(symbol_name)

                def __contains__(self, key):
                    _raise_on_peft_tensor_parallel_symbol_use(symbol_name)

                def __iter__(self):
                    _raise_on_peft_tensor_parallel_symbol_use(symbol_name)

                def __len__(self):
                    _raise_on_peft_tensor_parallel_symbol_use(symbol_name)

            value = _UnslothTensorParallelStyles()
        else:

            class _UnslothTensorParallelPlaceholder:
                def __init__(self, *args, **kwargs):
                    _raise_on_peft_tensor_parallel_symbol_use(symbol_name)

            value = _UnslothTensorParallelPlaceholder
            value.__name__ = f"UnslothTensorParallelPlaceholder{symbol_name}"

        setattr(value, _UNSLOTH_STUB_SENTINEL, True)
        setattr(tp_mod, symbol_name, value)

    for symbol in missing:
        _install_symbol_placeholder(symbol)

    return True


def _peft_stub_module_importable(name):
    """True iff ``import {name}`` would succeed without side effects."""
    if name in sys.modules and sys.modules[name] is not None:
        return True
    try:
        return importlib.util.find_spec(name) is not None
    except (ImportError, ValueError, ModuleNotFoundError):
        return False


def _make_peft_stub_module(fullname):
    import types as _types

    mod = _types.ModuleType(fullname)
    mod.__file__ = f"<unsloth stub: {fullname}>"
    mod.__package__ = fullname.rpartition(".")[0]
    setattr(mod, _UNSLOTH_STUB_SENTINEL, True)
    return mod


def _install_transformers_conversion_mapping_stub():
    """Stub the 3 symbols peft 0.19.x imports from this module at top level."""
    name = "transformers.conversion_mapping"
    existing = sys.modules.get(name)
    if existing is not None and getattr(existing, _UNSLOTH_STUB_SENTINEL, False):
        return existing

    mod = _make_peft_stub_module(name)

    # peft does ``.copy()`` + keyed assignment at module top; real dict suffices.
    mod._MODEL_TO_CONVERSION_PATTERN = {}

    def get_checkpoint_conversion_mapping(model_type, *args, **kwargs):
        # ``None`` = peft's "no conversion registered"; both callsites
        # early-return on it.
        return None

    def get_model_conversion_mapping(model, *args, **kwargs):
        return None

    mod.get_checkpoint_conversion_mapping = get_checkpoint_conversion_mapping
    mod.get_model_conversion_mapping = get_model_conversion_mapping

    sys.modules[name] = mod
    # Attach to parent so attribute-style access matches a real submodule.
    parent = sys.modules.get("transformers")
    if parent is not None and not hasattr(parent, "conversion_mapping"):
        try:
            parent.conversion_mapping = mod
        except Exception:
            # Frozen parent: sys.modules entry is enough for ``from ... import``.
            pass
    return mod


def _install_transformers_core_model_loading_stub():
    """Stub the 8 symbols peft 0.19.x imports from this module at top level.

    ``Concatenate`` and ``ConversionOps`` MUST be real classes (peft
    subclasses them at module top); the rest only appear in runtime
    ``isinstance`` / construction calls gated behind ``is_transformers_ge_v5``."""
    name = "transformers.core_model_loading"
    existing = sys.modules.get(name)
    if existing is not None and getattr(existing, _UNSLOTH_STUB_SENTINEL, False):
        return existing

    mod = _make_peft_stub_module(name)

    class ConversionOps:
        def convert(self, *args, **kwargs):  # pragma: no cover - inert stub
            raise NotImplementedError(
                "unsloth stub: transformers.core_model_loading.ConversionOps "
                "is a no-op on transformers <5. Upgrade transformers to v5+ "
                "to use peft.utils.transformers_weight_conversion at runtime."
            )

        @property
        def reverse_op(self):  # pragma: no cover - inert stub
            raise NotImplementedError

    class Concatenate(ConversionOps):
        def __init__(
            self,
            dim = 0,
            *args,
            **kwargs,
        ):
            self.dim = dim

    class MergeModulelist(ConversionOps):
        def __init__(self, *args, **kwargs):
            pass

    class Transpose(ConversionOps):
        def __init__(
            self,
            dim0 = 0,
            dim1 = 1,
            *args,
            **kwargs,
        ):
            self.dim0 = dim0
            self.dim1 = dim1

    class WeightConverter:
        def __init__(self, *args, **kwargs):
            # Accept any signature; upstream class evolves.
            self.args = args
            self.kwargs = kwargs

    class WeightRenaming:
        def __init__(
            self,
            source_patterns = None,
            target_patterns = None,
            *args,
            **kwargs,
        ):
            self.source_patterns = source_patterns
            self.target_patterns = target_patterns

    def dot_natural_key(key):
        return key

    def rename_source_key(original_key, renamings, converters):
        return original_key, None

    mod.ConversionOps = ConversionOps
    mod.Concatenate = Concatenate
    mod.MergeModulelist = MergeModulelist
    mod.Transpose = Transpose
    mod.WeightConverter = WeightConverter
    mod.WeightRenaming = WeightRenaming
    mod.dot_natural_key = dot_natural_key
    mod.rename_source_key = rename_source_key

    sys.modules[name] = mod
    parent = sys.modules.get("transformers")
    if parent is not None and not hasattr(parent, "core_model_loading"):
        try:
            parent.core_model_loading = mod
        except Exception:
            pass
    return mod


def fix_peft_transformers_weight_conversion_import():
    """Make ``from peft.utils import transformers_weight_conversion`` import
    cleanly on (peft 0.19.x, transformers 4.x) by stubbing the two missing
    transformers-v5 submodules. See header block above for details.

    Must run BEFORE ``patch_peft_weight_converter_compatibility`` -- that
    function's bare ``except (ImportError, AttributeError): return`` would
    otherwise silently no-op.

    No-op if peft / transformers missing, or if the peft module already
    imports cleanly. Idempotent and strictly additive (never overwrites a
    real ``transformers.conversion_mapping`` / ``core_model_loading``).

    Returns True if patched, False if no action needed, None if peft absent."""
    if importlib.util.find_spec("peft") is None:
        return None

    # Already importable? Either we patched, or transformers is v5+.
    try:
        importlib.import_module("peft.utils.transformers_weight_conversion")
        return False
    except ModuleNotFoundError as exc:
        # Only act on our specific drift class.
        missing = getattr(exc, "name", "") or ""
        if missing not in (
            "transformers.conversion_mapping",
            "transformers.core_model_loading",
        ):
            return False
    except ImportError as exc:
        # Older Python ImportError has no `.name`; string-match instead.
        msg = str(exc)
        if (
            "transformers.conversion_mapping" not in msg
            and "transformers.core_model_loading" not in msg
        ):
            return False

    # Need transformers loaded to attach stubs to its package.
    transformers_root = sys.modules.get("transformers")
    if transformers_root is None:
        try:
            transformers_root = importlib.import_module("transformers")
        except Exception:
            return False

    # Stub only the genuinely missing submodules; never clobber real ones.
    patched_any = False
    if not _peft_stub_module_importable("transformers.conversion_mapping"):
        _install_transformers_conversion_mapping_stub()
        patched_any = True

    if not _peft_stub_module_importable("transformers.core_model_loading"):
        _install_transformers_core_model_loading_stub()
        patched_any = True

    if not patched_any:
        # Real submodules present; failure was for some other reason.
        return False

    # Force a fresh import now that stubs are in place. Drop any cached
    # ``None`` entry first so importlib retries.
    pkg = "peft.utils.transformers_weight_conversion"
    if pkg in sys.modules and sys.modules[pkg] is None:
        del sys.modules[pkg]
    try:
        importlib.import_module(pkg)
    except Exception:
        # Other upstream drift; stubs stay installed so a later retry succeeds.
        return True

    logger.info(
        "Unsloth: stubbed transformers.conversion_mapping / "
        "transformers.core_model_loading so peft.utils."
        "transformers_weight_conversion imports cleanly on "
        "transformers <5."
    )
    return True


def patch_peft_weight_converter_compatibility():
    """Allow PEFT converter rebuilds on legacy converter constructors."""
    try:
        from peft.utils import transformers_weight_conversion as twc
    except (ImportError, AttributeError):
        return

    _patch_peft_moe_target_conversion(twc)

    if getattr(twc, "_unsloth_weight_converter_compat_patch", False):
        return

    import threading

    original_build = twc.build_peft_weight_mapping
    patch_lock = threading.RLock()

    def _patch_weight_converter_ctors(weight_conversions, patched):
        seen_classes = set()

        for conversion in weight_conversions:
            conversion_cls = conversion.__class__
            if conversion_cls in seen_classes:
                continue
            seen_classes.add(conversion_cls)

            original_init = conversion_cls.__init__
            params = inspect.signature(original_init).parameters
            supports_kwargs = any(p.kind == inspect.Parameter.VAR_KEYWORD for p in params.values())
            supports_distributed = "distributed_operation" in params
            supports_quantization = "quantization_operation" in params
            if supports_kwargs or (supports_distributed and supports_quantization):
                continue

            def _compat_init(
                self,
                *args,
                __original_init = original_init,
                __supports_distributed = supports_distributed,
                __supports_quantization = supports_quantization,
                **kwargs,
            ):
                unsupported = {}
                if not __supports_distributed and "distributed_operation" in kwargs:
                    unsupported["distributed_operation"] = kwargs.pop("distributed_operation")
                if not __supports_quantization and "quantization_operation" in kwargs:
                    unsupported["quantization_operation"] = kwargs.pop("quantization_operation")
                result = __original_init(self, *args, **kwargs)
                for name, value in unsupported.items():
                    if hasattr(self, name):
                        setattr(self, name, value)
                return result

            conversion_cls.__init__ = _compat_init
            patched.append((conversion_cls, original_init))

    @functools.wraps(original_build)
    def _build_peft_weight_mapping_compat(
        weight_conversions,
        adapter_name,
        peft_config = None,
    ):
        if not weight_conversions:
            return original_build(weight_conversions, adapter_name, peft_config)

        patched_classes = []
        with patch_lock:
            try:
                _patch_weight_converter_ctors(weight_conversions, patched_classes)
                return original_build(weight_conversions, adapter_name, peft_config)
            finally:
                for conversion_cls, original_init in patched_classes:
                    conversion_cls.__init__ = original_init

    twc.build_peft_weight_mapping = _build_peft_weight_mapping_compat
    twc._unsloth_weight_converter_compat_patch = True


def _patch_peft_moe_target_conversion(twc):
    """Keep PEFT 0.19 MoE conversion from rewriting explicit Unsloth targets."""
    if getattr(twc, "_unsloth_moe_target_conversion_patch", False):
        return

    original_convert_moe = getattr(twc, "_convert_peft_config_moe", None)
    if original_convert_moe is None:
        return

    @functools.wraps(original_convert_moe)
    def _convert_peft_config_moe_unsloth(peft_config, model_type: str) -> None:
        if getattr(peft_config, "target_parameters", None):
            return

        target_modules = getattr(peft_config, "target_modules", None)
        if isinstance(target_modules, str):
            if "." in target_modules:
                return
            return original_convert_moe(peft_config, model_type)

        if not target_modules:
            return original_convert_moe(peft_config, model_type)

        explicit_targets = {
            target for target in target_modules if isinstance(target, str) and "." in target
        }
        if not explicit_targets:
            return original_convert_moe(peft_config, model_type)

        bare_targets = set(target_modules) - explicit_targets
        if not bare_targets:
            return

        peft_config.target_modules = bare_targets
        original_convert_moe(peft_config, model_type)
        peft_config.target_modules = set(peft_config.target_modules or ()) | explicit_targets

    twc._convert_peft_config_moe = _convert_peft_config_moe_unsloth
    twc._unsloth_moe_target_conversion_patch = True


CAUSAL_CONV1D_BROKEN = False
_CAUSAL_CONV1D_PREFIX = "causal_conv1d"
_CAUSAL_CONV1D_BLOCKER_SENTINEL = "_unsloth_causal_conv1d_blocker"
VLLM_BROKEN = False
_VLLM_PREFIX = "vllm"
_VLLM_BLOCKER_SENTINEL = "_unsloth_vllm_blocker"
_ROCM_ENV_HINT_KEYS = (
    "ROCM_PATH",
    "ROCM_HOME",
    "HIP_PATH",
    "HSA_PATH",
    "HIP_VISIBLE_DEVICES",
    "ROCR_VISIBLE_DEVICES",
)
_ROCM_PATH_HINTS = (
    Path("/opt/rocm"),
    Path("/dev/kfd"),
    Path("/sys/module/amdgpu"),
)
_AMDGPU_ASIC_ID_TABLE_PATH_ENV = "AMDGPU_ASIC_ID_TABLE_PATH"
_AMDGPU_ASIC_ID_CANDIDATE_PATHS = (
    Path("/usr/share/libdrm/amdgpu.ids"),
    Path("/usr/local/share/libdrm/amdgpu.ids"),
    Path("/opt/rocm/share/libdrm/amdgpu.ids"),
    Path("/opt/amdgpu/share/libdrm/amdgpu.ids"),
)


def _log_rocm_detection(message):
    if UNSLOTH_ENABLE_LOGGING:
        logger.info(message)


@functools.lru_cache(1)
def _is_rocm_torch_build() -> bool:
    # Most official ROCm wheels include a local version suffix like +rocmX.Y.
    # Some custom/source builds do not, so we fall back to runtime hints.
    try:
        torch_version_raw = str(importlib_version("torch")).lower()
        if "rocm" in torch_version_raw:
            _log_rocm_detection("Unsloth: ROCm detection matched torch version tag (+rocm).")
            return True
    except Exception:
        pass

    # Environment hints commonly present on ROCm runtimes.
    for key in _ROCM_ENV_HINT_KEYS:
        value = os.environ.get(key, "")
        if isinstance(value, str) and value.strip():
            _log_rocm_detection(f"Unsloth: ROCm detection matched environment key `{key}`.")
            return True

    # Filesystem / driver hints for ROCm stacks.
    for path in _ROCM_PATH_HINTS:
        try:
            if path.exists():
                _log_rocm_detection(f"Unsloth: ROCm detection matched filesystem hint `{path}`.")
                return True
        except Exception:
            continue

    _log_rocm_detection("Unsloth: ROCm detection did not match any known hints.")
    return False


def _iter_amdgpu_asic_id_table_candidates():
    # Try torch-adjacent ids table paths first without importing torch.
    try:
        torch_spec = importlib.util.find_spec("torch")
    except Exception:
        torch_spec = None

    roots = []
    if torch_spec is not None:
        if torch_spec.origin:
            roots.append(Path(torch_spec.origin).resolve().parent)
        if torch_spec.submodule_search_locations:
            for location in torch_spec.submodule_search_locations:
                roots.append(Path(location).resolve())

    seen = set()
    for root in roots:
        for candidate in (
            root / "share" / "libdrm" / "amdgpu.ids",
            root.parent / "share" / "libdrm" / "amdgpu.ids",
            root.parent.parent / "share" / "libdrm" / "amdgpu.ids",
        ):
            candidate_str = str(candidate)
            if candidate_str in seen:
                continue
            seen.add(candidate_str)
            yield candidate

    for candidate in _AMDGPU_ASIC_ID_CANDIDATE_PATHS:
        candidate_str = str(candidate)
        if candidate_str in seen:
            continue
        seen.add(candidate_str)
        yield candidate


def configure_amdgpu_asic_id_table_path():
    # Honor an existing valid user-provided path.
    configured = os.environ.get(_AMDGPU_ASIC_ID_TABLE_PATH_ENV, "").strip()
    if configured:
        configured_path = Path(configured)
        try:
            if configured_path.is_file():
                return str(configured_path)
        except Exception:
            pass

    # Only attempt this on ROCm-like environments.
    if not _is_rocm_torch_build():
        return None

    for candidate in _iter_amdgpu_asic_id_table_candidates():
        try:
            if candidate.is_file():
                os.environ[_AMDGPU_ASIC_ID_TABLE_PATH_ENV] = str(candidate)
                if UNSLOTH_ENABLE_LOGGING:
                    logger.info(f"Unsloth: Set {_AMDGPU_ASIC_ID_TABLE_PATH_ENV}={candidate}")
                return str(candidate)
        except Exception:
            continue

    return None


# ---------------------------------------------------------------------------
# bitsandbytes Windows ROCm fix: cextension.py runs get_rocm_gpu_arch()
# (bnb >= 0.47) and get_rocm_warpsize() (0.49.x) at import, shelling out to
# rocminfo / hipinfo.exe via PATH. Neither is on PATH on Windows (AMD torch
# wheels put hipInfo.exe in venv Scripts), so every import logs ERROR +
# WARNING, ROCM_GPU_ARCH becomes "unknown", and warp size defaults to 64:
# wrong on RDNA (wave 32), breaking 4-bit blocksizes and
# ALLOW_PREQUANTIZED_MODELS. Upstream fix unmerged (bitsandbytes#1969), so a
# MetaPathFinder swaps both helpers for torch-device-props-first versions
# right after bitsandbytes.cuda_specs executes, before cextension reads
# them. Must run before `import unsloth_zoo` (imports bnb on ROCm).
# ---------------------------------------------------------------------------

_BNB_CUDA_SPECS_MODULE = "bitsandbytes.cuda_specs"
_BNB_ROCM_FIX_FINDER_SENTINEL = "_unsloth_bnb_rocm_fix_finder"
_BNB_ROCM_FIX_FUNCTION_FLAG = "__unsloth_bnb_rocm_fix__"


def _torch_rocm_device_props():
    """Device-0 props on a ROCm torch build with a visible GPU, else None.
    Never raises; bnb's own import initializes the device context anyway."""
    try:
        import torch

        if not getattr(getattr(torch, "version", None), "hip", None):
            return None
        if not torch.cuda.is_available():
            return None
        return torch.cuda.get_device_properties(0)
    except Exception:
        return None


def _iter_hipinfo_paths():
    """Yield existing hipInfo.exe paths: PATH, interpreter scripts dir (venv
    and conda layouts), then HIP SDK / AMD installer locations."""
    import shutil
    import sysconfig

    candidates = []
    try:
        resolved = shutil.which("hipinfo.exe")
        if resolved:
            candidates.append(resolved)
    except Exception:
        pass
    try:
        scripts_dir = sysconfig.get_path("scripts")
        if scripts_dir:
            candidates.append(os.path.join(scripts_dir, "hipInfo.exe"))
    except Exception:
        pass
    executable_dir = os.path.dirname(sys.executable or "")
    if executable_dir:
        candidates.append(os.path.join(executable_dir, "hipInfo.exe"))
        candidates.append(os.path.join(executable_dir, "Scripts", "hipInfo.exe"))
    for env_key in ("HIP_PATH", "ROCM_PATH"):
        root = os.environ.get(env_key, "").strip()
        if root:
            candidates.append(os.path.join(root, "bin", "hipInfo.exe"))
    rocm_root = os.path.join(os.environ.get("ProgramFiles", r"C:\Program Files"), "AMD", "ROCm")
    try:
        if os.path.isdir(rocm_root):
            for version_dir in sorted(os.listdir(rocm_root), reverse = True):
                candidates.append(os.path.join(rocm_root, version_dir, "bin", "hipInfo.exe"))
    except Exception:
        pass

    seen = set()
    for candidate in candidates:
        try:
            key = os.path.normcase(os.path.normpath(candidate))
            if key in seen:
                continue
            seen.add(key)
            if os.path.isfile(candidate):
                yield candidate
        except Exception:
            continue


def _run_hipinfo(hipinfo_path):
    """Run hipInfo.exe and return its stdout, or "" on any failure."""
    import subprocess
    try:
        result = subprocess.run(
            [hipinfo_path],
            capture_output = True,
            text = True,
            timeout = 15,
            creationflags = getattr(subprocess, "CREATE_NO_WINDOW", 0),
        )
        return result.stdout or ""
    except Exception as e:
        _log_rocm_detection(f"Unsloth: `{hipinfo_path}` failed: {e}")
        return ""


def _unsloth_get_rocm_gpu_arch():
    """Replaces bnb's get_rocm_gpu_arch: torch device props first (no
    subprocess), then hipInfo.exe by absolute path, then a quiet "unknown"."""
    try:
        import torch
        if not getattr(getattr(torch, "version", None), "hip", None):
            return "unknown"
    except Exception:
        return "unknown"
    props = _torch_rocm_device_props()
    if props is not None:
        try:
            # gcnArchName may carry feature flags, e.g. "gfx90a:sramecc+:xnack-"
            arch = str(props.gcnArchName).split(":")[0].strip()
            if arch.startswith("gfx"):
                return arch
        except Exception:
            pass
    for hipinfo_path in _iter_hipinfo_paths():
        match = re.search(r"gcnArchName:\s+gfx([a-zA-Z\d]+)", _run_hipinfo(hipinfo_path))
        if match:
            return "gfx" + match.group(1)
    _log_rocm_detection(
        "Unsloth: Could not detect the ROCm GPU architecture - bitsandbytes will see `unknown`."
    )
    return "unknown"


def _unsloth_get_rocm_warpsize():
    """Replaces bnb 0.49.x get_rocm_warpsize: upstream defaults to 64 when
    rocminfo is missing, wrong on RDNA (wave 32)."""
    try:
        import torch
        if not getattr(getattr(torch, "version", None), "hip", None):
            return 32  # upstream behavior: NVIDIA warp size is always 32
    except Exception:
        return 64  # upstream behavior: default to 64 on failure
    props = _torch_rocm_device_props()
    if props is not None:
        # torch 2.11 ROCm exposes warp_size; some builds used warpSize.
        for attribute_name in ("warp_size", "warpSize"):
            warp_size = getattr(props, attribute_name, None)
            if isinstance(warp_size, int) and warp_size in (32, 64):
                return warp_size
    for hipinfo_path in _iter_hipinfo_paths():
        match = re.search(r"^\s*warpSize:\s+(\d+)", _run_hipinfo(hipinfo_path), re.MULTILINE)
        if match and int(match.group(1)) in (32, 64):
            return int(match.group(1))
    _log_rocm_detection(
        "Unsloth: Could not detect the ROCm warp size - defaulting to 64 "
        "(bitsandbytes' own default)."
    )
    return 64


setattr(_unsloth_get_rocm_gpu_arch, _BNB_ROCM_FIX_FUNCTION_FLAG, True)
setattr(_unsloth_get_rocm_warpsize, _BNB_ROCM_FIX_FUNCTION_FLAG, True)


def _bnb_rocm_helper_is_broken(function):
    """True only for upstream's subprocess-only detectors; co_names works
    where getsource fails. Versions consulting torch props are untouched."""
    if function is None or not callable(function):
        return False
    if getattr(function, _BNB_ROCM_FIX_FUNCTION_FLAG, False):
        return False  # Already ours.
    try:
        function = inspect.unwrap(function)
    except Exception:
        pass
    code = getattr(function, "__code__", None)
    co_names = getattr(code, "co_names", ()) if code is not None else ()
    if not co_names:
        return False  # C function or opaque wrapper -- do not touch.
    if "get_device_properties" in co_names or "gcnArchName" in co_names:
        return False  # Fixed upstream -- no-op.
    return "subprocess" in co_names


def _patch_bnb_cuda_specs_module(module):
    """Swap broken ROCm detection helpers on an executed cuda_specs module.
    Returns True when the module ends up patched (now or previously)."""
    patched = False
    for attribute_name, replacement in (
        ("get_rocm_gpu_arch", _unsloth_get_rocm_gpu_arch),
        ("get_rocm_warpsize", _unsloth_get_rocm_warpsize),
    ):
        original = getattr(module, attribute_name, None)
        if getattr(original, _BNB_ROCM_FIX_FUNCTION_FLAG, False):
            patched = True  # Already ours.
            continue
        if not _bnb_rocm_helper_is_broken(original):
            continue
        setattr(module, attribute_name, replacement)
        patched = True
        logger.info(
            f"Unsloth: Patched bitsandbytes.cuda_specs.{attribute_name} - "
            f"avoids PATH-dependent subprocess GPU detection on Windows ROCm."
        )
    return patched


class _BnbCudaSpecsPatchLoader(importlib.abc.Loader):
    __slots__ = ("_loader",)

    def __init__(self, loader):
        self._loader = loader

    def create_module(self, spec):
        create_module = getattr(self._loader, "create_module", None)
        if create_module is None:
            return None
        return create_module(spec)

    def exec_module(self, module):
        self._loader.exec_module(module)
        # Patch after the module body ran, before cextension calls it. The
        # finder stays on sys.meta_path (same lifecycle as the blockers
        # above) so importlib.reload(bitsandbytes.cuda_specs) re-patches.
        try:
            _patch_bnb_cuda_specs_module(module)
        except Exception as e:
            _log_rocm_detection(f"Unsloth: bitsandbytes ROCm detection patch failed: {e}")

    def __getattr__(self, name):
        # Delegate get_source / get_filename etc. so introspection works.
        return getattr(self._loader, name)


class _BnbCudaSpecsPatchFinder(importlib.abc.MetaPathFinder):
    __slots__ = (_BNB_ROCM_FIX_FINDER_SENTINEL,)

    def __init__(self):
        setattr(self, _BNB_ROCM_FIX_FINDER_SENTINEL, True)

    def find_spec(
        self,
        fullname,
        path = None,
        target = None,
    ):
        if fullname != _BNB_CUDA_SPECS_MODULE:
            return None
        # Delegate to remaining finders (editable installs, frozen apps)
        # and wrap the loader that would actually be used.
        spec = None
        for finder in sys.meta_path:
            if finder is self or getattr(finder, _BNB_ROCM_FIX_FINDER_SENTINEL, False):
                continue
            finder_find_spec = getattr(finder, "find_spec", None)
            if finder_find_spec is None:
                continue
            try:
                spec = finder_find_spec(fullname, path, target)
            except Exception:
                spec = None
            if spec is not None:
                break
        if spec is None or spec.loader is None:
            return None
        if not hasattr(spec.loader, "exec_module"):
            return None  # Legacy loader -- let the stock machinery handle it.
        spec.loader = _BnbCudaSpecsPatchLoader(spec.loader)
        return spec


def _repair_imported_bitsandbytes_rocm_constants():
    """bnb imported before unsloth: noise already fired, but fix detectors
    and cached constants, incl. by-value ROCM_WARP_SIZE_64 copies."""
    cuda_specs = sys.modules.get(_BNB_CUDA_SPECS_MODULE)
    if cuda_specs is None:
        return
    if not _patch_bnb_cuda_specs_module(cuda_specs):
        return

    try:
        arch = cuda_specs.get_rocm_gpu_arch()
    except Exception:
        arch = "unknown"
    warp_size_64 = None
    get_rocm_warpsize = getattr(cuda_specs, "get_rocm_warpsize", None)
    if callable(get_rocm_warpsize):
        try:
            warp_size_64 = get_rocm_warpsize() == 64
        except Exception:
            warp_size_64 = None

    for module_name, module in list(sys.modules.items()):
        if module is None or module is cuda_specs:
            continue
        if module_name != "bitsandbytes" and not module_name.startswith("bitsandbytes."):
            continue
        try:
            if arch != "unknown" and getattr(module, "ROCM_GPU_ARCH", None) == "unknown":
                module.ROCM_GPU_ARCH = arch
            if warp_size_64 is not None and isinstance(
                getattr(module, "ROCM_WARP_SIZE_64", None), bool
            ):
                module.ROCM_WARP_SIZE_64 = warp_size_64
        except Exception:
            continue
    logger.info("Unsloth: Repaired bitsandbytes ROCm arch / warp-size constants in place.")


def fix_bitsandbytes_rocm_arch_detection():
    """Fix bnb's import-time ROCm arch / warp-size detection on Windows
    (see header above). No-op on non-Windows, non-ROCm, missing or
    upstream-fixed bnb. Idempotent. Opt out: UNSLOTH_DISABLE_BNB_ROCM_FIX=1."""
    if os.environ.get("UNSLOTH_DISABLE_BNB_ROCM_FIX", "0") == "1":
        return
    if sys.platform != "win32":
        return
    if not _is_rocm_torch_build():
        return

    # Already imported: prevention impossible, repair in place instead.
    if _BNB_CUDA_SPECS_MODULE in sys.modules:
        try:
            _repair_imported_bitsandbytes_rocm_constants()
        except Exception:
            pass
        return

    try:
        if importlib.util.find_spec("bitsandbytes") is None:
            return
    except Exception:
        return

    for finder in sys.meta_path:
        if getattr(finder, _BNB_ROCM_FIX_FINDER_SENTINEL, False):
            return  # Already installed -- idempotent.
    sys.meta_path.insert(0, _BnbCudaSpecsPatchFinder())
    _log_rocm_detection("Unsloth: Installed the bitsandbytes ROCm arch detection patch hook.")


def _is_causal_conv1d_name(module_name: str) -> bool:
    return module_name == _CAUSAL_CONV1D_PREFIX or module_name.startswith(
        _CAUSAL_CONV1D_PREFIX + "."
    )


def _is_vllm_name(module_name: str) -> bool:
    return module_name == _VLLM_PREFIX or module_name.startswith(_VLLM_PREFIX + ".")


def _resolve_module_name(module_name, package):
    if not isinstance(module_name, str):
        return module_name
    if module_name.startswith("."):
        try:
            return importlib.util.resolve_name(module_name, package)
        except Exception:
            return module_name
    return module_name


def _is_broken_causal_conv1d_error(error) -> bool:
    checked = set()
    current = error
    while current is not None and id(current) not in checked:
        checked.add(id(current))
        message = str(current).lower()
        if (
            ("causal_conv1d_cuda" in message and "undefined symbol" in message)
            or ("_zn3c103hip28c10_hip_check_implementation" in message)
            or ("causal_conv1d" in message and "undefined symbol" in message)
        ):
            return True
        current = getattr(current, "__cause__", None) or getattr(current, "__context__", None)
    return False


def _is_broken_vllm_error(error) -> bool:
    checked = set()
    current = error
    while current is not None and id(current) not in checked:
        checked.add(id(current))
        message = str(current).lower()
        if (
            ("vllm/_c" in message or "vllm._c" in message)
            and (
                "undefined symbol" in message
                or "cannot open shared object file" in message
                or ".so:" in message
            )
        ) or ("vllm" in message and "undefined symbol" in message):
            return True
        # Forced extension load raises the bare loader error (no "vllm._C"
        # wrapper); match any .so failure as callers feed only vLLM imports.
        if "cannot open shared object file" in message:
            return True
        current = getattr(current, "__cause__", None) or getattr(current, "__context__", None)
    return False


def _get_vllm_cuda_mismatch_message(error):
    """If the error is a CUDA version mismatch, return a helpful install message."""
    import re as _re

    checked = set()
    current = error
    wanted_cuda = None
    while current is not None and id(current) not in checked:
        checked.add(id(current))
        message = str(current)
        # Extract the CUDA version vllm was built for, e.g. "libcudart.so.12"
        match = _re.search(r"libcudart\.so\.(\d+)", message)
        if match:
            wanted_cuda = match.group(1)
            break
        current = getattr(current, "__cause__", None) or getattr(current, "__context__", None)
    if wanted_cuda is None:
        return None

    # Detect what CUDA version is actually available on the system
    system_cuda_display = None  # Human-readable, e.g. "13.0"
    system_cuda_tag = None  # For wheel URL, e.g. "130"
    try:
        import torch
        cuda_version = torch.version.cuda  # e.g. "13.0" or "12.8"
        if cuda_version:
            system_cuda_display = cuda_version
            system_cuda_tag = cuda_version.replace(".", "")[:3]  # "130" or "128"
    except Exception:
        pass

    if system_cuda_tag is None or system_cuda_tag.startswith(wanted_cuda):
        return None  # Not a mismatch or can't determine

    try:
        vllm_version = importlib_version("vllm").split("+")[0]
    except Exception:
        vllm_version = "VLLM_VERSION"

    cpu_arch = "x86_64"
    try:
        import platform
        cpu_arch = platform.machine()
    except Exception:
        pass

    return (
        f"Unsloth: vLLM was built for CUDA {wanted_cuda} but this system has "
        f"CUDA {system_cuda_display}. Please reinstall vLLM with the correct CUDA version:\n"
        f"\n"
        f"  uv pip install https://github.com/vllm-project/vllm/releases/download/"
        f"v{vllm_version}/vllm-{vllm_version}+cu{system_cuda_tag}-cp38-abi3-"
        f"manylinux_2_35_{cpu_arch}.whl"
    )


class _CausalConv1dImportBlockerLoader(importlib.abc.Loader):
    __slots__ = ("module_name",)

    def __init__(self, module_name):
        self.module_name = module_name

    def create_module(self, spec):
        return None

    def exec_module(self, module):
        raise ModuleNotFoundError(f"No module named '{self.module_name}'")


class _CausalConv1dImportBlockerFinder(importlib.abc.MetaPathFinder):
    __slots__ = (_CAUSAL_CONV1D_BLOCKER_SENTINEL,)

    def __init__(self):
        setattr(self, _CAUSAL_CONV1D_BLOCKER_SENTINEL, True)

    def find_spec(
        self,
        fullname,
        path = None,
        target = None,
    ):
        if not CAUSAL_CONV1D_BROKEN or not _is_causal_conv1d_name(fullname):
            return None
        return importlib.machinery.ModuleSpec(
            name = fullname,
            loader = _CausalConv1dImportBlockerLoader(fullname),
            is_package = fullname == _CAUSAL_CONV1D_PREFIX,
        )


class _VllmImportBlockerLoader(importlib.abc.Loader):
    __slots__ = ("module_name",)

    def __init__(self, module_name):
        self.module_name = module_name

    def create_module(self, spec):
        return None

    def exec_module(self, module):
        raise ModuleNotFoundError(f"No module named '{self.module_name}'")


class _VllmImportBlockerFinder(importlib.abc.MetaPathFinder):
    __slots__ = (_VLLM_BLOCKER_SENTINEL,)

    def __init__(self):
        setattr(self, _VLLM_BLOCKER_SENTINEL, True)

    def find_spec(
        self,
        fullname,
        path = None,
        target = None,
    ):
        if not VLLM_BROKEN or not _is_vllm_name(fullname):
            return None
        return importlib.machinery.ModuleSpec(
            name = fullname,
            loader = _VllmImportBlockerLoader(fullname),
            is_package = fullname == _VLLM_PREFIX,
        )


def _patch_find_spec_for_causal_conv1d():
    current_find_spec = importlib.util.find_spec
    if getattr(current_find_spec, "_unsloth_causal_conv1d_find_spec_patch", False):
        return

    def _blocked_find_spec(name, package = None):
        resolved_name = _resolve_module_name(name, package)
        if CAUSAL_CONV1D_BROKEN and isinstance(resolved_name, str):
            if _is_causal_conv1d_name(resolved_name):
                return None
        return current_find_spec(name, package)

    _blocked_find_spec._unsloth_causal_conv1d_find_spec_patch = True
    _blocked_find_spec._unsloth_original_find_spec = current_find_spec
    importlib.util.find_spec = _blocked_find_spec


def _patch_find_spec_for_vllm():
    current_find_spec = importlib.util.find_spec
    if getattr(current_find_spec, "_unsloth_vllm_find_spec_patch", False):
        return

    def _blocked_find_spec(name, package = None):
        resolved_name = _resolve_module_name(name, package)
        if VLLM_BROKEN and isinstance(resolved_name, str):
            if _is_vllm_name(resolved_name):
                return None
        return current_find_spec(name, package)

    _blocked_find_spec._unsloth_vllm_find_spec_patch = True
    _blocked_find_spec._unsloth_original_find_spec = current_find_spec
    importlib.util.find_spec = _blocked_find_spec


def _install_causal_conv1d_blocker():
    _patch_find_spec_for_causal_conv1d()
    for finder in sys.meta_path:
        if getattr(finder, _CAUSAL_CONV1D_BLOCKER_SENTINEL, False):
            return
    sys.meta_path.insert(0, _CausalConv1dImportBlockerFinder())


def _install_vllm_blocker():
    _patch_find_spec_for_vllm()
    for finder in sys.meta_path:
        if getattr(finder, _VLLM_BLOCKER_SENTINEL, False):
            return
    sys.meta_path.insert(0, _VllmImportBlockerFinder())


def _clear_causal_conv1d_modules():
    for module_name in list(sys.modules):
        if _is_causal_conv1d_name(module_name):
            sys.modules.pop(module_name, None)


def _clear_vllm_modules():
    for module_name in list(sys.modules):
        if _is_vllm_name(module_name):
            sys.modules.pop(module_name, None)


# vLLM's compiled extensions. A CUDA-major ABI break hits all of them, so
# probing the eagerly-loaded _C and its siblings reliably trips it.
_VLLM_COMPILED_EXTENSIONS = (
    "vllm._C",
    "vllm._C_stable_libtorch",
    "vllm._moe_C",
    "vllm._rocm_C",
)


def disable_broken_vllm(error = None):
    """Disable vLLM dynamically when its shared library is ABI-broken."""
    global VLLM_BROKEN
    if VLLM_BROKEN:
        _install_vllm_blocker()
        return True

    failure = error
    if failure is None:
        try:
            if importlib.util.find_spec("vllm") is None:
                return False
        except Exception:
            return False

        try:
            import vllm  # noqa: F401

            # Lazy vLLM lets a bare `import vllm` succeed even when an extension
            # is ABI-broken; force-load each to surface the .so failure here.
            # A missing one raises ModuleNotFoundError (skipped below).
            for _ext in _VLLM_COMPILED_EXTENSIONS:
                try:
                    importlib.import_module(_ext)
                except ModuleNotFoundError:
                    pass
            return False
        except Exception as import_error:
            failure = import_error

    if not _is_broken_vllm_error(failure):
        return False

    VLLM_BROKEN = True
    _clear_vllm_modules()
    _install_vllm_blocker()
    cuda_msg = _get_vllm_cuda_mismatch_message(failure)
    if cuda_msg:
        logger.warning(cuda_msg)
    else:
        logger.warning(
            "Unsloth: Detected broken vLLM binary extension; "
            "disabling vLLM imports and continuing import.\n"
            "Please reinstall via `uv pip install unsloth vllm torchvision torchaudio "
            "--torch-backend=auto`."
        )
    return True


def _disable_transformers_causal_conv1d():
    try:
        import transformers.utils.import_utils as tf_import_utils
    except Exception:
        return

    if hasattr(tf_import_utils, "is_causal_conv1d_available"):
        tf_import_utils.is_causal_conv1d_available = lambda: False

    for attr_name in (
        "_causal_conv1d_available",
        "_is_causal_conv1d_available",
    ):
        if hasattr(tf_import_utils, attr_name):
            setattr(tf_import_utils, attr_name, False)


def disable_broken_causal_conv1d():
    """Disable causal_conv1d dynamically when its shared library is ABI-broken.

    This mirrors Unsloth's FlashAttention fallback behavior: if importing causal_conv1d
    fails with a known binary symbol error, we disable it at startup so model imports do
    not hard-fail.
    """
    global CAUSAL_CONV1D_BROKEN
    if CAUSAL_CONV1D_BROKEN:
        _install_causal_conv1d_blocker()
        _disable_transformers_causal_conv1d()
        return

    try:
        if importlib.util.find_spec("causal_conv1d") is None:
            return
    except Exception:
        return

    try:
        import causal_conv1d  # noqa: F401
        return
    except Exception as error:
        if not _is_broken_causal_conv1d_error(error):
            return

    CAUSAL_CONV1D_BROKEN = True
    _clear_causal_conv1d_modules()
    _install_causal_conv1d_blocker()
    _disable_transformers_causal_conv1d()
    print(
        "Unsloth: Detected broken causal_conv1d binary; "
        "disabling causal_conv1d fast path and continuing import."
    )


_BNB_ROCM_DLL_RE = re.compile(r"libbitsandbytes_rocm(\d+)\.dll", re.IGNORECASE)


def _is_hip_torch_build():
    """True only when torch itself is a HIP/ROCm build. Env hints (HIP_PATH
    etc.) do not count: CUDA bitsandbytes raises at import when the ROCm
    override is set. Wheel tag first (no torch import); torch.version.hip
    fallback for source builds."""
    try:
        if "rocm" in str(importlib_version("torch")).lower():
            return True
    except Exception:
        pass
    try:
        import torch
        return bool(getattr(torch.version, "hip", None))
    except Exception:
        return False


def _detect_installed_bnb_rocm_version():
    """Highest installed ``libbitsandbytes_rocm<NN>.dll`` suffix ("72", "713")
    or ``None``. Listing order is unordered, so take the numeric max."""
    try:
        spec = importlib.util.find_spec("bitsandbytes")
    except Exception:
        return None
    if spec is None or not spec.submodule_search_locations:
        return None

    suffixes = []
    for pkg_dir in spec.submodule_search_locations:
        try:
            entries = os.listdir(pkg_dir)
        except Exception:
            continue
        for entry in entries:
            match = _BNB_ROCM_DLL_RE.fullmatch(entry)
            if match is not None:
                suffixes.append(match.group(1))
    if not suffixes:
        return None
    return max(suffixes, key = lambda value: int(value))


def maybe_set_windows_rocm_bnb_version():
    """Pin ``BNB_ROCM_VERSION`` from the installed wheel on Windows + ROCm torch.

    AMD's Windows wheel ships one ``libbitsandbytes_rocm<NN>.dll`` whose
    suffix can disagree with ``torch.version.hip`` (HIP 7.13 vs rocm72.dll),
    breaking the native 4-bit/8-bit paths. Pin the installed suffix before
    bitsandbytes is first imported.

    No-op unless ALL of: Windows, a real HIP torch build (env hints like
    HIP_PATH do not count), a ROCm DLL installed, and no explicit user value.
    Linux is untouched. Values seeded by Studio's venv sitecustomize.py
    (marked ``UNSLOTH_BNB_ROCM_VERSION_SOURCE=sitecustomize``) are
    redetectable defaults, not overrides; ``UNSLOTH_SKIP_BNB_ROCM_VERSION=1``
    opts out and drops a seeded default. Returns the value set, else None.
    """
    if sys.platform != "win32":
        return None
    if os.environ.get("UNSLOTH_SKIP_BNB_ROCM_VERSION") == "1":
        # Real opt-out: drop our seeded default (marker present); explicit
        # user values carry no marker and are kept.
        if os.environ.get("UNSLOTH_BNB_ROCM_VERSION_SOURCE") == "sitecustomize":
            os.environ.pop("BNB_ROCM_VERSION", None)
            os.environ.pop("UNSLOTH_BNB_ROCM_VERSION_SOURCE", None)
        return None
    if "BNB_ROCM_VERSION" in os.environ and (
        os.environ.get("UNSLOTH_BNB_ROCM_VERSION_SOURCE") != "sitecustomize"
    ):
        return None
    if not _is_hip_torch_build():
        return None
    version = _detect_installed_bnb_rocm_version()
    if version is None:
        return None
    os.environ["BNB_ROCM_VERSION"] = version
    os.environ["UNSLOTH_BNB_ROCM_VERSION_SOURCE"] = "detected"
    if UNSLOTH_ENABLE_LOGGING:
        logger.info(
            f"Unsloth: set BNB_ROCM_VERSION={version} "
            "(detected from the installed bitsandbytes ROCm wheel on Windows)."
        )
    return version


def patch_accelerate_recursively_apply():
    """
    Make Accelerate's recursive utilities tolerate Unsloth's EmptyLogits
    sentinel. recursively_apply returns the sentinel unchanged instead of
    raising TypeError, and find_device skips it while still finding real
    tensors, falling back to PartialState().device only for sentinel-only
    payloads. Both wrappers are idempotent and are propagated to every
    already imported accelerate namespace.
    """
    try:
        import accelerate.utils.operations as acc_ops
    except Exception:
        return

    original_recursively_apply = getattr(acc_ops, "recursively_apply", None)
    if original_recursively_apply is not None and not getattr(
        original_recursively_apply, "__unsloth_patched__", False
    ):

        @functools.wraps(original_recursively_apply)
        def _patched_recursively_apply(func, data, *args, **kwargs):
            if type(data).__name__ == "EmptyLogits":
                cls = type(data)
                if cls.__eq__ is object.__eq__:
                    # Debug mode compares gathered metadata across ranks with ==
                    cls.__eq__ = lambda self, other: type(other).__name__ == "EmptyLogits"
                return data
            return original_recursively_apply(func, data, *args, **kwargs)

        _patched_recursively_apply.__unsloth_patched__ = True

        for mod_name, mod in tuple(sys.modules.items()):
            if mod_name.startswith("accelerate") and mod is not None:
                if getattr(mod, "recursively_apply", None) is original_recursively_apply:
                    try:
                        setattr(mod, "recursively_apply", _patched_recursively_apply)
                    except Exception:
                        pass

    original_find_device = getattr(acc_ops, "find_device", None)
    if original_find_device is not None and not getattr(
        original_find_device, "__unsloth_patched__", False
    ):
        from collections.abc import Mapping

        @functools.wraps(original_find_device)
        def _patched_find_device(data):
            import torch

            found_sentinel = False

            def _search(obj):
                nonlocal found_sentinel
                if type(obj).__name__ == "EmptyLogits":
                    found_sentinel = True
                elif isinstance(obj, Mapping):
                    for value in obj.values():
                        device = _search(value)
                        if device is not None:
                            return device
                elif isinstance(obj, (tuple, list)):
                    for value in obj:
                        device = _search(value)
                        if device is not None:
                            return device
                elif isinstance(obj, torch.Tensor):
                    return obj.device
                return None

            device = _search(data)
            if device is None and found_sentinel:
                # Debug mode calls find_device(...).type on gather/broadcast inputs
                try:
                    from accelerate.state import PartialState
                    return PartialState().device
                except Exception:
                    pass
            return device

        _patched_find_device.__unsloth_patched__ = True

        for mod_name, mod in tuple(sys.modules.items()):
            if mod_name.startswith("accelerate") and mod is not None:
                if getattr(mod, "find_device", None) is original_find_device:
                    try:
                        setattr(mod, "find_device", _patched_find_device)
                    except Exception:
                        pass