fix: keep LoRA reloads working with PEFT 0.19 (#6748)

* fix: keep LoRA reloads working with PEFT 0.19

* test: exercise the PEFT tensor-parallel symbol extractor

* test: prove the full PEFT tensor-parallel seam

* fix: harden PEFT tensor-parallel shims

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* fix: fall back when PEFT tensor-parallel source inspection fails

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Lee Jackson <130007945+Imagineer99@users.noreply.github.com>
Co-authored-by: imagineer99 <samleejackson0@gmail.com>

tests/test_import_fixes_drift.py

@@ -23,6 +23,7 @@ import inspect
 import os
 import re
 import sys
+from pathlib import Path
 from importlib.metadata import version as importlib_version
 
 import pytest
@@ -593,6 +594,7 @@ def test_accelerate_gather_empty_logits_debug_mode_patch():
     orig_debug = state.debug
     orig_dist_type = state.distributed_type
     orig_num_processes = state.num_processes
+    orig_device = state.device
 
     state.debug = True
     state.distributed_type = DistributedType.MULTI_GPU
@@ -624,6 +626,8 @@ def test_accelerate_gather_empty_logits_debug_mode_patch():
                 side_effect = mock_gpu_broadcast,
             ),
         ):
+            state.device = torch.device("cpu")
+
             # Top-level EmptyLogits gathers to itself
             res = acc_ops.gather(e)
             assert res is e
@@ -656,6 +660,7 @@ def test_accelerate_gather_empty_logits_debug_mode_patch():
         state.debug = orig_debug
         state.distributed_type = orig_dist_type
         state.num_processes = orig_num_processes
+        state.device = orig_device
 
 
 def test_accelerate_patch_is_idempotent():
@@ -698,10 +703,8 @@ def test_accelerate_find_device_skips_empty_logits():
 
 def test_accelerate_patch_wired_into_gpu_init():
     """The patch must be installed at startup, not only importable."""
