ai-edge-torch-nightly 0.5.0.dev20250516__py3-none-any.whl → 0.5.0.dev20250518__py3-none-any.whl

ai_edge_torch/_convert/conversion.py

@@ -38,6 +38,7 @@ def _run_convert_passes(
   )
 
   passes = [
+      fx_passes.EliminateDeadCodePass(),
       fx_passes.OptimizeLayoutTransposesPass(),
       fx_passes.CanonicalizePass(),
       fx_passes.BuildAtenCompositePass(),

ai_edge_torch/_convert/fx_passes/__init__.py

@@ -17,6 +17,7 @@ from typing import Sequence, Union
 
 from ai_edge_torch._convert.fx_passes.build_aten_composite_pass import BuildAtenCompositePass
 from ai_edge_torch._convert.fx_passes.cast_inputs_bf16_to_f32_pass import CastInputsBf16ToF32Pass
+from ai_edge_torch._convert.fx_passes.eliminate_dead_code_pass import EliminateDeadCodePass
 from ai_edge_torch._convert.fx_passes.inject_mlir_debuginfo_pass import InjectMlirDebuginfoPass
 from ai_edge_torch._convert.fx_passes.optimize_layout_transposes_pass import OptimizeLayoutTransposesPass
 from ai_edge_torch._convert.fx_passes.remove_non_user_outputs_pass import RemoveNonUserOutputsPass

ai_edge_torch/_convert/fx_passes/eliminate_dead_code_pass.py

@@ -0,0 +1,40 @@
+# 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.
+# ==============================================================================
+"""Pass to eliminate dead code for ai-edge-torch conversion."""
+
+
+from ai_edge_torch import fx_infra
+import torch
+
+
+class EliminateDeadCodePass(fx_infra.PassBase):
+  """Eliminates dead code with dedicated rules for ai-edge-torch conversion."""
+
+  def call(self, graph_module: torch.fx.GraphModule):
+    def is_impure_node(node: torch.fx.Node):
+      # Starting from torch 2.7.0, random torch ops with
+      # _nondeterministic_seeded set are no longer considered pure. However,
+      # for conversion, unused random ops/tensors should still be removed.
+      if getattr(node.target, "_nondeterministic_seeded", False):
+        return False
+      return node.is_impure()
+
+    try:
+      graph_module.graph.eliminate_dead_code(is_impure_node)
+    except TypeError:
+      # eliminate_dead_code has no is_impure_node input in old torch versions.
+      pass
+
+    return fx_infra.PassResult(graph_module, True)

ai_edge_torch/generative/examples/amd_llama_135m/amd_llama_135m.py

@@ -51,9 +51,7 @@ def get_model_config(kv_cache_max_len: int = 1024) -> cfg.ModelConfig:
       activation=cfg.ActivationConfig(cfg.ActivationType.SILU),
       intermediate_size=2048,
   )
-  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,
@@ -69,7 +67,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,
   )
   return config
 

ai_edge_torch/generative/examples/amd_llama_135m/verify.py

@@ -15,15 +15,9 @@
 
 """Verifies the reauthored AMD-Llama-135M model."""
 
-import logging
-import pathlib
-
 from absl import app
 from absl import flags
