[pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

test_attach_combined_load_flags_detection.py

@@ -0,0 +1,100 @@
+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_load_in_4bit_bool_alone_activates(monkeypatch):
+    """Minimal contract: when only load_in_4bit=True is passed, the helper
+    activates without requiring model.is_loaded_in_4bit to also be True."""
+    import accelerate
+
+    called = {"n": 0}
+    monkeypatch.setattr(
+        accelerate,
+        "dispatch_model",
+        lambda *a, **kw: called.__setitem__("n", called["n"] + 1),
+    )
+    _, attach = _load_fns()
+    m = _FakeMod(params = [("w", "cuda:1")])  # no is_loaded_in_* attrs set
+    attach(
+        m,
+        load_in_4bit = True,
+        load_in_8bit = False,
+        offload_embedding = False,
+        fast_inference = False,
+    )
+    assert called["n"] == 1

test_attach_detects_is_loaded_in_8bit.py

@@ -0,0 +1,101 @@
+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_detects_bnb_via_is_loaded_in_8bit(monkeypatch):
+    """8-bit loads via quantization_config zero out load_in_*bit booleans;
+    detection must still fire via model.is_loaded_in_8bit attribute."""
+    import accelerate
+
+    called = {"n": 0}
+    monkeypatch.setattr(
+        accelerate,
+        "dispatch_model",
+        lambda *a, **kw: called.__setitem__("n", called["n"] + 1),
+    )
+    _, attach = _load_fns()
+    m = _FakeMod(params = [("w", "cuda:1")])
+    m.is_loaded_in_8bit = True
+    attach(
+        m,
+        load_in_4bit = False,
+        load_in_8bit = False,
+        offload_embedding = False,
+        fast_inference = False,
+    )
+    assert called["n"] == 1

test_attach_device_map_disk_entries.py

@@ -0,0 +1,103 @@
+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}"

test_attach_logger_error_no_misleading_warning.py

@@ -0,0 +1,110 @@
+import ast
+import warnings as _warnings
+from pathlib import Path
+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",
+    ]:
+        if p.exists():
+            return p
+    raise FileNotFoundError("vision.py not found")
+
+
+class _RaisingLogger:
+    def info(self, *a, **kw):
+        raise RuntimeError("simulated broken logging handler")
+
+
+def _load_fns(logger):
+    tree = ast.parse(_find_vision().read_text())
+    ns = {"torch": torch, "warnings": _warnings, "logger": logger}
+    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):
+        self._p = list(params or [])
+        self.hf_device_map = None
+
+    def named_parameters(self, recurse = True, remove_duplicate = False):
+        for n, d in self._p:
+            yield n, _P(d)
+
+    def parameters(self, recurse = True):
+        for _, p in self.named_parameters(recurse = recurse):
+            yield p
+
+    def named_buffers(self, recurse = True):
+        return iter([])
+
+    def named_children(self):
+        return iter([])
+
+
+def test_successful_dispatch_does_not_emit_misleading_warning_when_logger_raises(
+    monkeypatch,
+):
+    """When dispatch_model succeeds but the user's logger.info raises
+    (broken handler / strict test harness), the helper must NOT emit the
+    'Could not attach multi-device dispatch hooks automatically' warning,
+    because the hooks *are* installed."""
+    import accelerate
+
+    dispatch_called = {"n": 0}
+    monkeypatch.setattr(
+        accelerate,
+        "dispatch_model",
+        lambda *a, **kw: dispatch_called.__setitem__("n", dispatch_called["n"] + 1),
+    )
+    _, attach = _load_fns(_RaisingLogger())
+    m = _FakeMod(params = [("w", "cuda:1")])
+
+    with _warnings.catch_warnings(record = True) as caught:
+        _warnings.simplefilter("always")
+        try:
+            attach(
+                m,
+                load_in_4bit = True,
+                load_in_8bit = False,
+                offload_embedding = False,
+                fast_inference = False,
+            )
+        except RuntimeError:
+            # Post-fix: logger error may propagate (hooks are installed;
+            # surfacing a real logging misconfiguration is acceptable).
+            pass
+
+    assert dispatch_called["n"] == 1, "dispatch_model must have been called"
+    misleading = [
+        w
+        for w in caught
+        if "Could not attach multi-device dispatch hooks" in str(w.message)
+    ]
+    assert not misleading, (
+        f"Must not emit the 'Could not attach' warning when dispatch "
+        f"actually succeeded (got {[str(w.message) for w in misleading]})"
+    )

