diff --git a/test_attach_combined_load_flags_detection.py b/test_attach_combined_load_flags_detection.py new file mode 100644 index 0000000000..5872fa3c1a --- /dev/null +++ b/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 diff --git a/test_attach_detects_is_loaded_in_8bit.py b/test_attach_detects_is_loaded_in_8bit.py new file mode 100644 index 0000000000..6e51640968 --- /dev/null +++ b/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 diff --git a/test_attach_device_map_disk_entries.py b/test_attach_device_map_disk_entries.py new file mode 100644 index 0000000000..b06276c79f --- /dev/null +++ b/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}" diff --git a/test_attach_logger_error_no_misleading_warning.py b/test_attach_logger_error_no_misleading_warning.py new file mode 100644 index 0000000000..536dd39dd3 --- /dev/null +++ b/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]})" + ) diff --git a/test_attach_no_side_effects_on_early_exit.py b/test_attach_no_side_effects_on_early_exit.py new file mode 100644 index 0000000000..04f9c389a9 --- /dev/null +++ b/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" diff --git a/test_attach_restore_on_dispatch_exception.py b/test_attach_restore_on_dispatch_exception.py new file mode 100644 index 0000000000..7e320e5e91 --- /dev/null +++ b/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 diff --git a/test_infer_buffer_only_multi_device_recurses.py b/test_infer_buffer_only_multi_device_recurses.py new file mode 100644 index 0000000000..e67660c36e --- /dev/null +++ b/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 diff --git a/test_infer_params_and_buffers_same_module.py b/test_infer_params_and_buffers_same_module.py new file mode 100644 index 0000000000..bb9dd5bf07 --- /dev/null +++ b/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)} diff --git a/test_infer_three_level_deep_split.py b/test_infer_three_level_deep_split.py new file mode 100644 index 0000000000..f336afc190 --- /dev/null +++ b/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 diff --git a/test_infer_tied_params.py b/test_infer_tied_params.py new file mode 100644 index 0000000000..dbf93055b4 --- /dev/null +++ b/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)} diff --git a/test_infer_xpu_device_type_recursion.py b/test_infer_xpu_device_type_recursion.py new file mode 100644 index 0000000000..63e4d17ad9 --- /dev/null +++ b/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 diff --git a/test_llama_seq_class_hook_call_passes_false.py b/test_llama_seq_class_hook_call_passes_false.py new file mode 100644 index 0000000000..8fb1438200 --- /dev/null +++ b/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)})" diff --git a/test_vision_fastbasemodel_calls_helper.py b/test_vision_fastbasemodel_calls_helper.py new file mode 100644 index 0000000000..e23c75e482 --- /dev/null +++ b/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"