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

ai_edge_torch/generative/quantize/quant_recipes.py

@@ -0,0 +1,56 @@
+# 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_int8_weight_only_recipe() -> quant_config.QuantConfig:
+  return quant_config.QuantConfig(
+      generative_recipe=quant_recipe.GenerativeQuantRecipe(
+          default=quant_recipe_utils.create_layer_quant_int8_weight_only(),
+      )
+  )
+
+
+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/test_custom_dus.py

@@ -0,0 +1,107 @@
+# 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 suite of tests to validate the Dynamic Update Slice Custom Op."""
+
+from ai_edge_torch.generative.layers import kv_cache as kv_utils
+import ai_edge_torch.generative.layers.model_config as cfg
+import torch
+from torch import nn
+
+from absl.testing import absltest as googletest, parameterized
+
+
+def updated_slice_matches(buffer, update, index):
+  indexer = [slice(i, i + d) for i, d in zip(index, update.shape)]
+  buf = buffer[indexer]
+  return torch.allclose(buf, update)
+
+
+def intT(x):
+  return torch.tensor(x).int()
+
+
+class DUSMod(nn.Module):
+
+  def forward(self, buffer, update, index):
+    out = dynamic_update_slice(buffer, update, index)
+    out = out * 2
+    return out
+
+
+@googletest.skip('Enable this when odml_torch is default b/373387583')
+class TestCustomDUS(parameterized.TestCase):
+
+  @parameterized.named_parameters(
+      (
+          'DUS_whole_buffer',
+          torch.randn(1, 1280, 4, 64),
+          torch.randn([1, 1024, 4, 64]),
+          [intT(0), intT(0), intT(0), intT(0)],
+      ),
+      (
+          'DUS_kv_example',
+          torch.randn(2, 1280, 4, 64),
+          torch.randn([2, 1024, 4, 64]),
+          [intT(0), intT(0), intT(0), intT(0)],
+      ),
+      (
+          'DUS_3d',
+          torch.randn(2, 256, 4, 64),
+          torch.randn([2, 256, 2, 64]),
+          [intT(0), intT(0), intT(2), intT(0)],
+      ),
+      (
+          'DUS_3d_v2',
+          torch.randn(2, 256, 4, 64),
+          torch.randn([2, 256, 3, 64]),
+          [intT(0), intT(0), intT(1), intT(0)],
+      ),
+      (
+          'DUS_3d_v3',
+          torch.randn(6, 8, 32),
+          torch.randn([6, 3, 32]),
+          [intT(0), intT(5), intT(0)],
+      ),
+      (
+          'DUS_2d',
+          torch.randn(8, 32),
+          torch.randn([8, 12]),
+          [intT(0), intT(20)],
+      ),
+  )
+  def test_opcheck_dynamic_update_slice(self, buffer, update, indices):
+    torch.library.opcheck(dynamic_update_slice, (buffer, update, indices))
+    out = dynamic_update_slice(buffer, update, indices)
+    self.assertTrue(updated_slice_matches(out, update, indices))
+
+  def test_exported_program(self):
+    buffer = torch.randn(1, 1280, 4, 64)
+    update = torch.randn([1, 1024, 4, 64])
+    index = [intT(0), intT(0), intT(0), intT(0)]
+    dm = DUSMod()
+    ep = torch.export.export(dm, (buffer, update, index))
+    dus_in_exported_program = False
+    for node in ep.graph.nodes:
+      if node.op == 'call_function':
+        if node.target.__name__.startswith('dynamic_update_slice'):
+          dus_in_exported_program = True
+          break
+
+    self.assertTrue(dus_in_exported_program)
+
+
+if __name__ == '__main__':
+  googletest.main()

ai_edge_torch/generative/test/test_kv_cache.py

