ai-edge-torch-nightly 0.3.0.dev20240909__py3-none-any.whl → 0.3.0.dev20240913__py3-none-any.whl

ai_edge_torch/generative/test/test_model_conversion.py

@@ -12,16 +12,15 @@
 # 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
+
+"""Testing model conversion for a few gen-ai models."""
 
 import ai_edge_torch
 from ai_edge_torch import config as ai_edge_config
-from ai_edge_torch.generative.examples.gemma import gemma, gemma2
-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.test_models import toy_model_with_kv_cache
 from ai_edge_torch.generative.examples.tiny_llama import tiny_llama
-from ai_edge_torch.testing import model_coverage
+from ai_edge_torch.generative.layers import kv_cache
+from ai_edge_torch.generative.test import utils as test_utils
 import numpy as np
 import torch
 
@@ -49,22 +48,32 @@ class TestModelConversion(googletest.TestCase):
   )
   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).eval()
-    idx, input_pos = torch.tensor([[1]], dtype=torch.long), torch.tensor(
+    pytorch_model = toy_model_with_kv_cache.ToyModelWithKVCache(config).eval()
+    tokens, 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))
+    kv = kv_cache.KVCache.from_model_config(config)
+
+    edge_model = ai_edge_torch.convert(
+        pytorch_model,
+        sample_kwargs={
+            "tokens": tokens,
+            "input_pos": input_pos,
+            "kv_cache": kv,
+        },
+    )
     edge_model.set_interpreter_builder(
         self._interpreter_builder(edge_model.tflite_model())
     )
 
     self.assertTrue(
-        model_coverage.compare_tflite_torch(
+        test_utils.compare_tflite_torch(
             edge_model,
             pytorch_model,
-            (idx, input_pos),
-            num_valid_inputs=1,
+            tokens,
+            input_pos,
+            kv,
+            signature_name="serving_default",
             atol=1e-5,
             rtol=1e-5,
         )
@@ -77,22 +86,32 @@ class TestModelConversion(googletest.TestCase):
   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).eval()
-    idx, input_pos = torch.tensor([[1]], dtype=torch.long), torch.tensor(
+    pytorch_model = toy_model_with_kv_cache.ToyModelWithKVCache(config).eval()
+    tokens, 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))
+    kv = kv_cache.KVCache.from_model_config(config)
+
+    edge_model = ai_edge_torch.convert(
+        pytorch_model,
+        sample_kwargs={
+            "tokens": tokens,
+            "input_pos": input_pos,
+            "kv_cache": kv,
+        },
+    )
     edge_model.set_interpreter_builder(
         self._interpreter_builder(edge_model.tflite_model())
     )
 
     self.assertTrue(
-        model_coverage.compare_tflite_torch(
+        test_utils.compare_tflite_torch(
             edge_model,
             pytorch_model,
-            (idx, input_pos),
-            num_valid_inputs=1,
+            tokens,
+            input_pos,
+            kv,
+            signature_name="serving_default",
             atol=1e-5,
             rtol=1e-5,
         )
@@ -104,7 +123,7 @@ class TestModelConversion(googletest.TestCase):
   )
   def test_tiny_llama_multisig(self):
     config = tiny_llama.get_fake_model_config()
-    pytorch_model = tiny_llama.TinyLLamma(config).eval()
+    pytorch_model = tiny_llama.TinyLlama(config).eval()
 
     # prefill
     seq_len = 10
@@ -117,37 +136,56 @@ class TestModelConversion(googletest.TestCase):
     decode_token = torch.tensor([[1]], dtype=torch.long)
     decode_input_pos = torch.tensor([5], dtype=torch.int64)
 
+    kv = kv_cache.KVCache.from_model_config(config)
+
     edge_model = (
         ai_edge_torch.signature(
-            "prefill", pytorch_model, (prefill_tokens, prefill_input_pos)
+            "prefill",
+            pytorch_model,
+            sample_kwargs={
+                "tokens": prefill_tokens,
+                "input_pos": prefill_input_pos,
+                "kv_cache": kv,
+            },
+        )
+        .signature(
+            "decode",
+            pytorch_model,
+            sample_kwargs={
+                "tokens": decode_token,
+                "input_pos": decode_input_pos,
+                "kv_cache": kv,
+            },
         )
-        .signature("decode", pytorch_model, (decode_token, decode_input_pos))
         .convert()
     )
     edge_model.set_interpreter_builder(
         self._interpreter_builder(edge_model.tflite_model())
     )
 
-    copied_model = copy.deepcopy(pytorch_model)
-    copied_edge = copy.deepcopy(edge_model)
-
     self.assertTrue(
-        model_coverage.compare_tflite_torch(
+        test_utils.compare_tflite_torch(
             edge_model,
             pytorch_model,
-            (prefill_tokens, prefill_input_pos),
+            prefill_tokens,
+            prefill_input_pos,
+            kv,
             signature_name="prefill",
-            num_valid_inputs=1,
+            atol=1e-5,
+            rtol=1e-5,
         )
     )
 
     self.assertTrue(
-        model_coverage.compare_tflite_torch(
-            copied_edge,
-            copied_model,
-            (decode_token, decode_input_pos),
+        test_utils.compare_tflite_torch(
+            edge_model,
+            pytorch_model,
+            decode_token,
+            decode_input_pos,
+            kv,
             signature_name="decode",
-            num_valid_inputs=1,
+            atol=1e-5,
+            rtol=1e-5,
         )
     )
 

ai_edge_torch/generative/test/test_model_conversion_large.py

@@ -12,16 +12,16 @@
 # 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
+
+"""Testing model conversion for a few gen-ai models."""
 
 import ai_edge_torch
 from ai_edge_torch import config as ai_edge_config
-from ai_edge_torch.generative.examples.gemma import gemma, gemma2
-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
+from ai_edge_torch.generative.examples.gemma import gemma
+from ai_edge_torch.generative.examples.gemma import gemma2
+from ai_edge_torch.generative.examples.phi import phi2
+from ai_edge_torch.generative.layers import kv_cache
+from ai_edge_torch.generative.test import utils as test_utils
 import numpy as np
 import torch
 
@@ -55,18 +55,28 @@ class TestModelConversion(googletest.TestCase):
     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))
+    kv = kv_cache.KVCache.from_model_config(config)
+
+    edge_model = ai_edge_torch.convert(
+        model,
+        sample_kwargs={
+            "tokens": tokens,
+            "input_pos": input_pos,
+            "kv_cache": kv,
+        },
+    )
     edge_model.set_interpreter_builder(
         self._interpreter_builder(edge_model.tflite_model())
     )
 
     self.assertTrue(
-        model_coverage.compare_tflite_torch(
+        test_utils.compare_tflite_torch(
             edge_model,
             model,
-            (tokens, input_pos),
-            num_valid_inputs=1,
+            tokens,
+            input_pos,
+            kv,
+            signature_name="serving_default",
             atol=1e-2,
             rtol=1e-5,
         )
@@ -85,23 +95,31 @@ class TestModelConversion(googletest.TestCase):
     prefill_tokens = torch.full((1, 10), 0, dtype=torch.long, device="cpu")
     prefill_tokens[0, :4] = idx
     prefill_input_pos = torch.arange(0, 10)
+    kv = kv_cache.KVCache.from_model_config(config)
 
     edge_model = ai_edge_torch.signature(
-        "prefill", model, (prefill_tokens, prefill_input_pos)
+        "prefill",
+        model,
+        sample_kwargs={
+            "tokens": prefill_tokens,
+            "input_pos": prefill_input_pos,
+            "kv_cache": kv,
+        },
     ).convert()
     edge_model.set_interpreter_builder(
         self._interpreter_builder(edge_model.tflite_model())
     )
 
     self.assertTrue(
-        model_coverage.compare_tflite_torch(
+        test_utils.compare_tflite_torch(
             edge_model,
             model,
-            (prefill_tokens, prefill_input_pos),
+            prefill_tokens,
+            prefill_input_pos,
+            kv,
             signature_name="prefill",
-            num_valid_inputs=1,
-            atol=1e-2,
-            rtol=1e-5,
+            atol=1e-1,
+            rtol=1e-3,
         )
     )
 
@@ -117,18 +135,28 @@ class TestModelConversion(googletest.TestCase):
     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))
+    kv = kv_cache.KVCache.from_model_config(config)
+
+    edge_model = ai_edge_torch.convert(
+        pytorch_model,
+        sample_kwargs={
+            "tokens": tokens,
+            "input_pos": input_pos,
+            "kv_cache": kv,
+        },
+    )
     edge_model.set_interpreter_builder(
         self._interpreter_builder(edge_model.tflite_model())
     )
 
     self.assertTrue(
-        model_coverage.compare_tflite_torch(
+        test_utils.compare_tflite_torch(
             edge_model,
             pytorch_model,
-            (tokens, input_pos),
-            num_valid_inputs=1,
+            tokens,
+            input_pos,
+            kv,
+            signature_name="serving_default",
             atol=1e-3,
             rtol=1e-3,
         )

ai_edge_torch/generative/test/utils.py

@@ -0,0 +1,54 @@
+# 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."""
+
+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,
+):
+  """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)
+
+  input_kv_flatten = {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(),
+      **input_kv_flatten,
+  )
+
+  return np.allclose(
+      edge_output["logits"],
+      torch_output["logits"].detach().numpy(),
+      atol=atol,
+      rtol=rtol,
+  )

ai_edge_torch/generative/utilities/loader.py

@@ -221,7 +221,8 @@ class ModelLoader:
       converted_state: Dict[str, torch.Tensor],
   ):
     prefix = f"transformer_blocks.{idx}"
-    if config.ff_config.type == model_config.FeedForwardType.SEQUENTIAL:
+    ff_config = config.block_config(idx).ff_config
+    if 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(
@@ -230,7 +231,7 @@ class ModelLoader:
       converted_state[f"{prefix}.ff.w2.weight"] = state.pop(
           f"{ff_down_proj_name}.weight"
       )
-      if config.ff_config.use_bias:
+      if ff_config.use_bias:
         converted_state[f"{prefix}.ff.w1.bias"] = state.pop(
             f"{ff_up_proj_name}.bias"
         )
@@ -250,7 +251,7 @@ class ModelLoader:
       converted_state[f"{prefix}.ff.w1.weight"] = state.pop(
           f"{ff_gate_proj_name}.weight"
       )
-      if config.ff_config.use_bias:
+      if ff_config.use_bias:
         converted_state[f"{prefix}.ff.w3.bias"] = state.pop(
             f"{ff_up_proj_name}.bias"
         )
@@ -289,6 +290,7 @@ class ModelLoader:
       converted_state: Dict[str, torch.Tensor],
   ):
     prefix = f"transformer_blocks.{idx}"
+    attn_config = config.block_config(idx).attn_config
     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(
@@ -300,13 +302,13 @@ class ModelLoader:
       v_name = self._names.attn_value_proj.format(idx)
       converted_state[f"{prefix}.atten_func.qkv_projection.weight"] = (
           self._fuse_qkv(
-              config,
+              attn_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 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"
@@ -314,7 +316,7 @@ class ModelLoader:
       else:
         converted_state[f"{prefix}.atten_func.qkv_projection.bias"] = (
             self._fuse_qkv(
-                config,
+                attn_config,
                 state.pop(f"{q_name}.bias"),
                 state.pop(f"{k_name}.bias"),
                 state.pop(f"{v_name}.bias"),
@@ -325,7 +327,7 @@ class ModelLoader:
     converted_state[f"{prefix}.atten_func.output_projection.weight"] = (
         state.pop(f"{o_name}.weight")
     )
-    if config.attn_config.output_proj_use_bias:
+    if attn_config.output_proj_use_bias:
       converted_state[f"{prefix}.atten_func.output_projection.bias"] = (
           state.pop(f"{o_name}.bias")
       )
@@ -360,18 +362,16 @@ class ModelLoader:
 
   def _fuse_qkv(
       self,
-      config: model_config.ModelConfig,
+      attn_config: model_config.AttentionConfig,
       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.attn_config.head_dim * q_per_kv)
-      ks = torch.split(k, config.attn_config.head_dim)
-      vs = torch.split(v, config.attn_config.head_dim)
+    if attn_config.qkv_fused_interleaved:
+      q_per_kv = attn_config.num_heads // attn_config.num_query_groups
+      qs = torch.split(q, attn_config.head_dim * q_per_kv)
+      ks = torch.split(k, attn_config.head_dim)
+      vs = torch.split(v, attn_config.head_dim)
       cycled = [t for group in zip(qs, ks, vs) for t in group]
       return torch.cat(cycled)
     else:

ai_edge_torch/generative/utilities/t5_loader.py

@@ -279,7 +279,8 @@ class ModelLoader:
     prefix = additional_prefix + f"transformer_blocks.{idx}"
     if names.ff_up_proj is None or names.ff_down_proj is None:
       return
-    if config.ff_config.type == model_config.FeedForwardType.SEQUENTIAL:
+    ff_config = config.block_config(idx).ff_config
+    if ff_config.type == model_config.FeedForwardType.SEQUENTIAL:
       ff_up_proj_name = names.ff_up_proj.format(idx)
       ff_down_proj_name = names.ff_down_proj.format(idx)
       converted_state[f"{prefix}.ff.w1.weight"] = state.pop(
@@ -288,7 +289,7 @@ class ModelLoader:
       converted_state[f"{prefix}.ff.w2.weight"] = state.pop(
           f"{ff_down_proj_name}.weight"
       )
-      if config.ff_config.use_bias:
+      if ff_config.use_bias:
         converted_state[f"{prefix}.ff.w1.bias"] = state.pop(
             f"{ff_up_proj_name}.bias"
         )
@@ -309,7 +310,7 @@ class ModelLoader:
         converted_state[f"{prefix}.ff.w1.weight"] = state.pop(
             f"{ff_gate_proj_name}.weight"
         )
-        if config.ff_config.use_bias:
+        if ff_config.use_bias:
           converted_state[f"{prefix}.ff.w3.bias"] = state.pop(
               f"{ff_up_proj_name}.bias"
           )
@@ -337,20 +338,21 @@ class ModelLoader:
     ):
       return
     prefix = additional_prefix + f"transformer_blocks.{idx}"
+    attn_config = config.block_config(idx).attn_config
     q_name = names.attn_query_proj.format(idx)
     k_name = names.attn_key_proj.format(idx)
     v_name = names.attn_value_proj.format(idx)
     # model.encoder.transformer_blocks[0].atten_func.q_projection.weight
     if fuse_attention:
       converted_state[f"{prefix}.atten_func.attn.weight"] = self._fuse_qkv(
-          config,
+          attn_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 attn_config.qkv_use_bias:
         converted_state[f"{prefix}.atten_func.attn.bias"] = self._fuse_qkv(
-            config,
+            attn_config,
             state.pop(f"{q_name}.bias"),
             state.pop(f"{k_name}.bias"),
             state.pop(f"{v_name}.bias"),
@@ -365,7 +367,7 @@ class ModelLoader:
       converted_state[f"{prefix}.atten_func.v_projection.weight"] = state.pop(
           f"{v_name}.weight"
       )
-      if config.attn_config.qkv_use_bias:
+      if attn_config.qkv_use_bias:
         converted_state[f"{prefix}.atten_func.q_projection.bias"] = state.pop(
             f"{q_name}.bias"
         )
@@ -380,7 +382,7 @@ class ModelLoader:
     converted_state[f"{prefix}.atten_func.output_projection.weight"] = (
         state.pop(f"{o_name}.weight")
     )
-    if config.attn_config.output_proj_use_bias:
+    if attn_config.output_proj_use_bias:
       converted_state[f"{prefix}.atten_func.output_projection.bias"] = (
           state.pop(f"{o_name}.bias")
       )
@@ -402,6 +404,7 @@ class ModelLoader:
     ):
       return
     prefix = additional_prefix + f"transformer_blocks.{idx}"
+    attn_config = config.block_config(idx).attn_config
     q_name = names.cross_attn_query_proj.format(idx)
     k_name = names.cross_attn_key_proj.format(idx)
     v_name = names.cross_attn_value_proj.format(idx)
@@ -409,16 +412,16 @@ class ModelLoader:
     if fuse_attention:
       converted_state[f"{prefix}.cross_atten_func.attn.weight"] = (
           self._fuse_qkv(
-              config,
+              attn_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 attn_config.qkv_use_bias:
         converted_state[f"{prefix}.cross_atten_func.attn.bias"] = (
             self._fuse_qkv(
-                config,
+                attn_config,
                 state.pop(f"{q_name}.bias"),
                 state.pop(f"{k_name}.bias"),
                 state.pop(f"{v_name}.bias"),
@@ -434,7 +437,7 @@ class ModelLoader:
       converted_state[f"{prefix}.cross_atten_func.v_projection.weight"] = (
           state.pop(f"{v_name}.weight")
       )
-      if config.attn_config.qkv_use_bias:
+      if attn_config.qkv_use_bias:
         converted_state[f"{prefix}.cross_atten_func.q_projection.bias"] = (
             state.pop(f"{q_name}.bias")
         )
@@ -449,7 +452,7 @@ class ModelLoader:
     converted_state[f"{prefix}.cross_atten_func.output_projection.weight"] = (
         state.pop(f"{o_name}.weight")
     )
-    if config.attn_config.output_proj_use_bias:
+    if attn_config.output_proj_use_bias:
       converted_state[f"{prefix}.cross_atten_func.output_projection.bias"] = (
           state.pop(f"{o_name}.bias")
       )
@@ -496,16 +499,14 @@ class ModelLoader:
 
   def _fuse_qkv(
       self,
-      config: model_config.ModelConfig,
+      attn_config: model_config.AttentionConfig,
       q: torch.Tensor,
       k: torch.Tensor,
       v: torch.Tensor,
   ) -> torch.Tensor:
-    q_per_kv = (
-        config.attn_config.num_heads // config.attn_config.num_query_groups
-    )
-    qs = torch.split(q, config.attn_config.head_dim * q_per_kv)
-    ks = torch.split(k, config.attn_config.head_dim)
-    vs = torch.split(v, config.attn_config.head_dim)
+    q_per_kv = attn_config.num_heads // attn_config.num_query_groups
+    qs = torch.split(q, attn_config.head_dim * q_per_kv)
+    ks = torch.split(k, attn_config.head_dim)
+    vs = torch.split(v, attn_config.head_dim)
     cycled = [t for group in zip(qs, ks, vs) for t in group]
     return torch.cat(cycled)

ai_edge_torch/odml_torch/lowerings/__init__.py

@@ -16,6 +16,7 @@ from . import _basic
 from . import _batch_norm
 from . import _convolution
 from . import _jax_lowerings
+from . import _layer_norm
 from . import context
 from . import registry
 from . import utils