diff --git a/unsloth/models/llama.py b/unsloth/models/llama.py index 999711efdb..d39f2588ef 100644 --- a/unsloth/models/llama.py +++ b/unsloth/models/llama.py @@ -2478,6 +2478,16 @@ 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, + load_in_8bit = kwargs.get("load_in_8bit", False), + offload_embedding = False, + fast_inference = fast_inference, + ) elif not fast_inference: model = AutoModelForCausalLM.from_pretrained( model_name, @@ -2490,6 +2500,16 @@ class FastLlamaModel: attn_implementation = preferred_attn_impl, **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 = 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 90c93ea3f7..df371e00c8 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 @@ -97,6 +98,144 @@ __all__ = [ "FastBaseModel", ] + +def _infer_device_map_from_loaded_model(model): + """Build a compact device_map by inspecting actual parameter placements.""" + device_map = {} + + def _assign(module, prefix): + params = list(module.named_parameters(remove_duplicate = False)) + if not params: + bufs = list(module.named_buffers()) + if bufs: + device_map[prefix] = bufs[0][1].device + return + 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 + + +def _attach_bnb_multidevice_hooks( + model, load_in_4bit, load_in_8bit, offload_embedding, fast_inference +): + """ + Attach accelerate AlignDevicesHook on a bnb model loaded across multiple + devices (or a non-default device). No-op for single-GPU cuda:0, non-bnb, + vLLM, or already-dispatched models. + """ + if fast_inference: + return + is_bnb = ( + load_in_4bit + or load_in_8bit + or getattr(model, "is_loaded_in_4bit", False) + or getattr(model, "is_loaded_in_8bit", False) + or getattr(model, "quantization_method", None) == "bitsandbytes" + ) + if not is_bnb: + return + if offload_embedding: + return + if getattr(model, "hf_device_map", None) is not None: + return # already dispatched + + try: + all_devs = {p.device for p in model.parameters()} + except Exception as exc: + warnings.warn( + "Unsloth: Failed to determine device placement from model parameters, " + f"so multi-GPU hooks cannot be attached. ({type(exc).__name__}: {exc})", + RuntimeWarning, + stacklevel = 2, + ) + return + + cuda_devs = {d for d in all_devs if d.type == "cuda"} + if not cuda_devs: + return + + default_cuda = torch.device("cuda", 0) + if all_devs == {default_cuda}: + return + + try: + from accelerate import dispatch_model + except ImportError: + return # accelerate not available + + try: + inferred_map = _infer_device_map_from_loaded_model(model) + if not inferred_map: + return + + # bnb constructors reject _is_hf_initialized; strip before dispatch. + _extra_keys = ("_is_hf_initialized",) + _stripped = [] + for _, param in model.named_parameters(): + for key in _extra_keys: + if key in param.__dict__: + _stripped.append((param, key, param.__dict__.pop(key))) + + try: + # CUDA -> int index, non-CUDA -> type string ("cpu", "meta"). + device_map_int = { + k: (v.index if v.type == "cuda" else v.type) + if isinstance(v, torch.device) + else v + for k, v in inferred_map.items() + } + + # 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, + ) + 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 + + logger.info( + f"Unsloth: Attached accelerate AlignDevicesHook ({desc}) " + f"for bnb multi-GPU inference." + ) + except Exception as exc: + warnings.warn( + f"Unsloth: Could not attach multi-device dispatch hooks automatically " + f"({type(exc).__name__}: {exc}). " + "Cross-device inference may fail. Consider using a single GPU or " + "calling accelerate.dispatch_model() manually.", + RuntimeWarning, + stacklevel = 2, + ) + + global NUM_LOGITS_TO_KEEP NUM_LOGITS_TO_KEEP = dict() @@ -796,6 +935,14 @@ class FastBaseModel: # attn_implementation = attn_implementation, **kwargs, ) + # Attach dispatch hooks for bnb multi-device loads. + _attach_bnb_multidevice_hooks( + model, + load_in_4bit = load_in_4bit, + load_in_8bit = load_in_8bit, + offload_embedding = offload_embedding, + fast_inference = fast_inference, + ) if hasattr(model, "generate"): model.fast_generate = make_fast_generate_wrapper(model.generate) model.fast_generate_batches = error_out_no_vllm