test_attach_no_side_effects_on_early_exit.py

@@ -0,0 +1,88 @@
+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 _ObservableMod:
+    """Tracks whether any mutation happened on its parameters."""
+
+    def __init__(self, params):
+        self._p = [(n, _P(d)) for n, d in params]
+        self.hf_device_map = None
+
+    def named_parameters(self, recurse = True, remove_duplicate = False):
+        for n, p in self._p:
+            yield n, p
+
+    def parameters(self, recurse = True):
+        for _, p in self._p:
+            yield p
+
+    def named_buffers(self, recurse = True):
+        return iter([])
+
+    def named_children(self):
+        return iter([])
+
+
+def test_early_exit_does_not_strip_params(monkeypatch):
+    """When a guard triggers (fast_inference=True), the helper must return
+    before touching any parameter attributes (strip loop is never entered)."""
+    import accelerate
+
+    monkeypatch.setattr(accelerate, "dispatch_model", lambda *a, **kw: None)
+    _, attach = _load_fns()
+    m = _ObservableMod([("w", "cuda:1")])
+    p = m._p[0][1]
+    p._is_hf_initialized = "original"
+    p._other_attr = "keep"
+    attach(
+        m,
+        load_in_4bit = True,
+        load_in_8bit = False,
+        offload_embedding = False,
+        fast_inference = True,
+    )
+    assert p._is_hf_initialized == "original"
+    assert p._other_attr == "keep"

test_attach_restore_on_dispatch_exception.py

@@ -0,0 +1,89 @@
+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 _TrackMod:
+    def __init__(self, params):
+        self._p = [(n, _P(d)) for n, d in params]
+        self.hf_device_map = None
+
+    def named_parameters(self, recurse = True, remove_duplicate = False):
+        for n, p in self._p:
+            yield n, p
+
+    def parameters(self, recurse = True):
+        for _, p in self._p:
+            yield p
+
+    def named_buffers(self, recurse = True):
+        return iter([])
+
+    def named_children(self):
+        return iter([])
+
+
+def test_attach_restores_is_hf_initialized_after_dispatch_raises(monkeypatch):
+    """If dispatch_model raises, the inner finally must still restore the
+    stripped _is_hf_initialized attribute on every param."""
+    import accelerate
+
+    def boom(*a, **kw):
+        raise RuntimeError("dispatch blew up")
+
+    monkeypatch.setattr(accelerate, "dispatch_model", boom)
+    _, attach = _load_fns()
+    m = _TrackMod([("w", "cuda:1")])
+    p = next(iter(m._p))[1]
+    p._is_hf_initialized = True
+    with _warnings.catch_warnings():
+        _warnings.simplefilter("ignore")
+        attach(
+            m,
+            load_in_4bit = True,
+            load_in_8bit = False,
+            offload_embedding = False,
+            fast_inference = False,
+        )
+    assert p.__dict__.get("_is_hf_initialized") is True

test_infer_buffer_only_multi_device_recurses.py

