ai-edge-torch-nightly 0.3.0.dev20250123__py3-none-any.whl → 0.3.0.dev20250125__py3-none-any.whl

ai_edge_torch/_config.py

@@ -65,5 +65,14 @@ class _Config:
   def enable_group_norm_composite(self, value: bool):
     os.environ["ENABLE_GROUP_NORM_COMPOSITE"] = "y" if value else "n"
 
+  @property
+  def layout_optimize_partitioner(self) -> str:
+    """The algorithm to use for layout optimization."""
+    return os.environ.get("AIEDGETORCH_LAYOUT_OPTIMIZE_PARTITIONER", "DEFAULT")
+
+  @layout_optimize_partitioner.setter
+  def layout_optimize_partitioner(self, value: str):
+    os.environ["AIEDGETORCH_LAYOUT_OPTIMIZE_PARTITIONER"] = str(value).upper()
+
 
 config = _Config()

ai_edge_torch/_convert/fx_passes/optimize_layout_transposes_pass/layout_check.py

@@ -201,8 +201,14 @@ def _aten_group_norm_checker(node):
   return NHWCable(can_be=can_be, must_be=must_be)
 
 
-@nhwcable_node_checkers.register(aten.native_group_norm)
+@nhwcable_node_checkers.register(aten.native_group_norm.default)
 def _aten_native_group_norm_checker(node):
+  # aten.group_norm is removed from the decomp table, so aten.native_group_norm
+  # should never exist in the graph. However, torch 2.5.1 could ignore the
+  # decomp table updates, so still add this native_group_norm checker and
+  # rewriter to be safe.
+  # The checker and rewriter are the same as the ones for aten.group_norm.
+
   val = node.meta.get("val")
   if (
       not isinstance(val, (list, tuple))
@@ -210,13 +216,10 @@ def _aten_native_group_norm_checker(node):
       or not hasattr(val[0], "shape")
   ):
     return NHWCable(can_be=False, must_be=False)
-  if len(node.args) >= 3 and (
-      node.args[1] is not None or node.args[2] is not None
-  ):
-    # Disable NHWC rewriter due to precision issue with weight and bias.
-    # TODO(b/354780253): Re-enable NHWC rewriter with proper lowering.
-    return NHWCable(can_be=False, must_be=False)
-  return NHWCable(can_be=len(val[0].shape) == 4, must_be=False)
+
+  can_be = len(val[0].shape) == 4
+  must_be = can_be and ai_edge_torch.config.enable_group_norm_composite
+  return NHWCable(can_be=can_be, must_be=must_be)
 
 
 # ==== Ops must be NCHW

ai_edge_torch/_convert/fx_passes/optimize_layout_transposes_pass/layout_rewrite.py

@@ -391,34 +391,32 @@ def _aten_native_group_norm(node):
       eps: float,
       **kwargs,
   ):
-    input_reshaped = torch.reshape(
-        input,
-        [
-            batch_size,
-            flattened_inner_size,
-            num_groups,
-            num_channels // num_groups,
-        ],
-    )
-    reduction_dims = [1, 3]
-
-    biased_var, mean = torch.var_mean(
-        input_reshaped, dim=reduction_dims, unbiased=False, keepdim=True
+    is_composite_supported = (
+        ai_edge_torch.config.enable_group_norm_composite
+        and weight is not None
+        and bias is not None
     )
-    rstd = torch.rsqrt(biased_var + eps)
-
-    out = (input_reshaped - mean) * rstd
-    out = torch.reshape(out, input.shape)
 
-    if weight is not None:
-      out = out * weight
-    if bias is not None:
-      out = out + bias
+    builder = None
+    if is_composite_supported:
+      builder = StableHLOCompositeBuilder(
+          name="odml.group_norm",
+          attr={
+              "num_groups": num_groups,
+              "epsilon": eps,
+              "reduction_axes": [3],
+              "channel_axis": 3,
+          },
+      )
+      input, weight, bias = builder.mark_inputs(input, weight, bias)
 
-    mean = torch.squeeze(mean, reduction_dims)
-    rstd = torch.squeeze(rstd, reduction_dims)
+    input = utils.tensor_to_nchw(input)
+    output = aten.group_norm.default(input, num_groups, weight, bias, eps=eps)
+    output = utils.tensor_to_nhwc(output)
 
-    return out, mean, rstd
+    if builder is not None:
+      output = builder.mark_outputs(output)
+    return (output, None, None)
 
   node.target = native_group_norm
 

ai_edge_torch/_convert/fx_passes/optimize_layout_transposes_pass/pass_body.py

@@ -18,6 +18,7 @@ import operator
 import os
 from typing import Union
 
+import ai_edge_torch
 from ai_edge_torch import fx_infra
 from ai_edge_torch._convert.fx_passes.optimize_layout_transposes_pass import layout_check  # NOQA
 from ai_edge_torch._convert.fx_passes.optimize_layout_transposes_pass import layout_mark  # NOQA
@@ -261,10 +262,8 @@ class OptimizeLayoutTransposesPass(fx_infra.ExportedProgramPassBase):
     self.mark_const_nodes(exported_program)
 
     graph_module = exported_program.graph_module
-    partitioner = os.environ.get(
-        "AIEDGETORCH_LAYOUT_OPTIMIZE_PARTITIONER", None
-    )
-    if partitioner == "MINCUT":
+    partitioner = ai_edge_torch.config.layout_optimize_partitioner
+    if partitioner in ("MINCUT", "OPTIMAL"):
       graph_module = layout_partitioners.min_cut.partition(graph_module)
     elif partitioner == "GREEDY":
       graph_module = layout_partitioners.greedy.partition(graph_module)

ai_edge_torch/generative/examples/deepseek/__init__.py

@@ -0,0 +1,14 @@
+# 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.
+# ==============================================================================

ai_edge_torch/generative/examples/deepseek/convert_to_tflite.py

@@ -0,0 +1,80 @@
+# 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.
+# ==============================================================================
+
+"""Example of converting DeepSeek R1 distilled models to tflite model."""
+
+import os
+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 converter
+from ai_edge_torch.generative.utilities.model_builder import ExportConfig
+
+_CHECKPOINT_PATH = flags.DEFINE_string(
+    'checkpoint_path',
+    os.path.join(pathlib.Path.home(), 'Downloads/llm_data/deepseek'),
+    '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',
+    'deepseek',
+    'The prefix of the output tflite model name.',
+)
+_PREFILL_SEQ_LENS = flags.DEFINE_multi_integer(
+    'prefill_seq_lens',
+    (8, 64, 128, 256, 512, 1024),
+    'List of the maximum sizes of prefill input tensors.',
+)
+_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.',
+)
+_LORA_RANKS = flags.DEFINE_multi_integer(
+    'lora_ranks',
+    None,
+    'If set, the model will be converted with the provided list of LoRA ranks.',
+)
+
+
+def main(_):
+  pytorch_model = deepseek.build_model(
+      _CHECKPOINT_PATH.value, kv_cache_max_len=_KV_CACHE_MAX_LEN.value
+  )
+  converter.convert_to_tflite(
+      pytorch_model,
+      output_path=_OUTPUT_PATH.value,
+      output_name_prefix=_OUTPUT_NAME_PREFIX.value,
+      prefill_seq_len=_PREFILL_SEQ_LENS.value,
+      quantize=_QUANTIZE.value,
+      lora_ranks=_LORA_RANKS.value,
+      export_config=ExportConfig(),
+  )
+
+
+if __name__ == '__main__':
+  app.run(main)

