From 7b324a63a1c4427e6d867346eb6f84d768fa6984 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Thu, 16 Apr 2026 17:29:23 +0000 Subject: [PATCH] Fix review findings for PR #15 --- unsloth/models/llama.py | 2 +- unsloth/models/vision.py | 7 +++++++ 2 files changed, 8 insertions(+), 1 deletion(-) diff --git a/unsloth/models/llama.py b/unsloth/models/llama.py index 5088c6f8a0..c27a0f5382 100644 --- a/unsloth/models/llama.py +++ b/unsloth/models/llama.py @@ -2503,7 +2503,7 @@ class FastLlamaModel: load_in_4bit = load_in_4bit, load_in_8bit = kwargs.get("load_in_8bit", False), offload_embedding = False, - fast_inference = False, + fast_inference = fast_inference, ) 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 ffc47d813a..030db49168 100644 --- a/unsloth/models/vision.py +++ b/unsloth/models/vision.py @@ -207,9 +207,16 @@ def _attach_bnb_multidevice_hooks( } # force_hooks=True: install hooks even for single-device maps. + main_device = device_map_int.get("") + if main_device in (None, "cpu", "disk"): + main_device = next( + (d for d in device_map_int.values() if d not in ("cpu", "disk")), + None, + ) dispatch_model( model, device_map = device_map_int, + main_device = main_device, skip_keys = getattr(model, "_skip_keys_device_placement", None), force_hooks = True, )