@@ -0,0 +1,81 @@
+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",
+    ]:
+        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 _BufMod:
+    def __init__(self, buffers = None, children = None):
+        self._b = list(buffers or [])
+        self._c = list(children or [])
+        self.hf_device_map = None
+
+    def named_parameters(self, recurse = True, remove_duplicate = False):
+        if False:
+            yield
+
+    def parameters(self, recurse = True):
+        if False:
+            yield
+
+    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_buffer_only_subtree_with_multi_device_recurses():
+    """Param-less subtree whose buffers span multiple devices must NOT collapse
+    to the first buffer's device. It must recurse into children so each child
+    gets its own map entry keyed by its own prefix."""
+    infer, _ = _load_fns()
+    a = _BufMod(buffers = [("cache", "cuda:0")])
+    b = _BufMod(buffers = [("cache", "cuda:1")])
+    root = _BufMod(children = [("a", a), ("b", b)])
+    dm = infer(root)
+    assert dm.get("a") == torch.device("cuda", 0), dm
+    assert dm.get("b") == torch.device("cuda", 1), dm
+    assert "" not in dm, dm

test_infer_params_and_buffers_same_module.py

@@ -0,0 +1,89 @@
+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_infer_module_with_both_params_and_buffers():
+    """A leaf carrying BOTH a param and a buffer on the same device collapses
+    to a single entry; the buffer must not confuse the single-device path."""
+    infer, _ = _load_fns()
+    m = _FakeMod(
+        params = [("weight", "cuda:1")],
+        buffers = [("running_mean", "cuda:1"), ("running_var", "cuda:1")],
+    )
+    dm = infer(m)
+    assert dm == {"": torch.device("cuda", 1)}

test_infer_three_level_deep_split.py

@@ -0,0 +1,93 @@
+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_infer_three_level_deep_mixed():
+    """Split at the third level of nesting: the algorithm must recurse deep
+    enough to distinguish grandchildren on different devices."""
+    infer, _ = _load_fns()
+    g1 = _FakeMod(params = [("w", "cuda:0")])
+    g2 = _FakeMod(params = [("w", "cuda:1")])
+    level2 = _FakeMod(children = [("g1", g1), ("g2", g2)])
+    level1 = _FakeMod(children = [("l2", level2)])
+    root = _FakeMod(children = [("l1", level1)])
+    dm = infer(root)
+    assert dm.get("l1.l2.g1") == torch.device("cuda", 0)
+    assert dm.get("l1.l2.g2") == torch.device("cuda", 1)
+    # Intermediate levels that are mixed must NOT collapse prematurely
+    assert "l1" not in dm or len({dm.get("l1"), dm.get("l1.l2.g1")}) > 1

test_infer_tied_params.py

@@ -0,0 +1,77 @@
+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 _TiedMod:
+    """Emits the same parameter object under two different names to simulate
+    tied weights (lm_head.weight == embed.weight). With remove_duplicate=False
+    we yield both names; devices unioned must still be a single device."""
+
+    def __init__(self, dev):
+        self._shared = _P(dev)
+        self.hf_device_map = None
+
+    def named_parameters(self, recurse = True, remove_duplicate = False):
+        yield "embed.weight", self._shared
+        if not remove_duplicate:
+            yield "lm_head.weight", self._shared
+
+    def parameters(self, recurse = True):
+        yield self._shared
+
+    def named_buffers(self, recurse = True):
+        return iter([])
+
+    def named_children(self):
+        return iter([])
+
+
+def test_infer_tied_params_single_entry():
+    """Tied-weight models (same Parameter yielded twice under different names)
+    must still collapse to a single-device map entry."""
+    infer, _ = _load_fns()
+    m = _TiedMod("cuda:1")
+    dm = infer(m)
+    assert dm == {"": torch.device("cuda", 1)}

test_infer_xpu_device_type_recursion.py

