patch: fix EmptyLogits gathering in nested payloads and Accelerate recursively_apply (#6092)

* Fix EmptyLogits gathering in nested structure and patch recursively_apply on accelerator module

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

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

* Wire EmptyLogits Accelerate patch into startup and fix find_device, pickling, tests for PR #6092

- Call patch_accelerate_recursively_apply() in _gpu_init.py so real imports
  install it; previously it was only invoked by the tests
- Make both wrappers idempotent so repeated calls do not stack
- Rework find_device: skip EmptyLogits while still finding real tensors in any
  order, keep returning None for tensor-free payloads (AlignDevicesHook relies
  on None), fall back to PartialState().device only for sentinel-only payloads
- Give EmptyLogits stateless __reduce__ and drop the stomped pickle stubs on
  EMPTY_LOGITS so debug mode gather_object works in real distributed runs
- Put test tensors on PartialState().device so the debug mode test also passes
  on GPU machines, and add drift tests for startup wiring, idempotency and
  find_device ordering

Verified on 2x B200: ACCELERATE_DEBUG_MODE=1 torchrun gather/broadcast/pad of
sentinel and mixed payloads all pass, training losses unchanged, full drift
suite 25/25.

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

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

* Define EmptyLogits equality on the class for PR #6092

Gathered sentinel copies must compare equal in accelerate debug mode
regardless of whether the patched recursively_apply saw the sentinel first
in that process. Class body __eq__ requires restoring __hash__ explicitly.
Verified: 123 case simulation battery on accelerate 0.34.2 through latest,
2 process gloo CPU and NCCL GPU debug mode runs, drift suite 25/25.

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Daniel Han <danielhanchen@gmail.com>

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
 # ===========================================================================

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":

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

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):