ai-edge-torch-nightly 0.3.0.dev20250114__py3-none-any.whl

ai_edge_torch/generative/test/test_quantize.py

@@ -0,0 +1,183 @@
+# 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 ai_edge_torch
+from ai_edge_torch import config
+from ai_edge_torch.generative.examples.test_models import toy_model  # NOQA
+from ai_edge_torch.generative.quantize import quant_recipe
+from ai_edge_torch.generative.quantize import quant_recipe_utils
+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.quantize import quant_config
+from ai_edge_torch.testing import model_coverage
+import torch
+
+from absl.testing import absltest as googletest
+from absl.testing import parameterized
+
+
+class TestVerifyRecipes(parameterized.TestCase):
+  """Unit tests that check for model quantization recipes."""
+
+  @parameterized.parameters([
+      (Dtype.FP32, Dtype.FP32),
+      (Dtype.INT8, Dtype.INT8),
+      (Dtype.INT8, Dtype.FP16),
+      (Dtype.FP16, Dtype.INT8),
+      (Dtype.FP16, Dtype.FP16),
+  ])
+  def test_verify_invalid_recipes(
+      self,
+      activation,
+      weight,
+  ):
+    for m in Mode:
+      for a in Algorithm:
+        for g in Granularity:
+          with self.assertRaises(ValueError):
+            quant_recipe.LayerQuantRecipe(activation, weight, m, a, g).verify()
+
+  @parameterized.parameters([
+      (
+          Dtype.FP32,
+          Dtype.INT8,
+          Mode.DYNAMIC_RANGE,
+          Algorithm.MIN_MAX,
+          Granularity.CHANNELWISE,
+      ),
+      (
+          Dtype.FP32,
+          Dtype.INT8,
+          Mode.WEIGHT_ONLY,
+          Algorithm.MIN_MAX,
+          Granularity.CHANNELWISE,
+      ),
+      (
+          Dtype.FP32,
+          Dtype.FP16,
+          Mode.WEIGHT_ONLY,
+          Algorithm.FLOAT_CAST,
+          Granularity.NONE,
+      ),
+  ])
+  def test_verify_valid_recipes(
+      self,
+      activation,
+      weight,
+      mode,
+      algo,
+      granularity,
+  ):
+    quant_recipe.LayerQuantRecipe(
+        activation, weight, mode, algo, granularity
+    ).verify()
+
+
+class TestQuantizeConvert(parameterized.TestCase):
+  """Test conversion with quantization."""
+
+  def setUp(self):
+    super().setUp()
+    torch.manual_seed(0)
+    torch._dynamo.reset()
+
+  def _attention_int8_dynamic_recipe() -> quant_config.QuantConfig:
+    return quant_config.QuantConfig(
+        generative_recipe=quant_recipe.GenerativeQuantRecipe(
+            attention=quant_recipe_utils.create_layer_quant_int8_dynamic(),
+        )
+    )
+
+  def _feedforward_int8_dynamic_recipe() -> quant_config.QuantConfig:
+    return quant_config.QuantConfig(
+        generative_recipe=quant_recipe.GenerativeQuantRecipe(
+            feedforward=quant_recipe_utils.create_layer_quant_int8_dynamic(),
+        )
+    )
+
+  @parameterized.parameters([
+      (quant_recipes.full_fp16_recipe()),
+      (quant_recipes.full_int8_dynamic_recipe()),
+      (quant_recipes.full_int8_weight_only_recipe()),
+      (_attention_int8_dynamic_recipe()),
+      (_feedforward_int8_dynamic_recipe()),
+  ])
+  def test_quantize_convert_toy_sizes(self, quant_config):
+    config = toy_model.get_model_config()
+    pytorch_model = toy_model.ToySingleLayerModel(config)
+    idx = torch.unsqueeze(torch.arange(0, 100, dtype=torch.int), 0)
+    input_pos = torch.arange(0, 100, dtype=torch.int)
+
+    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),
+        "Quantized model isn't smaller than F32 model.",
+    )
+
+  def test_quantize_convert_toy_weight_sharing(self):
+    config = toy_model.get_model_config()
+    pytorch_model = toy_model.ToySingleLayerModelWeightSharing(config)
+    idx = torch.unsqueeze(torch.arange(0, 100, dtype=torch.int), 0)
+    input_pos = torch.arange(0, 100, dtype=torch.int)
+
+    quant_config = quant_recipes.full_int8_dynamic_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),
+        "Quantized model isn't smaller than F32 model.",
+    )
+
+  def test_quantize_convert_compare_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.int), 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__":
+  googletest.main()