@@ -0,0 +1,92 @@
+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 = dev if isinstance(dev, torch.device) else torch.device(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_infer_handles_xpu_device_recursion():
+    """XPU / custom-device types must recurse and be assigned correctly. The
+    infer function is device-type agnostic and should treat different xpu
+    indices as distinct devices for map-building purposes."""
+    infer, _ = _load_fns()
+    # xpu:0 and xpu:1 — avoids any cuda-specific codepath in infer
+    a = _FakeMod(params = [("w", torch.device("xpu", 0))])
+    b = _FakeMod(params = [("w", torch.device("xpu", 1))])
+    root = _FakeMod(children = [("a", a), ("b", b)])
+    dm = infer(root)
+    assert dm.get("a") == torch.device("xpu", 0)
+    assert dm.get("b") == torch.device("xpu", 1)
+    assert "" not in dm

test_llama_seq_class_hook_call_passes_false.py

@@ -0,0 +1,70 @@
+import ast
+from pathlib import Path
+
+
+def _find_llama():
+    for p in [
+        Path(__file__).resolve().parent / "unsloth" / "models" / "llama.py",
+        Path(__file__).resolve().parents[1] / "unsloth" / "models" / "llama.py",
+    ]:
+        if p.exists():
+            return p
+    raise FileNotFoundError("llama.py not found")
+
+
+def _collect_hook_calls_under(node):
+    """Yield every Call node invoking _attach_bnb_multidevice_hooks within the
+    AST subtree rooted at `node`."""
+    for sub in ast.walk(node):
+        if (
+            isinstance(sub, ast.Call)
+            and getattr(sub.func, "id", None) == "_attach_bnb_multidevice_hooks"
+        ):
+            yield sub
+
+
+def _kwarg_literal(call, name):
+    for kw in call.keywords:
+        if kw.arg == name and isinstance(kw.value, ast.Constant):
+            return kw.value.value
+        if kw.arg == name:
+            return kw.value
+    return None
+
+
+def test_seq_class_branch_passes_fast_inference_false():
+    """The AutoModelForSequenceClassification branch (gated by
+    `if num_labels is not None:`) must pass fast_inference=False to
+    _attach_bnb_multidevice_hooks. That branch never reaches vLLM, so
+    forwarding a truthy fast_inference would short-circuit hook install
+    and regress multi-GPU bnb seq-class inference."""
+    tree = ast.parse(_find_llama().read_text())
+
+    hook_calls = []
+    for if_node in ast.walk(tree):
+        if not isinstance(if_node, ast.If):
+            continue
+        test = if_node.test
+        is_num_labels_not_none = (
+            isinstance(test, ast.Compare)
+            and isinstance(test.left, ast.Name)
+            and test.left.id == "num_labels"
+            and len(test.ops) == 1
+            and isinstance(test.ops[0], ast.IsNot)
+            and len(test.comparators) == 1
+            and isinstance(test.comparators[0], ast.Constant)
+            and test.comparators[0].value is None
+        )
+        if not is_num_labels_not_none:
+            continue
+        for body_node in if_node.body:
+            hook_calls.extend(_collect_hook_calls_under(body_node))
+
+    assert (
+        hook_calls
+    ), "No _attach_bnb_multidevice_hooks call found under `if num_labels is not None:`"
+    for call in hook_calls:
+        v = _kwarg_literal(call, "fast_inference")
+        assert (
+            v is False
+        ), f"seq-class hook call must use fast_inference=False (got {ast.dump(call)})"

test_vision_fastbasemodel_calls_helper.py

@@ -0,0 +1,41 @@
+import ast
+from pathlib import Path
+
+
+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 test_vision_fastbasemodel_from_pretrained_calls_helper():
+    """FastBaseModel.from_pretrained must invoke _attach_bnb_multidevice_hooks
+    after the underlying model load so the inference hook path is reachable
+    via the base vision loader, not only via llama."""
+    src = _find_vision().read_text()
+    tree = ast.parse(src)
+    found = False
+    for cls in ast.walk(tree):
+        if not (isinstance(cls, ast.ClassDef) and cls.name == "FastBaseModel"):
+            continue
+        for fn in ast.walk(cls):
+            if not (isinstance(fn, ast.FunctionDef) and fn.name == "from_pretrained"):
+                continue
+            for node in ast.walk(fn):
+                if (
+                    isinstance(node, ast.Call)
+                    and getattr(node.func, "id", None)
+                    == "_attach_bnb_multidevice_hooks"
+                ):
+                    found = True
+                    break
+    assert (
+        found
+    ), "FastBaseModel.from_pretrained must call _attach_bnb_multidevice_hooks"