ai-edge-torch-nightly 0.3.0.dev20240926__py3-none-any.whl → 0.3.0.dev20240929__py3-none-any.whl

ai_edge_torch/generative/examples/gemma/gemma1.py

@@ -69,15 +69,10 @@ class Gemma(nn.Module):
     self.rope_cache = attn_utils.build_rope_cache(
         size=config.kv_cache_max,
         dim=int(attn_config.rotary_percentage * attn_config.head_dim),
-        base=10_000,
-        condense_ratio=1,
-        dtype=torch.float32,
-        device=torch.device("cpu"),
+        base=attn_config.rotary_base,
     )
     self.mask_cache = attn_utils.build_causal_mask_cache(
         size=config.kv_cache_max,
-        dtype=torch.float32,
-        device=torch.device("cpu"),
     )
     self.config = config
 
@@ -135,6 +130,7 @@ def get_model_config_2b(kv_cache_max_len: int = 1024) -> cfg.ModelConfig:
       num_heads=8,
       head_dim=256,
       num_query_groups=1,
+      rotary_base=10000,
       rotary_percentage=1.0,
   )
   ff_config = cfg.FeedForwardConfig(

ai_edge_torch/generative/examples/gemma/gemma2.py

@@ -109,21 +109,14 @@ class Gemma2(nn.Module):
     self.rope_cache = attn_utils.build_rope_cache(
         size=config.kv_cache_max,
         dim=int(attn_config.rotary_percentage * attn_config.head_dim),
-        base=10_000,
-        condense_ratio=1,
-        dtype=torch.float32,
-        device=torch.device("cpu"),
+        base=attn_config.rotary_base,
     )
     self.mask_cache = attn_utils.build_causal_mask_cache(
         size=config.kv_cache_max,
-        dtype=torch.float32,
-        device=torch.device("cpu"),
     )
     self.sliding_window_mask_cache = attn_utils.build_sliding_window_mask_cache(
         size=config.kv_cache_max,
         window_size=attn_config.sliding_window_size,
-        dtype=torch.float32,
-        device=torch.device("cpu"),
     )
     self.config = config
 
@@ -208,6 +201,7 @@ def get_model_config_2b(kv_cache_max_len: int = 1024) -> cfg.ModelConfig:
         num_heads=8,
         head_dim=256,
         num_query_groups=4,
+        rotary_base=10000,
         rotary_percentage=1.0,
         qkv_transpose_before_split=True,
         logit_softcap=50.0,

ai_edge_torch/generative/examples/gemma/verify_gemma2.py

@@ -20,7 +20,6 @@ from absl import app
 from absl import flags
 from ai_edge_torch.generative.examples.gemma import gemma2
 from ai_edge_torch.generative.examples.gemma import verify_util
-from ai_edge_torch.generative.utilities import verifier
 import kagglehub
 
 

ai_edge_torch/generative/examples/gemma/verify_util.py

@@ -15,7 +15,6 @@
 
 """Utility functions to verify the reauthored Gemma model."""
 
-import dataclasses
 import logging
 import os
 from typing import List, Tuple
@@ -27,26 +26,17 @@ from gemma import model as gemma_model
 import torch
 
 
-@dataclasses.dataclass
-class _Output:
-  logits: torch.Tensor
-
-
 class GemmaWrapper(verifier.ModelWrapper):
   """Gemma model wrapper for verification.
 
   Verifier calls model.forward() with maxium sequence length (1024) expecting
-  the output has 'logits' field while Gemma gets the input tokens with the
-  actual length and returns logits in a tuple.
+  the output is logits while Gemma gets the input tokens with the actual length
+  and returns logits in a tuple.
 
   Verifier runs tokenizer before model.generate() while Gemma runs the tokenizer
   inside model.generate().
   """
 
-  def __init__(self, model: torch.nn.Module, max_new_tokens: int):
-    super().__init__(model)
-    self.max_new_tokens = max_new_tokens
-
   def _get_actual_input_len(self, tokens: torch.Tensor) -> int:
     for i in range(tokens.shape[1]):
       if tokens[0, i] == 0:
@@ -63,7 +53,7 @@ class GemmaWrapper(verifier.ModelWrapper):
         (cache.clone(), cache.clone()) for _ in range(config.num_hidden_layers)
     ]
 
