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

ai_edge_torch/generative/utilities/verifier.py

@@ -0,0 +1,335 @@
+# 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.
+# ==============================================================================
+
+"""Common utility functions to verify the reauthored models."""
+
+import logging
+from typing import Any, List, Optional
+
+from ai_edge_torch.generative.layers import kv_cache as kv_utils
+from ai_edge_torch.generative.utilities.model_builder import ExportConfig
+import torch
+
+
+class ModelWrapper(torch.nn.Module):
+  """A wrapper for the model to be verified.
+
+  It unifies the interface of forward() and generate() of models for the
+  verification to call.
+  """
+
+  def __init__(self, model: torch.nn.Module):
+    """Initializes the wrapper.
+
+    Args:
+      model (torch.nn.Module): The model which might have different interfaces
+        of forward() and generate(). It could be a model built from HuggingFace
+        transformers, a regular PyTorch model, or a model re-authored with
+        ai_edge_torch Generative API.
+    """
+    super().__init__()
+    self.model = model
+    self.export_config = ExportConfig(output_logits_on_prefill=True)
+
+  def forward(
+      self, tokens: torch.Tensor, pixel_values: torch.Tensor = None
+  ) -> torch.Tensor:
+    """Gets output logits by forwarding the input tokens.
+
+    Args:
+      tokens (torch.Tensor): The input tokens to forward. Its dimension is
+        expected to be (batch_size=1, kv_cache_max_len).
+
+    Returns:
+      The output logits.
+    """
+    raise NotImplementedError("forward() is not implemented.")
+
+  def generate(
+      self,
+      prompts: torch.Tensor,
+      max_new_tokens: int,
+      pixel_values: torch.Tensor = None,
+  ) -> torch.IntTensor:
+    """Returns the response token IDs to the given prompts tensor.
+
+    The maximum number of tokens to generate might be set by subclasses.
+
+    Args:
+      prompts (torch.Tensor): The input token IDs to generate with. Its shape is
+        expected to be (batch_size=1, input_ids_len).
+      max_new_tokens (int): The maximum number of response token IDs to
+        generate.
+
+    Returns:
+      The tensor of response token IDs with shape of (batch_size=1,
+      response_ids_len).
+    """
+    raise NotImplementedError("generate() is not implemented.")
+
+
+class ReauthoredModelWrapper(ModelWrapper):
+  """A wrapper for the model reauthored with ai_edge_torch Generative API."""
+
+  def _init_kv_cache(self):
+    """Returns an initialized KV cache."""
+    return kv_utils.KVCache.from_model_config(self.model.config)
+
+  def _get_extra_args_for_forward(self) -> dict[str, Any]:
+    """Returns extra arguments for the forward() method."""
+    return {}
+
+  def _forward_with_kv_cache(
+      self,
+      tokens: torch.Tensor,
+      input_pos: torch.Tensor,
+      kv_cache: kv_utils.KVCache,
+      pixel_values: torch.Tensor,
+  ) -> tuple[torch.Tensor, kv_utils.KVCache]:
+    """Forwards the model and updates an external KV cache.
+
+    Args:
+      tokens (torch.Tensor): The input tokens to forward.
+      input_pos (torch.Tensor): The input positions to forward.
+      kv_cache (KVCache): The KV cache to forward.
+      pixel_values (torch.Tensor): The input pixel values to forward.
+
+    Returns:
+      The output logits and the updated KV cache.
+    """
+    extra_args = self._get_extra_args_for_forward()
+    if self.export_config is not None:
+      # Verification requires logit outputs on prefill for comparison.
+      if not self.export_config.output_logits_on_prefill:
+        raise ValueError("Verifier requires logit output on prefill.")
+      extra_args["export_config"] = self.export_config
+    if pixel_values is not None:
+      extra_args["pixel_values"] = pixel_values
+    output = self.model.forward(tokens, input_pos, kv_cache, **extra_args)
+    return output["logits"], output["kv_cache"]
+
+  def forward(
+      self, tokens: torch.Tensor, pixel_values: torch.Tensor = None
+  ) -> torch.Tensor:
+    input_pos = torch.arange(0, tokens.shape[1], dtype=torch.int)
+    logits, _ = self._forward_with_kv_cache(
+        tokens, input_pos, self._init_kv_cache(), pixel_values
+    )
+    return logits
+
+  def generate(
+      self,
+      prompts: torch.Tensor,
+      max_new_tokens: int,
+      pixel_values: torch.Tensor = None,
+      eos_token_id: Optional[int] = None,
+  ) -> torch.IntTensor:
+    input_ids = prompts[0].int().tolist()
+    tokens = torch.tensor([input_ids])
+    input_pos = torch.arange(0, tokens.shape[1], dtype=torch.int)
+    kv_cache = self._init_kv_cache()
+    for _ in range(max_new_tokens):
+      logits, kv_cache = self._forward_with_kv_cache(
+          tokens, input_pos, kv_cache, pixel_values
+      )
+      generated_token = logits[0][-1].argmax().item()
+      input_ids.append(generated_token)
+      if eos_token_id is not None and generated_token == eos_token_id:
+        break
+      tokens = torch.tensor([[generated_token]])
+      input_pos = torch.tensor([len(input_ids) - 1])
+      pixel_values = None  # Pass only for the first time.
+    return torch.tensor([input_ids])
+
+
+class TokenizerWrapper(torch.nn.Module):
+  """A wrapper for the tokenizer used for verification."""
+
+  def __init__(self, tokenizer: torch.nn.Module):
+    """Initializes the wrapper.
+
+    Args:
+      tokenizer (torch.nn.Module): The tokenizer to wrap.
+    """
+    super().__init__()
+    self.tokenizer = tokenizer
+
+  def encode(self, prompts: str) -> torch.Tensor:
+    """Encodes the prompts to token IDs."""
+    return self.tokenizer.encode(prompts, return_tensors="pt")
+
+  def decode(self, token_ids: torch.Tensor) -> str:
+    """Decodes the token IDs to a string."""
+    return self.tokenizer.decode(token_ids)
+
+
+def verify_with_input_ids(
+    original_model: ModelWrapper,
+    reauthored_model: ReauthoredModelWrapper,
+    input_ids: List[int],
+    kv_cache_max_len: int = 1024,
+    rtol: float = 1e-05,
+    atol: float = 1e-05,
+):
+  """Verifies if the model reauthored generates the same output of the oringal.
+
+  It compares only one outputs from the original and the reauthored model.
+
+  Args:
+    original_model (ModelWrapper): The original model.
+    reauthored_model (ReauthoredModelWrapper): The model reauthored with
+      ai_edge_torch Generative API.
+    input_ids (List[int]): The input token IDs to forward with.
+    kv_cache_max_len (int): The maximum sequence length of the KV cache.
+    rtol (float): The relative tolerance for the comparison.
+    atol (float): The absolute tolerance for the comparison.
+
+  Raises:
+    AssertError if the model reauthored fails to generate the same output of the
+      original.
+  """
+  tokens = torch.full((1, kv_cache_max_len), 0, dtype=torch.int, device="cpu")
+  tokens[0, : len(input_ids)] = torch.tensor([input_ids]).int()
+
+  logging.info("Forwarding the original model...")
+  outputs_original = original_model.forward(tokens)
+  logits_original = outputs_original[0, len(input_ids) - 1, :]
+  logging.info("logits_original: %s", logits_original)
+
+  logging.info("Forwarding the reauthored model...")
+  outputs_reauthored = reauthored_model.forward(tokens)
+  logits_reauthored = outputs_reauthored[0, len(input_ids) - 1, :]
+  logging.info("logits_reauthored: %s", logits_reauthored)
+
+  assert torch.allclose(
+      logits_original, logits_reauthored, rtol=rtol, atol=atol
+  )
+
+
+def verify_model_with_prompts(
+    original_model: ModelWrapper,
+    reauthored_model: ReauthoredModelWrapper,
+    tokenizer: TokenizerWrapper,
+    prompts: str,
+    max_new_tokens: int,
+):
+  """Verifies if the model reauthored generates the same answer of the oringal.
+
+  It compares an answer, i.e. multiple continuous outputs generated by the
+  original and the reauthored model.
+
+  Args:
+    original_model (ModelWrapper): The original model.
+    reauthored_model (ReauthoredModelWrapper): The model reauthored with
+      ai_edge_torch Generative API.
+    tokenizer (TokenizerWrapper): The tokenizer.
+    prompts (str): The input prompts to generate answers.
+    max_new_tokens (int): The maximum number of new tokens to generate.
+
+  Raises:
+    AssertError if the model reauthored fails to generate the same answer of the
+      original.
+  """
+  prompt_tokens = tokenizer.encode(prompts)
+
+  logging.info("Generating answer with the original model...")
+  outputs_original = original_model.generate(prompt_tokens, max_new_tokens)
+  response_original = tokenizer.decode(outputs_original[0])
+  logging.info("outputs_from_original_model: [[%s]]", response_original)
+
+  logging.info("Generating answer with the reauthored model...")
+  outputs_reauthored = reauthored_model.generate(
+      prompt_tokens,
+      max_new_tokens,
+      eos_token_id=getattr(tokenizer.tokenizer, "eos_token_id", None),
+  )
+  response_reauthored = tokenizer.decode(outputs_reauthored[0])
+  logging.info("outputs from reauthored model: [[%s]]", response_reauthored)
+
+  assert response_original == response_reauthored
+
+
+def verify_reauthored_model(
+    original_model: ModelWrapper,
+    reauthored_model: ReauthoredModelWrapper,
+    tokenizer: TokenizerWrapper,
+    generate_prompts: List[str],
+    max_new_tokens: int = 30,
+    forward_input_ids: List[List[int]] = [[1, 2, 3, 4]],
+    rtol: float = 1e-05,
+    atol: float = 1e-05,
+    continue_on_failure: bool = False,
+) -> bool:
+  """Verifies the reauthored model against the original model.
+
+  It verifies the reauthored model with two methods:
+  1. It compares the output of the original and the reauthored model with an
+     arbitrary input.
+  2. It compares the answer generated by the original and the reauthored model
+     with a prompt.
+
+  It prints out "PASS" or "FAILED" to the console. It returns True if all
+  verification passes, False otherwise.
+
+  Args:
+    original_model (ModelWrapper): The original model.
+    reauthored_model (ReauthoredModelWrapper): The model reauthored with
+      ai_edge_torch Generative API.
+    tokenizer (TokenizerWrapper): The tokenizer.
+    generate_prompts (List[str]): List of the input prompts to generate answers.
+    max_new_tokens (int): The maximum number of new tokens to generate.
+    forward_input_ids (List[torch.Tensor]): List if ihe input token IDs to
+      forward with.
+    rtol (float): The relative tolerance for the comparison.
+    atol (float): The absolute tolerance for the comparison.
+    continue_on_failure (bool): If True, it continues to verify the next prompt
+      or input IDs even if a previous one fails.
+  """
+  failure_count = 0
+
+  for input_ids in forward_input_ids:
+    logging.info("Verifying the reauthored model with input IDs: %s", input_ids)
+    try:
+      verify_with_input_ids(
+          original_model, reauthored_model, input_ids, rtol=rtol, atol=atol
+      )
+    except AssertionError as e:
+      logging.error("*** FAILED *** verify with input IDs: %s", input_ids)
+      failure_count += 1
+      if not continue_on_failure:
+        return False
+    else:
+      logging.info("*** PASSED *** verify with input IDs: %s", input_ids)
+
+  for prompts in generate_prompts:
+    logging.info("Verifying the reauthored model with prompts: %s", prompts)
+    try:
+      verify_model_with_prompts(
+          original_model, reauthored_model, tokenizer, prompts, max_new_tokens
+      )
+    except AssertionError as e:
+      logging.error("*** FAILED *** verify with prompts: %s", prompts)
+      failure_count += 1
+      if not continue_on_failure:
+        return False
+    else:
+      logging.info("*** PASSED *** verify with prompts: %s", prompts)
+
+  if failure_count == 0:
+    logging.info("*** PASSED *** verify_reauthored_model")
+    return True
+  else:
+    logging.error("*** FAILED *** verify_reauthored_model")
+    return False

