diff --git a/tests/test_import_fixes_drift.py b/tests/test_import_fixes_drift.py index 68c20fbe53..0bee68f940 100644 --- a/tests/test_import_fixes_drift.py +++ b/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 - import unsloth.import_fixes as import_fixes - - source = pathlib.Path(import_fixes.__file__).with_name("_gpu_init.py").read_text() + source = Path(__file__).resolve().parent.parent / "unsloth" / "_gpu_init.py" + source = source.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." diff --git a/tests/test_peft_tensor_parallel_compat.py b/tests/test_peft_tensor_parallel_compat.py new file mode 100644 index 0000000000..3dc00d46dc --- /dev/null +++ b/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 diff --git a/unsloth/_gpu_init.py b/unsloth/_gpu_init.py index 7f080336aa..e39b44488a 100644 --- a/unsloth/_gpu_init.py +++ b/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 diff --git a/unsloth/import_fixes.py b/unsloth/import_fixes.py index d979e5f22f..29e0c30dc1 100644 --- a/unsloth/import_fixes.py +++ b/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):