-  def forward(self, tokens: torch.Tensor) -> _Output:
+  def forward(self, tokens: torch.Tensor) -> torch.Tensor:
     """Forwards the model after reducing input tokens to the actual length."""
     actual_input_len = self._get_actual_input_len(tokens)
     input_pos = torch.arange(0, actual_input_len, dtype=torch.long)
@@ -79,28 +69,26 @@ class GemmaWrapper(verifier.ModelWrapper):
         top_ps=torch.tensor([1.0], dtype=torch.float),
         top_ks=torch.tensor([1], dtype=torch.long),
     )
-    return _Output(logits.float())
+    return logits
 
-  def generate(self, tokens: torch.Tensor) -> torch.Tensor:
+  def generate(
+      self, tokens: torch.Tensor, max_new_tokens: int
+  ) -> torch.IntTensor:
     """Generates the response after decoding the tokens into a string."""
     prompts = self.model.tokenizer.decode(tokens[0].tolist())
     response = self.model.generate(
-        prompts, device="cpu", output_len=self.max_new_tokens, top_k=1
+        prompts, device="cpu", output_len=max_new_tokens, top_k=1
     )
     return torch.tensor([self.model.tokenizer.encode(prompts + response)])
 
 
-class TokenizerWrapper(torch.nn.Module):
+class GemmaTokenizerWrapper(verifier.TokenizerWrapper):
   """Tokenizer wrapper for verification.
 
   Verifier expects the tokenizer to handle tokens in torch.Tensor while Gemma
   tokenizer expects tokens in a list.
   """
 
-  def __init__(self, tokenizer: torch.nn.Module):
-    super().__init__()
-    self.tokenizer = tokenizer
-
   def encode(self, text: str, **_) -> torch.Tensor:
     """Adds one more dimension to the output of the tokenizer."""
     return torch.tensor([self.tokenizer.encode(text)])
