diff --git a/tests/test_import_fixes_drift.py b/tests/test_import_fixes_drift.py index 6c75a4b886..c9f971c891 100644 --- a/tests/test_import_fixes_drift.py +++ b/tests/test_import_fixes_drift.py @@ -586,6 +586,164 @@ def test_accelerate_utils_imports_module_present(): ) +def test_accelerate_recursively_apply_empty_logits_patch(): + """Verify patch_accelerate_recursively_apply overrides recursively_apply to bypass EmptyLogits.""" + pytest.importorskip("accelerate") + + import accelerate.utils.operations as acc_ops + from unsloth.import_fixes import patch_accelerate_recursively_apply + + class EmptyLogits: + pass + + e = EmptyLogits() + patch_accelerate_recursively_apply() + + res = acc_ops.recursively_apply(lambda x: x, e, error_on_other_type = True) + assert res is e + + +def test_accelerate_gather_empty_logits_debug_mode_patch(): + """Verify gather and broadcast bypass EmptyLogits when debug mode is enabled.""" + pytest.importorskip("accelerate") + from accelerate.state import PartialState, DistributedType + import accelerate.utils.operations as acc_ops + from unsloth.import_fixes import patch_accelerate_recursively_apply + import unittest.mock as mock + import torch + + class EmptyLogits: + pass + + e = EmptyLogits() + patch_accelerate_recursively_apply() + + # Enable debug mode and mock distributed state + state = PartialState() + orig_debug = state.debug + orig_dist_type = state.distributed_type + orig_num_processes = state.num_processes + + state.debug = True + state.distributed_type = DistributedType.MULTI_GPU + state.num_processes = 2 + + # Mock gather_object to return [obj] * num_processes + def mock_gather_object(obj, *args, **kwargs): + return [obj] * state.num_processes + + # Mock _gpu_gather to recursively apply replication of tensors + def mock_gpu_gather(tensor, *args, **kwargs): + def _gather_one(t): + if t.ndim == 0: + t = t.clone()[None] + return torch.cat([t] * state.num_processes, dim = 0) + + return acc_ops.recursively_apply(_gather_one, tensor, error_on_other_type = True) + + # Mock _gpu_broadcast to return data unchanged + def mock_gpu_broadcast(data, *args, **kwargs): + return data + + try: + with ( + mock.patch( + "accelerate.utils.operations.gather_object", + side_effect = mock_gather_object, + ), + mock.patch("accelerate.utils.operations._gpu_gather", side_effect = mock_gpu_gather), + mock.patch( + "accelerate.utils.operations._gpu_broadcast", + side_effect = mock_gpu_broadcast, + ), + ): + # 1. Top-level EmptyLogits should gather correctly (returns e) + res = acc_ops.gather(e) + assert res is e + + # 2. Nested EmptyLogits alone + res_nested = acc_ops.gather([e]) + assert isinstance(res_nested, list) and res_nested[0] is e + + # 3. Mixed payload with real tensor and EmptyLogits + # Real tensor should be gathered (concatenated across processes). + # Tensors must live on state.device or the debug-mode device + # check fails on GPU machines. + real_tensor = torch.tensor([42], device = state.device) + payload = {"labels": real_tensor, "logits": e} + res_mixed = acc_ops.gather(payload) + + assert isinstance(res_mixed, dict) + assert res_mixed["logits"] is e + # Since num_processes = 2, it should be gathered to [42, 42] + assert torch.equal(res_mixed["labels"], torch.tensor([42, 42], device = state.device)) + + # 4. Broadcast with EmptyLogits + res_broadcast = acc_ops.broadcast(e) + assert res_broadcast is e + + # 5. Mixed payload with broadcast + res_broadcast_mixed = acc_ops.broadcast(payload) + assert isinstance(res_broadcast_mixed, dict) + assert res_broadcast_mixed["logits"] is e + assert torch.equal(res_broadcast_mixed["labels"], real_tensor) + finally: + state.debug = orig_debug + state.distributed_type = orig_dist_type + state.num_processes = orig_num_processes + + +def test_accelerate_patch_is_idempotent(): + """Calling patch_accelerate_recursively_apply twice must not stack wrappers.""" + pytest.importorskip("accelerate") + import accelerate.utils.operations as acc_ops + from unsloth.import_fixes import patch_accelerate_recursively_apply + + patch_accelerate_recursively_apply() + recursively_apply = acc_ops.recursively_apply + find_device = acc_ops.find_device + patch_accelerate_recursively_apply() + assert ( + acc_ops.recursively_apply is recursively_apply + ), "DRIFT DETECTED: recursively_apply was wrapped twice." + assert acc_ops.find_device is find_device, "DRIFT DETECTED: find_device was wrapped twice." + + +def test_accelerate_find_device_skips_empty_logits(): + """find_device must search past EmptyLogits and keep None for tensor-free data.""" + pytest.importorskip("accelerate") + import torch + import accelerate.utils.operations as acc_ops + from accelerate.state import PartialState + from unsloth.import_fixes import patch_accelerate_recursively_apply + + class EmptyLogits: + pass + + patch_accelerate_recursively_apply() + tensor = torch.tensor([1.0]) + # Sentinel first must not stop the search before the real tensor + assert acc_ops.find_device({"logits": EmptyLogits(), "labels": tensor}) == tensor.device + # Tensor-free payloads without the sentinel keep returning None + # (AlignDevicesHook relies on None to skip output device moves) + assert acc_ops.find_device({"a": 1}) is None + # Sentinel-only payloads fall back to the current device so that + # debug mode find_device(...).type does not raise AttributeError + assert acc_ops.find_device(EmptyLogits()) == PartialState().device + + +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() + 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." + ) + + # =========================================================================== # bitsandbytes -- ROCm arch / warp-size detection shape # =========================================================================== diff --git a/unsloth/_gpu_init.py b/unsloth/_gpu_init.py index 5da1e27a9d..86f675c281 100644 --- a/unsloth/_gpu_init.py +++ b/unsloth/_gpu_init.py @@ -189,6 +189,7 @@ from .import_fixes import ( fix_trl_vllm_ascend, fix_peft_transformers_weight_conversion_import, patch_peft_weight_converter_compatibility, + patch_accelerate_recursively_apply, ) fix_xformers_performance_issue() @@ -217,6 +218,7 @@ disable_broken_wandb() # build_peft_weight_mapping instead of being swallowed by its ImportError. fix_peft_transformers_weight_conversion_import() patch_peft_weight_converter_compatibility() +patch_accelerate_recursively_apply() del fix_xformers_performance_issue del fix_vllm_aimv2_issue @@ -240,6 +242,7 @@ del disable_torchcodec_if_broken del disable_broken_wandb del fix_peft_transformers_weight_conversion_import del patch_peft_weight_converter_compatibility +del patch_accelerate_recursively_apply # Torch 2.4 has including_emulation if DEVICE_TYPE == "cuda": diff --git a/unsloth/import_fixes.py b/unsloth/import_fixes.py index 695a6a577a..a100e11f0c 100644 --- a/unsloth/import_fixes.py +++ b/unsloth/import_fixes.py @@ -2584,3 +2584,93 @@ def maybe_set_windows_rocm_bnb_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 diff --git a/unsloth/models/_utils.py b/unsloth/models/_utils.py index 2f4e3a069e..6024a02c2c 100644 --- a/unsloth/models/_utils.py +++ b/unsloth/models/_utils.py @@ -2703,6 +2703,16 @@ class EmptyLogits: def __str__(self): return LOGITS_ERROR_STRING + def __reduce__(self): + # Stateless pickling so gather_object works on the sentinel + return (type(self), ()) + + def __eq__(self, other): + # Gathered copies must compare equal in accelerate debug mode + return type(other).__name__ == "EmptyLogits" + + __hash__ = object.__hash__ + EMPTY_LOGITS = EmptyLogits() functions = dir(torch.Tensor) @@ -2713,6 +2723,13 @@ for j, function in enumerate(functions): exec(f"EMPTY_LOGITS.{function} = raise_{j}", globals(), locals()) except: continue +# The loop above stomps pickle hooks with stubs returning None, which breaks +# gather_object on EMPTY_LOGITS in distributed runs. Restore default pickling. +for function in ("__reduce__", "__reduce_ex__", "__getstate__", "__setstate__"): + try: + delattr(EMPTY_LOGITS, function) + except Exception: + pass def validate_loftq_config(loftq_config, lora_dropout, bias, init_lora_weights, model):