diff --git a/unsloth/models/vision.py b/unsloth/models/vision.py index b4d5ccff23..57b0d02c3e 100644 --- a/unsloth/models/vision.py +++ b/unsloth/models/vision.py @@ -188,7 +188,14 @@ def _attach_bnb_multidevice_hooks( for p in model.parameters() if hasattr(p, "device") and p.device.type == "cuda" } - except Exception: + except Exception as exc: + import warnings + warnings.warn( + "Unsloth: Failed to determine CUDA devices from model parameters, " + f"so multi-GPU hooks cannot be attached. ({type(exc).__name__}: {exc})", + RuntimeWarning, + stacklevel = 2, + ) return if not cuda_devs: @@ -201,6 +208,7 @@ def _attach_bnb_multidevice_hooks( try: from accelerate import dispatch_model + from accelerate.hooks import AlignDevicesHook, add_hook_to_module from accelerate.utils import find_tied_parameters, retie_parameters except ImportError: return # accelerate not available @@ -235,6 +243,16 @@ def _attach_bnb_multidevice_hooks( # dispatch_model installs AlignDevicesHook on every module and # sub-module, exactly matching HF's own loading behavior. dispatch_model(model, device_map = device_map_int) + + # dispatch_model skips hook installation for single-device maps + # (accelerate fast-path). When all weights sit on a non-default + # GPU (e.g. cuda:1) the caller's inputs may still arrive on + # cuda:0, causing a device mismatch. Detect this and manually + # add a root-level AlignDevicesHook to route inputs. + if len(cuda_devs) == 1 and not hasattr(model, "_hf_hook"): + main_device = next(iter(cuda_devs)) + add_hook_to_module(model, AlignDevicesHook(execution_device = main_device)) + desc = f"{len(inferred_map)} block(s) across {len(cuda_devs)} device(s)" finally: # Restore stripped keys