Enable init from QEffAutoPeftModelForCausalLM with finite_adapters flag

Signed-off-by: Jou-An Chen <[email protected]>

Author: quic-jouachen

QEfficient/peft/auto.py

@@ -17,6 +17,7 @@
 from transformers import GenerationConfig, StoppingCriteria, StoppingCriteriaList
 from transformers.generation.streamers import BaseStreamer
 
+from QEfficient import QEffAutoLoraModelForCausalLM
 from QEfficient.base.modeling_qeff import QEFFBaseModel
 from QEfficient.base.onnx_transforms import FP16ClipTransform, OnnxTransform, SplitTensorsTransform
 from QEfficient.base.pytorch_transforms import PytorchTransform
@@ -152,7 +153,13 @@ def from_pretrained(cls, pretrained_name_or_path: str, *args, **kwargs):
         if kwargs.get("use_cache") is False:
             warnings.warn("Overriding to use_cache=True")
         kwargs["use_cache"] = True
-        obj = cls._from_pretrained(pretrained_name_or_path, *args, **kwargs)
+
+        if kwargs.pop("finite_adapters", False):  # initialize through finite_adapters class
+            obj = QEffAutoLoraModelForCausalLM.from_pretrained(
+                pretrained_model_name_or_path=pretrained_name_or_path, **kwargs
+            )
+        else:
+            obj = cls._from_pretrained(pretrained_name_or_path, *args, **kwargs)
         return obj
 
     def export(self, export_dir: Optional[str] = None) -> str:

examples/lora_models.py

@@ -7,7 +7,7 @@
 
 ## This example works on continuous batching with different lora adapters in the same batch ##
 
-from QEfficient import QEffAutoLoraModelForCausalLM
+from QEfficient import QEffAutoPeftModelForCausalLM
 from QEfficient.utils import load_hf_tokenizer
 
 base_model_name = "mistralai/Mistral-7B-v0.1"
@@ -18,17 +18,14 @@
 
 ## STEP 1 -- init base model
 
-# **Option1**: Download model weights from hugging face & Init it with QEffAuto model to apply QEff transforms
-# model_hf = AutoModelForCausalLM.from_pretrained(base_model_name)
-# qeff_model = QEffAutoLoraModelForCausalLM(model_hf, continuous_batching=True)
-
-# **Option2**: Initialize the model using from_pretrained() method
-qeff_model = QEffAutoLoraModelForCausalLM.from_pretrained(
-    pretrained_model_name_or_path=base_model_name, continuous_batching=True
+qeff_model = QEffAutoPeftModelForCausalLM.from_pretrained(
+    pretrained_name_or_path=base_model_name, continuous_batching=True, finite_adapters=True
 )
 
 # (alternative) non-cb initialization
-# qeff_model = QEffAutoLoraModelForCausalLM.from_pretrained(pretrained_model_name_or_path=base_model_name, continuous_batching=False)
+# qeff_model = QEffAutoPeftModelForCausalLM.from_pretrained(
+#     pretrained_name_or_path=base_model_name, continuous_batching=False, finite_adapters=True
+# )
 
 ## STEP 2 -- load adapter adapter
 adapter_id_gsm8k = qeff_model.load_adapter("predibase/gsm8k", "gsm8k")

tests/lora/test_lora_model.py

@@ -12,7 +12,7 @@
 from peft import LoraConfig
 from transformers import AutoConfig, AutoModelForCausalLM
 
-from QEfficient import QEffAutoLoraModelForCausalLM
+from QEfficient import QEffAutoLoraModelForCausalLM, QEffAutoPeftModelForCausalLM
 from QEfficient.utils import load_hf_tokenizer
 
 configs = [
@@ -74,6 +74,20 @@ def test_auto_lora_model_for_causal_lm_from_pretrained(base_model_name, adapter_
     assert len(qeff_model.active_adapter_to_id) == 0
 
 
+# test peft model initialization using from_pretrained approach
+@pytest.mark.parametrize("base_model_name,adapter_id_0,adapter_id_1", model_samples)
+def test_auto_peft_model_for_causal_lm_from_pretrained(base_model_name, adapter_id_0, adapter_id_1):
+    qeff_model = QEffAutoPeftModelForCausalLM.from_pretrained(
+        pretrained_name_or_path=base_model_name, finite_adapters=True
+    )
+
+    assert qeff_model.base_model_name == base_model_name
+    assert len(qeff_model.adapter_weights) == 0
+    assert len(qeff_model.adapter_configs) == 0
+    assert qeff_model.max_num_adapters == 0
+    assert len(qeff_model.active_adapter_to_id) == 0
+
+
 # test the init assertion for models that are not supported
 @pytest.mark.parametrize("base_model_name", ["distilbert/distilgpt2"])
 def test_auto_lora_model_for_causal_lm_init_from_unsupported_model(base_model_name):