ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Split: keep only 10 file(s)

Browse source
This commit is contained in:
Daniel Han 2026-04-20 23:19:40 +00:00
commit 9741e2ff0b
2 changed files with 26 additions and 30 deletions
Download patch file Download diff file Expand all files Collapse all files

9
unsloth/models/llama.py
View file

@ -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

47
unsloth/models/vision.py
View file

@ -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