Finite lorax support (#153)

* Initial commit for finite loras implementation

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

* Remove set delete adapter, add init assertion, update LinearMultiLoRA

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

* Fix base model inference index INTMAX issue

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

* Addressed review comments

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

* Rebase on PR116 and make API changes

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

* Enable init from QEffAutoPeftModelForCausalLM with finite_adapters flag

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

* Address review comments

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

* allow adapter_name passed as keyword argument, updated all finite lora tests to use single layer models

Signed-off-by: Onkar Chougule <[email protected]>

* added pytest on_qaic marker for lora test using AI_100 device

Signed-off-by: Onkar Chougule <[email protected]>

---------

Signed-off-by: Jou-An Chen <[email protected]>
Signed-off-by: Onkar Chougule <[email protected]>
Co-authored-by: Onkar Chougule <[email protected]>

Author: quic-jouachen

QEfficient/generation/text_generation_inference.py

@@ -230,6 +230,7 @@ def cloud_ai_100_exec_kv(
     stream: bool = True,
     write_io_dir: Optional[str] = None,
     automation=False,
+    prompt_to_lora_id_mapping: Optional[List[int]] = None,
 ):
     """
     This method generates output until ``eos`` or ``generation_len`` by executing the compiled ``qpc`` on ``Cloud AI 100`` Hardware cards.
@@ -277,6 +278,7 @@ def cloud_ai_100_exec_kv(
         stream=stream,
         write_io_dir=write_io_dir,
         full_batch_size=full_batch_size,
+        prompt_to_lora_id_mapping=prompt_to_lora_id_mapping,
     )
     if full_batch_size is None:
         exec_info = [
@@ -313,6 +315,7 @@ def __init__(
         qpc_path: str,
         prompt: List[str],
         full_batch_size: Optional[int] = None,
+        prompt_to_lora_id_mapping: Optional[List[int]] = None,
         ctx_len: Optional[int] = None,
         generation_len: Optional[int] = None,
         device_id: Optional[List[int]] = None,
@@ -342,6 +345,16 @@ def __init__(
             full_batch_size if full_batch_size else self._fetch_full_batch_size()
         )  # Check and fetch full batch size if CB is enabled
 
+        if prompt_to_lora_id_mapping:
+            self.prompt_to_lora_id_mapping_prefill = deque(prompt_to_lora_id_mapping)
+            if self.full_batch_size:
+                self.prompt_to_lora_id_mapping_decode = prompt_to_lora_id_mapping
+            else:
+                self.prompt_to_lora_id_mapping_decode = deque(prompt_to_lora_id_mapping)
+        else:
+            self.prompt_to_lora_id_mapping_prefill = None
+            self.prompt_to_lora_id_mapping_decode = None
+
         self.set_tokenizer_params()  # set tokenizer params
 
         # Initialize the storage variables.
@@ -464,6 +477,16 @@ def prepare_decode_inputs(self):
         if self.batch_index is not None:
             decode_inputs["batch_index"] = self.batch_index
 
+        if self.prompt_to_lora_id_mapping_decode:
+            if self.full_batch_size:
+                first_batch_lora_ids = [self.prompt_to_lora_id_mapping_decode[i] for i in range(self.full_batch_size)]
+                decode_inputs["lora_ids"] = np.array(first_batch_lora_ids, dtype=np.int64).reshape(
+                    self.full_batch_size, 1
+                )
+            else:
+                batch_lora_ids = [self.prompt_to_lora_id_mapping_decode.popleft() for i in range(self.batch_size)]
+                decode_inputs["lora_ids"] = np.array(batch_lora_ids, dtype=np.int64).reshape(self.batch_size, 1)
+
         return decode_inputs
 
     def _update_decode_input(self, outputs, position_ids, generation_len, decode_batch_id=None):
@@ -552,6 +575,15 @@ def run_prefill(self, prompt, generation_len, prefill_logit_bs=1, decode_batch_i
         if decode_batch_id is not None:
             inputs["batch_index"] = decode_batch_id
 
+        if self.prompt_to_lora_id_mapping_prefill:
+            if self.full_batch_size:
+                inputs["lora_ids"] = np.array(self.prompt_to_lora_id_mapping_prefill.popleft(), dtype=np.int64).reshape(
+                    1, 1
+                )
+            else:
+                batch_lora_ids = [self.prompt_to_lora_id_mapping_prefill.popleft() for i in range(self.batch_size)]
+                inputs["lora_ids"] = np.array(batch_lora_ids, dtype=np.int64).reshape(self.batch_size, 1)
+
         for i in range(num_chunks):
             chunk_inputs = inputs.copy()
             chunk_inputs["input_ids"] = inputs["input_ids"][
@@ -628,6 +660,12 @@ def run_continuous_batching_decode(self, prompt_queue, generation_len):
 
                         self.session.set_buffers({"logits": logits_out_placeholder})
                         decode_pause_time += perf_counter() - start
+
+                        if self.prompt_to_lora_id_mapping_decode:
+                            decode_inputs["lora_ids"][decode_batch_id] = self.prompt_to_lora_id_mapping_decode[
+                                batch_id_map[decode_batch_id]
+                            ]
+
                     else:
                         current_decode_ongoing[decode_batch_id] = False
                 else:
@@ -639,6 +677,7 @@ def run_continuous_batching_decode(self, prompt_queue, generation_len):
                     )
 
                     generated_id_current_index[decode_batch_id] += 1
+
         return decode_pause_time
 
     def run_decode(self, decode_inputs, generation_len):

QEfficient/peft/auto.py

@@ -12,7 +12,7 @@
 
 import numpy as np
 import torch
-from peft import AutoPeftModelForCausalLM, PeftModelForCausalLM, load_peft_weights
+from peft import AutoPeftModelForCausalLM, PeftConfig, PeftModelForCausalLM, load_peft_weights
 from torch import nn
 from transformers import GenerationConfig, StoppingCriteria, StoppingCriteriaList
 from transformers.generation.streamers import BaseStreamer
@@ -21,6 +21,7 @@
 from QEfficient.base.onnx_transforms import FP16ClipTransform, OnnxTransform, SplitTensorsTransform
 from QEfficient.base.pytorch_transforms import PytorchTransform
 from QEfficient.generation.cloud_infer import QAICInferenceSession
+from QEfficient.peft.lora import QEffAutoLoraModelForCausalLM
 from QEfficient.peft.onnx_transforms import AdapterWeightsToInputsTransform
 from QEfficient.peft.pytorch_transforms import PeftModelInputsTransform
 from QEfficient.transformers.pytorch_transforms import CustomOpsTransform, KVCacheTransform
@@ -38,6 +39,7 @@ class QEffAutoPeftModelForCausalLM(QEFFBaseModel):
 
     Args:
         :model (nn.Module): PyTorch model
+        :finite_adapters (bool): set True to enable finite adapter mode with QEffAutoLoraModelForCausalLM class. Please refer to QEffAutoLoraModelForCausalLM for API specification.
 
     .. code-block:: python
 
@@ -80,6 +82,9 @@ def __init__(self, model: nn.Module):
             for adapter_name in model.peft_config
         }
 
+    def __repr__(self) -> str:
+        return self.__class__.__name__ + "\n" + self.model.__repr__()
+
     @property
     def model_name(self) -> str:
         mname = self.model.get_base_model().__class__.__name__ + "-lora"
@@ -145,14 +150,31 @@ def from_pretrained(cls, pretrained_name_or_path: str, *args, **kwargs):
         """
         Args:
             :pretrained_name_or_path (str): Model card name from huggingface or local path to model directory.
+            :finite_adapters (bool): set True to enable finite adapter mode with QEffAutoLoraModelForCausalLM class. Please refer to QEffAutoLoraModelForCausalLM for API specification.
+            :adapter_name (str): Name used to identify loaded adapter.
             :args, kwargs: Additional arguments to pass to peft.AutoPeftModelForCausalLM.
         """
         if kwargs.get("full_batch_size"):
             raise NotImplementedError("Continuous batching currently not supported for PEFT models")
         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=PeftConfig.from_pretrained(
+                    pretrained_name_or_path
+                ).base_model_name_or_path,
+                **kwargs,
+            )
+            if adapter_name := kwargs.pop("adapter_name", None):
+                obj.load_adapter(pretrained_name_or_path, adapter_name=adapter_name)
+                return obj
+            if len(args) == 0 or not isinstance(list(args)[0], str):
+                raise TypeError("Required adapter name argument in string format")
+            obj.load_adapter(pretrained_name_or_path, list(args)[0])
+        else:
+            obj = cls._from_pretrained(pretrained_name_or_path, *args, **kwargs)
         return obj
 
     def export(self, export_dir: Optional[str] = None) -> str:

QEfficient/peft/lora/__init__.py

@@ -0,0 +1,12 @@
+# -----------------------------------------------------------------------------
+#
+# Copyright (c) 2024 Qualcomm Innovation Center, Inc. All rights reserved.
+# SPDX-License-Identifier: BSD-3-Clause
+#
+# ----------------------------------------------------------------------------
+
+from QEfficient.peft.lora.auto import QEffAutoLoraModelForCausalLM
+
+__all__ = [
+    "QEffAutoLoraModelForCausalLM",
+]