@@ -0,0 +1,120 @@
+# 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 suite of tests to validate KV Cache layer."""
+
+from ai_edge_torch.generative.layers import kv_cache as kv_utils
+import ai_edge_torch.generative.layers.model_config as cfg
+import torch
+
+from absl.testing import absltest as googletest
+
+
+class TestKVLayers(googletest.TestCase):
+
+  def _get_test_config(
+      self, num_layers, head_dim, num_query_groups, kv_cache_max_len
+  ):
+    attn_config = cfg.AttentionConfig(
+        num_heads=1, head_dim=head_dim, num_query_groups=num_query_groups
+    )
+    block_config = cfg.TransformerBlockConfig(
+        attn_config=attn_config, ff_config=None
+    )
+    config = cfg.ModelConfig(
+        kv_cache_max_len=kv_cache_max_len,
+        embedding_dim=head_dim,
+        block_configs=block_config,
+        num_layers=num_layers,
+        max_seq_len=None,
+        vocab_size=None,
+    )
+    return config
+
+  def test_cache_udpate(self):
+    N = 1
+    HEAD_DIM = 2
+    NUM_QG = 1
+    KV_LEN = 4
+    config = self._get_test_config(
+        num_layers=N,
+        head_dim=HEAD_DIM,
+        num_query_groups=NUM_QG,
+        kv_cache_max_len=KV_LEN,
+    )
+    kv = kv_utils.KVCache.from_model_config(config)
+    entry = kv.caches[0]
+    # single-slice update
+    input_pos = torch.tensor([1])
+    k_slice = v_slice = torch.full(
+        (1, 1, NUM_QG, HEAD_DIM), 5, dtype=torch.float
+    )
+    updated_entry = kv_utils.update(entry, input_pos, k_slice, v_slice)
+    self.assertEqual(
+        updated_entry.k_cache.numpy().flatten().tolist(),
+        [0, 0, 5, 5, 0, 0, 0, 0],
+    )
+    self.assertEqual(
+        updated_entry.v_cache.numpy().flatten().tolist(),
+        [0, 0, 5, 5, 0, 0, 0, 0],
+    )
+    # multi-slice update
+    input_pos = torch.tensor([0, 1])
+    k_slice = v_slice = torch.full(
+        (1, 2, NUM_QG, HEAD_DIM), 7, dtype=torch.float
+    )
+    updated_entry = kv_utils.update(entry, input_pos, k_slice, v_slice)
+    self.assertEqual(
+        updated_entry.k_cache.numpy().flatten().tolist(),
+        [7, 7, 7, 7, 0, 0, 0, 0],
+    )
+    self.assertEqual(
+        updated_entry.v_cache.numpy().flatten().tolist(),
+        [7, 7, 7, 7, 0, 0, 0, 0],
+    )
+
+  def test_serialization(self):
+    class TestModel(torch.nn.Module):
+
+      def forward(self, kv: kv_utils.KVCache) -> kv_utils.KVCache:
+        updated_kv_entries = [
+            kv_utils.KVCacheEntry(
+                torch.zeros_like(entry.k_cache), torch.zeros_like(entry.v_cache)
+            )
+            for entry in kv.caches
+        ]
+        return kv_utils.KVCache(updated_kv_entries)
+
+    N = 1
+    HEAD_DIM = 2
+    NUM_QG = 1
+    KV_LEN = 4
+    config = self._get_test_config(
+        num_layers=N,
+        head_dim=HEAD_DIM,
+        num_query_groups=NUM_QG,
+        kv_cache_max_len=KV_LEN,
+    )
+    kv = kv_utils.KVCache.from_model_config(config)
+    model = TestModel()
+    exported_program = torch.export.export(model, (kv,))
+    input_specs = exported_program.graph_signature.input_specs
+    self.assertEqual(len(input_specs), 2)
+    self.assertEqual(input_specs[0].arg.name, "kv_k_0")
+    self.assertEqual(input_specs[1].arg.name, "kv_v_0")
+
+
+if __name__ == "__main__":
+  googletest.main()

