From 4edb8943c2d387b0f28079e8352e2ee752e2abb2 Mon Sep 17 00:00:00 2001
From: Daniel Han <danielhanchen@gmail.com>
Date: Thu, 16 Apr 2026 15:58:59 +0000
Subject: [PATCH] Fix review findings for PR #15

---
 unsloth/models/llama.py  | 11 ++++++++++-
 unsloth/models/vision.py | 18 +++++++++++-------
 2 files changed, 21 insertions(+), 8 deletions(-)

diff --git a/unsloth/models/llama.py b/unsloth/models/llama.py
index cc021f7b97..7fb56f46e7 100644
--- a/unsloth/models/llama.py
+++ b/unsloth/models/llama.py
@@ -2478,6 +2478,15 @@ class FastLlamaModel:
                     and not _head.weight.is_floating_point()
                 ):
                     _head.to(dtype)
+            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 = False,
+            )
         elif not fast_inference:
             model = AutoModelForCausalLM.from_pretrained(
                 model_name,
@@ -2496,7 +2505,7 @@ class FastLlamaModel:
             _attach_bnb_multidevice_hooks(
                 model,
                 load_in_4bit = load_in_4bit,
-                load_in_8bit = False,
+                load_in_8bit = kwargs.get("load_in_8bit", False),
                 offload_embedding = False,
                 fast_inference = False,
             )
diff --git a/unsloth/models/vision.py b/unsloth/models/vision.py
index 1bdfa9b583..f67622b5da 100644
--- a/unsloth/models/vision.py
+++ b/unsloth/models/vision.py
@@ -75,6 +75,7 @@ import functools
 import os
 import gc
 import math
+import warnings
 from typing import Optional, Tuple, List, Union
 import re, inspect, sys
 import contextlib
@@ -107,7 +108,9 @@ def _infer_device_map_from_loaded_model(model):
         if not params:
             bufs = list(module.named_buffers())
             if bufs:
-                device_map[prefix] = bufs[0][1].device
+                buf_devs = {b.device for _, b in bufs}
+                if len(buf_devs) == 1:
+                    device_map[prefix] = next(iter(buf_devs))
             return
         devices = {p.device for _, p in params}
         if len(devices) == 1:
@@ -154,8 +157,6 @@ def _attach_bnb_multidevice_hooks(
             if hasattr(p, "device") and p.device.type == "cuda"
         }
     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})",
@@ -199,20 +200,23 @@ def _attach_bnb_multidevice_hooks(
             }
 
             # force_hooks=True: install hooks even for single-device maps.
-            dispatch_model(model, device_map = device_map_int, force_hooks = True)
+            dispatch_model(
+                model,
+                device_map = device_map_int,
+                skip_keys = getattr(model, "_skip_keys_device_placement", None),
+                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
 
-        print(
+        logger.info(
             f"Unsloth: Attached accelerate AlignDevicesHook ({desc}) "
             f"for bnb multi-GPU inference."
         )
     except Exception as exc:
-        import warnings
-
         warnings.warn(
             f"Unsloth: Could not attach multi-device dispatch hooks automatically "
             f"({type(exc).__name__}: {exc}). "