Initial commit for finite loras implementation

Signed-off-by: Jou-An Chen <quic_jouachen@quicinc.com>

Author: quic-jouachen

QEfficient/__init__.py

@@ -24,5 +24,6 @@
     "QEffAutoModel",
     "QEFFAutoModelForCausalLM",
     "QEffAutoPeftModelForCausalLM",
+    "QEffAutoLoraModelForCausalLM",
     "QEFFCommonLoader",
 ]

QEfficient/exporter/export_hf_to_cloud_ai_100.py

@@ -14,7 +14,7 @@
 from transformers import PreTrainedTokenizer, PreTrainedTokenizerFast
 
 import QEfficient
-from QEfficient.base.common import AUTO_MODEL_MAP_TO_MODEL_TYPE_MAP, QEFF_MODEL_TYPE, QEFFCommonLoader
+from QEfficient.base.common import QEFFCommonLoader
 from QEfficient.base.modeling_qeff import QEFFBaseModel
 from QEfficient.exporter.export_utils import export_onnx, fix_onnx_fp16, generate_input_files, run_model_on_ort
 from QEfficient.transformers.modeling_utils import get_lists_of_cb_qeff_models
@@ -149,6 +149,7 @@ def convert_to_cloud_kvstyle(
     tokenizer: Union[PreTrainedTokenizer, PreTrainedTokenizerFast],
     onnx_dir_path: str,
     seq_len: int,
+    max_num_adapters: int,
 ) -> str:
     """
     API to convert model with kv retention and export to ONNX.
@@ -181,7 +182,7 @@ def convert_to_cloud_kvstyle(
 
     # Decide path for saving exported ONNX files.
     model_name = export_kvstyle_transformed_model_to_onnx(
-        model_name, qeff_model.model, tokenizer, onnx_dir_path, seq_len
+        model_name, qeff_model.model, tokenizer, onnx_dir_path, seq_len, max_num_adapters
     )  # type: ignore
 
     # return the model path for automation.
@@ -195,6 +196,7 @@ def export_kvstyle_transformed_model_to_onnx(
     onnx_dir_path: str,
     seq_len: int,
     full_batch_size: Optional[int] = None,
+    max_num_adapters: Optional[int] = None,
 ) -> str:
     # Disabling requires_grad on all parameters
     for _, p in enumerate(transformed_model.parameters()):
@@ -213,6 +215,7 @@ def export_kvstyle_transformed_model_to_onnx(
         prompt_len=Constants.PROMPT_LEN,
         ctx_len=seq_len,
         full_batch_size=full_batch_size,
+        max_num_adapters=max_num_adapters,
     )
 
     inputs = input_handler.prepare_pytorch_inputs()
@@ -318,6 +321,7 @@ def export_for_cloud(
     onnx_dir_path: str,
     seq_length: int = Constants.SEQ_LEN,
     full_batch_size: Optional[int] = None,
+    max_num_adapters: Optional[int] = None,
 ) -> str:
     # Check if model architecture is supported for continuous batching.
     if full_batch_size and qeff_model.model.config.architectures[0] not in get_lists_of_cb_qeff_models.architectures:
@@ -326,19 +330,20 @@ def export_for_cloud(
         )
 
     # FIXME: move all this to class instead of here, and just call qeff_model.export here.
-    if AUTO_MODEL_MAP_TO_MODEL_TYPE_MAP.get(qeff_model.__class__, None) == QEFF_MODEL_TYPE.CAUSALLM:  # type: ignore
-        return export_lm_model_for_cloud(
-            model_name=model_name,
-            qeff_model=qeff_model,  # type: ignore
-            tokenizer=tokenizer,
-            onnx_dir_path=onnx_dir_path,
-            seq_length=seq_length,
-            full_batch_size=full_batch_size,
-        )
-    else:
-        raise NotImplementedError(
-            f"Only model type {QEFFAutoModelForCausalLM.__class__.__name__} is supported for export, got {type(qeff_model)}"
-        )
+    # if AUTO_MODEL_MAP_TO_MODEL_TYPE_MAP.get(qeff_model.__class__, None) == QEFF_MODEL_TYPE.CAUSALLM:  # type: ignore
+    return export_lm_model_for_cloud(
+        model_name=model_name,
+        qeff_model=qeff_model,  # type: ignore
+        tokenizer=tokenizer,
+        onnx_dir_path=onnx_dir_path,
+        seq_length=seq_length,
+        full_batch_size=full_batch_size,
+        max_num_adapters=max_num_adapters,
+    )
+    # else:
+    #     raise NotImplementedError(
+    #         f"Only model type {QEFFAutoModelForCausalLM.__class__.__name__} is supported for export, got {type(qeff_model)}"
+    #     )
 
 
 def export_lm_model_for_cloud(
@@ -348,6 +353,7 @@ def export_lm_model_for_cloud(
     onnx_dir_path: str,
     seq_length: int,
     full_batch_size: Optional[int] = None,
+    max_num_adapters: Optional[int] = None,
 ) -> str:
     if os.path.exists(onnx_dir_path):
         logger.warning(f"Overriding {onnx_dir_path}")
@@ -361,6 +367,7 @@ def export_lm_model_for_cloud(
             onnx_dir_path=onnx_dir_path,
             seq_len=seq_length,
             full_batch_size=full_batch_size,
+            max_num_adapters=max_num_adapters,
         )  # type: ignore
 
     else:
@@ -386,6 +393,7 @@ def qualcomm_efficient_converter(
     kv: bool = True,
     form_factor: str = "cloud",
     full_batch_size: Optional[int] = None,
+    max_num_adapters: Optional[int] = None,
 ) -> Tuple[str, str]:
     """
     This method is an alias for ``QEfficient.export``.