ai_edge_torch/generative/test/test_loader.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.
+# ==============================================================================
+# Testing weight loader utilities.
+
+import os
+import tempfile
+
+from ai_edge_torch.generative.examples.tiny_llama import tiny_llama
+from ai_edge_torch.generative.utilities import loader as loading_utils
+from ai_edge_torch.generative.utilities import model_builder
+import safetensors.torch
+import torch
+
+from absl.testing import absltest as googletest
+
+
+class TestLoader(googletest.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 = model_builder.DecoderOnlyModel(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__":
+  googletest.main()

ai_edge_torch/generative/test/test_lora.py

@@ -0,0 +1,147 @@
+# Copyright 2025 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 suite of tests to validate LoRA utilities."""
+
+from ai_edge_torch.generative.layers import lora as lora_utils
+import ai_edge_torch.generative.layers.model_config as cfg
+import torch
+from absl.testing import absltest as googletest
+from tensorflow.python.platform import resource_loader  # pylint: disable=g-direct-tensorflow-import
+
+
+class TestLora(googletest.TestCase):
+  """Tests for LoRA utilities."""
+
+  def test_safetensors_builder(self):
+    """Converts a safetensors file to a LoRA module."""
+
+    tensor_names = lora_utils.LoRATensorNames(
+        attn_query_w_a=(
+            "base_model.model.model.layers.{}.self_attn.q_proj.lora_A.weight"
+        ),
+        attn_query_w_b=(
+            "base_model.model.model.layers.{}.self_attn.q_proj.lora_B.weight"
+        ),
+        attn_key_w_a=(
+            "base_model.model.model.layers.{}.self_attn.k_proj.lora_A.weight"
+        ),
+        attn_key_w_b=(
+            "base_model.model.model.layers.{}.self_attn.k_proj.lora_B.weight"
+        ),
+        attn_value_w_a=(
+            "base_model.model.model.layers.{}.self_attn.v_proj.lora_A.weight"
+        ),
+        attn_value_w_b=(
+            "base_model.model.model.layers.{}.self_attn.v_proj.lora_B.weight"
+        ),
+        attn_output_w_a=(
+            "base_model.model.model.layers.{}.self_attn.o_proj.lora_A.weight"
+        ),
+        attn_output_w_b=(
+            "base_model.model.model.layers.{}.self_attn.o_proj.lora_B.weight"
+        ),
+    )
+
+    safetensors_file = resource_loader.get_path_to_datafile(
+        "fixtures/test_lora_rank16.safetensors"
+    )
+    config = self._get_test_config(
+        num_layers=1,
+        head_dim=8,
+        num_query_groups=1,
+        kv_cache_max_len=16,
+    )
+    lora = lora_utils.LoRA.from_safetensors(
+        safetensors_file,
+        scale=1.0,
+        lora_tensor_names=tensor_names,
+        config=config,
+    )
+    self.assertEqual(lora.get_rank(), 16)
+
+  def test_torch_export(self):
+    """Tests the export of the LoRA module."""
+
+    class TestModel(torch.nn.Module):
+
+      def forward(self, x: torch.Tensor, lora: lora_utils.LoRA) -> torch.Tensor:
+        x += lora_utils.apply_lora(x, lora.adapters[0].attention.query)
+        return x
+
+    n = 1
+    head_dim = 2
+    num_query_groups = 1
+    key_length = 4
+    config = self._get_test_config(
+        num_layers=n,
+        head_dim=head_dim,
+        num_query_groups=num_query_groups,
+        kv_cache_max_len=key_length,
+    )
+    inputs = torch.zeros((n, 1, head_dim))
+    lora = lora_utils.LoRA.zeros(rank=16, config=config)
+    model = TestModel()
+    exported_program = torch.export.export(model, (inputs, lora))
+    input_specs = exported_program.graph_signature.input_specs
+    # 9 inputs: 1 for x, 2 for query lora, 2 for key lora, 2 for value lora,
+    # 2 for output lora.
+    self.assertLen(input_specs, 9)
+    self.assertEqual(input_specs[0].arg.name, "x")
+    self.assertEqual(input_specs[1].arg.name, "lora_atten_q_a_prime_weight_0")
+    self.assertEqual(input_specs[2].arg.name, "lora_atten_q_b_prime_weight_0")
+    self.assertEqual(input_specs[3].arg.name, "lora_atten_k_a_prime_weight_0")
+    self.assertEqual(input_specs[4].arg.name, "lora_atten_k_b_prime_weight_0")
+    self.assertEqual(input_specs[5].arg.name, "lora_atten_v_a_prime_weight_0")
+    self.assertEqual(input_specs[6].arg.name, "lora_atten_v_b_prime_weight_0")
+    self.assertEqual(input_specs[7].arg.name, "lora_atten_o_a_prime_weight_0")
+    self.assertEqual(input_specs[8].arg.name, "lora_atten_o_b_prime_weight_0")
+
+  def test_lora_tflite_serialization(self):
+    """Tests the serialization of the LoRA module."""
+    config = self._get_test_config(
+        num_layers=2,
+        head_dim=8,
+        num_query_groups=1,
+        kv_cache_max_len=16,
+    )
+    lora = lora_utils.LoRA.random(rank=16, config=config)
+    flatbuffer_model = lora.to_tflite()
+    recovered_lora = lora_utils.LoRA.from_flatbuffers(flatbuffer_model)
+    self.assertEqual(lora, recovered_lora)
+
+  def _get_test_config(
+      self, num_layers, head_dim, num_query_groups, kv_cache_max_len
+  ):
+    """Returns a test model config."""
+    attn_config = cfg.AttentionConfig(
+        num_heads=1, head_dim=head_dim, num_query_groups=num_query_groups
+    )
+    block_config = cfg.TransformerBlockConfig(
+        attn_config=attn_config, ff_config=None
+    )
+    config = cfg.ModelConfig(
+        kv_cache_max_len=kv_cache_max_len,
+        embedding_dim=head_dim,
+        block_configs=block_config,
+        num_layers=num_layers,
+        max_seq_len=None,
+        vocab_size=None,
+    )
+    return config
+
+
+if __name__ == "__main__":
+  googletest.main()