+ +### Quantization + +Quantization can be done via the API exposed in [quantize](quantize/). To apply quantization, we need to create a configuration that fully expresses how the model should be quantized. This configuration is then passed into conversion, generating a quantized model. + +`quant_recipes.py` contains a list of recipes that are known to be well-supported during runtime. For the average user, this is a good starting point to select the quantization scheme that is best suited for your deployment needs. After identifying the target recipe, the model can be quantized as follows. This example is extracted from `generative/examples/quantize/example.py`. + +``` +quant_config = quant_recipes.full_linear_int8_dynamic_recipe() +edge_model = ai_edge_torch.convert( + model, (tokens, input_pos), quant_config=quant_config +) +``` + +Once converted, you will get a quantized `.tflite` model which will be ready for on-device execution. + +#### Supported schemes + +In the current release, the following schemes are supported: + +* Dynamic range quantization with FP32 activations and INT8 weights for linear ops +* FP16 quantization with FP16 weights and FP32 activations and computation for all ops + +These correspond to the available recipes in `quant_recipes.py` +
+ +### Convert PyTorch LLM to a TFLite model + +Once you re-author the model and validate its numerical accuracy, you can convert the `nn.Module` to TFLite format. Usually for LLMs, there are two entry functions (signatures) we can export: `prefill` and `decode`. Those two signatures only differ in the shape of arguments. + +For example, in the `generative/examples/test_models/toy_model_with_kv_cache.py`, you can define inputs for both signatures: + +Sample inputs for the `prefill` signature: +https://github.com/google-ai-edge/ai-edge-torch-archive/blob/1791dec62f1d3f60e7fe52138640d380f58b072d/ai_edge_torch/generative/examples/test_models/toy_model_with_kv_cache.py#L105-L108 + +Sample inputs for the `decode` signature: +https://github.com/google-ai-edge/ai-edge-torch-archive/blob/1791dec62f1d3f60e7fe52138640d380f58b072d/ai_edge_torch/generative/examples/test_models/toy_model_with_kv_cache.py#L111-L114 + +Then export the model to TFLite with: +https://github.com/google-ai-edge/ai-edge-torch-archive/blob/1791dec62f1d3f60e7fe52138640d380f58b072d/ai_edge_torch/generative/examples/test_models/toy_model_with_kv_cache.py#L133-L139 + +Please note that using the `prefill` and `decode` method conventions are required for easy integration into the Mediapipe LLM Inference API. +
+ +### End-to-End Inference Pipeline + +The model files typically only perform the core ML computation in the LLM pipeline. Deploying the full pipeline requires handling tokenization, sampling and any other pre or post-processing steps required by your system. There are two ways to deploy the converted LLMs on device as part of a full LLM Inference Pipeline. + +#### Use TFLite Runtime APIs + +The user needs to implement the entire LLM Pipeline themselves, and call TFLite Runtime APIs directly to invoke the model. A text generation pipeline typically requires a tokenizer/detokenizer and a sampler, in addition to model inference. The tokenizer converts the input text from a string to a list of integers. The `prefill` signature ingests the sequence of input tokens, and the `decode` signature is invoked to obtain a tensor of logits. The sampler selects a token based on the provided logits, and the decode loop is repeated autoregressively. Ultimately, the detokenizer maps the generated tokens back into human-readable text. + +This approach provides users with the most control. For example, they can implement streaming, get more control over system memory or implement advanced features such as constrained grammar decoding, speculative decoding etc. + +A very simple text generation pipeline based on a decoder-only-transformer is provided [here](https://github.com/google-ai-edge/ai-edge-torch-archive/blob/main/ai_edge_torch/generative/examples/c%2B%2B/text_generator_main.cc) for reference. Note that this example serves as a starting point, and users are expected to implement their own pipelines based on their model's specific requirements. + +#### Use MediaPipe LLM Inference API + +The [MediaPipe LLM Inference API](http://ai.google.dev/edge/mediapipe/solutions/genai/llm_inference) is a high-level API which supports LLM Inference using a prompt-in/prompt-out interface. While it supports some models "out of the box", you can also provide it LLMs converted via our Generative API, and get access to a simple high level interface with Java, and Swift bindings to easily integrate with Mobile Apps. It takes care of all the complexity of implementing the LLM pipeline under the hood, and makes deployment much easier. Unless, you want to explicitly control the pipeline, we would recommend using this for robustness, and ease of use. + +To deploy using the MP LLM Inference API, you need to +* Ensure you convert models using the expected convention of `prefill`, and `decode` functions in the examples. The pipeline only supports `SentencePiece` tokenizer, but it can support a wide variety of models. +* Bundle the converted TFLite files along with some other configurations such as start/stop tokens, tokenizer model etc. See [here](http://ai.google.dev/edge/mediapipe/solutions/genai/llm_inference#ai_edge_model_conversion) +* Once the bundle is created, you can easily invoke the pipeline using the mobile APIs [here](https://ai.google.dev/edge/mediapipe/solutions/genai/llm_inference/android#create_the_task). + +
+ +## Model visualization +### Install the Model Explorer package using the following command: +``` +pip install ai-edge-model-explorer +``` +Detailed install instruction can be found [here](https://github.com/google-ai-edge/model-explorer/wiki/1.-Installation). + +### Visualize the model using CLI +``` +model-explorer 'gemma-2b.tflite' +``` + +
+
+For an end-to-end example showing how to author, convert, quantize and execute, please refer to the steps [here](https://github.com/google-ai-edge/ai-edge-torch-archive/blob/main/ai_edge_torch/generative/examples/README.md)
++ +## What to expect + +### Future Roadmap +* Expanded accleration support on mobile, and web GPUs, and mobile NPUs. +* Advanced quantization approaches suitable for LLMs. +* Expanded support of models, including Diffusion models. +* LoRA support. + +### Known Issues +The following are known product issues we are actively working to fix. + +* The conversion, and serialization process is unoptimized for LLMs. It requires keeping multiple copies of the weights in memory for transformations, and serialization/deserialization. +* Runtime execution of the LLM in TFLite is missing some memory optimizations, and inefficient during memory unpacking on XNNPack. diff --git a/ai_edge_torch/generative/__init__.py b/ai_edge_torch/generative/__init__.py new file mode 100644 index 00000000..57b12003 --- /dev/null +++ b/ai_edge_torch/generative/__init__.py @@ -0,0 +1,14 @@ +# Copyright 2024 The AI Edge Torch Authors. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== diff --git a/ai_edge_torch/generative/examples/README.md b/ai_edge_torch/generative/examples/README.md new file mode 100644 index 00000000..aff62dec --- /dev/null +++ b/ai_edge_torch/generative/examples/README.md @@ -0,0 +1,80 @@ +# Example transformer models (decoder-only LLMs) +Here we provide a list of popular decoder-only LLMs composed via the transformer building blocks from this library. The main purpose is to demonstrate how to construct a new PyTorch LLM model from scratch using the AI Edge Torch Generative API, and convert it to TFLite format for on-device inference. + +## Gemma +Gemma is Google's open-source LLM. The model has both a 2B and 7B versions. See the [model's HuggingFace page](https://huggingface.co/docs/transformers/main/en/model_doc/gemma). The example we provide is Gemma 2B, and the checkpoint for the model can be found [here](https://huggingface.co/google/gemma-2b-it). + +## TinyLlama +[TinyLlama](https://github.com/jzhang38/TinyLlama) is a popular OSS smaller version of Meta's Llama2 model, with only 1.1B parameters. [HuggingFace checkpoint](https://huggingface.co/TinyLlama/TinyLlama-1.1B-Chat-v1.0). + +## Microsoft Phi-2 +Microsoft Phi-2 is also a decoder-only LLM with 2.7B parameters, see details on [HuggingFace](https://huggingface.co/microsoft/phi-2). + +## Overall workflow +To support a new LLM with the Edge Generative API, we need to go through the process of model (re)authoring, checkpoint mapping/loading, model quantization (via PT2E), model conversion to flatbuffer schema, model quality evaluation, benchmarking and on-device inference pipeline authoring. + +### Model (re)authoring +Model (re)authoring refers to the process of a few things: +1) Understanding the overall model architecture (encoder-decoder, decoder-only etc). +2) Compose the model using `ai_edge_torch` provided transformer building blocks. +For each of the example models, we have a model definition file (e.g. tiny_llama/tiny_llama.py) where a `nn.Module` is defined, with its layers and a forward function. There is also a `get_model_config` function which returns a `ModelConfig` instance with hyper-parameters such as embedding size, layer count etc. Finally, there is a `define_and_run` function which builds the model instance, and runs the forward pass with a few sample inputs. + +Here we use `TinyLlama` as an example to walk you through the authoring steps. + +#### Define model's structure +https://github.com/google-ai-edge/ai-edge-torch-archive/blob/e54638dd4a91ec09115f9ded1bd5540f3f1a4e68/ai_edge_torch/generative/examples/tiny_llama/tiny_llama.py#L43-L74 + +#### Define model's forward function +https://github.com/google-ai-edge/ai-edge-torch-archive/blob/e54638dd4a91ec09115f9ded1bd5540f3f1a4e68/ai_edge_torch/generative/examples/tiny_llama/tiny_llama.py#L79-L101 + +Now, you will have an `nn.Module` named `TinyLlama`, the next step is to restore the weights from orginal checkpoint into the new model. + +### Checkpoint mapping/loading +After the model is defined, we need to load the original trained weights to the +new model. This is needed because the `state_dict` of the new model will be +different from the original model's `state_dict`. There are helper functions in +place to simplify the `state_dict` mapping process (`utilities/loader.py`). +The user needs to provide a layer name tempelate (TensorNames) for the source +model. This tempelate is then used to create an updated `state_dict` that works +with the mapped model. The tensor map includes the following fields: +https://github.com/google-ai-edge/ai-edge-torch-archive/blob/3b753d80fdf00872baac523dc727b87b3dc271e7/ai_edge_torch/generative/utilities/loader.py#L120-L134 + +The fields that have a default value of `None` are optional and should only be +populated if they are relevant to the model architecture. For `TinyLlama`, we +will define the following `TENSOR_NAMES`: +https://github.com/google-ai-edge/ai-edge-torch-archive/blob/e54638dd4a91ec09115f9ded1bd5540f3f1a4e68/ai_edge_torch/generative/examples/tiny_llama/tiny_llama.py#L27-L40 + +With the `TensorNames` defined, a user can simply use the loading utils to load +an instance of the mapped model. For instance: + +``` +model = MappedModel(config) +loader = loading_utils.ModelLoader("path_to_checkpoint", TENSOR_NAMES) +loader.load(model) +``` + +Currently, `ModelLoader` supports PyTorch state dictionary and SafeTensors +checkpoints. We recommend testing the mapped model against your reference implementation +using a few input samples before proceeding to the conversion step. + +### Model conversion +In this step, we use the `ai_edge_torch`'s standard multi-signature conversion API to convert PyTorch `nn.Module` to a single TFLite flatbuffer for on-device execution. For example, in `tiny_llama/convert_to_tflite.py`, we use this python code to convert the `TinyLLama` model to a multi-signature TFLite model: +https://github.com/google-ai-edge/ai-edge-torch-archive/blob/3b753d80fdf00872baac523dc727b87b3dc271e7/ai_edge_torch/generative/examples/tiny_llama/convert_to_tflite.py#L22-L53 + +Once converted, you will get a `.tflite` model which will be ready for on-device execution. Note that the `.tflite` model generated uses static shapes. Inside the generated `.tflite` model, there will be two signatures defined (two entrypoints to the model): +1) `prefill`: taking 2 tensor inputs `prefill_tokens`, `prefill_input_pos`. With shape `(BATCH_SIZE, PREFILL_SEQ_LEN)` and `(PREFILL_SEQ_LEN)`. +2) `decode`: taking 2 tensor inputs `decode_token`, `decode_input_pos`. With shape `(1, 1)` and `(1)`. +To learn more about TFLite signatures, please refer to this [article](https://www.tensorflow.org/lite/guide/signatures). + +### Model quantization +To apply quantization, we need to create a configuration that fully expresses how the model should be quantized. This configuration is then passed into conversion, generating a quantized model. + +`quantize/quant_recipes.py` contains a list of recipes that are known to be well-supported during runtime. For the average user, this is a good starting point to select the quantization scheme that is best suited for your deployment needs. After identifying the target recipe, the model can be quantized as follows. This example is extracted from `generative/examples/quantize/example.py`. + +``` +quant_config = quant_recipes.full_linear_int8_dynamic_recipe() +edge_model = ai_edge_torch.convert( + model, (tokens, input_pos), quant_config=quant_config +) +``` +Once converted, you will get a quantized `.tflite` model which will be ready for on-device execution. diff --git a/ai_edge_torch/generative/examples/__init__.py b/ai_edge_torch/generative/examples/__init__.py new file mode 100644 index 00000000..57b12003 --- /dev/null +++ b/ai_edge_torch/generative/examples/__init__.py @@ -0,0 +1,14 @@ +# Copyright 2024 The AI Edge Torch Authors. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== diff --git a/ai_edge_torch/generative/examples/c++/BUILD b/ai_edge_torch/generative/examples/c++/BUILD new file mode 100644 index 00000000..87fc9b2e --- /dev/null +++ b/ai_edge_torch/generative/examples/c++/BUILD @@ -0,0 +1,44 @@ +# Copyright 2024 The AI Edge Torch Authors. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== + +load("@org_tensorflow//tensorflow/lite:build_def.bzl", "tflite_linkopts") + +licenses(["notice"]) + +package(default_visibility = [ + "//visibility:public", +]) + +cc_binary( + name = "text_generator_main", + srcs = [ + "text_generator_main.cc", + ], + linkopts = tflite_linkopts() + select({ + "@org_tensorflow//tensorflow:android": [ + "-pie", # Android 5.0 and later supports only PIE + "-lm", # some builtin ops, e.g., tanh, need -lm + ], + "//conditions:default": [], + }), + deps = [ + "@com_google_absl//absl/flags:flag", + "@com_google_absl//absl/flags:parse", + "@com_google_sentencepiece//:sentencepiece_processor", + "@org_tensorflow//tensorflow/lite:framework", + "@org_tensorflow//tensorflow/lite/experimental/genai:genai_ops", + "@org_tensorflow//tensorflow/lite/kernels:builtin_ops", + ], +) diff --git a/ai_edge_torch/generative/examples/c++/README.md b/ai_edge_torch/generative/examples/c++/README.md new file mode 100644 index 00000000..aba3e07a --- /dev/null +++ b/ai_edge_torch/generative/examples/c++/README.md @@ -0,0 +1,21 @@ +# AI Edge Examples + +This module offers illustrations of how to utilize and run exported models. The examples provided are designed to be concise and have limited dependencies on third-party libraries. Our intention is for developers to leverage these examples as a starting point for integrating the exported models with their unique model-specific pipelines and requirements. + +## Notes: + +* If compiling the examples to run on an Android device, you need to download Android NDK and SDK and set `$ANDROID_NDK_HOME` and `$ANDROID_HOME` environment variables. Please note that _bazel_ currently only supports NDK versions 19, 20, and 21. + +## Text Generation + +In `text_generator_main.cc`, we provide an example of running a decoder-only model end-to-end using TensorFlow Lite as our inference engine. + +To get started, you will need an exported model with two signatures: `prefill` and `decode`. The example takes in an input prompt, tokenizes it, "prefills" the model with the tokens, and decodes autoregressively with greedy sampling until a stop condition is met. Finally, it detokenizes the generated token IDs into text. + +It's important to note that while we use [SentencePiece](https://github.com/google/sentencepiece) as the tokenizer module in our example, it's not a requirement, and other tokenizers can be used as needed. Additionally, we're using a greedy sampling strategy, which simply takes an argmax over the output logits. There are many other options available that have been shown to generate better results. + +As an example, you can run `text_generator_main` for an exported Gemma model as follows: + +``` +bazel run -c opt //ai_edge_torch/generative/examples/c++:text_generator_main -- --tflite_model=PATH/gemma_it.tflite --sentencepiece_model=PATH/tokenizer.model --start_token="
![\"Open](\"https://colab.research.google.com/assets/colab-badge.svg\"){kind=icon}
\n",
+ ""
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 1,
+ "metadata": {
+ "id": "00e00b3b-d7ed-4e2e-815e-3addfc23c8f3"
+ },
+ "outputs": [],
+ "source": [
+ "# Copyright 2024 The AI Edge Torch Authors.\n",
+ "#\n",
+ "# Licensed under the Apache License, Version 2.0 (the \"License\");\n",
+ "# you may not use this file except in compliance with the License.\n",
+ "# You may obtain a copy of the License at\n",
+ "#\n",
+ "# http://www.apache.org/licenses/LICENSE-2.0\n",
+ "#\n",
+ "# Unless required by applicable law or agreed to in writing, software\n",
+ "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
+ "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
+ "# See the License for the specific language governing permissions and\n",
+ "# limitations under the License.\n",
+ "# ==============================================================================\n",
+ "# This is a simple colab showing how to re-author T5 (encoder-decoder) model,\n",
+ "# convert and run in a colab python environment."
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "a9bdc007e6ce"
+ },
+ "source": [
+ "Note: When running notebooks in this repository with Google Colab, some users may see\n",
+ "the following warning message:\n",
+ "\n",
+ "\n",
+ "\n",
+ "Please click `Restart Session` and run again."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {
+ "id": "a91d40b5-91f0-4c19-bdb4-a2f56fa1c5ff"
+ },
+ "outputs": [],
+ "source": [
+ "!pip install -r https://raw.githubusercontent.com/google-ai-edge/ai-edge-torch/main/requirements.txt\n",
+ "!pip install ai-edge-torch"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "WpEGGRs4FJo6"
+ },
+ "source": [
+ "## Download model checkpoint\n",
+ "First we download the T5 pytorch checkpoint from huggingface from https://huggingface.co/humarin/chatgpt_paraphraser_on_T5_base."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {
+ "id": "ivZdosGMowfl"
+ },
+ "outputs": [],
+ "source": [
+ "!curl -O -L https://huggingface.co/humarin/chatgpt_paraphraser_on_T5_base/resolve/main/pytorch_model.bin"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "R5aPvvxeF3bc"
+ },
+ "source": [
+ "## T5 Model Authoring and Conversion\n",
+ "Next, we import the T5 encoder/decoder implementation from `ai_edge_torch/generative/examples/t5`, and convert to TFLite with 2 signatures: `encode` and `decode`."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 2,
+ "metadata": {
+ "id": "YqPe7B8hwGh2"
+ },
+ "outputs": [],
+ "source": [
+ "import numpy as np\n",
+ "import torch\n",
+ "\n",
+ "import ai_edge_torch\n",
+ "from ai_edge_torch.generative.examples.t5 import t5\n",
+ "from ai_edge_torch.generative.quantize import quant_recipes\n",
+ "\n",
+ "\n",
+ "def convert_t5_to_tflite_multisig(checkpoint_path: str):\n",
+ " config = t5.get_model_config_t5()\n",
+ " # Temporarily disable HLFB until custom op issue is fixed.\n",
+ " config.enable_hlfb = False\n",
+ " embedding_layer = torch.nn.Embedding(\n",
+ " config.vocab_size, config.embedding_dim, padding_idx=0\n",
+ " )\n",
+ " t5_encoder_model = t5.build_t5_encoder_model(config, embedding_layer, checkpoint_path)\n",
+ " t5_decoder_model = t5.build_t5_decoder_model(config, embedding_layer, checkpoint_path)\n",
+ "\n",
+ " # encoder\n",
+ " seq_len = 512\n",
+ " prefill_e_tokens = torch.full((1, seq_len), 0, dtype=torch.long)\n",
+ " prompt_e_token = [1, 2, 3, 4, 5, 6]\n",
+ " prefill_e_tokens[0, : len(prompt_e_token)] = torch.tensor(\n",
+ " prompt_e_token, dtype=torch.long\n",
+ " )\n",
+ " prefill_e_input_pos = torch.arange(0, seq_len)\n",
+ " prefill_d_tokens = torch.full((1, seq_len), 0, dtype=torch.long)\n",
+ " prompt_d_token = [1, 2, 3, 4, 5, 6]\n",
+ " prefill_d_tokens[0, : len(prompt_d_token)] = torch.tensor(\n",
+ " prompt_d_token, dtype=torch.long\n",
+ " )\n",
+ " prefill_d_input_pos = torch.arange(0, seq_len)\n",
+ "\n",
+ " # decoder\n",
+ " decode_token = torch.tensor([[1]], dtype=torch.long)\n",
+ " decode_input_pos = torch.tensor([0], dtype=torch.int64)\n",
+ " decode_d_token = torch.tensor([[1]], dtype=torch.long)\n",
+ " decode_d_input_pos = torch.tensor([0], dtype=torch.int64)\n",
+ "\n",
+ " # Pad mask for self attention only on \"real\" tokens.\n",
+ " # Pad with `-inf` for any tokens indices that aren't desired.\n",
+ " pad_mask = torch.zeros([seq_len], dtype=torch.float32)\n",
+ " hidden_states = torch.zeros((1, 512, 768), dtype=torch.float32)\n",
+ " quant_config = quant_recipes.full_linear_int8_dynamic_recipe()\n",
+ "\n",
+ " edge_model = ai_edge_torch.signature(\n",
+ " 'encode',\n",
+ " t5_encoder_model.eval(),\n",
+ " (\n",
+ " prefill_e_tokens,\n",
+ " prefill_e_input_pos,\n",
+ " pad_mask,\n",
+ " ),\n",
+ " ).signature(\n",
+ " 'decode',\n",
+ " t5_decoder_model.eval(),\n",
+ " (\n",
+ " hidden_states,\n",
+ " decode_d_token,\n",
+ " decode_d_input_pos,\n",
+ " pad_mask,\n",
+ " ),\n",
+ " ).convert(quant_config=quant_config)\n",
+ "\n",
+ " edge_model.export('t5_encode_decode_2_sigs.tflite')\n",
+ " return edge_model\n"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "UlgHZtWIGhAc"
+ },
+ "source": [
+ "Finally, we call the convert function, this might take a few minutes to finish."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {
+ "id": "nMCuhDTawncf"
+ },
+ "outputs": [],
+ "source": [
+ "print('converting T5 to tflite.')\n",
+ "edge_model = convert_t5_to_tflite_multisig(\"/content/pytorch_model.bin\")"
+ ]
+ }
+ ],
+ "metadata": {
+ "colab": {
+ "name": "t5_conversion_colab.ipynb",
+ "toc_visible": true
+ },
+ "kernelspec": {
+ "display_name": "Python 3",
+ "name": "python3"
+ }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 0
+}
diff --git a/ai_edge_torch/generative/examples/t5/t5_lm_logits.pt b/ai_edge_torch/generative/examples/t5/t5_lm_logits.pt
new file mode 100644
index 00000000..2a41f0cb
Binary files /dev/null and b/ai_edge_torch/generative/examples/t5/t5_lm_logits.pt differ
diff --git a/ai_edge_torch/generative/examples/test_models/__init__.py b/ai_edge_torch/generative/examples/test_models/__init__.py
new file mode 100644
index 00000000..57b12003
--- /dev/null
+++ b/ai_edge_torch/generative/examples/test_models/__init__.py
@@ -0,0 +1,14 @@
+# Copyright 2024 The AI Edge Torch Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
diff --git a/ai_edge_torch/generative/examples/test_models/toy_model.py b/ai_edge_torch/generative/examples/test_models/toy_model.py
new file mode 100644
index 00000000..d5d013be
--- /dev/null
+++ b/ai_edge_torch/generative/examples/test_models/toy_model.py
@@ -0,0 +1,119 @@
+# Copyright 2024 The AI Edge Torch Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+# A toy example which has a single-layer transformer block.
+from typing import Tuple
+
+import numpy as np
+import torch
+import torch.nn as nn
+
+import ai_edge_torch
+from ai_edge_torch.generative.layers.attention import TransformerBlock
+import ai_edge_torch.generative.layers.attention_utils as attn_utils
+import ai_edge_torch.generative.layers.builder as builder
+import ai_edge_torch.generative.layers.model_config as cfg
+
+RoPECache = Tuple[torch.Tensor, torch.Tensor]
+KV_CACHE_MAX_LEN = 100
+
+
+class ToySingleLayerModel(torch.nn.Module):
+
+ def __init__(self, config: cfg.ModelConfig) -> None:
+ super().__init__()
+ self.lm_head = nn.Linear(
+ config.embedding_dim, config.vocab_size, bias=config.lm_head_use_bias
+ )
+ self.tok_embedding = nn.Embedding(config.vocab_size, config.embedding_dim)
+ self.transformer_block = TransformerBlock(config)
+ self.final_norm = builder.build_norm(
+ config.embedding_dim,
+ config.final_norm_config,
+ )
+ self.rope_cache = attn_utils.build_rope_cache(
+ size=config.max_seq_len,
+ dim=int(config.attn_config.rotary_percentage * config.head_dim),
+ base=10_000,
+ condense_ratio=1,
+ dtype=torch.float32,
+ device=torch.device('cpu'),
+ )
+ self.mask_cache = attn_utils.build_causal_mask_cache(
+ size=config.max_seq_len, dtype=torch.float32, device=torch.device('cpu')
+ )
+ self.config = config
+
+ @torch.inference_mode
+ def forward(self, idx: torch.Tensor, input_pos: torch.Tensor) -> torch.Tensor:
+ x = self.tok_embedding(idx)
+ cos, sin = self.rope_cache
+
+ cos = cos.index_select(0, input_pos)
+ sin = sin.index_select(0, input_pos)
+ mask = self.mask_cache.index_select(2, input_pos)
+ mask = mask[:, :, :, : self.config.max_seq_len]
+
+ x = self.transformer_block(x, (cos, sin), mask, input_pos)
+ x = self.final_norm(x)
+ return self.lm_head(x)
+
+
+def define_and_run() -> None:
+ attn_config = cfg.AttentionConfig(
+ num_heads=32, num_query_groups=4, rotary_percentage=1.0, enable_kv_cache=False
+ )
+ ff_config = cfg.FeedForwardConfig(
+ type=cfg.FeedForwardType.GATED,
+ activation=cfg.ActivationType.SILU,
+ intermediate_size=256,
+ )
+ norm_config = cfg.NormalizationConfig(type=cfg.NormalizationType.RMS_NORM)
+ config = cfg.ModelConfig(
+ vocab_size=400,
+ num_layers=1,
+ max_seq_len=KV_CACHE_MAX_LEN,
+ embedding_dim=128,
+ attn_config=attn_config,
+ ff_config=ff_config,
+ pre_attention_norm_config=norm_config,
+ pre_ff_norm_config=norm_config,
+ final_norm_config=norm_config,
+ )
+
+ model = ToySingleLayerModel(config)
+ idx = torch.unsqueeze(torch.arange(0, KV_CACHE_MAX_LEN), 0)
+ input_pos = torch.arange(0, KV_CACHE_MAX_LEN)
+ print('running an inference')
+ print(
+ model.forward(
+ idx,
+ input_pos,
+ )
+ )
+
+ # Convert model to tflite.
+ print('converting model to tflite')
+ edge_model = ai_edge_torch.convert(
+ model,
+ (
+ idx,
+ input_pos,
+ ),
+ )
+ edge_model.export('/tmp/toy_model.tflite')
+
+
+if __name__ == '__main__':
+ define_and_run()
diff --git a/ai_edge_torch/generative/examples/test_models/toy_model_with_kv_cache.py b/ai_edge_torch/generative/examples/test_models/toy_model_with_kv_cache.py
new file mode 100644
index 00000000..abc88669
--- /dev/null
+++ b/ai_edge_torch/generative/examples/test_models/toy_model_with_kv_cache.py
@@ -0,0 +1,143 @@
+# Copyright 2024 The AI Edge Torch Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+# A toy example which has basic transformer block (w/ KV-Cache).
+from typing import List, Tuple
+
+import numpy as np
+import torch
+import torch.nn as nn
+import torch_xla
+
+import ai_edge_torch
+from ai_edge_torch.generative.layers.attention import TransformerBlock
+import ai_edge_torch.generative.layers.attention_utils as attn_utils
+import ai_edge_torch.generative.layers.builder as builder
+import ai_edge_torch.generative.layers.model_config as cfg
+
+RoPECache = Tuple[torch.Tensor, torch.Tensor]
+
+
+class ToyModelWithKV(torch.nn.Module):
+
+ def __init__(self, config: cfg.ModelConfig) -> None:
+ super().__init__()
+ self.lm_head = nn.Linear(
+ config.embedding_dim, config.vocab_size, bias=config.lm_head_use_bias
+ )
+ self.tok_embedding = nn.Embedding(config.vocab_size, config.embedding_dim)
+ self.transformer_blocks = nn.ModuleList(
+ TransformerBlock(config) for _ in range(config.num_layers)
+ )
+ self.final_norm = builder.build_norm(
+ config.embedding_dim,
+ config.final_norm_config,
+ )
+ self.rope_cache = attn_utils.build_rope_cache(
+ size=config.max_seq_len,
+ dim=int(config.attn_config.rotary_percentage * config.head_dim),
+ base=10_000,
+ condense_ratio=1,
+ dtype=torch.float32,
+ device=torch.device('cpu'),
+ )
+ self.mask_cache = attn_utils.build_causal_mask_cache(
+ size=config.max_seq_len, dtype=torch.float32, device=torch.device('cpu')
+ )
+ self.config = config
+
+ @torch.inference_mode
+ def forward(self, idx: torch.Tensor, input_pos: torch.Tensor) -> torch.Tensor:
+ x = self.tok_embedding(idx)
+ cos, sin = self.rope_cache
+ cos = cos.index_select(0, input_pos)
+ sin = sin.index_select(0, input_pos)
+ mask = self.mask_cache.index_select(2, input_pos)
+ mask = mask[:, :, :, : self.config.max_seq_len]
+ for i, block in enumerate(self.transformer_blocks):
+ x = block(x, (cos, sin), mask, input_pos)
+ x = self.final_norm(x)
+ return self.lm_head(x)
+
+
+def _export_stablehlo_mlir(model, args):
+ ep = torch.export.export(model, args)
+ stablehlo_gm = torch_xla.stablehlo.exported_program_to_stablehlo(ep)
+ return stablehlo_gm.get_stablehlo_text()
+
+
+def get_model_config() -> cfg.ModelConfig:
+ attn_config = cfg.AttentionConfig(
+ num_heads=32, num_query_groups=4, rotary_percentage=1.0
+ )
+ ff_config = cfg.FeedForwardConfig(
+ type=cfg.FeedForwardType.GATED,
+ activation=cfg.ActivationType.SILU,
+ intermediate_size=256,
+ )
+ norm_config = cfg.NormalizationConfig(type=cfg.NormalizationType.RMS_NORM)
+ config = cfg.ModelConfig(
+ vocab_size=150,
+ num_layers=2,
+ max_seq_len=500,
+ embedding_dim=128,
+ attn_config=attn_config,
+ ff_config=ff_config,
+ pre_attention_norm_config=norm_config,
+ pre_ff_norm_config=norm_config,
+ final_norm_config=norm_config,
+ enable_hlfb=True,
+ )
+ return config
+
+
+def get_sample_prefill_inputs() -> Tuple[torch.Tensor, torch.Tensor]:
+ idx = torch.unsqueeze(torch.arange(0, 100), 0)
+ input_pos = torch.arange(0, 100)
+ return idx, input_pos
+
+
+def get_sample_decode_inputs() -> Tuple[torch.Tensor, torch.Tensor]:
+ idx = torch.tensor([[1]], dtype=torch.long)
+ input_pos = torch.tensor([10], dtype=torch.int64)
+ return idx, input_pos
+
+
+def define_and_run() -> None:
+ dump_mlir = False
+
+ config = get_model_config()
+ model = ToyModelWithKV(config)
+ print('running an inference')
+ idx, input_pos = get_sample_prefill_inputs()
+ decode_idx, decode_input_pos = get_sample_decode_inputs()
+ print(model.forward(idx, input_pos))
+
+ if dump_mlir:
+ mlir_text = _export_stablehlo_mlir(model, (idx, input_pos))
+ with open('/tmp/toy_model_with_kv.stablehlo.mlir', 'w') as f:
+ f.write(mlir_text)
+
+ # Convert model to tflite with 2 signatures (prefill + decode).
+ print('converting toy model to tflite with 2 signatures (prefill + decode)')
+ edge_model = (
+ ai_edge_torch.signature('prefill', model, (idx, input_pos))
+ .signature('decode', model, (decode_idx, decode_input_pos))
+ .convert()
+ )
+ edge_model.export('/tmp/toy_kv_cache.tflite')
+
+
+if __name__ == '__main__':
+ define_and_run()
diff --git a/ai_edge_torch/generative/examples/tiny_llama/__init__.py b/ai_edge_torch/generative/examples/tiny_llama/__init__.py
new file mode 100644
index 00000000..e69de29b
diff --git a/ai_edge_torch/generative/examples/tiny_llama/convert_to_tflite.py b/ai_edge_torch/generative/examples/tiny_llama/convert_to_tflite.py
new file mode 100644
index 00000000..21f1ae20
--- /dev/null
+++ b/ai_edge_torch/generative/examples/tiny_llama/convert_to_tflite.py
@@ -0,0 +1,66 @@
+# Copyright 2024 The AI Edge Torch Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+import os
+from pathlib import Path
+
+import torch
+
+import ai_edge_torch
+from ai_edge_torch.generative.examples.tiny_llama import tiny_llama
+from ai_edge_torch.generative.quantize import quant_recipes
+
+
+def convert_tiny_llama_to_tflite(
+ checkpoint_path: str,
+ prefill_seq_len: int = 512,
+ kv_cache_max_len: int = 1024,
+ quantize: bool = True,
+):
+ """An example method for converting TinyLlama model to multi-signature
+ tflite model.
+
+ Args:
+ checkpoint_path (str): The filepath to the model checkpoint, or directory holding the checkpoint.
+ prefill_seq_len (int, optional): The maximum size of prefill input tensor.
+ Defaults to 512.
+ kv_cache_max_len (int, optional): The maximum size of KV cache buffer,
+ including both prefill and decode. Defaults to 1024.
+ quantize (bool, optional): Whether the model should be quanized.
+ Defaults to True.
+ """
+ pytorch_model = tiny_llama.build_model(
+ checkpoint_path, kv_cache_max_len=kv_cache_max_len
+ )
+ # Tensors used to trace the model graph during conversion.
+ prefill_tokens = torch.full((1, prefill_seq_len), 0, dtype=torch.long)
+ prefill_input_pos = torch.arange(0, prefill_seq_len)
+ decode_token = torch.tensor([[0]], dtype=torch.long)
+ decode_input_pos = torch.tensor([0], dtype=torch.int64)
+
+ quant_config = quant_recipes.full_linear_int8_dynamic_recipe() if quantize else None
+ edge_model = (
+ ai_edge_torch.signature(
+ 'prefill', pytorch_model, (prefill_tokens, prefill_input_pos)
+ )
+ .signature('decode', pytorch_model, (decode_token, decode_input_pos))
+ .convert(quant_config=quant_config)
+ )
+ edge_model.export(f'/tmp/tiny_llama_seq{prefill_seq_len}_kv{kv_cache_max_len}.tflite')
+
+
+if __name__ == '__main__':
+ checkpoint_path = os.path.join(Path.home(), 'Downloads/llm_data/tiny_llama')
+ convert_tiny_llama_to_tflite(checkpoint_path)
diff --git a/ai_edge_torch/generative/examples/tiny_llama/tiny_llama.py b/ai_edge_torch/generative/examples/tiny_llama/tiny_llama.py
new file mode 100644
index 00000000..cde23667
--- /dev/null
+++ b/ai_edge_torch/generative/examples/tiny_llama/tiny_llama.py
@@ -0,0 +1,164 @@
+# Copyright 2024 The AI Edge Torch Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+# Example of building a TinyLlama model from the Edge Generative API layers.
+
+import os
+from pathlib import Path
+
+import numpy as np
+import torch
+import torch.nn as nn
+
+from ai_edge_torch.generative.layers.attention import TransformerBlock
+import ai_edge_torch.generative.layers.attention_utils as attn_utils
+import ai_edge_torch.generative.layers.builder as builder
+import ai_edge_torch.generative.layers.model_config as cfg
+import ai_edge_torch.generative.utilities.loader as loading_utils
+
+TENSOR_NAMES = loading_utils.ModelLoader.TensorNames(
+ ff_up_proj="model.layers.{}.mlp.up_proj",
+ ff_down_proj="model.layers.{}.mlp.down_proj",
+ ff_gate_proj="model.layers.{}.mlp.gate_proj",
+ attn_query_proj="model.layers.{}.self_attn.q_proj",
+ attn_key_proj="model.layers.{}.self_attn.k_proj",
+ attn_value_proj="model.layers.{}.self_attn.v_proj",
+ attn_output_proj="model.layers.{}.self_attn.o_proj",
+ pre_attn_norm="model.layers.{}.input_layernorm",
+ pre_ff_norm="model.layers.{}.post_attention_layernorm",
+ embedding="model.embed_tokens",
+ final_norm="model.norm",
+ lm_head="lm_head",
+)
+
+
+class TinyLLamma(nn.Module):
+
+ def __init__(self, config: cfg.ModelConfig):
+ super().__init__()
+
+ self.config = config
+ # Construct model layers.
+ self.lm_head = nn.Linear(
+ config.embedding_dim, config.vocab_size, bias=config.lm_head_use_bias
+ )
+ self.tok_embedding = nn.Embedding(
+ config.vocab_size, config.embedding_dim, padding_idx=0
+ )
+ self.transformer_blocks = nn.ModuleList(
+ TransformerBlock(config) for _ in range(config.num_layers)
+ )
+ self.final_norm = builder.build_norm(
+ config.embedding_dim,
+ config.final_norm_config,
+ )
+ self.rope_cache = attn_utils.build_rope_cache(
+ size=config.kv_cache_max,
+ dim=int(config.attn_config.rotary_percentage * config.head_dim),
+ base=10_000,
+ condense_ratio=1,
+ dtype=torch.float32,
+ device=torch.device("cpu"),
+ )
+ self.mask_cache = attn_utils.build_causal_mask_cache(
+ size=config.kv_cache_max, dtype=torch.float32, device=torch.device("cpu")
+ )
+ self.config = config
+
+ # The model's forward function takes in additional k/v cache tensors
+ # and returns the updated k/v cache tensors to the caller.
+ # This can be eliminated if we handle k/v cache updates inside the model itself.
+ @torch.inference_mode
+ def forward(self, idx: torch.Tensor, input_pos: torch.Tensor) -> torch.Tensor:
+ B, T = idx.size()
+ assert (
+ self.config.max_seq_len >= T
+ ), f"Cannot forward sequence of length {T}, max seq length is only {self.config.max_seq_len}"
+
+ cos, sin = self.rope_cache
+ cos = cos.index_select(0, input_pos)
+ sin = sin.index_select(0, input_pos)
+ mask = self.mask_cache.index_select(2, input_pos)
+ mask = mask[:, :, :, : self.config.kv_cache_max]
+
+ # forward the model itself
+ x = self.tok_embedding(idx) # token embeddings of shape (b, t, n_embd)
+
+ for i, block in enumerate(self.transformer_blocks):
+ x = block(x, (cos, sin), mask, input_pos)
+
+ x = self.final_norm(x)
+
+ res = self.lm_head(x) # (b, t, vocab_size)
+ return res
+
+
+def get_model_config(kv_cache_max_len: int = 1024) -> cfg.ModelConfig:
+ attn_config = cfg.AttentionConfig(
+ num_heads=32,
+ num_query_groups=4,
+ rotary_percentage=1.0,
+ )
+ ff_config = cfg.FeedForwardConfig(
+ type=cfg.FeedForwardType.GATED,
+ activation=cfg.ActivationType.SILU,
+ intermediate_size=5632,
+ )
+ norm_config = cfg.NormalizationConfig(type=cfg.NormalizationType.RMS_NORM)
+ config = cfg.ModelConfig(
+ vocab_size=32000,
+ num_layers=22,
+ max_seq_len=2048,
+ embedding_dim=2048,
+ kv_cache_max_len=kv_cache_max_len,
+ attn_config=attn_config,
+ ff_config=ff_config,
+ pre_attention_norm_config=norm_config,
+ pre_ff_norm_config=norm_config,
+ final_norm_config=norm_config,
+ enable_hlfb=True,
+ )
+ return config
+
+
+def get_fake_model_config_for_test() -> cfg.ModelConfig:
+ config = get_model_config()
+ config.vocab_size = 128
+ config.num_layers = 2
+ config.ff_config.intermediate_size = 256
+ return config
+
+
+def build_model(checkpoint_path, **kwargs) -> nn.Module:
+ config = get_model_config(**kwargs)
+ model = TinyLLamma(config)
+ loader = loading_utils.ModelLoader(checkpoint_path, TENSOR_NAMES)
+ loader.load(model)
+ return model
+
+
+def define_and_run() -> None:
+ kv_cache_max_len = 1024
+ checkpoint_path = os.path.join(Path.home(), "Downloads/llm_data/tiny_llama")
+ model = build_model(checkpoint_path, kv_cache_max_len=kv_cache_max_len)
+ idx = torch.from_numpy(np.array([[1, 2, 3, 4]]))
+ tokens = torch.full((1, kv_cache_max_len), 0, dtype=torch.long, device="cpu")
+ tokens[0, :4] = idx
+ input_pos = torch.arange(0, kv_cache_max_len)
+ print("running an inference")
+ print(model.forward(tokens, input_pos))
+
+
+if __name__ == "__main__":
+ define_and_run()
diff --git a/ai_edge_torch/generative/layers/README.md b/ai_edge_torch/generative/layers/README.md
new file mode 100644
index 00000000..ef06f663
--- /dev/null
+++ b/ai_edge_torch/generative/layers/README.md
@@ -0,0 +1,46 @@
+# ODML Transformer Layers
+Common Pytorch building blocks to re-author transformer models.
+
+## Attention layers
+`attention.py` and `attention_utils.py` contain common building blocks for the attention calculation, which is the key part of transformer models. You can use the abstractions provided here to compose your transformer model.
+
+These two files provide the following common Python helper functions:
+* `scaled_dot_product_attention`: helper function to compute scaled dot product attention on query, key and value tensors.
+* `scaled_dot_product_attention_with_hlfb`: same as `scaled_dot_product_attention` with the addition of HLFB (high-level function boundary) for improved performance.
+* `build_rope_cache`: pre-compute sin and cos values for Rotary Positional Embedding.
+* `build_causal_mask_cache`: build a cache for causal self attention mask.
+
+And also the following `nn.Module` classes:
+* `TransformerBlock`
+* `CausalSelfAttention`
+
+## Builder class for common layers
+In `builder.py`, it provides following helper functions:
+* `build_norm`: constructs different kinds of normalizers based on a config.
+* `build_ff`: constructs different kinds of feed forward layers, which includes Sequential or Gated.
+
+## Feed forward layer
+The library provides the following `nn.Modules` to represent feed forward layer.
+* `SequentialFeedForward`
+* `GatedFeedForward`
+
+## KV cache layer
+We provide a `nn.Module` KVCache to express the logic to update the cache. It also has internal logic to apply HLFB to ensure high-performance at runtime.
+
+## Model Configuration class
+Currently, the library provides the following configuration class for you to customize the transformer model:
+* `AttentionConfig`
+* `FeedForwardConfig`
+* `NormalizationConfig`
+* `ModelConfig`
+
+## Normalization layer
+`normalization.py` provides normalization modules currently not supported by Pytorch such as `RMSNorm`:
+* `RMSNorm`
+
+## RoPE Embedding
+`rotary_position_embedding.py` contains helper functions for applying RoPE to tensors.
+
+## High-Level function boundary for performance
+We introduce High-Level Function Boundary (HLFB) as a way of annotating performance-critical pieces of the model (e.g. `scaled_dot_product_attention`, or `KVCache`). HLFB allows the converter to lower the annotated blocks to performant TFLite custom ops. Following is an example of applying HLFB to `SDPA`:
+https://github.com/google-ai-edge/ai-edge-torch-archive/blob/3b753d80fdf00872baac523dc727b87b3dc271e7/ai_edge_torch/generative/layers/attention.py#L74-L122
diff --git a/ai_edge_torch/generative/layers/__init__.py b/ai_edge_torch/generative/layers/__init__.py
new file mode 100644
index 00000000..57b12003
--- /dev/null
+++ b/ai_edge_torch/generative/layers/__init__.py
@@ -0,0 +1,14 @@
+# Copyright 2024 The AI Edge Torch Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
diff --git a/ai_edge_torch/generative/layers/attention.py b/ai_edge_torch/generative/layers/attention.py
new file mode 100644
index 00000000..b6220e1e
--- /dev/null
+++ b/ai_edge_torch/generative/layers/attention.py
@@ -0,0 +1,288 @@
+# Copyright 2024 The AI Edge Torch Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+# Common building blocks for Attention layer.
+
+import math
+from typing import Optional, Tuple
+
+import torch
+from torch import nn
+import torch.nn.functional as F
+
+import ai_edge_torch.generative.layers.builder as builder
+from ai_edge_torch.generative.layers.kv_cache import KVCache
+import ai_edge_torch.generative.layers.model_config as cfg
+import ai_edge_torch.generative.layers.rotary_position_embedding as rotary_pos_emb
+from ai_edge_torch.hlfb import StableHLOCompositeBuilder
+
+
+def scaled_dot_product_attention(
+ q: torch.Tensor,
+ k: torch.Tensor,
+ v: torch.Tensor,
+ head_size: int,
+ mask: Optional[torch.Tensor] = None,
+ scale: Optional[float] = None,
+):
+ """Scaled dot product attention.
+
+ Args:
+ q (torch.Tensor): Query tensor, with shape [B, T, N, H].
+ k (torch.Tensor): Key tensor, with shape [B, T, KV_LEN, H].
+ v (torch.Tensor): Value tensor, with shape [B, T, KV_LEN, H].
+ head_size (int): head dimension.
+ mask (torch.Tensor): the optional mask tensor.
+
+ Returns:
+ The output tensor of scaled_dot_product_attention.
+ """
+
+ if scale is None:
+ scale = 1.0 / math.sqrt(head_size)
+
+ q = q.transpose(1, 2)
+ k = k.transpose(1, 2)
+ v = v.transpose(1, 2)
+ if q.size() != k.size():
+ # Handle the GQA case, where q.shape[1] % k.shape[1] == 0.
+ k = k.repeat_interleave(q.shape[1] // k.shape[1], dim=1)
+ v = v.repeat_interleave(q.shape[1] // v.shape[1], dim=1)
+ y = F.scaled_dot_product_attention(
+ q,
+ k,
+ v,
+ attn_mask=mask,
+ dropout_p=0.0,
+ is_causal=mask is None,
+ scale=scale,
+ )
+ return y.transpose(1, 2)
+
+
+def scaled_dot_product_attention_with_hlfb(
+ q: torch.Tensor,
+ k: torch.Tensor,
+ v: torch.Tensor,
+ head_size: int,
+ mask: Optional[torch.Tensor] = None,
+ scale: Optional[float] = None,
+):
+ """Scaled dot product attention with high-level function boundary enabled.
+
+ Args:
+ q (torch.Tensor): Query tensor, with shape [B, T, N, H].
+ k (torch.Tensor): Key tensor, with shape [B, T, KV_LEN, H].
+ v (torch.Tensor): Value tensor, with shape [B, T, KV_LEN, H].
+ head_size (int): head dimension.
+ mask (torch.Tensor): the optional mask tensor.
+
+ Returns:
+ The output tensor of scaled_dot_product_attention.
+ """
+
+ if scale is None:
+ scale = 1.0 / math.sqrt(head_size)
+
+ builder = StableHLOCompositeBuilder(
+ name="odml.scaled_dot_product_attention", attr={"scale": scale}
+ )
+ q, k, v, mask = builder.mark_inputs(q, k, v, mask)
+
+ q = q.transpose(1, 2)
+ k = k.transpose(1, 2)
+ v = v.transpose(1, 2)
+ if q.size() != k.size():
+ # Handle the GQA case, where q.shape[1] % k.shape[1] == 0.
+ k = k.repeat_interleave(q.shape[1] // k.shape[1], dim=1)
+ v = v.repeat_interleave(q.shape[1] // v.shape[1], dim=1)
+ y = F.scaled_dot_product_attention(
+ q,
+ k,
+ v,
+ attn_mask=mask,
+ dropout_p=0.0,
+ is_causal=mask is None,
+ scale=scale,
+ )
+
+ result = y.transpose(1, 2)
+ result = builder.mark_outputs(result)
+ return result
+
+
+class TransformerBlock(nn.Module):
+
+ def __init__(self, config: cfg.ModelConfig) -> None:
+ """Initialize an instance of the TransformerBlock.
+
+ Args:
+ config (cfg.ModelConfig): the configuration object
+ for this transformer block.
+ """
+
+ super().__init__()
+ self.pre_atten_norm = builder.build_norm(
+ config.embedding_dim, config.pre_attention_norm_config
+ )
+ self.atten_func = CausalSelfAttention(
+ config.embedding_dim,
+ config.attn_config,
+ config.kv_cache_max,
+ config.enable_hlfb,
+ )
+ self.pre_ff_norm = builder.build_norm(
+ config.embedding_dim, config.pre_ff_norm_config
+ )
+ self.ff = builder.build_ff(config.embedding_dim, config.ff_config)
+ self.config = config
+
+ def forward(
+ self,
+ x: torch.Tensor,
+ rope: Tuple[torch.Tensor, torch.Tensor],
+ mask: Optional[torch.Tensor] = None,
+ input_pos: Optional[torch.Tensor] = None,
+ ) -> torch.Tensor:
+ """Forward function of the TransformerBlock.
+
+ Args:
+ x (torch.Tensor): the input tensor.
+ rope (Tuple[torch.Tensor, torch.Tensor]): the input rope tensor.
+ mask (torch.Tensor): the optional mask tensor.
+ input_pos (torch.Tensor): the optional input position tensor.
+
+ Returns:
+ output activation from this transformer block.
+ """
+
+ if self.config.parallel_residual:
+ x_norm = self.pre_atten_norm(x)
+ attn_out = self.atten_func(x_norm, rope, mask, input_pos)
+ ff_out = self.ff(x_norm)
+ output = x + attn_out + ff_out
+ else:
+ x_norm = self.pre_atten_norm(x)
+ attn_out = self.atten_func(x_norm, rope, mask, input_pos)
+ x = x + attn_out
+ x_norm = self.pre_ff_norm(x)
+ output = x + self.ff(x_norm)
+
+ return output
+
+
+# CausalSelfAttention which can support MHQ, MQA or GQA.
+class CausalSelfAttention(nn.Module):
+
+ def __init__(
+ self,
+ dim: int,
+ config: cfg.AttentionConfig,
+ kv_cache_max: int,
+ enable_hlfb: bool,
+ ) -> None:
+ """Initialize an instance of CausalSelfAttention.
+
+ Args:
+ dim (int): causal attention's input/output dimmension.
+ config (cfg.AttentionConfig): attention specific configurations.
+ kv_cache_max (int): determines the size of the KV Cache buffer, if enabled.
+ enable_hlfb (bool): whether hlfb is enabled or not.
+ """
+ super().__init__()
+ self.head_dim = dim // config.num_heads
+ shape = (config.num_heads + 2 * config.num_query_groups) * self.head_dim
+ # Key, query, value projections for all heads.
+ self.qkv_projection = nn.Linear(dim, shape, bias=config.qkv_use_bias)
+ self.output_projection = nn.Linear(dim, dim, bias=config.output_proj_use_bias)
+ self.config = config
+ self.kv_cache = None
+
+ # Build a k/v cache with size (batch_size, kv_cache_max, n_heads, head_dim).
+ # Now only supports batch_size of 1.
+ # TODO(haoliang): support batch_size greater than 1.
+ if config.enable_kv_cache:
+ self.kv_cache = KVCache(
+ 1,
+ kv_cache_max,
+ config.num_query_groups,
+ self.head_dim,
+ enable_hlfb,
+ )
+
+ if enable_hlfb:
+ self.sdpa_func = scaled_dot_product_attention_with_hlfb
+ else:
+ self.sdpa_func = scaled_dot_product_attention
+
+ def forward(
+ self,
+ x: torch.Tensor,
+ rope: Tuple[torch.Tensor, torch.Tensor],
+ mask: Optional[torch.Tensor] = None,
+ input_pos: Optional[torch.Tensor] = None,
+ ) -> torch.Tensor:
+ """Forward function of the CausalSelfAttention layer.
+
+ Args:
+ x (torch.Tensor): the input tensor.
+ rope (Tuple[torch.Tensor, torch.Tensor]): the input rope tensor.
+ mask (torch.Tensor): the optional mask tensor.
+ input_pos (torch.Tensor): the optional input position tensor.
+
+ Returns:
+ output activation from this self attention layer.
+ """
+ # Batch size, sequence length, embedding dimensionality.
+ B, T, E = x.size()
+ assert B == 1, "Currently only batch_size = 1 is supported."
+
+ qkv = self.qkv_projection(x)
+
+ # Assemble into a number of query groups to support MHA, MQA and GQA.
+ q_per_kv = self.config.num_heads // self.config.num_query_groups
+ total_qkv = q_per_kv + 2 # Each group has >=1 queries, 1 key, and 1 value.
+ qkv = qkv.view(
+ B, T, self.config.num_query_groups, total_qkv, self.head_dim
+ ) # (B, T, num_query_groups, total_qkv, head_dim)
+
+ # Split batched computation into three.
+ q, k, v = qkv.split((q_per_kv, 1, 1), dim=-2)
+
+ q = q.reshape(B, T, -1, self.head_dim)
+ k = k.reshape(B, T, -1, self.head_dim)
+ v = v.reshape(B, T, -1, self.head_dim)
+
+ # Compute rotary positional embedding for query and key.
+ n_elem = int(self.config.rotary_percentage * self.head_dim)
+ cos, sin = rope
+ q_roped = rotary_pos_emb.apply_rope(
+ q[..., :n_elem], cos.repeat(1, 2), sin.repeat(1, 2)
+ )
+ k_roped = rotary_pos_emb.apply_rope(
+ k[..., :n_elem], cos.repeat(1, 2), sin.repeat(1, 2)
+ )
+ q = torch.cat((q_roped, q[..., n_elem:]), dim=-1)
+ k = torch.cat((k_roped, k[..., n_elem:]), dim=-1)
+
+ if self.kv_cache is not None:
+ # TODO(haoliang): Handle when execeeding max sequence length.
+ k, v = self.kv_cache.update_cache(input_pos, k, v)
+
+ y = self.sdpa_func(q, k, v, self.head_dim, mask=mask)
+ y = y.reshape(B, T, E)
+
+ # Compute the output projection.
+ y = self.output_projection(y)
+ return y
diff --git a/ai_edge_torch/generative/layers/attention_utils.py b/ai_edge_torch/generative/layers/attention_utils.py
new file mode 100644
index 00000000..73b919b6
--- /dev/null
+++ b/ai_edge_torch/generative/layers/attention_utils.py
@@ -0,0 +1,169 @@
+# Copyright 2024 The AI Edge Torch Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+# Common utility functions used with attention module.
+
+import math
+from typing import Tuple
+
+import torch
+
+
+def build_rope_cache(
+ size: int,
+ dim: int,
+ base: int = 10000,
+ condense_ratio: int = 1,
+ dtype: torch.dtype = torch.float32,
+ device: torch.device = None,
+) -> Tuple[torch.Tensor, torch.Tensor]:
+ """Precompute Rotary Positional Embedding Sin and Cos values for quick lookups
+ during the inference.
+
+ Args:
+ size (int): The size of the built cache.
+ dim (int): Each sequence's dimmension.
+ base (int, optional): Rope base value. Defaults to 10000.
+ condense_ratio (int, optional): The ratio by which sequence indicies are
+ condensed. Defaults to 1.
+ dtype (torch.dtype, optional): Output tensor's data type. Defaults to
+ torch.float32.
+ device (torch.device, optional): Output tensor's data type. Defaults to
+ None in which case "cpu" is used.
+
+ Returns:
+ Tuple[torch.Tensor, torch.Tensor]: Rope's Cosine and Sine waves.
+ """
+ if device is None:
+ device = torch.device('cpu')
+ theta = 1.0 / (base ** (torch.arange(0, dim, 2).float() / dim))
+ seq_idx = torch.arange(size) / condense_ratio
+ idx_theta = torch.outer(seq_idx, theta)
+ cos = torch.cos(idx_theta).to(dtype=dtype, device=device)
+ sin = torch.sin(idx_theta).to(dtype=dtype, device=device)
+ return cos, sin
+
+
+def build_causal_mask_cache(
+ size: int,
+ dtype: torch.dtype = torch.float32,
+ device: torch.device = None,
+) -> torch.Tensor:
+ """Build a cache for causal attention mask.
+
+ Args:
+ size (int): The size of the built mask cache.
+ dtype (torch.dtype, optional): Output tensor's data type. Defaults to
+ torch.float32.
+ device (torch.device, optional): Output tensor's data type. Defaults to
+ None in which case "cpu" is used.
+
+ Returns:
+ torch.Tensor: Causal attention mask.
+ """
+ if device is None:
+ device = torch.device('cpu')
+ mask = torch.full((size, size), float('-inf'), dtype=dtype, device=device)
+ return torch.triu(mask, diagonal=1).unsqueeze(0).unsqueeze(0)
+
+
+def relative_position_bucket(
+ relative_position: torch.Tensor,
+ bidirectional: bool,
+ num_buckets: int,
+ max_distance: int,
+) -> torch.Tensor:
+ """
+ Adapted from Mesh Tensorflow:
+ https://github.com/tensorflow/mesh/blob/0cb87fe07da627bf0b7e60475d59f95ed6b5be3d/mesh_tensorflow/transformer/transformer_layers.py#L593
+
+ Translate relative position to a bucket number for relative attention. The relative position is defined as
+ memory_position - query_position, i.e. the distance in tokens from the attending position to the attended-to
+ position. If bidirectional=False, then positive relative positions are invalid. We use smaller buckets for
+ small absolute relative_position and larger buckets for larger absolute relative_positions. All relative
+ positions >=max_distance map to the same bucket. All relative positions <=-max_distance map to the same bucket.
+ This should allow for more graceful generalization to longer sequences than the model has been trained on
+
+ Args:
+ relative_position: an int32 Tensor
+ bidirectional: a boolean - whether the attention is bidirectional
+ num_buckets: an integer for number of buckets.
+ max_distance: an integer for max distance.
+
+ Returns:
+ a Tensor with the same shape as relative_position, containing int32 values in the range [0, num_buckets)
+ """
+ relative_buckets = 0
+ if bidirectional:
+ num_buckets //= 2
+ relative_buckets += (relative_position > 0).to(torch.long) * num_buckets
+ relative_position = torch.abs(relative_position)
+ else:
+ relative_position = -torch.min(
+ relative_position, torch.zeros_like(relative_position)
+ )
+ # now relative_position is in the range [0, inf)
+
+ # half of the buckets are for exact increments in positions
+ max_exact = num_buckets // 2
+ is_small = relative_position < max_exact
+
+ # The other half of the buckets are for logarithmically bigger bins in positions up to max_distance
+ relative_position_if_large = max_exact + (
+ torch.log(relative_position.float() / max_exact)
+ / math.log(max_distance / max_exact)
+ * (num_buckets - max_exact)
+ ).to(torch.long)
+ relative_position_if_large = torch.min(
+ relative_position_if_large,
+ torch.full_like(relative_position_if_large, num_buckets - 1),
+ )
+
+ relative_buckets += torch.where(
+ is_small, relative_position, relative_position_if_large
+ )
+ return relative_buckets
+
+
+def build_relative_position_buckets(
+ query_length: int,
+ key_length: int,
+ bidirectional: bool = True,
+ num_buckets: int = 32,
+ max_distance: int = 128,
+) -> torch.Tensor:
+ """Relative position buckets for computing bias.
+
+ Args:
+ query_length: an integer of length of current query tensor.
+ key_length: an integer of length of current key tensor.
+ bidirectional: a boolean - whether the attention is bidirectional, default is True.
+ num_buckets: an integer for number of buckets, default is 32.
+ max_distance: an integer for max distance, default is 128.
+
+ Returns:
+ A torch.Tensor of computed relative position buckets.
+ """
+ context_position = torch.arange(query_length, dtype=torch.long)[:, None]
+ memory_position = torch.arange(key_length, dtype=torch.long)[None, :]
+ relative_position = (
+ memory_position - context_position
+ ) # shape (query_length, key_length)
+ rel_pos_bucket = relative_position_bucket(
+ relative_position, # shape (query_length, key_length)
+ bidirectional=bidirectional,
+ num_buckets=num_buckets,
+ max_distance=max_distance,
+ )
+ return rel_pos_bucket.unsqueeze(0).unsqueeze(0)
diff --git a/ai_edge_torch/generative/layers/builder.py b/ai_edge_torch/generative/layers/builder.py
new file mode 100644
index 00000000..6b12a274
--- /dev/null
+++ b/ai_edge_torch/generative/layers/builder.py
@@ -0,0 +1,103 @@
+# Copyright 2024 The AI Edge Torch Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+# Builder class for individual components.
+from torch import nn
+import torch.nn.functional as F
+
+import ai_edge_torch.generative.layers.feed_forward as feed_forward
+import ai_edge_torch.generative.layers.model_config as cfg
+import ai_edge_torch.generative.layers.normalization as normalization
+
+
+def build_norm(dim: int, config: cfg.NormalizationConfig):
+ """Builder function for normalizers.
+
+ Args:
+ dim (int): dimension of the input tensor.
+ config (`NormalizationConfig` object): the normalization configuration.
+
+ Returns:
+ The constructed `nn.Module` normalization layer.
+
+ Raises:
+ ValueError: If config's `layer_norm_type` is not supported.
+ """
+ if config.type == cfg.NormalizationType.NONE:
+ return lambda x: x
+ elif config.type == cfg.NormalizationType.RMS_NORM:
+ return normalization.RMSNorm(
+ dim,
+ eps=config.epsilon,
+ zero_centered_gamma=config.zero_centered,
+ )
+ elif config.type == cfg.NormalizationType.LAYER_NORM:
+ return nn.LayerNorm(dim, eps=config.epsilon)
+ else:
+ raise ValueError("Unsupported norm type.")
+
+
+def build_ff(dim: int, config: cfg.FeedForwardConfig):
+ """Builder function for Feed Forward. Supports `Sequential` and `Gated`.
+
+ Args:
+ dim (int): dimension of the input tensor.
+ config (`ModelConfig` object): the model configuration.
+
+ Returns:
+ The constructed `nn.Module` feedforward layer.
+
+ Raises:
+ ValueError: If config's `ff_type` is not supported.
+ """
+ ff_type = config.type
+ if ff_type == cfg.FeedForwardType.SEQUENTIAL:
+ ff_module = feed_forward.SequentialFeedForward
+ elif ff_type == cfg.FeedForwardType.GATED:
+ ff_module = feed_forward.GatedFeedForward
+ else:
+ raise ValueError("Unsupported feedforward type.")
+
+ activation = _get_activation(config.activation)
+
+ return ff_module(
+ dim=dim,
+ hidden_dim=config.intermediate_size,
+ activation=activation,
+ use_bias=config.use_bias,
+ )
+
+
+def _get_activation(type_: cfg.ActivationType):
+ """Get pytorch callable activation from the name.
+
+ Args:
+ name (string): activation's name.
+
+ Returns:
+ Activation function.
+
+ Raises:
+ ValueError: If activation name is not supported.
+ """
+ if type_ == cfg.ActivationType.SILU:
+ return F.silu
+ elif type_ == cfg.ActivationType.GELU:
+ return F.gelu
+ elif type_ == cfg.ActivationType.GELU_TANH:
+ return lambda x: F.gelu(x, approximate="tanh")
+ elif type_ == cfg.ActivationType.RELU:
+ return F.relu
+ else:
+ raise ValueError("Unsupported activation type.")
diff --git a/ai_edge_torch/generative/layers/feed_forward.py b/ai_edge_torch/generative/layers/feed_forward.py
new file mode 100644
index 00000000..b7f0a78e
--- /dev/null
+++ b/ai_edge_torch/generative/layers/feed_forward.py
@@ -0,0 +1,95 @@
+# Copyright 2024 The AI Edge Torch Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+# Common building blocks for FeedForward layers.
+
+from typing import Callable
+
+import torch
+from torch import nn
+import torch.nn.functional as F
+
+
+class SequentialFeedForward(nn.Module):
+ """Vanilla sequential Feedforward with customizable activation."""
+
+ def __init__(
+ self,
+ dim: int,
+ hidden_dim: int,
+ activation: Callable[[torch.Tensor], torch.Tensor],
+ use_bias=False,
+ ):
+ """Init function for feedforward layer.
+
+ Args:
+ dim(int): embedding size.
+ hidden_dim(int): hidden dim size of the feedforward layer.
+ activation(Callable): activation function used in this block.
+ use_bias(Boolean): whether to use bias. Default is false.
+ """
+ super().__init__()
+ self.act = activation
+ self.w1 = nn.Linear(dim, hidden_dim, bias=use_bias)
+ self.w2 = nn.Linear(hidden_dim, dim, bias=use_bias)
+
+ def forward(self, x):
+ """Forward pass for Feedforward layer.
+
+ Args:
+ x (torch.Tensor): the input tensor.
+
+ Returns:
+ torch.Tensor: output tensor after feedforward.
+ """
+ return self.w2(self.act(self.w1(x)))
+
+
+class GatedFeedForward(nn.Module):
+ """Gated Feedforward with customizable activation.
+
+ https://arxiv.org/pdf/2002.05202v1.pdf
+ """
+
+ def __init__(
+ self,
+ dim: int,
+ hidden_dim: int,
+ activation: Callable[[torch.Tensor], torch.Tensor],
+ use_bias=False,
+ ):
+ """Init function for feedforward layer.
+
+ Args:
+ dim(int): embedding size.
+ hidden_dim(int): hidden dim size of the feedforward layer.
+ activation(Callable): activation function used in this block.
+ use_bias(Boolean): whether to use bias. Default is false.
+ """
+ super().__init__()
+ self.act = activation
+ self.w1 = nn.Linear(dim, hidden_dim, bias=use_bias)
+ self.w2 = nn.Linear(hidden_dim, dim, bias=use_bias)
+ self.w3 = nn.Linear(dim, hidden_dim, bias=use_bias)
+
+ def forward(self, x):
+ """Forward pass for Feedforward layer.
+
+ Args:
+ x (torch.Tensor): the input tensor.
+
+ Returns:
+ torch.Tensor: output tensor after feedforward.
+ """
+ return self.w2(self.act(self.w1(x)) * self.w3(x))
diff --git a/ai_edge_torch/generative/layers/kv_cache.py b/ai_edge_torch/generative/layers/kv_cache.py
new file mode 100644
index 00000000..ae54cb7e
--- /dev/null
+++ b/ai_edge_torch/generative/layers/kv_cache.py
@@ -0,0 +1,83 @@
+# Copyright 2024 The AI Edge Torch Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+# `nn.Module` which implements a KV cache.
+
+import torch
+from torch import nn
+import torch_xla
+
+from ai_edge_torch.hlfb import StableHLOCompositeBuilder
+
+
+class KVCache(nn.Module):
+
+ def __init__(self, batch_size, kv_cache_max, n_heads, head_dim, enable_hlfb=False):
+ """Initializes the KVCache layer.
+
+ Args:
+ batch_size (int): batch size. Currently only batch size 1 is supported.
+ kv_cache_max (int): the max length of KV cache.
+ n_heads (int): number of kv heads.
+ head_dim (int): the head dimension size.
+ enable_hlfb (bool): whether hlfb is enabled or not.
+ """
+ super().__init__()
+ cache_shape = (batch_size, kv_cache_max, n_heads, head_dim)
+ self.register_buffer("k_cache", torch.zeros(cache_shape), persistent=False)
+ self.register_buffer("v_cache", torch.zeros(cache_shape), persistent=False)
+ self.enable_hlfb = enable_hlfb
+ self.kv_cache_max = kv_cache_max
+
+ def update_cache(self, input_pos, k_val, v_val):
+ """Update an entry in the KV cache.
+
+ Args:
+ input_pos (torch.Tensor): the input position.
+ k_val (torch.Tensor): the new `key` value.
+ v_val (torch.Tensor): the new `value` value.
+
+ Returns:
+ The updated key and value tensor.
+ """
+ if self.enable_hlfb:
+ return self.update_cache_with_hlfb(input_pos, k_val, v_val)
+
+ updated_k = self.k_cache.index_copy_(1, input_pos, k_val)
+ updated_v = self.v_cache.index_copy_(1, input_pos, v_val)
+ # Here we need a clone otherwise dynamo export will fail.
+ return torch.clone(updated_k), torch.clone(updated_v)
+
+ def update_cache_with_hlfb(self, input_pos, k_val, v_val):
+ """Update an entry in the KV cache and enable high-level function boundary.
+
+ Args:
+ input_pos (torch.Tensor): the input position.
+ k_val (torch.Tensor): the new `key` value.
+ v_val (torch.Tensor): the new `value` value.
+
+ Returns:
+ The updated key and value tensor.
+ """
+
+ builder = StableHLOCompositeBuilder(
+ name="odml.update_kv_cache", attr={"kv_cache_max": self.kv_cache_max}
+ )
+ k_cache, v_cache, input_pos, k_val, v_val = builder.mark_inputs(
+ self.k_cache, self.v_cache, input_pos, k_val, v_val
+ )
+ updated_k = k_cache.index_copy_(1, input_pos, k_val)
+ updated_v = v_cache.index_copy_(1, input_pos, v_val)
+ updated_k, updated_v = builder.mark_outputs(updated_k, updated_v)
+ return updated_k, updated_v
diff --git a/ai_edge_torch/generative/layers/model_config.py b/ai_edge_torch/generative/layers/model_config.py
new file mode 100644
index 00000000..f8796bc8
--- /dev/null
+++ b/ai_edge_torch/generative/layers/model_config.py
@@ -0,0 +1,135 @@
+# Copyright 2024 The AI Edge Torch Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+# Model configuration class.
+from dataclasses import dataclass
+from dataclasses import field
+import enum
+from typing import Optional
+
+
+@enum.unique
+class ActivationType(enum.Enum):
+ """Different activation functions supported by the default builder."""
+
+ LINEAR = enum.auto()
+ SILU = enum.auto()
+ GELU = enum.auto()
+ GELU_TANH = enum.auto()
+ RELU = enum.auto()
+
+
+@enum.unique
+class NormalizationType(enum.Enum):
+ """Different normalization functions"""
+
+ # No normalization is applied.
+ NONE = enum.auto()
+ RMS_NORM = enum.auto()
+ LAYER_NORM = enum.auto()
+
+
+@enum.unique
+class FeedForwardType(enum.Enum):
+ """Different variations of the Feed Forward module."""
+
+ # `output = linear(act(linear(x)))`.
+ SEQUENTIAL = enum.auto()
+ # `output = linear(act(linear(x)) * lienar(x))`.
+ GATED = enum.auto()
+
+
+@dataclass
+class AttentionConfig:
+ """Attention moduel's parameters."""
+
+ num_heads: int
+ # Used to determine number of groups in grouped query attention (GQA)
+ # https://arxiv.org/pdf/2305.13245.pdf
+ num_query_groups: Optional[int]
+ # Percentage of Rotary Positional Embedding added Q and K projections.
+ rotary_percentage: Optional[float] = None
+ # Whether to use bias with Query, Key, and Value projection.
+ qkv_use_bias: bool = False
+ # Whether to use bias with attention output projection.
+ output_proj_use_bias: bool = False
+ enable_kv_cache: bool = True
+ relative_attention_num_buckets: int = 0
+ relative_attention_max_distance: int = 0
+
+
+@dataclass
+class FeedForwardConfig:
+ """FeedForward module's parameters."""
+
+ type: FeedForwardType
+ activation: ActivationType
+ intermediate_size: int
+ use_bias: bool = False
+
+
+@dataclass
+class NormalizationConfig:
+ """Normalizater parameters."""
+
+ type: NormalizationType = NormalizationType.NONE
+ epsilon: float = 1e-5
+ zero_centered: bool = False
+
+
+@dataclass
+class ModelConfig:
+ """Base configurations for building a transformer architecture."""
+
+ vocab_size: int
+ num_layers: int
+ max_seq_len: int
+ embedding_dim: int
+
+ attn_config: AttentionConfig
+ ff_config: FeedForwardConfig
+ # The normalization applied to attention's input.
+ pre_attention_norm_config: NormalizationConfig = field(
+ default_factory=NormalizationConfig
+ )
+ # The normalization applied to feed forward's input.
+ pre_ff_norm_config: NormalizationConfig = field(default_factory=NormalizationConfig)
+ # The normalization applied before LM head.
+ final_norm_config: NormalizationConfig = field(default_factory=NormalizationConfig)
+
+ # If set to True, only pre_attention_norm is applied to the input and the
+ # decode's output is computed as `output = input + attn_out + ff_out` where
+ # attention and feed forward are called with pre_attention_norm's output.
+ parallel_residual: bool = False
+ # Use bias term within LLM's HEAD.
+ lm_head_use_bias: bool = False
+ # Whether to turn on high-level function boundary.
+ enable_hlfb: bool = False
+
+ # The maximum sequence length of the KV cache. Should not exceed max_seq_len.
+ kv_cache_max_len: int = 0
+
+ # The Attention computation will include relative positional bias.
+ relative_attention: bool = False
+
+ @property
+ def kv_cache_max(self) -> int:
+ if self.kv_cache_max_len > 0:
+ return self.kv_cache_max_len
+ else:
+ return self.max_seq_len
+
+ @property
+ def head_dim(self) -> int:
+ return self.embedding_dim // self.attn_config.num_heads
diff --git a/ai_edge_torch/generative/layers/normalization.py b/ai_edge_torch/generative/layers/normalization.py
new file mode 100644
index 00000000..86e90daa
--- /dev/null
+++ b/ai_edge_torch/generative/layers/normalization.py
@@ -0,0 +1,62 @@
+# Copyright 2024 The AI Edge Torch Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+# Common normalization layers.
+
+import torch
+
+
+# Implementation for RMSNorm from: https://arxiv.org/abs/1910.07467
+class RMSNorm(torch.nn.Module):
+
+ def __init__(self, dim: int, eps: float = 1e-6, zero_centered_gamma=False):
+ """
+ Initialize the RMSNorm layer.
+
+ Args:
+ dim (int): dimension of the input tensor.
+ eps (float): A small float value to ensure numerical stability (default: 1e-6).
+ """
+ super().__init__()
+ self.eps = eps
+ self.weight = torch.nn.Parameter(torch.ones(dim))
+ self.zero_centered_gamma = zero_centered_gamma
+
+ def _norm(self, x):
+ """
+ Apply RMSNorm normalization.
+
+ Args:
+ x (torch.Tensor): input tensor.
+
+ Returns:
+ torch.Tensor: The normalized output tensor.
+ """
+ return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)
+
+ def forward(self, x):
+ """
+ Running the forward pass of RMSNorm layer.
+
+ Args:
+ x (torch.Tensor): input tensor.
+
+ Returns:
+ torch.Tensor: output tensor after applying RMSNorm.
+ """
+ output = self._norm(x.float()).type_as(x)
+ if self.zero_centered_gamma:
+ return output * (1 + self.weight)
+ else:
+ return output * self.weight
diff --git a/ai_edge_torch/generative/layers/rotary_position_embedding.py b/ai_edge_torch/generative/layers/rotary_position_embedding.py
new file mode 100644
index 00000000..9e282b04
--- /dev/null
+++ b/ai_edge_torch/generative/layers/rotary_position_embedding.py
@@ -0,0 +1,36 @@
+# Copyright 2024 The AI Edge Torch Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+# Implementation for Rotary Position embedding. https://arxiv.org/pdf/2104.09864.pdf
+import torch
+
+
+def apply_rope(x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor:
+ """Computes rotary positional embedding.
+
+ Args:
+ x(torch.Tensor): the input tensor.
+ cos(torch.Tensor): cosine value for the rope.
+ sin(torch.Tensor): sin value for the rope.
+
+ Returns:
+ output tensor of RoPE.
+ """
+ x = x.transpose(1, 2)
+ head_size = x.size(-1)
+ x1 = x[..., : head_size // 2] # (B, nh, T, hs/2)
+ x2 = x[..., head_size // 2 :] # (B, nh, T, hs/2)
+ rotated = torch.cat((-x2, x1), dim=-1) # (B, nh, T, hs)
+ roped = (x * cos) + (rotated * sin)
+ return roped.transpose(1, 2).type_as(x)
diff --git a/ai_edge_torch/generative/quantize/README.md b/ai_edge_torch/generative/quantize/README.md
new file mode 100644
index 00000000..10336de5
--- /dev/null
+++ b/ai_edge_torch/generative/quantize/README.md
@@ -0,0 +1,24 @@
+# Quantization for the AI Edge Torch Generative API
+
+## Typical usage
+
+To apply quantization, we need to create a configuration that fully expresses how the model should be quantized. This configuration is then passed into conversion, generating a quantized model.
+
+`quant_recipes.py` contains a list of recipes that are known to be well-supported during runtime. For the average user, this is a good starting point to select the quantization scheme that is best suited for your deployment needs. After identifying the target recipe, the model can be quantized as follows. This example is extracted from `generative/examples/quantize/example.py`.
+
+```
+quant_config = quant_recipes.full_linear_int8_dynamic_recipe()
+edge_model = ai_edge_torch.convert(
+ model, (tokens, input_pos), quant_config=quant_config
+)
+```
+Once converted, you will get a quantized `.tflite` model which will be ready for on-device deployment.
+
+## Supported schemes
+
+In the current release, the following schemes are supported:
+
+* Dynamic range quantization with FP32 activations and INT8 weights for linear ops
+* FP16 quantization with FP16 weights and FP32 activations and computation for all ops
+
+These correspond to the available recipes in `quant_recipes.py`
diff --git a/ai_edge_torch/generative/quantize/__init__.py b/ai_edge_torch/generative/quantize/__init__.py
new file mode 100644
index 00000000..57b12003
--- /dev/null
+++ b/ai_edge_torch/generative/quantize/__init__.py
@@ -0,0 +1,14 @@
+# Copyright 2024 The AI Edge Torch Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
diff --git a/ai_edge_torch/generative/quantize/example.py b/ai_edge_torch/generative/quantize/example.py
new file mode 100644
index 00000000..24ca0a8d
--- /dev/null
+++ b/ai_edge_torch/generative/quantize/example.py
@@ -0,0 +1,45 @@
+# Copyright 2024 The AI Edge Torch Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+import numpy as np
+import torch
+
+import ai_edge_torch
+from ai_edge_torch.generative.examples.gemma import gemma
+from ai_edge_torch.generative.quantize import quant_recipes
+
+
+def main():
+ # Build a PyTorch model as usual
+ config = gemma.get_fake_model_config_2b_for_test()
+ model = gemma.Gemma(config)
+ idx = torch.from_numpy(np.array([[1, 2, 3, 4]]))
+ tokens = torch.full((1, 10), 0, dtype=torch.long, device="cpu")
+ tokens[0, :4] = idx
+ input_pos = torch.arange(0, 10)
+
+ # Create a quantization recipe to be applied to the model
+ quant_config = quant_recipes.full_linear_int8_dynamic_recipe()
+ print(quant_config)
+
+ # Convert with quantization
+ edge_model = ai_edge_torch.convert(
+ model, (tokens, input_pos), quant_config=quant_config
+ )
+ edge_model.export("/tmp/gemma_2b_quantized.tflite")
+
+
+if __name__ == "__main__":
+ main()
diff --git a/ai_edge_torch/generative/quantize/quant_attrs.py b/ai_edge_torch/generative/quantize/quant_attrs.py
new file mode 100644
index 00000000..bbdceba3
--- /dev/null
+++ b/ai_edge_torch/generative/quantize/quant_attrs.py
@@ -0,0 +1,66 @@
+# Copyright 2024 The AI Edge Torch Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+import enum
+
+
+@enum.unique
+class Dtype(enum.Enum):
+ """Data types and precision of tensors."""
+
+ FP32 = enum.auto()
+ FP16 = enum.auto()
+ INT8 = enum.auto()
+
+
+@enum.unique
+class Algorithm(enum.Enum):
+ """Algorithm used to calculate quantization parameters.
+
+ Attributes:
+ MIN_MAX: Maps the min/max of floating point space to the min/max of
+ quantized space and quantize uniformly.
+ """
+
+ MIN_MAX = enum.auto()
+
+
+@enum.unique
+class Mode(enum.Enum):
+ """Mode of quantization.
+
+ Attributes:
+ DYNAMIC_RANGE: Quantize activations during runtime and weights statically to
+ perform computation in integers.
+ WEIGHT_ONLY: Quantize weights statically and dequantize during runtime to
+ perform computation in floating points.
+ """
+
+ DYNAMIC_RANGE = enum.auto()
+ WEIGHT_ONLY = enum.auto()
+
+
+@enum.unique
+class Granularity(enum.Enum):
+ """Granularity of quantization parameters.
+
+ Attributes:
+ NONE: Granularity not applicable to this quantization scheme.
+ CHANNELWISE: Or per-channel quantization. Each channel of relevant tensors
+ is quantized independently of one another.
+ """
+
+ NONE = enum.auto()
+ CHANNELWISE = enum.auto()
diff --git a/ai_edge_torch/generative/quantize/quant_recipe.py b/ai_edge_torch/generative/quantize/quant_recipe.py
new file mode 100644
index 00000000..86d7a6b1
--- /dev/null
+++ b/ai_edge_torch/generative/quantize/quant_recipe.py
@@ -0,0 +1,106 @@
+# Copyright 2024 The AI Edge Torch Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+from dataclasses import dataclass
+import enum
+from typing import Optional
+
+from ai_edge_torch.generative.quantize import quant_attrs
+from ai_edge_torch.generative.quantize import supported_schemes
+
+
+@dataclass
+class LayerQuantRecipe:
+ """Quantization recipe for a single Edge Generative API layer (e.g. Attention).
+
+ Generic layer-scoped quantization recipe that specifies how this layer should
+ be quantized by the Edge Generative API. This is applicable to layers implemented
+ in ai_edge_torch/generative/layers/. Combinations of attributes that are not
+ supported during runtime will be detected when .verify() is called.
+
+ Attributes:
+ activation_dtype: Desired data type of activation tensors.
+ weight_dtype: Desired data type of weight tensors.
+ mode: Type of quantization.
+ algorithm: Algorithm for calculating quantization parameters.
+ granularity: Granularity of quantization.
+ """
+
+ activation_dtype: quant_attrs.Dtype
+ weight_dtype: quant_attrs.Dtype
+ mode: quant_attrs.Mode
+ algorithm: quant_attrs.Algorithm
+ granularity: quant_attrs.Granularity
+
+ def __str__(self):
+ return (
+ f'(a:{self.activation_dtype.name}, '
+ f'w:{self.weight_dtype.name}, '
+ f'{self.mode.name}, '
+ f'{self.algorithm.name}, '
+ f'{self.granularity.name})'
+ )
+
+ __repr__ = __str__
+
+ def verify(self):
+ """Checks if all attributes configured are supported in runtime.
+
+ Raises:
+ ValueError: If any attributes are incompatible.
+ """
+ is_valid = False
+ for supported in supported_schemes.get_supported_layer_schemes():
+ if (
+ self.activation_dtype == supported[0]
+ and self.weight_dtype == supported[1]
+ and self.mode == supported[2]
+ and self.algorithm == supported[3]
+ and self.granularity == supported[4]
+ ):
+ is_valid = True
+ break
+
+ if not is_valid:
+ raise ValueError(
+ 'Unsupported LayerQuantRecipe configuration. See get_supported_recipe_matrix()'
+ )
+
+
+@dataclass
+class TransformerQuantRecipe:
+ """Quantization recipe for a model composed of the Edge Generative API layers.
+
+ Attributes:
+ default: The quantization recipe for global scope of the model.
+ """
+
+ default: Optional[LayerQuantRecipe] = None
+
+ def __str__(self):
+ return f"""TransformerQuantRecipe(
+ Default: {self.default}
+)"""
+
+ __repr__ = __str__
+
+ def verify(self):
+ """Checks if the recipe configured can be supported in runtime.
+
+ Raises:
+ ValueError: If the recipe configured is invalid or unsupported.
+ """
+ if self.default is not None:
+ self.default.verify()
diff --git a/ai_edge_torch/generative/quantize/quant_recipe_utils.py b/ai_edge_torch/generative/quantize/quant_recipe_utils.py
new file mode 100644
index 00000000..441b86eb
--- /dev/null
+++ b/ai_edge_torch/generative/quantize/quant_recipe_utils.py
@@ -0,0 +1,51 @@
+# Copyright 2024 The AI Edge Torch Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+"""Helper functions to construct custom quantization recipes.
+
+These are intended for more advanced users who want to configure their own
+quantization recipes. For pre-constructed recipes, use `quant_recipes.py` instead.
+
+Typical usage example:
+
+1. Applying a single layer recipe to the entire model
+
+ quant_recipe.TransformerQuantRecipe(
+ default=quant_recipe_utils.create_layer_quant_int8_dynamic()
+ )
+"""
+
+from ai_edge_torch.generative.quantize import quant_attrs
+from ai_edge_torch.generative.quantize import quant_recipe
+
+
+def create_layer_quant_int8_dynamic() -> quant_recipe.LayerQuantRecipe:
+ return quant_recipe.LayerQuantRecipe(
+ activation_dtype=quant_attrs.Dtype.FP32,
+ weight_dtype=quant_attrs.Dtype.INT8,
+ mode=quant_attrs.Mode.DYNAMIC_RANGE,
+ algorithm=quant_attrs.Algorithm.MIN_MAX,
+ granularity=quant_attrs.Granularity.CHANNELWISE,
+ )
+
+
+def create_layer_quant_fp16() -> quant_recipe.LayerQuantRecipe:
+ return quant_recipe.LayerQuantRecipe(
+ activation_dtype=quant_attrs.Dtype.FP32,
+ weight_dtype=quant_attrs.Dtype.FP16,
+ mode=quant_attrs.Mode.WEIGHT_ONLY,
+ algorithm=quant_attrs.Algorithm.MIN_MAX,
+ granularity=quant_attrs.Granularity.NONE,
+ )
diff --git a/ai_edge_torch/generative/quantize/quant_recipes.py b/ai_edge_torch/generative/quantize/quant_recipes.py
new file mode 100644
index 00000000..5c36118f
--- /dev/null
+++ b/ai_edge_torch/generative/quantize/quant_recipes.py
@@ -0,0 +1,48 @@
+# Copyright 2024 The AI Edge Torch Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+"""Helper functions to create common and supported quantization recipes.
+
+These recipes will work with models created with the Edge Generative API only.
+Assume Transformer architecture congruent with
+ai_edge_torch/generative/layers/model_config.py:ModelConfig.
+
+Typical usage example:
+
+ quant_config = quant_recipes.full_linear_int8_dynamic_recipe()
+ edge_model = ai_edge_torch.convert(
+ model, (tokens, input_pos), quant_config=quant_config
+ )
+"""
+
+from ai_edge_torch.generative.quantize import quant_recipe
+from ai_edge_torch.generative.quantize import quant_recipe_utils
+from ai_edge_torch.quantize import quant_config
+
+
+def full_linear_int8_dynamic_recipe() -> quant_config.QuantConfig:
+ return quant_config.QuantConfig(
+ transformer_recipe=quant_recipe.TransformerQuantRecipe(
+ default=quant_recipe_utils.create_layer_quant_int8_dynamic()
+ )
+ )
+
+
+def full_fp16_recipe() -> quant_config.QuantConfig:
+ return quant_config.QuantConfig(
+ transformer_recipe=quant_recipe.TransformerQuantRecipe(
+ default=quant_recipe_utils.create_layer_quant_fp16()
+ )
+ )
diff --git a/ai_edge_torch/generative/quantize/supported_schemes.py b/ai_edge_torch/generative/quantize/supported_schemes.py
new file mode 100644
index 00000000..4086f44d
--- /dev/null
+++ b/ai_edge_torch/generative/quantize/supported_schemes.py
@@ -0,0 +1,31 @@
+# Copyright 2024 The AI Edge Torch Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+
+def get_supported_layer_schemes():
+ """List of layer-scoped quantization schemes supported in runtime.
+
+ Returns:
+ List of tuple(activation_dtype, weight_dtype, mode, algorithm, granularity).
+ """
+ from ai_edge_torch.generative.quantize.quant_attrs import Algorithm as _a
+ from ai_edge_torch.generative.quantize.quant_attrs import Dtype as _t
+ from ai_edge_torch.generative.quantize.quant_attrs import Granularity as _g
+ from ai_edge_torch.generative.quantize.quant_attrs import Mode as _m
+
+ return [
+ (_t.FP32, _t.INT8, _m.DYNAMIC_RANGE, _a.MIN_MAX, _g.CHANNELWISE),
+ (_t.FP32, _t.FP16, _m.WEIGHT_ONLY, _a.MIN_MAX, _g.NONE),
+ ]
diff --git a/ai_edge_torch/generative/screenshots/gemma-tflite.png b/ai_edge_torch/generative/screenshots/gemma-tflite.png
new file mode 100644
index 00000000..ac47ef0a
Binary files /dev/null and b/ai_edge_torch/generative/screenshots/gemma-tflite.png differ
diff --git a/ai_edge_torch/generative/test/__init__.py b/ai_edge_torch/generative/test/__init__.py
new file mode 100644
index 00000000..57b12003
--- /dev/null
+++ b/ai_edge_torch/generative/test/__init__.py
@@ -0,0 +1,14 @@
+# Copyright 2024 The AI Edge Torch Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
diff --git a/ai_edge_torch/generative/test/test_model_conversion.py b/ai_edge_torch/generative/test/test_model_conversion.py
new file mode 100644
index 00000000..07ab485c
--- /dev/null
+++ b/ai_edge_torch/generative/test/test_model_conversion.py
@@ -0,0 +1,201 @@
+# Copyright 2024 The AI Edge Torch Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+# Testing model conversion for a few gen-ai models.
+import copy
+import os
+import tempfile
+import unittest
+
+import numpy as np
+import torch
+
+import ai_edge_torch
+from ai_edge_torch.generative.examples.gemma import gemma
+from ai_edge_torch.generative.examples.phi2 import phi2
+from ai_edge_torch.generative.examples.test_models import toy_model_with_kv_cache # NOQA
+from ai_edge_torch.generative.examples.tiny_llama import tiny_llama
+from ai_edge_torch.testing import model_coverage
+
+
+class TestModelConversion(unittest.TestCase):
+ """Unit tests that check for model conversion and correctness."""
+
+ def test_toy_model_with_kv_cache(self):
+ self.skipTest("b/338288901")
+ config = toy_model_with_kv_cache.get_model_config()
+ pytorch_model = toy_model_with_kv_cache.ToyModelWithKV(config)
+ idx, input_pos = torch.tensor([[1]], dtype=torch.long), torch.tensor(
+ [10], dtype=torch.int64
+ )
+
+ edge_model = ai_edge_torch.convert(pytorch_model, (idx, input_pos))
+
+ self.assertTrue(
+ model_coverage.compare_tflite_torch(
+ edge_model,
+ pytorch_model,
+ (idx, input_pos),
+ num_valid_inputs=1,
+ atol=1e-5,
+ rtol=1e-5,
+ )
+ )
+
+ def test_toy_model_with_kv_cache_with_hlfb(self):
+ self.skipTest("b/338288901")
+ config = toy_model_with_kv_cache.get_model_config()
+ config.enable_hlfb = True
+ pytorch_model = toy_model_with_kv_cache.ToyModelWithKV(config)
+ idx, input_pos = torch.tensor([[1]], dtype=torch.long), torch.tensor(
+ [10], dtype=torch.int64
+ )
+
+ edge_model = ai_edge_torch.convert(pytorch_model, (idx, input_pos))
+
+ self.assertTrue(
+ model_coverage.compare_tflite_torch(
+ edge_model,
+ pytorch_model,
+ (idx, input_pos),
+ num_valid_inputs=1,
+ atol=1e-5,
+ rtol=1e-5,
+ )
+ )
+
+ def test_tiny_llama(self):
+ self.skipTest("b/338288901")
+ config = tiny_llama.get_fake_model_config_for_test()
+ pytorch_model = tiny_llama.TinyLLamma(config)
+
+ idx = torch.from_numpy(np.array([[1, 2, 3, 4]]))
+ tokens = torch.full((1, 10), 0, dtype=torch.long, device="cpu")
+ tokens[0, :4] = idx
+ input_pos = torch.arange(0, 10)
+
+ edge_model = ai_edge_torch.convert(pytorch_model, (tokens, input_pos))
+
+ self.assertTrue(
+ model_coverage.compare_tflite_torch(
+ edge_model,
+ pytorch_model,
+ (tokens, input_pos),
+ num_valid_inputs=1,
+ atol=1e-5,
+ rtol=1e-5,
+ )
+ )
+
+ def test_tiny_llama_multisig(self):
+ self.skipTest("b/338288901")
+ config = tiny_llama.get_fake_model_config_for_test()
+ pytorch_model = tiny_llama.TinyLLamma(config)
+
+ # prefill
+ seq_len = 10
+ prefill_tokens = torch.full((1, seq_len), 0, dtype=torch.long, device="cpu")
+ prompt_token = torch.from_numpy(np.array([1, 2, 3, 4]))
+ prefill_tokens[0, : len(prompt_token)] = prompt_token
+ prefill_input_pos = torch.arange(0, seq_len)
+
+ # decode
+ decode_token = torch.tensor([[1]], dtype=torch.long)
+ decode_input_pos = torch.tensor([5], dtype=torch.int64)
+
+ edge_model = (
+ ai_edge_torch.signature(
+ "prefill", pytorch_model, (prefill_tokens, prefill_input_pos)
+ )
+ .signature("decode", pytorch_model, (decode_token, decode_input_pos))
+ .convert()
+ )
+
+ # For the pytorch model, the KV cache is a persistent state internal to the model, and it
+ # will be shared for prefill and decode. However, for tflite, currently we can't share
+ # kv-cache between the two signatures. prefill will change the content in kv-cache,
+ # but it won't be readable by the decode tflite model. This means the output of running `decode` after
+ # running `prefill` in pytorch will be different from the output of running `decode` after `prefill` via ai_edge_torch.
+ copied_model = copy.deepcopy(pytorch_model)
+
+ self.assertTrue(
+ model_coverage.compare_tflite_torch(
+ edge_model,
+ pytorch_model,
+ (prefill_tokens, prefill_input_pos),
+ signature_name="prefill",
+ num_valid_inputs=1,
+ )
+ )
+
+ self.assertTrue(
+ model_coverage.compare_tflite_torch(
+ edge_model,
+ copied_model,
+ (decode_token, decode_input_pos),
+ signature_name="decode",
+ num_valid_inputs=1,
+ )
+ )
+
+ def test_gemma(self):
+ self.skipTest("b/338288901")
+ config = gemma.get_fake_model_config_2b_for_test()
+ model = gemma.Gemma(config)
+
+ idx = torch.from_numpy(np.array([[1, 2, 3, 4]]))
+ tokens = torch.full((1, 10), 0, dtype=torch.long, device="cpu")
+ tokens[0, :4] = idx
+ input_pos = torch.arange(0, 10)
+
+ edge_model = ai_edge_torch.convert(model, (tokens, input_pos))
+
+ # TODO(talumbau, haoliang): debug numerical diff.
+ self.assertTrue(
+ model_coverage.compare_tflite_torch(
+ edge_model,
+ model,
+ (tokens, input_pos),
+ num_valid_inputs=1,
+ atol=1e-2,
+ rtol=1e-5,
+ )
+ )
+
+ def test_phi2(self):
+ self.skipTest("b/338288901")
+ config = phi2.get_fake_model_config_for_test()
+ pytorch_model = phi2.Phi2(config)
+
+ idx = torch.from_numpy(np.array([[1, 2, 3, 4]]))
+ tokens = torch.full((1, 10), 0, dtype=torch.long, device="cpu")
+ tokens[0, :4] = idx
+ input_pos = torch.arange(0, 10)
+
+ edge_model = ai_edge_torch.convert(pytorch_model, (tokens, input_pos))
+
+ self.assertTrue(
+ model_coverage.compare_tflite_torch(
+ edge_model,
+ pytorch_model,
+ (tokens, input_pos),
+ num_valid_inputs=1,
+ atol=1e-5,
+ rtol=1e-5,
+ )
+ )
+
+
+if __name__ == "__main__":
+ unittest.main()
diff --git a/ai_edge_torch/generative/test/test_quantize.py b/ai_edge_torch/generative/test/test_quantize.py
new file mode 100644
index 00000000..7941fa43
--- /dev/null
+++ b/ai_edge_torch/generative/test/test_quantize.py
@@ -0,0 +1,109 @@
+# Copyright 2024 The AI Edge Torch Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+import unittest
+
+from parameterized import parameterized
+import torch
+
+import ai_edge_torch
+from ai_edge_torch.generative.examples.test_models import toy_model_with_kv_cache # NOQA
+from ai_edge_torch.generative.quantize import quant_recipe
+from ai_edge_torch.generative.quantize import quant_recipes
+from ai_edge_torch.generative.quantize.quant_attrs import Algorithm
+from ai_edge_torch.generative.quantize.quant_attrs import Dtype
+from ai_edge_torch.generative.quantize.quant_attrs import Granularity
+from ai_edge_torch.generative.quantize.quant_attrs import Mode
+from ai_edge_torch.testing import model_coverage
+
+
+class TestVerifyRecipes(unittest.TestCase):
+ """Unit tests that check for model quantization recipes."""
+
+ @parameterized.expand(
+ [
+ (Dtype.FP32, Dtype.FP32, Mode.DYNAMIC_RANGE),
+ (Dtype.INT8, Dtype.INT8, Mode.DYNAMIC_RANGE),
+ (Dtype.INT8, Dtype.FP16, Mode.DYNAMIC_RANGE),
+ (Dtype.FP16, Dtype.INT8, Mode.DYNAMIC_RANGE),
+ (Dtype.FP32, Dtype.FP32, Mode.WEIGHT_ONLY),
+ (Dtype.INT8, Dtype.INT8, Mode.WEIGHT_ONLY),
+ (Dtype.FP16, Dtype.INT8, Mode.WEIGHT_ONLY),
+ (Dtype.INT8, Dtype.FP16, Mode.WEIGHT_ONLY),
+ (Dtype.FP16, Dtype.FP16, Mode.WEIGHT_ONLY),
+ ]
+ )
+ def test_verify_invalid_recipes(
+ self,
+ activation,
+ weight,
+ mode,
+ algo=Algorithm.MIN_MAX,
+ granularity=Granularity.CHANNELWISE,
+ ):
+ with self.assertRaises(ValueError):
+ quant_recipe.LayerQuantRecipe(
+ activation, weight, mode, algo, granularity
+ ).verify()
+
+ @parameterized.expand(
+ [
+ (Dtype.FP32, Dtype.INT8, Mode.DYNAMIC_RANGE, Granularity.CHANNELWISE),
+ (Dtype.FP32, Dtype.FP16, Mode.WEIGHT_ONLY, Granularity.NONE),
+ ]
+ )
+ def test_verify_valid_recipes(
+ self,
+ activation,
+ weight,
+ mode,
+ granularity,
+ algo=Algorithm.MIN_MAX,
+ ):
+ quant_recipe.LayerQuantRecipe(activation, weight, mode, algo, granularity).verify()
+
+
+class TestQuantizeConvert(unittest.TestCase):
+ """Test conversion with quantization."""
+
+ def test_quantize_convert_toy(self):
+ self.skipTest("b/338288901")
+ config = toy_model_with_kv_cache.get_model_config()
+ pytorch_model = toy_model_with_kv_cache.ToyModelWithKV(config)
+ idx, input_pos = torch.tensor([[1]], dtype=torch.long), torch.tensor(
+ [10], dtype=torch.int64
+ )
+
+ quant_config = quant_recipes.full_fp16_recipe()
+ quantized_model = ai_edge_torch.convert(
+ pytorch_model, (idx, input_pos), quant_config=quant_config
+ )
+ float_model = ai_edge_torch.convert(pytorch_model, (idx, input_pos))
+
+ self.assertLess(len(quantized_model._tflite_model), len(float_model._tflite_model))
+ self.assertTrue(
+ model_coverage.compare_tflite_torch(
+ quantized_model,
+ pytorch_model,
+ (idx, input_pos),
+ num_valid_inputs=1,
+ atol=1e-3,
+ rtol=1e-3,
+ )
+ )
+
+
+if __name__ == "__main__":
+ unittest.main()
diff --git a/ai_edge_torch/generative/utilities/__init__.py b/ai_edge_torch/generative/utilities/__init__.py
new file mode 100644
index 00000000..a27919bd
--- /dev/null
+++ b/ai_edge_torch/generative/utilities/__init__.py
@@ -0,0 +1,15 @@
+# Copyright 2024 The AI Edge Torch Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+# This module contains common utility functions.
diff --git a/ai_edge_torch/generative/utilities/loader.py b/ai_edge_torch/generative/utilities/loader.py
new file mode 100644
index 00000000..873e935d
--- /dev/null
+++ b/ai_edge_torch/generative/utilities/loader.py
@@ -0,0 +1,290 @@
+# Copyright 2024 The AI Edge Torch Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+# Common utility functions for data loading etc.
+from dataclasses import dataclass
+import glob
+import os
+from typing import Callable, Dict
+
+from safetensors import safe_open
+import torch
+
+from ai_edge_torch.generative.layers import model_config
+
+
+def load_safetensors(full_path: str):
+ """Loads safetensors into a single state dictionary.
+
+ Args:
+ full_path (string): the directory that contains the safetensor files.
+
+ Returns:
+ A state dictionary contating loaded tensors.
+
+ Raises:
+ ValueError: If no tensors are loaded from the provided directory or file.
+ """
+ pattern = (
+ os.path.join(full_path, "*.safetensors")
+ if os.path.isdir(full_path)
+ else full_path
+ )
+ files = []
+ for file in glob.glob(pattern):
+ files.append(file)
+
+ tensors = {}
+ for file in files:
+ with safe_open(file, framework="pt") as fp:
+ for k in fp.keys():
+ assert k not in tensors
+ tensors[k] = fp.get_tensor(k)
+
+ if not tensors:
+ raise ValueError("Failed to load SafeTensors.")
+ return tensors
+
+
+def load_pytorch_statedict(full_path: str):
+ """Loads state dictionary binaries into a single state dictionary.
+
+ Args:
+ full_path (string): the directory that contains the bin files.
+
+ Returns:
+ A state dictionary contating loaded tensors.
+
+ Raises:
+ ValueError: If no tensors are loaded from the provided directory or file.
+ """
+ pattern = os.path.join(full_path, "*.bin") if os.path.isdir(full_path) else full_path
+ files = []
+ for file in glob.glob(pattern):
+ files.append(file)
+
+ tensors = {}
+ for file in files:
+ this_file_tensors = torch.load(file)
+ for k in this_file_tensors:
+ assert k not in tensors
+ tensors.update(this_file_tensors)
+
+ if not tensors:
+ raise ValueError("Failed to load torch bin files.")
+ return tensors
+
+
+class ModelLoader:
+ """A utility class for loading and converting model checkpoints to the
+ Edge Generative API layer format.
+ """
+
+ @dataclass
+ class TensorNames:
+ attn_query_proj: str
+ attn_key_proj: str
+ attn_value_proj: str
+ attn_output_proj: str
+
+ ff_up_proj: str
+ ff_down_proj: str
+ ff_gate_proj: str = None
+
+ pre_attn_norm: str = None
+ pre_ff_norm: str = None
+ embedding: str = None
+ final_norm: str = None
+ lm_head: str = None
+
+ def __init__(self, file_name: str, names: TensorNames) -> None:
+ """ModelLoader constructor. Can be used to load multiple models of the same
+ type.
+
+ Args:
+ file_name (str): Path to the checkpoint. Can be a directory or an
+ exact file.
+ names (TensorNames): An instance of `TensorNames` to determine mappings.
+ """
+ self._file_name = file_name
+ self._names = names
+ self._loader = self._get_loader()
+
+ def load(self, model: torch.nn.Module, strict: bool = True):
+ """Load the model from the checkpoint
+
+ Args:
+ model (torch.nn.Module): The pytorch model that needs to be loaded.
+ strict (bool, optional): Whether the converted keys are strictly
+ matched. Defaults to True.
+
+ Raises:
+ ValueError: If conversion results in unmapped tensors and strict mode is
+ enabled.
+ """
+ state = self._loader(self._file_name)
+ converted_state = dict()
+ if self._names.embedding is not None:
+ converted_state["tok_embedding.weight"] = state.pop(
+ f"{self._names.embedding}.weight"
+ )
+ if self._names.lm_head is not None:
+ converted_state["lm_head.weight"] = state.pop(f"{self._names.lm_head}.weight")
+ if model.config.lm_head_use_bias:
+ converted_state["lm_head.bias"] = state.pop(f"{self._names.lm_head}.bias")
+ if self._names.final_norm is not None:
+ final_norm_name = self._names.final_norm
+ converted_state["final_norm.weight"] = state.pop(f"{final_norm_name}.weight")
+ if f"{final_norm_name}.bias" in state:
+ converted_state["final_norm.bias"] = state.pop(f"{final_norm_name}.bias")
+
+ for i in range(model.config.num_layers):
+ self._map_norm(i, model.config, state, converted_state)
+ self._map_feedforward(i, model.config, state, converted_state)
+ self._map_attention(i, model.config, state, converted_state)
+
+ if strict and state:
+ raise ValueError(
+ f"Failed to map all tensor. Remaing tensor are: {list(state.keys())}"
+ )
+ model.load_state_dict(converted_state, strict=strict)
+
+ def _get_loader(self) -> Callable[[str], Dict[str, torch.Tensor]]:
+ """A best effort method for finding appropriate state loader.
+
+ Raises:
+ ValueError: If it fails to find an appropriate loader.
+
+ Returns:
+ Callable[[str], Dict[str, torch.Tensor]]: State loader to be used.
+ """
+ if os.path.isdir(self._file_name):
+ if glob.glob(os.path.join(self._file_name, "*.safetensors")):
+ return load_safetensors
+ if glob.glob(os.path.join(self._file_name, "*.bin")):
+ return load_pytorch_statedict
+
+ if self._file_name.endswith(".safetensors"):
+ return load_safetensors
+
+ if self._file_name.endswith(".bin"):
+ return load_pytorch_statedict
+
+ raise ValueError(f"File format not supported.")
+
+ def _map_feedforward(
+ self,
+ idx: int,
+ config: model_config.ModelConfig,
+ state: Dict[str, torch.Tensor],
+ converted_state: Dict[str, torch.Tensor],
+ ):
+ prefix = f"transformer_blocks.{idx}"
+ if config.ff_config.type == model_config.FeedForwardType.SEQUENTIAL:
+ ff_up_proj_name = self._names.ff_up_proj.format(idx)
+ ff_down_proj_name = self._names.ff_down_proj.format(idx)
+ converted_state[f"{prefix}.ff.w1.weight"] = state.pop(f"{ff_up_proj_name}.weight")
+ converted_state[f"{prefix}.ff.w2.weight"] = state.pop(
+ f"{ff_down_proj_name}.weight"
+ )
+ if config.ff_config.use_bias:
+ converted_state[f"{prefix}.ff.w1.bias"] = state.pop(f"{ff_up_proj_name}.bias")
+ converted_state[f"{prefix}.ff.w2.bias"] = state.pop(f"{ff_down_proj_name}.bias")
+ else:
+ ff_up_proj_name = self._names.ff_up_proj.format(idx)
+ ff_down_proj_name = self._names.ff_down_proj.format(idx)
+ ff_gate_proj_name = self._names.ff_gate_proj.format(idx)
+ converted_state[f"{prefix}.ff.w3.weight"] = state.pop(f"{ff_up_proj_name}.weight")
+ converted_state[f"{prefix}.ff.w2.weight"] = state.pop(
+ f"{ff_down_proj_name}.weight"
+ )
+ converted_state[f"{prefix}.ff.w1.weight"] = state.pop(
+ f"{ff_gate_proj_name}.weight"
+ )
+ if config.ff_config.use_bias:
+ converted_state[f"{prefix}.ff.w3.bias"] = state.pop(f"{ff_up_proj_name}.bias")
+ converted_state[f"{prefix}.ff.w2.bias"] = state.pop(f"{ff_down_proj_name}.bias")
+ converted_state[f"{prefix}.ff.w1.bias"] = state.pop(f"{ff_gate_proj_name}.bias")
+
+ def _map_attention(
+ self,
+ idx: int,
+ config: model_config.ModelConfig,
+ state: Dict[str, torch.Tensor],
+ converted_state: Dict[str, torch.Tensor],
+ ):
+ prefix = f"transformer_blocks.{idx}"
+ q_name = self._names.attn_query_proj.format(idx)
+ k_name = self._names.attn_key_proj.format(idx)
+ v_name = self._names.attn_value_proj.format(idx)
+ converted_state[f"{prefix}.atten_func.attn.weight"] = self._fuse_qkv(
+ config,
+ state.pop(f"{q_name}.weight"),
+ state.pop(f"{k_name}.weight"),
+ state.pop(f"{v_name}.weight"),
+ )
+ if config.attn_config.qkv_use_bias:
+ converted_state[f"{prefix}.atten_func.attn.bias"] = self._fuse_qkv(
+ config,
+ state.pop(f"{q_name}.bias"),
+ state.pop(f"{k_name}.bias"),
+ state.pop(f"{v_name}.bias"),
+ )
+
+ o_name = self._names.attn_output_proj.format(idx)
+ converted_state[f"{prefix}.atten_func.proj.weight"] = state.pop(f"{o_name}.weight")
+ if config.attn_config.output_proj_use_bias:
+ converted_state[f"{prefix}.atten_func.proj.bias"] = state.pop(f"{o_name}.bias")
+
+ def _map_norm(
+ self,
+ idx: int,
+ config: model_config.ModelConfig,
+ state: Dict[str, torch.Tensor],
+ converted_state: Dict[str, torch.Tensor],
+ ):
+ prefix = f"transformer_blocks.{idx}"
+ if self._names.pre_attn_norm is not None:
+ pre_attn_norm_name = self._names.pre_attn_norm.format(idx)
+ converted_state[f"{prefix}.pre_atten_norm.weight"] = state.pop(
+ f"{pre_attn_norm_name}.weight"
+ )
+ if f"{pre_attn_norm_name}.bias" in state:
+ converted_state[f"{prefix}.pre_atten_norm.bias"] = state.pop(
+ f"{pre_attn_norm_name}.bias"
+ )
+
+ if self._names.pre_ff_norm is not None:
+ pre_ff_norm_name = self._names.pre_ff_norm.format(idx)
+ converted_state[f"{prefix}.pre_ff_norm.weight"] = state.pop(
+ f"{pre_ff_norm_name}.weight"
+ )
+ if f"{pre_ff_norm_name}.bias" in state:
+ converted_state[f"{prefix}.pre_ff_norm.bias"] = state.pop(
+ f"{pre_ff_norm_name}.bias"
+ )
+
+ def _fuse_qkv(
+ self,
+ config: model_config.ModelConfig,
+ q: torch.Tensor,
+ k: torch.Tensor,
+ v: torch.Tensor,
+ ) -> torch.Tensor:
+ q_per_kv = config.attn_config.num_heads // config.attn_config.num_query_groups
+ qs = torch.split(q, config.head_dim * q_per_kv)
+ ks = torch.split(k, config.head_dim)
+ vs = torch.split(v, config.head_dim)
+ cycled = [t for group in zip(qs, ks, vs) for t in group]
+ return torch.cat(cycled)
diff --git a/ai_edge_torch/generative/utilities/t5_loader.py b/ai_edge_torch/generative/utilities/t5_loader.py
new file mode 100644
index 00000000..54aab3e1
--- /dev/null
+++ b/ai_edge_torch/generative/utilities/t5_loader.py
@@ -0,0 +1,467 @@
+# Copyright 2024 The AI Edge Torch Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+# Common utility functions for data loading etc.
+from dataclasses import dataclass
+import glob
+import os
+from typing import Callable, Dict
+
+from safetensors import safe_open
+import torch
+
+from ai_edge_torch.generative.layers import model_config
+
+
+def load_safetensors(full_path: str):
+ """Loads safetensors into a single state dictionary.
+
+ Args:
+ full_path (string): the safetensor filename or directory that contains the
+ safetensor files.
+
+ Returns:
+ A state dictionary contating loaded tensors.
+
+ Raises:
+ ValueError: If no tensors are loaded from the provided directory or file.
+ """
+ pattern = (
+ os.path.join(full_path, "*.safetensors")
+ if os.path.isdir(full_path)
+ else full_path
+ )
+ files = []
+ for file in glob.glob(pattern):
+ files.append(file)
+
+ tensors = {}
+ for file in files:
+ with safe_open(file, framework="pt") as fp:
+ for k in fp.keys():
+ assert k not in tensors
+ tensors[k] = fp.get_tensor(k)
+
+ if not tensors:
+ raise ValueError("Failed to load SafeTensors.")
+ return tensors
+
+
+def load_pytorch_statedict(full_path: str):
+ """Loads state dictionary binaries into a single state dictionary.
+
+ Args:
+ full_path (string): the bin filename or directory that contains the bin
+ files.
+
+ Returns:
+ A state dictionary contating loaded tensors.
+
+ Raises:
+ ValueError: If no tensors are loaded from the provided directory or file.
+ """
+ pattern = os.path.join(full_path, "*.bin") if os.path.isdir(full_path) else full_path
+ files = []
+ for file in glob.glob(pattern):
+ files.append(file)
+
+ tensors = {}
+ for file in files:
+ this_file_tensors = torch.load(file, map_location=torch.device("cpu"))
+ for k in this_file_tensors:
+ assert k not in tensors
+ tensors.update(this_file_tensors)
+
+ if not tensors:
+ raise ValueError("Failed to load torch bin files.")
+ return tensors
+
+
+class ModelLoader:
+ """A utility class for loading and converting model checkpoints to ODML
+ transformer layer format.
+ """
+
+ @dataclass
+ class TensorNames:
+ attn_query_proj: str = None
+ attn_key_proj: str = None
+ attn_value_proj: str = None
+ attn_output_proj: str = None
+ relative_attn_bias: str = None
+
+ cross_attn_query_proj: str = None
+ cross_attn_key_proj: str = None
+ cross_attn_value_proj: str = None
+ cross_attn_output_proj: str = None
+
+ ff_up_proj: str = None
+ ff_down_proj: str = None
+ ff_gate_proj: str = None
+
+ pre_attn_norm: str = None
+ pre_cross_attn_norm: str = None
+ pre_ff_norm: str = None
+ embedding: str = None
+ final_norm: str = None
+ lm_head: str = None
+
+ def __init__(self, file_name: str, names: TensorNames) -> None:
+ """ModelLoader constructor. Can be used to load multiple models of the same
+ type.
+
+ Args:
+ file_name (str): Path to the checkpoint. Can be a directory or an
+ exact file.
+ names (TensorNames): An instance of `TensorNames` to determine mappings.
+ """
+ self._file_name = file_name
+ self._names = names
+ self._loader = self._get_loader()
+
+ def load(
+ self, model: torch.nn.Module, strict: bool = True, fuse_attention: bool = True
+ ):
+ """Load the model from the checkpoint
+
+ Args:
+ model (torch.nn.Module): The pytorch model that needs to be loaded.
+ strict (bool, optional): Whether the converted keys are strictly
+ matched. Defaults to True.
+
+ Raises:
+ ValueError: If conversion results in unmapped tensors and strict mode is
+ enabled.
+ """
+ state = self._loader(self._file_name)
+
+ if isinstance(self._names, ModelLoader.TensorNames):
+ converted_state = self._do_load(
+ model, state, self._names, fuse_attention=fuse_attention
+ )
+ elif isinstance(self._names, dict):
+ converted_state = {}
+ for additional_prefix, names in self._names.items():
+ local_converted_state = self._do_load(
+ model,
+ state,
+ self._names[additional_prefix],
+ additional_prefix,
+ fuse_attention=fuse_attention,
+ )
+ converted_state.update(local_converted_state)
+ else:
+ raise ValueError(f"Unkown type for names: {type(self._names)}")
+
+ if strict and state:
+ raise ValueError(
+ f"Failed to map all tensor. Remaining tensor are: {list(state.keys())}"
+ )
+ model.load_state_dict(converted_state, strict=strict)
+
+ def _do_load(self, model, state, names, additional_prefix="", fuse_attention=True):
+ """Load the model from the checkpoint
+
+ Args:
+ model (torch.nn.Module): The pytorch model that needs to be loaded.
+ state (Dict[str, torch.Tensor]): The pytorch state dictionary
+ names (TensorNames]): The TensorNames for the model we are loading.
+
+ Returns:
+ Dict[str, torch.Tensor]: Map of name to tensor for loading.
+ """
+ converted_state = dict()
+ if names.embedding is not None:
+ converted_state["tok_embedding.weight"] = state.pop(f"{names.embedding}.weight")
+ if names.lm_head is not None:
+ converted_state["lm_head.weight"] = state.pop(f"{names.lm_head}.weight")
+ if model.config.lm_head_use_bias:
+ converted_state["lm_head.bias"] = state.pop(f"{names.lm_head}.bias")
+ if names.final_norm is not None:
+ final_norm_name = names.final_norm
+ prefix = additional_prefix
+ converted_state[f"{prefix}final_norm.weight"] = state.pop(
+ f"{final_norm_name}.weight"
+ )
+ if f"{final_norm_name}.bias" in state:
+ converted_state["final_norm.bias"] = state.pop(f"{final_norm_name}.bias")
+
+ if names.relative_attn_bias:
+ rel_attn_name = names.relative_attn_bias
+ prefix = additional_prefix + f"transformer_blocks.0"
+ converted_state[f"{prefix}.atten_func.relative_attention_bias.weight"] = (
+ state.pop(f"{rel_attn_name}.weight")
+ )
+
+ for i in range(model.config.num_layers):
+ self._map_norm(i, model.config, state, converted_state, names, additional_prefix)
+ self._map_feedforward(
+ i, model.config, state, converted_state, names, additional_prefix
+ )
+ self._map_attention(
+ i,
+ model.config,
+ state,
+ converted_state,
+ names,
+ additional_prefix,
+ fuse_attention,
+ )
+ self._map_cross_attention(
+ i,
+ model.config,
+ state,
+ converted_state,
+ names,
+ additional_prefix,
+ fuse_attention,
+ )
+
+ return converted_state
+
+ def _get_loader(self) -> Callable[[str], Dict[str, torch.Tensor]]:
+ """A best effort method for finding appropriate state loader.
+
+ Raises:
+ ValueError: If it fails to find an appropriate loader.
+
+ Returns:
+ Callable[[str], Dict[str, torch.Tensor]]: State loader to be used.
+ """
+ if os.path.isdir(self._file_name):
+ if glob.glob(os.path.join(self._file_name, "*.safetensors")):
+ return load_safetensors
+ if glob.glob(os.path.join(self._file_name, "*.bin")):
+ return load_pytorch_statedict
+
+ if self._file_name.endswith(".safetensors"):
+ return load_safetensors
+
+ if self._file_name.endswith(".bin"):
+ return load_pytorch_statedict
+
+ raise ValueError(f"File format not supported.")
+
+ def _map_feedforward(
+ self,
+ idx: int,
+ config: model_config.ModelConfig,
+ state: Dict[str, torch.Tensor],
+ converted_state: Dict[str, torch.Tensor],
+ names: TensorNames,
+ additional_prefix: str = "",
+ ):
+ prefix = additional_prefix + f"transformer_blocks.{idx}"
+ if names.ff_up_proj is None or names.ff_down_proj is None:
+ return
+ if config.ff_config.type == model_config.FeedForwardType.SEQUENTIAL:
+ ff_up_proj_name = names.ff_up_proj.format(idx)
+ ff_down_proj_name = names.ff_down_proj.format(idx)
+ converted_state[f"{prefix}.ff.w1.weight"] = state.pop(f"{ff_up_proj_name}.weight")
+ converted_state[f"{prefix}.ff.w2.weight"] = state.pop(
+ f"{ff_down_proj_name}.weight"
+ )
+ if config.ff_config.use_bias:
+ converted_state[f"{prefix}.ff.w1.bias"] = state.pop(f"{ff_up_proj_name}.bias")
+ converted_state[f"{prefix}.ff.w2.bias"] = state.pop(f"{ff_down_proj_name}.bias")
+ else:
+ if names.ff_gate_proj is not None:
+ ff_up_proj_name = names.ff_up_proj.format(idx)
+ ff_down_proj_name = names.ff_down_proj.format(idx)
+ ff_gate_proj_name = names.ff_gate_proj.format(idx)
+ converted_state[f"{prefix}.ff.w3.weight"] = state.pop(
+ f"{ff_up_proj_name}.weight"
+ )
+ converted_state[f"{prefix}.ff.w2.weight"] = state.pop(
+ f"{ff_down_proj_name}.weight"
+ )
+ converted_state[f"{prefix}.ff.w1.weight"] = state.pop(
+ f"{ff_gate_proj_name}.weight"
+ )
+ if config.ff_config.use_bias:
+ converted_state[f"{prefix}.ff.w3.bias"] = state.pop(f"{ff_up_proj_name}.bias")
+ converted_state[f"{prefix}.ff.w2.bias"] = state.pop(
+ f"{ff_down_proj_name}.bias"
+ )
+ converted_state[f"{prefix}.ff.w1.bias"] = state.pop(
+ f"{ff_gate_proj_name}.bias"
+ )
+
+ def _map_attention(
+ self,
+ idx: int,
+ config: model_config.ModelConfig,
+ state: Dict[str, torch.Tensor],
+ converted_state: Dict[str, torch.Tensor],
+ names: TensorNames,
+ additional_prefix: str = "",
+ fuse_attention: bool = True,
+ ):
+ if (
+ names.attn_query_proj is None
+ or names.attn_key_proj is None
+ or names.attn_value_proj is None
+ ):
+ return
+ prefix = additional_prefix + f"transformer_blocks.{idx}"
+ q_name = names.attn_query_proj.format(idx)
+ k_name = names.attn_key_proj.format(idx)
+ v_name = names.attn_value_proj.format(idx)
+ # model.encoder.transformer_blocks[0].atten_func.q.weight
+ if fuse_attention:
+ converted_state[f"{prefix}.atten_func.attn.weight"] = self._fuse_qkv(
+ config,
+ state.pop(f"{q_name}.weight"),
+ state.pop(f"{k_name}.weight"),
+ state.pop(f"{v_name}.weight"),
+ )
+ if config.attn_config.qkv_use_bias:
+ converted_state[f"{prefix}.atten_func.attn.bias"] = self._fuse_qkv(
+ config,
+ state.pop(f"{q_name}.bias"),
+ state.pop(f"{k_name}.bias"),
+ state.pop(f"{v_name}.bias"),
+ )
+ else:
+ converted_state[f"{prefix}.atten_func.q.weight"] = state.pop(f"{q_name}.weight")
+ converted_state[f"{prefix}.atten_func.k.weight"] = state.pop(f"{k_name}.weight")
+ converted_state[f"{prefix}.atten_func.v.weight"] = state.pop(f"{v_name}.weight")
+ if config.attn_config.qkv_use_bias:
+ converted_state[f"{prefix}.atten_func.q.bias"] = state.pop(f"{q_name}.bias")
+ converted_state[f"{prefix}.atten_func.k.bias"] = state.pop(f"{k_name}.bias")
+ converted_state[f"{prefix}.atten_func.v.bias"] = state.pop(f"{v_name}.bias")
+
+ o_name = names.attn_output_proj.format(idx)
+ converted_state[f"{prefix}.atten_func.proj.weight"] = state.pop(f"{o_name}.weight")
+ if config.attn_config.output_proj_use_bias:
+ converted_state[f"{prefix}.atten_func.proj.bias"] = state.pop(f"{o_name}.bias")
+
+ def _map_cross_attention(
+ self,
+ idx: int,
+ config: model_config.ModelConfig,
+ state: Dict[str, torch.Tensor],
+ converted_state: Dict[str, torch.Tensor],
+ names: TensorNames,
+ additional_prefix: str = "",
+ fuse_attention: bool = True,
+ ):
+ if (
+ names.cross_attn_query_proj is None
+ or names.cross_attn_key_proj is None
+ or names.cross_attn_value_proj is None
+ ):
+ return
+ prefix = additional_prefix + f"transformer_blocks.{idx}"
+ q_name = names.cross_attn_query_proj.format(idx)
+ k_name = names.cross_attn_key_proj.format(idx)
+ v_name = names.cross_attn_value_proj.format(idx)
+
+ if fuse_attention:
+ converted_state[f"{prefix}.cross_atten_func.attn.weight"] = self._fuse_qkv(
+ config,
+ state.pop(f"{q_name}.weight"),
+ state.pop(f"{k_name}.weight"),
+ state.pop(f"{v_name}.weight"),
+ )
+ if config.attn_config.qkv_use_bias:
+ converted_state[f"{prefix}.cross_atten_func.attn.bias"] = self._fuse_qkv(
+ config,
+ state.pop(f"{q_name}.bias"),
+ state.pop(f"{k_name}.bias"),
+ state.pop(f"{v_name}.bias"),
+ )
+ else:
+ converted_state[f"{prefix}.cross_atten_func.q.weight"] = state.pop(
+ f"{q_name}.weight"
+ )
+ converted_state[f"{prefix}.cross_atten_func.k.weight"] = state.pop(
+ f"{k_name}.weight"
+ )
+ converted_state[f"{prefix}.cross_atten_func.v.weight"] = state.pop(
+ f"{v_name}.weight"
+ )
+ if config.attn_config.qkv_use_bias:
+ converted_state[f"{prefix}.cross_atten_func.q.bias"] = state.pop(
+ f"{q_name}.bias"
+ )
+ converted_state[f"{prefix}.cross_atten_func.k.bias"] = state.pop(
+ f"{k_name}.bias"
+ )
+ converted_state[f"{prefix}.cross_atten_func.v.bias"] = state.pop(
+ f"{v_name}.bias"
+ )
+
+ o_name = names.cross_attn_output_proj.format(idx)
+ converted_state[f"{prefix}.cross_atten_func.proj.weight"] = state.pop(
+ f"{o_name}.weight"
+ )
+ if config.attn_config.output_proj_use_bias:
+ converted_state[f"{prefix}.cross_atten_func.proj.bias"] = state.pop(
+ f"{o_name}.bias"
+ )
+
+ def _map_norm(
+ self,
+ idx: int,
+ config: model_config.ModelConfig,
+ state: Dict[str, torch.Tensor],
+ converted_state: Dict[str, torch.Tensor],
+ names: TensorNames,
+ additional_prefix: str = "",
+ ):
+ prefix = additional_prefix + f"transformer_blocks.{idx}"
+ if names.pre_attn_norm is not None:
+ pre_attn_norm_name = names.pre_attn_norm.format(idx)
+ converted_state[f"{prefix}.atten_func.pre_atten_norm.weight"] = state.pop(
+ f"{pre_attn_norm_name}.weight"
+ )
+ if f"{pre_attn_norm_name}.bias" in state:
+ converted_state[f"{prefix}.atten_func.pre_atten_norm.bias"] = state.pop(
+ f"{pre_attn_norm_name}.bias"
+ )
+
+ if names.pre_cross_attn_norm:
+ pre_cross_attn_norm_name = names.pre_cross_attn_norm.format(idx)
+ converted_state[f"{prefix}.cross_atten_func.pre_atten_norm.weight"] = state.pop(
+ f"{pre_cross_attn_norm_name}.weight"
+ )
+ if f"{pre_cross_attn_norm_name}.bias" in state:
+ converted_state[f"{prefix}.cross_atten_func.pre_atten_norm.bias"] = state.pop(
+ f"{pre_cross_attn_norm_name}.bias"
+ )
+
+ if names.pre_ff_norm is not None:
+ pre_ff_norm_name = names.pre_ff_norm.format(idx)
+ converted_state[f"{prefix}.pre_ff_norm.weight"] = state.pop(
+ f"{pre_ff_norm_name}.weight"
+ )
+ if f"{pre_ff_norm_name}.bias" in state:
+ converted_state[f"{prefix}.pre_ff_norm.bias"] = state.pop(
+ f"{pre_ff_norm_name}.bias"
+ )
+
+ def _fuse_qkv(
+ self,
+ config: model_config.ModelConfig,
+ q: torch.Tensor,
+ k: torch.Tensor,
+ v: torch.Tensor,
+ ) -> torch.Tensor:
+ q_per_kv = config.attn_config.num_heads // config.attn_config.num_query_groups
+ qs = torch.split(q, config.head_dim * q_per_kv)
+ ks = torch.split(k, config.head_dim)
+ vs = torch.split(v, config.head_dim)
+ cycled = [t for group in zip(qs, ks, vs) for t in group]
+ return torch.cat(cycled)
diff --git a/ai_edge_torch/hlfb/__init__.py b/ai_edge_torch/hlfb/__init__.py
new file mode 100644
index 00000000..2f1c13d8
--- /dev/null
+++ b/ai_edge_torch/hlfb/__init__.py
@@ -0,0 +1,16 @@
+# Copyright 2024 The AI Edge Torch Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+from torch_xla.experimental.mark_pattern_utils import StableHLOCompositeBuilder
diff --git a/ai_edge_torch/hlfb/mark_pattern/__init__.py b/ai_edge_torch/hlfb/mark_pattern/__init__.py
new file mode 100644
index 00000000..86a44f9a
--- /dev/null
+++ b/ai_edge_torch/hlfb/mark_pattern/__init__.py
@@ -0,0 +1,139 @@
+# Copyright 2024 The AI Edge Torch Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+import copy
+from typing import Any
+import uuid
+
+import torch
+from torch_xla.experimental import xla_marker
+
+from ai_edge_torch.hlfb.mark_pattern.pattern import Pattern
+from ai_edge_torch.hlfb.mark_pattern.pattern import ScalarAttrTracker # NOQA
+
+
+@torch._dynamo.assume_constant_result
+def _get_uuid() -> str:
+ return uuid.uuid4().hex
+
+
+# TODO: Move to a general fx utils file.
+def _prepose_placeholder_nodes(graph: torch.fx.Graph):
+ nodes = [node for node in graph.nodes if node.op == "placeholder"] + [
+ node for node in graph.nodes if node.op != "placeholder"
+ ]
+
+ for a, b in zip(nodes, nodes[1:]):
+ if a.next is not b:
+ a.append(b)
+ return graph
+
+
+def _insert_marker(
+ graph_module: torch.fx.GraphModule,
+ node: torch.fx.Node,
+ name: str,
+ pos: int,
+ id: str,
+ is_input: bool,
+ attr: dict[str, Any] = None,
+):
+ attr = xla_marker.serialize_composite_attr(attr) if attr else None
+ with graph_module.graph.inserting_after(node):
+ new_node = graph_module.graph.call_function(
+ torch.ops.xla.mark_tensor,
+ args=(node,),
+ kwargs={
+ "name": name,
+ "pos": pos,
+ "id": id,
+ "is_input": is_input,
+ "attr": attr,
+ },
+ )
+
+ new_node.meta = node.meta
+ return new_node
+
+
+def mark_pattern(
+ graph_module: torch.fx.GraphModule,
+ pattern: Pattern,
+) -> torch.fx.GraphModule:
+ """Mark all existences of pattern graph in the GraphModule with fx pattern matching.
+ The marked subgraphs will be lowered in StableHLO composite ops.
+ Args:
+ graph_module (torch.fx.GraphModule): GraphModule to be matched and marked.
+ pattern (ai_edge_torch.hlfb.mark_pattern.Pattern): Pattern to match.
+ Returns:
+ The modified graph_module with additional marker ops in graph.
+ """
+ # Create a copy of graph_module and sanitize it for pattern matching.
+ graph_module_to_match = copy.deepcopy(graph_module)
+ for n, m in zip(graph_module.graph.nodes, graph_module_to_match.graph.nodes):
+ m.meta["ORIGINAL_NODE"] = n
+
+ # Sanitize graph_module to match in the same way as pattern's graph_module.
+ graph_module_to_match = passes.remove_clone_ops(graph_module_to_match)
+
+ match_with_attrs = pattern.match(graph_module_to_match)
+
+ for match, attr in match_with_attrs:
+ match_id = _get_uuid()
+
+ # NOTE: Current graph rewriter (_insert_marker) does not work perfectly
+ # with continuous matches e.g. matching (a + b) on (w + x + y + z). The
+ # rewritten results may be undetermined with false negative - some
+ # matches may not be marked in the lowering, while the marked ones would
+ # always be correct.
+ # TODO(cnchan): completely support mark_pattern on continuous matches.
+ for i, pattern_input_node in enumerate(pattern.input_nodes):
+ input_node = match.nodes_map[pattern_input_node]
+ new_input_node = _insert_marker(
+ graph_module,
+ input_node.meta["ORIGINAL_NODE"],
+ name=pattern.name,
+ pos=i,
+ id=match_id,
+ is_input=True,
+ )
+
+ # Only replace input by the marker node for those nodes used in the pattern.
+ in_pattern_nodes = set(match.nodes_map.values())
+ for user in input_node.users.keys():
+ if user in in_pattern_nodes:
+ user.meta["ORIGINAL_NODE"].replace_input_with(
+ input_node.meta["ORIGINAL_NODE"], new_input_node
+ )
+
+ for i, pattern_output_node in enumerate(pattern.output_nodes):
+ output_node = match.nodes_map[pattern_output_node]
+ new_output_node = _insert_marker(
+ graph_module,
+ output_node.meta["ORIGINAL_NODE"],
+ name=pattern.name,
+ pos=i,
+ id=match_id,
+ is_input=False,
+ attr=attr, # torch_xla internal: only output marker needs attr.
+ )
+ output_node.meta["ORIGINAL_NODE"].replace_all_uses_with(new_output_node)
+ new_output_node.update_arg(0, output_node.meta["ORIGINAL_NODE"])
+
+ graph_module.graph.eliminate_dead_code()
+ _prepose_placeholder_nodes(graph_module.graph)
+
+ graph_module.graph.lint()
+ graph_module.recompile()
+ return graph_module
diff --git a/ai_edge_torch/hlfb/mark_pattern/passes.py b/ai_edge_torch/hlfb/mark_pattern/passes.py
new file mode 100644
index 00000000..dbae4f54
--- /dev/null
+++ b/ai_edge_torch/hlfb/mark_pattern/passes.py
@@ -0,0 +1,42 @@
+# Copyright 2024 The AI Edge Torch Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+import torch
+
+
+def remove_clone_ops(gm: torch.fx.GraphModule):
+ # torch export adds additional aten.clone nodes to produce contiguous in memory tensors
+ # depending on tensor sizes for runtime efficiency. However, these unpredictable clone
+ # nodes can break the pattern matching. Thus remove all clones in model and pattern graphs.
+ for node in gm.graph.nodes:
+ if node.op == "call_function" and node.name.startswith("clone"):
+ node.replace_all_uses_with(node.args[0])
+ gm.graph.erase_node(node)
+
+ gm.graph.lint()
+ gm.recompile()
+ return gm
+
+
+def remove_dangling_args(gm: torch.fx.GraphModule):
+ nodes_to_erase = []
+ for node in gm.graph.nodes:
+ if node.op == "placeholder" and len(node.users) == 0:
+ nodes_to_erase.append(node)
+ for node in nodes_to_erase:
+ gm.graph.erase_node(node)
+
+ gm.graph.lint()
+ gm.recompile()
+ return gm
diff --git a/ai_edge_torch/hlfb/mark_pattern/pattern.py b/ai_edge_torch/hlfb/mark_pattern/pattern.py
new file mode 100644
index 00000000..ddca823d
--- /dev/null
+++ b/ai_edge_torch/hlfb/mark_pattern/pattern.py
@@ -0,0 +1,260 @@
+# Copyright 2024 The AI Edge Torch Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+import copy
+import dataclasses
+from typing import Any, Callable, Optional, Union
+
+import torch
+from torch.export.graph_signature import TensorArgument
+from torch.fx import Graph
+from torch.fx import GraphModule
+from torch.fx.passes.utils.matcher_utils import InternalMatch
+from torch.fx.passes.utils.matcher_utils import SubgraphMatcher
+
+from ai_edge_torch.hlfb.mark_pattern import passes
+
+
+def _are_equal(x: Any, y: Any) -> bool:
+ if type(x) != type(y):
+ return False
+ if type(x) in [int, str]:
+ return x == y
+ if isinstance(x, float):
+ rel_tol = 1e-07
+ abs_tol = 0.0
+ return abs(x - y) <= max(rel_tol * max(abs(x), abs(y)), abs_tol)
+ if isinstance(x, list):
+ if len(x) != len(y):
+ return False
+ return all([_are_equal(a, b) for a, b in zip(x, y)])
+
+ raise Exception(f"Cannot compare type: {type(x)}")
+
+
+@dataclasses.dataclass
+class ScalarAttrTracker:
+ """ScalarAttrTracker is used to track the occurrence of a pattern's
+ scalar arg/attr in the pattern decomposed graph. Since a scalar attr
+ to the pattern can be transformed and turned into a/some ops' scalar
+ arg in the decomposed graph, it would be hard to programmatically get
+ the attr value from the pattern match. With the tracker and tracking info,
+ we could target the position of the decomposed op's scalar arg derived
+ from the pattern arg/attr and retrieve the value from the InternalMatch.
+
+ Args:
+ name (str): name of the attr to track.
+ pattern_arg_pos (int): the index of the attr to track in the pattern's
+ export_args.
+ transform (Callable): the transform function used when targeting the
+ occurrence of the attr value in the decomposed graph. An attr value
+ may be transformed during the decomposition and appear as a derived
+ value.
+ inverse_transform (Callable): the inverse transform function that maps
+ the transformed value back to the original attr value.
+ """
+
+ attr_name: str
+ pattern_arg_pos: int
+ transform: Callable = lambda x: x
+ inverse_transform: Callable = lambda x: x
+ _source_targets: list[tuple[Any, Any]] = dataclasses.field(default_factory=list)
+
+ def track(self, *sources):
+ """Register magic values to track the (transformed) attr values in
+ the pattern decomposed graph.
+ """
+ for source in sources:
+ target = self.transform(source)
+ if not _are_equal(self.inverse_transform(target), source):
+ raise Exception(f"Invalid transform/inverse_transform for {self.attr_name}")
+ self._source_targets.append([source, target])
+ return self
+
+
+@dataclasses.dataclass
+class ScalarAttrLocation:
+ attr_name: str
+ node_name: str
+ pos: Union[int, str]
+ _tracker: ScalarAttrTracker
+
+ @property
+ def index(self):
+ return self.pos if isinstance(self.pos, int) else None
+
+ @property
+ def key(self):
+ return self.pos if isinstance(self.pos, str) else None
+
+
+def _find_scalar_attr(
+ pattern_module: torch.nn.Module, export_args: tuple[Any], tracker: ScalarAttrTracker
+) -> ScalarAttrLocation:
+ scalar_loc_intersections = None
+ for source, target in tracker._source_targets:
+ track_args = list(export_args)
+ track_args[tracker.pattern_arg_pos] = source
+ ep = torch.export.export(pattern_module, tuple(track_args))
+
+ scalar_locs = set()
+ nodes = ep.graph_module.graph.nodes
+ for n in nodes:
+ for arg_pos, arg in enumerate(n.args):
+ if type(arg) == type(target) and arg == target:
+ scalar_locs.add((n.name, arg_pos))
+ for attr, val in n.kwargs.items():
+ if type(val) == type(target) and val == target:
+ scalar_locs.add((n.name, attr))
+
+ if scalar_loc_intersections is None:
+ scalar_loc_intersections = scalar_locs
+ else:
+ scalar_loc_intersections = scalar_loc_intersections & scalar_locs
+
+ if not scalar_loc_intersections:
+ break
+
+ if not scalar_loc_intersections:
+ return None
+ # Choose any occurrence as the attr provider
+ node_name, pos = scalar_loc_intersections.pop()
+ return ScalarAttrLocation(tracker.attr_name, node_name, pos, _tracker=tracker)
+
+
+class Pattern:
+
+ def __init__(
+ self,
+ name: str,
+ module: Union[Callable, torch.nn.Module],
+ export_args: tuple[Any],
+ *,
+ attr_builder: Callable[
+ ["Pattern", GraphModule, InternalMatch], Optional[dict[str, Any]]
+ ] = None,
+ scalar_attr_trackers: list[ScalarAttrTracker] = None,
+ ):
+ """The PyTorch computation pattern to match against a model.
+
+ Args:
+ name (str): the name of the pattern. It would be propagated to
+ the `name` attr in StableHLO composite ops for the matched
+ model subgraphs in the lowering.
+ module (torch.nn.Module or Callable): the PyTorch computation.
+ export_args (tuple[Any]): the args used to export the pattern module
+ with torch.export.export. If export_args contains non-tensor
+ Python scalars, there must be a corresponding attr tracker
+ in `scalar_attr_trackers` for each scalar arg.
+ attr_builder (Callable[[Pattern, GraphModule, InternalMatch], Optional[dict[str, Any]]]):
+ the callable that produces the a scalar attrs dict, which would be
+ propagated to `attr` in StableHLO composite ops for the matched
+ model subgraphs in the lowering.
+ scalar_attr_trackers (list[ScalarAttrTracker]): the trackers
+ for scalar args in `export_args`, which are used to track
+ the attr occurrence(s) and retrieve their values from the
+ matched subgraph.
+ """
+ if not isinstance(module, torch.nn.Module):
+
+ class PatternModule(torch.nn.Module):
+
+ def __init__(self, func):
+ super().__init__()
+ self.func = func
+
+ def forward(self, *args, **kwargs):
+ return self.func(*args, **kwargs)
+
+ module = PatternModule(module).eval()
+
+ self.name = name
+ self.exported_program = torch.export.export(module, export_args)
+ self.graph_module = self.exported_program.graph_module
+ self.attr_builder = attr_builder
+ self._scalar_attr_trackers = scalar_attr_trackers if scalar_attr_trackers else []
+
+ # Sanitize graph_module for more precise pattern matching.
+ # The graph_module to match against this pattern should apply equivalent
+ # sanitization.
+ self.graph_module = passes.remove_clone_ops(self.graph_module)
+ self.graph_module = passes.remove_dangling_args(self.graph_module)
+
+ self._scalar_attr_locations = []
+ for tracker in self._scalar_attr_trackers:
+ self._scalar_attr_locations.append(
+ _find_scalar_attr(module, export_args, tracker)
+ )
+
+ # Builds list of ordered input and output nodes.
+ self.graph_nodes_map = {}
+ for node in self.graph_module.graph.nodes:
+ self.graph_nodes_map[node.name] = node
+
+ self.input_nodes = tuple(
+ self.graph_nodes_map[spec.arg.name]
+ for spec in self.exported_program.graph_signature.input_specs
+ if isinstance(spec.arg, TensorArgument)
+ )
+ self.output_nodes = tuple(
+ self.graph_nodes_map[spec.arg.name]
+ for spec in self.exported_program.graph_signature.output_specs
+ )
+
+ def register_attr_builder(self, attr_builder):
+ self.attr_builder = attr_builder
+ return attr_builder
+
+ def match(
+ self,
+ graph_module: GraphModule,
+ ) -> list[tuple[InternalMatch, dict[str, Any]]]:
+ matcher = SubgraphMatcher(
+ self.graph_module.graph,
+ match_output=False,
+ match_placeholder=False,
+ remove_overlapping_matches=True,
+ ignore_literals=True,
+ )
+ matches = matcher.match(graph_module.graph)
+
+ match_with_attrs = []
+ # Graph traversal must be done in the reverser order (from SubgraphMatcher).
+ for match in matches[::-1]:
+ if self.attr_builder is not None:
+ attrs = self.attr_builder(self, graph_module, match)
+ else:
+ attrs = {}
+
+ for loc in self._scalar_attr_locations:
+ attrs[loc.attr_name] = self._get_attr_value_from_pattern_match(match, loc)
+
+ attrs = attrs if attrs else None
+ match_with_attrs.append((match, attrs))
+ return match_with_attrs
+
+ def _get_attr_value_from_pattern_match(
+ self,
+ match: InternalMatch,
+ loc: ScalarAttrLocation,
+ ):
+ matched_val = None
+ for k, v in match.nodes_map.items():
+ if k.name == loc.node_name:
+ if loc.index:
+ matched_val = v.args[loc.index]
+ elif loc.key in v.kwargs.keys():
+ matched_val = v.kwargs[loc.key]
+ attr_val = loc._tracker.inverse_transform(matched_val)
+ return attr_val
diff --git a/ai_edge_torch/hlfb/test/__init__.py b/ai_edge_torch/hlfb/test/__init__.py
new file mode 100644
index 00000000..57b12003
--- /dev/null
+++ b/ai_edge_torch/hlfb/test/__init__.py
@@ -0,0 +1,14 @@
+# Copyright 2024 The AI Edge Torch Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
diff --git a/ai_edge_torch/hlfb/test/test_mark_pattern.py b/ai_edge_torch/hlfb/test/test_mark_pattern.py
new file mode 100644
index 00000000..00c0280e
--- /dev/null
+++ b/ai_edge_torch/hlfb/test/test_mark_pattern.py
@@ -0,0 +1,133 @@
+# Copyright 2024 The AI Edge Torch Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+import unittest
+
+import torch
+import torch_xla
+
+from ai_edge_torch.hlfb import mark_pattern
+
+
+def _export_stablehlo_mlir(model, args=None):
+ if not isinstance(model, torch.export.ExportedProgram):
+ ep = torch.export.export(model, args)
+ else:
+ ep = model
+ stablehlo_gm = torch_xla.stablehlo.exported_program_to_stablehlo(ep)
+ return stablehlo_gm.get_stablehlo_text()
+
+
+class TestMarkPattern(unittest.TestCase):
+
+ def test_mark_pattern(self):
+
+ class TestModel(torch.nn.Module):
+
+ def forward(self, x):
+ return x * x + x + x
+
+ pattern = mark_pattern.Pattern(
+ "test.add",
+ lambda a, b: a + b,
+ export_args=(torch.rand(2, 2), torch.rand(2, 2)),
+ )
+
+ model = TestModel().eval()
+ args = (torch.rand(20, 20),)
+ exported_program = torch.export.export(model, args)
+ mark_pattern.mark_pattern(exported_program.graph_module, pattern)
+ mlir = _export_stablehlo_mlir(exported_program)
+
+ self.assertEqual(mlir.count('stablehlo.composite "test.add"'), 2)
+
+ def test_mark_pattern_with_attr_builder(self):
+ class TestModel(torch.nn.Module):
+
+ def forward(self, x):
+ return x * x * x + x - x * x + x
+
+ pattern = mark_pattern.Pattern(
+ "test.add",
+ lambda a, b: a + b,
+ export_args=(torch.rand(2, 2), torch.rand(2, 2)),
+ attr_builder=lambda *args: {"alias": "test.test_add"},
+ )
+
+ model = TestModel().eval()
+ args = (torch.rand(20, 20),)
+ exported_program = torch.export.export(model, args)
+ mark_pattern.mark_pattern(exported_program.graph_module, pattern)
+ mlir = _export_stablehlo_mlir(exported_program)
+
+ self.assertEqual(mlir.count('stablehlo.composite "test.add"'), 2)
+ self.assertEqual(mlir.count('composite_attributes = {alias = "test.test_add"}'), 2)
+
+ def test_mark_pattern_with_scalar_attr_tracker(self):
+ class TestModel(torch.nn.Module):
+
+ def forward(self, x):
+ r = x
+ for idx in range(5):
+ r = torch.nn.LogSoftmax(dim=idx % 2)(r) * x
+ return r
+
+ pattern = mark_pattern.Pattern(
+ "test.log_softmax",
+ lambda x, dim: torch.nn.functional.log_softmax(x, dim=dim),
+ export_args=(torch.rand(10, 10, 10), 1),
+ scalar_attr_trackers=[
+ mark_pattern.ScalarAttrTracker("dim", pattern_arg_pos=1)
+ .track(0)
+ .track(1)
+ .track(2),
+ ],
+ )
+
+ model = TestModel().eval()
+ args = (torch.rand(10, 10),)
+ exported_program = torch.export.export(model, args)
+ mark_pattern.mark_pattern(exported_program.graph_module, pattern)
+ mlir = _export_stablehlo_mlir(exported_program)
+
+ self.assertEqual(mlir.count('stablehlo.composite "test.log_softmax"'), 5)
+ self.assertEqual(mlir.count("composite_attributes = {dim = 0 : i64}"), 3)
+ self.assertEqual(mlir.count("composite_attributes = {dim = 1 : i64}"), 2)
+
+ def test_mark_tangent_model_and_pattern_input(self):
+ class TestModel(torch.nn.Module):
+
+ def forward(self, x, y):
+ z = torch.ops.aten.relu(x)
+ z = z + y
+ return z
+
+ pattern = mark_pattern.Pattern(
+ "test.relu",
+ lambda x: torch.ops.aten.relu(x),
+ export_args=(torch.rand(2, 2),),
+ )
+
+ model = TestModel().eval()
+ args = (torch.rand(20, 20), torch.rand(20, 20))
+ exported_program = torch.export.export(model, args)
+ mark_pattern.mark_pattern(exported_program.graph_module, pattern)
+ mlir = _export_stablehlo_mlir(exported_program)
+
+ self.assertEqual(mlir.count('stablehlo.composite "test.relu'), 1)
+
+
+if __name__ == "__main__":
+ unittest.main()
diff --git a/ai_edge_torch/hlfb/test/test_stablehlo_composite_builder.py b/ai_edge_torch/hlfb/test/test_stablehlo_composite_builder.py
new file mode 100644
index 00000000..61e9b960
--- /dev/null
+++ b/ai_edge_torch/hlfb/test/test_stablehlo_composite_builder.py
@@ -0,0 +1,270 @@
+# Copyright 2024 The AI Edge Torch Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+import math
+import unittest
+
+import torch
+import torch.nn.functional as F
+import torch_xla
+
+from ai_edge_torch.hlfb import StableHLOCompositeBuilder
+
+
+def _export_stablehlo_mlir(model, args):
+ ep = torch.export.export(model, args)
+ stablehlo_gm = torch_xla.stablehlo.exported_program_to_stablehlo(ep)
+ return stablehlo_gm.get_stablehlo_text()
+
+
+class TestStableHLOCompositeBuilder(unittest.TestCase):
+
+ def test_build_composite(self):
+ class SampleModel(torch.nn.Module):
+
+ def forward(self, x):
+ builder = StableHLOCompositeBuilder(name="test.plus_two")
+ y = x + 1
+ y = builder.mark_inputs(y)
+ z = y + 2
+ z = builder.mark_outputs(z)
+ return z
+
+ mlir = _export_stablehlo_mlir(SampleModel().eval(), (torch.rand((2, 2)),))
+ self.assertEqual(mlir.count('stablehlo.composite "test.plus_two"'), 1)
+
+ def test_build_multiple_composites(self):
+ class SampleModel(torch.nn.Module):
+
+ def plus_one(self, x: torch.Tensor):
+ builder = StableHLOCompositeBuilder("test.plus_one")
+ x = builder.mark_inputs(x)
+ y = x + 1
+ y = builder.mark_outputs(y)
+ return y
+
+ def plus_two(self, x: torch.Tensor):
+ builder = StableHLOCompositeBuilder("test.plus_two")
+ x = builder.mark_inputs(x)
+ y = x + 2
+ y = builder.mark_outputs(y)
+ return y
+
+ def forward(self, x):
+ x = self.plus_two(x)
+ x = x + 3
+ x = self.plus_one(x)
+ x = x + 4
+ x = self.plus_two(x)
+ return x
+
+ mlir = _export_stablehlo_mlir(SampleModel().eval(), (torch.rand((2, 2)),))
+ self.assertEqual(mlir.count('stablehlo.composite "test.plus_one"'), 1)
+ self.assertEqual(mlir.count('stablehlo.composite "test.plus_two"'), 2)
+
+ def test_build_composite_with_attr(self):
+ class SampleModel(torch.nn.Module):
+
+ def __init__(self):
+ super().__init__()
+
+ def log_softmax(self, x: torch.Tensor, dim: int):
+ builder = StableHLOCompositeBuilder(name="test.log_softmax", attr={"dim": dim})
+ x = builder.mark_inputs(x)
+ y = torch.nn.functional.log_softmax(x, dim=dim)
+ y = builder.mark_outputs(y)
+ return y
+
+ def forward(self, x):
+ x = x + 1
+ x = self.log_softmax(x, 0)
+ x = self.log_softmax(x, 1)
+ return x
+
+ mlir = _export_stablehlo_mlir(SampleModel().eval(), (torch.rand((2, 2)),))
+ self.assertEqual(mlir.count('stablehlo.composite "test.log_softmax"'), 2)
+ self.assertEqual(mlir.count("composite_attributes = {dim = 0 : i64}"), 1)
+ self.assertEqual(mlir.count("composite_attributes = {dim = 1 : i64}"), 1)
+
+ def test_build_composite_with_mix_type_attrs(self):
+ class SampleModel(torch.nn.Module):
+
+ def __init__(self):
+ super().__init__()
+
+ def log_softmax(self, x: torch.Tensor, dim: int):
+ builder = StableHLOCompositeBuilder(
+ name="test.log_softmax",
+ attr={
+ "dim": dim,
+ "source": "torch.nn",
+ "version": 1.0,
+ },
+ )
+ x = builder.mark_inputs(x)
+ y = torch.nn.functional.log_softmax(x, dim=dim)
+ y = builder.mark_outputs(y)
+ return y
+
+ def forward(self, x):
+ x = x + 1
+ x = self.log_softmax(x, 0)
+ return x
+
+ mlir = _export_stablehlo_mlir(SampleModel().eval(), (torch.rand((2, 2)),))
+ self.assertEqual(mlir.count('stablehlo.composite "test.log_softmax"'), 1)
+ self.assertEqual(
+ mlir.count(
+ 'composite_attributes = {dim = 0 : i64, source = "torch.nn", version = 1.000000e+00 : f32}'
+ ),
+ 1,
+ )
+
+ def test_sdpa_composite(self):
+ class SDPAModel(torch.nn.Module):
+
+ def scaled_dot_product_attention(
+ self,
+ q: torch.Tensor,
+ k: torch.Tensor,
+ v: torch.Tensor,
+ head_size: int,
+ mask: torch.Tensor,
+ ):
+ builder = StableHLOCompositeBuilder("test.scaled_dot_product_attention")
+ q, k, v, mask = builder.mark_inputs(q, k, v, mask)
+
+ scale = 1.0 / math.sqrt(head_size)
+
+ q = q.transpose(1, 2)
+ k = k.transpose(1, 2)
+ v = v.transpose(1, 2)
+ y = F.scaled_dot_product_attention(
+ q,
+ k,
+ v,
+ attn_mask=mask,
+ dropout_p=0.0,
+ is_causal=mask is None,
+ scale=scale,
+ )
+ result = y.transpose(1, 2)
+ result = builder.mark_outputs(result)
+ return result
+
+ def forward(self, q, k, v, mask):
+ x = self.scaled_dot_product_attention(
+ q,
+ k,
+ v,
+ 8,
+ mask,
+ )
+ return x
+
+ query = torch.rand(1, 1, 32, 4)
+ key = torch.rand(1, 500, 1, 4)
+ value = torch.rand(1, 500, 1, 4)
+ mask = torch.rand(1, 1, 1, 500)
+
+ mlir = _export_stablehlo_mlir(
+ SDPAModel().eval(),
+ (query, key, value, mask),
+ )
+ self.assertEqual(
+ mlir.count('stablehlo.composite "test.scaled_dot_product_attention"'), 1
+ )
+
+ def test_sdpa_composite_with_attr(self):
+ class SDPAModel(torch.nn.Module):
+
+ def scaled_dot_product_attention(
+ self,
+ q: torch.Tensor,
+ k: torch.Tensor,
+ v: torch.Tensor,
+ head_size: int,
+ include_captanh: bool,
+ ):
+ builder = StableHLOCompositeBuilder(
+ name="test.scaled_dot_product_attention",
+ attr={"include_captanh": include_captanh},
+ )
+ q, k, v = builder.mark_inputs(q, k, v)
+
+ scale = 1.0 / math.sqrt(head_size)
+
+ q = q.transpose(1, 2)
+ k = k.transpose(1, 2)
+ v = v.transpose(1, 2)
+ y = F.scaled_dot_product_attention(
+ q,
+ k,
+ v,
+ attn_mask=None,
+ dropout_p=0.0,
+ is_causal=True,
+ scale=scale,
+ )
+ result = y.transpose(1, 2)
+ result = builder.mark_outputs(result)
+ return result
+
+ def forward(self, q, k, v):
+ x = self.scaled_dot_product_attention(q, k, v, 8, True)
+ y = self.scaled_dot_product_attention(q, k, v, 8, False)
+ return x + y
+
+ query = torch.rand(1, 1, 32, 4)
+ key = torch.rand(1, 500, 1, 4)
+ value = torch.rand(1, 500, 1, 4)
+ mlir = _export_stablehlo_mlir(
+ SDPAModel().eval(),
+ (query, key, value),
+ )
+ self.assertEqual(
+ mlir.count('stablehlo.composite "test.scaled_dot_product_attention"'), 2
+ )
+ self.assertEqual(mlir.count("composite_attributes = {include_captanh = true}"), 1)
+ self.assertEqual(mlir.count("composite_attributes = {include_captanh = false}"), 1)
+
+ def test_build_composite_with_multiple_inputs_outputs(self):
+ class SampleModel(torch.nn.Module):
+
+ def mimo_sample(self, a, b, c):
+ builder = StableHLOCompositeBuilder(name="test.mimo_sample")
+
+ a, b, c = builder.mark_inputs(a, b, c)
+ x = a + b + c
+ y = (a - b) * x
+ z = (c + 1.0) * a
+ x, y, z = builder.mark_outputs(x, y, z)
+
+ result = x + y * z
+ return result
+
+ def forward(self, a, b, c):
+ x = self.mimo_sample(a, b, c)
+ x = self.mimo_sample(a, b, x)
+ x = self.mimo_sample(x, x, c)
+ return x
+
+ mlir = _export_stablehlo_mlir(
+ SampleModel().eval(), (torch.rand(2), torch.rand(2), torch.rand(2))
+ )
+ self.assertEqual(mlir.count('stablehlo.composite "test.mimo_sample"'), 3)
+
+
+if __name__ == "__main__":
+ unittest.main()
diff --git a/ai_edge_torch/model.py b/ai_edge_torch/model.py
new file mode 100644
index 00000000..27632887
--- /dev/null
+++ b/ai_edge_torch/model.py
@@ -0,0 +1,134 @@
+# Copyright 2024 The AI Edge Torch Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+"""Represents an ai_edge_torch model.
+
+PyTorch models can be converted to this representation through `ai_edge_torch.convert`.
+"""
+from __future__ import annotations
+
+import abc
+
+import numpy as np
+import numpy.typing as npt
+import tensorflow as tf
+
+from ai_edge_torch.convert import conversion_utils as cutils
+
+
+class Model(abc.ABC):
+ """Represents and edge model."""
+
+ @abc.abstractmethod
+ def __call__(
+ self, *args: npt.ArrayLike, signature_name: str = cutils.DEFAULT_SIGNATURE_NAME
+ ) -> npt.ArrayLike | tuple[npt.ArrayLike]:
+ raise NotImplementedError()
+
+ @abc.abstractmethod
+ def export(self, path: str):
+ raise NotImplementedError()
+
+ @staticmethod
+ def load(path: str) -> TfLiteModel:
+ tflite_model = TfLiteModel.load(path)
+ if tflite_model:
+ return tflite_model
+
+ raise ValueError(f'File format in {path} cannot be deserialized.')
+
+
+class TfLiteModel(Model):
+ """An edge model which uses tflite under-the-hood."""
+
+ def __init__(self, tflite_model):
+ """Initializes the TfLiteModel instance using a TFLite serialized object.
+
+ Args:
+ tflite_model: A TFlite serialized object.
+ """
+ self._tflite_model = tflite_model
+
+ def __call__(
+ self, *args: npt.ArrayLike, signature_name: str = cutils.DEFAULT_SIGNATURE_NAME
+ ) -> npt.ArrayLike | tuple[npt.ArrayLike]:
+ """Runs inference on the edge model using the provided arguments.
+
+ Args:
+ *args: The arguments to be passed to the model for inference.
+ signature_name: The name of the signature to be used for inference.
+ The default signature is used if not provided.
+ """
+ interpreter = tf.lite.Interpreter(model_content=self._tflite_model)
+ interpreter.allocate_tensors()
+
+ signature_list = interpreter.get_signature_list()
+ if signature_name not in signature_list:
+ raise ValueError(
+ f"Invalid signature name provided. Available signatures: {', '.join(signature_list.keys())}"
+ )
+
+ try:
+ runner = interpreter.get_signature_runner(signature_name)
+ except ValueError as exception:
+ if 'Invalid signature_key provided.' in str(exception):
+ raise ValueError(
+ f'Invalid signature key provided. Available signatures: {list(signature_list.keys())}'
+ )
+ else:
+ raise exception
+
+ if len(signature_list[signature_name]['inputs']) != len(args):
+ raise ValueError(
+ f"The model requires {len(signature_list[signature_name]['inputs'])} arguments but {len(args)} was provided."
+ )
+
+ # Gather the input dictionary based on the signature.
+ inputs = {f'args_{idx}': args[idx] for idx in range(len(args))}
+ outputs = runner(**inputs)
+
+ return (
+ outputs['output_0']
+ if len(outputs) == 1
+ else [outputs[f'output_{idx}'] for idx in range(len(outputs))]
+ )
+
+ def export(self, path: str) -> None:
+ """Serializes the edge model to disk.
+
+ Args:
+ path: The path to file to which the model is serialized.
+ """
+ with open(path, 'wb') as file_handle:
+ file_handle.write(self._tflite_model)
+
+ @staticmethod
+ def load(path: str) -> TfLiteModel | None:
+ """Returns an edge (tflite) model by reading it from the disk.
+
+ Args:
+ str: The path to the model.
+ """
+ with open(path, 'rb') as file_handle:
+ model_content = file_handle.read()
+
+ # Check if this is indeed a tflite model:
+ try:
+ interpreter = tf.lite.Interpreter(model_content=model_content)
+ interpreter.get_signature_list()
+ except:
+ return None
+
+ return TfLiteModel(model_content)
diff --git a/ai_edge_torch/quantize/__init__.py b/ai_edge_torch/quantize/__init__.py
new file mode 100644
index 00000000..f2ccff56
--- /dev/null
+++ b/ai_edge_torch/quantize/__init__.py
@@ -0,0 +1,16 @@
+# Copyright 2024 The AI Edge Torch Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+from .pt2e_quantizer import PT2EQuantizer
diff --git a/ai_edge_torch/quantize/pt2e_quantizer.py b/ai_edge_torch/quantize/pt2e_quantizer.py
new file mode 100644
index 00000000..c429c38a
--- /dev/null
+++ b/ai_edge_torch/quantize/pt2e_quantizer.py
@@ -0,0 +1,438 @@
+# Copyright 2024 The AI Edge Torch Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+from __future__ import annotations
+
+import copy
+import functools
+from typing import Any, Callable, Dict, List, Optional, Set
+
+import torch
+from torch.ao.quantization.fake_quantize import FusedMovingAvgObsFakeQuantize
+from torch.ao.quantization.observer import HistogramObserver
+from torch.ao.quantization.observer import MinMaxObserver
+from torch.ao.quantization.observer import MovingAverageMinMaxObserver
+from torch.ao.quantization.observer import MovingAveragePerChannelMinMaxObserver # NOQA
+from torch.ao.quantization.observer import PerChannelMinMaxObserver
+from torch.ao.quantization.observer import PlaceholderObserver
+from torch.ao.quantization.qconfig import _ObserverOrFakeQuantizeConstructor
+from torch.ao.quantization.quantizer import FixedQParamsQuantizationSpec
+from torch.ao.quantization.quantizer import QuantizationSpec
+from torch.ao.quantization.quantizer import Quantizer
+from torch.fx import Node
+import torch.nn.functional as F
+
+from ai_edge_torch.quantize.pt2e_quantizer_utils import _convert_scalars_to_attrs # NOQA
+from ai_edge_torch.quantize.pt2e_quantizer_utils import OP_TO_ANNOTATOR
+from ai_edge_torch.quantize.pt2e_quantizer_utils import OperatorConfig
+from ai_edge_torch.quantize.pt2e_quantizer_utils import OperatorPatternType
+from ai_edge_torch.quantize.pt2e_quantizer_utils import propagate_annotation
+from ai_edge_torch.quantize.pt2e_quantizer_utils import QuantizationConfig
+
+__all__ = [
+ "PT2EQuantizer",
+ "get_symmetric_quantization_config",
+]
+
+
+def _supported_symmetric_quantized_operators() -> Dict[str, List[OperatorPatternType]]:
+ supported_operators: Dict[str, List[OperatorPatternType]] = {
+ # Both conv and linear should be able to handle relu + hardtanh fusion since
+ # those are clamp ops
+ "conv2d": [
+ [torch.nn.Conv2d, torch.nn.ReLU],
+ [torch.nn.Conv2d, F.relu],
+ [F.conv2d, torch.nn.ReLU],
+ [F.conv2d, F.relu],
+ ],
+ "linear": [[torch.nn.Linear], [F.linear]],
+ "add": [[torch.add]],
+ "max_pool2d": [[torch.nn.MaxPool2d], [F.max_pool2d]],
+ "adaptive_avg_pool2d": [
+ [torch.nn.AdaptiveAvgPool2d],
+ [F.adaptive_avg_pool2d],
+ ],
+ }
+ return copy.deepcopy(supported_operators)
+
+
+def _get_supported_symmetric_config_and_operators() -> List[OperatorConfig]:
+ supported_config_and_operators: List[OperatorConfig] = []
+ for quantization_config in [
+ get_symmetric_quantization_config(),
+ get_symmetric_quantization_config(is_qat=True),
+ get_symmetric_quantization_config(is_per_channel=True),
+ get_symmetric_quantization_config(is_per_channel=True, is_qat=True),
+ ]:
+ ops = _supported_symmetric_quantized_operators()
+ for pattern_list in ops.values():
+ supported_config_and_operators.append(
+ OperatorConfig(quantization_config, pattern_list)
+ )
+ return copy.deepcopy(supported_config_and_operators)
+
+
+@functools.lru_cache
+def get_symmetric_quantization_config(
+ is_per_channel: bool = False,
+ is_qat: bool = False,
+ is_dynamic: bool = False,
+):
+ if is_qat:
+ if is_dynamic:
+ raise NotImplementedError("dynamic quantization for qat is not yet implemented.")
+ act_observer_or_fake_quant_ctr = FusedMovingAvgObsFakeQuantize
+ else:
+ if is_dynamic:
+ act_observer_or_fake_quant_ctr = PlaceholderObserver # type: ignore[assignment]
+ else:
+ act_observer_or_fake_quant_ctr = HistogramObserver # type: ignore[assignment]
+
+ act_quantization_spec = QuantizationSpec(
+ dtype=torch.int8,
+ quant_min=-128,
+ quant_max=127,
+ qscheme=torch.per_tensor_affine,
+ is_dynamic=is_dynamic,
+ observer_or_fake_quant_ctr=act_observer_or_fake_quant_ctr.with_args(eps=2**-12),
+ )
+ qscheme = (
+ torch.per_channel_symmetric if is_per_channel else torch.per_tensor_symmetric
+ )
+ weight_observer_or_fake_quant_ctr: _ObserverOrFakeQuantizeConstructor = MinMaxObserver
+ if is_qat:
+ weight_observer_or_fake_quant_ctr = FusedMovingAvgObsFakeQuantize
+ elif is_per_channel:
+ weight_observer_or_fake_quant_ctr = PerChannelMinMaxObserver
+
+ extra_args: Dict[str, Any] = {"eps": 2**-12}
+ if is_qat:
+ if qscheme == torch.per_tensor_symmetric:
+ extra_args["observer"] = MovingAverageMinMaxObserver
+ else:
+ extra_args["observer"] = MovingAveragePerChannelMinMaxObserver # type: ignore[dict-item]
+ weight_quantization_spec = QuantizationSpec(
+ dtype=torch.int8,
+ quant_min=-127,
+ quant_max=127,
+ qscheme=qscheme,
+ ch_axis=0,
+ is_dynamic=False,
+ observer_or_fake_quant_ctr=weight_observer_or_fake_quant_ctr.with_args(
+ **extra_args
+ ),
+ )
+
+ bias_quantization_spec = None
+
+ # Some TFLite ops (e.g. Logistic, Softmax) have fixed qparams requirements
+ fixed_qparams_spec = FixedQParamsQuantizationSpec(
+ dtype=torch.int8,
+ scale=1 / 256,
+ zero_point=-128,
+ quant_min=-128,
+ quant_max=127,
+ qscheme=torch.per_tensor_affine,
+ )
+
+ if is_dynamic:
+ # Only valid for TFLite downstream to have no input activation quantization
+ # because dynamic quantization should be legalized to TFLite DRQ kernels
+ # which calculate quantization parameters during runtime inside the kernels
+ quantization_config = QuantizationConfig(
+ None,
+ None,
+ weight_quantization_spec,
+ bias_quantization_spec,
+ None,
+ is_qat,
+ True,
+ )
+ else:
+ quantization_config = QuantizationConfig(
+ act_quantization_spec,
+ act_quantization_spec,
+ weight_quantization_spec,
+ bias_quantization_spec,
+ fixed_qparams_spec,
+ is_qat,
+ False,
+ )
+ return quantization_config
+
+
+def _get_supported_config_and_operators() -> List[OperatorConfig]:
+ return _get_supported_symmetric_config_and_operators()
+
+
+def _get_module_name_filter(module_name: str):
+ """Get the module_name_filter function for a given module name, the filter accepts
+ a node and checks if the node comes from a module that has certain module name
+
+ For example:
+ node: linear_op = call_function[...](...) # comes from a module with name blocks.sub.linear1
+
+
+ >> module_name_filter = _get_module_name_filter("blocks.sub")
+ >> print(module_name_filter(node))
+ True # the node is from "blocks.sub" based on the fully qualified name "blocks.sub.linear1"
+ """
+
+ def module_name_filter(n: Node) -> bool:
+ # example: {
+ # 'L__self___sub': ("L['self'].sub", \n",
+ ""
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "id": "00e00b3b-d7ed-4e2e-815e-3addfc23c8f3",
+ "metadata": {
+ "id": "23dd0a0eba89"
+ },
+ "outputs": [],
+ "source": [
+ "# Copyright 2024 The AI Edge Torch Authors.\n",
+ "#\n",
+ "# Licensed under the Apache License, Version 2.0 (the \"License\");\n",
+ "# you may not use this file except in compliance with the License.\n",
+ "# You may obtain a copy of the License at\n",
+ "#\n",
+ "# http://www.apache.org/licenses/LICENSE-2.0\n",
+ "#\n",
+ "# Unless required by applicable law or agreed to in writing, software\n",
+ "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
+ "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
+ "# See the License for the specific language governing permissions and\n",
+ "# limitations under the License.\n",
+ "# =============================================================================="
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "a9bdc007e6ce"
+ },
+ "source": [
+ "Note: When running notebooks in this repository with Google Colab, some users may see\n",
+ "the following warning message:\n",
+ "\n",
+ "\n",
+ "\n",
+ "Please click `Restart Session` and run again."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "id": "a91d40b5-91f0-4c19-bdb4-a2f56fa1c5ff",
+ "metadata": {
+ "id": "2b09cc13a5c1"
+ },
+ "outputs": [],
+ "source": [
+ "!pip install -r https://raw.githubusercontent.com/google-ai-edge/ai-edge-torch/main/requirements.txt\n",
+ "!pip install ai-edge-torch"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "id": "f744e7c3-e360-4f3a-8d75-74759265b4aa",
+ "metadata": {
+ "id": "2027d669fce6"
+ },
+ "outputs": [],
+ "source": [
+ "import numpy as np\n",
+ "import ai_edge_torch\n",
+ "import torch\n",
+ "import torchvision"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "id": "cec203fc-7b6d-41bf-9716-4b76af45b019",
+ "metadata": {
+ "id": "2bf24a1bd28e"
+ },
+ "source": [
+ "# Sample PyTorch Model\n",
+ "\n",
+ "Instantiate `resnet18` as a sample model from PyTorch's `torchvision` package. We also provide it with a sample input and execute it directly via PyTorch."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "id": "c4ad2105-ce72-4f00-8f4d-74b0d505d422",
+ "metadata": {
+ "id": "c96810c259a9"
+ },
+ "outputs": [],
+ "source": [
+ "resnet18 = torchvision.models.resnet18(torchvision.models.ResNet18_Weights.IMAGENET1K_V1).eval()\n",
+ "sample_inputs = (torch.randn(1, 3, 224, 224),)\n",
+ "torch_output = resnet18(*sample_inputs)"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "id": "efbc9364-e0ce-4213-a0a7-07b0b6a264ae",
+ "metadata": {
+ "id": "ba2ad90ae477"
+ },
+ "source": [
+ "# Conversion\n",
+ "The `convert` function provided by the `ai_edge_torch` package allows conversion from a PyTorch model to an on-device model. The conversion process also requires a model's sample input for tracing and shape inference.\n",
+ "\n",
+ "**Note**: The source PyTorch model needs to be compliant with `torch.export` introduced in PyTorch 2.1.0 ."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "id": "508e4a18-dd21-410c-bf65-ecb062d4d3ba",
+ "metadata": {
+ "id": "26a208b29579"
+ },
+ "outputs": [],
+ "source": [
+ "edge_model = ai_edge_torch.convert(resnet18, sample_inputs)"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "id": "35ee138e-f93d-4b47-a698-27f985fb2d3a",
+ "metadata": {
+ "id": "f61e660adb9f"
+ },
+ "source": [
+ "# Inference\n",
+ "Get outputs from inference with the TFLite runtime by directly calling the edge_model with the inputs. Many of the details of [TFLite inference in Python](https://www.tensorflow.org/lite/guide/inference#load_and_run_a_model_in_python) are abstracted away with this API."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "id": "0b9860f6-0dc4-41ca-ac31-6a177e89f4c3",
+ "metadata": {
+ "id": "d53042b5e46a"
+ },
+ "outputs": [],
+ "source": [
+ "edge_output = edge_model(*sample_inputs)"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "id": "b9b6c6ca-1ebf-4011-b5cd-86a5be666f1c",
+ "metadata": {
+ "id": "7862f0d68600"
+ },
+ "source": [
+ "# Validation\n",
+ "Here, we make sure that the output generated by the on-device prepared model created by `ai_edge_torch` matches the output generated by PyTorch."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "id": "f2cd200f-aa9e-4eb2-94cf-9d683a01ded8",
+ "metadata": {
+ "id": "ea6b6914879e"
+ },
+ "outputs": [],
+ "source": [
+ "if np.allclose(torch_output.detach().numpy(), edge_output, atol=1e-5):\n",
+ " print(\"Inference result with Pytorch and TfLite was within tolerance\")\n",
+ "else:\n",
+ " print(\"Something wrong with Pytorch --> TfLite\")"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "id": "5ee2c9f3-585a-43ff-a9ef-82e0e7b58dc3",
+ "metadata": {
+ "id": "83468e71907a"
+ },
+ "source": [
+ "# Serialization\n",
+ "The on-device prepared model also provides an `export` interface which can be used to serialize the model. This serializes the model as a TFLite Flatbuffers file."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "id": "4df580b7-ff62-4d24-b99c-ae0c43b47a88",
+ "metadata": {
+ "id": "942812454807"
+ },
+ "outputs": [],
+ "source": [
+ "edge_model.export('resnet.tflite')\n",
+ "\n",
+ "# Download the tflite flatbuffer which can be used with the existing TfLite APIs.\n",
+ "# from google.colab import files\n",
+ "# files.download('resnet.tflite')"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "id": "92d06de3-2a33-4d9c-bdc0-8128379f1d6d",
+ "metadata": {
+ "id": "52027ca7613f"
+ },
+ "source": [
+ "# Visualization\n",
+ "The TFLite flatbuffer can be visualized using the AI Edge Model Explorer."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "id": "6d04ede7-5e72-41ba-b4a8-4debe7e12507",
+ "metadata": {
+ "id": "1c5cc28c58de"
+ },
+ "outputs": [],
+ "source": [
+ "!pip install ai-edge-model-explorer\n",
+ "\n",
+ "import model_explorer\n",
+ "model_explorer.visualize('resnet.tflite')"
+ ]
+ }
+ ],
+ "metadata": {
+ "colab": {
+ "name": "getting_started.ipynb",
+ "toc_visible": true
+ },
+ "kernelspec": {
+ "display_name": "Python 3",
+ "name": "python3"
+ }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 0
+}
diff --git a/format.sh b/format.sh
new file mode 100755
index 00000000..33e2e87d
--- /dev/null
+++ b/format.sh
@@ -0,0 +1,22 @@
+#!/bin/bash
+# Copyright 2024 The AI Edge Torch Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+# A helper script to format code. Must be called from repo's root.
+#
+
+set -ex
+
+pyink --pyink-use-majority-quotes --pyink-indentation 2 ./
+isort --profile google --multi-line 7 --skip .downloads ./
diff --git a/requirements.txt b/requirements.txt
new file mode 100644
index 00000000..d99ffcfb
--- /dev/null
+++ b/requirements.txt
@@ -0,0 +1,15 @@
+scipy
+numpy
+tabulate
+safetensors
+--pre
+tf-nightly==2.17.0.dev20240509
+-f https://download.pytorch.org/whl/nightly/torch_nightly.html
+torch==2.4.0.dev20240429+cpu
+-f https://download.pytorch.org/whl/nightly/torch_nightly.html
+torchvision==0.19.0.dev20240429+cpu
+-f https://download.pytorch.org/whl/nightly/torch_nightly.html
+torchaudio==2.2.0.dev20240429+cpu
+torch_xla@https://storage.googleapis.com/pytorch-xla-releases/wheels/tpuvm/torch_xla-nightly+20240429-cp311-cp311-linux_x86_64.whl; python_version=="3.11"
+torch_xla@https://storage.googleapis.com/pytorch-xla-releases/wheels/tpuvm/torch_xla-nightly+20240429-cp310-cp310-linux_x86_64.whl; python_version=="3.10"
+torch_xla@https://storage.googleapis.com/pytorch-xla-releases/wheels/tpuvm/torch_xla-nightly+20240429-cp39-cp39-linux_x86_64.whl; python_version=="3.9"
\ No newline at end of file
diff --git a/run_tests.sh b/run_tests.sh
new file mode 100755
index 00000000..97352616
--- /dev/null
+++ b/run_tests.sh
@@ -0,0 +1,18 @@
+#!/usr/bin/env bash
+# Copyright 2024 The AI Edge Torch Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+SCRIPT_DIR=$( cd -- "$( dirname -- "${BASH_SOURCE[0]}" )" &> /dev/null && pwd )
+PYTHONPATH=$SCRIPT_DIR:$PYTHONPATH python -m pytest $SCRIPT_DIR -n auto
diff --git a/setup.py b/setup.py
new file mode 100644
index 00000000..0b9ca240
--- /dev/null
+++ b/setup.py
@@ -0,0 +1,69 @@
+# Copyright 2024 The AI Edge Torch Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+import pathlib
+
+from setuptools import find_packages
+from setuptools import setup
+
+here = pathlib.Path(__file__).parent.resolve()
+
+# Get the long description from the README file
+long_description = """
+Library that supports converting PyTorch models into a .tflite format, which can
+then be run with TensorFlow Lite and MediaPipe. This enables applications for
+Android, iOS and IOT that can run models completely on-device.
+
+More details are in the project's [GitHub repository](https://github.com/google-ai-edge/ai-edge-torch).
+""".lstrip()
+
+setup(
+ name="ai-edge-torch",
+ version="0.1.1",
+ description="Supporting PyTorch models with the Google AI Edge TFLite runtime.",
+ long_description=long_description,
+ long_description_content_type="text/markdown",
+ url="https://github.com/google-ai-edge/ai-edge-torch",
+ classifiers=[
+ "Development Status :: 4 - Beta",
+ "Intended Audience :: Developers",
+ "Intended Audience :: Education",
+ "Intended Audience :: Science/Research",
+ "License :: OSI Approved :: Apache Software License",
+ "Programming Language :: Python :: 3",
+ "Programming Language :: Python :: 3 :: Only",
+ "Programming Language :: Python :: 3.9",
+ "Programming Language :: Python :: 3.10",
+ "Programming Language :: Python :: 3.11",
+ "Topic :: Scientific/Engineering",
+ "Topic :: Scientific/Engineering :: Mathematics",
+ "Topic :: Scientific/Engineering :: Artificial Intelligence",
+ "Topic :: Software Development",
+ "Topic :: Software Development :: Libraries",
+ "Topic :: Software Development :: Libraries :: Python Modules",
+ ],
+ keywords="On-Device ML, AI, Google, TFLite, PyTorch, LLMs, GenAI",
+ packages=find_packages(
+ include=["ai_edge_torch*"],
+ ),
+ python_requires=">=3.9, <3.12",
+ install_requires=[
+ "numpy",
+ "scipy",
+ "safetensors",
+ "tabulate",
+ "torch==2.4.*", # 2.4.0 stable release does not exist. Force using torch nightly.
+ ],
+)
diff --git a/test/__init__.py b/test/__init__.py
new file mode 100644
index 00000000..e69de29b
diff --git a/test/image_classification/colab/mobile_vit_mpt.ipynb b/test/image_classification/colab/mobile_vit_mpt.ipynb
new file mode 100644
index 00000000..04fa5b1c
--- /dev/null
+++ b/test/image_classification/colab/mobile_vit_mpt.ipynb
@@ -0,0 +1,602 @@
+{
+ "cells": [
+ {
+ "cell_type": "code",
+ "execution_count": 13,
+ "metadata": {
+ "id": "r4lisalb-A5R"
+ },
+ "outputs": [],
+ "source": [
+ "# Copyright 2024 The AI Edge Torch Authors.\n",
+ "#\n",
+ "# Licensed under the Apache License, Version 2.0 (the \"License\");\n",
+ "# you may not use this file except in compliance with the License.\n",
+ "# You may obtain a copy of the License at\n",
+ "#\n",
+ "# http://www.apache.org/licenses/LICENSE-2.0\n",
+ "#\n",
+ "# Unless required by applicable law or agreed to in writing, software\n",
+ "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
+ "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
+ "# See the License for the specific language governing permissions and\n",
+ "# limitations under the License.\n",
+ "# =============================================================================="
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "LwrH6f2sGJ6U"
+ },
+ "source": [
+ "This Colab demonstrates how to convert a PyTorch [MobileViT](https://huggingface.co/docs/transformers/en/model_doc/mobilevit#overview) model to a TensorFlow Lite model using the ai_edge_torch library. It also guides you through running the converted model with MediaPipe's Image Classification Task."
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "Mzf2MdHoG-9c"
+ },
+ "source": [
+ "# Prerequisites"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "hux_Gsc_G4nl"
+ },
+ "source": [
+ "First install all dependencies."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {
+ "id": "l-9--DWON236"
+ },
+ "outputs": [],
+ "source": [
+ "!pip install -r https://raw.githubusercontent.com/google-ai-edge/ai-edge-torch/main/requirements.txt\n",
+ "!pip install ai-edge-torch\n",
+ "!pip install transformers pillow requests matplotlib mediapipe"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "IUMh9GRk17fV"
+ },
+ "source": [
+ "Then download and read the test image."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 37,
+ "metadata": {
+ "id": "lfdgp-4Id51J"
+ },
+ "outputs": [
+ {
+ "name": "stdout",
+ "output_type": "stream",
+ "text": [
+ " % Total % Received % Xferd Average Speed Time Time Time Current\n",
+ " Dload Upload Total Spent Left Speed\n",
+ "100 144k 100 144k 0 0 482k 0 --:--:-- --:--:-- --:--:-- 482k\n"
+ ]
+ },
+ {
+ "data": {
+ "text/plain": [
+ "O1W8P3y8`%gf<_8Ku--1 +n8P6OiynqB?sR@*eP#&lBXI1