ai-edge-torch-nightly 0.2.0.dev20240714__py3-none-any.whl

ai_edge_torch/generative/quantize/quant_recipe.py

@@ -0,0 +1,151 @@
+# 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
+from typing import Optional, Union
+
+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 GenerativeQuantRecipe:
+  """Quantization recipe for a model composed of the Edge Generative API layers.
+
+  Some layers can be specified with different `LayerQuantRecipe` for each block by
+  providing a dictionary keyed by the TransformerBlock index, e.g. attention
+  and feedforward. For example,
+
+  ```
+  default = LayerQuantRecipeA
+  attention = { 2: LayerQuantRecipeB }
+  feedforward = { 3: LayerQuantRecipeC }
+  ```
+
+  will apply LayerQuantRecipeA to the entire model, overriden by
+  LayerQuantRecipeB for the TransformerBlock[2].attention layer and
+  LayerQuantRecipeC for the TransformerBlock[3].feedforward layer. Any config
+  with invalid indices will be ignored.
+
+  Attributes:
+    default: The quantization recipe for global scope of the model.
+    embedding: Recipe for the embedding table.
+    attention: Recipe for the attention blocks. This could be specified with
+      different LayerQuantRecipe for each block by providing a dictionary
+      keyed by the TransformerBlock index.
+    feedforward: Recipe for the feedforward layers. This could be specified with
+      different LayerQuantRecipe for each block by providing a dictionary
+      keyed by the TransformerBlock index.
+  """
+
+  default: Optional[LayerQuantRecipe] = None
+  embedding: Optional[LayerQuantRecipe] = None
+  attention: Union[
+      Optional[LayerQuantRecipe], Optional[dict[int, LayerQuantRecipe]]
+  ] = None
+  feedforward: Union[
+      Optional[LayerQuantRecipe], Optional[dict[int, LayerQuantRecipe]]
+  ] = None
+
+  def __str__(self):
+    return f"""GenerativeQuantRecipe(
+  Default: {self.default}
+  Embedding: {self.embedding}
+  Attention: {self.attention}
+  Feedforward: {self.feedforward}
+)"""
+
+  __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()
+    if self.embedding is not None:
+      self.embedding.verify()
+    if self.attention is not None:
+      if isinstance(self.attention, dict):
+        for recipe in self.attention.values():
+          recipe.verify()
+      else:
+        self.attention.verify()
+    if self.feedforward is not None:
+      if isinstance(self.feedforward, dict):
+        for recipe in self.feedforward.values():
+          recipe.verify()
+      else:
+        self.feedforward.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.GenerativeQuantRecipe(
+    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.FLOAT_CAST,
+      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_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_int8_dynamic_recipe() -> quant_config.QuantConfig:
+  return quant_config.QuantConfig(
+      generative_recipe=quant_recipe.GenerativeQuantRecipe(
+          default=quant_recipe_utils.create_layer_quant_int8_dynamic(),
+      )
+  )
+
+
+def full_fp16_recipe() -> quant_config.QuantConfig:
+  return quant_config.QuantConfig(
+      generative_recipe=quant_recipe.GenerativeQuantRecipe(
+          default=quant_recipe_utils.create_layer_quant_fp16()
+      )
+  )

ai_edge_torch/generative/quantize/supported_schemes.py

@@ -0,0 +1,32 @@
+# 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.INT8, _m.WEIGHT_ONLY, _a.MIN_MAX, _g.CHANNELWISE),
+      (_t.FP32, _t.FP16, _m.WEIGHT_ONLY, _a.FLOAT_CAST, _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/loader_test.py

@@ -0,0 +1,80 @@
+# 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 weight loader utilities.
+
+import os
+import tempfile
+import unittest
+
+import safetensors.torch
+import torch
+
+from ai_edge_torch.generative.examples.tiny_llama import tiny_llama
+from ai_edge_torch.generative.utilities import loader as loading_utils
+
+
+class TestLoader(unittest.TestCase):
+  """Unit tests that check weight loader."""
+
+  def test_load_safetensors(self):
+    with tempfile.TemporaryDirectory() as temp_dir:
+      file_path = os.path.join(temp_dir, "test.safetensors")
+      test_data = {"weight": torch.randn(20, 10), "bias": torch.randn(20)}
+      safetensors.torch.save_file(test_data, file_path)
+
+      loaded_tensors = loading_utils.load_safetensors(file_path)
+      self.assertIn("weight", loaded_tensors)
+      self.assertIn("bias", loaded_tensors)
+
+  def test_load_statedict(self):
+    with tempfile.TemporaryDirectory() as temp_dir:
+      file_path = os.path.join(temp_dir, "test.pt")
+      model = torch.nn.Linear(10, 5)
+      state_dict = model.state_dict()
+      torch.save(state_dict, file_path)
+
+      loaded_tensors = loading_utils.load_pytorch_statedict(file_path)
+      self.assertIn("weight", loaded_tensors)
+      self.assertIn("bias", loaded_tensors)
+
+  def test_model_loader(self):
+    with tempfile.TemporaryDirectory() as temp_dir:
+      file_path = os.path.join(temp_dir, "test.safetensors")
+      test_weights = {
+          "lm_head.weight": torch.randn((32000, 2048)),
+          "model.embed_tokens.weight": torch.randn((32000, 2048)),
+          "model.layers.0.input_layernorm.weight": torch.randn((2048,)),
+          "model.layers.0.mlp.down_proj.weight": torch.randn((2048, 5632)),
+          "model.layers.0.mlp.gate_proj.weight": torch.randn((5632, 2048)),
+          "model.layers.0.mlp.up_proj.weight": torch.randn((5632, 2048)),
+          "model.layers.0.post_attention_layernorm.weight": torch.randn((2048,)),
+          "model.layers.0.self_attn.k_proj.weight": torch.randn((256, 2048)),
+          "model.layers.0.self_attn.o_proj.weight": torch.randn((2048, 2048)),
+          "model.layers.0.self_attn.q_proj.weight": torch.randn((2048, 2048)),
+          "model.layers.0.self_attn.v_proj.weight": torch.randn((256, 2048)),
+          "model.norm.weight": torch.randn((2048,)),
+      }
+      safetensors.torch.save_file(test_weights, file_path)
+      cfg = tiny_llama.get_model_config()
+      cfg.num_layers = 1
+      model = tiny_llama.TinyLLamma(cfg)
+
+      loader = loading_utils.ModelLoader(file_path, tiny_llama.TENSOR_NAMES)
+      # if returns successfully, it means all the tensors were initiallized.
+      loader.load(model, strict=True)
+
+
+if __name__ == "__main__":
+  unittest.main()

ai_edge_torch/generative/test/test_model_conversion.py

@@ -0,0 +1,235 @@
+# 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):
+    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))
+
+    # TODO(b/338288901): re-enable test to check output tensors.
+    skip_output_check = True
+    if skip_output_check is False:
+      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_multi_batches(self):
+    config = toy_model_with_kv_cache.get_model_config()
+    config.batch_size = 2
+    pytorch_model = toy_model_with_kv_cache.ToyModelWithKV(config)
+    idx, input_pos = torch.tensor([[1], [2]], dtype=torch.long), torch.tensor(
+        [10], dtype=torch.int64
+    )
+
+    edge_model = ai_edge_torch.convert(pytorch_model, (idx, input_pos))
+
+    # TODO(b/338288901): re-enable test to check output tensors.
+    skip_output_check = True
+    if skip_output_check is False:
+      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):
+    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))
+
+    # TODO(b/338288901): re-enable test to check output tensors.
+    skip_output_check = True
+    if skip_output_check is False:
+      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))
+
+    # TODO(b/338288901): re-enable test to check output tensors.
+    skip_output_check = True
+    if skip_output_check is False:
+      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):
+    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()
+    )
+
+    # TODO(b/338288901): re-enable test to check output tensors.
+    skip_output_check = True
+    if skip_output_check is False:
+      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(b/338288901): re-enable test to check output tensors.
+    skip_output_check = True
+    if skip_output_check is False:
+      # 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))
+
+    # TODO(b/338288901): re-enable test to check output tensors.
+    skip_output_check = True
+    if skip_output_check is False:
+      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()