-from ai_edge_torch.generative.examples.amd_llama_135m import amd_llama_135m
-from ai_edge_torch.generative.utilities import transformers_verifier
-from ai_edge_torch.generative.utilities import verifier
-import transformers
+from ai_edge_torch.generative.examples.amd_llama_135m import verify_util
 
 
 _PROMPTS = flags.DEFINE_multi_string(
@@ -39,32 +33,10 @@ _MAX_NEW_TOKENS = flags.DEFINE_integer(
 
 
 def main(_):
-  checkpoint = "amd/AMD-Llama-135m"
-  logging.info("Loading the original model from: %s", checkpoint)
-  original_model = transformers.AutoModelForCausalLM.from_pretrained(
-      checkpoint, trust_remote_code=True
-  )
-
-  # 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 = amd_llama_135m.build_model(str(reauthored_checkpoint))
-
-  logging.info("Loading the tokenizer from: %s", checkpoint)
-  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,
+  verify_util.verify_amd_llama_135m(
+      "amd/AMD-Llama-135m",
       max_new_tokens=_MAX_NEW_TOKENS.value,
-      atol=1e-04,
+      prompts=_PROMPTS.value,
   )
 
 

ai_edge_torch/generative/examples/amd_llama_135m/verify_util.py

@@ -0,0 +1,76 @@
+# 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.
+# ==============================================================================
+"""Utils for verifying the AMD-Llama-135M model."""
+import logging
+import os
+import pathlib
+
+from ai_edge_torch.generative.examples.amd_llama_135m import amd_llama_135m
+from ai_edge_torch.generative.utilities import loader
+from ai_edge_torch.generative.utilities import transformers_verifier
+from ai_edge_torch.generative.utilities import verifier
+import transformers
+
+
+DEFAULT_PROMPTS = ["Tell me a story?\nOnce upon a time"]
+
+
+def verify_amd_llama_135m(
+    checkpoint_dir: str,
+    weight_filename: str = "model.safetensors",
+    max_new_tokens: int = 30,
+    initialize_from_local: bool = True,
+    prompts: list[str] | None = None,
+) -> bool:
+  """Verifies the reauthored AMD-Llama-135M model with a custom loader."""
+  logging.info("Loading the original model from: %s", checkpoint_dir)
+  original_model = transformers.AutoModelForCausalLM.from_pretrained(
+      checkpoint_dir
+  )
+
+  logging.info("Building the reauthored model from: %s", checkpoint_dir)
+  custom_loader = (
+      None
+      if initialize_from_local
+      else loader.get_custom_loader("", "safetensors")
+  )
+
+  if initialize_from_local:
+    # Locate the cached dir.
+    cached_config_file = transformers.utils.cached_file(
+        checkpoint_dir, transformers.utils.CONFIG_NAME
+    )
+    reauthored_checkpoint = pathlib.Path(cached_config_file).parent
+  else:
+    reauthored_checkpoint = os.path.join(checkpoint_dir, weight_filename)
+
+  logging.info("Building the reauthored model from: %s", reauthored_checkpoint)
+  reauthored_model = amd_llama_135m.build_model(
+      checkpoint_path=reauthored_checkpoint,
+      custom_loader=custom_loader,
+  )
+
+  logging.info("Loading the tokenizer from: %s", checkpoint_dir)
+  tokenizer = transformers.AutoTokenizer.from_pretrained(checkpoint_dir)
+  return verifier.verify_reauthored_model(
+      original_model=transformers_verifier.TransformersModelWrapper(
+          original_model
+      ),
+      reauthored_model=verifier.ReauthoredModelWrapper(reauthored_model),
+      tokenizer=verifier.TokenizerWrapper(tokenizer),
+      generate_prompts=DEFAULT_PROMPTS if prompts is None else prompts,
+      max_new_tokens=max_new_tokens,
+      atol=1e-04,
+  )

ai_edge_torch/generative/examples/deepseek/deepseek.py

@@ -53,9 +53,7 @@ def get_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,
@@ -72,7 +70,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,
   )
   return config
 

ai_edge_torch/generative/examples/deepseek/verify.py

@@ -15,15 +15,9 @@
 
 """Verifies the reauthored DeepSeek R1 distilled 1.5B model."""
 
-import logging
-import pathlib
-
 from absl import app
 from absl import flags
-from ai_edge_torch.generative.examples.deepseek import deepseek
-from ai_edge_torch.generative.utilities import transformers_verifier
-from ai_edge_torch.generative.utilities import verifier
-import transformers
+from ai_edge_torch.generative.examples.deepseek import verify_util
 
 
 _PROMPTS = flags.DEFINE_multi_string(
@@ -39,30 +33,10 @@ _MAX_NEW_TOKENS = flags.DEFINE_integer(
 
 
 def main(_):
-  checkpoint = "deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B"
-  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 = deepseek.build_model(str(reauthored_checkpoint))
-
-  logging.info("Loading the tokenizer from: %s", checkpoint)
-  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,
+  verify_util.verify_deepseek_r1_distill_1_5b(
+      "deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B",
       max_new_tokens=_MAX_NEW_TOKENS.value,
-      atol=1e-04,
+      prompts=_PROMPTS.value,
   )
 
 

ai_edge_torch/generative/examples/deepseek/verify_util.py

@@ -0,0 +1,76 @@
+# 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.
+# ==============================================================================
+"""Utils for verifying the DeepSeek R1 distilled 1.5B model."""
+import logging
+import os
+import pathlib
+
+from ai_edge_torch.generative.examples.deepseek import deepseek
+from ai_edge_torch.generative.utilities import loader
+from ai_edge_torch.generative.utilities import transformers_verifier
+from ai_edge_torch.generative.utilities import verifier
+import transformers
+
+
+DEFAULT_PROMPTS = ["What is the meaning of life?"]
+
+
+def verify_deepseek_r1_distill_1_5b(
+    checkpoint_dir: str,
+    weight_filename: str = "model.safetensors",
+    max_new_tokens: int = 30,
+    initialize_from_local: bool = True,
+    prompts: list[str] | None = None,
+) -> bool:
+  """Verifies the reauthored DeepSeek R1 distilled 1.5B model with a custom loader."""
+  logging.info("Loading the original model from: %s", checkpoint_dir)
+  original_model = transformers.AutoModelForCausalLM.from_pretrained(
+      checkpoint_dir
+  )
+
+  logging.info("Building the reauthored model from: %s", checkpoint_dir)
+  custom_loader = (
+      None
+      if initialize_from_local
+      else loader.get_custom_loader("", "safetensors")
+  )
+
+  if initialize_from_local:
+    # Locate the cached dir.
+    cached_config_file = transformers.utils.cached_file(
+        checkpoint_dir, transformers.utils.CONFIG_NAME
+    )
+    reauthored_checkpoint = pathlib.Path(cached_config_file).parent
+  else:
+    reauthored_checkpoint = os.path.join(checkpoint_dir, weight_filename)
+
+  logging.info("Building the reauthored model from: %s", reauthored_checkpoint)
+  reauthored_model = deepseek.build_model(
+      checkpoint_path=reauthored_checkpoint,
+      custom_loader=custom_loader,
+  )
+
+  logging.info("Loading the tokenizer from: %s", checkpoint_dir)
+  tokenizer = transformers.AutoTokenizer.from_pretrained(checkpoint_dir)
+  return verifier.verify_reauthored_model(
+      original_model=transformers_verifier.TransformersModelWrapper(
+          original_model
+      ),
+      reauthored_model=verifier.ReauthoredModelWrapper(reauthored_model),
+      tokenizer=verifier.TokenizerWrapper(tokenizer),
+      generate_prompts=DEFAULT_PROMPTS if prompts is None else prompts,
+      max_new_tokens=max_new_tokens,
+      atol=1e-04,
+  )

ai_edge_torch/generative/examples/gemma/gemma1.py

@@ -65,10 +65,7 @@ def get_model_config_2b(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,
-      enable_hlfb=True,
+      type=cfg.NormalizationType.RMS_NORM, epsilon=1e-6, zero_centered=True
   )
   block_config = cfg.TransformerBlockConfig(
       attn_config=attn_config,
@@ -87,7 +84,6 @@ def get_model_config_2b(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/gemma/gemma2.py

@@ -233,10 +233,7 @@ def get_model_config_2b(kv_cache_max_len: int = 1024) -> cfg.ModelConfig:
     The model config for a Gemma 2B model.
   """
   norm_config = cfg.NormalizationConfig(
-      type=cfg.NormalizationType.RMS_NORM,
-      epsilon=1e-6,
-      zero_centered=True,
-      enable_hlfb=True,
+      type=cfg.NormalizationType.RMS_NORM, epsilon=1e-6, zero_centered=True
   )
   ff_config = cfg.FeedForwardConfig(
       type=cfg.FeedForwardType.GATED,
@@ -284,7 +281,6 @@ def get_model_config_2b(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/gemma/verify_util.py

@@ -17,11 +17,13 @@
 
 import logging
 import os
-from typing import List, Tuple
+from typing import Callable, Dict, List, Tuple
 
+from ai_edge_torch.generative.examples.gemma import gemma1
 from ai_edge_torch.generative.examples.gemma import gemma2
 import ai_edge_torch.generative.layers.attention_utils as attn_utils
 import ai_edge_torch.generative.layers.kv_cache as kv_utils
+from ai_edge_torch.generative.utilities import loader
 from ai_edge_torch.generative.utilities import verifier
 from gemma import config as gemma_config
 from gemma import model as gemma_model
@@ -107,6 +109,7 @@ def verify_reauthored_gemma_model(
     generate_prompts: List[str],
     forward_input_ids: List[List[int]],
     weight_filename: str = "model.ckpt",
+    custom_loader: Callable[[str], Dict[str, torch.Tensor]] | None = None,
     tokenizer_filename: str = "tokenizer.model",
     max_new_tokens: int = 20,
     mask_as_input: bool = False,
@@ -125,7 +128,14 @@ def verify_reauthored_gemma_model(
 
   logging.info("Loading the original model from: %s", checkpoint)
   original_model = gemma_model.GemmaForCausalLM(config).eval()
-  original_model.load_weights(os.path.join(checkpoint, weight_filename))
+  checkpoint_path = os.path.join(checkpoint, weight_filename)
+  if custom_loader is None:
+    original_model.load_weights(checkpoint_path)
+  else:
+    original_model.load_state_dict(
+        custom_loader(checkpoint_path)["model_state_dict"],
+        strict=False,
+    )
 
   return verifier.verify_reauthored_model(
       original_model=GemmaWrapper(original_model),
@@ -144,27 +154,62 @@ def verify_reauthored_gemma_model(
 
 
 def verify_gemma2(
-    gemma2_model_path: str,
+    checkpoint_dir: str,
+    weight_filename: str,
     prompts: List[str],
     max_new_tokens: int,
     mask_as_input: bool = False,
     kv_layout: kv_utils.KVLayout = kv_utils.KV_LAYOUT_DEFAULT,
+    custom_loader: Callable[[str], Dict[str, torch.Tensor]] | None = None,
 ) -> bool:
   """Verifies the reauthored Gemma2 model.
 
   Return True if the verification passes, False otherwise.
   """
-  logging.info("Building the reauthored model from: %s", gemma2_model_path)
-  reauthored_model = gemma2.build_2b_model(gemma2_model_path)
+  checkpoint_path = os.path.join(checkpoint_dir, weight_filename)
+  logging.info("Building the reauthored model from: %s", checkpoint_path)
+  reauthored_model = gemma2.build_2b_model(checkpoint_path, custom_loader)
 
   return verify_reauthored_gemma_model(
-      checkpoint=gemma2_model_path,
+      checkpoint=checkpoint_dir,
       variant="2b-v2",
       reauthored_model=reauthored_model,
       generate_prompts=prompts,
       forward_input_ids=[[2, 651, 9456, 576, 573, 3520, 3858, 603, 235248]],
+      weight_filename=weight_filename,
+      custom_loader=custom_loader,
       max_new_tokens=max_new_tokens,
       mask_as_input=mask_as_input,
       kv_layout=kv_layout,
       atol=1e-04,
   )
+
+
+def verify_gemma1_with_custom_loader(checkpoint_dir: str) -> bool:
+  """Verifies the reauthored Gemma1 model with a custom loader."""
+  weight_filename = "gemma-2b-it.ckpt"
+  checkpoint_path = os.path.join(checkpoint_dir, weight_filename)
+  custom_loader = loader.get_custom_loader(checkpoint_path)
+  reauthored_model = gemma1.build_2b_model(checkpoint_path, custom_loader)
+  return verify_reauthored_gemma_model(
+      checkpoint=checkpoint_dir,
+      variant="2b",
+      reauthored_model=reauthored_model,
+      weight_filename=weight_filename,
+      custom_loader=custom_loader,
+      generate_prompts=["What is the meaning of life?"],
+      forward_input_ids=[[1, 2, 3, 4]],
+      max_new_tokens=30,
+  )
+
+
+def verify_gemma2_with_custom_loader(checkpoint_dir: str) -> bool:
+  """Verifies the reauthored Gemma2 model with a custom loader."""
+  return verify_gemma2(
+      checkpoint_dir=checkpoint_dir,
+      weight_filename="model.ckpt",
+      prompts=["What is the meaning of life?"],
+      max_new_tokens=30,
+      mask_as_input=True,
+      custom_loader=loader.get_custom_loader("", checkpoint_format="pt"),
+  )

ai_edge_torch/generative/examples/gemma3/decoder.py

@@ -149,8 +149,12 @@ class Decoder(nn.Module):
           cache_len=attention_mask.shape[-1],
           sliding_window_size=sliding_window_size,
       )
-      # Combine masks using logical AND (min in this case).
-      combined_mask = torch.min(attention_mask, sliding_mask)
+      # Expand sliding_mask to match attention_mask's dimensions
+      # (e.g., [B, 1, seq_len, cache_len]).
+      # Assuming the head dimension is dim 1 for attention_mask.
+      expanded_sliding_mask = sliding_mask.unsqueeze(1)
+      # Combine masks using logical AND (min ensures -inf propagates).
+      combined_mask = torch.min(attention_mask, expanded_sliding_mask)
       return combined_mask
     return attention_mask
 
@@ -161,9 +165,9 @@ class Decoder(nn.Module):
       sliding_window_size: int,
   ) -> torch.Tensor:
     """Creates mask for sliding window attention (PyTorch)."""
-    cache_positions = torch.tensor(
-        [i for i in range(cache_len)], dtype=torch.int32
-    )
+    # Use torch.arange to create a tensor with a range of integers in a
+    # Dynamo-friendly way.
+    cache_positions = torch.arange(cache_len, dtype=torch.int32)
     cache_positions = cache_positions.view(1, 1, -1)  # [1, 1, cache_len]
     segment_pos_expanded = segment_pos.clone().unsqueeze(-1)  # [B, seq_len, 1]
 
@@ -329,10 +333,7 @@ def get_decoder_config_1b(kv_cache_max_len: int = 2048) -> cfg.ModelConfig:
     The model config for a Gemma 1B model.
   """
   norm_config = cfg.NormalizationConfig(
-      type=cfg.NormalizationType.RMS_NORM,
-      epsilon=1e-6,
-      zero_centered=True,
-      enable_hlfb=True,
+      type=cfg.NormalizationType.RMS_NORM, epsilon=1e-6, zero_centered=True,
   )
   ff_config = cfg.FeedForwardConfig(
       type=cfg.FeedForwardType.GATED,
@@ -379,7 +380,6 @@ def get_decoder_config_1b(kv_cache_max_len: int = 2048) -> 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=None,
   )
   return config

ai_edge_torch/generative/examples/gemma3/gemma3.py

@@ -158,9 +158,7 @@ def get_fake_model_config(**kwargs) -> Gemma3MMConfig:
       image_projection_scale=128**0.5,
       image_projection_use_bias=False,
       mm_norm_config=cfg.NormalizationConfig(
-          type=cfg.NormalizationType.LAYER_NORM,
-          epsilon=1e-6,
-          enable_hlfb=True,
+          type=cfg.NormalizationType.LAYER_NORM, epsilon=1e-6
       ),
       mm_extra_tokens=32,
   )

ai_edge_torch/generative/examples/gemma3/image_encoder.py

@@ -98,9 +98,7 @@ def get_image_encoder_config() -> cfg.ModelConfig:
       output_proj_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
   )
   ff_config = cfg.FeedForwardConfig(
       type=cfg.FeedForwardType.SEQUENTIAL,
@@ -123,7 +121,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,
       num_mm_tokens_per_image=256,
   )
   return config

ai_edge_torch/generative/examples/gemma3/verify_util.py

@@ -22,6 +22,7 @@ from typing import Callable, Dict, List, Optional, Tuple
 from ai_edge_torch.generative.examples.gemma3 import gemma3
 from ai_edge_torch.generative.layers import kv_cache as kv_utils
 import ai_edge_torch.generative.layers.attention_utils as attn_utils
+from ai_edge_torch.generative.utilities import loader
 from ai_edge_torch.generative.utilities import verifier
 from gemma import config as gemma_config
 from gemma import model as gemma_model
@@ -260,3 +261,15 @@ def verify_gemma3(
       custom_loader=custom_loader,
       atol=1e-04,
   )
+
+
+def verify_gemma3_with_custom_loader(checkpoint: str) -> bool:
+  """Verifies the reauthored Gemma3 model with a custom loader."""
+  return verify_gemma3(
+      checkpoint=checkpoint,
+      prompts=["What is the meaning of life?"],
+      max_new_tokens=30,
+      variant="1b",
+      weight_filename="model.ckpt",
+      custom_loader=loader.get_custom_loader("", checkpoint_format="pt"),
+  )

ai_edge_torch/generative/examples/hammer/hammer.py

@@ -45,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,
@@ -63,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