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

ai_edge_torch/generative/test/test_quantize.py

@@ -0,0 +1,162 @@
+# 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  # 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
+
+
+class TestVerifyRecipes(unittest.TestCase):
+  """Unit tests that check for model quantization recipes."""
+
+  @parameterized.expand(
+      [
+          (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.expand(
+      [
+          (
+              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(unittest.TestCase):
+  """Test conversion with quantization."""
+
+  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.expand(
+      [
+          (quant_recipes.full_fp16_recipe(), 0.65),
+          (quant_recipes.full_int8_dynamic_recipe(), 0.47),
+          (_attention_int8_dynamic_recipe(), 0.89),
+          (_feedforward_int8_dynamic_recipe(), 0.72),
+      ]
+  )
+  def test_quantize_convert_toy_sizes(self, quant_config, expected_compression):
+    config = toy_model.get_model_config()
+    pytorch_model = toy_model.ToySingleLayerModel(config)
+    idx = torch.unsqueeze(torch.arange(0, 100), 0)
+    input_pos = torch.arange(0, 100)
+
+    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.assertAlmostEqual(
+        len(quantized_model._tflite_model) / len(float_model._tflite_model),
+        expected_compression,
+        delta=0.01,
+    )
+
+  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.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.

ai_edge_torch/generative/utilities/loader.py

@@ -0,0 +1,328 @@
+# 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, List, Tuple
+
+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.
+  """
+  files = []
+  patterns = []
+  if os.path.isdir(full_path):
+    patterns.append(os.path.join(full_path, "*.bin"))
+    patterns.append(os.path.join(full_path, "*.pt"))
+  else:
+    patterns.append(full_path)
+  for pattern in patterns:
+    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 = None
+    attn_key_proj: str = None
+    attn_value_proj: str = None
+    attn_fused_qkv_proj: str = None
+    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_ff_norm: str = None
+    embedding: str = None
+    embedding_position: 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
+  ) -> Tuple[List[str], List[str]]:
+    """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.
+
+    Returns:
+        missing_keys (List[str]): a list of str containing the missing keys.
+        unexpected_keys (List[str]): a list of str containing the unexpected keys.
+
+    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.embedding_position is not None:
+        converted_state["tok_embedding_position"] = state.pop(
+            f"{self._names.embedding_position}"
+        )
+    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())}"
+      )
+    return 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")) or glob.glob(
+          os.path.join(self._file_name, "*.pt")
+      ):
+        return load_pytorch_statedict
+
+    if self._file_name.endswith(".safetensors"):
+      return load_safetensors
+
+    if self._file_name.endswith(".bin") or self._file_name.endswith(".pt"):
+      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}"
+    if self._names.attn_fused_qkv_proj:
+      fused_qkv_name = self._names.attn_fused_qkv_proj.format(idx)
+      converted_state[f"{prefix}.atten_func.qkv_projection.weight"] = state.pop(
+          f"{fused_qkv_name}.weight"
+      )
+    else:
+      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.qkv_projection.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:
+      if self._names.attn_fused_qkv_proj:
+        converted_state[f"{prefix}.atten_func.qkv_projection.bias"] = state.pop(
+            f"{fused_qkv_name}.bias"
+        )
+      else:
+        converted_state[f"{prefix}.atten_func.qkv_projection.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.output_projection.weight"] = state.pop(
+        f"{o_name}.weight"
+    )
+    if config.attn_config.output_proj_use_bias:
+      converted_state[f"{prefix}.atten_func.output_projection.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:
+    if config.attn_config.qkv_fused_interleaved:
+      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)
+    else:
+      return torch.cat([q, k, v], dim=0)