From d6b758035af2e527232776f1140b048bfa873fd4 Mon Sep 17 00:00:00 2001
From: Daniel Han <danielhanchen@gmail.com>
Date: Thu, 16 Apr 2026 14:41:44 +0000
Subject: [PATCH] fix: address review findings -- force_hooks, CPU device
 conversion, llama.py path

Three fixes from 13-reviewer consensus:

1. Use dispatch_model(force_hooks=True) instead of manual AlignDevicesHook
   fallback. This properly handles single-device non-default GPU maps
   (e.g. all weights on cuda:1) with io_same_device routing, and
   eliminates the need for the separate add_hook_to_module call.

2. Convert non-CUDA devices (cpu, meta) to type strings in device_map_int.
   dispatch_model uses string equality (device == "cpu") internally, so
   torch.device("cpu") must become "cpu" not stay as an object.

3. Add _attach_bnb_multidevice_hooks call to FastLlamaModel.from_pretrained
   in llama.py. The default 4-bit FastLanguageModel path routes through
   llama.py:2482, not FastBaseModel, so the hook was never called for the
   most common use case (Llama, Gemma, Mistral, etc.).

Also removes redundant retie_parameters call (dispatch_model already
calls it internally) and unused imports.
---
 unsloth/models/llama.py  | 11 +++++++++++
 unsloth/models/vision.py | 33 ++++++++++-----------------------
 2 files changed, 21 insertions(+), 23 deletions(-)

diff --git a/unsloth/models/llama.py b/unsloth/models/llama.py
index 425df1c084..98f25c7dc7 100644
--- a/unsloth/models/llama.py
+++ b/unsloth/models/llama.py
@@ -2490,6 +2490,17 @@ class FastLlamaModel:
                 attn_implementation = preferred_attn_impl,
                 **kwargs,
             )
+            # Attach AlignDevicesHook for bnb multi-device / non-default-device
+            # loads.  The bnb loading path places weights on CUDA devices but
+            # never calls dispatch_model, so no hooks are installed.
+            from unsloth.models.vision import _attach_bnb_multidevice_hooks
+            _attach_bnb_multidevice_hooks(
+                model,
+                load_in_4bit      = load_in_4bit,
+                load_in_8bit      = False,
+                offload_embedding = False,
+                fast_inference    = False,
+            )
             model.fast_generate = make_fast_generate_wrapper(model.generate)
             model.fast_generate_batches = None
         else:
diff --git a/unsloth/models/vision.py b/unsloth/models/vision.py
index 284b29f45b..76a74631b1 100644
--- a/unsloth/models/vision.py
+++ b/unsloth/models/vision.py
@@ -209,8 +209,6 @@ 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
 
@@ -219,9 +217,6 @@ def _attach_bnb_multidevice_hooks(
         if not inferred_map:
             return
 
-        # Preserve tied-parameter references before dispatching.
-        tied_params = find_tied_parameters(model)
-
         # accelerate's set_module_tensor_to_device grabs param.__dict__ and
         # passes it as **kwargs to the parameter class constructor.  HF
         # Transformers adds _is_hf_initialized to parameter __dict__ which
@@ -235,36 +230,28 @@ def _attach_bnb_multidevice_hooks(
                     _stripped.append((param, key, param.__dict__.pop(key)))
 
         try:
-            # Convert device_map values from torch.device to int (HF convention)
+            # Convert device_map values to the format dispatch_model expects:
+            # CUDA devices -> int index, non-CUDA devices -> type string.
+            # dispatch_model uses string equality (device == "cpu") internally,
+            # so torch.device("cpu") must become "cpu", not stay as an object.
             device_map_int = {
-                k: v.index if isinstance(v, torch.device) and v.type == "cuda" else v
+                k: (v.index if v.type == "cuda" else v.type)
+                   if isinstance(v, torch.device) else v
                 for k, v in inferred_map.items()
             }
 
             # 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)
-                )
-
+            # force_hooks=True ensures hooks are installed even for single-
+            # device maps (e.g. all weights on cuda:1), where the default
+            # fast-path would skip hook installation entirely.
+            dispatch_model(model, device_map = device_map_int, force_hooks = True)
             desc = f"{len(inferred_map)} block(s) across {len(cuda_devs)} device(s)"
         finally:
             # Restore stripped keys
             for param, key, val in _stripped:
                 param.__dict__[key] = val
 
-        # Retie parameters that hook installation may have unlinked.
-        retie_parameters(model, tied_params)
-
         print(
             f"Unsloth: Attached accelerate AlignDevicesHook ({desc}) "
             f"for bnb multi-GPU inference."