@@ -133,10 +121,11 @@ def verify_reauthored_gemma_model(
   original_model.load_weights(os.path.join(checkpoint, weight_filename))
 
   verifier.verify_reauthored_model(
-      original_model=GemmaWrapper(original_model, max_new_tokens),
-      reauthored_model=reauthored_model,
-      tokenizer=TokenizerWrapper(original_model.tokenizer),
+      original_model=GemmaWrapper(original_model),
+      reauthored_model=verifier.ReauthoredModelWrapper(reauthored_model),
+      tokenizer=GemmaTokenizerWrapper(original_model.tokenizer),
       generate_prompts=generate_prompts,
+      max_new_tokens=max_new_tokens,
       forward_input_ids=forward_input_ids,
       rtol=rtol,
       atol=atol,

ai_edge_torch/generative/examples/llama/__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/examples/llama/convert_3b_to_tflite.py

@@ -0,0 +1,68 @@
+# 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.
+# ==============================================================================
+
+"""Example of converting Llama 3.2 3B model to multi-signature tflite model."""
+
+import os
+import pathlib
+
+from absl import app
+from absl import flags
+from ai_edge_torch.generative.examples.llama import llama
+from ai_edge_torch.generative.utilities import converter
+
+_CHECKPOINT_PATH = flags.DEFINE_string(
+    'checkpoint_path',
+    os.path.join(pathlib.Path.home(), 'Downloads/llm_data/llama'),
+    'The path to the model checkpoint, or directory holding the checkpoint.',
+)
+_TFLITE_PATH = flags.DEFINE_string(
+    'tflite_path',
+    '/tmp/',
+    'The tflite file path to export.',
+)
+_PREFILL_SEQ_LEN = flags.DEFINE_integer(
+    'prefill_seq_len',
+    1024,
+    'The maximum size of prefill input tensor.',
+)
+_KV_CACHE_MAX_LEN = flags.DEFINE_integer(
+    'kv_cache_max_len',
+    1280,
+    'The maximum size of KV cache buffer, including both prefill and decode.',
+)
+_QUANTIZE = flags.DEFINE_bool(
+    'quantize',
+    True,
+    'Whether the model should be quantized.',
+)
+
+
+def main(_):
+  pytorch_model = llama.build_3b_model(
+      _CHECKPOINT_PATH.value, kv_cache_max_len=_KV_CACHE_MAX_LEN.value
+  )
+  quant_suffix = 'q8' if _QUANTIZE.value else 'f32'
+  output_filename = f'llama_3b_{quant_suffix}_seq{_PREFILL_SEQ_LEN.value}_ekv{_KV_CACHE_MAX_LEN.value}.tflite'
+  converter.convert_to_tflite(
+      pytorch_model,
+      tflite_path=os.path.join(_TFLITE_PATH.value, output_filename),
+      prefill_seq_len=_PREFILL_SEQ_LEN.value,
+      quantize=_QUANTIZE.value,
+  )
+
+
+if __name__ == '__main__':
+  app.run(main)

ai_edge_torch/generative/examples/llama/convert_to_tflite.py

@@ -0,0 +1,68 @@
+# 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.
+# ==============================================================================
+
+"""Example of converting Llama 3.2 1B model to multi-signature tflite model."""
+
+import os
+import pathlib
+
+from absl import app
+from absl import flags
+from ai_edge_torch.generative.examples.llama import llama
+from ai_edge_torch.generative.utilities import converter
+
+_CHECKPOINT_PATH = flags.DEFINE_string(
+    'checkpoint_path',
+    os.path.join(pathlib.Path.home(), 'Downloads/llm_data/llama'),
+    'The path to the model checkpoint, or directory holding the checkpoint.',
+)
+_TFLITE_PATH = flags.DEFINE_string(
+    'tflite_path',
+    '/tmp/',
+    'The tflite file path to export.',
+)
+_PREFILL_SEQ_LEN = flags.DEFINE_integer(
+    'prefill_seq_len',
+    1024,
+    'The maximum size of prefill input tensor.',
+)
+_KV_CACHE_MAX_LEN = flags.DEFINE_integer(
+    'kv_cache_max_len',
+    1280,
+    'The maximum size of KV cache buffer, including both prefill and decode.',
+)
+_QUANTIZE = flags.DEFINE_bool(
+    'quantize',
+    True,
+    'Whether the model should be quantized.',
+)
+
+
+def main(_):
+  pytorch_model = llama.build_model(
+      _CHECKPOINT_PATH.value, kv_cache_max_len=_KV_CACHE_MAX_LEN.value
+  )
+  quant_suffix = 'q8' if _QUANTIZE.value else 'f32'
+  output_filename = f'llama_{quant_suffix}_seq{_PREFILL_SEQ_LEN.value}_ekv{_KV_CACHE_MAX_LEN.value}.tflite'
+  converter.convert_to_tflite(
+      pytorch_model,
+      tflite_path=os.path.join(_TFLITE_PATH.value, output_filename),
+      prefill_seq_len=_PREFILL_SEQ_LEN.value,
+      quantize=_QUANTIZE.value,
+  )
+
+
+if __name__ == '__main__':
+  app.run(main)

ai_edge_torch/generative/examples/llama/llama.py

@@ -0,0 +1,204 @@
+# 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.
+# ==============================================================================
+
+"""Example of building Llama 3.2 models."""
+
+import copy
+import math
+from typing import Tuple
+
+from ai_edge_torch.generative.examples.tiny_llama import tiny_llama
+import ai_edge_torch.generative.layers.model_config as cfg
+import ai_edge_torch.generative.utilities.loader as loading_utils
+import torch
+from torch import nn
+
+TENSOR_NAMES = copy.copy(tiny_llama.TENSOR_NAMES)
+# SmolLM re-uses the embedding as the head projection layer.
+TENSOR_NAMES.lm_head = None
+
+
+def _build_llama3_rope_cache(
+    size: int,
+    dim: int,
+    base: int,
+    condense_ratio: int,
+    dtype: torch.dtype,
+    device: torch.device,
+    factor: float,
+    low_freq_factor: float,
+    high_freq_factor: float,
+    max_seq_len: int,
+) -> Tuple[torch.Tensor, torch.Tensor]:
+  """Precomputes Rotary Positional Embeddings for Llama 3.2 model.
+
+  It's a modified version of attn_utils.build_rope_cache with additional
+  arguments for Llama 3.2 model. It precomputes Rotary Positional Embedding Sin
+  and Cos values with scaling factors for quick lookup during the inference.
+
+  Reference:
+  https://github.com/huggingface/transformers/blob/main/src/transformers/modeling_rope_utils.py#L307
+
+  Args:
+      size (int): The size of the built cache.
+      dim (int): Each sequence's dimmension.
+      base (int, optional): Rope base value.
+      condense_ratio (int, optional): The ratio by which sequence indicies are
+        condensed.
+      dtype (torch.dtype, optional): Output tensor's data type.
+      device (torch.device, optional): Output tensor's data type.
+      factor (float): Factor to scale theta down for tokens in long range in the
+        sequence.
+      low_freq_factor (float): Factor to determine if tokens are in long range
+        in the sequence.
+      high_freq_factor (float): Factor to determine if tokens are in short range
+        in the sequence.
+      max_seq_len (int): The original token sequence length before extending
+        ROPE to support longer sequence.
+
+  Returns:
+      Tuple[torch.Tensor, torch.Tensor]: Rope's Cosine and Sine waves.
+  """
+  theta = 1.0 / (base ** (torch.arange(0, dim, 2).float() / dim))
+  low_freq_wavelen = max_seq_len / low_freq_factor
+  high_freq_wavelen = max_seq_len / high_freq_factor
+  wavelen = 2 * math.pi / theta
+  # wavelen < high_freq_wavelen: do nothing
+  # wavelen > low_freq_wavelen: divide by factor
+  theta = torch.where(wavelen > low_freq_wavelen, theta / factor, theta)
+  # otherwise: interpolate between the two, using a smooth factor
+  smooth_factor = (max_seq_len / wavelen - low_freq_factor) / (
+      high_freq_factor - low_freq_factor
+  )
+  smoothed_theta = (1 - smooth_factor) * theta / factor + smooth_factor * theta
+  is_medium = ~(wavelen < high_freq_wavelen) * ~(wavelen > low_freq_wavelen)
+  theta = torch.where(is_medium, smoothed_theta, theta)
+
+  seq_idx = torch.arange(size) / condense_ratio
+  idx_theta = torch.outer(seq_idx, theta)
+  cos = torch.cos(idx_theta).to(dtype=dtype, device=device)
+  sin = torch.sin(idx_theta).to(dtype=dtype, device=device)
+  return cos, sin
+
+
+class Llama(tiny_llama.TinyLlama):
+  """A Llama model built from the Edge Generative API layers.
+
+  Llama 3.2 shares the same architecture as TinyLlama except ROPE calculation.
+  """
+
+  def __init__(self, config: cfg.ModelConfig):
+    super().__init__(config)
+    # Llama 3.2 re-uses the embedding as the head projection layer.
+    self.lm_head.weight.data = self.tok_embedding.weight.data
+    # Llama has only one block config.
+    attn_config = self.config.block_config(0).attn_config
+    self.rope_cache = _build_llama3_rope_cache(
+        size=self.config.kv_cache_max,
+        dim=int(attn_config.rotary_percentage * attn_config.head_dim),
+        base=attn_config.rotary_base,
+        condense_ratio=1,
+        dtype=torch.float32,
+        device=torch.device("cpu"),
+        factor=32.0,
+        low_freq_factor=1.0,
+        high_freq_factor=4.0,
+        max_seq_len=self.config.max_seq_len,
+    )
+
+
+def get_model_config(kv_cache_max_len: int = 1024) -> cfg.ModelConfig:
+  """Returns the model config for a Llama 3.2-1B model.
+
+  Args:
+    kv_cache_max_len (int): The maximum sequence length of the KV cache. Default
+      is 1024.
+
+  Returns:
+    The model config for a SmolLM model.
+  """
+  attn_config = cfg.AttentionConfig(
+      num_heads=32,
+      head_dim=64,
+      num_query_groups=8,
+      rotary_base=500000,
+      rotary_percentage=1.0,
+  )
+  ff_config = cfg.FeedForwardConfig(
+      type=cfg.FeedForwardType.GATED,
+      activation=cfg.ActivationConfig(cfg.ActivationType.SILU),
+      intermediate_size=8192,
+  )
+  norm_config = cfg.NormalizationConfig(type=cfg.NormalizationType.RMS_NORM)
+  block_config = cfg.TransformerBlockConfig(
+      attn_config=attn_config,
+      ff_config=ff_config,
+      pre_attention_norm_config=norm_config,
+      post_attention_norm_config=norm_config,
+  )
+  config = cfg.ModelConfig(
+      vocab_size=128256,
+      num_layers=16,
+      max_seq_len=8192,
+      embedding_dim=2048,
+      kv_cache_max_len=kv_cache_max_len,
+      block_configs=block_config,
+      final_norm_config=norm_config,
+      enable_hlfb=True,
+  )
+  return config
+
+
+def get_3b_model_config(kv_cache_max_len: int = 1024) -> cfg.ModelConfig:
+  """Returns the model config for a Llama 3.2-3B model."""
+  config = get_model_config(kv_cache_max_len)
+  # Llama 3.2 has only one block config.
+  attn_config = config.block_config(0).attn_config
+  attn_config.num_heads = 24
+  attn_config.head_dim = 128
+  config.num_layers = 28
+  config.embedding_dim = 3072
+  return config
+
+
+def get_fake_model_config(**kwargs) -> cfg.ModelConfig:
+  config = get_model_config(**kwargs)
+  config.vocab_size = 128
+  config.num_layers = 2
+  # SmolLM has only one block config.
+  config.block_config(0).ff_config.intermediate_size = 64
+  return config
+
+
+def build_model(checkpoint_path: str, **kwargs) -> nn.Module:
+  config = get_model_config(**kwargs)
+  model = Llama(config)
+  loader = loading_utils.ModelLoader(checkpoint_path, TENSOR_NAMES)
+  # Since embedding and lm-head use the same weight, we need to set strict
+  # to False.
+  loader.load(model, strict=False)
+  model.eval()
+  return model
+
+
+def build_3b_model(checkpoint_path: str, **kwargs) -> nn.Module:
+  config = get_3b_model_config(**kwargs)
+  model = Llama(config)
+  loader = loading_utils.ModelLoader(checkpoint_path, TENSOR_NAMES)
+  # Since embedding and lm-head use the same weight, we need to set strict
+  # to False.
+  loader.load(model, strict=False)
+  model.eval()
+  return model

ai_edge_torch/generative/examples/llama/verify.py

@@ -0,0 +1,73 @@
+# 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.
+# ==============================================================================
+
+"""Verifies the reauthored Llama 3.2-1B model."""
+
+import logging
+import pathlib
+
+from absl import app
+from absl import flags
+from ai_edge_torch.generative.examples.llama import llama
+from ai_edge_torch.generative.utilities import transformers_verifier
+from ai_edge_torch.generative.utilities import verifier
+import transformers
+
+
+_PROMPTS = flags.DEFINE_multi_string(
+    "prompts",
+    "What is the meaning of life?",
+    "The input prompts to generate answers.",
+)
+_MAX_NEW_TOKENS = flags.DEFINE_integer(
+    "max_new_tokens",
+    30,
+    "The maximum size of the generated tokens.",
+)
+
+
+def main(_):
+  checkpoint = "meta-llama/Llama-3.2-1B-Instruct"
+  logging.info("Loading the original model from: %s", checkpoint)
+  original_model = transformers.AutoModelForCausalLM.from_pretrained(checkpoint)
+
+  # Locate the cached dir.
+  cached_config_file = transformers.utils.cached_file(
+      checkpoint, transformers.utils.CONFIG_NAME
+  )
+  reauthored_checkpoint = pathlib.Path(cached_config_file).parent
+  logging.info("Building the reauthored model from: %s", reauthored_checkpoint)
+  reauthored_model = llama.build_model(reauthored_checkpoint)
+
+  logging.info("Loading the tokenizer from: %s", checkpoint)
+  # Llama tokenizer_config.json sets a fast tokenizer class explicitly,
+  # "PreTrainedTokenizerFast". It works only when the fast tokenizer is
+  # available.
+  tokenizer = transformers.AutoTokenizer.from_pretrained(checkpoint)
+
+  verifier.verify_reauthored_model(
+      original_model=transformers_verifier.TransformersModelWrapper(
+          original_model
+      ),
+      reauthored_model=verifier.ReauthoredModelWrapper(reauthored_model),
+      tokenizer=verifier.TokenizerWrapper(tokenizer),
+      generate_prompts=_PROMPTS.value,
+      max_new_tokens=_MAX_NEW_TOKENS.value,
+      atol=1e-04,
+  )
+
+
+if __name__ == "__main__":
+  app.run(main)

ai_edge_torch/generative/examples/llama/verify_3b.py

@@ -0,0 +1,73 @@
+# 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.
+# ==============================================================================
+
+"""Verifies the reauthored Llama 3.2-3B model."""
+
+import logging
+import pathlib
+
+from absl import app
+from absl import flags
+from ai_edge_torch.generative.examples.llama import llama
+from ai_edge_torch.generative.utilities import transformers_verifier
+from ai_edge_torch.generative.utilities import verifier
+import transformers
+
+
+_PROMPTS = flags.DEFINE_multi_string(
+    "prompts",
+    "What is the meaning of life?",
+    "The input prompts to generate answers.",
+)
+_MAX_NEW_TOKENS = flags.DEFINE_integer(
+    "max_new_tokens",
+    30,
+    "The maximum size of the generated tokens.",
+)
+
+
+def main(_):
+  checkpoint = "meta-llama/Llama-3.2-3B-Instruct"
+  logging.info("Loading the original model from: %s", checkpoint)
+  original_model = transformers.AutoModelForCausalLM.from_pretrained(checkpoint)
+
+  # Locate the cached dir.
+  cached_config_file = transformers.utils.cached_file(
+      checkpoint, transformers.utils.CONFIG_NAME
+  )
+  reauthored_checkpoint = pathlib.Path(cached_config_file).parent
+  logging.info("Building the reauthored model from: %s", reauthored_checkpoint)
+  reauthored_model = llama.build_3b_model(reauthored_checkpoint)
+
+  logging.info("Loading the tokenizer from: %s", checkpoint)
+  # Llama tokenizer_config.json sets a fast tokenizer class explicitly,
+  # "PreTrainedTokenizerFast". It works only when the fast tokenizer is
+  # available.
+  tokenizer = transformers.AutoTokenizer.from_pretrained(checkpoint)
+
+  verifier.verify_reauthored_model(
+      original_model=transformers_verifier.TransformersModelWrapper(
+          original_model
+      ),
+      reauthored_model=verifier.ReauthoredModelWrapper(reauthored_model),
+      tokenizer=verifier.TokenizerWrapper(tokenizer),
+      generate_prompts=_PROMPTS.value,
+      max_new_tokens=_MAX_NEW_TOKENS.value,
+      atol=1e-04,
+  )
+
+
+if __name__ == "__main__":
+  app.run(main)

ai_edge_torch/generative/examples/openelm/openelm.py

@@ -68,15 +68,10 @@ class OpenELM(nn.Module):
     self.rope_cache = attn_utils.build_rope_cache(
         size=config.kv_cache_max,
         dim=int(attn_config.rotary_percentage * attn_config.head_dim),
-        base=10_000,
-        condense_ratio=1,
-        dtype=torch.float32,
-        device=torch.device("cpu"),
+        base=attn_config.rotary_base,
     )
     self.mask_cache = attn_utils.build_causal_mask_cache(
         size=config.kv_cache_max,
-        dtype=torch.float32,
-        device=torch.device("cpu"),
     )
     self.config = config
 
@@ -154,6 +149,7 @@ def get_model_config(kv_cache_max_len: int = 1024) -> cfg.ModelConfig:
             num_heads=num_heads[idx],
             head_dim=128,
             num_query_groups=num_query_groups[idx],
+            rotary_base=10000,
             rotary_percentage=1.0,
             qkv_transpose_before_split=True,
             query_norm_config=norm_config,