From 9741e2ff0bd974a8904862665a3d82bdccefc520 Mon Sep 17 00:00:00 2001
From: Daniel Han <danielhanchen@gmail.com>
Date: Mon, 20 Apr 2026 23:19:40 +0000
Subject: [PATCH] Split: keep only 10 file(s)

---
 unsloth/models/llama.py  |  9 ++++++--
 unsloth/models/vision.py | 47 ++++++++++++++++------------------------
 2 files changed, 26 insertions(+), 30 deletions(-)

diff --git a/unsloth/models/llama.py b/unsloth/models/llama.py
index c27a0f5382..d39f2588ef 100644
--- a/unsloth/models/llama.py
+++ b/unsloth/models/llama.py
@@ -76,7 +76,7 @@ from transformers.modeling_attn_mask_utils import (
 )
 from ..kernels import *
 from ..tokenizer_utils import *
-from .vision import FastBaseModel, _attach_bnb_multidevice_hooks
+from .vision import FastBaseModel
 
 # Final patching code
 from transformers.models.llama.modeling_llama import (
@@ -2478,6 +2478,9 @@ class FastLlamaModel:
                     and not _head.weight.is_floating_point()
                 ):
                     _head.to(dtype)
+            # Attach dispatch hooks for bnb multi-device loads.
+            from unsloth.models.vision import _attach_bnb_multidevice_hooks
+
             _attach_bnb_multidevice_hooks(
                 model,
                 load_in_4bit = load_in_4bit,
@@ -2498,12 +2501,14 @@ class FastLlamaModel:
                 **kwargs,
             )
             # Attach dispatch hooks for bnb multi-device loads.
+            from unsloth.models.vision import _attach_bnb_multidevice_hooks
+
             _attach_bnb_multidevice_hooks(
                 model,
                 load_in_4bit = load_in_4bit,
                 load_in_8bit = kwargs.get("load_in_8bit", False),
                 offload_embedding = False,
-                fast_inference = fast_inference,
+                fast_inference = False,
             )
             model.fast_generate = make_fast_generate_wrapper(model.generate)
             model.fast_generate_batches = None
diff --git a/unsloth/models/vision.py b/unsloth/models/vision.py
index 61e0822aa5..df371e00c8 100644
--- a/unsloth/models/vision.py
+++ b/unsloth/models/vision.py
@@ -104,39 +104,30 @@ def _infer_device_map_from_loaded_model(model):
     device_map = {}
 
     def _assign(module, prefix):
-        subtree_devs = {
-            p.device for _, p in module.named_parameters(remove_duplicate = False)
-        }
-        if not subtree_devs:
+        params = list(module.named_parameters(remove_duplicate = False))
+        if not params:
             bufs = list(module.named_buffers())
             if bufs:
-                buf_devs = {b.device for _, b in bufs}
-                if len(buf_devs) == 1:
-                    device_map[prefix] = next(iter(buf_devs))
-                else:
-                    for child_name, child in module.named_children():
-                        child_prefix = (
-                            f"{prefix}.{child_name}" if prefix else child_name
-                        )
-                        _assign(child, child_prefix)
+                device_map[prefix] = bufs[0][1].device
             return
-        if len(subtree_devs) == 1:
-            device_map[prefix] = next(iter(subtree_devs))
-            return
-        for child_name, child in module.named_children():
-            child_prefix = f"{prefix}.{child_name}" if prefix else child_name
-            _assign(child, child_prefix)
-        local_devs = {
-            p.device
-            for _, p in module.named_parameters(
-                recurse = False,
-                remove_duplicate = False,
-            )
-        }
-        if local_devs and len(local_devs) == 1:
-            device_map[prefix] = next(iter(local_devs))
+        devices = {p.device for _, p in params}
+        if len(devices) == 1:
+            device_map[prefix] = next(iter(devices))
+        else:
+            for child_name, child in module.named_children():
+                child_prefix = f"{prefix}.{child_name}" if prefix else child_name
+                _assign(child, child_prefix)
+            for pname, param in module.named_parameters(remove_duplicate = False):
+                if "." not in pname:
+                    full = f"{prefix}.{pname}" if prefix else pname
+                    if not any(
+                        full == k or full.startswith(k + ".") for k in device_map
+                    ):
+                        device_map[full] = param.device
 
     _assign(model, "")
+    if "" in device_map and len(device_map) > 1:
+        device_map.pop("")
     return device_map