ai_edge_torch/generative/test/utils.py

@@ -0,0 +1,82 @@
+# 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 utils for testing."""
+
+import logging
+
+from ai_edge_torch import model
+from ai_edge_torch.generative.layers import kv_cache as kv_utils
+from ai_edge_torch.lowertools import common_utils
+import numpy as np
+import torch
+from torch.utils import _pytree as pytree
+
+
+def compare_tflite_torch(
+    edge_model: model.Model,
+    torch_model: torch.nn.Module,
+    tokens: torch.Tensor,
+    input_pos: torch.Tensor,
+    kv_cache: kv_utils.KVCache,
+    signature_name: str,
+    atol: float = 1e-5,
+    rtol: float = 1e-5,
+    **kwargs,
+) -> bool:
+  """Compares torch models and TFLite models."""
+  values, spec = pytree.tree_flatten({"kv_cache": kv_cache})
+  flat_names = common_utils.flat_dict_names(spec.children_specs, spec.context)
+  torch_output = torch_model(tokens, input_pos, kv_cache, **kwargs)
+
+  if "pixel_values" in kwargs:
+    kwargs["pixel_values"] = kwargs["pixel_values"].numpy()
+  kwargs.update({k: v.numpy() for k, v in zip(flat_names, values)})
+  edge_output = edge_model(
+      signature_name=signature_name,
+      tokens=tokens.numpy(),
+      input_pos=input_pos.numpy(),
+      **kwargs,
+  )
+
+  return compare_logits(
+      edge_output["logits"], torch_output["logits"].detach().numpy(), atol, rtol
+  )
+
+
+def compare_logits(
+    edge_logits: np.ndarray,
+    torch_logits: dict[str, torch.Tensor],
+    atol: float = 1e-5,
+    rtol: float = 1e-5,
+) -> bool:
+  """Compares logits from edge model and torch model."""
+  if np.allclose(edge_logits, torch_logits, rtol, atol, equal_nan=True):
+    return True
+
+  logging.info("edge_logits: %s", edge_logits)
+  logging.info("torch_logits: %s", torch_logits)
+
+  orig_atol = atol
+  while rtol < 1:
+    atol = orig_atol
+    while atol < 1:
+      if np.allclose(edge_logits, torch_logits, rtol, atol, equal_nan=True):
+        logging.info("Got allclose true with atol=%s, rtol=%s", atol, rtol)
+        return False
+      atol *= 10
+    rtol *= 10
+  logging.info("allclose failed with reasonable atol and rtol.")
+  return False

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.

ai_edge_torch/generative/utilities/converter.py

