From ab18542b7b90348423ea4df7fb7f884b4508dcc1 Mon Sep 17 00:00:00 2001
From: Daniel Han <danielhanchen@gmail.com>
Date: Thu, 16 Apr 2026 14:20:50 +0000
Subject: [PATCH] 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.
---
 unsloth/models/vision.py | 20 +++++++++++++++++++-
 1 file changed, 19 insertions(+), 1 deletion(-)

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