ai_edge_torch/generative/examples/deepseek/deepseek.py

@@ -0,0 +1,92 @@
+# 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.
+# ==============================================================================
+
+"""Example of building DeepSeek R1 distilled models."""
+
+import ai_edge_torch.generative.layers.model_config as cfg
+from ai_edge_torch.generative.utilities import model_builder
+from torch import nn
+
+TENSOR_NAMES = model_builder.TENSOR_NAMES_WITH_SEPARATE_LM_HEAD
+
+
+class DeepSeekDistillQwen(model_builder.DecoderOnlyModel):
+  """A DeepSeek distilled model based on Qwen."""
+  pass
+
+
+def get_model_config(kv_cache_max_len: int = 1024) -> cfg.ModelConfig:
+  """Returns the model config for a Qwen 2.5 3B 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=12,
+      head_dim=128,
+      num_query_groups=2,
+      rotary_base=10000,
+      rotary_percentage=1.0,
+      qkv_use_bias=True,
+  )
+  ff_config = cfg.FeedForwardConfig(
+      type=cfg.FeedForwardType.GATED,
+      activation=cfg.ActivationConfig(cfg.ActivationType.SILU),
+      intermediate_size=8960,
+  )
+  norm_config = cfg.NormalizationConfig(
+      type=cfg.NormalizationType.RMS_NORM,
+      epsilon=1e-06,
+  )
+  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=151936,
+      num_layers=28,
+      max_seq_len=4096,
+      embedding_dim=1536,
+      kv_cache_max_len=kv_cache_max_len,
+      block_configs=block_config,
+      final_norm_config=norm_config,
+      lm_head_share_weight_with_embedding=False,
+      enable_hlfb=True,
+  )
+  return config
+
+
+def get_fake_model_config(**kwargs) -> cfg.ModelConfig:
+  config = get_model_config(**kwargs)
+  config.vocab_size = 128
+  config.num_layers = 2
+  # DeepSeek-R1-Distill-Qwen 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:
+  return model_builder.build_decoder_only_model(
+      checkpoint_path=checkpoint_path,
+      config=get_model_config(**kwargs),
+      tensor_names=TENSOR_NAMES,
+      model_class=DeepSeekDistillQwen,
+  )

ai_edge_torch/generative/examples/deepseek/verify.py

@@ -0,0 +1,70 @@
+# 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.
+# ==============================================================================
+
+"""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
+
+
+_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 = "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(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,
+      max_new_tokens=_MAX_NEW_TOKENS.value,
+      atol=1e-04,
+  )
+
+
+if __name__ == "__main__":
+  app.run(main)

ai_edge_torch/generative/examples/stable_diffusion/convert_to_tflite.py

@@ -85,6 +85,7 @@ def convert_stable_diffusion_to_tflite(
       clip.TENSOR_NAMES,
   )
   loader.load(clip_model, strict=False)
+  clip_model.eval()
 
   diffusion_model = diffusion.Diffusion(
       diffusion.get_model_config(batch_size=2, device_type=_DEVICE_TYPE.value)
@@ -93,6 +94,7 @@ def convert_stable_diffusion_to_tflite(
       diffusion_ckpt_path, diffusion.TENSOR_NAMES
   )
   diffusion_loader.load(diffusion_model, strict=False)
+  diffusion_model.eval()
 
   decoder_model = decoder.Decoder(
       decoder.get_model_config(device_type=_DEVICE_TYPE.value)
@@ -101,6 +103,7 @@ def convert_stable_diffusion_to_tflite(
       decoder_ckpt_path, decoder.TENSOR_NAMES
   )
   decoder_loader.load(decoder_model, strict=False)
+  decoder_model.eval()
 
   # TODO(yichunk): enable image encoder conversion
   # if encoder_ckpt_path is not None:

ai_edge_torch/generative/layers/experimental/__init__.py

@@ -0,0 +1,14 @@
+# 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.
+#