@@ -466,6 +474,7 @@ def qualcomm_efficient_converter(
             onnx_dir_path=onnx_dir_path,
             seq_length=seq_length,
             full_batch_size=full_batch_size,
+            max_num_adapters=max_num_adapters,
         )
         return onnx_dir_path, generated_onnx_model_path
     else:

QEfficient/exporter/export_utils.py

@@ -83,6 +83,8 @@ def export_onnx(
             dynamic_axes[iname] = {0: dynamic_axis_past_key, 2: "ctx_len"}
         elif iname == "batch_index":
             dynamic_axes[iname] = {0: "batch_size"}
+        elif iname == "lora_ids":
+            dynamic_axes[iname] = {0: "batch_size"}
 
     if "past_key.0" in input_names and "attention_mask" in input_names:
         dynamic_axes["attention_mask"] = {0: "batch_size", 1: "ctx_len"}

QEfficient/generation/text_generation_inference.py

@@ -230,6 +230,7 @@ def cloud_ai_100_exec_kv(
     write_io_dir: Optional[str] = None,
     automation=False,
     full_batch_size: Optional[int] = None,
+    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,13 @@ 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)
+            self.prompt_to_lora_id_mapping_decode = 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.
@@ -460,6 +470,10 @@ 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 and self.full_batch_size is not None:
+            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)
+
         return decode_inputs
 
     def _update_decode_input(self, outputs, position_ids, generation_len, decode_batch_id=None):
@@ -547,6 +561,11 @@ 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:
+            inputs["lora_ids"] = np.array(self.prompt_to_lora_id_mapping_prefill.popleft(), dtype=np.int64).reshape(
+                1, 1
+            )
+
         for i in range(num_chunks):
             chunk_inputs = inputs.copy()
             chunk_inputs["input_ids"] = inputs["input_ids"][
@@ -634,6 +653,12 @@ def run_continuous_batching_decode(self, prompt_queue, generation_len):
                     )
 
                     generated_id_current_index[decode_batch_id] += 1
+
+                    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]
+                        ]
+
         return decode_pause_time
 
     def run_decode(self, decode_inputs, generation_len):

QEfficient/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.lora.auto import QEffAutoLoraModelForCausalLM
+
+__all__ = [
+    "QEffAutoLoraModelForCausalLM",
+]