fix: handle single-device non-default GPU and warn on parameter errors

dispatch_model skips AlignDevicesHook 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.

Also replace the silent except-and-return on CUDA device collection
with a RuntimeWarning so users can debug failures.
This commit is contained in:
Daniel Han 2026-04-16 14:20:50 +00:00
commit ab18542b7b

View file

@ -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