@@ -0,0 +1,215 @@
+# 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 model conversion."""
+
+import os
+from typing import Optional, Union
+from ai_edge_torch._convert import converter as converter_utils
+from ai_edge_torch.generative.layers import lora as lora_utils
+import ai_edge_torch.generative.layers.kv_cache as kv_utils
+import ai_edge_torch.generative.layers.model_config as cfg
+from ai_edge_torch.generative.quantize import quant_recipes
+from ai_edge_torch.generative.utilities.model_builder import ExportConfig
+import torch
+
+
+class ExportableModule(torch.nn.Module):
+
+  def __init__(self, module, **extra_kwargs):
+    super().__init__()
+    self.module = module
+    self.extra_kwargs = extra_kwargs
+
+  def forward(self, *export_args, **export_kwargs):
+    full_kwargs = {**export_kwargs, **self.extra_kwargs}
+    return self.module(*export_args, **full_kwargs)
+
+
+def convert_to_tflite(
+    pytorch_model: torch.nn.Module,
+    output_path: str,
+    output_name_prefix: str,
+    prefill_seq_len: Union[int, list[int]],
+    pixel_values_size: torch.Size = None,
+    quantize: bool = True,
+    config: cfg.ModelConfig = None,
+    lora_ranks: Optional[list[int]] = None,
+    export_config: ExportConfig = None,
+):
+  """Converts a nn.Module model to multi-signature tflite model.
+
+  A PyTorch model will be converted to a tflite model with several signatures:
+    * "prefill_[prefill_seq_len]" (or "prefill" if only one prefill_seq_len is
+        passed),
+    * "prefill_[preill_seq_len]_pixel" (or "prefill_pixel" if only one
+        prefill_seq_len is passed) if num_pixel_values > 0, and
+    * "decode".
+
+  "prefill_[prefill_seq_len]" (or "prefill" if only one prefill_seq_len is
+  passed) signature takes as a sample input:
+    * a tensor of shape [1, prefill_seq_len] of token sequence,
+    * a tensor of shape [1, prefill_seq_len] of token positions, and
+    * an external KV cache.
+
+  If num_pixel_values > 0, "prefill_[prefill_seq_len]_pixel" (or "prefill_pixel"
+  if only one prefill_seq_len is passed) signature takes as a sample input:
+    * a tensor of shape [1, prefill_seq_len] of token sequence,
+    * a tensor of shape [1, prefill_seq_len] of token positions,
+    * an external KV cache, and
+    * a tensor of shape [1, num_pixel_values] of pixel values.
+
+  "decode" signature takes as a sample input:
+    * a tensor of shape [1, 1] of token sequence,
+    * a tensor of shape [1, 1] of the token position, and
+    * an external KV cache.
+
+  The final tflite model will be exported to tflite_path.
+
+  Args:
+      pytorch_model (torch.nn.Module): PyTorch model to convert to tflite.
+      output_path (str): The path to export the tflite model.
+      output_name_prefix (str): The prefix of the tflite model name.
+      prefill_seq_len (Union[int, list[int]]): The prefill sequence length to
+        use. If a list, the model will have multiple prefill signatures.
+      pixel_values_size (torch.Size, optional): The size of pixel values to pass
+        to the model. If None, the model is not expected to take pixel values.
+      quantize (bool, optional): Whether the model should be quanized. Defaults
+        to True.
+      config (cfg.ModelConfig, optional): The model config used to configure KV
+        cache. If None, it uses the config of the pytorch_model.
+      lora_ranks (list[int], optional): The ranks of the LORA layers. If None,
+        no LoRA signatures will be added.
+  """
+  # pylint: disable=protected-access
+  torch._dynamo.config.cache_size_limit = 64
+
+  config = config if config else pytorch_model.config
+  prefill_seq_lens = (
+      [prefill_seq_len] if isinstance(prefill_seq_len, int) else prefill_seq_len
+  )
+  loras = [None]
+  if lora_ranks is not None:
+    for rank in lora_ranks:
+      lora = lora_utils.LoRA.zeros(rank, config)
+      loras.append(lora)
+
+  quant_suffix = 'q8' if quantize else 'f32'
+  kv_size = config.kv_cache_max_len
+  lora_suffix = (
+      '' if not lora_ranks else f'_lora{",".join(map(str, lora_ranks))}'
+  )
+  output_filename = (
+      f'{output_name_prefix}_{quant_suffix}_ekv{kv_size}{lora_suffix}.tflite'
+  )
+  output_file = os.path.join(output_path, output_filename)
+
+  _export_helper(
+      pytorch_model,
+      output_file,
+      prefill_seq_lens,
+      pixel_values_size,
+      quantize,
+      config,
+      loras,
+      export_config,
+  )
+
+
+def _export_helper(
+    pytorch_model: torch.nn.Module,
+    output_file: str,
+    prefill_seq_lens: list[int],
+    pixel_values_size: torch.Size,
+    quantize: bool,
+    config: cfg.ModelConfig,
+    loras: list[None | lora_utils.LoRA],
+    export_config: ExportConfig,
+):
+  """Helper function to export a model to tflite."""
+  prefill_tokens_list = []
+  prefill_input_pos_list = []
+  for seq_len in prefill_seq_lens:
+    prefill_tokens_list.append(torch.full((1, seq_len), 0, dtype=torch.int))
+    prefill_input_pos_list.append(torch.arange(0, seq_len, dtype=torch.int))
+
+  prefill_pixel_values = (
+      torch.full((1,) + pixel_values_size, 0, dtype=torch.float32)
+      if pixel_values_size
+      else None
+  )
+
+  decode_token = torch.tensor([[0]], dtype=torch.int)
+  decode_input_pos = torch.tensor([0], dtype=torch.int)
+  kv = kv_utils.KVCache.from_model_config(config)
+
+  quant_config = quant_recipes.full_int8_dynamic_recipe() if quantize else None
+
+  # For export, we create a module that captures any non-exportable,
+  # arugments, e.g. the generation config object.
+  mod = ExportableModule(pytorch_model, export_config=export_config)
+
+  converter = converter_utils.Converter()
+  for lora in loras:
+    for i in range(len(prefill_seq_lens)):
+      prefill_seq_len = prefill_seq_lens[i]
+      prefill_tokens = prefill_tokens_list[i]
+      prefill_input_pos = prefill_input_pos_list[i]
+      if i == 0 and len(prefill_seq_lens) == 1:
+        prefill_signature_name = 'prefill'
+      else:
+        prefill_signature_name = f'prefill_{prefill_seq_len}'
+
+      sample_kwargs = {
+          'tokens': prefill_tokens,
+          'input_pos': prefill_input_pos,
+          'kv_cache': kv,
+      }
+      if lora is not None:
+        prefill_signature_name += f'_lora_r{lora.get_rank()}'
+        sample_kwargs['lora'] = lora
+
+      converter.add_signature(
+          prefill_signature_name,
+          mod,
+          sample_kwargs=sample_kwargs,
+      )
+
+      if prefill_pixel_values is not None:
+        converter.add_signature(
+            prefill_signature_name + '_pixel',
+            mod,
+            sample_kwargs={
+                **sample_kwargs,
+                'pixel_values': prefill_pixel_values,
+            },
+        )
+
+    sample_kwargs = {
+        'tokens': decode_token,
+        'input_pos': decode_input_pos,
+        'kv_cache': kv,
+    }
+    if lora is not None:
+      sample_kwargs['lora'] = lora
+
+    converter.add_signature(
+        'decode' if lora is None else f'decode_lora_r{lora.get_rank()}',
+        mod,
+        sample_kwargs=sample_kwargs,
+    )
+
+  edge_model = converter.convert(quant_config=quant_config)
+  edge_model.export(output_file)

