ai-edge-torch-nightly 0.1.dev202405131930__py3-none-any.whl

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()

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,
+  )

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()
+      )
+  )

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),
+  ]

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.
+# ==============================================================================

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()

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()

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.