ai_edge_torch/hlfb/__init__.py

@@ -0,0 +1,16 @@
+# 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.
+# ==============================================================================
+
+from ai_edge_torch.lowertools import StableHLOCompositeBuilder

ai_edge_torch/hlfb/mark_pattern/__init__.py

@@ -0,0 +1,153 @@
+# 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.
+# ==============================================================================
+import copy
+from typing import Any
+import uuid
+
+from ai_edge_torch import lowertools
+from ai_edge_torch.hlfb.mark_pattern import fx_utils
+from ai_edge_torch.hlfb.mark_pattern import pattern as pattern_module
+import torch
+
+
+@torch._dynamo.assume_constant_result
+def _get_uuid() -> str:
+  return uuid.uuid4().hex
+
+
+# TODO: Move to a general fx utils file.
+def _prepose_placeholder_nodes(graph: torch.fx.Graph):
+  nodes = [node for node in graph.nodes if node.op == "placeholder"] + [
+      node for node in graph.nodes if node.op != "placeholder"
+  ]
+
+  for a, b in zip(nodes, nodes[1:]):
+    if a.next is not b:
+      a.append(b)
+  return graph
+
+
+def _insert_marker(
+    graph_module: torch.fx.GraphModule,
+    node: torch.fx.Node,
+    name: str,
+    pos: int,
+    id: str,
+    is_input: bool,
+    attr: dict[str, Any] = None,
+):
+  attr = lowertools.serialize_composite_attr(attr) if attr else None
+  with graph_module.graph.inserting_after(node):
+    new_node = graph_module.graph.call_function(
+        lowertools.mark_tensor_op,
+        args=(node,),
+        kwargs={
+            "name": name,
+            "pos": pos,
+            "id": id,
+            "is_input": is_input,
+            "attr": attr,
+        },
+    )
+
+  new_node.meta = node.meta
+  return new_node
+
+
+def mark_pattern(
+    graph_module: torch.fx.GraphModule,
+    pattern: pattern_module.Pattern,
+) -> torch.fx.GraphModule:
+  """Mark all existences of pattern graph in the GraphModule with fx pattern matching.
+
+  The marked subgraphs will be lowered in StableHLO composite ops.
+
+  Args:
+    graph_module (torch.fx.GraphModule): GraphModule to be matched and marked.
+    pattern (ai_edge_torch.hlfb.mark_pattern.Pattern): Pattern to match.
+
+  Returns:
+    The modified graph_module with additional marker ops in graph.
+  """
+  # Create a copy of graph_module and sanitize it for pattern matching.
+  graph_module_to_match = copy.deepcopy(graph_module)
+  for n, m in zip(graph_module.graph.nodes, graph_module_to_match.graph.nodes):
+    m.meta["ORIGINAL_NODE"] = n
+
+  # Sanitize graph_module to match in the same way as pattern's graph_module.
+  graph_module_to_match = fx_utils.remove_clone_ops(graph_module_to_match)
+
+  match_with_attrs = pattern.match(graph_module_to_match)
+
+  for match, attr in match_with_attrs:
+    match_id = _get_uuid()
+
+    # NOTE: Current graph rewriter (_insert_marker) does not work perfectly
+    # with continuous matches e.g. matching (a + b) on (w + x + y + z). The
+    # rewritten results may be undetermined with false negative - some
+    # matches may not be marked in the lowering, while the marked ones would
+    # always be correct.
+    # TODO(cnchan): completely support mark_pattern on continuous matches.
+    for i, pattern_input_node in enumerate(pattern.input_nodes):
+      input_node = match.nodes_map[pattern_input_node]
+      new_input_node = _insert_marker(
+          graph_module,
+          input_node.meta["ORIGINAL_NODE"],
+          name=pattern.name,
+          pos=i,
+          id=match_id,
+          is_input=True,
+      )
+
+      # Only replace input by the marker node for those nodes used in the
+      # pattern.
+      in_pattern_nodes = set(match.nodes_map.values())
+      for user in input_node.users.keys():
+        if user not in in_pattern_nodes:
+          continue
+
+        user.meta["ORIGINAL_NODE"].replace_input_with(
+            input_node.meta["ORIGINAL_NODE"], new_input_node
+        )
+        # Pattern matching graph sanitization may remove clone ops, which means
+        # the user's input in the original graph may be a clone op. When
+        # replacing the input with the marker node, we need to further try
+        # replacing the input of the clone op that connects to the user.
+        for original_user_input in user.meta["ORIGINAL_NODE"].all_input_nodes:
+          if fx_utils.is_clone_op(original_user_input):
+            original_user_input.replace_input_with(
+                input_node.meta["ORIGINAL_NODE"], new_input_node
+            )
+
+    for i, pattern_output_node in enumerate(pattern.output_nodes):
+      output_node = match.nodes_map[pattern_output_node]
+      new_output_node = _insert_marker(
+          graph_module,
+          output_node.meta["ORIGINAL_NODE"],
+          name=pattern.name,
+          pos=i,
+          id=match_id,
+          is_input=False,
+          attr=attr,  # torch_xla internal: only output marker needs attr.
+      )
+      output_node.meta["ORIGINAL_NODE"].replace_all_uses_with(new_output_node)
+      new_output_node.update_arg(0, output_node.meta["ORIGINAL_NODE"])
+
+  graph_module.graph.eliminate_dead_code()
+  _prepose_placeholder_nodes(graph_module.graph)
+
+  graph_module.graph.lint()
+  graph_module.recompile()
+  return graph_module

