ai-edge-torch-nightly 0.5.0.dev20250514__py3-none-any.whl → 0.5.0.dev20250515__py3-none-any.whl

ai_edge_torch/generative/examples/gemma3/convert_gemma3_to_tflite.py

@@ -19,11 +19,19 @@ from absl import app
 from ai_edge_torch.generative.examples.gemma3 import gemma3
 from ai_edge_torch.generative.utilities import converter
 from ai_edge_torch.generative.utilities import export_config
+from ai_edge_torch.generative.utilities import loader
 
 flags = converter.define_conversion_flags(
     'gemma3-1b', default_mask_as_input=True, default_transpose_kv_cache=True
 )
 
+_CUSTOM_CHECKPOINT_LOADER = flags.DEFINE_bool(
+    'custom_checkpoint_loader',
+    False,
+    'If true, the conversion script will use a custom checkpoint loader which'
+    ' will read a checkpoint from a remote source.',
+)
+
 _MODEL_SIZE = flags.DEFINE_string(
     'model_size',
     '1b',
@@ -32,10 +40,16 @@ _MODEL_SIZE = flags.DEFINE_string(
 
 
 def main(_):
+  custom_loader = None
+  if flags.FLAGS.custom_checkpoint_loader:
+    # If loading from a remote source, try to get a custom loader first.
+    custom_loader = loader.get_custom_loader(flags.FLAGS.checkpoint_path)
+
   if _MODEL_SIZE.value == '1b':
     pytorch_model = gemma3.build_model_1b(
         flags.FLAGS.checkpoint_path,
         kv_cache_max_len=flags.FLAGS.kv_cache_max_len,
+        custom_loader=custom_loader,
     )
   else:
     raise ValueError(f'Unsupported model size: {_MODEL_SIZE.value}')

ai_edge_torch/generative/layers/normalization.py

@@ -28,6 +28,8 @@ class RMSNorm(torch.nn.Module):
       dim: int,
       eps: float = 1e-6,
       zero_centered_gamma=False,
+      with_scale: bool = False,
+      scale_shift: float = 1.0,
       enable_hlfb: bool = False,
   ):
     """Initialize the RMSNorm layer.
@@ -37,13 +39,22 @@ class RMSNorm(torch.nn.Module):
       eps (float): A small float value to ensure numerical stability (default:
         1e-6).
       zero_centered_gamma (bool): Whether or not gamma has an offset.
+      with_scale (bool): Whether or not to use a scale parameter.
+      scale_shift (float): The shift to apply to the scale parameter.
       enable_hlfb (bool): use HLFB in the op.
     """
     super().__init__()
+    self.dim = dim
     self.enable_hlfb = enable_hlfb
     self.eps = eps
-    self.weight = torch.nn.Parameter(torch.ones(dim))
+    self.weight = torch.nn.Parameter(torch.ones(dim), requires_grad=False)
     self.zero_centered_gamma = zero_centered_gamma
+    self.with_scale = with_scale
+    if with_scale:
+      self.scale = torch.nn.Parameter(
+          torch.zeros((dim,), dtype=torch.float32), requires_grad=False
+      )
+    self.scale_shift = scale_shift
 
   def _norm(self, x):
     """Apply RMSNorm normalization.
@@ -70,14 +81,20 @@ class RMSNorm(torch.nn.Module):
     else:
       w = self.weight
 
+    final_scale = (
+        self.scale + self.scale_shift
+        if self.with_scale
+        else torch.ones((self.dim,), dtype=torch.float32)
+    )
     if self.enable_hlfb:
       return rms_norm_with_hlfb(
           x,
           w,
           self.eps,
+          final_scale,
       )
     else:
-      output = self._norm(x.float()).type_as(x)
+      output = self._norm(x.float()).type_as(x) * final_scale
       return output * w
 
 
@@ -104,8 +121,8 @@ class GroupNorm(torch.nn.Module):
     self.enable_hlfb = enable_hlfb
     self.group_num = group_num
     self.eps = eps
-    self.weight = torch.nn.Parameter(torch.empty(dim))
-    self.bias = torch.nn.Parameter(torch.empty(dim))
+    self.weight = torch.nn.Parameter(torch.empty(dim), requires_grad=False)
+    self.bias = torch.nn.Parameter(torch.empty(dim), requires_grad=False)
 
   def forward(self, x):
     """Running the forward pass of GroupNorm layer.
@@ -140,8 +157,8 @@ class LayerNorm(torch.nn.Module):
     self.enable_hlfb = enable_hlfb
     self.normalized_shape = (dim,)
     self.eps = eps
-    self.weight = torch.nn.Parameter(torch.empty(dim))
-    self.bias = torch.nn.Parameter(torch.empty(dim))
+    self.weight = torch.nn.Parameter(torch.empty(dim), requires_grad=False)
+    self.bias = torch.nn.Parameter(torch.empty(dim), requires_grad=False)
 
   def forward(self, x):
     """Running the forward pass of LayerNorm layer.
@@ -165,6 +182,7 @@ def rms_norm_with_hlfb(
     x: torch.Tensor,
     w: torch.Tensor,
     eps: float,
+    final_scale: torch.Tensor,
 ):
   """RMS Normalization with high-level function boundary enabled.
 
@@ -172,6 +190,7 @@ def rms_norm_with_hlfb(
     x (torch.Tensor): Input tensor for RMS Normalization, with BCHW shape.
     w (torch.Tensor): The learned parameter tensor for normalization.
     eps (float): A small float value to ensure numerical stability.
+    final_scale (torch.Tensor): The final scale to apply to the normalization.
 
   Returns:
     The output tensor of RMS Normalization.
@@ -185,7 +204,7 @@ def rms_norm_with_hlfb(
   def _norm(x):
     return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + eps)
 
-  output = _norm(x.float()).type_as(x)
+  output = _norm(x.float()).type_as(x) * final_scale
   out = output * w
 
   out = builder.mark_outputs(out)

ai_edge_torch/generative/layers/normalization_test.py

@@ -0,0 +1,73 @@
+# 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.
+# ==============================================================================
+"""Tests for normalization layers."""
+
+from ai_edge_torch.generative.layers import normalization
+import torch
+from absl.testing import absltest as googletest
+from absl.testing import parameterized
+
+
+class NormalizationTest(parameterized.TestCase):
+
+  @parameterized.named_parameters(
+      dict(
+          testcase_name="rms_norm_test_1",
+          model_dim=10,
+          with_scale=False,
+          scale_shift=1.0,
+          enable_hlfb=False,
+          expected_values=torch.ones((10,), dtype=torch.float32),
+      ),
+      dict(
+          testcase_name="rms_norm_test_2",
+          model_dim=10,
+          with_scale=True,
+          scale_shift=2.0,
+          enable_hlfb=False,
+          expected_values=torch.ones((10,), dtype=torch.float32) * 2.0,
+      ),
+      dict(
+          testcase_name="rms_norm_test_3",
+          model_dim=10,
+          with_scale=True,
+          scale_shift=2.0,
+          enable_hlfb=True,
+          expected_values=torch.ones((10,), dtype=torch.float32) * 2.0,
+      ),
+  )
+  def test_rms_norm(
+      self,
+      model_dim: int,
+      with_scale: bool,
+      scale_shift: float,
+      enable_hlfb: bool,
+      expected_values: torch.Tensor,
+  ):
+    rms_norm = normalization.RMSNorm(
+        dim=model_dim,
+        with_scale=with_scale,
+        scale_shift=scale_shift,
+        enable_hlfb=enable_hlfb,
+    )
+
+    x = torch.ones((model_dim,), dtype=torch.float32)
+    out = rms_norm(x)
+    self.assertEqual(out.shape, (model_dim,))
+    self.assertTrue(torch.allclose(out, expected_values))
+
+
+if __name__ == "__main__":
+  googletest.main()

ai_edge_torch/generative/utilities/loader.py

@@ -19,10 +19,36 @@ import os
 from typing import Callable, Dict, List, Tuple
 
 from ai_edge_torch.generative.layers import model_config
+import safetensors
 from safetensors import safe_open
 import torch
 
 
+def get_custom_loader(
+    checkpoint_path: str,
+) -> Callable[[str], Dict[str, torch.Tensor]]:
+  """Returns a custom loader for the given checkpoint path.
+
+  Those customer loaders can either support state dictionary or safetensors, and
+  the actual data might be fetched from a remote source.
+
+  Args:
+    checkpoint_path (string): The path to the checkpoint.
+
+  Returns:
+    Callable[[str], Dict[str, torch.Tensor]]: The custom loader.
+
+  Raises:
+    ValueError: If the checkpoint format is not supported.
+  """
+
+  if os.path.splitext(checkpoint_path)[1] in [".bin", ".pt", ".ckpt"]:
+    return lambda path: torch.load(path, weights_only=True)
+  if checkpoint_path.endswith(".safetensors"):
+    return safetensors.torch.load_file
+  raise ValueError(f"Unsupported checkpoint format: {checkpoint_path}")
+
+
 def load_safetensors(full_path: str):
   """Loads safetensors into a single state dictionary.
 

ai_edge_torch/version.py

@@ -13,4 +13,4 @@
 # limitations under the License.
 # ==============================================================================
 
-__version__ = "0.5.0.dev20250514"
+__version__ = "0.5.0.dev20250515"

{ai_edge_torch_nightly-0.5.0.dev20250514.dist-info → ai_edge_torch_nightly-0.5.0.dev20250515.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: ai-edge-torch-nightly
-Version: 0.5.0.dev20250514
+Version: 0.5.0.dev20250515
 Summary: Supporting PyTorch models with the Google AI Edge TFLite runtime.
 Home-page: https://github.com/google-ai-edge/ai-edge-torch
 Keywords: On-Device ML,AI,Google,TFLite,PyTorch,LLMs,GenAI

{ai_edge_torch_nightly-0.5.0.dev20250514.dist-info → ai_edge_torch_nightly-0.5.0.dev20250515.dist-info}/RECORD

@@ -2,7 +2,7 @@ ai_edge_torch/__init__.py,sha256=8sPR_5uXJA4NEE0nIwNdSl-ADOJEoR8hAgYvBQDY70Y,120
 ai_edge_torch/_config.py,sha256=AiqhbcheF7j_ozIGDLC89k1we95aVgFDa-tR6h7UI0s,2529
 ai_edge_torch/conftest.py,sha256=r0GTrhMRhlmOGrrkvumHN8hkmyug6WvF60vWq8wRIBI,758
 ai_edge_torch/model.py,sha256=wxjSFq_rBSxSqbUE8E8EJTCkgvgaRLjq_ZuAM-IZpCU,5606
-ai_edge_torch/version.py,sha256=ZvSDZpKkUslpMEN4pPp4xI6n8g3mHZMdfIcYeWth5Dg,706
+ai_edge_torch/version.py,sha256=QVmEdwoLJem1gNQul_CoRyfqOc1Ljjy48x9GmKmuAOU,706
 ai_edge_torch/_convert/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
 ai_edge_torch/_convert/conversion.py,sha256=QVugYVfbyaeBgSKKbhFzHG5oXA7t3M-40JcpcdSu6W8,5436
 ai_edge_torch/_convert/conversion_utils.py,sha256=Sr8qXVcTwc-ZnZmK7yxVrIOOp1S_vNrwzC0zUvLTI2o,2160
@@ -68,7 +68,7 @@ ai_edge_torch/generative/examples/gemma/verify_gemma1.py,sha256=ip-Gmk4CI5f0GWSd
 ai_edge_torch/generative/examples/gemma/verify_gemma2.py,sha256=jhiyinOqPt5ZZjEadDRZt_wY5fiLSCpMo54PcxFaL_Q,1789
 ai_edge_torch/generative/examples/gemma/verify_util.py,sha256=n7f2nF6Lin_tDvPs0JVldsuaBzo7pAwi5YAHAhlIxQg,6139
 ai_edge_torch/generative/examples/gemma3/__init__.py,sha256=JaAnrFoXTl3RJX97XspklkTyqOHVyAgRJsZtzNDd10c,671
-ai_edge_torch/generative/examples/gemma3/convert_gemma3_to_tflite.py,sha256=JLXXn2mFEBs4DlHH_O6hpEG9KInJqsCdWy3DrgUjT1c,1827
+ai_edge_torch/generative/examples/gemma3/convert_gemma3_to_tflite.py,sha256=SsiK9xKCyboi5y-HdoFSN02QxRo0XabyzotUq46zO0E,2357
 ai_edge_torch/generative/examples/gemma3/decoder.py,sha256=shdgLzKDUi0vyNOAsrIVAEFb3Adltsri6Rx1-wxzVf4,15089
 ai_edge_torch/generative/examples/gemma3/gemma3.py,sha256=ZorRtnbElWsctcA0nEbfwjx0C578voF7fjFEvWSR5Ck,6582
 ai_edge_torch/generative/examples/gemma3/image_encoder.py,sha256=uRoLoBWzFtQz5wFZfPCxbkvZsgPAqSkUUsV3977GbYc,5184
@@ -166,7 +166,8 @@ ai_edge_torch/generative/layers/feed_forward_test.py,sha256=8ZGy79BBpsyS6yKKDEKr
 ai_edge_torch/generative/layers/kv_cache.py,sha256=b-7shzDaKexmvQF7P3SiAmIz4ZofjYWv3m5u71GojsA,10460
 ai_edge_torch/generative/layers/lora.py,sha256=hsvWLLOnW7HQ0AysOZu30x_cetMquDd1tjfyLz8HCSU,17892
 ai_edge_torch/generative/layers/model_config.py,sha256=X_gjN5524DCDBNXsX5GrOBlkKM4UHzj_RfdCD0-VOxQ,8572
-ai_edge_torch/generative/layers/normalization.py,sha256=MbwH-n80Fob5YvjBzdqDjBizMHLzSJGYRDdbD-rL5C0,6174
+ai_edge_torch/generative/layers/normalization.py,sha256=ijwCpi22NLX-Sygwy5sK9l9WjGvbPIhZvVwoBAonWAo,7014
+ai_edge_torch/generative/layers/normalization_test.py,sha256=zwurZly-TgFxdgVVdpzu9vCpcLbd5RYt_gKg9Lfg1jI,2248
 ai_edge_torch/generative/layers/rotary_position_embedding.py,sha256=975zR202MdIrILJ7blceAcxrNqX1ZCN0ECKG1gz-bV8,2655
 ai_edge_torch/generative/layers/scaled_dot_product_attention.py,sha256=2_AgwENsaOgaxgiSqgoj0V0JzQ09dFtP_nBhX-lJK2g,5648
 ai_edge_torch/generative/layers/scaled_dot_product_attention_test.py,sha256=c6JBMQsq9XeMmR1XvGEIidNsoh-YIvichXo2LwVHgr4,3301
@@ -194,7 +195,7 @@ ai_edge_torch/generative/test/utils.py,sha256=tF6aCfAGJnc9dmzCnZCEOuKNVimfWOqscv
 ai_edge_torch/generative/utilities/__init__.py,sha256=-_jxnnFnCgnTU4oTm4MnRsvL5lqhomBNdFBbqfmfHPo,720
 ai_edge_torch/generative/utilities/converter.py,sha256=4zcDlhgCQQyLylH8NLgVjnelou2pW6HWJHBFYsFyHuw,15020
 ai_edge_torch/generative/utilities/export_config.py,sha256=5IvR3grlMd4mWO5c_Y4x9Fk1b1xa57MzlYNE8XUaN28,2049
-ai_edge_torch/generative/utilities/loader.py,sha256=nw2REQ9sGWDwphShfRqNFICYmwIjqLp6bDcwVmsNTtg,14067
+ai_edge_torch/generative/utilities/loader.py,sha256=tSiew77hB_zyn6rpcfegSg1zrriqHSz63KjV9_llBxg,14893
 ai_edge_torch/generative/utilities/model_builder.py,sha256=tBfOcsI_NcneggHqkCSydYN3ZgmkzPc6nW0AJrA81wI,6461
 ai_edge_torch/generative/utilities/moonshine_loader.py,sha256=_RpFabSqtGH5PHiP3_1f6QfO14qMADUxr_HGRlVDFB0,4891
 ai_edge_torch/generative/utilities/stable_diffusion_loader.py,sha256=dqPD9qRXEWtU3ombslOC-BE2l_dMwHoCNu7NsIJhsso,36158
@@ -251,8 +252,8 @@ ai_edge_torch/testing/__init__.py,sha256=_yGgvnBZWb7T3IN3mc4x1sS4vM96HZwM8pwIcPG
 ai_edge_torch/testing/export.py,sha256=k5mGDGzwc23Z4zaIVDs8CNh-oOt64gsf9MS9NjhbPy4,3293
 ai_edge_torch/testing/model_coverage/__init__.py,sha256=5P8J6Zk5YYtDvTBucFvB9NGSRI7Gw_24WnrbhXgycEE,765
 ai_edge_torch/testing/model_coverage/model_coverage.py,sha256=UPB448aMDUyC0HNYVqio2rcJPnDN0tBQMP08J6vPYew,4718
-ai_edge_torch_nightly-0.5.0.dev20250514.dist-info/LICENSE,sha256=z8d0m5b2O9McPEK1xHG_dWgUBT6EfBDz6wA0F7xSPTA,11358
-ai_edge_torch_nightly-0.5.0.dev20250514.dist-info/METADATA,sha256=4_d1LvNhvXOHKlqYZDcBYSLdYDmoGvWMgCK5PJasNiU,2074
-ai_edge_torch_nightly-0.5.0.dev20250514.dist-info/WHEEL,sha256=tZoeGjtWxWRfdplE7E3d45VPlLNQnvbKiYnx7gwAy8A,92
-ai_edge_torch_nightly-0.5.0.dev20250514.dist-info/top_level.txt,sha256=5KXRaF2hwkApYxf7Y8y_tVb9aulGTlbOoNdbx1aKRkE,14
-ai_edge_torch_nightly-0.5.0.dev20250514.dist-info/RECORD,,
+ai_edge_torch_nightly-0.5.0.dev20250515.dist-info/LICENSE,sha256=z8d0m5b2O9McPEK1xHG_dWgUBT6EfBDz6wA0F7xSPTA,11358
+ai_edge_torch_nightly-0.5.0.dev20250515.dist-info/METADATA,sha256=FmCPouaJYszNPCOfgIx8WGFkGv5LrqR6_OGpciU2eKc,2074
+ai_edge_torch_nightly-0.5.0.dev20250515.dist-info/WHEEL,sha256=tZoeGjtWxWRfdplE7E3d45VPlLNQnvbKiYnx7gwAy8A,92
+ai_edge_torch_nightly-0.5.0.dev20250515.dist-info/top_level.txt,sha256=5KXRaF2hwkApYxf7Y8y_tVb9aulGTlbOoNdbx1aKRkE,14
+ai_edge_torch_nightly-0.5.0.dev20250515.dist-info/RECORD,,