81 lines
3 KiB
Python
81 lines
3 KiB
Python
import ast
|
|
from pathlib import Path
|
|
import logging
|
|
import warnings as _warnings
|
|
import torch
|
|
|
|
|
|
def _find_vision():
|
|
for p in [
|
|
Path(__file__).resolve().parent / "unsloth" / "models" / "vision.py",
|
|
Path(__file__).resolve().parents[1] / "unsloth" / "models" / "vision.py",
|
|
Path("/mnt/disks/unslothai/ubuntu/workspace_25/github_review/unsloth-pr-5053-staging-3/unsloth/models/vision.py"),
|
|
]:
|
|
if p.exists():
|
|
return p
|
|
raise FileNotFoundError("vision.py not found")
|
|
|
|
|
|
def _load_fns():
|
|
tree = ast.parse(_find_vision().read_text())
|
|
ns = {"torch": torch, "warnings": _warnings, "logger": logging.getLogger("test")}
|
|
for node in tree.body:
|
|
if isinstance(node, ast.FunctionDef) and node.name in {
|
|
"_infer_device_map_from_loaded_model",
|
|
"_attach_bnb_multidevice_hooks",
|
|
}:
|
|
exec(compile(ast.Module(body=[node], type_ignores=[]), str(_find_vision()), "exec"), ns)
|
|
return ns["_infer_device_map_from_loaded_model"], ns["_attach_bnb_multidevice_hooks"]
|
|
|
|
|
|
class _P:
|
|
def __init__(self, dev):
|
|
self.device = torch.device(dev) if isinstance(dev, str) else dev
|
|
|
|
|
|
class _FakeMod:
|
|
def __init__(self, params=None, buffers=None, children=None, hf_device_map=None):
|
|
self._p = list(params or [])
|
|
self._b = list(buffers or [])
|
|
self._c = list(children or [])
|
|
self.hf_device_map = hf_device_map
|
|
|
|
def named_parameters(self, recurse=True, remove_duplicate=False):
|
|
for n, d in self._p:
|
|
yield n, _P(d)
|
|
if recurse:
|
|
for cn, cm in self._c:
|
|
for pn, pp in cm.named_parameters(recurse=True, remove_duplicate=remove_duplicate):
|
|
yield f"{cn}.{pn}", pp
|
|
|
|
def parameters(self, recurse=True):
|
|
for _, p in self.named_parameters(recurse=recurse):
|
|
yield p
|
|
|
|
def named_buffers(self, recurse=True):
|
|
for n, d in self._b:
|
|
yield n, _P(d)
|
|
if recurse:
|
|
for cn, cm in self._c:
|
|
for bn, bb in cm.named_buffers(recurse=True):
|
|
yield f"{cn}.{bn}", bb
|
|
|
|
def named_children(self):
|
|
yield from self._c
|
|
|
|
|
|
def test_attach_main_device_skips_cpu_and_disk_candidates(monkeypatch):
|
|
"""When inferred_map values mix cpu + gpu, main_device fallback must skip
|
|
non-device entries. Verifies the iter-4 `d not in ("cpu", "disk")` filter
|
|
handles both string constants."""
|
|
import accelerate
|
|
rec = {}
|
|
monkeypatch.setattr(accelerate, "dispatch_model", lambda model, **kw: rec.update(kw))
|
|
_, attach = _load_fns()
|
|
# First entry is cpu; fallback must find the cuda:1 entry instead.
|
|
a = _FakeMod(params=[("w", "cpu")])
|
|
b = _FakeMod(params=[("w", "cuda:1")])
|
|
m = _FakeMod(children=[("a", a), ("b", b)])
|
|
attach(m, load_in_4bit=True, load_in_8bit=False, offload_embedding=False, fast_inference=False)
|
|
md = rec.get("main_device")
|
|
assert md == 1, f"main_device must skip cpu/disk strings, got {md!r}"
|