ai_edge_torch/hlfb/mark_pattern/fx_utils.py

@@ -0,0 +1,69 @@
+# 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.
+# ==============================================================================
+"""FX graph utilities for pattern matching clean ups."""
+
+import torch
+
+
+def is_clone_op(node: torch.fx.Node) -> bool:
+  """Checks if the node is a clone op."""
+  return (
+      node.op == "call_function" and node.target == torch.ops.aten.clone.default
+  )
+
+
+def remove_clone_ops(gm: torch.fx.GraphModule):
+  """Removes clone ops from the graph.
+
+  torch export adds additional aten.clone nodes to produce contiguous in memory
+  tensors depending on tensor sizes for runtime efficiency. However, these
+  unpredictable clone nodes can break the pattern matching. Thus remove all
+  clones in model and pattern graphs.
+
+  Args:
+    gm: The graph module to remove clone ops from.
+
+  Returns:
+    The graph module with clone ops removed.
+  """
+  for node in gm.graph.nodes:
+    if is_clone_op(node):
+      node.replace_all_uses_with(node.args[0])
+      gm.graph.erase_node(node)
+
+  gm.graph.lint()
+  gm.recompile()
+  return gm
+
+
+def remove_dangling_args(gm: torch.fx.GraphModule):
+  """Removes dangling args from the graph.
+
+  Args:
+    gm: The graph module to remove dangling args from.
+
+  Returns:
+    The graph module with dangling args removed.
+  """
+  nodes_to_erase = []
+  for node in gm.graph.nodes:
+    if node.op == "placeholder" and len(node.users) == 0:
+      nodes_to_erase.append(node)
+  for node in nodes_to_erase:
+    gm.graph.erase_node(node)
+
+  gm.graph.lint()
+  gm.recompile()
+  return gm