Use loader device selection for EoRA adapter generation (#2800)

Author: Qubitium

gptqmodel/models/auto.py

@@ -68,7 +68,7 @@
 )
 from ..utils.hub import list_repo_files  # noqa: E402
 from ..utils.model import find_modules  # noqa: E402
-from ..utils.torch import CPU, torch_empty_cache  # noqa: E402
+from ..utils.torch import torch_empty_cache  # noqa: E402
 from .base import BaseQModel  # noqa: E402
 from .definitions.afmoe import AfMoeQModel  # noqa: E402
 from .definitions.apertus import ApertusQModel  # noqa: E402
@@ -686,6 +686,7 @@ def generate(
             # pass-through vars for load()
             trust_remote_code: bool = False,
             dtype: Optional[Union[str, torch.dtype]] = None,
+            device: Optional[Union[str, torch.device]] = None,
         ):
             if not adapter or not isinstance(adapter, Lora):
                 raise ValueError(f"Adapter: expected `adapter` type to be `Lora`: actual = `{adapter}`.")
@@ -696,7 +697,7 @@ def generate(
             quantized_model = GPTQModel.load(
                 model_id_or_path=quantized_model_id_or_path,
                 backend=BACKEND.GPTQ_TORCH,
-                device=CPU,
+                device=device,
                 trust_remote_code=trust_remote_code,
                 dtype=dtype,
             )
@@ -715,7 +716,7 @@ def generate(
                 backend=BACKEND.GPTQ_TORCH,
                 trust_remote_code=trust_remote_code,
                 dtype=dtype,
-                device=CPU,
+                device=device,
             )
 
             log.info("Model: Adapter generation started")

tests/test_auto_adapter_generate.py

@@ -0,0 +1,73 @@
+# GPU=-1
+from types import SimpleNamespace
+
+from gptqmodel.adapter.adapter import Lora
+from gptqmodel.models import auto
+
+
+class _FakeNativeModel:
+    def __init__(self):
+        self.generate_calls = []
+
+    def _eora_generate(self, **kwargs):
+        self.generate_calls.append(kwargs)
+
+
+def _run_adapter_generate(tmp_path, monkeypatch, *, device):
+    load_calls = []
+    find_modules_calls = []
+
+    quantized_model = SimpleNamespace(quantize_config="qcfg", model="quantized-model")
+    native_model = _FakeNativeModel()
+
+    def fake_load(cls, model_id_or_path, *args, **kwargs):
+        load_calls.append((model_id_or_path, kwargs.copy()))
+        if model_id_or_path == "quantized":
+            return quantized_model
+        if model_id_or_path == "native":
+            return native_model
+        raise AssertionError(f"unexpected load target: {model_id_or_path}")
+
+    monkeypatch.setattr(auto.GPTQModel, "load", classmethod(fake_load))
+    monkeypatch.setattr(
+        auto,
+        "find_modules",
+        lambda module, layers: find_modules_calls.append((module, layers)) or {"module": object()},
+    )
+    monkeypatch.setattr(auto, "torch_empty_cache", lambda: None)
+
+    adapter = Lora(path=str(tmp_path / "adapter"), rank=8)
+    kwargs = {
+        "adapter": adapter,
+        "model_id_or_path": "native",
+        "quantized_model_id_or_path": "quantized",
+        "calibration_dataset": ["sample"],
+    }
+    if device is not None:
+        kwargs["device"] = device
+
+    auto.GPTQModel.adapter.generate(**kwargs)
+
+    return load_calls, find_modules_calls, native_model.generate_calls
+
+
+def test_adapter_generate_defaults_to_loader_device_selection(tmp_path, monkeypatch):
+    load_calls, find_modules_calls, generate_calls = _run_adapter_generate(
+        tmp_path,
+        monkeypatch,
+        device=None,
+    )
+
+    assert [kwargs["device"] for _, kwargs in load_calls] == [None, None]
+    assert find_modules_calls == [("quantized-model", [auto.TorchLinear])]
+    assert generate_calls[0]["quantized_modules"].keys() == {"module"}
+
+
+def test_adapter_generate_forwards_explicit_device(tmp_path, monkeypatch):
+    load_calls, _, _ = _run_adapter_generate(
+        tmp_path,
+        monkeypatch,
+        device="cuda:2",
+    )
+
+    assert [kwargs["device"] for _, kwargs in load_calls] == ["cuda:2", "cuda:2"]