ai-edge-torch-nightly 0.3.0.dev20250114__py3-none-any.whl

ai_edge_torch/generative/examples/openelm/openelm.py

@@ -0,0 +1,127 @@
+# 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 an OpenELM model."""
+
+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
+from torch import nn
+
+TENSOR_NAMES = loading_utils.ModelLoader.TensorNames(
+    ff_up_proj="transformer.layers.{}.ffn.proj_1",
+    ff_down_proj="transformer.layers.{}.ffn.proj_2",
+    attn_fused_qkv_proj="transformer.layers.{}.attn.qkv_proj",
+    attn_query_norm="transformer.layers.{}.attn.q_norm",
+    attn_key_norm="transformer.layers.{}.attn.k_norm",
+    attn_output_proj="transformer.layers.{}.attn.out_proj",
+    pre_attn_norm="transformer.layers.{}.attn_norm",
+    pre_ff_norm="transformer.layers.{}.ffn_norm",
+    embedding="transformer.token_embeddings",
+    final_norm="transformer.norm",
+    lm_head=None,
+)
+
+
+class OpenELM(model_builder.DecoderOnlyModel):
+  """An OpenELM model built from the Edge Generative API layers."""
+  pass
+
+
+def get_model_config(kv_cache_max_len: int = 1024) -> cfg.ModelConfig:
+  """Returns the model config for an OpenELM model.
+
+  Args:
+    kv_cache_max_len (int): The maximum sequence length of the KV cache. Default
+      is 1024.
+
+  Returns:
+    The model config for an OpenELM model.
+  """
+  norm_config = cfg.NormalizationConfig(
+      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
+
+  def make_divisible(v, d):
+    """Ensures that all layers have a channel number that is divisible by d."""
+    new_v = int(v + d / 2) // d * d
+    # Make sure that round down does not go down by more than 10%.
+    if new_v < 0.9 * v:
+      new_v += d
+    return new_v
+
+  # The way to get intermediate size is from
+  # https://huggingface.co/apple/OpenELM-3B/blob/main/modeling_openelm.py
+  def get_intermediate_size(idx: int) -> int:
+    return make_divisible((0.5 + 0.1 * idx) * 3072, 256)
+
+  def get_block_config(idx: int) -> cfg.TransformerBlockConfig:
+    return cfg.TransformerBlockConfig(
+        attn_config=cfg.AttentionConfig(
+            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,
+            key_norm_config=norm_config,
+        ),
+        ff_config=cfg.FeedForwardConfig(
+            type=cfg.FeedForwardType.SEQUENTIAL,
+            activation=cfg.ActivationConfig(cfg.ActivationType.SILU_GLU),
+            intermediate_size=get_intermediate_size(idx),
+            pre_ff_norm_config=norm_config,
+        ),
+        pre_attention_norm_config=norm_config,
+    )
+
+  num_layers = 36
+  config = cfg.ModelConfig(
+      vocab_size=32000,
+      num_layers=num_layers,
+      max_seq_len=2048,
+      embedding_dim=3072,
+      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
+
+
+def get_fake_model_config(kv_cache_max_len: int = 128) -> cfg.ModelConfig:
+  config = get_model_config(kv_cache_max_len)
+  config.vocab_size = 128
+  config.num_layers = 2
+  config.max_seq_len = 2 * kv_cache_max_len
+  config.embedding_dim = 128
+  config.block_configs = config.block_configs[: config.num_layers]
+  for block_config in config.block_configs:
+    block_config.attn_config.num_heads = 3
+    block_config.attn_config.head_dim = 64
+    block_config.ff_config.intermediate_size = 128
+  return config
+
+
+def build_model(checkpoint_path: str, **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,
+  )

ai_edge_torch/generative/examples/openelm/verify.py

@@ -0,0 +1,71 @@
+# 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 OpenELM-3B model."""
+
+import logging
+import pathlib
+from absl import app
+from absl import flags
+from ai_edge_torch.generative.examples.openelm import openelm
+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 = "apple/OpenELM-3B"
+  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 = openelm.build_model(reauthored_checkpoint)
+
+  tokenizer_checkpoint = "meta-llama/Llama-2-7b-hf"
+  logging.info("Loading the tokenizer from: %s", tokenizer_checkpoint)
+  tokenizer = transformers.AutoTokenizer.from_pretrained(tokenizer_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,
+  )
+
+
+if __name__ == "__main__":
+  app.run(main)

ai_edge_torch/generative/examples/paligemma/__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/paligemma/convert_to_tflite.py

@@ -0,0 +1,95 @@
+# 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 a PaliGemma model to multi-signature tflite model.
+
+DISCLAIMER: It works only with ODML Torch conversion backend. Refer to
+https://github.com/google-ai-edge/ai-edge-torch/blob/main/docs/pytorch_converter/README.md#use-odml-torch-conversion-backend-experimental.
+"""
+
+import os
+import pathlib
+
+from absl import app
+from absl import flags
+from ai_edge_torch.generative.examples.paligemma import paligemma
+from ai_edge_torch.generative.utilities import converter
+from ai_edge_torch.generative.utilities.model_builder import ExportConfig
+import torch
+
+_VERSION = flags.DEFINE_enum(
+    'version',
+    '2',
+    ['1', '2'],
+    'The version of PaliGemma model to verify.',
+)
+_CHECKPOINT_PATH = flags.DEFINE_string(
+    'checkpoint_path',
+    os.path.join(pathlib.Path.home(), 'Downloads/llm_data/paligemma2-3b-224'),
+    'The path to the model checkpoint, or directory holding the checkpoint.',
+)
+_OUTPUT_PATH = flags.DEFINE_string(
+    'output_path',
+    '/tmp/',
+    'The path to export the tflite model.',
+)
+_OUTPUT_NAME_PREFIX = flags.DEFINE_string(
+    'output_name_prefix',
+    'paligemma',
+    'The prefix of the output tflite model name.',
+)
+_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.',
+)
+_PIXEL_VALUES_SIZE = flags.DEFINE_multi_integer(
+    'pixel_values_size',
+    [3, 224, 224],
+    'The size of prefill pixel values except the batch dimension.',
+)
+_QUANTIZE = flags.DEFINE_bool(
+    'quantize',
+    True,
+    'Whether the model should be quantized.',
+)
+
+
+def main(_):
+  pytorch_model = paligemma.build_model(
+      _CHECKPOINT_PATH.value,
+      version=int(_VERSION.value),
+      kv_cache_max_len=_KV_CACHE_MAX_LEN.value,
+  )
+
+  converter.convert_to_tflite(
+      pytorch_model,
+      output_path=_OUTPUT_PATH.value,
+      output_name_prefix=f'{_OUTPUT_NAME_PREFIX.value}_{_VERSION.value}',
+      prefill_seq_len=_PREFILL_SEQ_LEN.value,
+      pixel_values_size=torch.Size(_PIXEL_VALUES_SIZE.value),
+      quantize=_QUANTIZE.value,
+      config=pytorch_model.config.decoder_config,
+      export_config=ExportConfig(),
+  )
+
+
+if __name__ == '__main__':
+  app.run(main)

ai_edge_torch/generative/examples/paligemma/decoder.py

@@ -0,0 +1,151 @@
+# 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 a decoder of PaliGemma 3B model which is Gemma1."""
+
+from typing import Optional
+
+from ai_edge_torch.generative.layers import kv_cache as kv_utils
+import ai_edge_torch.generative.layers.model_config as cfg
+import ai_edge_torch.generative.layers.rotary_position_embedding as rotary_pos_emb
+from ai_edge_torch.generative.utilities import model_builder
+import ai_edge_torch.generative.utilities.loader as loading_utils
+import torch
+
+TENSOR_NAMES = loading_utils.ModelLoader.TensorNames(
+    ff_up_proj="language_model.model.layers.{}.mlp.up_proj",
+    ff_down_proj="language_model.model.layers.{}.mlp.down_proj",
+    ff_gate_proj="language_model.model.layers.{}.mlp.gate_proj",
+    attn_query_proj="language_model.model.layers.{}.self_attn.q_proj",
+    attn_key_proj="language_model.model.layers.{}.self_attn.k_proj",
+    attn_value_proj="language_model.model.layers.{}.self_attn.v_proj",
+    attn_output_proj="language_model.model.layers.{}.self_attn.o_proj",
+    pre_attn_norm="language_model.model.layers.{}.input_layernorm",
+    post_attn_norm="language_model.model.layers.{}.post_attention_layernorm",
+    embedding="language_model.model.embed_tokens",
+    final_norm="language_model.model.norm",
+    lm_head=None,
+)
+
+
+class Decoder(model_builder.DecoderOnlyModel):
+  """A decoder of PaliGemma 3B model which is Gemma1.
+
+  Besides a tensor of text token IDs, forward() can also take a tensor of
+  embeddings which may include text or image or both.
+  """
+
+  @torch.inference_mode
+  def forward(
+      self,
+      tokens: torch.Tensor,
+      input_pos: torch.Tensor,
+      kv_cache: kv_utils.KVCache,
+      input_embeds: torch.Tensor = None,
+      mask: Optional[torch.Tensor] = None,
+      export_config: Optional[model_builder.ExportConfig] = None,
+      called_by_generate: bool = True,
+  ) -> dict[torch.Tensor, kv_utils.KVCache]:
+    if input_embeds is None:
+      return super().forward(tokens, input_pos, kv_cache)
+
+    assert input_embeds is not None
+
+    repo_pos = input_pos + 1  # PaliGemma position is 1-based.
+    # ROPE parameters for all attn_configs are the same. Take the first one.
+    attn_config = self.config.block_config(0).attn_config
+    n_elem = int(attn_config.rotary_percentage * attn_config.head_dim)
+    rope = rotary_pos_emb.build_rope(repo_pos, n_elem, attn_config.rotary_base)
+
+    # The first part of input_embeds are image embeddings. Diagonal causal mask
+    # doesn't work here.
+    embeds_len = input_embeds.shape[1]
+    if mask is None:
+      mask = torch.zeros(embeds_len, self.config.kv_cache_max)
+      mask[:, embeds_len:] = float("-inf")
+
+    return self._forward_with_embeds(
+        input_embeds, rope, mask, input_pos, kv_cache
+    )
+
+
+def get_decoder_config(kv_cache_max_len: int = 1024) -> cfg.ModelConfig:
+  """Returns the model config for the decoder of a PaliGemma 3B model.
+
+  Args:
+    kv_cache_max_len (int): The maximum sequence length of the KV cache. Default
+      is 1024.
+
+  Returns:
+    The model config for the decoder of a PaliGemma 3B model.
+  """
+  attn_config = cfg.AttentionConfig(
+      num_heads=8,
+      head_dim=256,
+      num_query_groups=1,
+      rotary_base=10000,
+      rotary_percentage=1.0,
+  )
+  ff_config = cfg.FeedForwardConfig(
+      type=cfg.FeedForwardType.GATED,
+      activation=cfg.ActivationConfig(cfg.ActivationType.GELU_TANH),
+      intermediate_size=16384,
+  )
+  norm_config = cfg.NormalizationConfig(
+      type=cfg.NormalizationType.RMS_NORM,
+      epsilon=1e-6,
+      zero_centered=True,
+  )
+  block_config = cfg.TransformerBlockConfig(
+      attn_config=attn_config,
+      ff_config=ff_config,
+      pre_attention_norm_config=norm_config,
+      post_attention_norm_config=norm_config,
+  )
+  embedding_dim = 2048
+  config = cfg.ModelConfig(
+      vocab_size=257216,
+      num_layers=18,
+      max_seq_len=8192,
+      embedding_dim=embedding_dim,
+      embedding_scale=embedding_dim**0.5,
+      kv_cache_max_len=kv_cache_max_len,
+      block_configs=block_config,
+      final_norm_config=norm_config,
+      lm_head_use_bias=False,
+      enable_hlfb=True,
+  )
+  return config
+
+
+def get_fake_decoder_config(kv_cache_max_len: int = 128) -> cfg.ModelConfig:
+  config = get_decoder_config(kv_cache_max_len)
+  # PaliGemma decoder has only one block config.
+  config.block_config(0).ff_config.intermediate_size = 128
+  config.vocab_size = 128
+  config.num_layers = 2
+  config.max_seq_len = 2 * kv_cache_max_len
+  config.embedding_dim = 128
+  config.embedding_scale = 128**0.5
+  return config
+
+
+def build_decoder(checkpoint_path: str, **kwargs) -> torch.nn.Module:
+  return model_builder.build_decoder_only_model(
+      checkpoint_path=checkpoint_path,
+      config=get_decoder_config(**kwargs),
+      tensor_names=TENSOR_NAMES,
+      model_class=Decoder,
+  )

ai_edge_torch/generative/examples/paligemma/decoder2.py

@@ -0,0 +1,177 @@
+# 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 a decoder of PaliGemma2 3B model which is Gemma2."""
+
+from typing import Optional
+
+from ai_edge_torch.generative.examples.gemma import gemma2
+from ai_edge_torch.generative.layers import kv_cache as kv_utils
+import ai_edge_torch.generative.layers.model_config as cfg
+import ai_edge_torch.generative.layers.rotary_position_embedding as rotary_pos_emb
+from ai_edge_torch.generative.utilities import model_builder
+import ai_edge_torch.generative.utilities.loader as loading_utils
+import torch
+
+TENSOR_NAMES = loading_utils.ModelLoader.TensorNames(
+    ff_up_proj="language_model.model.layers.{}.mlp.up_proj",
+    ff_down_proj="language_model.model.layers.{}.mlp.down_proj",
+    ff_gate_proj="language_model.model.layers.{}.mlp.gate_proj",
+    attn_query_proj="language_model.model.layers.{}.self_attn.q_proj",
+    attn_key_proj="language_model.model.layers.{}.self_attn.k_proj",
+    attn_value_proj="language_model.model.layers.{}.self_attn.v_proj",
+    attn_output_proj="language_model.model.layers.{}.self_attn.o_proj",
+    pre_attn_norm="language_model.model.layers.{}.input_layernorm",
+    post_attn_norm="language_model.model.layers.{}.post_attention_layernorm",
+    pre_ff_norm="language_model.model.layers.{}.pre_feedforward_layernorm",
+    post_ff_norm="language_model.model.layers.{}.post_feedforward_layernorm",
+    embedding="language_model.model.embed_tokens",
+    final_norm="language_model.model.norm",
+    lm_head=None,
+)
+
+
+class Decoder2(gemma2.Gemma2):
+  """A decoder of PaliGemma2 3B model which is Gemma2.
+
+  Besides a tensor of text token IDs, forward() can also take a tensor of
+  embeddings which may include text or image or both.
+  """
+
+  @torch.inference_mode
+  def forward(
+      self,
+      tokens: torch.Tensor,
+      input_pos: torch.Tensor,
+      kv_cache: kv_utils.KVCache,
+      input_embeds: torch.Tensor = None,
+      mask: Optional[torch.Tensor] = None,
+      export_config: Optional[model_builder.ExportConfig] = None,
+      called_by_generate: bool = True,
+  ) -> dict[torch.Tensor, kv_utils.KVCache]:
+    if input_embeds is None:
+      return super().forward(tokens, input_pos, kv_cache)
+
+    assert input_embeds is not None
+
+    repo_pos = input_pos + 1  # PaliGemma2 position is 1-based.
+    # ROPE parameters for all attn_configs are the same. Take the first one.
+    attn_config = self.config.block_config(0).attn_config
+    n_elem = int(attn_config.rotary_percentage * attn_config.head_dim)
+    rope = rotary_pos_emb.build_rope(repo_pos, n_elem, attn_config.rotary_base)
+
+    if mask is None:
+      if called_by_generate:
+        # PaliGemma2 generate() use a diagonal causal mask even with image embeds.
+        mask = [
+            self.get_attention_mask(
+                self.config.block_config(i).attn_config.attn_type, input_pos
+            )
+            for i in range(self.config.num_layers)
+        ]
+      else:
+        # By default, don't mask image embeds with a diagonal causal mask.
+        embeds_len = input_embeds.shape[1]
+        mask = torch.zeros(embeds_len, self.config.kv_cache_max)
+        mask[:, embeds_len:] = float("-inf")
+        mask = [mask] * self.config.num_layers
+
+    return self._forward_with_embeds(
+        input_embeds, rope, mask, input_pos, kv_cache, export_config
+    )
+
+
+def get_decoder2_config(kv_cache_max_len: int = 1024) -> cfg.ModelConfig:
+  """Returns the model config for the decoder of a PaliGemma 3B model.
+
+  Args:
+    kv_cache_max_len (int): The maximum sequence length of the KV cache. Default
+      is 1024.
+
+  Returns:
+    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,
+  )
+  ff_config = cfg.FeedForwardConfig(
+      type=cfg.FeedForwardType.GATED,
+      activation=cfg.ActivationConfig(cfg.ActivationType.GELU_TANH),
+      intermediate_size=9216,
+      pre_ff_norm_config=norm_config,
+      post_ff_norm_config=norm_config,
+  )
+
+  def get_block_config(idx: int) -> cfg.TransformerBlockConfig:
+    attn_config = cfg.AttentionConfig(
+        num_heads=8,
+        head_dim=256,
+        num_query_groups=4,
+        rotary_base=10000,
+        rotary_percentage=1.0,
+        logit_softcap=50.0,
+        sliding_window_size=4096,
+        attn_type=(
+            cfg.AttentionType.GLOBAL
+            if idx % 2 == 0
+            else cfg.AttentionType.LOCAL_SLIDING
+        ),
+    )
+    return cfg.TransformerBlockConfig(
+        attn_config=attn_config,
+        ff_config=ff_config,
+        pre_attention_norm_config=norm_config,
+        post_attention_norm_config=norm_config,
+    )
+
+  num_layers = 26
+  embedding_dim = 2304
+  config = cfg.ModelConfig(
+      vocab_size=257216,
+      num_layers=num_layers,
+      max_seq_len=8192,
+      embedding_dim=embedding_dim,
+      embedding_scale=embedding_dim**0.5,
+      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,
+      lm_head_use_bias=False,
+      enable_hlfb=True,
+      final_logit_softcap=30.0,
+  )
+  return config
+
+
+def get_fake_decoder2_config(kv_cache_max_len: int = 128) -> cfg.ModelConfig:
+  config = get_decoder2_config(kv_cache_max_len)
+  # PaliGemma2 decoder has only one block config.
+  config.block_config(0).ff_config.intermediate_size = 128
+  config.vocab_size = 128
+  config.num_layers = 2
+  config.max_seq_len = 2 * kv_cache_max_len
+  config.embedding_dim = 128
+  config.embedding_scale = 128**0.5
+  return config
+
+
+def build_decoder2(checkpoint_path: str, **kwargs) -> torch.nn.Module:
+  return model_builder.build_decoder_only_model(
+      checkpoint_path=checkpoint_path,
+      config=get_decoder2_config(**kwargs),
+      tensor_names=TENSOR_NAMES,
+      model_class=Decoder2,
+  )