-    import pathlib
+    source = Path(__file__).resolve().parent.parent / "unsloth" / "_gpu_init.py"
-    import unsloth.import_fixes as import_fixes
+    source = source.read_text()
-
-    source = pathlib.Path(import_fixes.__file__).with_name("_gpu_init.py").read_text()
     assert "patch_accelerate_recursively_apply()" in source, (
         "DRIFT DETECTED: patch_accelerate_recursively_apply is defined but "
         "never called in _gpu_init.py, so real imports never install it."

tests/test_peft_tensor_parallel_compat.py

@@ -0,0 +1,252 @@
+import importlib.machinery
+import importlib.util
+from pathlib import Path
+import sys
+import types
+
+import pytest
+
+
+REPO_ROOT = Path(__file__).resolve().parents[1]
+IMPORT_FIXES = REPO_ROOT / "unsloth" / "import_fixes.py"
+
+
+def _load_import_fixes():
+    spec = importlib.util.spec_from_file_location("_unsloth_import_fixes_peft_tp", IMPORT_FIXES)
+    module = importlib.util.module_from_spec(spec)
+    spec.loader.exec_module(module)
+    return module
+
+
+def _install_fake_module(
+    name,
+    *,
+    is_package = False,
+    attrs = None,
+):
+    module = types.ModuleType(name)
+    module.__spec__ = importlib.machinery.ModuleSpec(name, loader = None, is_package = is_package)
+    if is_package:
+        module.__path__ = []
+    if attrs:
+        for k, v in attrs.items():
+            setattr(module, k, v)
+    sys.modules[name] = module
+    return module
+
+
+@pytest.fixture(autouse = True)
+def _restore_import_fixtures():
+    keep = {
+        "transformers",
+        "transformers.integrations",
+        "transformers.integrations.tensor_parallel",
+        "peft",
+        "peft.utils",
+        "peft.utils.save_and_load",
+    }
+    snapshot = {name: sys.modules.get(name) for name in keep}
+    yield
+    for name, value in snapshot.items():
+        if value is None:
+            sys.modules.pop(name, None)
+        else:
+            sys.modules[name] = value
+
+
+def _install_fake_transformers_tensor_parallel(existing):
+    transformers = _install_fake_module("transformers", is_package = True)
+    integrations = _install_fake_module("transformers.integrations", is_package = True)
+    setattr(transformers, "integrations", integrations)
+
+    tp = _install_fake_module(
+        "transformers.integrations.tensor_parallel",
+        attrs = existing,
+    )
+    setattr(integrations, "tensor_parallel", tp)
+    return tp
+
+
+def _fake_peft_shard_state_dict_for_tp():
+    from transformers.integrations.tensor_parallel import (
+        ALL_PARALLEL_STYLES,
+        ColwiseParallel,
+        EmbeddingParallel,
+        RowwiseParallel,
+    )
+    return (
+        ALL_PARALLEL_STYLES,
+        ColwiseParallel,
+        EmbeddingParallel,
+        RowwiseParallel,
+    )
+
+
+def _install_fake_peft_tensor_parallel_import():
+    peft = _install_fake_module("peft", is_package = True)
+    utils = _install_fake_module("peft.utils", is_package = True)
+    setattr(peft, "utils", utils)
+
+    save_and_load = _install_fake_module("peft.utils.save_and_load")
+    save_and_load._maybe_shard_state_dict_for_tp = _fake_peft_shard_state_dict_for_tp
+    setattr(utils, "save_and_load", save_and_load)
+    return save_and_load
+
+
+def test_missing_tensor_parallel_symbol_import_succeeds_after_fix(monkeypatch):
+    module = _load_import_fixes()
+
+    _install_fake_peft_tensor_parallel_import()
+    _install_fake_transformers_tensor_parallel(
+        {
+            "ColwiseParallel": object,
+            "RowwiseParallel": object,
+        }
+    )
+
+    with pytest.raises(ImportError):
+        _fake_peft_shard_state_dict_for_tp()
+
+    assert module.fix_peft_transformers_tensor_parallel_import_compat() is True
+
+    import transformers.integrations.tensor_parallel as patched
+
+    all_parallel_styles, _, embedding_parallel, _ = _fake_peft_shard_state_dict_for_tp()
+
+    assert patched == sys.modules["transformers.integrations.tensor_parallel"]
+    assert embedding_parallel is getattr(patched, "EmbeddingParallel")
+    assert getattr(embedding_parallel, "__unsloth_stub__", False)
+    with pytest.raises(NotImplementedError, match = "ALL_PARALLEL_STYLES"):
+        all_parallel_styles["rowwise"]
+    with pytest.raises(NotImplementedError, match = "ALL_PARALLEL_STYLES"):
+        "rowwise" in all_parallel_styles
+    with pytest.raises(NotImplementedError, match = "ALL_PARALLEL_STYLES"):
+        iter(all_parallel_styles)
+    with pytest.raises(NotImplementedError, match = "ALL_PARALLEL_STYLES"):
+        len(all_parallel_styles)
+
+
+def test_existing_embedding_parallel_is_not_replaced(monkeypatch):
+    module = _load_import_fixes()
+
+    class RealEmbeddingParallel:
+        pass
+
+    tp_mod = _install_fake_transformers_tensor_parallel(
+        {
+            "EmbeddingParallel": RealEmbeddingParallel,
+            "ColwiseParallel": object,
+        }
+    )
+
+    monkeypatch.setattr(
+        module,
+        "_extract_peft_tensor_parallel_imported_symbols",
+        lambda: ("ALL_PARALLEL_STYLES", "ColwiseParallel", "EmbeddingParallel", "RowwiseParallel"),
+    )
+    assert module.fix_peft_transformers_tensor_parallel_import_compat() is True
+
+    assert tp_mod.EmbeddingParallel is RealEmbeddingParallel
+    assert not getattr(tp_mod.EmbeddingParallel, "__unsloth_stub__", False)
+
+
+def test_missing_tensor_parallel_module_is_not_created(monkeypatch):
+    module = _load_import_fixes()
+    previous_tp_module = sys.modules.pop("transformers.integrations.tensor_parallel", None)
+
+    original_find_spec = module.importlib.util.find_spec
+
+    def fake_find_spec(name):
+        if name == "transformers.integrations.tensor_parallel":
+            return None
+        return original_find_spec(name)
+
+    monkeypatch.setattr(module.importlib.util, "find_spec", fake_find_spec)
+    monkeypatch.setattr(
+        module,
+        "_extract_peft_tensor_parallel_imported_symbols",
+        lambda: ("EmbeddingParallel",),
+    )
+
+    assert module.fix_peft_transformers_tensor_parallel_import_compat() is None
+
+    try:
+        spec = importlib.util.find_spec("transformers.integrations.tensor_parallel")
+    except ModuleNotFoundError as exc:
+        assert exc.name == "transformers"
+    else:
+        assert spec is None
+    assert "transformers.integrations.tensor_parallel" not in sys.modules
+
+
+def test_placeholder_raises_on_real_use(monkeypatch):
+    module = _load_import_fixes()
+
+    tp_mod = _install_fake_transformers_tensor_parallel({})
+    monkeypatch.setattr(
+        module,
+        "_extract_peft_tensor_parallel_imported_symbols",
+        lambda: ("EmbeddingParallel",),
+    )
+    assert module.fix_peft_transformers_tensor_parallel_import_compat() is True
+
+    with pytest.raises(NotImplementedError, match = "EmbeddingParallel"):
+        tp_mod.EmbeddingParallel()
+
+    assert getattr(tp_mod.EmbeddingParallel, "__unsloth_stub__", False)
+
+
+def test_symbol_extractor_falls_back_when_parse_returns_no_identifiers(monkeypatch):
+    module = _load_import_fixes()
+
+    _install_fake_peft_tensor_parallel_import()
+    monkeypatch.setattr(
+        module.inspect,
+        "getsource",
+        lambda _: "from transformers.integrations.tensor_parallel import ()",
+    )
+
+    assert module._extract_peft_tensor_parallel_imported_symbols() == (
+        "ALL_PARALLEL_STYLES",
+        "ColwiseParallel",
+        "EmbeddingParallel",
+        "RowwiseParallel",
+    )
+
+
+def test_symbol_extractor_falls_back_when_getsource_fails(monkeypatch):
+    module = _load_import_fixes()
+
+    _install_fake_peft_tensor_parallel_import()
+    monkeypatch.setattr(
+        module.inspect,
+        "getsource",
+        lambda _: (_ for _ in ()).throw(ValueError("no source")),
+    )
+
+    assert module._extract_peft_tensor_parallel_imported_symbols() == (
+        "ALL_PARALLEL_STYLES",
+        "ColwiseParallel",
+        "EmbeddingParallel",
+        "RowwiseParallel",
+    )
+
+
+def test_tensor_parallel_import_module_not_found_returns_none(monkeypatch):
+    module = _load_import_fixes()
+    _install_fake_peft_tensor_parallel_import()
+
+    monkeypatch.setattr(
+        module.importlib.util,
+        "find_spec",
+        lambda name: object() if name == "transformers.integrations.tensor_parallel" else None,
+    )
+
+    def _raise_missing(name):
+        exc = ModuleNotFoundError(name)
+        exc.name = name
+        raise exc
+
+    monkeypatch.setattr(module.importlib, "import_module", _raise_missing)
+
+    assert module.fix_peft_transformers_tensor_parallel_import_compat() is None

unsloth/_gpu_init.py

@@ -188,6 +188,7 @@ from .import_fixes import (
     disable_torchcodec_if_broken,
     disable_broken_wandb,
     fix_trl_vllm_ascend,
+    fix_peft_transformers_tensor_parallel_import_compat,
     fix_peft_transformers_weight_conversion_import,
     patch_peft_weight_converter_compatibility,
     patch_accelerate_recursively_apply,
@@ -219,6 +220,7 @@ disable_broken_wandb()
 # Must run before patch_peft_weight_converter_compatibility: stubs the
 # transformers v5 submodules peft 0.19.x imports, so the next patch can wrap
 # build_peft_weight_mapping instead of being swallowed by its ImportError.
+fix_peft_transformers_tensor_parallel_import_compat()
 fix_peft_transformers_weight_conversion_import()
 patch_peft_weight_converter_compatibility()
 patch_accelerate_recursively_apply()
@@ -244,6 +246,7 @@ del patch_vllm_for_notebooks
 del patch_torchcodec_audio_decoder
 del disable_torchcodec_if_broken
 del disable_broken_wandb
+del fix_peft_transformers_tensor_parallel_import_compat
 del fix_peft_transformers_weight_conversion_import
 del patch_peft_weight_converter_compatibility
 del patch_accelerate_recursively_apply

unsloth/import_fixes.py

@@ -1492,6 +1492,146 @@ def disable_broken_wandb():
 # 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 expects from ``transformers.integrations.tensor_parallel``.
+
+    The supported PEFT import line is the one in
+    ``peft.utils.save_and_load._maybe_shard_state_dict_for_tp`` for fast
+    LoRA adapter checkpoints. Parse that source to avoid stale hard-coded
+    symbol lists.
+    """
+    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():
+    """Preserve existing ``transformers.integrations.tensor_parallel`` objects, then add
+    lightweight placeholders for symbols that PEFT expects but this transformers
+    build omits.
+
+    Returns ``True`` when patched, ``False`` when no patch is needed, and
+    ``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):