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:
parent
8603233829
commit
ab18542b7b
1 changed files with 19 additions and 1 deletions
|
|
@ -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
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue