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

ai_edge_torch/generative/examples/hammer/hammer.py

@@ -15,8 +15,10 @@
 
 """Example of building Hammer 2.1 models."""
 
+from typing import Callable, Dict
 import ai_edge_torch.generative.layers.model_config as cfg
 from ai_edge_torch.generative.utilities import model_builder
+import torch
 from torch import nn
 
 TENSOR_NAMES = model_builder.TENSOR_NAMES
@@ -43,9 +45,7 @@ def get_1_5b_model_config(kv_cache_max_len: int = 1024) -> cfg.ModelConfig:
       intermediate_size=8960,
   )
   norm_config = cfg.NormalizationConfig(
-      type=cfg.NormalizationType.RMS_NORM,
-      epsilon=1e-06,
-      enable_hlfb=True,
+      type=cfg.NormalizationType.RMS_NORM, epsilon=1e-06
   )
   block_config = cfg.TransformerBlockConfig(
       attn_config=attn_config,
@@ -61,7 +61,6 @@ def get_1_5b_model_config(kv_cache_max_len: int = 1024) -> cfg.ModelConfig:
       kv_cache_max_len=kv_cache_max_len,
       block_configs=block_config,
       final_norm_config=norm_config,
-      enable_hlfb=True,
   )
   return config
 
@@ -89,19 +88,29 @@ def get_fake_model_config(**kwargs) -> cfg.ModelConfig:
   return config
 
 
-def build_1_5b_model(checkpoint_path: str, **kwargs) -> nn.Module:
+def build_1_5b_model(
+    checkpoint_path: str,
+    custom_loader: Callable[[str], Dict[str, torch.Tensor]] = None,
+    **kwargs
+) -> nn.Module:
   return model_builder.build_decoder_only_model(
       checkpoint_path=checkpoint_path,
       config=get_1_5b_model_config(**kwargs),
       tensor_names=TENSOR_NAMES,
       model_class=Hammer,
+      custom_loader=custom_loader,
   )
 
 
-def build_0_5b_model(checkpoint_path: str, **kwargs) -> nn.Module:
+def build_0_5b_model(
+    checkpoint_path: str,
+    custom_loader: Callable[[str], Dict[str, torch.Tensor]] = None,
+    **kwargs
+) -> nn.Module:
   return model_builder.build_decoder_only_model(
       checkpoint_path=checkpoint_path,
       config=get_0_5b_model_config(**kwargs),
       tensor_names=TENSOR_NAMES,
       model_class=Hammer,
+      custom_loader=custom_loader,
   )

ai_edge_torch/generative/examples/llama/convert_to_tflite.py

@@ -19,6 +19,7 @@ from absl import app
 from ai_edge_torch.generative.examples.llama import llama
 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('llama')
@@ -37,8 +38,13 @@ _BUILDER = {
 
 
 def main(_):
+  checkpoint_path = flags.FLAGS.checkpoint_path
   pytorch_model = _BUILDER[_MODEL_SIZE.value](
-      flags.FLAGS.checkpoint_path, kv_cache_max_len=flags.FLAGS.kv_cache_max_len
+      checkpoint_path,
+      custom_loader=loader.maybe_get_custom_loader(
+          checkpoint_path, flags.FLAGS.custom_checkpoint_loader
+      ),
+      kv_cache_max_len=flags.FLAGS.kv_cache_max_len,
   )
   converter.convert_to_tflite(
       pytorch_model,

ai_edge_torch/generative/examples/llama/llama.py

@@ -17,7 +17,7 @@
 
 from functools import partial
 import math
-from typing import Tuple
+from typing import Callable, Dict, Tuple
 
 import ai_edge_torch.generative.layers.model_config as cfg
 from ai_edge_torch.generative.utilities import model_builder
@@ -121,9 +121,7 @@ def get_1b_model_config(kv_cache_max_len: int = 1024) -> cfg.ModelConfig:
       activation=cfg.ActivationConfig(cfg.ActivationType.SILU),
       intermediate_size=8192,
   )
-  norm_config = cfg.NormalizationConfig(
-      type=cfg.NormalizationType.RMS_NORM, enable_hlfb=True,
-  )
+  norm_config = cfg.NormalizationConfig(type=cfg.NormalizationType.RMS_NORM)
   block_config = cfg.TransformerBlockConfig(
       attn_config=attn_config,
       ff_config=ff_config,
@@ -152,7 +150,6 @@ def get_1b_model_config(kv_cache_max_len: int = 1024) -> cfg.ModelConfig:
       kv_cache_max_len=kv_cache_max_len,
       block_configs=block_config,
       final_norm_config=norm_config,
-      enable_hlfb=True,
       build_rope=build_rope,
   )
   return config
@@ -180,19 +177,38 @@ def get_fake_model_config(**kwargs) -> cfg.ModelConfig:
 
 
 def _build_model(
-    checkpoint_path: str, config: cfg.ModelConfig
+    checkpoint_path: str,
+    config: cfg.ModelConfig,
+    custom_loader: Callable[[str], Dict[str, torch.Tensor]] = None,
 ) -> torch.nn.Module:
   return model_builder.build_decoder_only_model(
       checkpoint_path=checkpoint_path,
       config=config,
       tensor_names=TENSOR_NAMES,
       model_class=Llama,
+      custom_loader=custom_loader,
   )
 
 
-def build_1b_model(checkpoint_path: str, **kwargs) -> torch.nn.Module:
-  return _build_model(checkpoint_path, get_1b_model_config(**kwargs))
+def build_1b_model(
+    checkpoint_path: str,
+    custom_loader: Callable[[str], Dict[str, torch.Tensor]] = None,
+    **kwargs
+) -> torch.nn.Module:
+  return _build_model(
+      checkpoint_path,
+      get_1b_model_config(**kwargs),
+      custom_loader=custom_loader,
+  )
 
 
-def build_3b_model(checkpoint_path: str, **kwargs) -> torch.nn.Module:
-  return _build_model(checkpoint_path, get_3b_model_config(**kwargs))
+def build_3b_model(
+    checkpoint_path: str,
+    custom_loader: Callable[[str], Dict[str, torch.Tensor]] = None,
+    **kwargs
+) -> torch.nn.Module:
+  return _build_model(
+      checkpoint_path,
+      get_3b_model_config(**kwargs),
+      custom_loader=custom_loader,
+  )

ai_edge_torch/generative/examples/moonshine/convert_moonshine_to_tflite.py

@@ -22,7 +22,6 @@ from absl import app
 from absl import flags
 import ai_edge_torch
 from ai_edge_torch.generative.examples.moonshine import moonshine
-from ai_edge_torch.generative.utilities import converter
 import torch
 
 _CHECKPOINT_PATH = flags.DEFINE_string(

ai_edge_torch/generative/examples/openelm/convert_to_tflite.py

@@ -19,13 +19,19 @@ from absl import app
 from ai_edge_torch.generative.examples.openelm import openelm
 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("openelm")
 
 
 def main(_):
+  checkpoint_path = flags.FLAGS.checkpoint_path
   pytorch_model = openelm.build_model(
-      flags.FLAGS.checkpoint_path, kv_cache_max_len=flags.FLAGS.kv_cache_max_len
+      checkpoint_path,
+      custom_loader=loader.maybe_get_custom_loader(
+          checkpoint_path, flags.FLAGS.custom_checkpoint_loader
+      ),
+      kv_cache_max_len=flags.FLAGS.kv_cache_max_len,
   )
   converter.convert_to_tflite(
       pytorch_model,

ai_edge_torch/generative/examples/openelm/openelm.py

@@ -15,9 +15,11 @@
 
 """Example of building an OpenELM model."""
 
+from typing import Callable, Dict
 import ai_edge_torch.generative.layers.model_config as cfg
 from ai_edge_torch.generative.utilities import model_builder
 import ai_edge_torch.generative.utilities.loader as loading_utils
+import torch
 from torch import nn
 
 TENSOR_NAMES = loading_utils.ModelLoader.TensorNames(
@@ -51,7 +53,7 @@ def get_model_config(kv_cache_max_len: int = 1024) -> cfg.ModelConfig:
     The model config for an OpenELM model.
   """
   norm_config = cfg.NormalizationConfig(
-      type=cfg.NormalizationType.RMS_NORM, epsilon=1e-6, enable_hlfb=True
+      type=cfg.NormalizationType.RMS_NORM, epsilon=1e-6
   )
   num_heads = [12] * 4 + [16] * 14 + [20] * 12 + [24] * 6
   num_query_groups = [3] * 4 + [4] * 14 + [5] * 12 + [6] * 6
@@ -99,7 +101,6 @@ def get_model_config(kv_cache_max_len: int = 1024) -> cfg.ModelConfig:
       kv_cache_max_len=kv_cache_max_len,
       block_configs=[get_block_config(i) for i in range(num_layers)],
       final_norm_config=norm_config,
-      enable_hlfb=True,
   )
   return config
 
@@ -118,10 +119,15 @@ def get_fake_model_config(kv_cache_max_len: int = 128) -> cfg.ModelConfig:
   return config
 
 
-def build_model(checkpoint_path: str, **kwargs) -> nn.Module:
+def build_model(
+    checkpoint_path: str,
+    custom_loader: Callable[[str], Dict[str, torch.Tensor]] = None,
+    **kwargs
+) -> nn.Module:
   return model_builder.build_decoder_only_model(
       checkpoint_path=checkpoint_path,
       config=get_model_config(**kwargs),
       tensor_names=TENSOR_NAMES,
       model_class=OpenELM,
+      custom_loader=custom_loader,
   )

ai_edge_torch/generative/examples/paligemma/convert_to_tflite.py

@@ -19,6 +19,7 @@ from absl import app
 from ai_edge_torch.generative.examples.paligemma import paligemma
 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
 import torch
 
 flags = converter.define_conversion_flags('paligemma2-3b-224')
@@ -32,9 +33,13 @@ _VERSION = flags.DEFINE_enum(
 
 
 def main(_):
+  checkpoint_path = flags.FLAGS.checkpoint_path
   pytorch_model = paligemma.build_model(
-      flags.FLAGS.checkpoint_path,
+      checkpoint_path,
       version=int(_VERSION.value),
+      custom_loader=loader.maybe_get_custom_loader(
+          checkpoint_path, flags.FLAGS.custom_checkpoint_loader
+      ),
       kv_cache_max_len=flags.FLAGS.kv_cache_max_len,
   )
 

ai_edge_torch/generative/examples/paligemma/decoder.py

@@ -110,9 +110,7 @@ def get_decoder_config(kv_cache_max_len: int = 1024) -> cfg.ModelConfig:
       intermediate_size=16384,
   )
   norm_config = cfg.NormalizationConfig(
-      type=cfg.NormalizationType.RMS_NORM,
-      epsilon=1e-6,
-      zero_centered=True,
+      type=cfg.NormalizationType.RMS_NORM, epsilon=1e-6, zero_centered=True
   )
   block_config = cfg.TransformerBlockConfig(
       attn_config=attn_config,
@@ -131,7 +129,6 @@ def get_decoder_config(kv_cache_max_len: int = 1024) -> cfg.ModelConfig:
       block_configs=block_config,
       final_norm_config=norm_config,
       lm_head_use_bias=False,
-      enable_hlfb=True,
   )
   return config
 

ai_edge_torch/generative/examples/paligemma/decoder2.py

@@ -93,9 +93,7 @@ def get_decoder2_config(kv_cache_max_len: int = 1024) -> cfg.ModelConfig:
     The model config for the decoder of a PaliGemma 3B model.
   """
   norm_config = cfg.NormalizationConfig(
-      type=cfg.NormalizationType.RMS_NORM,
-      epsilon=1e-6,
-      zero_centered=True,
+      type=cfg.NormalizationType.RMS_NORM, epsilon=1e-6, zero_centered=True
   )
   ff_config = cfg.FeedForwardConfig(
       type=cfg.FeedForwardType.GATED,
@@ -139,7 +137,6 @@ def get_decoder2_config(kv_cache_max_len: int = 1024) -> cfg.ModelConfig:
       block_configs=[get_block_config(i) for i in range(num_layers)],
       final_norm_config=norm_config,
       lm_head_use_bias=False,
-      enable_hlfb=True,
       final_logit_softcap=30.0,
   )
   return config

ai_edge_torch/generative/examples/paligemma/image_encoder.py

@@ -66,7 +66,8 @@ class SiglipVisionEncoder(nn.Module):
         config.image_embedding.image_size // config.image_embedding.patch_size
     ) ** 2
     self.tok_embedding_position = nn.Parameter(
-        torch.zeros((num_patches, config.embedding_dim))
+        torch.zeros((num_patches, config.embedding_dim)),
+        requires_grad=False,
     )
 
     self.transformer_blocks = nn.ModuleList(
@@ -117,9 +118,7 @@ def get_image_encoder_config() -> cfg.ModelConfig:
       use_bias=True,
   )
   norm_config = cfg.NormalizationConfig(
-      type=cfg.NormalizationType.LAYER_NORM,
-      epsilon=1e-6,
-      enable_hlfb=True,
+      type=cfg.NormalizationType.LAYER_NORM, epsilon=1e-6
   )
   block_config = cfg.TransformerBlockConfig(
       attn_config=attn_config,
@@ -136,7 +135,6 @@ def get_image_encoder_config() -> cfg.ModelConfig:
       image_embedding=image_embedding_config,
       block_configs=block_config,
       final_norm_config=norm_config,
-      enable_hlfb=True,
   )
   return config
 

ai_edge_torch/generative/examples/paligemma/paligemma.py

@@ -16,7 +16,7 @@
 """Example of building a full-stack of PaliGemma model."""
 
 import dataclasses
-from typing import Optional
+from typing import Callable, Dict, Optional
 
 from ai_edge_torch.generative.examples.paligemma import decoder
 from ai_edge_torch.generative.examples.paligemma import decoder2
@@ -139,7 +139,12 @@ def get_fake_model_config(get_decoder_config, **kwargs) -> PaliGemmaConfig:
   )
 
 
-def build_model(checkpoint_path: str, version: int = 2, **kwargs) -> PaliGemma:
+def build_model(
+    checkpoint_path: str,
+    version: int = 2,
+    custom_loader: Callable[[str], Dict[str, torch.Tensor]] = None,
+    **kwargs,
+) -> PaliGemma:
   if version == 1:
     decoder_class = decoder.Decoder
     decoder_tensor_names = decoder.TENSOR_NAMES
@@ -153,15 +158,17 @@ def build_model(checkpoint_path: str, version: int = 2, **kwargs) -> PaliGemma:
   model = PaliGemma(config, decoder_class)
   # Load the parameters of image encoder.
   loader = loading_utils.ModelLoader(
-      checkpoint_path, image_encoder.TENSOR_NAMES
+      checkpoint_path, image_encoder.TENSOR_NAMES, custom_loader
   )
   loader.load(model.image_encoder, strict=False)
   # Load the parameters of decoder.
-  loader = loading_utils.ModelLoader(checkpoint_path, decoder_tensor_names)
+  loader = loading_utils.ModelLoader(
+      checkpoint_path, decoder_tensor_names, custom_loader
+  )
   loader.load(model.decoder, strict=False)
 
   # Load the parameters of image projection.
-  loader = loading_utils.ModelLoader(checkpoint_path, None)
+  loader = loading_utils.ModelLoader(checkpoint_path, None, custom_loader)
   state = loader.get_state()
   converted_state = dict()
   converted_state["weight"] = state.pop(f"{PROJECTION_TENSOR_NAME}.weight")

ai_edge_torch/generative/examples/paligemma/verify.py

@@ -21,6 +21,7 @@ from absl import app
 from absl import flags
 from ai_edge_torch.generative.examples.paligemma import paligemma
 from ai_edge_torch.generative.layers import kv_cache
+from ai_edge_torch.generative.utilities import transformers_verifier
 from ai_edge_torch.generative.utilities import verifier
 import kagglehub
 from PIL import Image
@@ -39,10 +40,15 @@ _IMAGE_URL = flags.DEFINE_string(
     "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/transformers/tasks/car.jpg?download=true",
     "The image URI to encode.",
 )
-_PROMPTS = flags.DEFINE_string(
-    "prompts",
+_PROMPTS_WITH_IMAGE = flags.DEFINE_string(
+    "prompts_with_image",
     "<image><bos>describe en",
-    "The input prompts to generate answers.",
+    "The input prompts to generate answers with an image.",
+)
+_PROMPTS_TEXT_ONLY = flags.DEFINE_multi_string(
+    "prompts_text_only",
+    "What is the meaning of life?",
+    "The input prompts to generate answers only with text.",
 )
 _MAX_NEW_TOKENS = flags.DEFINE_integer(
     "max_new_tokens",
@@ -84,6 +90,7 @@ def main(_):
   reauthored_model = paligemma.build_model(
       reauthored_checkpoint, version=int(_VERSION.value)
   )
+  wrapped_reauthored_model = ReauthoredPaliGemmaWrapper(reauthored_model)
 
   logging.info("Loading the processor from: %s", checkpoint)
   # It works only when GemmaTokenizerFast is available. In some environments,
@@ -91,9 +98,25 @@ def main(_):
   # sentencepiece model file properly.
   processor = transformers.AutoProcessor.from_pretrained(checkpoint)
 
+  logging.info("Verifying with text-only prompts...")
+  verifier.verify_reauthored_model(
+      original_model=transformers_verifier.TransformersModelWrapper(
+          original_model
+      ),
+      reauthored_model=wrapped_reauthored_model,
+      tokenizer=verifier.TokenizerWrapper(processor.tokenizer),
+      generate_prompts=_PROMPTS_TEXT_ONLY.value,
+      max_new_tokens=_MAX_NEW_TOKENS.value,
+      verify_inputs=False,  # Numeric check not working. Disable it for now.
+      atol=1e-04,
+  )
+
+  logging.info("Verifying with image input...")
   logging.info("Loading the image from: %s", _IMAGE_URL.value)
   image = Image.open(requests.get(_IMAGE_URL.value, stream=True).raw)
-  inputs = processor(text=_PROMPTS.value, images=image, return_tensors="pt")
+  inputs = processor(
+      text=_PROMPTS_WITH_IMAGE.value, images=image, return_tensors="pt"
+  )
 
   logging.info("Verifying the reauthored model with model.forward()...")
   logging.info("Forwarding the original model...")
@@ -104,7 +127,6 @@ def main(_):
   logging.info("outputs_original: %s", outputs_original)
 
   logging.info("Forwarding the reauthored model...")
-  wrapped_reauthored_model = ReauthoredPaliGemmaWrapper(reauthored_model)
   outputs_reauthored = wrapped_reauthored_model.forward(
       tokens=inputs["input_ids"],
       pixel_values=inputs["pixel_values"],

ai_edge_torch/generative/examples/phi/convert_phi3_to_tflite.py

@@ -19,13 +19,19 @@ from absl import app
 from ai_edge_torch.generative.examples.phi import phi3
 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("phi3")
 
 
 def main(_):
+  checkpoint_path = flags.FLAGS.checkpoint_path
   pytorch_model = phi3.build_model(
-      flags.FLAGS.checkpoint_path, kv_cache_max_len=flags.FLAGS.kv_cache_max_len
+      checkpoint_path,
+      custom_loader=loader.maybe_get_custom_loader(
+          checkpoint_path, flags.FLAGS.custom_checkpoint_loader
+      ),
+      kv_cache_max_len=flags.FLAGS.kv_cache_max_len,
   )
   converter.convert_to_tflite(
       pytorch_model,

ai_edge_torch/generative/examples/phi/convert_phi4_to_tflite.py

@@ -19,13 +19,19 @@ from absl import app
 from ai_edge_torch.generative.examples.phi import phi4
 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("phi4")
 
 
 def main(_):
+  checkpoint_path = flags.FLAGS.checkpoint_path
   pytorch_model = phi4.build_model(
-      flags.FLAGS.checkpoint_path, kv_cache_max_len=flags.FLAGS.kv_cache_max_len
+      checkpoint_path,
+      custom_loader=loader.maybe_get_custom_loader(
+          checkpoint_path, flags.FLAGS.custom_checkpoint_loader
+      ),
+      kv_cache_max_len=flags.FLAGS.kv_cache_max_len,
   )
   converter.convert_to_tflite(
       pytorch_model,

ai_edge_torch/generative/examples/phi/convert_to_tflite.py

@@ -20,13 +20,19 @@ from absl import app
 from ai_edge_torch.generative.examples.phi import phi2
 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("phi2")
 
 
 def main(_):
+  checkpoint_path = flags.FLAGS.checkpoint_path
   pytorch_model = phi2.build_model(
-      flags.FLAGS.checkpoint_path, kv_cache_max_len=flags.FLAGS.kv_cache_max_len
+      checkpoint_path,
+      custom_loader=loader.maybe_get_custom_loader(
+          checkpoint_path, flags.FLAGS.custom_checkpoint_loader
+      ),
+      kv_cache_max_len=flags.FLAGS.kv_cache_max_len,
   )
   converter.convert_to_tflite(
       pytorch_model,

ai_edge_torch/generative/examples/phi/phi2.py

@@ -15,9 +15,11 @@
 
 """Example of building a Phi-2 model."""
 
+from typing import Callable, Dict
 import ai_edge_torch.generative.layers.model_config as cfg
 from ai_edge_torch.generative.utilities import model_builder
 import ai_edge_torch.generative.utilities.loader as loading_utils
+import torch
 from torch import nn
 
 TENSOR_NAMES = loading_utils.ModelLoader.TensorNames(
@@ -64,9 +66,7 @@ def get_model_config(kv_cache_max_len: int = 1024) -> cfg.ModelConfig:
       intermediate_size=10240,
       use_bias=True,
   )
-  norm_config = cfg.NormalizationConfig(
-      type=cfg.NormalizationType.LAYER_NORM, enable_hlfb=True
-  )
+  norm_config = cfg.NormalizationConfig(type=cfg.NormalizationType.LAYER_NORM)
   block_config = cfg.TransformerBlockConfig(
       attn_config=attn_config,
       ff_config=ff_config,
@@ -83,7 +83,6 @@ def get_model_config(kv_cache_max_len: int = 1024) -> cfg.ModelConfig:
       final_norm_config=norm_config,
       lm_head_use_bias=True,
       lm_head_share_weight_with_embedding=False,
-      enable_hlfb=True,
   )
   return config
 
@@ -98,10 +97,15 @@ def get_fake_model_config(kv_cache_max_len: int = 128) -> cfg.ModelConfig:
   return config
 
 
-def build_model(checkpoint_path: str, **kwargs) -> nn.Module:
+def build_model(
+    checkpoint_path: str,
+    custom_loader: Callable[[str], Dict[str, torch.Tensor]] = None,
+    **kwargs
+) -> nn.Module:
   return model_builder.build_decoder_only_model(
       checkpoint_path=checkpoint_path,
       config=get_model_config(**kwargs),
       tensor_names=TENSOR_NAMES,
       model_class=Phi2,
+      custom_loader=custom_loader,
   )

ai_edge_torch/generative/examples/phi/phi3.py

@@ -17,7 +17,7 @@
 
 from functools import partial
 import math
-from typing import Tuple
+from typing import Callable, Dict, Tuple
 
 import ai_edge_torch.generative.layers.model_config as cfg
 from ai_edge_torch.generative.utilities import model_builder
@@ -162,9 +162,7 @@ def get_model_config(kv_cache_max_len: int = 1024) -> cfg.ModelConfig:
       activation=cfg.ActivationConfig(cfg.ActivationType.SILU_GLU),
       intermediate_size=8192,
   )
-  norm_config = cfg.NormalizationConfig(
-      type=cfg.NormalizationType.RMS_NORM, enable_hlfb=True,
-  )
+  norm_config = cfg.NormalizationConfig(type=cfg.NormalizationType.RMS_NORM)
   block_config = cfg.TransformerBlockConfig(
       attn_config=attn_config,
       ff_config=ff_config,
@@ -192,7 +190,6 @@ def get_model_config(kv_cache_max_len: int = 1024) -> cfg.ModelConfig:
       block_configs=block_config,
       final_norm_config=norm_config,
       lm_head_share_weight_with_embedding=False,
-      enable_hlfb=True,
       build_rope=build_rope,
   )
   return config
@@ -208,11 +205,16 @@ def get_fake_model_config(kv_cache_max_len: int = 128) -> cfg.ModelConfig:
   return config
 
 
-def build_model(checkpoint_path: str, **kwargs) -> torch.nn.Module:
+def build_model(
+    checkpoint_path: str,
+    custom_loader: Callable[[str], Dict[str, torch.Tensor]] = None,
+    **kwargs
+) -> torch.nn.Module:
   """Instantiates the model instance and load checkpoint if provided."""
   return model_builder.build_decoder_only_model(
       checkpoint_path=checkpoint_path,
       config=get_model_config(**kwargs),
       tensor_names=TENSOR_NAMES,
       model_class=Phi3_5Mini,
+      custom_loader=custom_loader,
   )

ai_edge_torch/generative/examples/phi/phi4.py

@@ -17,7 +17,7 @@
 
 from functools import partial
 import math
-from typing import Tuple
+from typing import Callable, Dict, Tuple
 
 import ai_edge_torch.generative.layers.model_config as cfg
 from ai_edge_torch.generative.utilities import model_builder
@@ -112,9 +112,7 @@ def get_model_config(kv_cache_max_len: int = 1024) -> cfg.ModelConfig:
       activation=cfg.ActivationConfig(cfg.ActivationType.SILU_GLU),
       intermediate_size=8192,
   )
-  norm_config = cfg.NormalizationConfig(
-      type=cfg.NormalizationType.RMS_NORM, enable_hlfb=True
-  )
+  norm_config = cfg.NormalizationConfig(type=cfg.NormalizationType.RMS_NORM)
   block_config = cfg.TransformerBlockConfig(
       attn_config=attn_config,
       ff_config=ff_config,
@@ -141,7 +139,6 @@ def get_model_config(kv_cache_max_len: int = 1024) -> cfg.ModelConfig:
       embedding_dim=3072,
       block_configs=block_config,
       final_norm_config=norm_config,
-      enable_hlfb=True,
       build_rope=build_rope,
   )
   return config
@@ -157,11 +154,16 @@ def get_fake_model_config(kv_cache_max_len: int = 128) -> cfg.ModelConfig:
   return config
 
 
-def build_model(checkpoint_path: str, **kwargs) -> torch.nn.Module:
+def build_model(
+    checkpoint_path: str,
+    custom_loader: Callable[[str], Dict[str, torch.Tensor]] = None,
+    **kwargs
+) -> torch.nn.Module:
   """Instantiates the model instance and load checkpoint if provided."""
   return model_builder.build_decoder_only_model(
       checkpoint_path=checkpoint_path,
       config=get_model_config(**kwargs),
       tensor_names=TENSOR_NAMES,
       model_class=Phi4Mini,
+      custom_loader=custom_loader,
   )

ai_edge_torch/generative/examples/qwen/convert_to_tflite.py

@@ -19,6 +19,7 @@ from absl import app
 from ai_edge_torch.generative.examples.qwen import qwen
 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('qwen')
 
@@ -37,8 +38,13 @@ _BUILDER = {
 
 
 def main(_):
+  checkpoint_path = flags.FLAGS.checkpoint_path
   pytorch_model = _BUILDER[_MODEL_SIZE.value](
-      flags.FLAGS.checkpoint_path, kv_cache_max_len=flags.FLAGS.kv_cache_max_len
+      checkpoint_path,
+      custom_loader=loader.maybe_get_custom_loader(
+          checkpoint_path, flags.FLAGS.custom_checkpoint_loader
+      ),
+      kv_cache_max_len=flags.FLAGS.kv_cache_max_len,
   )
   converter.convert_to_tflite(
       pytorch_model,