ai_edge_torch/generative/utilities/dynamic_update_slice.py

@@ -0,0 +1,56 @@
+# 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 Sequence
+from ai_edge_torch.odml_torch import lowerings
+from jax._src.lib.mlir import ir
+from jax._src.lib.mlir.dialects import hlo as stablehlo
+import torch
+
+
+# Use torch.library.custom_op to define a new custom operator.
+# TODO: Update impl for multiple non-trivial start_indices
+@torch.library.custom_op("ai_edge_torch::dynamic_update_slice", mutates_args=())
+def dynamic_update_slice(
+    in_tensor: torch.Tensor,
+    update: torch.Tensor,
+    start_indices: Sequence[torch.Tensor],
+) -> torch.Tensor:
+  compare_size = torch.tensor(in_tensor.size()) == torch.tensor(update.size())
+  mismatch = torch.nonzero(~compare_size, as_tuple=False)
+  dim = mismatch[0].item() if len(mismatch) > 0 else 0
+  start = start_indices[dim].item()
+  end = start + update.shape[dim]
+  indices = torch.arange(start, end).to(torch.long)
+  return in_tensor.index_copy(dim, indices, update)
+
+
+# Use register_fake to add a ``FakeTensor`` kernel for the operator
+@dynamic_update_slice.register_fake
+def _(in_tensor, update, start_indices):
+  return in_tensor.clone().detach()
+
+
+@lowerings.lower(torch.ops.ai_edge_torch.dynamic_update_slice)
+def _dynamic_update_slice_lower(
+    lctx,
+    in_tensor: ir.Value,
+    update: ir.Value,
+    start_indices: Sequence[ir.Value],
+):
+  return stablehlo.dynamic_update_slice(in_tensor, update, start_indices)