code/src/mfflora/_compat.py

"""Environment compat shims.

Import this module *before* peft / transformers in any script that uses peft.
It disables the bitsandbytes dispatch path inside peft because the local bnb
install is broken (incompatible with installed Triton 3.5.1's missing
``triton.ops`` module).
"""

from __future__ import annotations

import os

# Make sure Triton has a writable cache; many transformers import paths trigger
# Triton autotune which fails on a non-writable home dir.
os.environ.setdefault("TRITON_CACHE_DIR", "/tmp/triton_cache")
os.makedirs(os.environ["TRITON_CACHE_DIR"], exist_ok=True)


def disable_bnb_in_peft() -> None:
    """Force peft to behave as if bitsandbytes is not installed.

    ``peft/__init__.py`` transitively imports ``peft.tuners.lora.model``,
    which captures ``is_bnb_available`` by name. So we have to patch *both*
    the ``import_utils`` namespace (where the function lives) and the
    ``lora.model`` namespace (where it has already been bound).
    """

    try:
        import peft.import_utils as iu

        iu.is_bnb_available = lambda: False  # type: ignore[assignment]
        iu.is_bnb_4bit_available = lambda: False  # type: ignore[assignment]
        # Patch already-imported submodules.
        for mod_name in (
            "peft.tuners.lora.model",
            "peft.tuners.lora.layer",
            "peft.tuners.lora.config",
            "peft.tuners.tuners_utils",
        ):
            mod = __import__(mod_name, fromlist=["_"])
            if hasattr(mod, "is_bnb_available"):
                mod.is_bnb_available = lambda: False  # type: ignore[assignment]
            if hasattr(mod, "is_bnb_4bit_available"):
                mod.is_bnb_4bit_available = lambda: False  # type: ignore[assignment]
    except ImportError:
        pass


disable_bnb_in_peft()