Add review tests for PR #5053

test_attach_combined_load_flags_detection.py

@@ -0,0 +1,76 @@
+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,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 _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,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",
+        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_no_side_effects_on_early_exit.py

@@ -0,0 +1,69 @@
+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,70 @@
+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_params_and_buffers_same_module.py

@@ -0,0 +1,76 @@
+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,80 @@
+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,65 @@
+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,79 @@
+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_vision_fastbasemodel_calls_helper.py

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